diff --git a/.github/scripts/agent-guides-drive.sh b/.github/scripts/agent-guides-drive.sh index 9b85b20177..f4189a159e 100755 --- a/.github/scripts/agent-guides-drive.sh +++ b/.github/scripts/agent-guides-drive.sh @@ -154,7 +154,12 @@ raw_env() { # $1 = var name -> value (one shlex-quote layer stripped) # writers as a side effect (it writes each agent's relocated session config). parse_connect() { local raw="$LOGS_DIR/connect-${AGENT}.txt" - if ! unsloth start "$AGENT" --no-launch --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then + # CONNECT_YOLO=1 adds --yolo. opencode/openclaw gate tool approval through their + # config (which now prompts by default), so the file-edit test opts into auto-approval + # here, the same intent as claude/codex's per-call bypass flags. + local yolo=() + [ -n "${CONNECT_YOLO:-}" ] && yolo=(--yolo) + if ! unsloth start "$AGENT" --no-launch "${yolo[@]}" --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then cat_redacted "$raw" guide_fail "'unsloth start ${AGENT} --no-launch' exited non-zero" fi @@ -371,7 +376,7 @@ case "$MODE" in hermes) patch_hermes_tools none invoke_via_connect "$OUT" -z "$PROMPT" ;; openclaw) patch_openclaw_agent notools - invoke_via_connect "$OUT" agent --local --agent ci \ + CONNECT_CMD_OVERRIDE=openclaw invoke_via_connect "$OUT" agent --local --agent ci \ --model "unsloth/${UNSLOTH_MODEL_ID}" --message "$PROMPT" ;; *) invoke_via_connect "$OUT" "$PROMPT" ;; esac @@ -394,7 +399,10 @@ case "$MODE" in T2='Run hello.py with python and show me the exact output.' # The start.py recipe writers + crosscheck must see the repo; run them - # from the repo root BEFORE cd-ing into the scratch work dir. + # from the repo root BEFORE cd-ing into the scratch work dir. opencode/openclaw + # gate tool approval through their config (prompting by default), so file-edit + # opts them into auto-approval to run edits/commands headlessly. + case "$AGENT" in opencode|openclaw) CONNECT_YOLO=1 ;; esac parse_connect crosscheck_contract # File-edit needs real tools, so we cannot zero them as in connection. @@ -441,7 +449,7 @@ case "$MODE" in fi ;; opencode) invoke_via_connect "$out" run "$prompt" ;; hermes) invoke_via_connect "$out" -z "$prompt" ;; - openclaw) invoke_via_connect "$out" agent --local --agent ci \ + openclaw) CONNECT_CMD_OVERRIDE=openclaw invoke_via_connect "$out" agent --local --agent ci \ --model "unsloth/${UNSLOTH_MODEL_ID}" --message "$prompt" ;; *) invoke_via_connect "$out" "$prompt" ;; esac @@ -519,6 +527,154 @@ case "$MODE" in echo "[claude] attribution A/B OK (suppressed HIT, header=1 MISS)" ;; + # ── resume: does a launched agent's session survive exit and resume? ──── + # Unlike the other modes, this drives the real LAUNCH path (`unsloth start + # ...`, the interactive default), not the --no-launch recipe. That + # path relocates each agent's home to a throwaway temp dir wiped on exit, so + # a session cannot be resumed -- unless --persist routes it to the stable + # Unsloth agents dir instead. We run one headless turn per pass and check + # whether the turn left a session in a persistent store (deterministic, no + # reliance on the model recalling anything), for a baseline pass and a + # --persist pass, and assert the expected split for this agent. + resume) + CODEWORD="PLATYPUS7" + T1="Remember this codeword for later: ${CODEWORD}. Reply with just the word OK." + T2="What codeword did I ask you to remember? Reply with just that word." + WORK="$WORKDIR_BASE/${AGENT}-resume" + + # STABLE_HOME: the stable dir that --no-launch (and --persist) relocate to. + # Read it from a --no-launch probe (which also writes the agent's config + # there). codex/pi relocate their whole home/HOME here; opencode/claude keep + # their session data in a fixed user dir, so STABLE_HOME stays empty for them. + parse_connect + case "$AGENT" in + codex) STABLE_HOME="$(raw_env CODEX_HOME)" ;; + pi) STABLE_HOME="$(raw_env HOME)" ;; + *) STABLE_HOME="" ;; + esac + + # The persistent stores a session would land in if it were NOT wiped. We + # count files here before/after each turn; a positive delta means the + # session persisted (is resumable), zero means it went to a wiped temp dir. + resume_tracked_dirs() { + case "$AGENT" in + codex) printf '%s\n' "$HOME/.codex" ;; + opencode) printf '%s\n' "$HOME/.local/share/opencode" "$HOME/.config/opencode" ;; + claude) printf '%s\n' "$HOME/.claude" ;; + pi) printf '%s\n' "$HOME/.pi" ;; + *) : ;; + esac + [ -n "$STABLE_HOME" ] && printf '%s\n' "$STABLE_HOME" + } + count_session_files() { + local total=0 d n + while IFS= read -r d; do + [ -n "$d" ] && [ -d "$d" ] || continue + n="$(find "$d" -type f 2>/dev/null | wc -l)"; total=$((total + n)) + done < <(resume_tracked_dirs) + echo "$total" + } + + # The headless first-turn subcommand per agent (mirrors file-edit's map), + # forwarded verbatim through the launch path as passthrough args. + set_t1_cmd() { + case "$AGENT" in + claude) T1_CMD=("${CLAUDE_CONNECT_FLAGS[@]}" -p "$T1") ;; + codex) T1_CMD=(exec "$T1") ;; + opencode) T1_CMD=(run "$T1") ;; + pi) T1_CMD=(-p "$T1") ;; + *) guide_fail "resume mode does not cover agent '$AGENT'" ;; + esac + } + + # Run one headless turn through the launch path. $1=outfile, $2="" or + # "--persist", rest = the agent subcommand. --yolo auto-approves so no tool + # prompt can hang; --api-key attaches to the already-served CI model. + launch_turn() { + local out="$1" rflag="$2"; shift 2 + local flag=(); [ -n "$rflag" ] && flag=("$rflag") + run_timed "$out" unsloth start "$AGENT" "${flag[@]}" --yolo \ + --api-key "$UNSLOTH_API_KEY" "$@" + local rc=$? + redact "$out" + return "$rc" + } + + # One pass: fresh work dir, one planting turn, set RESULT to PERSISTED/WIPED + # from the session-store delta. Runs in the main shell (not a command + # substitution) so a hang's guide_fail actually fails the job and the + # progress lines reach the CI log. $1 = "" (baseline) or "--persist". + RESULT="" + run_pass() { + local rflag="$1" label="baseline" + [ -n "$rflag" ] && label="resume" + rm -rf "$WORK"; mkdir -p "$WORK" + set_t1_cmd + local out="$LOGS_DIR/${AGENT}-resume-${label}.txt" + local before after rc + before="$(count_session_files)" + pushd "$WORK" >/dev/null || guide_fail "could not enter work dir $WORK" + launch_turn "$out" "$rflag" "${T1_CMD[@]}"; rc=$? + popd >/dev/null || true + after="$(count_session_files)" + echo "[$AGENT] ${label}: session files ${before} -> ${after} (rc=${rc})" + # The turn must succeed for the delta to mean anything: an agent that writes a + # session file then errors would otherwise be misread as PERSISTED. Mirror the + # file-edit mode and fail the pass on a non-zero launch (the flagship codex recall + # below stays WARN-only, driven by its own launch_turn calls). + [ "$rc" -eq 0 ] || { echo "[$AGENT] ${label} transcript (tail):"; tail -30 "$out" 2>/dev/null || true; \ + guide_fail "resume ${label} turn for ${AGENT} exited non-zero (rc=${rc})"; } + if [ "$after" -gt "$before" ]; then RESULT="PERSISTED"; else RESULT="WIPED"; fi + } + + run_pass ""; BASELINE="$RESULT" + # Only the temp-dir agents (codex/pi) need the --persist pass to prove the fix. + # opencode/claude persist either way, so the baseline already proves it and a + # second full CPU turn only risks a timeout; skip it for them. + case "$AGENT" in + codex|pi) run_pass "--persist"; RESUME="$RESULT" ;; + *) RESUME="n/a (persists either way)" ;; + esac + + # Expected: codex/pi relocate their whole home to the temp dir, so a plain + # launch is WIPED and only --persist PERSISTS. opencode/claude keep their + # session data in a fixed user dir, so the baseline already PERSISTS. + case "$AGENT" in + codex|pi) EXPECT_BASELINE="WIPED" ;; + opencode|claude) EXPECT_BASELINE="PERSISTED" ;; + esac + + echo "──────────────────────────────────────────────" + echo "[$AGENT] RESUME EXPERIMENT" + echo " baseline (unsloth start ${AGENT}): ${BASELINE} (expected ${EXPECT_BASELINE})" + echo " with --persist (unsloth start ${AGENT} --persist): ${RESUME}" + echo "──────────────────────────────────────────────" + + [ "$BASELINE" = "$EXPECT_BASELINE" ] \ + || guide_fail "baseline resume behavior for ${AGENT} was ${BASELINE}, expected ${EXPECT_BASELINE}" + case "$AGENT" in + codex|pi) + [ "$RESUME" = "PERSISTED" ] \ + || guide_fail "--persist did not persist ${AGENT}'s session (got ${RESUME}); the session dir is still not stable" ;; + esac + + # Flagship behavioral proof (codex only, WARN-only): after a --persist plant, + # resume the session and check the model actually recalls the codeword. A + # miss is not a failure (the CI model is small); the mechanism gate above is + # the real assertion. + if [ "$AGENT" = "codex" ]; then + rm -rf "$WORK"; mkdir -p "$WORK" + ( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-plant.txt" "--persist" exec "$T1" ) || true + ( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-recall.txt" "--persist" exec resume --last "$T2" ) || true + if grep -q "$CODEWORD" "$LOGS_DIR/codex-resume-recall.txt" 2>/dev/null; then + echo "[codex] behavioral recall HIT: resumed session remembered ${CODEWORD}" + else + echo "::warning::[codex] behavioral recall MISS (small CI model); mechanism gate still passed" + fi + fi + echo "[$AGENT] resume OK" + ;; + *) echo "agent-guides-drive.sh: unknown mode '$MODE'" >&2 exit 2 diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index 7978a200c0..1bb4c2bb58 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -364,6 +364,9 @@ jobs: tests/utils/test_attention_masks.py \ tests/utils/test_trunc_normal_patch.py \ tests/python/test_fast_language_model_text_only.py \ + tests/test_bad_mappings_redirect.py \ + tests/test_prefetch_snapshot_scope.py \ + tests/test_gemma_2b_mapper_key.py \ --deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap' # The deselected test monkeypatches flash_attn_varlen_func, which is # only bound on the module when `flash_attn` is importable. flash_attn diff --git a/.github/workflows/local-agent-guides-ci.yml b/.github/workflows/local-agent-guides-ci.yml index 47f75dc1ba..25796bd5cf 100644 --- a/.github/workflows/local-agent-guides-ci.yml +++ b/.github/workflows/local-agent-guides-ci.yml @@ -471,6 +471,176 @@ jobs: redacted-configs/ retention-days: 7 + # ═════════════════════════════════════════════════════════════════════ + # Job: resume + # Does a conversation started with `unsloth start ` survive exit + # and resume? This drives the REAL launch path (not the --no-launch + # recipe the other jobs use). A plain launch relocates the agent home to + # a temp dir wiped on exit, so codex/pi cannot resume; --persist routes the + # session to the stable Unsloth agents dir so it persists. opencode/claude + # keep their session data in a fixed user dir, so they persist either way. + # Dispatch-only: it is an end-to-end experiment, not a PR gate. + # ═════════════════════════════════════════════════════════════════════ + resume: + name: resume (${{ matrix.agent }}) + if: github.event_name == 'workflow_dispatch' + runs-on: ubuntu-latest + timeout-minutes: 60 + strategy: + fail-fast: false + matrix: + # codex/pi relocate their whole home (resume broken without --persist); + # opencode/claude keep session data in a fixed dir (resume already works). + # One agent from each class proves the split end to end; openclaw/hermes + # share codex's relocation mechanism and are covered by the unit tests. + agent: [codex, opencode, claude, pi] + env: + GGUF_REPO: unsloth/gemma-4-E4B-it-GGUF + GGUF_FILE: gemma-4-E4B-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18904' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux deps for llama.cpp prebuilt + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore GGUF model file + id: cache-gguf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Download GGUF if cache miss + id: download-gguf + if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p gguf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache + + - name: Save GGUF model file + if: always() && steps.download-gguf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Serve unsloth run --disable-tools (gemma-4-E4B) + run: | + unsloth studio reset-password + bash .github/scripts/serve-unsloth-run.sh \ + --gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \ + --port "$STUDIO_PORT" --log-dir logs \ + --extra "--seed $UNSLOTH_SEED --temp 0" \ + --health-timeout 900 + + - name: Preflight the agent's API dialect (class-a isolation) + env: + AGENT: ${{ matrix.agent }} + run: | + set -uo pipefail + B="$UNSLOTH_BASE_URL"; K="$UNSLOTH_API_KEY" + preflight_fail() { + echo "::error::[server/API regression] agent=$AGENT: $* (preflight failed BEFORE install/connect). Endpoint contract lives in studio/backend/routes/**."; + exit 1 + } + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/models" \ + -H "Authorization: Bearer $K") || true + [ "$code" = "200" ] || preflight_fail "/v1/models returned HTTP $code" + case "$AGENT" in + claude) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/messages" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/messages returned HTTP $code" + ;; + codex) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/responses" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"input\":\"Hi\",\"max_output_tokens\":16,\"stream\":true}") || true + [ "$code" = "200" ] || preflight_fail "/v1/responses returned HTTP $code" + ;; + *) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/chat/completions" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/chat/completions returned HTTP $code" + ;; + esac + echo "preflight OK for $AGENT" + + - name: Install agent CLI (class-b isolation) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-install.sh "$AGENT" + + - name: Resume experiment (launch path) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-drive.sh resume "$AGENT" + + - name: Collect server logs (debug) + if: always() + run: | + mkdir -p logs/studio-logs + cp -r "$HOME/.unsloth/studio/logs/." logs/studio-logs/ 2>/dev/null || true + if [ -n "${UNSLOTH_API_KEY:-}" ]; then + grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do + sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true + done + fi + + - name: Stop Studio + if: always() + run: | + if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then + kill "${UNSLOTH_SERVER_PID}" 2>/dev/null || true + fi + sleep 2 + ss -tln 2>/dev/null | grep ":${STUDIO_PORT}" || true + + - name: Upload logs + if: always() + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: resume-${{ matrix.agent }}-log + path: | + logs/ + agent-workdir/ + redacted-configs/ + retention-days: 7 + # ═════════════════════════════════════════════════════════════════════ # Job 3: prompt-cache # (a) curl 2-turn /v1/chat/completions: assert turn-2 cached_tokens > 0 diff --git a/.github/workflows/lockfile-audit.yml b/.github/workflows/lockfile-audit.yml index 9c28e21672..aaf258d615 100644 --- a/.github/workflows/lockfile-audit.yml +++ b/.github/workflows/lockfile-audit.yml @@ -60,11 +60,11 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 5 steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 with: persist-credentials: false - - uses: actions/setup-python@v5 + - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0 with: python-version: '3.12' diff --git a/.github/workflows/ossf.yml b/.github/workflows/ossf.yml new file mode 100644 index 0000000000..f9a270540f --- /dev/null +++ b/.github/workflows/ossf.yml @@ -0,0 +1,78 @@ +# This workflow uses actions that are not certified by GitHub. They are provided +# by a third-party and are governed by separate terms of service, privacy +# policy, and support documentation. + +name: Scorecard supply-chain security +on: + # For Branch-Protection check. Only the default branch is supported. See + # https://github.com/ossf/scorecard/blob/main/docs/checks.md#branch-protection + branch_protection_rule: + # To guarantee Maintained check is occasionally updated. See + # https://github.com/ossf/scorecard/blob/main/docs/checks.md#maintained + schedule: + - cron: '21 20 * * 0' + push: + branches: [ "main" ] + +# Declare default permissions as read only. +permissions: read-all + +jobs: + analysis: + name: Scorecard analysis + runs-on: ubuntu-latest + # `publish_results: true` only works when run from the default branch. conditional can be removed if disabled. + if: github.event.repository.default_branch == github.ref_name || github.event_name == 'pull_request' + permissions: + # Needed to upload the results to code-scanning dashboard. + security-events: write + # Needed to publish results and get a badge (see publish_results below). + id-token: write + # Uncomment the permissions below if installing in a private repository. + # contents: read + # actions: read + + steps: + - name: "Checkout code" + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + persist-credentials: false + + - name: "Run analysis" + uses: ossf/scorecard-action@f49aabe0b5af0936a0987cfb85d86b75731b0186 # v2.4.1 + with: + results_file: results.sarif + results_format: sarif + # (Optional) "write" PAT token. Uncomment the `repo_token` line below if: + # - you want to enable the Branch-Protection check on a *public* repository, or + # - you are installing Scorecard on a *private* repository + # To create the PAT, follow the steps in https://github.com/ossf/scorecard-action?tab=readme-ov-file#authentication-with-fine-grained-pat-optional. + # repo_token: ${{ secrets.SCORECARD_TOKEN }} + + # Public repositories: + # - Publish results to OpenSSF REST API for easy access by consumers + # - Allows the repository to include the Scorecard badge. + # - See https://github.com/ossf/scorecard-action#publishing-results. + # For private repositories: + # - `publish_results` will always be set to `false`, regardless + # of the value entered here. + publish_results: true + + # (Optional) Uncomment file_mode if you have a .gitattributes with files marked export-ignore + # file_mode: git + + # Upload the results as artifacts (optional). Commenting out will disable uploads of run results in SARIF + # format to the repository Actions tab. + - name: "Upload artifact" + uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1 + with: + name: SARIF file + path: results.sarif + retention-days: 5 + + # Upload the results to GitHub's code scanning dashboard (optional). + # Commenting out will disable upload of results to your repo's Code Scanning dashboard + - name: "Upload to code-scanning" + uses: github/codeql-action/upload-sarif@v3 + with: + sarif_file: results.sarif diff --git a/.github/workflows/security-audit.yml b/.github/workflows/security-audit.yml index 0ef2ad1e9d..1275d12216 100644 --- a/.github/workflows/security-audit.yml +++ b/.github/workflows/security-audit.yml @@ -2,8 +2,8 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. # Multi-language supply-chain audit. Triggers: -# - PRs touching any dependency manifest (Python / npm / Cargo) or -# this workflow file, +# - PRs touching any dependency manifest (Python / npm / Cargo), a +# scanner or its allowlist baseline, or this workflow file, # - push to main / pip, # - nightly @ 04:13 UTC so newly-published advisories surface even # when no PR opens, @@ -57,7 +57,9 @@ on: - 'studio/src-tauri/Cargo.lock' - 'pyproject.toml' - 'scripts/scan_packages.py' + - 'scripts/scan_packages_baseline.json' - 'scripts/scan_npm_packages.py' + - 'scripts/scan_npm_packages_baseline.json' - '.github/workflows/security-audit.yml' push: branches: [main, pip] diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml index aebf90380a..f540c11da4 100644 --- a/.github/workflows/studio-inference-smoke.yml +++ b/.github/workflows/studio-inference-smoke.yml @@ -444,6 +444,8 @@ jobs: python - <<'PY' import json import os + import time + import urllib.error import urllib.request BASE = os.environ["BASE_URL"] @@ -464,8 +466,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) def post_sse(path, body, *, timeout = 600): """POST a streaming request and accumulate the assistant @@ -938,6 +956,8 @@ jobs: import base64 import json import os + import time + import urllib.error import urllib.request from openai import OpenAI from anthropic import Anthropic @@ -956,8 +976,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) # ── 1. response_format = json_object (JSON mode) ───────────── # llama.cpp's HTTP server supports OpenAI-compatible JSON diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml index d562294d42..03c0a8580d 100644 --- a/.github/workflows/studio-mac-inference-smoke.yml +++ b/.github/workflows/studio-mac-inference-smoke.yml @@ -430,6 +430,8 @@ jobs: python - <<'PY' import json import os + import time + import urllib.error import urllib.request BASE = os.environ["BASE_URL"] @@ -450,8 +452,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) def post_sse(path, body, *, timeout = 600): """POST a streaming request and accumulate the assistant @@ -825,6 +843,8 @@ jobs: import base64 import json import os + import time + import urllib.error import urllib.request from openai import OpenAI from anthropic import Anthropic @@ -848,8 +868,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) # ── 1. response_format = json_object (JSON mode) ───────────── # llama.cpp's HTTP server supports OpenAI-compatible JSON diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index 8186c07211..0453c9212a 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -634,6 +634,8 @@ jobs: python - <<'PY' import json import os + import time + import urllib.error import urllib.request BASE = os.environ["BASE_URL"] @@ -656,8 +658,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) def post_sse(path, body, *, timeout = 600): body = {**body, "stream": True} @@ -1063,6 +1081,8 @@ jobs: import base64 import json import os + import time + import urllib.error import urllib.request from openai import OpenAI from anthropic import Anthropic @@ -1082,8 +1102,24 @@ jobs: "Content-Type": "application/json", }, ) - with urllib.request.urlopen(req, timeout = timeout) as resp: - return resp.status, json.loads(resp.read().decode()) + # Shared CI runners stall sporadically, so retry transport-level + # failures only; HTTP status errors surface immediately. Bounded + # to fit the job's timeout-minutes: short probes get 3 full + # attempts, long probes one retry capped at 300s (a healthy + # server answers a retry quickly; a stalled one never does). + attempts = 3 if timeout <= 300 else 2 + for attempt in range(attempts): + try: + t = timeout if attempt == 0 else min(timeout, 300) + with urllib.request.urlopen(req, timeout = t) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError: + raise + except (TimeoutError, ConnectionError, urllib.error.URLError) as exc: + if attempt == attempts - 1: + raise + print(f"[retry] {path}: {exc!r}", flush = True) + time.sleep(15) # ── 1. response_format = json_object (JSON mode) ───────────── status, data = post("/v1/chat/completions", { @@ -1334,42 +1370,75 @@ jobs: try { Add-MpPreference -ExclusionPath $p -ErrorAction Stop } catch { } } - - name: Hide Visual Studio + CMake (simulate a host with no build tools) + - name: Prepare no-build-tools simulation shell: pwsh run: | $ErrorActionPreference = 'Stop' - # A Program Files dir can hold a transient handle (Defender / MSBuild node) - # so Rename-Item intermittently fails with "Access is denied"; retry to ride it out. - function Rename-WithRetry($Path, $NewName) { - for ($i = 1; $i -le 6; $i++) { - try { Rename-Item -LiteralPath $Path -NewName $NewName -ErrorAction Stop; return } - catch { if ($i -eq 6) { throw }; Start-Sleep -Seconds 3 } + $root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools' + $pf = Join-Path $root 'ProgramFiles' + $pfx86 = Join-Path $root 'ProgramFilesx86' + New-Item -ItemType Directory -Force -Path $pf, $pfx86 | Out-Null + + $blocked = [System.Collections.Generic.HashSet[string]]::new([StringComparer]::OrdinalIgnoreCase) + foreach ($tool in @('cmake', 'cl.exe')) { + foreach ($cmd in (Get-Command $tool -All -ErrorAction SilentlyContinue)) { + if ($cmd.Source) { + $dir = Split-Path -Parent $cmd.Source + if ($dir) { + [void] $blocked.Add( + [Environment]::ExpandEnvironmentVariables($dir).Trim().Trim('"').TrimEnd('\')) + } + } } } - # Rename the Visual Studio install roots (incl. the Installer that holds - # vswhere.exe) so Find-VsBuildTools' vswhere + filesystem scan both miss. - foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) { - if (Test-Path -LiteralPath $d) { - Rename-WithRetry $d ((Split-Path $d -Leaf) + '.vsoff') - Write-Host "Hid VS: $d" - } + # Normalized comparison so registry spellings (trailing slash, + # unexpanded %VAR%) still match. + function Test-Blocked([string]$p) { + $n = [Environment]::ExpandEnvironmentVariables($p).Trim().Trim('"').TrimEnd('\') + return $blocked.Contains($n) } - # Surgically rename each cmake executable on PATH (not its parent dir -- - # cmake can share a dir with other shims) so Get-Command cmake fails. - $hidden = @() - foreach ($c in (Get-Command cmake -All -ErrorAction SilentlyContinue)) { - if ($c.Source -and (Test-Path -LiteralPath $c.Source)) { - Rename-WithRetry $c.Source ((Split-Path $c.Source -Leaf) + '.off') - $hidden += $c.Source - Write-Host "Hid cmake: $($c.Source)" - } + + $pathParts = $env:Path -split [IO.Path]::PathSeparator | + Where-Object { $_ -and -not (Test-Blocked $_) } + $noBuildToolsPath = $pathParts -join [IO.Path]::PathSeparator + + # install.ps1's Refresh-SessionPath and setup.ps1's Refresh-Environment + # rebuild the session Path from these scopes mid-install, so filter + # them too. Originals are saved for the cleanup step. + foreach ($scope in @('Machine', 'User')) { + $orig = [Environment]::GetEnvironmentVariable('Path', $scope) + if (-not $orig) { continue } + Set-Content -LiteralPath (Join-Path $root "orig-path-$scope.txt") -Value $orig -NoNewline + $kept = ($orig -split ';' | Where-Object { $_ -and -not (Test-Blocked $_) }) -join ';' + [Environment]::SetEnvironmentVariable('Path', $kept, $scope) + Write-Host "Filtered $scope Path scope." + } + + "NO_BUILD_TOOLS_PROGRAMFILES=$pf" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PROGRAMFILES_X86=$pfx86" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PATH<&1 | Tee-Object -FilePath logs/install.log @@ -1480,19 +1553,19 @@ jobs: [ -n "$CONTENT" ] && [ "$CONTENT" != "null" ] || { echo "::error::empty completion"; exit 1; } echo "Inference OK without Visual Studio: $CONTENT" - - name: Restore Visual Studio + CMake + - name: Clean no-build-tools simulation if: always() shell: pwsh run: | - foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) { - $off = "$d.vsoff" - if (Test-Path -LiteralPath $off) { Rename-Item -LiteralPath $off -NewName (Split-Path $d -Leaf); Write-Host "Restored $d" } - } - if ($env:HIDDEN_CMAKE) { - foreach ($src in ($env:HIDDEN_CMAKE -split '\|')) { - if ($src -and (Test-Path -LiteralPath "$src.off")) { Rename-Item -LiteralPath "$src.off" -NewName (Split-Path $src -Leaf) } + $root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools' + foreach ($scope in @('Machine', 'User')) { + $saved = Join-Path $root "orig-path-$scope.txt" + if (Test-Path -LiteralPath $saved) { + [Environment]::SetEnvironmentVariable('Path', (Get-Content -LiteralPath $saved -Raw), $scope) + Write-Host "Restored $scope Path scope." } } + Remove-Item -LiteralPath $root -Recurse -Force -ErrorAction SilentlyContinue - name: Stop Studio if: always() @@ -1540,21 +1613,34 @@ jobs: with: python-version: '3.12' - - name: Hide Visual Studio + - name: Prepare no-build-tools simulation shell: pwsh run: | $ErrorActionPreference = 'Stop' - # Retry the rename: a Program Files dir can hold a transient handle that - # makes Rename-Item intermittently fail with "Access is denied". - function Rename-WithRetry($Path, $NewName) { - for ($i = 1; $i -le 6; $i++) { - try { Rename-Item -LiteralPath $Path -NewName $NewName -ErrorAction Stop; return } - catch { if ($i -eq 6) { throw }; Start-Sleep -Seconds 3 } + $root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools' + $pf = Join-Path $root 'ProgramFiles' + $pfx86 = Join-Path $root 'ProgramFilesx86' + New-Item -ItemType Directory -Force -Path $pf, $pfx86 | Out-Null + + $blocked = [System.Collections.Generic.HashSet[string]]::new([StringComparer]::OrdinalIgnoreCase) + foreach ($tool in @('cmake', 'cl.exe')) { + foreach ($cmd in (Get-Command $tool -All -ErrorAction SilentlyContinue)) { + if ($cmd.Source) { + $dir = Split-Path -Parent $cmd.Source + if ($dir) { [void] $blocked.Add($dir) } + } } } - foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) { - if (Test-Path -LiteralPath $d) { Rename-WithRetry $d ((Split-Path $d -Leaf) + '.vsoff'); Write-Host "Hid VS: $d" } - } + + $pathParts = $env:Path -split [IO.Path]::PathSeparator | + Where-Object { $_ -and -not $blocked.Contains($_) } + $noBuildToolsPath = $pathParts -join [IO.Path]::PathSeparator + + "NO_BUILD_TOOLS_PROGRAMFILES=$pf" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PROGRAMFILES_X86=$pfx86" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PATH< /tmp/resolve.json || { - echo "::error::resolver exited non-zero"; cat /tmp/resolve.json || true; exit 1; } - cat /tmp/resolve.json - echo "Prebuilt resolver ran with no Visual Studio present." + if ($LASTEXITCODE -ne 0) { Write-Host "::error::pip install huggingface_hub failed"; exit 1 } + python studio/install_llama_prebuilt.py --resolve-prebuilt latest --output-format json > resolve.json + if ($LASTEXITCODE -ne 0) { + Write-Host "::error::resolver exited non-zero" + if (Test-Path resolve.json) { Get-Content resolve.json } + exit 1 + } + Get-Content resolve.json + Write-Host "Prebuilt resolver ran with no Visual Studio present." - - name: Restore Visual Studio + - name: Clean no-build-tools simulation if: always() shell: pwsh run: | - foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) { - $off = "$d.vsoff" - if (Test-Path -LiteralPath $off) { Rename-Item -LiteralPath $off -NewName (Split-Path $d -Leaf); Write-Host "Restored $d" } - } + Remove-Item -LiteralPath (Join-Path $env:GITHUB_WORKSPACE 'no-build-tools') -Recurse -Force -ErrorAction SilentlyContinue # ── folded from studio-setup-ps1-vs2026.yml: setup.ps1 unit tests + real-VS detection + vcredist ── pester: @@ -1610,6 +1705,13 @@ jobs: - name: Install Pester v5 shell: pwsh run: | + # PSGallery is intermittently absent from the repository list on GitHub's Windows + # runners, which makes `Set-PSRepository PSGallery` fail with "No repository with the + # name 'PSGallery' was found." Re-register the default gallery first so the policy + # change and module install below always have a repository to target. + if (-not (Get-PSRepository -Name PSGallery -ErrorAction SilentlyContinue)) { + Register-PSRepository -Default -ErrorAction SilentlyContinue + } Set-PSRepository PSGallery -InstallationPolicy Trusted Install-Module Pester -MinimumVersion 5.5.0 -Force -SkipPublisherCheck -Scope CurrentUser Import-Module Pester -MinimumVersion 5.5.0 diff --git a/.github/workflows/studio-windows-update-smoke.yml b/.github/workflows/studio-windows-update-smoke.yml index 888b3d70a3..5b92f1a3e0 100644 --- a/.github/workflows/studio-windows-update-smoke.yml +++ b/.github/workflows/studio-windows-update-smoke.yml @@ -6,9 +6,9 @@ # windows-latest runner: # # 1. install.ps1 --local --no-torch installs Studio AND auto-fetches -# the prebuilt llama.cpp Windows binary (llama-bNNNN-bin-win-cpu- -# x64 from ggml-org/llama.cpp). Hitting the source-build fallback -# is treated as an Unsloth bug -- Studio must always pick the +# the prebuilt llama.cpp Windows binary (app--windows-x64-cpu +# from unslothai/llama.cpp). Hitting the source-build fallback is +# treated as an Unsloth bug -- Studio must always pick the # prebuilt on Windows. # 2. unsloth studio update --local is idempotent. Two consecutive # runs both report "prebuilt up to date and validated", no diff --git a/.github/workflows/version-compat-ci.yml b/.github/workflows/version-compat-ci.yml index e492d21e99..6becccc90a 100644 --- a/.github/workflows/version-compat-ci.yml +++ b/.github/workflows/version-compat-ci.yml @@ -285,6 +285,92 @@ jobs: tests/vllm_compat/test_extended_module_imports.py \ -v --tb=short + # Fake-CUDA GRPO/SFT/DPO patch run against REAL TRL (latest + main). Unlike + # the static symbol/source greps above, this drives unsloth's actual + # source-transform patchers (models/rl.py + rl_replacements.py) on a CPU-only + # runner under the tests/conftest.py spoof harness -- no GPU, no training. + # Catches structural TRL drift the greps miss (e.g. TRL 1.7.0's 2->3-tuple + # per-token-logps return, restructured PEFT ref-adapter block) by asserting + # the generated Unsloth trainer still satisfies the transform contracts. + grpo-fake-run: + name: GRPO fake-run (latest + main TRL, CPU spoof) + runs-on: ubuntu-latest + timeout-minutes: 18 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + path: unsloth + - name: Clone unsloth-zoo @ main + run: | + for attempt in 1 2 3; do + rm -rf "$RUNNER_TEMP/unsloth-zoo" + if git clone --depth=1 https://github.com/unslothai/unsloth-zoo \ + "$RUNNER_TEMP/unsloth-zoo"; then + break + fi + if [ "$attempt" -eq 3 ]; then + echo "::error::git clone unsloth-zoo failed after 3 attempts" + exit 1 + fi + delay=$((5 * attempt)) + echo "::warning::clone failed (attempt $attempt/3), retrying in ${delay}s..." + sleep "$delay" + done + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - name: Install CPU torch + ecosystem + TRL latest + run: | + python -m pip install --upgrade pip + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \ + 'torch>=2.4,<2.11' 'torchvision<0.26' 'torchcodec<0.10' + # Ecosystem floors unsloth needs; TRL itself is installed last so it + # can pull the transformers/peft it requires. + pip install \ + 'transformers>=4.57' 'peft>=0.18.0' 'accelerate>=1.0' 'datasets>=3.4,<5' \ + 'bitsandbytes>=0.45.5' sentencepiece protobuf safetensors numpy 'pytest>=8' \ + 'huggingface_hub>=0.34' tqdm packaging psutil triton Pillow + pip install --upgrade trl + pip install --no-deps -e "$RUNNER_TEMP/unsloth-zoo" + pip install --no-deps -e ./unsloth + - name: Fake-run vs TRL latest + env: + UNSLOTH_IS_PRESENT: '1' + UNSLOTH_COMPILE_DISABLE: '1' + # Disable dynamo/inductor at the process level, before conftest.py's early + # `import unsloth`, so the GRPO hot path never compiles on the GPU-less runner + # (defense in depth; the CPU fake-train also flips this at runtime). + TORCHDYNAMO_DISABLE: '1' + TORCH_COMPILE_DISABLE: '1' + PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python + run: | + cd unsloth + python -c "import trl; print('Resolved TRL', trl.__version__)" + PYTHONPATH=. python -m pytest \ + tests/version_compat/test_trl_grpo_fake_run.py \ + tests/version_compat/test_trl_fake_train_cpu.py \ + -v --tb=short + # `main` is scheduled/dispatch-only so PR jobs stay fast and a bleeding-edge + # TRL break does not red every PR. github.event_name is valid in a step if. + - name: Fake-run vs TRL main (scheduled / dispatch only) + if: ${{ github.event_name != 'pull_request' }} + env: + UNSLOTH_IS_PRESENT: '1' + UNSLOTH_COMPILE_DISABLE: '1' + TORCHDYNAMO_DISABLE: '1' + TORCH_COMPILE_DISABLE: '1' + PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python + run: | + pip install --upgrade "git+https://github.com/huggingface/trl" + cd unsloth + python -c "import trl; print('Resolved TRL', trl.__version__)" + PYTHONPATH=. python -m pytest \ + tests/version_compat/test_trl_grpo_fake_run.py \ + tests/version_compat/test_trl_fake_train_cpu.py \ + -v --tb=short + # Daily-only: same suites but with --strict on importable upstream # tags. Schedule-only so PR jobs stay fast; cron tolerates a flake. daily-fresh-fetch: diff --git a/.gitignore b/.gitignore index 9f7d4b8c60..39ca2226ca 100644 --- a/.gitignore +++ b/.gitignore @@ -11,6 +11,8 @@ outputs/ exports/ /datasets/ studio/backend/assets/datasets/ +# Generated async worker / reviewer transcripts (never part of the product). +studio/backend/async_task_outputs/ unsloth_training_checkpoints/ *.gguf *.safetensors diff --git a/README.md b/README.md index e3fd4e6980..849ee2e87b 100644 --- a/README.md +++ b/README.md @@ -212,7 +212,7 @@ By default `unsloth studio` binds to `127.0.0.1` (this machine only). To reach i ```bash unsloth studio --secure -p 8888 ``` -- `-H 0.0.0.0`: bind the raw port on all network interfaces, reachable from anywhere on the network. Only use this on a trusted network. +- `-H 0.0.0.0`: bind the raw port on all network interfaces, reachable from anywhere on the network. This also starts a public Cloudflare quick tunnel by default, which publishes an internet-reachable `https://*.trycloudflare.com` URL even behind a firewall. Both the raw port and the tunnel expose Studio beyond this machine, so only use this on a network you trust; pass `--no-cloudflare` to drop the public link while keeping the network bind. ```bash unsloth studio -H 0.0.0.0 -p 8888 ``` diff --git a/install.ps1 b/install.ps1 index 3bf9317960..4618029111 100644 --- a/install.ps1 +++ b/install.ps1 @@ -473,6 +473,17 @@ function Install-UnslothStudio { param( [Parameter(Mandatory = $true)][ScriptBlock]$Command ) + # Installer-pinned index installs (torch) must beat an inherited uv mirror + # (#6898): when the command pins an index, clear every uv index env var so + # it wins, then restore in finally. Other installs keep the user's mirror. + $savedUvIndex = $null + if ($Command.ToString() -match '--default-index') { + $savedUvIndex = @{} + foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') { + $savedUvIndex[$n] = [Environment]::GetEnvironmentVariable($n) + Remove-Item "Env:$n" -ErrorAction SilentlyContinue + } + } $prevEap = $ErrorActionPreference $ErrorActionPreference = "Continue" try { @@ -492,6 +503,7 @@ function Install-UnslothStudio { return [int]$LASTEXITCODE } finally { $ErrorActionPreference = $prevEap + if ($savedUvIndex) { foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } } } } } @@ -2254,7 +2266,7 @@ exit 0 if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below @@ -2268,7 +2280,7 @@ exit 0 } } } else { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -2299,7 +2311,7 @@ exit 0 # ABI-incompatible torchvision/torchaudio on AMD's per-arch index. $visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } $audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec $visionSpec $audioSpec } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $torchSpec $visionSpec $audioSpec } if ($torchInstallExit -ne 0) { # Transient AMD-index failure: fall back to a CPU base so the install # still completes; Studio setup retries ROCm afterwards. @@ -2312,7 +2324,7 @@ exit 0 # torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU # torch>= range, so without it uv would keep the ROCm build and only swap # the companions -- a mismatched venv the flavor-repair block won't fix. - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $CpuFallbackIndexUrl } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $CpuFallbackIndexUrl } if ($torchInstallExit -ne 0) { Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) @@ -2329,7 +2341,7 @@ exit 0 } else { Write-TauriLog "STEP" "Installing PyTorch" substep "installing PyTorch ($TorchIndexUrl)..." - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl } if ($torchInstallExit -ne 0) { Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) @@ -2341,7 +2353,7 @@ exit 0 if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic } @@ -2353,7 +2365,7 @@ exit 0 } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } } else { $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -2381,7 +2393,7 @@ exit 0 Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.6.7" "unsloth>=2026.6.9" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.2" "unsloth>=2026.7.2" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) @@ -2412,7 +2424,7 @@ exit 0 # keeps a stale torch==X+cpu against a CUDA index and setup.ps1 then loops on # "torch cpu != required cuXXX". Reinstall the right triplet when a GPU build is # expected: CUDA from $TorchIndexUrl, ROCm from $ROCmIndexUrl (repo.amd.com gfx* - # is a PEP 503 index uv resolves via --index-url, same URL the fresh ROCm install + # is a PEP 503 index uv resolves via --default-index, same URL the fresh ROCm install # above uses). --no-torch / CPU-only hosts (expected cpu) are no-ops. if (-not $SkipTorch) { $expectedTorchTag = Get-ExpectedTorchFlavorTag -TorchIndexUrl $TorchIndexUrl -ROCmIndexUrl $ROCmIndexUrl @@ -2428,7 +2440,7 @@ exit 0 $visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } $audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow" - $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec } + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec } if ($torchFixExit -ne 0) { Write-Host "[ERROR] Failed to reinstall PyTorch with the correct ROCm build (exit code $torchFixExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit) @@ -2440,7 +2452,7 @@ exit 0 } elseif ($expectedTorchTag -ne 'rocm') { # CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet. substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow" - $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } if ($torchFixExit -ne 0) { Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit) diff --git a/install.sh b/install.sh index c16648aa73..fc571e4917 100755 --- a/install.sh +++ b/install.sh @@ -159,6 +159,12 @@ run_maybe_quiet() { run_install_cmd() { _label="$1" shift + # Installer-pinned index installs (torch) must beat an inherited uv mirror + # (#6898): when we pass --default-index, neutralize every uv index env var so + # the pinned index wins. Other installs keep the user's mirror. + case " $* " in + *" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL "$@" ;; + esac if _is_verbose; then "$@" && return 0 _rc=$? @@ -1446,8 +1452,14 @@ if [ "$_NO_TORCH_FLAG" = true ] || [ "$MAC_INTEL" = true ]; then SKIP_TORCH=true fi +# Apple Silicon: exclude broken mlx-lm 0.31.3 (QK-norm load regression for +# gemma4 / qwen3_5; mlx-lm #1242). A curl-piped install has no overrides file +# and skips the guarded MLX step (SKIP_STUDIO_BASE=1), so this is the only cover. +_MLX_LM_EXCLUDE_ARG="" + # Apple Silicon: override mlx-vlm / mlx-lm's transformers pin (see overrides file). if [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then + _MLX_LM_EXCLUDE_ARG="mlx-lm!=0.31.3" _OVERRIDES_FILE="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)/studio/backend/requirements/single-env/overrides-darwin-arm64.txt" if [ -f "$_OVERRIDES_FILE" ]; then # uv splits UV_OVERRIDE on whitespace, so a repo path with whitespace @@ -1481,6 +1493,81 @@ elif [ "$OS" = "macos" ]; then fi tauri_diag_marker "$_TAURI_INITIAL_GPU_BRANCH" "none" +# AMD GPU name from the Windows host via WMI, or empty. Discrete cards aren't in +# /proc/cpuinfo, so ask Windows. Cached ("-" = negative), self-contained, bounded +# to 10s. Defined here so the reroute below can use it before _run_bounded exists. +_WSL_AMD_GPU_NAME_CACHE="" +_wsl_amd_gpu_name() { + if [ -n "$_WSL_AMD_GPU_NAME_CACHE" ]; then + [ "$_WSL_AMD_GPU_NAME_CACHE" = "-" ] && return 1 + printf '%s' "$_WSL_AMD_GPU_NAME_CACHE"; return 0 + fi + command -v powershell.exe >/dev/null 2>&1 || { _WSL_AMD_GPU_NAME_CACHE="-"; return 1; } + _wag_ps="(Get-CimInstance Win32_VideoController | Where-Object { \$_.Name -match 'AMD|Radeon' } | Select-Object -First 1).Name" + if command -v timeout >/dev/null 2>&1; then + _wag_n="$(timeout 10 powershell.exe -NoProfile -Command "$_wag_ps" 2>/dev/null | tr -d '\r\n\000')" + else + _wag_n="$(powershell.exe -NoProfile -Command "$_wag_ps" 2>/dev/null | tr -d '\r\n\000')" + fi + if [ -n "$_wag_n" ]; then _WSL_AMD_GPU_NAME_CACHE="$_wag_n"; printf '%s' "$_wag_n"; return 0; fi + _WSL_AMD_GPU_NAME_CACHE="-"; return 1 +} + +# ── Bounded command runner ── +# Runs a command under a 10s timeout when the `timeout` binary is available, +# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during +# driver init or after a reset) from hanging the installer: a timed-out probe +# exits nonzero and is treated exactly like a failed probe. No-op semantics on +# hosts without `timeout` (e.g. macOS) or when the probe is healthy. +_run_bounded() { + if command -v timeout >/dev/null 2>&1; then + timeout 10 "$@" + else + "$@" + fi +} + +# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every +# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to +# the AMD card). Unset means all devices visible. nvidia-smi ignores this env +# var, so the probes below cannot see the distinction on their own. +_cvd_hides_nvidia() { + [ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1 + _cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]') + [ -z "$_cvd_trim" ] || [ "$_cvd_trim" = "-1" ] +} + +# ── NVIDIA usable-GPU helper ── +# Returns 0 (true) if an NVIDIA GPU is present and usable. +# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs, +# which the NVIDIA driver populates on Linux regardless of nvidia-smi state +# -- handles PATH gaps, subprocess timeouts, and driver init races that +# could otherwise cause nvidia-smi to fail and silence NVIDIA detection. +# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches +# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs. +_has_usable_nvidia_gpu() { + if _cvd_hides_nvidia; then + return 1 + fi + _nvsmi="" + if command -v nvidia-smi >/dev/null 2>&1; then + _nvsmi="nvidia-smi" + elif [ -x "/usr/bin/nvidia-smi" ]; then + _nvsmi="/usr/bin/nvidia-smi" + fi + if [ -n "$_nvsmi" ]; then + if _run_bounded "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then + return 0 + fi + fi + # Fallback: NVIDIA driver exposes one subdir per GPU under this path. + if [ -d /proc/driver/nvidia/gpus ] && \ + [ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then + return 0 + fi + return 1 +} + # Strix Halo ROCm-on-WSL only targets Ubuntu 24.04. On a newer distro (e.g. 26.04) # with a 24.04 distro present, re-run the install there and stop; else fall through # to CPU + the `wsl --install` hint below (never auto-create a distro). Runs before @@ -1491,7 +1578,15 @@ _maybe_reroute_strixhalo_to_2404() { [ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0 [ "${UNSLOTH_WSL_REROUTED:-0}" = "1" ] && return 0 [ -e /dev/dxg ] || return 0 - grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null || return 0 + # A usable NVIDIA GPU (common on hybrid AMD+NVIDIA hosts) means the CUDA path works on + # this distro, so don't reroute for AMD. _has_usable_nvidia_gpu (moved above) honors + # CUDA_VISIBLE_DEVICES=""/-1 and the /proc/driver/nvidia fallback for PATH/timeout gaps. + if _has_usable_nvidia_gpu; then return 0; fi + # Strix APUs show in /proc/cpuinfo; discrete cards don't, so also try WMI. Either reroutes. + if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \ + && ! _wsl_amd_gpu_name >/dev/null 2>&1; then + return 0 + fi # Already ROCm-on-WSL? leave a working GPU alone, whatever the version. if [ -e /opt/rocm/lib/librocdxg.so ] || [ -e /opt/rocm/lib64/librocdxg.so ]; then return 0 @@ -1965,61 +2060,6 @@ _has_amd_rocm_gpu() { return 1 } -# ── Bounded command runner ── -# Runs a command under a 10s timeout when the `timeout` binary is available, -# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during -# driver init or after a reset) from hanging the installer: a timed-out probe -# exits nonzero and is treated exactly like a failed probe. No-op semantics on -# hosts without `timeout` (e.g. macOS) or when the probe is healthy. -_run_bounded() { - if command -v timeout >/dev/null 2>&1; then - timeout 10 "$@" - else - "$@" - fi -} - -# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every -# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to -# the AMD card). Unset means all devices visible. nvidia-smi ignores this env -# var, so the probes below cannot see the distinction on their own. -_cvd_hides_nvidia() { - [ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1 - _cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]') - [ -z "$_cvd_trim" ] || [ "$_cvd_trim" = "-1" ] -} - -# ── NVIDIA usable-GPU helper ── -# Returns 0 (true) if an NVIDIA GPU is present and usable. -# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs, -# which the NVIDIA driver populates on Linux regardless of nvidia-smi state -# -- handles PATH gaps, subprocess timeouts, and driver init races that -# could otherwise cause nvidia-smi to fail and silence NVIDIA detection. -# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches -# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs. -_has_usable_nvidia_gpu() { - if _cvd_hides_nvidia; then - return 1 - fi - _nvsmi="" - if command -v nvidia-smi >/dev/null 2>&1; then - _nvsmi="nvidia-smi" - elif [ -x "/usr/bin/nvidia-smi" ]; then - _nvsmi="/usr/bin/nvidia-smi" - fi - if [ -n "$_nvsmi" ]; then - if _run_bounded "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then - return 0 - fi - fi - # Fallback: NVIDIA driver exposes one subdir per GPU under this path. - if [ -d /proc/driver/nvidia/gpus ] && \ - [ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then - return 0 - fi - return 1 -} - # ── Detect GPU and choose PyTorch index URL ── # Mirrors Get-TorchIndexUrl in install.ps1. # On CPU-only machines this returns the cpu index, avoiding the solver @@ -2199,9 +2239,9 @@ _expected_torch_flavor_tag() { esac } -# Whether index ($1) supports a plain --index-url reinstall. pytorch.org cuXXX / +# Whether index ($1) supports a plain --default-index reinstall. pytorch.org cuXXX / # rocmX.Y AND the repo.amd.com gfx* indexes are all PEP 503 simple indexes that uv -# resolves (torch + every transitive dep) via --index-url -- the same URLs the +# resolves (torch + every transitive dep) via --default-index -- the same URLs the # fresh-install paths above already use -- so a stale wheel is auto-repairable. # Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall. # Lowercase the leaf first so a cased leaf (e.g. gfx120X-all) is recognised. @@ -2432,19 +2472,19 @@ _persist_rocm_wsl_dropin() { fi } +# _wsl_amd_gpu_name is defined earlier so both the reroute and this bootstrap can use it. _maybe_bootstrap_rocm_wsl() { [ "${OS:-}" = "wsl" ] || return 0 [ "${SKIP_TORCH:-false}" = "false" ] || return 0 [ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0 # Leave any already-usable GPU completely alone (NVIDIA, or working ROCm). if _has_usable_nvidia_gpu; then return 0; fi - # "Usable ROCm" here = rocminfo enumerates the gfx1151 agent. Don't use the - # generic _has_amd_rocm_gpu: its broad gfx match accepts "gfx11-generic" and - # would skip this bootstrap while the real GPU is still unusable. awk consumes - # all input, so rocminfo isn't SIGPIPE'd like `grep -q` would under pipefail. + # Usable ROCm = rocminfo enumerates a real GPU agent: gfx[1-9] (excludes gfx000, + # the CPU agent) and not the "gfx11-generic" fallback. awk consumes all input so + # rocminfo isn't SIGPIPE'd like `grep -q` under pipefail. _ensure_rocm_probe_env if command -v rocminfo >/dev/null 2>&1 && \ - rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx1151/{found=1} END{exit !found}'; then + rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9]/ && !/generic/{found=1} END{exit !found}'; then # rocminfo may work only via the transient env _ensure_rocm_probe_env # just set, which dies with the installer. Persist the drop-in so login # shells (Studio, llama.cpp) inherit it -- else a reinstall over an @@ -2454,9 +2494,12 @@ _maybe_bootstrap_rocm_wsl() { fi # WSL GPU passthrough device must exist (present on any WSL2 GPU host). [ -e /dev/dxg ] || return 0 - # Only Strix Halo (gfx1151): rocminfo can't tell us the arch yet, so match - # the CPU model string WSL exposes (e.g. "AMD Ryzen AI Max+ ... Radeon 8060S"). - grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null || return 0 + # Strix APUs show in /proc/cpuinfo (the CPU model); discrete cards don't, so also + # ask the Windows host. Either signal suffices; the bootstrap detects arch from rocminfo. + if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \ + && ! _wsl_amd_gpu_name >/dev/null 2>&1; then + return 0 + fi command -v bash >/dev/null 2>&1 || return 0 # Fast path: already configured (librocdxg present) but launched from a @@ -2474,7 +2517,8 @@ _maybe_bootstrap_rocm_wsl() { fi echo "" - substep "Detected AMD Strix Halo (Radeon 8000S) in WSL with no ROCm runtime yet." "$C_WARN" + _rw_gpu="$(_wsl_amd_gpu_name 2>/dev/null || true)"; [ -n "$_rw_gpu" ] || _rw_gpu="an AMD GPU" + substep "Detected ${_rw_gpu} in WSL with no ROCm runtime yet." "$C_WARN" substep "Setting up ROCm-on-WSL (ROCm 7.2 + librocdxg) automatically to enable this GPU." substep "One-time, uses sudo and a large download. (skip: re-run with UNSLOTH_SKIP_ROCM_WSL_SETUP=1)" @@ -2852,7 +2896,7 @@ if [ "$_MIGRATED" = true ]; then # to prevent transitive torch resolution. run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" + "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" # Resolve pydantic WITH deps so pip pins pydantic-core to the # matching version (no-torch-runtime.txt below is --no-deps). # All transitive deps are torch-free. @@ -2863,9 +2907,11 @@ if [ "$_MIGRATED" = true ]; then run_install_cmd_retry "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT" fi else + # Pin mlx-lm away from 0.31.3 here too: a curl-piped migration has no + # overrides file, so UV_OVERRIDE is unset and this positional is the only cover. run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" + "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" ${_MLX_LM_EXCLUDE_ARG:-} fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2888,7 +2934,7 @@ if [ "$_MIGRATED" = true ]; then substep "repairing ROCm torch (overwritten by dependency resolution)..." run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \ - --index-url "$TORCH_INDEX_URL" \ + --default-index "$TORCH_INDEX_URL" \ --force-reinstall # torch was actually reinstalled from $TORCH_INDEX_URL now, so the # marker should record it (the preserved-torch case above must not). @@ -3017,7 +3063,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \ - --index-url "$TORCH_INDEX_URL" + --default-index "$TORCH_INDEX_URL" else substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..." # Record the ACTUAL wheel source for the torch-index marker: this @@ -3049,18 +3095,18 @@ elif [ -n "$TORCH_INDEX_URL" ]; then substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \ - --index-url "$TORCH_INDEX_URL" + --default-index "$TORCH_INDEX_URL" fi else substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN" run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \ - --index-url "$TORCH_INDEX_URL" + --default-index "$TORCH_INDEX_URL" fi else substep "installing PyTorch ($TORCH_INDEX_URL)..." run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \ - --index-url "$TORCH_INDEX_URL" + --default-index "$TORCH_INDEX_URL" fi # AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths). # Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm @@ -3081,7 +3127,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" + "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" # Same pydantic-with-deps trick as the migrated branch. run_install_cmd_retry "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic @@ -3099,7 +3145,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" + --upgrade-package unsloth "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -3108,7 +3154,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" else run_install_cmd_retry "install unsloth" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth -- "$PACKAGE_NAME" + --upgrade-package unsloth -- "$PACKAGE_NAME" ${_MLX_LM_EXCLUDE_ARG:-} fi # AMD ROCm: repair torch if the unsloth/unsloth-zoo install pulled in # CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1. @@ -3120,7 +3166,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then substep "repairing ROCm torch (overwritten by dependency resolution)..." run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \ - --index-url "$TORCH_INDEX_URL" \ + --default-index "$TORCH_INDEX_URL" \ --force-reinstall # The repair reinstalled torch from $TORCH_INDEX_URL (the generic # ROCm index), not the Radeon --find-links repo, so record THAT as @@ -3138,7 +3184,7 @@ else tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.6.7" "unsloth>=2026.6.9" --torch-backend=auto + run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.2" "unsloth>=2026.7.2" --torch-backend=auto substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -3162,14 +3208,14 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true) _installed_torch_tag="" [ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver") - # Repair when flavor is wrong AND the index is plain --index-url reinstallable + # Repair when flavor is wrong AND the index is plain --default-index reinstallable # (cuXXX / rocmX.Y / repo.amd.com gfx*); an unknown mirror leaf -> warn only. if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \ && [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..." run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \ - --index-url "$TORCH_INDEX_URL" \ + --default-index "$TORCH_INDEX_URL" \ --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio # torch was re-landed from $TORCH_INDEX_URL, so record it even on a # migrated venv whose flavor was genuinely wrong (the gfx-switch case @@ -3184,7 +3230,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN" substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN" substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN" - substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" \"$TORCHVISION_CONSTRAINT\" \"$TORCHAUDIO_CONSTRAINT\" --index-url $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN" + substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" \"$TORCHVISION_CONSTRAINT\" \"$TORCHAUDIO_CONSTRAINT\" --default-index $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN" fi fi fi diff --git a/pyproject.toml b/pyproject.toml index 844ead2454..2b79121c82 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,7 +73,7 @@ triton = [ ] huggingfacenotorch = [ - "unsloth_zoo>=2026.6.7", + "unsloth_zoo>=2026.7.2", "wheel>=0.42.0", "packaging", "numpy", @@ -94,7 +94,7 @@ huggingfacenotorch = [ ] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.6.7", + "unsloth_zoo>=2026.7.2", "torchvision", "unsloth[triton]", ] @@ -579,7 +579,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2026.6.7", + "unsloth_zoo>=2026.7.2", "packaging", "tyro", "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0", diff --git a/scripts/install_rocm_wsl_strixhalo.sh b/scripts/install_rocm_wsl_strixhalo.sh index 5ef9ee386a..aa560fc432 100644 --- a/scripts/install_rocm_wsl_strixhalo.sh +++ b/scripts/install_rocm_wsl_strixhalo.sh @@ -3,13 +3,14 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. # # ────────────────────────────────────────────────────────────────────────────── -# Enable ROCm-on-WSL for AMD Strix Halo (Radeon 8060S / gfx1151) +# Enable ROCm-on-WSL for AMD GPUs (Strix Halo/Point APUs AND discrete Radeon RX +# 7000/9000). Verified on gfx1151 (Radeon 8060S) and gfx1200 (Radeon RX 9060 XT). # ────────────────────────────────────────────────────────────────────────────── -# install.sh already routes gfx1151 to the right ROCm wheels once a ROCm runtime -# is present; what it does NOT do is install AMD's ROCm userspace + the WSL DXG -# bridge. This helper automates that Linux-side prerequisite on Ubuntu 24.04 -# WSL2 and is invoked by install.sh when it sees a Strix Halo APU in WSL (via -# /dev/dxg) but no ROCm runtime yet. Fully idempotent (re-run just re-verifies). +# install.sh routes the detected arch to the right ROCm wheels once a runtime exists; +# what it does NOT do is install AMD's ROCm userspace + the WSL DXG bridge (librocdxg). +# This helper does that Linux-side prerequisite on Ubuntu 24.04 WSL2, invoked by +# install.sh when it sees an AMD GPU via /dev/dxg but no ROCm yet. Arch-agnostic: the +# arch is auto-detected from rocminfo (override UNSLOTH_WSL_GFX=gfx1200). Idempotent. # # Manual, admin-gated Windows prerequisite: an AMD Adrenalin driver with # production ROCDXG/WSL support (26.2.2+). install.ps1 offers to update it. Once @@ -34,10 +35,12 @@ set -euo pipefail # ── Tunables (override via env) ────────────────────────────────────────────── ROCM_VER="${UNSLOTH_WSL_ROCM_VER:-7.2.1}" # ROCm release to install -GFX="gfx1151" +# GPU arch: empty = auto-detect from rocminfo after install (override UNSLOTH_WSL_GFX=gfx1200). +# The ROCm + librocdxg setup is arch-agnostic; only verify + the smoke test need the arch. +GFX="${UNSLOTH_WSL_GFX:-}" LIBROCDXG_REF="${UNSLOTH_LIBROCDXG_REF:-develop}" # ROCm/librocdxg git ref to build -# AMD's gfx1151 wheel index (same one install.sh uses); only for the smoke test. -TORCH_INDEX="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}/${GFX}/" +# AMD's wheel index for the (optional) smoke test; resolved after arch detection. +TORCH_INDEX="" # Optional torch smoke test (throwaway venv). OFF by default: install.sh installs # torch itself into the real venv right after, so a duplicate download is wasteful. SMOKE_TEST="${UNSLOTH_WSL_SMOKE_TEST:-0}" @@ -220,12 +223,12 @@ $SUDO ldconfig say "Persisting ROCm-on-WSL environment" _envfile="/etc/profile.d/unsloth-rocm-wsl.sh" $SUDO tee "$_envfile" >/dev/null <>> Unsloth ROCm-on-WSL (gfx1151) >>> +# >>> Unsloth ROCm-on-WSL >>> export HSA_ENABLE_DXG_DETECTION=1 export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 export PATH="${ROCM_DIR}/bin:\${PATH}" export LD_LIBRARY_PATH="${ROCM_DIR}/lib:\${LD_LIBRARY_PATH:-}" -# <<< Unsloth ROCm-on-WSL (gfx1151) <<< +# <<< Unsloth ROCm-on-WSL <<< EOF # also drop into ~/.bashrc for interactive shells if [ -n "${HOME:-}" ] && ! grep -q "Unsloth ROCm-on-WSL" "${HOME}/.bashrc" 2>/dev/null; then @@ -237,32 +240,50 @@ export PATH="${ROCM_DIR}/bin:${PATH}" export LD_LIBRARY_PATH="${ROCM_DIR}/lib:${LD_LIBRARY_PATH:-}" # ── Step 5: verify the runtime enumerates the GPU ──────────────────────────── -say "Verifying rocminfo sees ${GFX}" +say "Verifying rocminfo enumerates the GPU over DXG" # Capture rocminfo into a var BEFORE grepping: piping into `grep -q` SIGPIPEs # rocminfo on first match, which under `set -o pipefail` turns a successful match -# into a pipeline failure. Match the gfx1151 ISA "Name:" agent exactly (not a -# broad gfx1[0-9]) so a generic fallback ISA or unrelated RDNA GPU can't pass. +# into a pipeline failure. _rocminfo_out="$(rocminfo 2>/dev/null || true)" -if ! printf '%s\n' "$_rocminfo_out" | grep -qE "Name:[[:space:]]*${GFX}([^0-9]|$)"; then +# GPU agents advertise an ISA "Name: gfxNNNN". Match gfx[1-9] (excludes gfx000, the CPU +# agent), drop the "gfx*-generic" fallback ISA, and take the first real GPU arch. +_detected_gfx="$(printf '%s\n' "$_rocminfo_out" | grep -E 'Name:[[:space:]]*gfx[1-9]' | grep -v 'generic' | grep -oE 'gfx[1-9][0-9a-z]*' | head -1 || true)" +if [ -z "$_detected_gfx" ]; then printf '%s\n' "$_rocminfo_out" | head -25 >&2 || true - die "rocminfo did not enumerate a ${GFX} GPU agent. Most common cause: the Windows AMD driver predates production ROCDXG -- update Adrenalin (install.ps1 offers this), reboot, and re-run." + die "rocminfo did not enumerate any GPU agent. Most common cause: the Windows AMD driver predates production ROCDXG -- update Adrenalin (install.ps1 offers this), reboot, and re-run." fi +# Honour a caller-pinned arch (sanity-check via a consuming grep, not grep -q: under +# pipefail -q would SIGPIPE printf on large output and misreport the arch); else adopt. +if [ -n "$GFX" ] && ! printf '%s\n' "$_rocminfo_out" | grep -E "Name:[[:space:]]*${GFX}([^0-9]|$)" >/dev/null; then + die "rocminfo enumerated '${_detected_gfx}' but not the requested UNSLOTH_WSL_GFX='${GFX}'." +fi +GFX="${GFX:-$_detected_gfx}" # Display-only summary: best-effort (|| true) so head's early pipe-close under # `set -o pipefail` can't fail the bootstrap after verification already passed. printf '%s\n' "$_rocminfo_out" | grep -E 'Marketing Name|Device Type|Compute Unit' | grep -iE "Radeon|GPU|Compute" | head -3 || true note "ROCm-on-WSL runtime is live for ${GFX}." -# ── Step 6 (optional): torch smoke test from the gfx1151 index ─────────────── +# ── Step 6 (optional): torch smoke test from AMD's per-arch wheel index ─────── if [ "$SMOKE_TEST" = "1" ]; then say "Smoke-testing PyTorch on ${GFX} (throwaway venv)" + # Map the detected arch to AMD's repo.amd.com wheel family index. + case "$GFX" in + gfx1200|gfx1201) _fam="gfx120X-all" ;; + gfx1100|gfx1101|gfx1102|gfx1103) _fam="gfx110X-all" ;; + *) _fam="$GFX" ;; # gfx1150/gfx1151/gfx90a: own index + esac + TORCH_INDEX="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}/${_fam}/" _venv="${HOME}/.unsloth/rocm-smoketest" rm -rf "$_venv"; python3 -m venv "$_venv" "$_venv/bin/pip" install --quiet --upgrade pip - # gfx1151 index is primary (torch + triton); PyPI only an extra for pure-py + # AMD arch index is primary (torch + triton); PyPI only an extra for pure-py # deps. The constraint keeps pip on the ROCm wheel, not a newer PyPI CUDA torch. "$_venv/bin/pip" install --index-url "$TORCH_INDEX" \ --extra-index-url https://pypi.org/simple "$TORCH_CONSTRAINT" || \ die "torch install from ${TORCH_INDEX} failed." + # WSL: torch's bundled ROCr must load the DXG bridge -- drop librocdxg into torch/lib. + _tlib="$("$_venv/bin/python" -c 'import torch,os;print(os.path.join(os.path.dirname(torch.__file__),"lib"))' 2>/dev/null || true)" + [ -d "$_tlib" ] && cp -f "${ROCM_DIR}"/lib/librocdxg.so* "$_tlib"/ 2>/dev/null || true "$_venv/bin/python" - <<'PY' import torch ok = torch.cuda.is_available() diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json index 046566d148..27fa801a2a 100644 --- a/scripts/scan_packages_baseline.json +++ b/scripts/scan_packages_baseline.json @@ -39,7 +39,7 @@ "file": "botocore/utils.py", "check": "Reads credential paths AND makes network calls", "severity": "CRITICAL", - "evidence": "Creds: L3551: CACHE_DIR = os.path.expanduser(os.path.join('~', '.aws', 'boto', 'cache')) | L3721: return os.path.expanduser(os.path.join('~', '.aws', 'login', 'cache'))\nNetwork: L32: from urllib.request import getproxies, proxy_bypass", + "evidence": "Creds: L3551: CACHE_DIR = os.path.expanduser(os.path.join('~', '.aws', 'boto', 'cache')) | L3719: return os.path.expanduser(os.path.join('~', '.aws', 'login', 'cache'))\nNetwork: L32: from urllib.request import getproxies, proxy_bypass", "evidence_hash": "2d691bc373ab872aad23c744104596ba6d0d9f3b35aa101c7edbff4429b174c1" }, { @@ -55,23 +55,23 @@ "file": "datasets/utils/file_utils.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L443: while True: sha256:feba37d77721aa658e1786d2e4b67de76fefe1ceeb3ce8529d361c5241778eea", - "evidence_hash": "2e458563dec752d0a9896c9685d368d9906867110db315ab751e3eb6ec63f51c" + "evidence": "L441: while True: sha256:ce92e38c17c524815e1f9055be77235028c1e68e41b45cbfe9c8f1b867a205da", + "evidence_hash": "cb36281d28a975d101121c0702ee05eeee470879520d39a8be552129333f514d" }, { "package": "datasets", "file": "datasets/utils/file_utils.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L441: while True: sha256:ce92e38c17c524815e1f9055be77235028c1e68e41b45cbfe9c8f1b867a205da", - "evidence_hash": "cb36281d28a975d101121c0702ee05eeee470879520d39a8be552129333f514d" + "evidence": "L443: while True: sha256:feba37d77721aa658e1786d2e4b67de76fefe1ceeb3ce8529d361c5241778eea", + "evidence_hash": "2e458563dec752d0a9896c9685d368d9906867110db315ab751e3eb6ec63f51c" }, { "package": "diffusers", "file": "diffusers/utils/import_utils.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L1015: return importlib.import_module(\".\" + module_name, self.__name__)", + "evidence": "L1052: return importlib.import_module(\".\" + module_name, self.__name__)", "evidence_hash": "e584ecfdb097d9482bb19cd3992813bc1a119cfd4c40af14748bafe22900d91e" }, { @@ -79,7 +79,7 @@ "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, stream=True, timeout=DIFFUSERS_REQUEST_TIMEOUT).raw)", + "evidence": "Env: L236: value = os.environ[key]\nNetwork: L691: response = requests.get(arry, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L712: response = requests.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L731: image = PIL.Image.open(requests.get(image, stream=True, timeout=DIFFUSERS_REQUEST_TIMEOUT).raw)", "evidence_hash": "671190a6106c6ee9674e5e5942dc0940e1d2f8c78d5faf674413c2345b783fd9" }, { @@ -90,12 +90,20 @@ "evidence": "Archive: L317: a['TarFileType'] = tarfile.open(fileobj=_fileW,mode='w')\nNetwork: L330: x['SocketType'] = _socket = socket.socket()", "evidence_hash": "894862e547cf91b90cd6e4b495db3fb05b7490ef0d63de7e795a7e3d9447d850" }, + { + "package": "fastapi", + "file": "fastapi/routing.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L586: while True: sha256:251135b5ebfdd1248916449f32262575e003ef64382501c65b7e4061d67bda45", + "evidence_hash": "365aef4449c8089753d9398417cd76ab762cef547d75db70d87bca9c0b550ab5" + }, { "package": "fastmcp-slim", "file": "fastmcp/cli/apps_dev.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L1340: with tarfile.open(fileobj=io.BytesIO(data), mode=\"r:gz\") as tar:\nNetwork: L1291: with httpx.Client(timeout=30.0) as client: | L1305: with httpx.Client(timeout=30.0) as client: | L1335: with httpx.Client(timeout=30.0) as 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": "Archive: L1353: with tarfile.open(fileobj=io.BytesIO(data), mode=\"r:gz\") as tar:\nNetwork: L1304: with httpx.Client(timeout=30.0) as client: | L1318: with httpx.Client(timeout=30.0) as client: | L1348: with httpx.Client(timeout=30.0) as client: | L1549: client = httpx.AsyncClient(\nL1550: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1551: ) | L1713: async with httpx.AsyncClient(trust_env=False) as client: | L1781: with socket.socket(family, socket.SOCK_STREAM) as s:", "evidence_hash": "73a7a72013e9f800627ea07e6dbc3beeb8c905a6a5480c8fd896f0063173d25c" }, { @@ -103,8 +111,8 @@ "file": "fastmcp/cli/apps_dev.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "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" + "evidence": "FS: L637: history.replaceState(null, \"\", url); sha256:17068ba5bfed62c3a3007ec8bf3e0ea41ef6529b9e6112064d9afb3be9231436\nNetwork: L1304: with httpx.Client(timeout=30.0) as client: | L1318: with httpx.Client(timeout=30.0) as client: | L1348: with httpx.Client(timeout=30.0) as client: | L1549: client = httpx.AsyncClient(\nL1550: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1551: ) | L1713: async with httpx.AsyncClient(trust_env=False) as client: | L1781: with socket.socket(family, socket.SOCK_STREAM) as s:", + "evidence_hash": "e5325edfada6499540e6f0c24a0868979d275522e2b6a180aa9b5dd3280681b4" }, { "package": "fonttools", @@ -132,19 +140,35 @@ }, { "package": "huggingface-hub", - "file": "huggingface_hub/hf_api.py", + "file": "huggingface_hub/_sandbox.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L3746: while True: sha256:0c73ed1a7447120b112c063b14e720c6695bc11d00eb6b912cd0f10dc3e29b31", - "evidence_hash": "22f50b930e44146c5350bb99e6e6ebb09feea9bf1e899e407bedc4ffaf06721b" + "evidence": "L1179: while True: sha256:33ceddf9e42aae207e891e97808c518e92a0b27ab60e4326256717bfb25a3a38", + "evidence_hash": "802fd41d8bb17bf425e99d128c0351c820103a5efb74690a4086e542a71437b8" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/_sandbox.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L83: d=/tmp/.sbx-server\nL84: if command -v wget >/dev/null 2>&1; then wget -q --header \"Authorization: Bearer $SBX_DL_TOKEN\" -O \"$d\" \"$SBX_SERVER_URL\"\nL85: elif command -v curl >/dev/null 2>&1; then curl -fsSL -H \"Authorization: Bearer $SBX_DL_TOKEN\" -o \"$d\" \"$SBX_SERVER_URL\"\nL86: else cp \"$SBX_SERVER_MOUNT/sbx-server\" \"$d\"; fi\nL87: chmod +x \"$d\"", + "evidence_hash": "6908a3fe328fa94ee22a119998d6ad07cfa1ba4efa2628acf240f4204fd76e22" }, { "package": "huggingface-hub", "file": "huggingface_hub/hf_api.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L4600: while True: sha256:f4a851312a1832efe1b435aa1275a82184e19cc3f47e2cd244373d56c11de272", - "evidence_hash": "dc8fcf44788e32f42d1cc2eb0e2deb55eb2dbf2c3a55909a7d503e450f45e602" + "evidence": "L4677: while True: sha256:04afb38843e4125d1476f3f04bdad0edf1f63f8d75ad49a713b13e4bc68612fb", + "evidence_hash": "18877a2502c862b46a5d7e33fa7c39ab4ef32da7e1b07f596fd455f4376770c6" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/hf_api.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L3746: while True: sha256:0c73ed1a7447120b112c063b14e720c6695bc11d00eb6b912cd0f10dc3e29b31", + "evidence_hash": "22f50b930e44146c5350bb99e6e6ebb09feea9bf1e899e407bedc4ffaf06721b" }, { "package": "huggingface-hub", @@ -159,8 +183,8 @@ "file": "huggingface_hub/utils/_http.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L443: while True: sha256:0ab4fed32d3af10f361963371f681923481377508a405b5d8770cef75f859168", - "evidence_hash": "1484f6b92f41c427ba8cbc7c4695a94975fea683dfa83aa510b4b0e982be4721" + "evidence": "L462: while True: sha256:c75d1ee228cf7703a8c28551d649395a1f89f69a3aba69413f5bbcbd10c31958", + "evidence_hash": "d4d5f83fed39b87898cf776d5dad0bf1a6388a932f5fb7997d1070b50e46213e" }, { "package": "huggingface-hub", @@ -218,6 +242,22 @@ "evidence": "L5: import socket sha256:915068303029fa5806199f256fb74504c65f253f9aee8ea23d8e384bb772b1c7", "evidence_hash": "30be130f165f418dfd37b144c5ae333de184b95f828ab8bd4010a67b84a5f814" }, + { + "package": "multiprocess", + "file": "multiprocess/tests/__init__.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L3355: os.dup2(conn.fileno(), i) | L3387: \"test needs os.dup2()\") | L3405: os.dup2(fd, newfd) | L19: import socket sha256:26a745abdc7e89da28ab943394234d8ccb415e805477c3cc1f7d4766341a4c4c", + "evidence_hash": "a6b9bb85e9bb6682ab0dea4f95fd9266e8802f118c76d86dd87f7ab5864872cf" + }, + { + "package": "multiprocess", + "file": "multiprocess/tests/__init__.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L3521: os.dup2(conn.fileno(), i) | L3553: \"test needs os.dup2()\") | L3571: os.dup2(fd, newfd) | L20: import socket sha256:07d2933301c0dbeeb6e42381687827d8dd7cfd7471986c559ca64283d5ae6e24", + "evidence_hash": "db1f4ca69865ec3911d7450fe11d212b817139deda21cd7a4ee32d547a8dc452" + }, { "package": "numba", "file": "numba/pycc/decorators.py", @@ -231,7 +271,7 @@ "file": "numba/tests/support.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L1021: os.dup2(w, fd) | L1026: os.dup2(save, fd)", + "evidence": "L1016: os.dup2(w, fd) | L1021: os.dup2(save, fd)", "evidence_hash": "fea7aa03d48bf0f4386302fa444984c4f5dfc772cfec3f1df199fd33a52eec10" }, { @@ -298,13 +338,21 @@ "evidence": "Env: L105: token = os.environ.get(\"AWS_BEARER_TOKEN_BEDROCK\") | L150: environment_token = os.environ.get(\"AWS_BEARER_TOKEN_BEDROCK\")\nNetwork: L415: http_client: httpx.Client | None = None, | L531: http_client: httpx.Client | None = None, | L649: http_client: httpx.AsyncClient | None = None, | L767: http_client: httpx.AsyncClient | None = None,", "evidence_hash": "92dbec8ccd79c1e0bc41e93cdd0bdbb091220616c6a1352873196e9dda6bd85c" }, + { + "package": "openai", + "file": "openai/resources/beta/responses/responses.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L3999: while True: sha256:df298b6eaf3416589b79f4ef283f8fb76e54d505bfda8840673f8e6419117e2e", + "evidence_hash": "10ce5cb5a7097fcff4042ddcfb4802edda60aa4b7b113c8b926a52ddb76f78c2" + }, { "package": "openai", "file": "openai/resources/beta/threads/runs/runs.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L1074: while True: sha256:ef6d59a4a10b73a5af491f10af2885b7a309fda9468eb0f9572d19558d3ceb9f", - "evidence_hash": "43c03b55fedcbc980e5e6649c3c4493729128d280cc868349ab9590908ea5f99" + "evidence": "L1053: while True: sha256:973bb1aeca2e17e022872dc343a1bf5d8fe33bfa59fe01e2f8fe875522db5bce", + "evidence_hash": "24626e4aa53047a515ead563b42c07c43f73a2c5b82978fa59f58ffc2859e19b" }, { "package": "openai", @@ -319,7 +367,7 @@ "file": "openai/resources/responses/responses.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L3803: while True: sha256:1ce0b5a388c747945cdfda1a71b77afdfd03ae840d7aa9fa62f02eb00aa5e29f", + "evidence": "L3950: while True: sha256:1ce0b5a388c747945cdfda1a71b77afdfd03ae840d7aa9fa62f02eb00aa5e29f", "evidence_hash": "6de300ebb5e6e17cb51c89cbcdf08515a44655182f0776f0908a9d1043ebbcd7" }, { @@ -495,16 +543,16 @@ "file": "sklearn/datasets/_openml.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L100: while True: sha256:270363bb66980201e477f9b94886e4023f7a3d21b5ce026b7603a8c249a50c5b", - "evidence_hash": "53edbe07c312d459068d38e537b5114e65685ac3d4487b0423fa4542b5df20fe" + "evidence": "L100: while True: sha256:1f05a1b4fdd843b309634f583cb5e919866ef38ec5aa0b7d8a66ac8820655594", + "evidence_hash": "69597a64e5670a0f9a3c2aafc0bde4160f6170a9e2dc38f2c413cfa8d22ad193" }, { "package": "scikit-learn", "file": "sklearn/datasets/_openml.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L100: while True: sha256:1f05a1b4fdd843b309634f583cb5e919866ef38ec5aa0b7d8a66ac8820655594", - "evidence_hash": "69597a64e5670a0f9a3c2aafc0bde4160f6170a9e2dc38f2c413cfa8d22ad193" + "evidence": "L100: while True: sha256:270363bb66980201e477f9b94886e4023f7a3d21b5ce026b7603a8c249a50c5b", + "evidence_hash": "53edbe07c312d459068d38e537b5114e65685ac3d4487b0423fa4542b5df20fe" }, { "package": "scikit-learn", @@ -642,6 +690,14 @@ "evidence": "Base64: L1211: content = base64.b64decode(data)\nSubprocess: L2692: subprocess.run(\nL2693: cmd.split(), capture_output=True, text=True, check=True\nL2694: ) | L2995: cmd_output = subprocess.run(\nL2996: (\"openssl\", \"sha512\", filename), capture_output=True, text=True\nL2997: ) | L3707: out = subprocess.check_output(\nL3708: [\"ldd\", os.path.join(search, file)]\nL3709: ) | L3791: jobs.append(functools.partial(subprocess.check_call, cmd)) | L3876: subprocess.check_call(\nL3877: shlex.split(halide_cmd_gen.get_command_line())\nL3878: ) | L4336: subprocess.check_output(\nL4337: cmd_parts, stderr=subprocess.STDOUT, env=os.environ\nL4338: ) | L4591: output = subprocess.check_output(\nL4592: cmd_parts,\nL4593: stderr=subprocess.STDOUT,\nL4594: text=True,\nL4595: env=os.environ,\nL4596: )", "evidence_hash": "c09774087b702a6c5d6e2e85d9239c7c241ec938fbe9c0153e8f0b5c0710389b" }, + { + "package": "torch", + "file": "torch/_inductor/codecache.py", + "check": "base64 decode + subprocess execution (staged payload)", + "severity": "CRITICAL", + "evidence": "Base64: L1727: content = base64.b64decode(data)\nSubprocess: L3270: subprocess.run(\nL3271: cmd, capture_output=True, text=True, check=True\nL3272: ) | L3583: cmd_output = subprocess.run(\nL3584: (\"openssl\", \"sha512\", filename), capture_output=True, text=True\nL3585: ) | L4338: out = subprocess.check_output(\nL4339: [\"ldd\", os.path.join(search, file)]\nL4340: ) | L4422: jobs.append(functools.partial(subprocess.check_call, cmd)) | L4507: subprocess.check_call(\nL4508: shlex.split(halide_cmd_gen.get_command_line())\nL4509: ) | L4992: subprocess.check_output(\nL4993: cmd_parts, stderr=subprocess.STDOUT, env=os.environ\nL4994: ) | L5247: output = subprocess.check_output(\nL5248: cmd_parts,\nL5249: stderr=subprocess.STDOUT,\nL5250: text=True,\nL5251: env=os.environ,\nL5252: )", + "evidence_hash": "87f77b5f51cb84fe9950fdeeb90fe8710e1b863100e90b5e2cfb228a725bee06" + }, { "package": "torch", "file": "torch/ao/__init__.py", @@ -695,7 +751,7 @@ "file": "torch/testing/_internal/common_utils.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L4770: env = os.environ.copy()\nNetwork: L4832: with request.urlopen(url, timeout=15) as f1, open(path, 'wb' if binary else 'w') as f2: | L4850: with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:", + "evidence": "Env: L4900: env = os.environ.copy()\nNetwork: L4962: with request.urlopen(url, timeout=15) as f1, open(path, 'wb' if binary else 'w') as f2: | L4980: with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:", "evidence_hash": "704a851b9d68c9b885b9e15538bd7e96f03875503b618fe6f126c4438edd7386" }, { @@ -706,6 +762,14 @@ "evidence": "L32: import socket sha256:89faaaa8bc908e02dad73fd59b2b481fa91189c84b39b556c2766e71d2783bf3", "evidence_hash": "3d23d77ace91812a07cb9508cf352185d154176e8e8c8b9b28fa92cdbcfe0d53" }, + { + "package": "torch", + "file": "torch/testing/_internal/common_utils.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L32: import socket sha256:ba439cbf568b194872f1d974c02b0487e51f677b67e379400522d0992600bd2d", + "evidence_hash": "88e98b227573997f86eedea8e885a407b0dd549d46d4a3f0b840ec5aafe66865" + }, { "package": "torchvision", "file": "torchvision/datasets/utils.py", @@ -743,8 +807,8 @@ "file": "transformers/testing_utils.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L1663: while True: sha256:969e911d30c37a279ad915fb8c3d2d0a3f5705a7eb82ae6e00687388b68bbe65", - "evidence_hash": "2aa8e94baa805d599720a16afee6f08976482e301333e619e6c343389498ad15" + "evidence": "L1577: while True: sha256:2c6152f9da685f728e58d39dfc1827bc794f52606f56983bf38b5c6d0857cd5b", + "evidence_hash": "cdada67f3327237f00838a6750a4908dfaf76b9ab30c1352495c340d4fbd15c9" }, { "package": "transformers", @@ -759,15 +823,15 @@ "file": "transformers/testing_utils.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L1577: while True: sha256:2c6152f9da685f728e58d39dfc1827bc794f52606f56983bf38b5c6d0857cd5b", - "evidence_hash": "cdada67f3327237f00838a6750a4908dfaf76b9ab30c1352495c340d4fbd15c9" + "evidence": "L1699: while True: sha256:969e911d30c37a279ad915fb8c3d2d0a3f5705a7eb82ae6e00687388b68bbe65", + "evidence_hash": "2aa8e94baa805d599720a16afee6f08976482e301333e619e6c343389498ad15" }, { "package": "transformers", "file": "transformers/testing_utils.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: 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": "Env: L288: value = os.environ[key] | L304: value = os.environ[key] | L2165: env = os.environ.copy() | L2287: for k in list(os.environ.keys()):\nNetwork: L2597: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:", "evidence_hash": "73ff16aee09cf163fb3a7a04dfa2cf610595bde2f19460a579397695f728e3f4" }, { @@ -799,16 +863,16 @@ "file": "trl/extras/vllm_client.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L146: while True: sha256:2beedc742e1f085eaa10fd3bc40be97d2331d21887ef1b9ccdfa2150a184edfe", - "evidence_hash": "1540dffaaa053780e953e04c11d9c6b9c74b91cb60f3e6d87451ba7fe7db46db" + "evidence": "L152: while True: sha256:93e7d409e300af445376e6defbe2d0241aa19ecf63ed41b780fbb91c7d09856f", + "evidence_hash": "208838617172de61bca201d2a1bbeb5aa5aaa55feb1a1069cf39214673a7d6d1" }, { "package": "trl", "file": "trl/extras/vllm_client.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L152: while True: sha256:93e7d409e300af445376e6defbe2d0241aa19ecf63ed41b780fbb91c7d09856f", - "evidence_hash": "208838617172de61bca201d2a1bbeb5aa5aaa55feb1a1069cf39214673a7d6d1" + "evidence": "L146: while True: sha256:2beedc742e1f085eaa10fd3bc40be97d2331d21887ef1b9ccdfa2150a184edfe", + "evidence_hash": "1540dffaaa053780e953e04c11d9c6b9c74b91cb60f3e6d87451ba7fe7db46db" }, { "package": "trl", @@ -866,6 +930,14 @@ "evidence": "Crypto: L294: r\"|\\b(?:xprv|xpub|bc1|0x[a-fA-F0-9]{40})\\b\",\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as resp:", "evidence_hash": "278ff15b0b702d37d7f0b30a1e55a31bf2b11883685718a47478fbb5ce7f5212" }, + { + "package": "unsloth-zoo", + "file": "scripts/scan_packages.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\", | L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\", sha256:78268349021e21bedcd2eaaa5b4a71b0de1d52e023ada914dfdc09515ee1aad8", + "evidence_hash": "590fe1c96c442fbea5eb8642650257bc0b0199e919b9bacdb11dfa767b6fe839" + }, { "package": "unsloth-zoo", "file": "tests/security/fixtures/_build.py", @@ -919,16 +991,16 @@ "file": "tests/test_mlx_save_export_regressions.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L164: temporary_location=\"/tmp/ignored\", sha256:78837e80d48e872ef191aaacfe5e1c621a98a20df486a70a41d1a932d074a5b3", - "evidence_hash": "dd11376e664d0d7e7f4cc4baf57eacd4b7ae7b03222dce3912ce68b63dbfca1e" + "evidence": "L165: temporary_location=\"/tmp/ignored\", sha256:9f8502377b19666288b28399633dfc6740a64d0cb70ad1615e38b1269f94bf37", + "evidence_hash": "b7262d6e58f2ebad961dd3e64ca6c32bba356b5044d7a642d7dbd36a58cb6c81" }, { "package": "unsloth-zoo", "file": "tests/test_quantize_gguf_q2_k_l.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L67: input_gguf=\"/tmp/in.gguf\", sha256:32532cadc357beee1009f4e86481bdbe60a0b7bf47f6bb022b05ec1b8e15aed0", - "evidence_hash": "49f5b67379de17178f21a9bc93b79d6b94a70ecbdd16de86574934aac30a071d" + "evidence": "L67: input_gguf=\"/tmp/in.gguf\", sha256:06789b55e8f31426c233f37ff7d3729cc9e1f61c0829abd2c00c39216c63c7ad", + "evidence_hash": "ad4913d9099eb9b70e09d6860b242eb5f48c67e46d9bf4ae35c1c38a267d753b" }, { "package": "unsloth-zoo", @@ -951,7 +1023,7 @@ "file": "unsloth_zoo/llama_cpp.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "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": "Archive: L938: with tarfile.open(archive_path, \"r:gz\") as archive:\nNetwork: L691: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1699: response = requests.get(\nL1700: LLAMA_CPP_CONVERT_FILE, timeout = (10, 120)\nL1701: ) | L2873: check = requests.get(llama_cpp_chat_file, timeout = 5)", "evidence_hash": "b9f3b1652349fa8ef9ac2d1715978aca1e1632165851a00a2698dd47189e410c" }, { @@ -959,7 +1031,7 @@ "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()) | 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": "Env: L125: keynames = \"\\n\" + \"\\n\".join(os.environ.keys()) | L683: token = os.environ.get(\"GH_TOKEN\") or os.environ.get(\"GITHUB_TOKEN\")\nNetwork: L691: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1699: response = requests.get(\nL1700: LLAMA_CPP_CONVERT_FILE, timeout = (10, 120)\nL1701: ) | L2873: check = requests.get(llama_cpp_chat_file, timeout = 5)", "evidence_hash": "9cd0b1bb59c7eb1d814d7636dfd167c34f265eb7c4521a9d88b2bdcfd535b926" }, { @@ -1002,6 +1074,14 @@ "evidence": "Obfusc: L87: __import__(name)\nExec: L735: exec(\"\"\"exec _code_ in _globs_, _locs_\"\"\")", "evidence_hash": "3cb7d8247dea7dd3d7b21ededc0181c58c50099aeb73c9138a286f3d1ad92d4f" }, + { + "package": "cffi", + "file": "cffi/_cffi_gen_src.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L52: compiled = compile(source=pysrc, filename=filename, mode='exec')\nExec: L53: exec(compiled, globs, globs)", + "evidence_hash": "c429e4c977a61db6b7c717b5a552fce74eda622213e49eb5467a3782fd746fb9" + }, { "package": "cffi", "file": "cffi/setuptools_ext.py", @@ -1127,7 +1207,7 @@ "file": "numba/tests/support.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L879: __import__(modname)\nExec: L813: eval(co, globs, ns)", + "evidence": "Obfusc: L874: __import__(modname)\nExec: L808: eval(co, globs, ns)", "evidence_hash": "649a7d750f903478243b0bcb9e8020521b505fc7fedc5b696ec01f4efc096109" }, { @@ -1159,7 +1239,7 @@ "file": "numba/tests/test_np_functions.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L7106: exec(compile(funcstr, '', 'exec'), globals(), dct)\nExec: L7106: exec(compile(funcstr, '', 'exec'), globals(), dct)", + "evidence": "Obfusc: L7118: exec(compile(funcstr, '', 'exec'), globals(), dct)\nExec: L7118: exec(compile(funcstr, '', 'exec'), globals(), dct)", "evidence_hash": "9e81164131d16056fb56ad3cd11b8d129d1ff4f5855031e8b501e0335d5c14ed" }, { @@ -1175,16 +1255,16 @@ "file": "numpy/testing/_private/utils.py", "check": "Anti-analysis/sandbox evasion + suspicious behavior", "severity": "HIGH", - "evidence": "Anti: L2777: original_trace = sys.gettrace() | L2779: sys.settrace(None) | L2782: sys.settrace(original_trace)\nSubprocess: L1478: output = subprocess.run(cmd, capture_output=True, text=True)\nExec: L1346: exec(astr, dict) | L1632: exec(code, globs, locs)", - "evidence_hash": "27468a6828101c6c026ae25aca8aa90ef485fd62b2c8f0967479edae9c965844" + "evidence": "Anti: L2788: original_trace = sys.gettrace() | L2790: sys.settrace(None) | L2793: sys.settrace(original_trace)\nSubprocess: L1486: output = subprocess.run(cmd, capture_output=True, text=True) | L2889: res = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True,\nL2890: errors=\"replace\", **kwargs)\nExec: L1352: exec(astr, dict) | L1640: exec(code, globs, locs)", + "evidence_hash": "9c6961817e5b1751e572dfe0858286703bb835870ecdfd6a7a9fdd8372a5dd2b" }, { "package": "numpy", "file": "numpy/testing/_private/utils.py", "check": "Anti-analysis/sandbox evasion + suspicious behavior", "severity": "HIGH", - "evidence": "Anti: 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" + "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", @@ -1199,7 +1279,7 @@ "file": "PIL/Image.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L422: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), []) | L490: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), [])\nExec: L3772: def eval(image: Image, *args: Callable[[int], float]) -> Image:", + "evidence": "Obfusc: L422: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), []) | L490: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), [])\nExec: L3776: def eval(image: Image, *args: Callable[[int], float]) -> Image:", "evidence_hash": "c2c1e7ae44e15862caf8de549d09db7b35e93282450f07ef61aaf5450a408c13" }, { @@ -1255,7 +1335,7 @@ "file": "setuptools/_distutils/compilers/C/base.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L1286: __import__(module_name)\nExec: L1113: if lib_type not in eval(expected):", + "evidence": "Obfusc: L1287: __import__(module_name)\nExec: L1114: if lib_type not in eval(expected):", "evidence_hash": "368651e9818ed2d1bb009027d3bcfbf94ae30639c0882a6c2bddde97b8c4f1e5" }, { @@ -1271,7 +1351,7 @@ "file": "setuptools/tests/config/test_pyprojecttoml.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L364: \"setup.py\": \"__import__('setuptools').setup(include_package_data=False)\",\nExec: L98: \"__main__.py\": \"def exec(): print('hello')\",", + "evidence": "Obfusc: L387: \"setup.py\": \"__import__('setuptools').setup(include_package_data=False)\",\nExec: L98: \"__main__.py\": \"def exec(): print('hello')\",", "evidence_hash": "067d41014f72a61d8b4adf25f3659d1f66a0e909f732223f48837aa7684df4e6" }, { @@ -1279,7 +1359,7 @@ "file": "setuptools/tests/test_editable_install.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L120: SETUP_SCRIPT_STUB = \"__import__('setuptools').setup()\"\nExec: L449: exec(finder, loc, loc)", + "evidence": "Obfusc: L120: SETUP_SCRIPT_STUB = \"__import__('setuptools').setup()\"\nExec: L447: exec(finder, loc, loc)", "evidence_hash": "a78d7f5af7eb4ba92656cda258c195b92f6337c585c97d0823e47a9d4a2eb15d" }, { @@ -1322,12 +1402,20 @@ "evidence": "Obfusc: L919: c = compile(funcstr, filename, 'exec')\nExec: L163: module = eval(import_command) | L170: exec(import_command, {}, namespace) | L903: exec(ln, {}, namespace) | L909: exec(ln, {}, namespace) | L920: exec(c, namespace, funclocals)", "evidence_hash": "ab4f5819576a70038301668b8f3e4a781c4b757b146117d5d93eab1896a5a6cd" }, + { + "package": "tensorboard", + "file": "tensorboard/plugins/projector/tf_projector_plugin/projector_binary.js", + "check": "Python wheel ships large JS bundle (uncommon; manually review)", + "severity": "HIGH", + "evidence": "sha256: 53c38430766be25dc672a30846ac3b9eba86aee35eb0746785ec012647c7d9a2", + "evidence_hash": "2c6384e8115a6d5dacf1f84d8f724832d8dc59feb442bb98ffae0857c0ccb381" + }, { "package": "torch", "file": "torch/_dynamo/bytecode_debugger.py", "check": "Anti-analysis/sandbox evasion + suspicious behavior", "severity": "HIGH", - "evidence": "Anti: L1048: self._old_trace = sys.gettrace() | L1049: sys.settrace(self._settrace_callback) | L1106: sys.settrace(self._old_trace)\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": "Anti: L1052: self._old_trace = sys.gettrace() | L1053: sys.settrace(self._settrace_callback) | L1113: sys.settrace(self._old_trace)\nExec: L684: result = eval(arg, frame_globals, eval_locals) | L709: result = eval(cmd, frame_globals, eval_locals) | L717: exec(cmd, frame_globals, eval_locals)", "evidence_hash": "dc2afd1769d357c15b69802bd2799fafa059c0b1dcdd4937528fb5b601962f1b" }, { @@ -1343,7 +1431,7 @@ "file": "torch/fx/experimental/rewriter.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L46: code = compile(dest_ast, \"\", \"exec\")\nExec: L49: exec(code, globals_dict)", + "evidence": "Obfusc: L44: code = compile(dest_ast, \"\", \"exec\")\nExec: L47: exec(code, globals_dict)", "evidence_hash": "76374f96feed416eec390458843621f33524cfb8d93ef0f3eb4cb1b47d0ad748" }, { @@ -1359,7 +1447,7 @@ "file": "torch/package/package_importer.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L602: def __import__(self, name, globals=None, locals=None, fromlist=(), level=0):\nExec: L412: exec(code, ns)", + "evidence": "Obfusc: L599: def __import__(self, name, globals=None, locals=None, fromlist=(), level=0):\nExec: L412: exec(code, ns)", "evidence_hash": "c7c0650f0c74a086d224112f77ee76634b8f47afc047ce27fee8c7fc45560512" }, { @@ -1391,7 +1479,7 @@ "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: L1158: assert ppl == pytest.approx(__import__(\"math\").exp(2.5))\nExec: L1136: def eval(self):", "evidence_hash": "c409327ef6420cc0c7224506fcb82b11bbc9838a6f2f97c9c2cfc00a40c4cdbf" }, { @@ -1407,7 +1495,7 @@ "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": "Obfusc: L1013: _mod = __import__(model_location, fromlist=items) | L4295: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\\nL4296: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4297: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nL4298: f' \"-____-\" Trainable parameters = {get_model_param_count(model, trainable_only=True):,} of {get_model_param_count(model):,} ({get_model_param_count(model, trainable_only=True)/get_model_p sha256:9e832c1e7b2815aa44dc42638d821cb9df22550db88d85bd4bfb29c13f984b53 | L4296: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4297: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nL4298: f' \"-____-\" Trainable parameters = {get_model_param_count(model, trainable_only=True):,} of {get_model_param_count(model):,} ({get_model_param_count(model, trainable_only=True)/get_model_p sha256:9e832c1e7b2815aa44dc42638d821cb9df22550db88d85bd4bfb29c13f984b53 | L4295: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\\nL4296: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4297: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nExec: L612: if eval(_dtype) is not None: | L613: dtype = eval(_dtype) | L955: _modeling_file = eval(model_location) | L1255: f = eval(f\"{model_location}.{module}\") | L1563: exec(f\"def raise_{j}(*args, **kwargs): print('{function}')\", globals(), locals()) | L1564: try: exec(f\"EMPTY_LOGITS.{function} = raise_{j}\", globals(), locals()) | L2699: exec(f\"import {parent}\", locals(), globals()) | L2830: dir(eval(parent)), | L2834: exec(f\"{parent}.{child}.forward = forward\", globals(), locals()) | L2908: module = eval(f\"modeling_file.{module}\") | L2935: inner_class = eval(f\"modeling_file.{inner_class}\") | L3065: exec(f\"from timm.layers.norm_act import {norm}\") | L3073: forward = eval(norm).forward | L3079: exec(f\"timm.layers.norm_act.{norm}.forward = forward\") | L3096: exec(f\"from timm.models._efficientnet_blocks import {block}\") | L3104: forward = eval(block).forward | L3110: exec(f\"timm.models._efficientnet_blocks.{block}.forward = forward\") | L3389: exec(f\"import {model_location}\", globals()) | L3392: modeling_file = eval(model_location) | L3405: exec(\nL3406: \"model_logger.addFilter(HideLoggingMessage('`use_cache`'))\", globals(), locals()\nL3407: ) | L3409: exec(\nL3410: \"model_logger.addFilter(HideLoggingMessage('compile_config'))\",\nL3411: globals(),\nL3412: locals(),\nL3413: ) | L3564: source = eval(f\"modeling_file.{module}\") | L3578: source = eval(f\"modeling_file.{module}\") | L3679: source = eval(f\"modeling_file.{module}\") | L3717: source = eval(f\"{model_location}.{module}\") | L3788: source = eval(f\"{model_location}.{module}\") | L3836: source = eval(f\"{model_location}.{module}\") | L4058: source = eval(f\"{model_location}.{module}\") | L4069: exec(\nL4070: f\"{model_location}.{module}._update_causal_mask = no_update_causal_mask\",\nL4071: globals(),\nL4072: ) | L4135: source = eval(f\"{model_location}.{module}\") | L4176: module_cls = eval(f\"{model_location}.{module}\") | L4213: module_cls = eval(f\"{model_location}.{module}\") | L4280: exec(\nL4281: \"from transformers.trainer import (\" + \", \".join(x for x in good_items) + \")\",\nL4282: globals(),\nL4283: ) | L4345: exec(inner_training_loop, globals()) | L4353: function = eval(f\"{model_location}.{module}\") | L4431: function = eval(f\"{model_location}.{module}\") | L4566: source = eval(f\"{model_location}.torch\") | L4573: function = eval(f\"source.nn.{module}\") | L4632: exec(\nL4633: f\"{model_location}.torch.nn.{module}.forward = forward\",\nL4634: globals(),\nL4635: locals(),\nL4636: ) | L4638: exec(\nL4639: f\"{model_location}.nn.{module}.forward = forward\",\nL4640: globals(),\nL4641: locals(),\nL4642: ) | L4646: exec(\nL4647: f\"combined_module.torch.nn.{module}.forward = forward\",\nL4648: globals(),\nL4649: locals(),\nL4650: ) | L4652: exec(\nL4653: f\"combined_module.nn.{module}.forward = forward\",\nL4654: globals(),\nL4655: locals(),\nL4656: ) | L4673: exec(\nL4674: f\"{model_location}.{module} = combined_module.{module}\",\nL4675: globals(),\nL4676: locals(),\nL4677: ) | L4687: check_dicts = dir(eval(f\"{model_location}\")) | L4689: item = eval(f\"{model_location}.{check}\") | L4699: exec(\nL4700: f\"{model_location}.{check}['{key}'] = combined_module.{replaced_class}\",\nL4701: globals(),\nL4702: locals(),\nL4703: )", "evidence_hash": "ec1875fd32d00fe885e566ebda75163e46e838ca31020abb57e0991892c2bdf7" }, { @@ -1423,8 +1511,8 @@ "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" + "evidence": "Obfusc: L2869: _mod = __import__(module_name, fromlist=[\"_\"])\nExec: L148: mx.eval(model.parameters()) | L180: mx.eval(model.parameters()) | L732: mx.eval(model.parameters()) | L733: mx.eval(mx.distributed.all_sum(mx.array(1.0), stream=mx.cpu)) | L799: mx.eval(model.parameters()) | L802: mx.eval(mx.distributed.all_sum(mx.array(1.0), stream=mx.cpu)) | L2673: model.eval() | L3256: mx.eval(model.parameters()) | L3372: mx.eval(module.weight) | L5666: mx.eval(model.parameters()) | L5716: mx.eval(model.parameters()) | L5859: mx.eval(model.parameters())", + "evidence_hash": "7b44760032c5df6d379ccfdd0bff3d23f857f64e08210fa0fba8d2881d457634" }, { "package": "unsloth-zoo", @@ -1439,7 +1527,7 @@ "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": "Obfusc: L4015: module = __import__('transformers', fromlist=[model_class_name])\nExec: L3897: exec(f\"from transformers.modeling_utils import ({', '.join(functions)})\", locals(), globals()) | L3943: exec(save_pretrained, globals(), functions)", "evidence_hash": "530b2383acd9fe8330aa65cd0bf86164aaacd47770e7c8d0752195bee36396ec" }, { @@ -1449,46 +1537,6 @@ "severity": "HIGH", "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/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index b4c908b0cb..22a21a2ebc 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -564,6 +564,12 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]] for n, tids in b["module_import_targets"].items(): if tids & after_used: continue # resolved -> fine + # `from __future__ import ...` is a compiler directive, not a runtime + # binding: the name (`annotations`, ...) is never loaded, so it can never + # "resolve" to a use. Skip it so a legitimately-added future import + # (e.g. `annotations` for lazy PEP 604 `X | None` on py3.9) is not flagged. + if all(t.startswith("from:__future__:") for t in tids): + continue newly_added = bool(tids - before_module_targets) was_used_before = bool(tids & before_used) if newly_added or was_used_before: @@ -588,9 +594,23 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]] # package object and only *add* submodule attributes (e.g. adding # `import urllib.error` next to `import urllib.request`). Nothing the name # resolved to before is lost, so no reference is re-pointed -- skip it. + # + # A deliberate *relocation* is also benign and must not block: when a name + # keeps its spelling but its import source is moved A -> B in THIS diff (the + # old `from A import x` is removed at module level and a new `from B import x` + # is added), the swap is intentional, not a silent re-point to a pre-existing + # different object. This mirrors the relocation tolerance already applied to + # TARGET-MISSING. The dangerous case -- the name now resolving to a target + # that already existed before (shadow/clash) -- is NOT exempted. + removed_module_targets = before_module_targets - after_module_targets for key, tafter in b["target_by_use"].items(): tbefore = a["target_by_use"].get(key) if tbefore and tbefore != tafter and (tbefore - tafter): + lost = tbefore - tafter + gained = tafter - tbefore + relocated = lost <= removed_module_targets and gained <= added_module_targets + if relocated: + continue findings.append( ( "BLOCKER", diff --git a/studio/backend/assets/configs/inference_defaults.json b/studio/backend/assets/configs/inference_defaults.json index 1c7a409bc1..0633f80bbc 100644 --- a/studio/backend/assets/configs/inference_defaults.json +++ b/studio/backend/assets/configs/inference_defaults.json @@ -235,6 +235,13 @@ "min_p": 0.1, "repetition_penalty": 1.0 }, + "deepseek-v4": { + "temperature": 1.0, + "top_p": 1.0, + "top_k": -1, + "min_p": 0.0, + "repetition_penalty": 1.0 + }, "deepseek-r1": { "temperature": 0.6, "top_p": 0.95, @@ -394,7 +401,7 @@ "phi-4", "phi-3", "mistral-nemo", "mistral-small", "mistral-large", "magistral", "ministral", "devstral", "pixtral", - "deepseek-r1", "deepseek-v3", "deepseek-ocr", + "deepseek-v4", "deepseek-r1", "deepseek-v3", "deepseek-ocr", "glm-5", "glm-4", "nemotron", "minimax-m2.7", "minimax-m2.5", "minimax", diff --git a/studio/backend/core/inference/__init__.py b/studio/backend/core/inference/__init__.py index 2faf70bb79..ad78157418 100644 --- a/studio/backend/core/inference/__init__.py +++ b/studio/backend/core/inference/__init__.py @@ -7,13 +7,16 @@ Inference submodule - backend for model loading and generation. The default get_inference_backend() returns an InferenceOrchestrator that delegates to a subprocess. The original InferenceBackend runs inside the subprocess and can be imported directly from .inference when needed. + +Public names are resolved lazily (PEP 562): importing this package -- or a +dependency-light leaf like ``core.inference.chat_eos`` -- must NOT eagerly pull +the orchestrator / llama_cpp import chain (httpx, subprocess plumbing, the ML +backend and its Studio dependencies). Those load only when a public name is +actually accessed, so standalone helpers stay unit-testable without the full +inference stack. """ -from .orchestrator import InferenceOrchestrator, get_inference_backend -from .llama_cpp import LlamaCppBackend - -# Expose InferenceOrchestrator as InferenceBackend for backward compat. -InferenceBackend = InferenceOrchestrator +from typing import TYPE_CHECKING __all__ = [ "InferenceBackend", @@ -21,3 +24,33 @@ __all__ = [ "get_inference_backend", "LlamaCppBackend", ] + +# name -> (submodule, attribute); InferenceBackend aliases InferenceOrchestrator. +_LAZY_ATTRS = { + "InferenceOrchestrator": ("orchestrator", "InferenceOrchestrator"), + "InferenceBackend": ("orchestrator", "InferenceOrchestrator"), + "get_inference_backend": ("orchestrator", "get_inference_backend"), + "LlamaCppBackend": ("llama_cpp", "LlamaCppBackend"), +} + + +def __getattr__(name): + try: + submodule, attr = _LAZY_ATTRS[name] + except KeyError: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None + from importlib import import_module + + value = getattr(import_module(f"{__name__}.{submodule}"), attr) + globals()[name] = value # cache so later access skips __getattr__ + return value + + +def __dir__(): + return sorted(set(globals()) | set(__all__)) + + +if TYPE_CHECKING: # keep static analysers / IDEs aware of the lazy names + from .llama_cpp import LlamaCppBackend + from .orchestrator import InferenceOrchestrator, get_inference_backend + InferenceBackend = InferenceOrchestrator diff --git a/studio/backend/core/inference/_vulkan_probe.py b/studio/backend/core/inference/_vulkan_probe.py new file mode 100644 index 0000000000..706346daad --- /dev/null +++ b/studio/backend/core/inference/_vulkan_probe.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Standalone free-VRAM probe for the bundled ggml Vulkan backend. + +Run in a short-lived subprocess (``python _vulkan_probe.py ``) so the +Vulkan instance never lives in the long-running backend process. Loads the +bundled ggml Vulkan backend from ```` and prints one +``\\t\\t\\t`` line per device to stdout. +Indices are ggml's own Vulkan device ordinals, which need not match nvidia-smi +order. ``is_igpu`` (from ggml's device type) is ``1`` for an integrated GPU +sharing system RAM. ``total_bytes`` is the device-local heap; the reader uses +it to reserve absolute headroom on a discrete card (parity with the CUDA/ROCm +fit) and ignores it for an iGPU, whose "VRAM" is shared system RAM. + +Uses only the standard library so it stays runnable as a bare script. +""" + +import ctypes +import os +import sys + +# ggml_backend_dev_type enum (ggml-backend.h): CPU=0, GPU=1, IGPU=2, ... +_GGML_BACKEND_DEVICE_TYPE_IGPU = 2 + + +def _igpu_flags(base, lib, count: int) -> list[bool]: + """Per-device integrated-GPU flags via ggml's backend registry. + + The Vulkan reg enumerates devices in the same order as + ``ggml_backend_vk_get_device_memory`` (each context uses ``ctx->device = + i``), so reg index == device ordinal. Returns all-False on any failure so + the reader never over-caps a discrete card. + """ + flags = [False] * count + try: + lib.ggml_backend_vk_reg.restype = ctypes.c_void_p + lib.ggml_backend_vk_reg.argtypes = [] + base.ggml_backend_reg_dev_count.restype = ctypes.c_size_t + base.ggml_backend_reg_dev_count.argtypes = [ctypes.c_void_p] + base.ggml_backend_reg_dev_get.restype = ctypes.c_void_p + base.ggml_backend_reg_dev_get.argtypes = [ctypes.c_void_p, ctypes.c_size_t] + base.ggml_backend_dev_type.restype = ctypes.c_int + base.ggml_backend_dev_type.argtypes = [ctypes.c_void_p] + + reg = lib.ggml_backend_vk_reg() + if not reg: + return flags + dev_count = base.ggml_backend_reg_dev_count(reg) + for i in range(min(count, dev_count)): + dev = base.ggml_backend_reg_dev_get(reg, i) + if dev: + flags[i] = base.ggml_backend_dev_type(dev) == _GGML_BACKEND_DEVICE_TYPE_IGPU + except Exception: + # Best-effort: any failure degrades to "discrete" so the memory + # readings still get through instead of crashing the probe. + pass + return flags + + +def main() -> int: + if len(sys.argv) < 2: + return 0 + bindir = sys.argv[1] + + # Hold add_dll_directory's handle for the rest of main() (the documented + # idiom) so bindir stays on the search path while the sibling ggml DLLs + # resolve below. + _dll_dir = None + if sys.platform == "win32": + base_name, vk_name = "ggml-base.dll", "ggml-vulkan.dll" + try: + _dll_dir = os.add_dll_directory(bindir) + except Exception: + pass + else: + base_name, vk_name = "libggml-base.so", "libggml-vulkan.so" + + # RTLD_GLOBAL exposes ggml-base's symbols to ggml-vulkan on POSIX. getattr + # falls back to 0 where the flag doesn't exist (Windows CDLL ignores mode). + _rtld_global = getattr(ctypes, "RTLD_GLOBAL", 0) + try: + base = ctypes.CDLL(os.path.join(bindir, base_name), mode = _rtld_global) + lib = ctypes.CDLL(os.path.join(bindir, vk_name), mode = _rtld_global) + except OSError as e: + print(f"ggml-vulkan load failed: {e}", file = sys.stderr) + return 1 + + lib.ggml_backend_vk_get_device_count.restype = ctypes.c_int + lib.ggml_backend_vk_get_device_count.argtypes = [] + lib.ggml_backend_vk_get_device_memory.restype = None + lib.ggml_backend_vk_get_device_memory.argtypes = [ + ctypes.c_int, + ctypes.POINTER(ctypes.c_size_t), + ctypes.POINTER(ctypes.c_size_t), + ] + + count = lib.ggml_backend_vk_get_device_count() + igpu = _igpu_flags(base, lib, count) + rows = [] + for i in range(count): + free, total = ctypes.c_size_t(0), ctypes.c_size_t(0) + lib.ggml_backend_vk_get_device_memory(i, ctypes.byref(free), ctypes.byref(total)) + rows.append("%d\t%d\t%d\t%d" % (i, free.value, int(igpu[i]), total.value)) + sys.stdout.write("\n".join(rows)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/studio/backend/core/inference/anthropic_compat.py b/studio/backend/core/inference/anthropic_compat.py index 7b572a28ff..3c7a4cb182 100644 --- a/studio/backend/core/inference/anthropic_compat.py +++ b/studio/backend/core/inference/anthropic_compat.py @@ -258,6 +258,10 @@ class AnthropicStreamEmitter: self._open_tool_use_id: Optional[str] = None self._open_tool_args_sent: bool = False self._prev_text: str = "" + # Net minus in the text emitted to the client. Tracked + # from emitted deltas (not _prev_text, which a final bare shrink clobbers) + # so an unclosed reasoning-only block can be balanced before close. + self._open_think_tags: int = 0 self._usage: dict = {} def start( @@ -317,6 +321,7 @@ class AnthropicStreamEmitter: """Close any open block and emit message_delta + message_stop.""" events = [] if self._text_block_open or self._open_tool_call_id is not None: + events.extend(self._close_open_think()) events.append(self._close_block()) self._open_tool_call_id = None self._open_tool_use_id = None @@ -344,12 +349,33 @@ class AnthropicStreamEmitter: ) return events + def _close_open_think(self) -> list[str]: + """Emit a ```` delta when the streamed text left a ```` + open. This emitter diffs cumulative snapshots and drops the generator's + final bare shrink, so a reasoning-only reply would otherwise end on an + unclosed tag. Mirrors the chat route's reasoning extractor, which closes + the block on finish; balances the block before it is closed.""" + if not self._text_block_open or self._open_think_tags <= 0: + return [] + self._open_think_tags = 0 + return [ + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self.block_index, + "delta": {"type": "text_delta", "text": ""}, + }, + ) + ] + def _handle_content(self, event: dict) -> list[str]: cumulative = event.get("text", "") new_text = cumulative[len(self._prev_text) :] self._prev_text = cumulative if not new_text: return [] + self._open_think_tags += new_text.count("") - new_text.count("") if not self._text_block_open: events = self._open_text_block() else: @@ -374,6 +400,7 @@ class AnthropicStreamEmitter: events = [] if self._text_block_open: + events.extend(self._close_open_think()) events.append(self._close_block()) # Defensive: close a stale open tool_use block before starting another. elif self._open_tool_call_id is not None: @@ -452,6 +479,7 @@ class AnthropicStreamEmitter: events.extend(self._open_text_block()) # Reset text tracking for the next synthesis turn self._prev_text = "" + self._open_think_tags = 0 return events def _open_text_block(self) -> list[str]: diff --git a/studio/backend/core/inference/chat_eos.py b/studio/backend/core/inference/chat_eos.py new file mode 100644 index 0000000000..2a5d0db228 --- /dev/null +++ b/studio/backend/core/inference/chat_eos.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Resolve a chat model's assistant-turn-end stop tokens. + +Some checkpoints set eos_token_id to a bare document terminator (Qwen3.5 ships +config eos ``<|endoftext|>`` though chat turns end with ``<|im_end|>``, and its +small chat variants ship no generation_config), so generation runs past the turn +and loops -- re-emitting tool calls or hallucinating ``<|im_start|>`` turns. + +Turn-end markers are derived from the tokenizer's ``chat_template`` (the tokens it +actually uses to end a turn), not raw vocab membership: a base/coder model can +carry ChatML control tokens in a shared vocab without using them, and a loader +may have synced ``eos_token`` to the document terminator. Dependency-light (no +torch / unsloth) so it is unit-testable without the full inference stack. +""" + +from typing import Optional + +# Canonical assistant-turn-end markers per chat family. +_CHAT_TURN_END_TOKENS = ( + "<|im_end|>", # ChatML: Qwen, Yi + "<|eot_id|>", # Llama 3.x + "<|eom_id|>", # Llama 3.x tool turns + "", # Gemma + "", # Gemma-4 + "<|end|>", # Phi + "<|end_of_turn|>", # OpenChat / Starling (barred, distinct from Gemma's) +) +# harmony/gpt-oss uses <|end|> as a channel delimiter, not the turn end, and has +# its own streamer, so its eos is left untouched. +_HARMONY_MARKERS = ("<|channel|>", "<|constrain|>") + + +def _eos_id_set(eos_token_id) -> set: + if isinstance(eos_token_id, (list, tuple)): + return {int(t) for t in eos_token_id if t is not None} + if eos_token_id is not None: + return {int(eos_token_id)} + return set() + + +def _collect_template_text(chat_template) -> str: + """Flatten a tokenizer ``chat_template`` into one scannable string. + + Usually the template is a single jinja string, but multi-variant models + (e.g. Hermes-3: a ``default`` plus a ``tool_use`` template) expose it as a + ``{name: template}`` dict -- or, as stored in tokenizer_config.json, a list + of ``{"name": ..., "template": ...}`` dicts. Scanning only the ``str`` case + would skip turn-end detection for those valid models, so gather every string + leaf (variant names are harmless: they never contain the markers). + """ + if isinstance(chat_template, str): + return chat_template + if isinstance(chat_template, dict): + values = chat_template.values() + elif isinstance(chat_template, (list, tuple)): + values = chat_template + else: + return "" + parts = [_collect_template_text(v) for v in values] + return "\n".join(p for p in parts if p) + + +def resolve_chat_turn_end_eos_ids_using(template_tokenizer, id_tokenizer) -> list: + """eos of ``id_tokenizer`` plus any canonical turn-end marker the + ``template_tokenizer``'s chat_template uses, resolved to ids on ``id_tokenizer`` -- + the tokenizer generation actually uses. + + Pass the same tokenizer for both at load time. After a mapped ``get_chat_template`` + pass the MAPPED tokenizer as ``template_tokenizer`` (it carries the effective + template) and the ORIGINAL generation tokenizer as ``id_tokenizer``: a mapped + template registered ``map_eos_token=True`` can hand back a tokenizer whose vocab + folds the turn-end token onto the doc-eos id, and generate_stream re-reads the + original tokenizer, so resolving ids on the mapped tokenizer would store the wrong + (doc-eos) id and let generation run past the real turn marker.""" + ids = _eos_id_set(getattr(id_tokenizer, "eos_token_id", None)) + template = _collect_template_text(getattr(template_tokenizer, "chat_template", None)) + if not template or any(h in template for h in _HARMONY_MARKERS): + return sorted(ids) + unk = getattr(id_tokenizer, "unk_token_id", None) + for marker in _CHAT_TURN_END_TOKENS: + if marker in template: + try: + tid = id_tokenizer.convert_tokens_to_ids(marker) + except Exception: + tid = None + if tid is not None and tid != unk and int(tid) >= 0: + ids.add(int(tid)) + return sorted(ids) + + +def resolve_chat_turn_end_eos_ids(tokenizer) -> list: + """tokenizer.eos plus any canonical turn-end marker the model's chat_template + actually uses. Cheap (convert_tokens_to_ids per marker, no get_vocab); intended + to be resolved once at load. Returns eos unchanged for harmony templates.""" + return resolve_chat_turn_end_eos_ids_using(tokenizer, tokenizer) + + +def chat_eos_repair(current_eos, turn_end_ids) -> Optional[list]: + """Merged eos_token_id list, or None if ``current_eos`` already covers every + resolved turn-end id. Used to repair a model's generation_config at load so + every ``.generate()`` path (vision, tool loops) stops at the turn boundary.""" + if not turn_end_ids: + return None + current_set = _eos_id_set(current_eos) + if set(turn_end_ids) <= current_set: + return None + return sorted(current_set | set(turn_end_ids)) diff --git a/studio/backend/core/inference/chat_template_helpers.py b/studio/backend/core/inference/chat_template_helpers.py index b85e9c348a..897db8262d 100644 --- a/studio/backend/core/inference/chat_template_helpers.py +++ b/studio/backend/core/inference/chat_template_helpers.py @@ -3,11 +3,96 @@ """ Dependency-light wrapper around tokenizer.apply_chat_template with a kwarg -fallback for templates that reject reasoning/tools args. +fallback for templates that reject reasoning/tools args, plus the shared +native-chat-template fallback used by the transformers and MLX backends. """ +import copy +import json +import logging from typing import Optional +_THINK_OPEN = "" +_THINK_CLOSE = "" + + +def detect_think_prefill(prompt: Optional[str], special_tokens = None) -> str: + """Return the trailing open ```` prefill of a rendered prompt. + + Reasoning templates (Qwen3.6, DeepSeek-R1-style) end the generation + prompt with ``\\n`` so the model starts reasoning immediately. + Because that opening tag is part of the *prompt*, skip_prompt streaming + never emits it, and the frontend's ````/```` parser shows + the reasoning as plain text instead of a thinking block. (The GGUF path + is unaffected: llama-server's reasoning parser returns + ``reasoning_content``, which gets re-wrapped in think tags.) + + Returns the exact prompt tail to re-emit at the start of the generated + stream (e.g. ``"\\n"``), or ``""`` when the prompt does not end + with an open think block, including the ``enable_thinking=False`` case + where templates prefill an already-closed ``\\n\\n``. + + ``special_tokens`` is the tokenizer's special-token list. If ```` + is one, the streamer's skip_special_tokens strips the model's closing tag, + so re-emitting the open would leave an unclosed block that swallows the + answer. In that case return ``""`` and fall back to plain text. + """ + if not prompt: + return "" + open_idx = prompt.rfind(_THINK_OPEN) + if open_idx == -1: + return "" + tail = prompt[open_idx:] + if _THINK_CLOSE in tail or tail.strip() != _THINK_OPEN: + return "" + if special_tokens and _THINK_CLOSE in set(special_tokens): + return "" + return tail + + +logger = logging.getLogger(__name__) + + +def _normalize_tool_call_arguments(messages: list) -> list: + """Coerce each assistant ``tool_calls[].function.arguments`` from a JSON + string to a dict. + + The OpenAI wire format carries ``arguments`` as a JSON string, but some chat + templates (e.g. the stricter Qwen tool templates shipped with mlx-community + checkpoints) iterate ``arguments.items()`` and raise + ``TypeError: Can only get item pairs from a mapping.`` on the string form + when a prior tool call is re-rendered on the next turn. A dict works on both + strict and lenient templates, so parse the string; leave non-JSON or non-dict + values untouched. Returns the original list unchanged when nothing needed + coercing (no copy).""" + mutated = False + out: list = [] + for msg in messages: + tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else None + if not tool_calls: + out.append(msg) + continue + new_calls = [] + msg_changed = False + for call in tool_calls: + fn = call.get("function") if isinstance(call, dict) else None + args = fn.get("arguments") if isinstance(fn, dict) else None + if isinstance(args, str): + try: + parsed = json.loads(args) + except (ValueError, TypeError): + parsed = None + if isinstance(parsed, dict): + call = {**call, "function": {**fn, "arguments": parsed}} + msg_changed = True + new_calls.append(call) + if msg_changed: + out.append({**msg, "tool_calls": new_calls}) + mutated = True + else: + out.append(msg) + return out if mutated else messages + def apply_chat_template_for_generation( tokenizer, @@ -38,21 +123,209 @@ def apply_chat_template_for_generation( attempts.append(dict(reasoning_kwargs)) attempts.append({}) - last_exc: Optional[Exception] = None - for kwargs in attempts: + def _render(msgs: list) -> str: + last_exc: Optional[Exception] = None + for kwargs in attempts: + try: + return tokenizer.apply_chat_template( + msgs, + tokenize = False, + add_generation_prompt = True, + **kwargs, + ) + except TypeError as e: + last_exc = e + continue + except Exception as e: + last_exc = e + break + if last_exc is not None: + raise last_exc + raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result") + + try: + return _render(messages) + except Exception: + # Strict tool templates reject the JSON-string ``arguments`` form via + # TypeError or a broad Jinja raise_exception, so retry with dicts coerced. + # Original messages render first, so working templates stay byte-identical. + normalized = _normalize_tool_call_arguments(messages) + if normalized is messages: + raise + return _render(normalized) + + +def render_native_template( + *, + model_info: dict, + active_model_name: Optional[str], + messages: list, + tools: list, + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + preserve_thinking: Optional[bool] = None, + apply_fn = None, + hf_token: Optional[str] = None, +) -> Optional[str]: + """Render ``messages`` + ``tools`` with the model's NATIVE chat template. + + Some Unsloth override templates (e.g. ``mistral``, ``gemma-4``) do not emit + the ``tools`` schema, so a tool-calling turn silently stops advertising tools. + The native template ships in the model repo and carries the family's + tool-calling syntax. It is loaded straight from the repo (bypassing any + override on the live tokenizer) and cached on ``model_info``. Returns the + rendered prompt only if the native template actually emits the tools (render + differs with vs without tools); otherwise ``None``. + + ``hf_token`` is the token the model was loaded with -- passed to the repo load + so a gated/private model's native template can still be fetched (otherwise the + fallback fails silently and keeps the override prompt that dropped tools). + + ``trust_remote_code`` is sourced from ``model_info`` (the value the model was + actually loaded with) rather than a call-site argument, so the native-template + reload uses exactly the consent already granted at load. A custom-code tokenizer + repo raises in ``AutoTokenizer.from_pretrained`` unless ``trust_remote_code`` is + passed, so without this the fallback fails silently and keeps the tool-dropping + prompt for a model the user already consented to run remote code for. For a LoRA + adapter the reload targets the base model, whose remote code was gated and loaded + under the same stored flag, so re-passing it executes no unconsented code. + """ + # ``apply_fn`` lets a backend inject its own render; defaults to the module helper. + if apply_fn is None: + apply_fn = apply_chat_template_for_generation + native_tpl = model_info.get("native_chat_template") + if native_tpl is None: + # A LoRA adapter's native template lives on the base model, not the adapter id. + template_source = model_info.get("base_model") or active_model_name + # Re-use the load-time trust_remote_code so a custom-code tokenizer repo can + # instantiate its class (the stored flag already covers template_source). + trust_remote_code = bool(model_info.get("trust_remote_code", False)) try: - return tokenizer.apply_chat_template( - messages, - tokenize = False, - add_generation_prompt = True, - **kwargs, + from transformers import AutoTokenizer + nt = AutoTokenizer.from_pretrained( + template_source, + token = hf_token if hf_token and hf_token.strip() else None, + trust_remote_code = trust_remote_code, ) - except TypeError as e: - last_exc = e - continue - except Exception as e: - last_exc = e - break - if last_exc is not None: - raise last_exc - raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result") + native_tpl = nt.chat_template or False + except Exception as exc: + logger.warning( + "Could not load native chat template for '%s': %s", + template_source, + exc, + ) + # A failed fetch is not "no template": leave the sentinel unset so the next + # call retries (caching False would pin the tool-dropping override). + return None + model_info["native_chat_template"] = native_tpl + if not native_tpl: + return None + + tokenizer = model_info.get("tokenizer") or model_info.get("processor") + if tokenizer is None: + return None + tokenizer = getattr(tokenizer, "tokenizer", tokenizer) + # Render on a shallow copy: mutating the shared tokenizer.chat_template (outside the + # generation lock) races concurrent requests. + try: + render_tokenizer = copy.copy(tokenizer) + render_tokenizer.chat_template = native_tpl + except Exception as exc: + logger.warning( + "Could not clone tokenizer for native-template render of '%s': %s", + active_model_name, + exc, + ) + return None + try: + with_tools = apply_fn( + render_tokenizer, + messages, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + ) + no_tools = apply_fn( + render_tokenizer, + messages, + tools = None, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + ) + except Exception as exc: + logger.warning( + "Native-template tool render failed for '%s': %s", + active_model_name, + exc, + ) + return None + return with_tools if with_tools != no_tools else None + + +def render_with_native_template_fallback( + *, + formatted_prompt: str, + tokenizer, + model_info: dict, + active_model_name: Optional[str], + messages: list, + tools: Optional[list], + enable_thinking: Optional[bool] = None, + reasoning_effort: Optional[str] = None, + preserve_thinking: Optional[bool] = None, + apply_fn = None, + hf_token: Optional[str] = None, +) -> str: + """Return ``formatted_prompt``, swapping in a native-template render when an + override template dropped the ``tools`` schema. + + If ``tools`` were requested but the live render is identical with and without + them (detected by comparison, robust against tool names in the system prompt), + re-render with the model's native template. Shared by the transformers and MLX + backends so both advertise tools consistently. ``hf_token`` is forwarded so a + gated/private model's native template can still be fetched.""" + if not tools: + return formatted_prompt + if apply_fn is None: + apply_fn = apply_chat_template_for_generation + # Probe whether the live template dropped the schema. A tools-requiring template + # can raise here; on any error keep the valid tools prompt rather than lose it. + try: + probe_no_tools = apply_fn( + tokenizer, + messages, + tools = None, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + ) + except Exception as exc: + logger.warning( + "No-tools probe failed for '%s'; keeping the existing tools prompt: %s", + active_model_name, + exc, + ) + return formatted_prompt + if formatted_prompt != probe_no_tools: + return formatted_prompt # template already emits the tools schema + native_prompt = render_native_template( + model_info = model_info, + active_model_name = active_model_name, + messages = messages, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + apply_fn = apply_fn, + hf_token = hf_token, + ) + if native_prompt: + logger.info( + "Override template for '%s' dropped tool schemas; using the model's " + "native template for this tool-calling turn.", + active_model_name, + ) + return native_prompt + return formatted_prompt diff --git a/studio/backend/core/inference/defaults.py b/studio/backend/core/inference/defaults.py index b64605e16f..a1d03c03e0 100644 --- a/studio/backend/core/inference/defaults.py +++ b/studio/backend/core/inference/defaults.py @@ -8,6 +8,7 @@ import utils.hardware.hardware as hw DEFAULT_MODELS_GGUF = [ "unsloth/Qwen3.6-27B-MTP-GGUF", "unsloth/Qwen3.6-35B-A3B-MTP-GGUF", + "unsloth/DeepSeek-V4-Flash-GGUF", "unsloth/gemma-4-E2B-it-GGUF", "unsloth/gemma-4-E4B-it-GGUF", "unsloth/gemma-4-31B-it-GGUF", @@ -27,6 +28,7 @@ DEFAULT_MODELS_GGUF = [ DEFAULT_MODELS_STANDARD = [ "unsloth/Qwen3.6-27B-MTP-GGUF", "unsloth/Qwen3.6-35B-A3B-MTP-GGUF", + "unsloth/DeepSeek-V4-Flash-GGUF", "unsloth/gemma-4-E2B-it-GGUF", "unsloth/gemma-4-E4B-it-GGUF", "unsloth/gemma-4-31B-it-GGUF", diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 4dca4db768..7e69e05124 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -27,6 +27,11 @@ from utils.hardware import ( from core.inference.audio_codecs import AudioCodecManager from core.inference.runtime_context import runtime_context_length from core.inference.message_content import content_to_text +from core.inference.chat_eos import ( + chat_eos_repair, + resolve_chat_turn_end_eos_ids_using, +) +from core.inference.presence_penalty import _make_presence_penalty_processor from io import StringIO import structlog from loggers import get_logger @@ -210,6 +215,50 @@ class InferenceBackend: # API uses -1 to disable top-k; transformers uses 0. return 0 if top_k < 0 else top_k + def _resolve_chat_eos(self, model_name: str) -> None: + """Resolve this chat model's assistant-turn-end stop tokens once at load, + cache them in model_info, and repair generation_config so every + ``.generate()`` path stops at the turn boundary. + + Some checkpoints (e.g. Qwen3.5 / Qwen3.6 small chat models) end turns with + ``<|im_end|>`` but ship ``config.eos_token_id = <|endoftext|>`` and no + ``generation_config.json``, so paths that read ``generation_config`` (the + vision path, tool loops) run past the turn and loop. Turn-end markers are + derived from the chat_template (see chat_eos.resolve_chat_turn_end_eos_ids), + so base/coder models and harmony templates are left untouched. + """ + info = self.models.get(model_name) or {} + model = info.get("model") + container = info.get("tokenizer") + tokenizer = getattr(container, "tokenizer", container) # unwrap processors + if model is None or tokenizer is None: + return + # Vision models carry the chat_template on the processor, not the inner + # tokenizer. Read markers from whichever has one, but resolve ids on the + # generation tokenizer, else the vision path misses the turn-end token. + template_source = container if getattr(container, "chat_template", None) else tokenizer + try: + turn_end_ids = resolve_chat_turn_end_eos_ids_using(template_source, tokenizer) + except Exception as e: # never block a load on eos resolution + logger.warning("Chat turn-end eos resolution failed for %s: %s", model_name, e) + return + info["chat_turn_end_eos_ids"] = turn_end_ids + + gen = getattr(model, "generation_config", None) + if gen is None: + return + repaired = chat_eos_repair(gen.eos_token_id, turn_end_ids) + if repaired is None: + return + previous = gen.eos_token_id + gen.eos_token_id = repaired + logger.info( + "Repaired generation_config.eos_token_id for %s: %s -> %s", + model_name, + previous, + repaired, + ) + def load_model( self, config: ModelConfig, @@ -221,6 +270,9 @@ class InferenceBackend: gpu_ids: Optional[list[int]] = None, ) -> bool: """Load any model: base, LoRA adapter, text, or vision.""" + # Keep the token so the native-template fallback can fetch a + # gated model's repo template later during generation. + self._hf_token = hf_token # GGUF uses max_seq_length=0 as "model default"; Unsloth crashes on it. if max_seq_length <= 0: max_seq_length = 2048 @@ -231,6 +283,8 @@ class InferenceBackend: # Already loaded? if model_name in self.models and self.models[model_name].get("model"): logger.info(f"Model {model_name} already loaded") + if hf_token: + self.models[model_name]["hf_token"] = hf_token self.active_model_name = model_name return True @@ -246,6 +300,14 @@ class InferenceBackend: ) self.models[model_name] = { + # Per-model token: the native-template fallback must use the + # token this model was loaded with, not whichever loaded last. + "hf_token": hf_token, + # Per-model consent: the native-template reload must re-use the + # exact trust_remote_code this model (and a LoRA's base) was loaded + # with, so a custom-code tokenizer repo can be re-fetched without + # executing any code the user did not already consent to. + "trust_remote_code": trust_remote_code, "is_vision": config.is_vision, "is_lora": config.is_lora, "is_audio": config.is_audio, @@ -496,6 +558,7 @@ class InferenceBackend: max_seq_length, ) + self._resolve_chat_eos(model_name) self._load_chat_template_info(model_name) self.active_model_name = model_name @@ -766,9 +829,11 @@ class InferenceBackend: preserve_thinking: Optional[bool] = None, max_tool_iterations: int = 25, auto_heal_tool_calls: bool = True, + nudge_tool_calls: Optional[bool] = None, tool_call_timeout: int = 300, session_id: Optional[str] = None, rag_scope: Optional[dict] = None, + presence_penalty: float = 0.0, ): """Run an agentic tool loop on top of ``generate_chat_response``. @@ -802,6 +867,7 @@ class InferenceBackend: enable_thinking = enable_thinking, reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, + presence_penalty = presence_penalty, ) initial = list(messages) @@ -815,6 +881,7 @@ class InferenceBackend: execute_tool = execute_tool, cancel_event = cancel_event, auto_heal_tool_calls = auto_heal_tool_calls, + nudge_tool_calls = nudge_tool_calls, max_tool_iterations = max_tool_iterations, tool_call_timeout = tool_call_timeout, session_id = session_id, @@ -837,12 +904,14 @@ class InferenceBackend: enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Generate response for text or vision models (lock held by background thread). ``tools`` / ``enable_thinking`` / ``reasoning_effort`` / ``preserve_thinking`` are forwarded into ``apply_chat_template`` so templates that understand them (Qwen3, Llama 3.1+, gpt-oss harmony) advertise tool schemas / reasoning controls. + ``presence_penalty`` matches the GGUF sampling path (0 disables it). """ yield from self._generate_chat_response_inner( messages = messages, @@ -859,6 +928,7 @@ class InferenceBackend: enable_thinking = enable_thinking, reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, + presence_penalty = presence_penalty, ) def _generate_chat_response_inner( @@ -878,6 +948,7 @@ class InferenceBackend: enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Inner generation logic, called by generate_chat_response and generate_with_adapter_control. @@ -917,6 +988,7 @@ class InferenceBackend: max_new_tokens, repetition_penalty, cancel_event = cancel_event, + presence_penalty = presence_penalty, ) return else: @@ -946,6 +1018,22 @@ class InferenceBackend: tokenizer, chat_template = template_name, ) + # The mapper installs the effective template only now, at generate + # time, so re-resolve and UNION into the load-time cache (never + # overwrite). get_chat_template can return a remapped tokenizer + # (turn-end folded onto doc-eos) while generate_stream reads the + # original, so take marker strings from the mapped template but + # resolve their ids on the original. + try: + _gen_tok = model_info.get("tokenizer") or tokenizer + refreshed = resolve_chat_turn_end_eos_ids_using( + getattr(tokenizer, "tokenizer", tokenizer), + getattr(_gen_tok, "tokenizer", _gen_tok), + ) + existing = model_info.get("chat_turn_end_eos_ids") or [] + model_info["chat_turn_end_eos_ids"] = sorted(set(existing) | set(refreshed)) + except Exception as e: + logger.warning(f"Could not refresh chat turn-end eos after template: {e}") else: logger.info( f"No registered Unsloth template for {self.active_model_name}, using tokenizer default" @@ -975,6 +1063,27 @@ class InferenceBackend: reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, ) + + # If tools were requested but the (possibly overridden) template ignored + # them, fall back to the model's native template (shared with MLX). + from core.inference.chat_template_helpers import ( + render_with_native_template_fallback, + ) + + formatted_prompt = render_with_native_template_fallback( + formatted_prompt = formatted_prompt, + tokenizer = tokenizer, + model_info = model_info, + active_model_name = self.active_model_name, + messages = template_messages, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + apply_fn = self._apply_chat_template_for_generation, + hf_token = model_info.get("hf_token"), + ) + logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...") except Exception as e: logger.error(f"Error applying chat template: {e}") @@ -992,6 +1101,7 @@ class InferenceBackend: repetition_penalty, cancel_event = cancel_event, _adapter_state = _adapter_state, + presence_penalty = presence_penalty, ) def _generate_vision_response( @@ -1006,6 +1116,7 @@ class InferenceBackend: max_new_tokens, repetition_penalty, cancel_event = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Handle vision model generation with true token-by-token streaming.""" model_info = self.models[self.active_model_name] @@ -1067,13 +1178,22 @@ class InferenceBackend: add_special_tokens = False, return_tensors = "pt", ).to(model.device) + prompt_text = input_text else: # Text-only path for a vision model formatted_prompt = self.format_chat_prompt(messages, system_prompt) inputs = raw_tokenizer(formatted_prompt, return_tensors = "pt").to(model.device) + prompt_text = formatted_prompt # Stream with TextIteratorStreamer + background thread try: + from core.inference.chat_template_helpers import detect_think_prefill + + # Re-emit an open prefill swallowed by skip_prompt (see + # generate_stream). + think_prefix = detect_think_prefill( + prompt_text, getattr(raw_tokenizer, "all_special_tokens", None) + ) from transformers import TextIteratorStreamer import threading @@ -1095,6 +1215,14 @@ class InferenceBackend: top_k = top_k, min_p = min_p, ) + # Presence penalty (GGUF parity) for VLM chat. + _vision_input_ids = inputs.get("input_ids") if hasattr(inputs, "get") else None + if _vision_input_ids is not None: + _pp = _make_presence_penalty_processor( + presence_penalty, int(_vision_input_ids.shape[1]) + ) + if _pp is not None: + generation_kwargs["logits_processor"] = _pp err: dict[str, str] = {} @@ -1114,7 +1242,11 @@ class InferenceBackend: thread = threading.Thread(target = generate_fn) thread.start() - output = "" + output = think_prefix + # Emit the prefilled before the first token so the block + # renders during prompt prefill (which can take seconds). + if think_prefix: + yield think_prefix from queue import Empty generation_complete = False @@ -1323,11 +1455,13 @@ class InferenceBackend: repetition_penalty: float = 1.0, cancel_event = None, _adapter_state = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Generate a streaming text response (text models only). _adapter_state: if not None, the background thread toggles adapters before model.generate(), under _generation_lock. + ``presence_penalty`` matches the GGUF sampling path via a logits processor (0 disables it). """ if not self.active_model_name: yield "Error: No active model" @@ -1346,6 +1480,16 @@ class InferenceBackend: from transformers import TextIteratorStreamer import threading + from core.inference.chat_template_helpers import detect_think_prefill + + # skip_prompt swallows an open prefilled by the template; + # re-emit it so the frontend can render the thinking block. + # gpt-oss emits its own tags via HarmonyTextStreamer. + think_prefix = ( + "" + if self._is_gpt_oss_model() + else detect_think_prefill(prompt, getattr(tokenizer, "all_special_tokens", None)) + ) # gpt-oss models: HarmonyTextStreamer parses the multi-channel # harmony protocol into tags @@ -1382,11 +1526,18 @@ class InferenceBackend: min_p = min_p, repetition_penalty = repetition_penalty, do_sample = temperature > 0, - eos_token_id = tokenizer.eos_token_id, + # Resolved once at load (chat_template-derived turn-end tokens). + eos_token_id = model_info.get("chat_turn_end_eos_ids") or tokenizer.eos_token_id, pad_token_id = tokenizer.eos_token_id if tokenizer.pad_token_id is None else tokenizer.pad_token_id, ) + # Presence penalty (GGUF parity); prompt_len excludes prompt tokens. + _pp = _make_presence_penalty_processor( + presence_penalty, int(inputs["input_ids"].shape[1]) + ) + if _pp is not None: + generation_kwargs["logits_processor"] = _pp if cancel_event is not None: from transformers.generation.stopping_criteria import ( StoppingCriteria, @@ -1422,7 +1573,11 @@ class InferenceBackend: thread = threading.Thread(target = generate_fn) thread.start() - output = "" + output = think_prefix + # Emit the prefilled before the first token so the block + # renders during prompt prefill (which can take seconds). + if think_prefix: + yield think_prefix from queue import Empty generation_complete = False diff --git a/studio/backend/core/inference/llama_admission.py b/studio/backend/core/inference/llama_admission.py new file mode 100644 index 0000000000..b6a939c87b --- /dev/null +++ b/studio/backend/core/inference/llama_admission.py @@ -0,0 +1,368 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Admission control for local llama-server generation requests. + +The helpers in this module deliberately know nothing about FastAPI, SSE, or the +OpenAI-compatible route shape. They only coordinate how many upstream generation +requests may be active for one llama-server backend and provide a cancellable +FIFO queue for excess requests. +""" + +from __future__ import annotations + +import asyncio +import os +import threading +from collections import deque +from dataclasses import dataclass +from typing import Deque, Optional + + +ADMISSION_CONTROL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL" +ADMISSION_QUEUE_TIMEOUT_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_QUEUE_TIMEOUT" +ADMISSION_KEEPALIVE_INTERVAL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_KEEPALIVE_INTERVAL" +ADMISSION_MAX_QUEUE_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE" + +DEFAULT_ADMISSION_ENABLED = True +DEFAULT_ADMISSION_QUEUE_TIMEOUT_S = None +DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S = 5.0 +DEFAULT_ADMISSION_MAX_QUEUE = 64 + + +@dataclass(frozen = True) +class LlamaAdmissionConfig: + enabled: bool = DEFAULT_ADMISSION_ENABLED + queue_timeout_s: Optional[float] = DEFAULT_ADMISSION_QUEUE_TIMEOUT_S + keepalive_interval_s: float = DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S + max_queue: Optional[int] = DEFAULT_ADMISSION_MAX_QUEUE + + +@dataclass(frozen = True) +class LlamaAdmissionSnapshot: + key: str + capacity: int + active: int + queued: int + + +class LlamaAdmissionError(Exception): + def __init__( + self, + message: str, + *, + snapshot: Optional[LlamaAdmissionSnapshot] = None, + ): + super().__init__(message) + self.snapshot = snapshot + + +class LlamaAdmissionQueueFull(LlamaAdmissionError): + pass + + +class LlamaAdmissionTimeout(LlamaAdmissionError): + pass + + +class LlamaAdmissionCancelled(LlamaAdmissionError): + pass + + +def _bool_env(name: str, default: bool) -> bool: + value = os.environ.get(name) + if value is None or not value.strip(): + return default + value = value.strip().lower() + if value in {"1", "true", "yes", "on"}: + return True + if value in {"0", "false", "no", "off"}: + return False + return default + + +def _optional_positive_float_env(name: str, default: Optional[float]) -> Optional[float]: + value = os.environ.get(name) + if value is None or not value.strip(): + return default + try: + parsed = float(value.strip()) + except ValueError: + return default + return parsed if parsed > 0 else None + + +def _positive_float_env(name: str, default: float) -> float: + value = os.environ.get(name) + if value is None or not value.strip(): + return default + try: + parsed = float(value.strip()) + except ValueError: + return default + return parsed if parsed > 0 else default + + +def _optional_positive_int_env(name: str, default: Optional[int]) -> Optional[int]: + value = os.environ.get(name) + if value is None or not value.strip(): + return default + try: + parsed = int(value.strip()) + except ValueError: + return default + return parsed if parsed > 0 else None + + +def llama_admission_config_from_env() -> LlamaAdmissionConfig: + return LlamaAdmissionConfig( + enabled = _bool_env(ADMISSION_CONTROL_ENV, DEFAULT_ADMISSION_ENABLED), + queue_timeout_s = _optional_positive_float_env( + ADMISSION_QUEUE_TIMEOUT_ENV, + DEFAULT_ADMISSION_QUEUE_TIMEOUT_S, + ), + keepalive_interval_s = _positive_float_env( + ADMISSION_KEEPALIVE_INTERVAL_ENV, + DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S, + ), + max_queue = _optional_positive_int_env( + ADMISSION_MAX_QUEUE_ENV, + DEFAULT_ADMISSION_MAX_QUEUE, + ), + ) + + +@dataclass +class _Waiter: + loop: asyncio.AbstractEventLoop + future: asyncio.Future + cancelled: bool = False + granted_lease: Optional["LlamaAdmissionLease"] = None + + +class LlamaAdmissionLease: + def __init__(self, queue: Optional["LlamaAdmissionQueue"]): + self._queue = queue + self._released = False + self._release_lock = threading.Lock() + + def release(self) -> None: + queue = None + with self._release_lock: + if self._released: + return + self._released = True + queue = self._queue + if queue is not None: + queue.release() + + async def __aenter__(self) -> "LlamaAdmissionLease": + return self + + async def __aexit__(self, *_args) -> None: + self.release() + + +class LlamaAdmissionReservation: + def __init__( + self, + *, + queue: Optional["LlamaAdmissionQueue"], + lease: Optional[LlamaAdmissionLease] = None, + waiter: Optional[_Waiter] = None, + snapshot: Optional[LlamaAdmissionSnapshot] = None, + ): + self._queue = queue + self._lease = lease + self._waiter = waiter + self.snapshot = snapshot + + @property + def is_cancelled(self) -> bool: + return self._lease is None and self._waiter is None + + def lease_nowait(self) -> Optional[LlamaAdmissionLease]: + if self._lease is not None: + return self._lease + if self._waiter is None or not self._waiter.future.done(): + return None + if self._waiter.future.cancelled(): + self._waiter.cancelled = True + self._waiter = None + return None + self._lease = self._waiter.future.result() + self._waiter = None + return self._lease + + async def wait(self, timeout_s: float) -> Optional[LlamaAdmissionLease]: + lease = self.lease_nowait() + if lease is not None: + return lease + if self._waiter is None: + return None + waiter = self._waiter + try: + await asyncio.wait_for(asyncio.shield(waiter.future), timeout = timeout_s) + except asyncio.CancelledError: + if waiter.future.cancelled(): + waiter.cancelled = True + if self._waiter is waiter: + self._waiter = None + return None + raise + return self.lease_nowait() + + def cancel(self) -> None: + lease = self.lease_nowait() + if lease is not None: + lease.release() + self._lease = None + return + if self._queue is not None and self._waiter is not None: + self._queue.cancel(self._waiter) + self._waiter = None + + def snapshot_now(self) -> Optional[LlamaAdmissionSnapshot]: + if self._queue is None: + return self.snapshot + return self._queue.snapshot() + + +class LlamaAdmissionQueue: + def __init__(self, key: str): + self.key = key + self._lock = threading.Lock() + self._active = 0 + self._capacity = 1 + self._waiters: Deque[_Waiter] = deque() + + def reserve(self, *, capacity: int, config: LlamaAdmissionConfig) -> LlamaAdmissionReservation: + capacity = max(1, int(capacity or 1)) + if not config.enabled: + return LlamaAdmissionReservation( + queue = None, + lease = LlamaAdmissionLease(None), + snapshot = LlamaAdmissionSnapshot(self.key, capacity, 0, 0), + ) + + loop = asyncio.get_running_loop() + with self._lock: + self._capacity = capacity + self._prune_waiters_locked() + self._grant_waiters_locked() + if self._active < self._capacity and not self._waiters: + self._active += 1 + return LlamaAdmissionReservation( + queue = self, + lease = LlamaAdmissionLease(self), + snapshot = self._snapshot_locked(), + ) + if config.max_queue is not None and len(self._waiters) >= config.max_queue: + raise LlamaAdmissionQueueFull( + "llama-server generation queue is full", + snapshot = self._snapshot_locked(), + ) + waiter = _Waiter( + loop = loop, + future = loop.create_future(), + ) + self._waiters.append(waiter) + return LlamaAdmissionReservation( + queue = self, + waiter = waiter, + snapshot = self._snapshot_locked(), + ) + + def release(self) -> None: + with self._lock: + if self._active > 0: + self._active -= 1 + self._grant_waiters_locked() + + def cancel(self, waiter: _Waiter) -> None: + lease_to_release = None + with self._lock: + waiter.cancelled = True + try: + self._waiters.remove(waiter) + except ValueError: + pass + if waiter.granted_lease is not None: + lease_to_release = waiter.granted_lease + waiter.granted_lease = None + if not waiter.future.done(): + waiter.loop.call_soon_threadsafe(waiter.future.cancel) + if lease_to_release is not None: + lease_to_release.release() + + def snapshot(self) -> LlamaAdmissionSnapshot: + with self._lock: + self._prune_waiters_locked() + return self._snapshot_locked() + + def is_idle(self) -> bool: + with self._lock: + self._prune_waiters_locked() + return self._active == 0 and not self._waiters + + def _grant_waiters_locked(self) -> None: + self._prune_waiters_locked() + while self._waiters and self._active < self._capacity: + waiter = self._waiters.popleft() + if waiter.cancelled or waiter.future.done(): + continue + self._active += 1 + lease = LlamaAdmissionLease(self) + waiter.granted_lease = lease + waiter.loop.call_soon_threadsafe(self._deliver_lease, waiter, lease) + + def _deliver_lease(self, waiter: _Waiter, lease: LlamaAdmissionLease) -> None: + if waiter.cancelled or waiter.future.done(): + waiter.granted_lease = None + if not waiter.future.done(): + waiter.future.cancel() + lease.release() + return + try: + waiter.future.set_result(lease) + waiter.granted_lease = None + except asyncio.InvalidStateError: + waiter.granted_lease = None + lease.release() + + def _prune_waiters_locked(self) -> None: + self._waiters = deque( + waiter for waiter in self._waiters if not waiter.cancelled and not waiter.future.done() + ) + + def _snapshot_locked(self) -> LlamaAdmissionSnapshot: + return LlamaAdmissionSnapshot( + key = self.key, + capacity = self._capacity, + active = self._active, + queued = len(self._waiters), + ) + + +_QUEUES_LOCK = threading.Lock() +_QUEUES: dict[str, LlamaAdmissionQueue] = {} + + +def get_llama_admission_queue(key: str) -> LlamaAdmissionQueue: + with _QUEUES_LOCK: + queue = _QUEUES.get(key) + if queue is None: + queue = LlamaAdmissionQueue(key) + _QUEUES[key] = queue + # base_url carries a fresh ephemeral port on every model load, so + # each load registers a new key. Drop the now-idle queues from prior + # loads so the registry can't grow without bound on a long-running + # server. Queues with in-flight requests are kept until they drain. + for stale_key in [k for k in _QUEUES if k != key and _QUEUES[k].is_idle()]: + del _QUEUES[stale_key] + return queue + + +def reset_llama_admission_queues() -> None: + with _QUEUES_LOCK: + _QUEUES.clear() diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 3ccfc5cdfe..b06c6eb5cb 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -38,9 +38,35 @@ from core.inference.llama_server_args import ( strip_shadowing_flags, strip_split_mode_only, ) + +# Share strip / signal constants with the multi-format parser so BUFFERING also +# catches Llama-3 / Mistral / Gemma 4 (legacy helper only knew / "list[str]": return out -# ── Pre-compiled patterns for plan-without-action re-prompt ── -# Forward-looking intent signals: the model is describing what it *will* -# do rather than giving a final answer. -_INTENT_SIGNAL = re.compile( - r"(?i)(" - # Direct intent ("I'll ...", "Let me ...", straight + curly apostrophes). - # Excludes "I can"/"I should"/"I want to"/"let's" (common in answers). - # Negative lookahead drops negated forms ("I will not") so a refusal - # doesn't trigger a re-prompt. - r"\b(i['\u2019](ll|m going to|m gonna)|i am (going to|gonna)|i will|i shall|let me|allow me)\b(?!\s+(?:not|never)\b)" - r"|" - # Step/plan framing: "First ...", "Step 1:", "Here's my plan" - r"\b(?:first\b|step \d+:?|here['\u2019]?s (?:my |the |a )?(?:plan|approach))" - r"|" - # "Now I" / "Next I" patterns - r"\b(?:now i|next i)\b" - r")" -) -_MAX_REPROMPTS = 1 +# Plan-without-action re-prompt state (intent signal, caps, message) now lives +# in tool_call_parser, imported above under its old aliases. # Default max_tokens to the effective context when known. The floor is high # enough for reasoning-heavy GGUFs and max_tokens-omitting API clients. @@ -231,7 +244,10 @@ _DEFAULT_FIRST_TOKEN_TIMEOUT_S = 1200.0 # 20 min # is exempt because it needs immediate artifact feedback. _PROVISIONAL_ARGS_MIN_CHARS = 256 _DEFAULT_STREAM_STALL_TIMEOUT_S = 120.0 # 2 min -_REPROMPT_MAX_CHARS = 2000 +# Cap tool calls from a single TEXTUAL-fallback turn (mirrors the safetensors +# loop). Structured delta.tool_calls are grammar-bounded by llama-server; text +# parsed from content is not, so one runaway turn could fan out unbounded. +_MAX_TOOL_CALLS_PER_TURN = 8 _FORCED_REPEAT_PLAN_SIGNAL = re.compile( r"\b(?:i\s+will|i'll|let\s+me|going\s+to|need\s+to|call|use|run|search|fetch|render)\b", re.I, @@ -242,9 +258,70 @@ _FINAL_ANSWER_SIGNAL = re.compile( ) -def _is_short_intent_without_action(text: str) -> bool: - stripped = text.strip() - return 0 < len(stripped) < _REPROMPT_MAX_CHARS and _INTENT_SIGNAL.search(stripped) is not None +def _gguf_active_tool_names(active_tools: list[dict]) -> list[str]: + names = [ + (tool.get("function") or {}).get("name") + for tool in (active_tools or []) + if isinstance(tool, dict) and isinstance(tool.get("function"), dict) + ] + return [name for name in names if name] + + +# Rehearsal NAME chars (word + hyphen, matching the parser); the lookbehind excludes the +# Mistral [CALL_ID]...[ARGS] shape. +_GGUF_REHEARSAL_ARGS_RE = re.compile(r"(? int: + """Index of the first ``NAME[ARGS]`` whose NAME is an active tool, else -1. A + bare/inactive-name ``foo[ARGS]`` in prose is not a call; mirrors the safetensors + ``_earliest_tool_signal`` name-gating (no unrestricted GGUF mode).""" + active = set(_gguf_active_tool_names(active_tools)) + if not active: + return -1 + for m in _GGUF_REHEARSAL_ARGS_RE.finditer(text): + if m.group(1) in active: + return m.start() + return -1 + + +def _gguf_has_genuine_tool_signal(text: str, signals, active_tools: list[dict]) -> bool: + """True when ``text`` holds a genuine tool-call boundary for one of ``signals``. + + Unambiguous markers (````, ``[TOOL_CALLS]``, ``= 0: + return True + continue + if sig in text: + return True + return False + + +def _is_rehearsal_prefix(stripped: str, active_tools: list[dict]) -> bool: + """True if ``stripped`` is a (possibly partial) prefix of ``NAME[ARGS]`` for an + active tool -- the bare tool name arriving in its own chunk before ``[ARGS]{...}``. + Mirrors the safetensors loop so the split rehearsal call is not streamed.""" + if not stripped or any(ch.isspace() for ch in stripped): + return False + for name in _gguf_active_tool_names(active_tools): + if stripped == name or f"{name}[ARGS]".startswith(stripped): + return True + return False + + +def _held_rehearsal_tail_len(text: str, active_tools: list[dict]) -> int: + """Length of a trailing bare tool-name token that may be a split rehearsal call + (``...web_search`` with ``[ARGS]{...}`` still to arrive), so STREAMING can hold it + instead of leaking the name. Returns 0 for ordinary prose. Mirrors safetensors.""" + i = len(text) + while i > 0 and not text[i - 1].isspace(): + i -= 1 + tail = text[i:] + return len(tail) if tail and _is_rehearsal_prefix(tail, active_tools) else 0 def _should_suppress_forced_no_tool_output(text: str) -> bool: @@ -304,6 +381,19 @@ def _probe_dns_dead(host: str = "huggingface.co", timeout: float = 2.0) -> bool: return True if result[0] is None else result[0] +def _hf_env_offline() -> bool: + """True when an HF offline env var is set to any truthy value (1/true/yes/on). + + Mirrors utils.models.model_config._env_offline so a user-set HF_HUB_OFFLINE=true + (not just "1") still routes through the local-cache reuse path below. + """ + try: + from utils.models.model_config import _env_offline + return _env_offline() + except Exception: + return os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in {"1", "true", "yes", "on"} + + @contextlib.contextmanager def _hf_offline_if_dns_dead(): """Set HF_HUB_OFFLINE for this block only when DNS to huggingface.co fails; @@ -545,6 +635,13 @@ _TOOL_TEMPLATE_MARKERS = ( "'role' == 'tool'", 'message.role == "tool"', "message.role == 'tool'", + # DeepSeek: no top-level ``{% if tools %}`` block; it gates emission on + # ``message['role'] == 'tool'`` plus ``message['tool_calls'] is defined``. + "message['role'] == 'tool'", + 'message["role"] == "tool"', + "message['tool_calls']", + 'message["tool_calls"]', + "tool_calls is defined", ) @@ -605,6 +702,16 @@ def detect_reasoning_flags( else [] ) if effort_levels: + # DeepSeek-V4's encoder accepts reasoning_effort {'high', 'max'} but its + # template only branches on 'max', so the literal scan misses 'high'. Add it + # (matched on whole repo-name segments, so 'deepseek-v40' won't false-match) + # to expose the full none/high/max ladder instead of none/max. + segments = re.split(r"[-_.]", (model_identifier or "").lower().split("/")[-1]) + is_dsv4 = "deepseek4" in segments or any( + a == "deepseek" and b == "v4" for a, b in zip(segments, segments[1:]) + ) + if is_dsv4 and "high" not in effort_levels: + effort_levels = sorted(set(effort_levels) | {"high"}, key = _REASONING_EFFORT_SCALE.index) # GLM-5.2-style: an enable_thinking on/off gate PLUS a reasoning_effort # level among a discrete set (e.g. 'high' | 'max'). Distinct from # gpt-oss (reasoning_effort only, no on/off gate) and Qwen @@ -747,6 +854,112 @@ def _gguf_snapshot_files(snapshot: Path) -> list[str]: ] +def _cached_hf_snapshot_file( + repo_id: str, + filename: str, + *, + expected_size: Optional[int] = None, +) -> Optional[str]: + """Return a cached snapshot file even when HF's current-ref probe misses it.""" + if not filename: + return None + parts = [part for part in filename.replace("\\", "/").split("/") if part] + if not parts or any(part in (".", "..") for part in parts): + return None + try: + from utils.models.model_config import _iter_hf_cache_snapshots + for snap in _iter_hf_cache_snapshots(repo_id): + candidate = snap.joinpath(*parts) + if not candidate.is_file(): + continue + if expected_size: + try: + if candidate.stat().st_size < expected_size: + continue + except OSError: + continue + return str(candidate) + except Exception as e: + logger.debug("Snapshot cache lookup failed for %s/%s: %s", repo_id, filename, e) + return None + + +def _snapshot_has_all_shards( + main_path: str, main_filename: str, shards: Iterable[str], expected_sizes: dict[str, int] +) -> bool: + """True when every shard sits beside ``main_path`` in the same cache snapshot. + + llama.cpp loads a split GGUF by resolving its siblings from the main shard's + directory, so a cached main shard is only safe to reuse when the rest of the + set is co-located; otherwise the caller must fetch the whole set together. + """ + root = Path(main_path) + for _ in [part for part in main_filename.replace("\\", "/").split("/") if part]: + root = root.parent + for shard in shards: + parts = [part for part in shard.replace("\\", "/").split("/") if part] + if not parts or any(part in (".", "..") for part in parts): + return False + sibling = root.joinpath(*parts) + try: + if not sibling.is_file(): + return False + expected = expected_sizes.get(shard) + if expected and sibling.stat().st_size < expected: + return False + except OSError: + return False + return True + + +def _resolve_repo_id_casing(hf_repo: str) -> str: + """Map a requested repo id to its cached canonical casing, or return it unchanged. + + A case-variant request (for example a lowercased id) resolves to the + canonical-cased cache directory so the main GGUF and its companions + (mmproj / MTP drafter) all read the same cache entry. Returns ``hf_repo`` + unchanged when resolution is unavailable or errors. + """ + try: + from utils.paths import resolve_cached_repo_id_case + return resolve_cached_repo_id_case(hf_repo) + except Exception: + return hf_repo + + +def _cached_colocated_split_main( + repo_id: str, main_filename: str, shards: Iterable[str], expected_sizes: dict[str, int] +) -> Optional[str]: + """Main-shard path from a cache snapshot that also holds every sibling shard. + + A newer snapshot may hold only the first shard while an older snapshot has the + complete split set. ``_cached_hf_snapshot_file`` would return that newer partial + main and the co-location check would then force a refetch, so scan snapshots for + one where the whole set is present and return that main path instead. None when + no snapshot holds the full set. + """ + main_parts = [part for part in main_filename.replace("\\", "/").split("/") if part] + if not main_parts or any(part in (".", "..") for part in main_parts): + return None + try: + from utils.models.model_config import _iter_hf_cache_snapshots + for snap in _iter_hf_cache_snapshots(repo_id): + main_path = snap.joinpath(*main_parts) + if not main_path.is_file(): + continue + expected_main = expected_sizes.get(main_filename) + try: + if expected_main and main_path.stat().st_size < expected_main: + continue + except OSError: + continue + if _snapshot_has_all_shards(str(main_path), main_filename, shards, expected_sizes): + return str(main_path) + except Exception as e: + logger.debug("Co-located split snapshot lookup failed for %s: %s", repo_id, e) + return None + + def _gguf_extra_shards(files: Iterable[str], first_shard: str) -> list[str]: m = _SHARD_FULL_RE.match(first_shard) if not m: @@ -1227,6 +1440,50 @@ def _backfill_usage_from_timings(usage, timings): return out +def _vulkan_lib_filename() -> str: + return "ggml-vulkan.dll" if sys.platform == "win32" else "libggml-vulkan.so" + + +# Host RAM to leave free on an integrated GPU, matching llama.cpp's own --fit +# margin (default 1024 MiB per device). ggml reports an iGPU's "VRAM" as shared +# system RAM, so hold back the same margin rather than inventing a larger one. +_IGPU_HOST_RESERVE_MIB = 1024 + + +def _apply_igpu_host_reserve_mib(free_mib: int, is_igpu: bool) -> int: + """Reserve host headroom on an integrated (shared-memory) Vulkan GPU. + + An iGPU's reported free "VRAM" is really free system RAM, so sizing + context/offload against all of it would push the host into swap or the OOM + killer. Leave the same margin llama.cpp's --fit uses. ``is_igpu`` comes from + ggml's device type, so a discrete card is never touched; only ever reduces. + """ + if not is_igpu: + return free_mib + return max(0, free_mib - _IGPU_HOST_RESERVE_MIB) + + +def _llama_lib_dir(binary: str) -> Path: + # The installer exposes llama-server as a top-level entrypoint into build/bin/, + # where the ggml backend libs live, so callers looking for sibling libs (Vulkan + # detection, LD_LIBRARY_PATH, probe bindir) need the real dir. It is normally a + # symlink (resolve() reaches build/bin), but create_exec_entrypoint falls back to + # a shell wrapper (exec "$(dirname "$0")/build/bin/llama-server" "$@") when it + # cannot symlink, and resolve() stops at the wrapper file. Follow the wrapper's + # exec target too, so a wrapper-based install still finds build/bin. + resolved = Path(binary).resolve() + try: + with open(resolved, "rb") as _f: + _head = _f.read(256) + if _head.startswith(b"#!"): + _m = re.search(r'exec "\$\(dirname "\$0"\)/([^"]+)"', _head.decode("utf-8", "ignore")) + if _m: + return (resolved.parent / _m.group(1)).resolve().parent + except OSError: + pass + return resolved.parent + + def _is_external_link(path: Path) -> bool: """True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink or a Windows directory junction / reparse point. Such a link resolves into @@ -1284,6 +1541,7 @@ class LlamaCppBackend: self._context_length: Optional[int] = None self._effective_context_length: Optional[int] = None self._max_context_length: Optional[int] = None + self._effective_parallel_slots: int = 1 self._chat_template: Optional[str] = None self._chat_template_override: Optional[str] = None self._supports_reasoning: bool = False @@ -1469,6 +1727,15 @@ class LlamaCppBackend: """Return the effective context length the server is running at.""" return self._effective_context_length or self._context_length + @property + def effective_parallel_slots(self) -> int: + """Return the serving-slot count the active llama-server actually uses.""" + try: + slots = int(getattr(self, "_effective_parallel_slots", 1)) + except (TypeError, ValueError): + slots = 1 + return max(1, slots) + @property def max_context_length(self) -> Optional[int]: """Return the largest context that fits on this hardware at load time. @@ -1485,6 +1752,16 @@ class LlamaCppBackend: """Return the model's native context length from GGUF metadata.""" return self._context_length + def _commit_effective_parallel_slots(self, n_parallel: int) -> None: + try: + slots = int(n_parallel) + except (TypeError, ValueError): + slots = 1 + self._effective_parallel_slots = max(1, slots) + + def _reset_effective_parallel_slots(self) -> None: + self._effective_parallel_slots = 1 + @staticmethod def _read_rss_bytes(pid: int) -> Optional[int]: """Resident set size of ``pid`` in bytes, from /proc//status (Linux). @@ -1660,9 +1937,13 @@ class LlamaCppBackend: # 'low' effort the way gpt-oss does (those models genuinely # cannot disable). thinking_off = enable_thinking is False or reasoning_effort == "none" - if enable_thinking is not None or reasoning_effort == "none": + # A named effort level implies thinking on, so emit enable_thinking + # even if the caller sent only reasoning_effort (else the template + # defaults it off and the requested level never renders). + effort_on = reasoning_effort in self._reasoning_effort_levels + if enable_thinking is not None or reasoning_effort == "none" or effort_on: kwargs["enable_thinking"] = not thinking_off - if not thinking_off and reasoning_effort in self._reasoning_effort_levels: + if not thinking_off and effort_on: kwargs["reasoning_effort"] = reasoning_effort elif self._reasoning_style == "reasoning_effort": if reasoning_effort in ("none", "low", "medium", "high"): @@ -1686,6 +1967,13 @@ class LlamaCppBackend: return False return self._supports_tools + @property + def supports_tool_passthrough(self) -> bool: + # supports_tools is forced off for DiffusionGemma (its agentic loop drops the + # per-step canvas frames), but client passthrough skips that loop, so it uses + # the real _supports_tools. + return self._supports_tools + @property def cache_type_kv(self) -> Optional[str]: return self._cache_type_kv @@ -2058,6 +2346,30 @@ class LlamaCppBackend: return total + @staticmethod + def _is_vulkan_backend(binary: Optional[str] = None) -> bool: + """True if the installed llama.cpp build is Vulkan-only. + + The official prebuilts are single-backend, so the Vulkan ggml lib next + to llama-server identifies a Vulkan build. Keeps the free-memory probe + and GPU pin in ggml's Vulkan device-index space. For a custom + multi-backend build with a CUDA or HIP ggml lib alongside Vulkan, defer + to that backend (torch-usable, better-understood probe/pin). + """ + binary = binary or LlamaCppBackend._find_llama_server_binary() + if not binary: + return False + lib_dir = _llama_lib_dir(binary) + if not (lib_dir / _vulkan_lib_filename()).is_file(): + return False + for _backend in ("cuda", "hip"): + sibling = ( + f"ggml-{_backend}.dll" if sys.platform == "win32" else f"libggml-{_backend}.so" + ) + if (lib_dir / sibling).is_file(): + return False + return True + @staticmethod def _resolve_visible_physical_ids() -> Optional[list[int]]: """Physical GPU ids behind the active visibility mask (HIP/ROCR/CUDA on @@ -2220,11 +2532,42 @@ class LlamaCppBackend: return True @staticmethod - def _get_gpu_free_memory() -> list[tuple[int, int]]: + def _visible_devices_mask(env_name: str) -> Optional[set[int]]: + """Physical indices a ``*_VISIBLE_DEVICES`` mask permits, or None if unset. + + ``if x.strip()`` filters trailing-comma masks ("0,1,"); an empty mask + ("") yields an empty set (all devices hidden), distinct from an unset + var (None, no mask). Used by the nvidia-smi probe. + """ + raw = os.environ.get(env_name) + if raw is None: + return None + try: + return set(int(x.strip()) for x in raw.split(",") if x.strip()) + except ValueError: + return None + + @staticmethod + def _vulkan_pin_args(gpu_indices: Optional[Iterable[int]]) -> list[str]: + """``--device Vulkan,...`` to pin a Vulkan launch to selected GPUs. + + The indices are ggml's compact Vulkan ordinals (as _get_gpu_free_memory + reports and the registry names ``Vulkan``). Pin by that name, NOT via + GGML_VK_VISIBLE_DEVICES: ggml parses that env var in the raw + vkEnumeratePhysicalDevices space (before dropping CPU/llvmpipe devices + and deduplicating ICDs), so a compact ordinal there could select a + different physical device or the CPU rasterizer. + """ + if not gpu_indices: + return [] + return ["--device", ",".join(f"Vulkan{i}" for i in gpu_indices)] + + @staticmethod + def _get_gpu_free_memory(binary: Optional[str] = None) -> list[tuple[int, int]]: """Query free memory per GPU. Returns ``(gpu_index, free_mib)`` sorted by index; empty if no supported GPU is reachable. Thin wrapper over ``_get_gpu_memory`` for callers that only need free VRAM.""" - return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory()] + return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory(binary)] @staticmethod def _apple_metal_memory_budget_bytes() -> int: @@ -2255,7 +2598,7 @@ class LlamaCppBackend: return int(rec_bytes * _APPLE_UNIFIED_MEMORY_FRACTION) @staticmethod - def _get_gpu_memory() -> list[tuple[int, int, int]]: + def _get_gpu_memory(binary: Optional[str] = None) -> list[tuple[int, int, int]]: """Query free AND total memory per GPU. Order: @@ -2267,9 +2610,18 @@ class LlamaCppBackend: probe returned [] on AMD) and NVIDIA hosts missing ``nvidia-smi`` from PATH. + On a Vulkan build the ggml Vulkan probe is authoritative, so the indices + are ggml's compact Vulkan ordinals (the space the pin selects via + ``--device Vulkan``). It reports ``total`` for discrete cards and 0 + for an iGPU (shared RAM) so the fit falls back to free*frac there. + Otherwise nvidia-smi / torch cover NVIDIA + AMD ROCm. + Returns (gpu_index, free_mib, total_mib) sorted by index; empty if no - supported GPU is reachable. ``total`` lets the fit reserve absolute headroom. + supported GPU is reachable. """ + binary = binary or LlamaCppBackend._find_llama_server_binary() + if LlamaCppBackend._is_vulkan_backend(binary): + return LlamaCppBackend._get_gpu_free_memory_vulkan(binary) # ── NVIDIA via nvidia-smi ──────────────────────────────────── try: result = subprocess.run( @@ -2285,16 +2637,7 @@ class LlamaCppBackend: **_windows_hidden_subprocess_kwargs(), ) if result.returncode == 0: - allowed: Optional[set[int]] = None - cvd = os.environ.get("CUDA_VISIBLE_DEVICES") - if cvd is not None: - try: - # `if x.strip()` filters trailing-comma masks ("0,1,"). - # Empty mask (CVD="") yields an empty set -> all GPUs - # filtered out, per codebase convention. - allowed = set(int(x.strip()) for x in cvd.split(",") if x.strip()) - except ValueError: - pass + allowed = LlamaCppBackend._visible_devices_mask("CUDA_VISIBLE_DEVICES") gpus: list[tuple[int, int, int]] = [] for line in result.stdout.strip().splitlines(): parts = [p.strip() for p in line.split(",")] @@ -2359,6 +2702,91 @@ class LlamaCppBackend: logger.debug(f"torch GPU probe failed: {e}") return [] + @staticmethod + def _get_gpu_free_memory_vulkan(binary: Optional[str] = None) -> list[tuple[int, int, int]]: + """Query free (and total) VRAM per device via the bundled ggml Vulkan backend. + + Loads ``libggml-vulkan`` in a short-lived subprocess (no Vulkan instance + in this process) and returns (device_index, free_mib, total_mib) sorted + by index. The index is ggml's compact Vulkan ordinal -- the one the + registry names ``Vulkan`` and load_model pins with ``--device``, + NOT the raw ``GGML_VK_VISIBLE_DEVICES`` space. A user-set + ``GGML_VK_VISIBLE_DEVICES`` is honored by ggml (passed through), so the + list already reflects it. iGPUs leave a host-RAM margin (see + ``_apply_igpu_host_reserve_mib``) and report total 0; discrete cards pass + their real total through. [] when no Vulkan build or device is reachable. + """ + binary = binary or LlamaCppBackend._find_llama_server_binary() + if not binary: + return [] + binary_dir = _llama_lib_dir(binary) + if not (binary_dir / _vulkan_lib_filename()).is_file(): + return [] + + env = child_env_without_native_path_secret() + # Pass any inherited GGML_VK_VISIBLE_DEVICES through to ggml unchanged so + # the probe enumerates the same device list the launch will, named + # Vulkan0..N in the compact order reported here and pinned by that name + # via --device -- probe, mask, and pin stay in one index space. Do NOT + # filter the mask in Python: ggml parses the env var in raw + # vkEnumeratePhysicalDevices space while this probe reports the compact + # post-filter ordinal, so a Python filter would compare mismatched spaces. + if sys.platform != "win32": + # Let the loader resolve sibling ggml libs next to the binary. + existing_ld = env.get("LD_LIBRARY_PATH", "") + env["LD_LIBRARY_PATH"] = ( + f"{binary_dir}:{existing_ld}" if existing_ld else str(binary_dir) + ) + probe_script = Path(__file__).with_name("_vulkan_probe.py") + try: + result = subprocess.run( + [sys.executable, str(probe_script), str(binary_dir)], + capture_output = True, + text = True, + timeout = 15, + env = env, + **_windows_hidden_subprocess_kwargs(), + ) + if result.returncode != 0: + logger.debug( + f"vulkan GPU probe exited {result.returncode}: {result.stderr.strip()}" + ) + return [] + except Exception as e: + logger.debug(f"vulkan GPU probe failed: {e}") + return [] + + gpus: list[tuple[int, int, int]] = [] + for line in result.stdout.strip().splitlines(): + parts = line.split("\t") + if len(parts) != 4: + continue + try: + idx = int(parts[0]) + free_mib = int(parts[1]) // (1024 * 1024) + is_igpu = parts[2] == "1" + # iGPU "total" is shared RAM, not a VRAM budget -> keep 0 so the + # fit stays on free*frac (the host reserve below is its + # headroom); a discrete card passes its real total through. + total_mib = 0 if is_igpu else int(parts[3]) // (1024 * 1024) + except ValueError: + continue + capped = _apply_igpu_host_reserve_mib(free_mib, is_igpu) + if capped < free_mib: + logger.info( + f"Vulkan device VK{idx} is an integrated GPU sharing system " + f"RAM; reserving {free_mib - capped}MiB host headroom " + f"({free_mib}->{capped}MiB usable)" + ) + gpus.append((idx, capped, total_mib)) + gpus.sort(key = lambda g: g[0]) + if gpus: + logger.info( + "Vulkan GPU memory detected: " + + ", ".join(f"VK{idx}={free}MiB" for idx, free, _total in gpus) + ) + return gpus + @staticmethod def _available_system_memory_mib() -> Optional[int]: """Available system RAM in MiB (psutil, then /proc/meminfo), or None if @@ -2587,7 +3015,8 @@ class LlamaCppBackend: def _llama_server_env_for_binary(binary: str) -> dict[str, str]: """Build a subprocess env that lets llama-server resolve native libs.""" env = child_env_without_native_path_secret() - binary_dir = str(Path(binary).parent) + # _llama_lib_dir resolves the llama-server symlink to the real build/bin. + binary_dir = str(_llama_lib_dir(binary)) if sys.platform == "win32": # Ordering: see _build_windows_path_dirs. #5106. @@ -3048,6 +3477,12 @@ class LlamaCppBackend: _CTX_COMPUTE_BYTES_PER_EMBD = 2.25 # quantized KV, regular attention (dequant scratch) _CTX_COMPUTE_BYTES_PER_EMBD_MLA = 1.25 # quantized KV, MLA (compressed attn: measured 0.94x) _CTX_COMPUTE_F16_MASK_SAFETY = 1.5 # f16/bf16/f32 KV: KQ mask only (n_ubatch*2 B/tok) + # DeepSeek-V4 (deepseek4): its lightning indexer + sparse attention reserve a large + # context-scaling compute buffer the rates above miss (present even with an f16 + # cache). Measured on UD-Q4_K_XL (ub=512): ~2 GiB at 16k -> ~65.5 GiB at 1M. Without + # it auto-fit commits the full 1M train context, OOMs the reserve, and spills to CPU. + _DSV4_CTX_COMPUTE_FLAT_BYTES = 2 * 1024**3 # ctx-independent indexer scratch + _DSV4_CTX_COMPUTE_BYTES_PER_TOK = 72000 # per token at ub=512 (~72 GiB at 1M) def _estimate_compute_buffer_bytes( self, @@ -3097,6 +3532,14 @@ class LlamaCppBackend: if n_embd <= 0 or n_ctx <= 0: return 0 ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH)) + if getattr(self, "_architecture", None) == "deepseek4": + # DSV4 indexer/CSA buffer (see constants): flat + linear, ub-scaled. Fires + # for any KV type -- the indexer scratch is present even with an f16 cache. + ub_scale = ub / self._DEFAULT_N_UBATCH + return int( + self._DSV4_CTX_COMPUTE_FLAT_BYTES + + self._DSV4_CTX_COMPUTE_BYTES_PER_TOK * n_ctx * ub_scale + ) if _kv_bytes_per_elem(cache_type_kv) < 2.0: # Quantized cache: the dequant scratch dominates and scales with n_embd. # MLA (compressed KV) needs far less of it: measured 0.94 x n_embd on @@ -3737,7 +4180,11 @@ class LlamaCppBackend: # Auto-size (0): the visual server probes the largest context that fits this GPU's VRAM # (capped at the training context). An explicit in-range n_ctx overrides it. maxtok = n_ctx if (n_ctx and 0 < n_ctx <= 65536) else 0 - gpu = os.environ.get("DG_GPU", "0") + # No visible CUDA GPU: a genuine CPU host, or a GPU host masked with + # CUDA_VISIBLE_DEVICES="" to force CPU serving. Keep the visual-server child + # CPU-masked (empty --gpu) so the shim does not re-expose GPU 0 via its default. + cpu_only = self._effective_gpu_count() == 0 + gpu = "" if cpu_only else os.environ.get("DG_GPU", "0") cmd = list(shim_cmd) + [ "--gguf", @@ -3757,6 +4204,12 @@ class LlamaCppBackend: # refuses to load unless UNSLOTH_IS_PRESENT is set (normally by `import # unsloth`). The shim never imports unsloth, so set it here as unsloth does. env["UNSLOTH_IS_PRESENT"] = "1" + # The shim's `import unsloth_zoo` aborts in get_device_type() ("needs a GPU") + # when no accelerator is visible, even though it only drives the CPU + # visual-server binary and does no torch GPU work. Allow the CPU device so the + # runner starts; the visual server still runs on the CPU llama.cpp build. + if cpu_only: + env.setdefault("UNSLOTH_ALLOW_CPU", "1") env["DG_VISUAL_BIN"] = visual_bin env["DG_GPU"] = gpu # The file-override shim imports its sibling visual_engine; put its dir on PYTHONPATH. @@ -3877,6 +4330,15 @@ class LlamaCppBackend: "Install it with: pip install huggingface_hub" ) + resolved_hf_repo = _resolve_repo_id_casing(hf_repo) + if resolved_hf_repo != hf_repo: + logger.info( + "Using cached repo_id casing '%s' for requested '%s'", + resolved_hf_repo, + hf_repo, + ) + hf_repo = resolved_hf_repo + # Resolve the filename from the variant gguf_filename = None gguf_extra_shards: list[str] = [] @@ -3922,10 +4384,12 @@ class LlamaCppBackend: # Check disk space; fall back to a smaller variant if needed all_gguf_files = [gguf_filename] + gguf_extra_shards + expected_sizes: dict[str, int] = {} try: from huggingface_hub import get_paths_info, try_to_load_from_cache path_infos = list(get_paths_info(hf_repo, all_gguf_files, token = hf_token)) + expected_sizes = {p.path: p.size for p in path_infos if p.size} total_bytes = sum((p.size or 0) for p in path_infos) # Subtract bytes already in the HF cache so we only preflight @@ -3934,7 +4398,26 @@ class LlamaCppBackend: # cold whenever free disk is below the full weight footprint, # even though nothing needs downloading. already_cached_bytes = 0 - if not force: + # Cross-snapshot / case-variant cache reuse is offline-only (see the download + # path below); online, hf_hub_download fetches the current revision and + # resumes partials, so an old snapshot must not be counted as cached here or + # the preflight would under-count the download and skip the disk fallback. + offline = _hf_env_offline() + # A split GGUF whose shards are not co-located in a single snapshot is + # refetched as a whole set later, so it must not be counted as cached here. + split_needs_refetch = False + if offline and not force and gguf_extra_shards: + # Scan all snapshots for one that holds the whole set co-located, so a + # newer snapshot with only the first shard does not mask an older + # complete one and needlessly trip the disk fallback. + if ( + _cached_colocated_split_main( + hf_repo, gguf_filename, gguf_extra_shards, expected_sizes + ) + is None + ): + split_needs_refetch = True + if not force and not split_needs_refetch: for p in path_infos: if not p.size: continue @@ -3942,6 +4425,15 @@ class LlamaCppBackend: cached_path = try_to_load_from_cache(hf_repo, p.path) except Exception: cached_path = None + if ( + not (isinstance(cached_path, str) and os.path.exists(cached_path)) + and offline + ): + cached_path = _cached_hf_snapshot_file( + hf_repo, + p.path, + expected_size = p.size, + ) if isinstance(cached_path, str) and os.path.exists(cached_path): try: on_disk = os.path.getsize(cached_path) @@ -4004,6 +4496,13 @@ class LlamaCppBackend: ) else: gguf_extra_shards = [] + # Record the fallback's size so the later cache-reuse probe can + # size-verify it; only for a single-file fallback, since + # _find_smallest_fitting_variant returns the whole-variant size + # and using that as the first shard's expected size would reject + # a valid cached first shard of a split fallback. + if not gguf_extra_shards: + expected_sizes[fallback_file] = fallback_size else: raise RuntimeError( f"Not enough disk space to download any variant. " @@ -4023,25 +4522,45 @@ class LlamaCppBackend: raise RuntimeError("Cancelled") dl_start = time.monotonic() # Xet primary, HTTP fallback on stall; per-file so finished shards stay cached. - local_path = hf_hub_download_with_xet_fallback( - hf_repo, - gguf_filename, - hf_token, - cancel_event = cancel_event, - on_status = lambda m: logger.info(m), - force_download = force, - ) - for shard in gguf_extra_shards: - if cancel_event.is_set(): - raise RuntimeError("Cancelled") - logger.info(f"Resolving GGUF shard: {shard}") - hf_hub_download_with_xet_fallback( + local_path = None + # Reuse a cached copy from another snapshot / case-variant repo dir only when + # offline. Online, fall through to hf_hub_download so its revision/etag check + # fetches the current file (and resumes a partial) instead of serving a stale + # same-name blob from an older revision. + if not force and _hf_env_offline(): + if gguf_extra_shards: + # A split GGUF must load every shard from one snapshot; reuse only a + # snapshot that holds the whole set co-located, scanning past a newer + # snapshot that has just the first shard while an older one is complete. + local_path = _cached_colocated_split_main( + hf_repo, gguf_filename, gguf_extra_shards, expected_sizes + ) + else: + local_path = _cached_hf_snapshot_file( + hf_repo, + gguf_filename, + expected_size = expected_sizes.get(gguf_filename), + ) + if local_path is None: + local_path = hf_hub_download_with_xet_fallback( hf_repo, - shard, + gguf_filename, hf_token, cancel_event = cancel_event, + on_status = lambda m: logger.info(m), force_download = force, ) + for shard in gguf_extra_shards: + if cancel_event.is_set(): + raise RuntimeError("Cancelled") + logger.info(f"Resolving GGUF shard: {shard}") + hf_hub_download_with_xet_fallback( + hf_repo, + shard, + hf_token, + cancel_event = cancel_event, + force_download = force, + ) except Exception as e: if isinstance(e, RuntimeError) and "Cancelled" in str(e): raise @@ -4119,6 +4638,17 @@ class LlamaCppBackend: if target is None or cancel_event.is_set(): return None + # Offline, resolve the companion straight from the cache snapshot that + # holds it. resolve_cached_repo_id_case can return a partial lower-case + # spelling when any dir exists under the requested casing, so calling + # hf_hub_download with hf_repo would miss the canonical file and silently + # drop the companion. _cached_hf_snapshot_file scans every case variant. + if _hf_env_offline(): + cached = _cached_hf_snapshot_file(hf_repo, target) + if cached: + logger.info("Resolved %s from local HF cache: %s", label, cached) + return cached + try: logger.info(f"Downloading {label}: {hf_repo}/{target}") # Same policy; companions are best-effort (caller below swallows failures to None). @@ -4167,6 +4697,29 @@ class LlamaCppBackend: cancel_event = cancel_event, ) + def _cached_repo_mtp_drafter(self, hf_repo: str) -> Optional[str]: + """A drafter already in this repo's local HF cache, reused offline when a + fresh copy can't be fetched. Prefers a repo-root ``mtp-*.gguf`` across all + cached snapshots; else an existing ``MTP/`` copy (any precision -- the + target verifies every drafted token). None if none is cached.""" + try: + from utils.models.model_config import _iter_hf_cache_snapshots + + roots: list[Path] = [] + subdirs: list[Path] = [] + for snap in _iter_hf_cache_snapshots(hf_repo): # newest first + for f in sorted(_gguf_snapshot_files(snap)): + if _is_companion_gguf_path(f) and "mmproj" not in f.lower(): + (roots if "/" not in f else subdirs).append(snap / f) + # Keep snapshot order (newest first), root before any MTP/ copy, so a + # newer main GGUF pairs with the newest cached drafter, not a stale one. + for cand in roots + subdirs: + if cand.is_file(): + return str(cand) + except Exception as e: + logger.debug("Cached MTP drafter lookup failed for %s: %s", hf_repo, e) + return None + def _download_mtp( self, *, @@ -4183,11 +4736,25 @@ class LlamaCppBackend: are intentionally skipped. Returns the local path, or None. """ + # Offline, reuse any drafter already on disk (a fresh copy can't be + # fetched). Online, _download_companion_gguf/hf_hub_download reuse the + # current cached file and refetch a changed one, so skip the probe here + # rather than pair new weights with a stale draft. + if _hf_env_offline(): + cached = self._cached_repo_mtp_drafter(hf_repo) + if cached: + logger.info(f"Reusing cached MTP drafter (offline): {cached}") + return cached + def _pick_mtp(candidates: list[str]) -> Optional[str]: + # Root-level only: MTP/ subdir copies now share the mtp- prefix but + # are explicit-selection, not auto-fetch (they'd sort ahead of root). mtp_files = sorted( f for f in candidates - if f.lower().endswith(".gguf") and Path(f).name.lower().startswith("mtp-") + if f.lower().endswith(".gguf") + and "/" not in f + and Path(f).name.lower().startswith("mtp-") ) return mtp_files[0] if mtp_files else None @@ -4889,6 +5456,7 @@ class LlamaCppBackend: # Resolve llama-server now but defer a not-found error: a block-diffusion # GGUF uses the diffusion runner, and its arch is only known after the header. binary = self._find_llama_server_binary() + is_vulkan_backend = self._is_vulkan_backend(binary) # ── Phase 2: download (NO lock held, so cancel can proceed) ── # mtp_draft_path arrives set for local Gemma loads (detected @@ -4897,6 +5465,19 @@ class LlamaCppBackend: # dead; cleanup runs even on exception so a transient hiccup # can't quarantine future loads. if hf_repo: + # Resolve the requested repo id to its cached canonical casing once, + # up front, so the main GGUF and its companions (mmproj / MTP drafter) + # all resolve from the same cache entry. Otherwise a case-variant + # request resolves the main file from the canonical cache dir while the + # companions keep the requested casing and miss the cached files. + _resolved_repo = _resolve_repo_id_casing(hf_repo) + if _resolved_repo != hf_repo: + logger.info( + "Using cached repo_id casing '%s' for requested '%s'", + _resolved_repo, + hf_repo, + ) + hf_repo = _resolved_repo with _hf_offline_if_dns_dead(): model_path = self._download_gguf( hf_repo = hf_repo, @@ -5115,7 +5696,8 @@ class LlamaCppBackend: model_size = gguf_size + mmproj_size # 2-tuple gpus for existing logic + a total map for the absolute # per-GPU headroom (correct when the GPU is already partly used). - _gpu_mem = self._get_gpu_memory() + # Pass binary so a Vulkan build probes ggml's Vulkan ordinals. + _gpu_mem = self._get_gpu_memory(binary) gpus = [(idx, free) for idx, free, _t in _gpu_mem] total_by_idx = {idx: total for idx, _f, total in _gpu_mem} @@ -5888,7 +6470,12 @@ class LlamaCppBackend: # cap, not the ROCm-reported VRAM, is the real ceiling); refuse an # oversize load the OS would otherwise kill mid-flight. Base model # only: an optional MTP drafter is dropped by the MTP-drop fallback. - if model_size is not None and self._amd_apu_wants_unified_memory(gpu_indices): + # CUDA/ROCm ids only; a Vulkan build's gpu_indices are ggml ordinals. + if ( + model_size is not None + and not is_vulkan_backend + and self._amd_apu_wants_unified_memory(gpu_indices) + ): _ram_msg = self._apu_ram_shortfall_message( model_size, self._available_system_memory_mib() ) @@ -6151,6 +6738,12 @@ class LlamaCppBackend: ", ".join(unsupported_cache_flags), ) + # Vulkan pins via --device (a cmd arg, unlike the env-based + # CUDA/ROCm pin below), emitted BEFORE user extras so llama.cpp's + # last-wins parsing lets a user --device override Studio's pick. + if is_vulkan_backend and gpu_indices is not None: + cmd += LlamaCppBackend._vulkan_pin_args(gpu_indices) + # User pass-through args go last so llama.cpp's last-wins parsing # lets the user override Studio's auto-set flags. Already # validated by the route via validate_extra_args(). @@ -6202,23 +6795,25 @@ class LlamaCppBackend: env.setdefault("OMP_NUM_THREADS", "2") # AMD unified-memory APUs (gfx1150/gfx1151): let llama.cpp use - # shared system RAM. setdefault so a user value wins. - if self._amd_apu_wants_unified_memory(gpu_indices): + # shared system RAM. setdefault so a user value wins. Not on Vulkan + # (nor DC below): gpu_indices are ggml ordinals, not CUDA/ROCm ids. + if not is_vulkan_backend and self._amd_apu_wants_unified_memory(gpu_indices): env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1") logger.info("AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1") # DC NVIDIA GPUs: FP32 accum (+ P2P / launch queues for multi-GPU). # See _apply_datacenter_env; opt out with UNSLOTH_DISABLE_DC_TUNING=1. - if self._apply_datacenter_env(env, gpu_indices): + if not is_vulkan_backend and self._apply_datacenter_env(env, gpu_indices): multi_gpu = self._effective_gpu_count(gpu_indices) > 1 logger.info( f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})" ) # Pin to selected GPU(s). On ROCm, narrowing only - # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full - # set, so set HIP_VISIBLE_DEVICES too. - if gpu_indices is not None: + # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full set, so + # set HIP_VISIBLE_DEVICES too. Vulkan is pinned via --device + # (above), not here. + if gpu_indices is not None and not is_vulkan_backend: pinned = ",".join(str(i) for i in gpu_indices) env["CUDA_VISIBLE_DEVICES"] = pinned try: @@ -6612,6 +7207,7 @@ class LlamaCppBackend: ) self._healthy = True + self._commit_effective_parallel_slots(n_parallel) # Commit caller intent only after _healthy=True so a failed start # can't poison the next inheritance check. None keeps prior, [] @@ -7149,6 +7745,7 @@ class LlamaCppBackend: self._context_length = None self._effective_context_length = None self._max_context_length = None + self._reset_effective_parallel_slots() self._chat_template = None self._chat_template_override = None self._supports_reasoning = False @@ -7204,6 +7801,7 @@ class LlamaCppBackend: # Stop the watchdog before a deliberate kill so a planned reload/unload # isn't seen as a crash; a real crash never routes through here. self._stop_mtp_crash_watchdog() + self._reset_effective_parallel_slots() if self._process is None: return try: @@ -7881,12 +8479,17 @@ class LlamaCppBackend: # ── Message building (OpenAI format) ────────────────────────── @staticmethod - def _parse_tool_calls_from_text(content: str, *, allow_incomplete: bool = True) -> list[dict]: - """Thin wrapper around the shared parser in tool_call_parser - so safetensors and llama_cpp pick up the same fixes.""" + def _parse_tool_calls_from_text( + content: str, + *, + allow_incomplete: bool = True, + enabled_tool_names: Optional[set] = None, + ) -> list[dict]: + """Wrapper around the shared parser; ``enabled_tool_names`` gates the markerless bare-JSON form.""" return _shared_parse_tool_calls_from_text( content, allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, ) @staticmethod @@ -8040,7 +8643,7 @@ class LlamaCppBackend: ): """Open one streaming POST and let cancel interrupt prefill or reads.""" if cancel_event is not None and cancel_event.is_set(): - raise GeneratorExit + raise _LlamaStreamCancelled _cancel_closed = threading.Event() _response_ref: list = [None] @@ -8085,13 +8688,13 @@ class LlamaCppBackend: ) as response: _response_ref[0] = response if cancel_event is not None and cancel_event.is_set(): - raise GeneratorExit + raise _LlamaStreamCancelled yield response return except (httpx.RequestError, RuntimeError): # Response was closed by the cancel watcher if cancel_event is not None and cancel_event.is_set(): - raise GeneratorExit + raise _LlamaStreamCancelled raise finally: _cancel_closed.set() @@ -8294,6 +8897,8 @@ class LlamaCppBackend: "finish_reason": _metadata_finish_reason, } + except _LlamaStreamCancelled: + return except httpx.ConnectError as e: # Server already down. If this was an MTP+tensor crash, recover by # reloading without MTP (scheduled in the background) and fail this @@ -8353,6 +8958,7 @@ class LlamaCppBackend: preserve_thinking: Optional[bool] = None, max_tool_iterations: int = 25, auto_heal_tool_calls: bool = True, + nudge_tool_calls: Optional[bool] = None, tool_call_timeout: int = 300, session_id: Optional[str] = None, rag_scope: Optional[dict] = None, @@ -8392,12 +8998,23 @@ class LlamaCppBackend: _reasoning_started_at: Optional[float] = None _reasoning_summary_emitted = False + # Gate telling a genuine NAME[ARGS] rehearsal from inactive-name prose; built from the + # ORIGINAL tools list so a spent one-shot still reads as a tool name. None = no gate. + _enabled_names_gate = set(_gguf_active_tool_names(tools)) if tools else None + # Detection must see the same names as the strip gate (ORIGINAL list, incl. a spent + # one-shot), else its repeat is stripped but never drained and the turn ends blank. + _detect_tools = list(tools or []) + def _reasoning_summary_event(started_at: float) -> dict: return { "type": "reasoning_summary", "duration_ms": round((time.monotonic() - started_at) * 1000.0), } + # Enabled-name gate for the markerless Gemma strip (disabled/example + # names stay visible). Set per iteration; None = pre-loop name-agnostic. + _enabled_tool_names = None + def _strip_tool_markup( text: str, *, @@ -8406,14 +9023,42 @@ class LlamaCppBackend: ) -> str: if not (auto_heal_tool_calls or force): return text - return strip_tool_call_markup(text, final = final) + # Delegate to the shared parser-side strip so the GGUF cleanup covers every family the + # parser promotes (Llama <|python_tag|>, Mistral [TOOL_CALLS], bare rehearsal, function + # XML, Gemma) and stays aligned with detection; tool_healing's strip omits the loop-only + # forms (python_tag / Mistral name) and would leak them into display. + return _shared_strip_tool_markup( + text, final = final, enabled_tool_names = _enabled_names_gate + ) def _strip_tool_markup_streaming(text: str, *, force: bool = False) -> str: if not (auto_heal_tool_calls or force): return text - for pat in _TOOL_ALL_PATS: - text = pat.sub("", text) - return text + + def _seg(segment: str, is_last: bool) -> str: + # Same scan order as the parser's _strip_segment (seg_final -> is_last): balanced + # strips first (nested JSON removed whole; literal markup inside a value is that + # call's data), then the guarded function-XML / GLM scans, then the regex arms + # (DeepSeek / Kimi / closed forms). EOS-anchored tail arms run only on the last + # segment (a bare ``foo[ARGS]`` before is prose). Rehearsal + markerless + # strips are name-gated on the ORIGINAL list (strip/detect aligned). + seg = _strip_mistral_closed_calls(segment) + seg = _strip_bracket_tag_calls(seg, enabled_tool_names = _enabled_names_gate) + if is_last: + seg = _strip_gemma_wrapperless_calls(seg, _enabled_names_gate) + seg = _strip_function_xml_calls(seg, final = is_last) + seg = _strip_glm_calls(seg, final = is_last) + pats = _PARSER_TOOL_ALL_PATS if is_last else _PARSER_TOOL_CLOSED_PATS + for pat in pats: + seg = pat.sub("", seg) + if is_last: + seg = apply_tool_strip_patterns( + seg, [_REHEARSAL_TAIL_STRIP_RE], enabled_tool_names = _enabled_names_gate + ) + return seg + + # Preserve think blocks verbatim (a rehearsed call inside one must not be deleted). + return strip_outside_think(text, _seg) def _build_metadata_event(usage, timings, finish_reason): """Final usage+timings metadata event for the given pass, merging its @@ -8449,13 +9094,38 @@ class LlamaCppBackend: } def _flush_reasoning_and_buffer(): - """Append buffered reasoning (as a block) then the held + """Close a live-streamed block (or emit the buffered reasoning + as one block if it never streamed), then append the held content_buffer to the cumulative display text.""" - nonlocal cumulative_display - if reasoning_accum: + nonlocal cumulative_display, in_thinking + if in_thinking: + cumulative_display += "" + in_thinking = False + elif reasoning_accum: cumulative_display += "" + reasoning_accum + "" cumulative_display += content_buffer + def _close_streamed_think() -> bool: + """Close a live-streamed before a tool call drains, so + consumers without a reasoning extractor (Anthropic) get a balanced + block. Returns True when the caller should yield the result.""" + nonlocal cumulative_display, in_thinking, _last_emitted + if not in_thinking: + return False + cumulative_display += "" + in_thinking = False + if len(cumulative_display) > len(_last_emitted) and not _suppress_visible_output: + _last_emitted = cumulative_display + return True + return False + + def _looks_like_enabled_bare_json(text: str, enabled_tool_names: set) -> bool: + """True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False.""" + probe = strip_llama3_leading_sentinels(text.lstrip()) + if not (probe.startswith("{") and ('"name"' in probe or '"function"' in probe)): + return False + return strip_leading_bare_json_call(probe, enabled_tool_names) != probe + tool_controller = ToolLoopController( tools = tools, auto_heal_tool_calls = auto_heal_tool_calls, @@ -8469,18 +9139,21 @@ class LlamaCppBackend: ) _MAX_BUFFER_CHARS = 32 + # Hold a leading ``{`` well past the 32-char XML cap until it balances (mirrors safetensors). + _MAX_BARE_JSON_BUFFER = 16384 _append_budget_exhausted_nudge = True # RAG: cap knowledge-base searches per assistant turn. The controller is # tool-agnostic, so this gate stays in the loop. _kb_search_count = 0 # ── Re-prompt on plan-without-action ───────────────── - # When the model describes what it intends to do (forward-looking - # language) without calling a tool, re-prompt once. Only triggers on - # responses signaling intent/planning -- a direct answer like "4" or - # "Hello!" won't match. Pattern compiled at module level - # (_INTENT_SIGNAL). + # Model describes intent without calling a tool: re-prompt once. A + # direct answer ("4", "Hello!") won't match. Pattern shared with the + # safetensors loop (tool_call_parser.INTENT_SIGNAL). _reprompt_count = 0 + # Gates ``max_tool_iterations`` on real tool turns (not the enlarged range) so reserved + # re-prompt slots don't extend the budget. Mirrors the safetensors guard. + _tool_iters_done = 0 _forced_tool_call_pending = False # Reserve extra iterations for re-prompts so they don't consume the @@ -8489,12 +9162,21 @@ class LlamaCppBackend: for iteration in range(max_tool_iterations + _extra): if cancel_event is not None and cancel_event.is_set(): return + # Whether this turn ran a tool; a no-op-only turn stays False and doesn't consume budget. + _turn_executed_real_tool = False active_tools = tool_controller.active_tools() if not active_tools: _append_budget_exhausted_nudge = False break - _tool_xml_signals = TOOL_XML_SIGNALS + # Gate the markerless bare-JSON form on enabled names so an ordinary JSON answer isn't misread as a call. + _enabled_tool_names = { + (tool.get("function") or {}).get("name") + for tool in active_tools + if (tool.get("function") or {}).get("name") + } + # Shared signal tuple so GGUF BUFFERING wakes on every format the parser knows (like safetensors). + _tool_xml_signals = _SHARED_TOOL_XML_SIGNALS # Build payload -- stream: True so we detect tool signals # in the first 1-2 chunks without a non-streaming penalty. @@ -8624,6 +9306,10 @@ class LlamaCppBackend: # the structured tool call. has_structured_tc = True detect_state = _S_DRAINING + # Close the reasoning prefix before the tool card + # (mirrors the is_match path). + if _close_streamed_think(): + yield {"type": "content", "text": cumulative_display} for tc_d in tc_deltas: idx = tc_d.get("index", 0) if idx not in tool_calls_acc: @@ -8709,17 +9395,17 @@ class LlamaCppBackend: continue # ── Reasoning tokens ── - # Yield only in STREAMING. In BUFFERING and - # DRAINING, accumulate silently so we don't - # corrupt the consumer's prev_text tracker - # (routes/inference.py never resets it - # between tool iterations). + # Stream live except while DRAINING: reasoning is + # orthogonal to tool detection (content_buffer + # only), and the route resets prev_text on + # tool_start, so the block stays a + # monotonic prefix like the no-tool path. reasoning = delta.get("reasoning_content", "") if reasoning: if _reasoning_started_at is None: _reasoning_started_at = time.monotonic() reasoning_accum += reasoning - if detect_state == _S_STREAMING: + if detect_state != _S_DRAINING: if not in_thinking: cumulative_display += "" in_thinking = True @@ -8752,12 +9438,18 @@ class LlamaCppBackend: in_thinking = False cumulative_display += token cleaned = _strip_tool_markup_streaming(cumulative_display) - if len(cleaned) > len(_last_emitted): - _last_emitted = cleaned + # Hold a trailing bare active-tool-name (split rehearsal) + # until [ARGS] arrives; released by later prose or stream end. + _hold = _held_rehearsal_tail_len(cleaned, _detect_tools) + _emit = ( + cleaned[: len(cleaned) - _hold] if _hold else cleaned + ) + if len(_emit) > len(_last_emitted): + _last_emitted = _emit if not _suppress_visible_output: yield { "type": "content", - "text": cleaned, + "text": _emit, } elif detect_state == _S_BUFFERING: @@ -8766,7 +9458,8 @@ class LlamaCppBackend: if not stripped_buf: continue - # Check tool signal prefixes. + # Bracket tags arrive mid-buffer, so substring-check too; + # ``[ARGS]`` counts only as a regex-matched NAME[ARGS]. is_prefix = False is_match = False for sig in _tool_xml_signals: @@ -8776,14 +9469,91 @@ class LlamaCppBackend: if sig.startswith(stripped_buf): is_prefix = True break + if sig == "[ARGS]": + # Active NAME[ARGS] only; inactive-name prose + # is gated out, not drained/parsed. + if ( + _gguf_rehearsal_signal_pos( + stripped_buf, _detect_tools + ) + >= 0 + ): + is_match = True + break + elif sig.startswith("[") and sig in stripped_buf: + is_match = True + break - if is_match: + # Split rehearsal: hold the bare name until + # its [ARGS] arrives and matches above. + is_rehearsal_prefix = False + if ( + not is_match + and not is_prefix + and _is_rehearsal_prefix(stripped_buf, _detect_tools) + ): + is_prefix = True + is_rehearsal_prefix = True + + # Signal-less call shapes (mirror the safetensors + # loop): Llama-3.2 bare {"name":..} and Gemma + # call:NAME{...} would otherwise stream raw. + _hold_buffer = False + # Whole buffer is the call (no visible prefix) -- drain silently. + _drain_silently = False + if not is_match and not is_prefix: + _bare = strip_llama3_leading_sentinels(stripped_buf) + if _bare.startswith("{"): + if _balanced_brace_end(_bare, 0) is None: + if len(stripped_buf) < _MAX_BARE_JSON_BUFFER: + _hold_buffer = True + elif _looks_like_enabled_bare_json( + _bare, _enabled_tool_names + ): + # Oversized still-open enabled call: drain + # rather than leak; a giant ordinary JSON + # answer still streams. + _drain_silently = True + elif self._parse_tool_calls_from_text( + content_buffer, + allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, + ): + _drain_silently = True + elif ( + "call:".startswith(stripped_buf) + or _GEMMA_BARE_TC_PREFIX_RE.match(stripped_buf) + is not None + or _GEMMA_BARE_TC_RE.match(stripped_buf) is not None + ): + # Whitespace-tolerant like the parser. + if _GEMMA_BARE_TC_RE.match(stripped_buf): + _drain_silently = True + elif len(stripped_buf) < _MAX_BUFFER_CHARS: + _hold_buffer = True + + if _drain_silently: + # The buffered content IS the call; drain it + # without yielding. A live prefix is + # separate from it -- close that. + detect_state = _S_DRAINING + if _close_streamed_think(): + yield { + "type": "content", + "text": cumulative_display, + } + elif is_match: # Tool signal -- flush any visible # prefix before DRAINING so the # route sends it before tool_start. + # Use the final strip (all families incl. Llama + # <|python_tag|> / Mistral name): the buffer holds + # the whole call, so a streaming closed-only strip + # would leak its open-ended markup as display text. _flush_reasoning_and_buffer() - cleaned = _strip_tool_markup_streaming( + cleaned = _strip_tool_markup( cumulative_display, + final = True, force = True, ) if len(cleaned) > len(_last_emitted): @@ -8794,7 +9564,15 @@ class LlamaCppBackend: "text": cleaned, } detect_state = _S_DRAINING - elif is_prefix and len(stripped_buf) < _MAX_BUFFER_CHARS: + elif _hold_buffer or ( + is_prefix + and ( + is_rehearsal_prefix + or len(stripped_buf) < _MAX_BUFFER_CHARS + ) + ): + # A rehearsal prefix is self-bounded; the buffer + # cap must not cut long MCP names short. pass # keep buffering else: # Not a tool -- flush buffer @@ -8805,12 +9583,20 @@ class LlamaCppBackend: cleaned = _strip_tool_markup( cumulative_display, ) - if len(cleaned) > len(_last_emitted): - _last_emitted = cleaned + # Same trailing-name hold as STREAMING for this + # first flush out of BUFFERING. + _hold = _held_rehearsal_tail_len(cleaned, _detect_tools) + _emit = ( + cleaned[: len(cleaned) - _hold] + if _hold + else cleaned + ) + if len(_emit) > len(_last_emitted): + _last_emitted = _emit if not _suppress_visible_output: yield { "type": "content", - "text": cleaned, + "text": _emit, } except json.JSONDecodeError: @@ -8821,7 +9607,18 @@ class LlamaCppBackend: # ── Resolve BUFFERING at stream end ── if detect_state == _S_BUFFERING: stripped_buf = content_buffer.lstrip() - if stripped_buf and any(s in stripped_buf for s in _tool_xml_signals): + # A held bare-JSON fragment has no XML signal; route it to DRAINING (the signal-only + # gate below would flush the raw JSON to the user). + _bare_eos = strip_llama3_leading_sentinels(stripped_buf) + # Gate on enabled names so an ordinary JSON answer isn't routed to DRAINING and dropped. + _is_bare_tc = bool(active_tools) and _looks_like_enabled_bare_json( + _bare_eos, _enabled_tool_names + ) + if stripped_buf and _gguf_has_genuine_tool_signal( + stripped_buf, _tool_xml_signals, _detect_tools + ): + detect_state = _S_DRAINING + elif _is_bare_tc: detect_state = _S_DRAINING elif content_accum or reasoning_accum: detect_state = _S_STREAMING @@ -8837,7 +9634,9 @@ class LlamaCppBackend: ), } elif reasoning_accum and not has_content_tokens: - # Reasoning-only reply: show it as plain text. + # Reasoning-only reply: show it as the main response, + # not a thinking block (mirrors the no-tool path; the + # route's extractor closes the streamed ). if _reasoning_started_at is not None and not _reasoning_summary_emitted: _reasoning_summary_emitted = True yield _reasoning_summary_event(_reasoning_started_at) @@ -8848,20 +9647,26 @@ class LlamaCppBackend: "text": cumulative_display, } else: + # Held buffer was no tool signal and no enabled bare-JSON call: a leading ``{`` is an + # ordinary JSON answer and must be shown; any other partial-markup prefix is dropped. + _held = strip_llama3_leading_sentinels(content_buffer.lstrip()) + if _held.startswith("{") and not _suppress_visible_output: + yield {"type": "content", "text": _held} return # ── STREAMING path: no tool call ── if detect_state == _S_STREAMING: - # Safety net: check for XML tool signals in content. The + # Safety net: re-parse the full content for tool calls. The # route layer resets prev_text on tool_start, so post-tool # synthesis streams correctly even if content was emitted # before the tool XML. - _safety_tc = None - if any(s in content_accum for s in _tool_xml_signals): - _safety_tc = self._parse_tool_calls_from_text( - content_accum, - allow_incomplete = auto_heal_tool_calls, - ) + # Unconditional (not gated on _tool_xml_signals): bare-JSON and Gemma wrapper-less + # calls carry no XML signal, so a signal gate would let them slip past. + _safety_tc = self._parse_tool_calls_from_text( + content_accum, + allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, + ) if not _safety_tc: # ── Re-prompt on plan-without-action ── # If the model described its intent (forward-looking @@ -8879,8 +9684,10 @@ class LlamaCppBackend: r"(?i)\brender[_\s-]?html\b", _stripped, ) + # None keeps the default-on re-prompt; False disables it. if ( auto_heal_tool_calls + and (nudge_tool_calls is None or nudge_tool_calls) and active_tools and not _render_html_already_done_intent and _reprompt_count < _MAX_REPROMPTS @@ -8909,12 +9716,7 @@ class LlamaCppBackend: conversation.append( { "role": "user", - "content": ( - "You have access to enabled tools. If a tool is needed to satisfy " - "the user's request or complete the action you described, call " - f"{tool_hint} now. If no tool is needed, provide the final answer " - "and follow the user's requested format." - ), + "content": _reprompt_to_act_message(tool_hint), } ) # Accumulate tokens and timing from this iteration. @@ -8946,6 +9748,12 @@ class LlamaCppBackend: "type": "content", "text": forced_visible_text, } + elif not _suppress_visible_output: + # Turn ended as a plain answer (no [ARGS] followed): the held + # rehearsal tail is real prose, release it. + _final_clean = _strip_tool_markup_streaming(cumulative_display) + if len(_final_clean) > len(_last_emitted): + yield {"type": "content", "text": _final_clean} # Content was already streamed. Yield metadata. yield {"type": "status", "text": ""} @@ -8978,10 +9786,13 @@ class LlamaCppBackend: for i in sorted(tool_calls_acc) if (tool_calls_acc[i].get("function", {}).get("name", "").strip()) ] or None - if not tool_calls and any(s in content_accum for s in _tool_xml_signals): + if not tool_calls: + # Unconditional re-parse: we only reach DRAINING when the buffer looked like a + # call, and bare-JSON / Gemma wrapper-less calls carry no XML signal to gate on. tool_calls = self._parse_tool_calls_from_text( content_accum, allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, ) if tool_calls and not has_structured_tc: content_text = _strip_tool_markup( @@ -8989,6 +9800,11 @@ class LlamaCppBackend: final = True, force = True, ) + # ``_strip_tool_markup`` only knows XML; also drop a leading bare-JSON call so the + # executed call isn't replayed as text or next-turn history. + content_text = strip_leading_bare_json_call( + content_text, _enabled_tool_names + ) if tool_calls: logger.info( f"Parsed {len(tool_calls)} tool call(s) from " @@ -9002,6 +9818,13 @@ class LlamaCppBackend: if content_accum: # Strip leaked tool-call XML before yielding. content_accum = _strip_tool_markup(content_accum, final = True) + # A truncated bare-JSON call has no XML markup to strip and didn't parse. With + # Auto-Heal on, drop a leading ENABLED-tool fragment (ordinary JSON answers untouched); + # off keeps it visible per the strict contract. + if content_accum and active_tools and auto_heal_tool_calls: + content_accum = strip_leading_bare_json_call( + content_accum, _enabled_tool_names + ) if content_accum: yield {"type": "content", "text": content_accum} _meta = _build_metadata_event( @@ -9019,6 +9842,29 @@ class LlamaCppBackend: _accumulated_predicted_ms += _it.get("predicted_ms", 0) _accumulated_predicted_n += _it.get("predicted_n", 0) + # Collapse exact-duplicate calls and cap the count for the TEXTUAL + # fallback (mirrors the safetensors loop; see _MAX_TOOL_CALLS_PER_TURN). + if tool_calls and not has_structured_tc and len(tool_calls) > 1: + _seen_keys: set = set() + _deduped: list = [] + for _tc in tool_calls: + _fn = _tc.get("function", {}) or {} + _key = (_fn.get("name", ""), str(_fn.get("arguments", ""))) + if _key in _seen_keys: + continue + _seen_keys.add(_key) + _deduped.append(_tc) + if len(_deduped) >= _MAX_TOOL_CALLS_PER_TURN: + break + if len(_deduped) != len(tool_calls): + logger.info( + "GGUF textual fallback: collapsed %d repeated tool call(s) " + "in one turn to %d", + len(tool_calls), + len(_deduped), + ) + tool_calls = _deduped + # disable_parallel_tool_use: execute only the first tool call # this turn. Truncate before building assistant_msg so the # conversation stays consistent and extra calls are never executed. @@ -9144,6 +9990,8 @@ class LlamaCppBackend: _kb_search_count += 1 completion = tool_controller.record_result(decision, result) resolved_provisional_tool_call_ids.add(decision.tool_call_id) + # A tool ran this turn, so it counts against the caller's budget. + _turn_executed_real_tool = True yield completion.tool_end_event() conversation.append(completion.tool_message()) @@ -9167,8 +10015,16 @@ class LlamaCppBackend: if tool_controller.force_final_answer or not tool_controller.active_tools(): _append_budget_exhausted_nudge = False break + # Count only real tool turns against the cap so reserved re-prompt slots can't become + # extra tool rounds; a no-op correction turn doesn't consume budget (GGUF parity). + if _turn_executed_real_tool: + _tool_iters_done += 1 + if _tool_iters_done >= max_tool_iterations: + break continue + except _LlamaStreamCancelled: + return except httpx.ConnectError: # Mark unresolved provisional cards as failed before raising. for _pid, _pname in provisional_started_tool_calls.items(): @@ -9351,6 +10207,8 @@ class LlamaCppBackend: if _meta is not None: yield _meta + except _LlamaStreamCancelled: + return except httpx.ConnectError: raise RuntimeError("Lost connection to llama-server") except Exception as e: diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 5c7799152f..6287b184a6 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -5,6 +5,7 @@ Drop-in replacement for InferenceBackend — same interface, uses mlx-lm/mlx-vlm instead of torch/transformers for model loading and generation. """ +import os import threading from typing import Optional, Generator from core.inference.runtime_context import runtime_context_length @@ -41,6 +42,87 @@ def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps): } +def _mlx_distributed_rank_size(group = None): + """Return ``(rank, world_size)`` for an optional MLX distributed group.""" + if group is None: + return 0, 1 + rank = int(group.rank()) + world_size = int(group.size()) + if world_size < 1: + raise ValueError(f"Invalid MLX distributed world_size={world_size}.") + if rank < 0 or rank >= world_size: + raise ValueError(f"Invalid MLX distributed rank={rank} for world_size={world_size}.") + return rank, world_size + + +def _mlx_distributed_backend_from_env(): + if os.environ.get("MLX_JACCL_COORDINATOR") and os.environ.get("MLX_IBV_DEVICES"): + return "jaccl" + return None + + +def _init_mlx_distributed(): + """Initialize MLX distributed state, falling back to singleton metadata.""" + import mlx.core as mx + + group = None + rank = 0 + world_size = 1 + distributed = getattr(mx, "distributed", None) + init = getattr(distributed, "init", None) if distributed is not None else None + if callable(init): + backend = _mlx_distributed_backend_from_env() + if backend is None: + group = init() + else: + try: + group = init(backend = backend) + except TypeError: + group = init() + if group is not None: + rank, world_size = _mlx_distributed_rank_size(group) + return group, rank, world_size + + +def _make_mlx_presence_penalty_processor(penalty: float): + """Presence penalty as an mlx_lm/mlx_vlm logits processor, matching the safetensors path. + + generate_step calls processors as ``fn(tokens, logits)`` with ``tokens`` the + full running sequence; the first call is prompt-only, so latch that length + and penalize only after it. + """ + state = {"prompt_len": None} + + def _processor(tokens, logits): + if state["prompt_len"] is None: + # First call is prompt-only; latch its length. + state["prompt_len"] = int(tokens.shape[0]) + return logits + generated = tokens[state["prompt_len"] :] + if generated.size == 0: + return logits + import mlx.core as mx + + vocab = logits.shape[-1] + # Bound ids to [0, vocab) before indexing logits: MLX does no bounds + # checking and out-of-bounds indexing is undefined behavior (crash / + # corruption), unlike torch's harmless negative wrap. MLX also lacks + # boolean-mask filtering, so out-of-range/negative ids route to a + # scratch slot at index vocab (dropped before the subtract) that never + # collides with a real token: real ids (including 0) are penalized + # once, strays ignored. + valid = (generated >= 0) & (generated < vocab) + safe = mx.where(valid, generated, vocab).astype(mx.int32) + # Scatter penalty into a (vocab + 1)-wide mask: duplicate ids are + # idempotent (presence applies once per token); scratch column dropped. + mask = mx.zeros((vocab + 1,), dtype = logits.dtype) + mask[safe] = penalty + logits = logits - mask[:vocab] + return logits + + return _processor + + class MLXInferenceBackend: def __init__(self): self.models = {} @@ -49,7 +131,7 @@ class MLXInferenceBackend: self.loaded_local_models = [] self.device = "mlx" self._generation_lock = threading.Lock() - # usage/timings of the latest generation; shipped on gen_done. + # usage/timings of the latest generation, shipped on gen_done. self.last_generation_stats = None self._model = None @@ -57,6 +139,9 @@ class MLXInferenceBackend: self._processor = None self._is_vlm = False self._config = {} + self._distributed_group = None + self._distributed_rank = 0 + self._distributed_world_size = 1 # Recorded for unload to release pinned memory back to the OS. self._memory_limits_applied = {} @@ -101,16 +186,26 @@ class MLXInferenceBackend: trust_remote_code = False, gpu_ids = None, dtype = None, + parallel_mode = None, + distributed_group = None, ) -> bool: import mlx.core as mx + # Keep the token so the native-template fallback can fetch a gated + # model's repo template during generation. + self._hf_token = hf_token model_name = config.identifier if hasattr(config, "identifier") else str(config) is_vision = getattr(config, "is_vision", False) + distributed_rank, distributed_size = _mlx_distributed_rank_size(distributed_group) + is_distributed = distributed_group is not None and distributed_size > 1 + self._distributed_group = distributed_group + self._distributed_rank = distributed_rank + self._distributed_world_size = distributed_size - # GGUF guard. GGUF models are served by llama-server in the parent - # process, not mlx-lm here. Reaching this with is_gguf=True means the - # route's first detection flaked (transient HF Hub) but the subprocess - # re-detected GGUF; raise loudly instead of a cryptic mlx_lm error. + # GGUF guard: GGUF is served by llama-server in the parent process, + # not mlx-lm. Reaching here with is_gguf=True means the route's + # detection flaked but the subprocess re-detected GGUF; raise loudly + # instead of a cryptic mlx_lm error. if getattr(config, "is_gguf", False): raise RuntimeError( f"MLXInferenceBackend cannot load GGUF model '{model_name}': " @@ -129,11 +224,26 @@ class MLXInferenceBackend: is_lora = getattr(config, "is_lora", False) logger.info( - "Loading %s via %s (is_lora=%s)", + "Loading %s via %s (is_lora=%s, distributed=%s, rank=%s/%s, mode=%s)", model_name, "mlx-vlm" if is_vision else "mlx-lm", is_lora, + is_distributed, + distributed_rank, + distributed_size, + parallel_mode, ) + if is_distributed and parallel_mode not in ("pipeline", "tensor"): + raise ValueError( + "Unsloth: distributed MLX inference requires parallel_mode='pipeline' " + "or parallel_mode='tensor'." + ) + if is_distributed and is_lora: + raise ValueError( + "Unsloth: distributed MLX inference for LoRA adapter repos " + "is not supported yet. Merge/export the adapter into an MLX model " + "before distributed inference." + ) try: from unsloth_zoo.mlx.loader import FastMLXModel @@ -143,14 +253,23 @@ class MLXInferenceBackend: "(unsloth_zoo.mlx.loader). Reinstall via install.sh on Apple Silicon." ) from e + load_kwargs = { + "max_seq_length": max_seq_length, + "dtype": dtype, + "load_in_4bit": load_in_4bit, + "token": hf_token, + "trust_remote_code": trust_remote_code, + "text_only": False if is_vision else True, + } + if is_distributed: + if parallel_mode == "pipeline": + load_kwargs["pipeline_group"] = distributed_group + else: + load_kwargs["tensor_group"] = distributed_group + model, tokenizer_or_processor = FastMLXModel.from_pretrained( model_name, - max_seq_length = max_seq_length, - dtype = dtype, - load_in_4bit = load_in_4bit, - token = hf_token, - trust_remote_code = trust_remote_code, - text_only = False if is_vision else True, + **load_kwargs, ) if is_vision: @@ -168,18 +287,25 @@ class MLXInferenceBackend: self.active_model_name = model_name self.models[model_name] = { + # Per-model token for the native-template fallback (matches transformers). + "hf_token": hf_token, + # Per-model trust_remote_code reused by the native-template reload (matches transformers). + "trust_remote_code": trust_remote_code, "model": self._model, "tokenizer": self._tokenizer, "processor": self._processor, "is_vision": is_vision, "is_lora": getattr(config, "is_lora", False), + # For a LoRA adapter the native chat template lives on the base model. + "base_model": getattr(config, "base_model", None) + if getattr(config, "is_lora", False) + else None, "is_audio": False, "audio_type": None, "has_audio_input": False, "context_length": runtime_context_length(self._model, max_seq_length), } - # Capture chat_template_info so the worker IPC reply ships it back and - # the route layer classifies capabilities like the other paths. + # Capture chat_template_info for the worker IPC reply and route capability classification. self._populate_chat_template_info(model_name) logger.info("Model %s loaded successfully", model_name) @@ -237,6 +363,9 @@ class MLXInferenceBackend: self._model = None self._tokenizer = None self._processor = None + self._distributed_group = None + self._distributed_rank = 0 + self._distributed_world_size = 1 if self.active_model_name == model_name: self.active_model_name = None gc.collect() @@ -264,12 +393,12 @@ class MLXInferenceBackend: max_new_tokens = 256, repetition_penalty = 1.0, cancel_event = None, - # Reasoning / tool kwargs forwarded by the route + worker; rendered via - # apply_chat_template_for_generation like the transformers path. + # Reasoning / tool kwargs, rendered via apply_chat_template_for_generation (transformers parity). tools = None, enable_thinking = None, reasoning_effort = None, preserve_thinking = None, + presence_penalty = 0.0, ) -> Generator[str, None, None]: if self._model is None: raise RuntimeError("No model loaded") @@ -277,7 +406,6 @@ class MLXInferenceBackend: # Reset so a failed run cannot surface stale stats. self.last_generation_stats = None - # Build messages with system prompt full_messages = [] if system_prompt: full_messages.append({"role": "system", "content": system_prompt}) @@ -294,7 +422,6 @@ class MLXInferenceBackend: {"type": "text", "text": content}, ] elif isinstance(content, list): - # Prepend image if not already present has_image = any( p.get("type") == "image" for p in content if isinstance(p, dict) ) @@ -317,6 +444,7 @@ class MLXInferenceBackend: enable_thinking = enable_thinking, reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, + presence_penalty = presence_penalty, ) else: yield from self._generate_text( @@ -332,6 +460,7 @@ class MLXInferenceBackend: enable_thinking = enable_thinking, reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, + presence_penalty = presence_penalty, ) def _generate_text( @@ -349,12 +478,15 @@ class MLXInferenceBackend: enable_thinking = None, reasoning_effort = None, preserve_thinking = None, + presence_penalty = 0.0, ): from mlx_lm import stream_generate from mlx_lm.sample_utils import make_sampler, make_logits_processors from core.inference.chat_template_helpers import ( apply_chat_template_for_generation, + detect_think_prefill, + render_with_native_template_fallback, ) prompt = apply_chat_template_for_generation( @@ -368,6 +500,34 @@ class MLXInferenceBackend: if prompt is None: raise RuntimeError("apply_chat_template returned None — tokenizer may be incompatible") + # Parity with the transformers backend: if the template dropped the + # requested tools, fall back to the native template so MLX text models + # keep advertising them. self._tokenizer is this entry's tokenizer, so + # probe and native render share a renderer. (VLM renders via the + # processor for image tokens and is not wired here.) + model_info = self.models.get(self.active_model_name, {}) + prompt = render_with_native_template_fallback( + formatted_prompt = prompt, + tokenizer = self._tokenizer, + model_info = model_info, + active_model_name = self.active_model_name, + messages = messages, + tools = tools, + enable_thinking = enable_thinking, + reasoning_effort = reasoning_effort, + preserve_thinking = preserve_thinking, + hf_token = model_info.get("hf_token"), + ) + + # An open prefilled by the template lives in the prompt, not + # the generated tokens; re-emit it so the frontend renders the block. + think_prefix = detect_think_prefill( + prompt, getattr(self._tokenizer, "all_special_tokens", None) + ) + # Emit it before the first token so the block renders during prefill. + if think_prefix: + yield think_prefix + sampler = make_sampler( temp = temperature, top_p = top_p, @@ -375,15 +535,21 @@ class MLXInferenceBackend: min_p = float(min_p or 0.0), min_tokens_to_keep = 1, ) - # Only build a logits processor for a non-trivial repetition penalty. - logits_processors = None + # Repetition and/or presence penalty processors (GGUF/safetensors parity). + logits_processors = [] if repetition_penalty is not None and float(repetition_penalty) not in ( 0.0, 1.0, ): - logits_processors = make_logits_processors( - repetition_penalty = float(repetition_penalty), + logits_processors.extend( + make_logits_processors( + repetition_penalty = float(repetition_penalty), + ) ) + if presence_penalty: + logits_processors.append(_make_mlx_presence_penalty_processor(float(presence_penalty))) + if not logits_processors: + logits_processors = None token_ids = [] logger.info( @@ -410,12 +576,11 @@ class MLXInferenceBackend: ): final_response = response token_ids.append(response.token) - # Decode full sequence with skip_special_tokens cumulative = self._tokenizer.decode( token_ids, skip_special_tokens = True, ) - yield cumulative + yield think_prefix + cumulative if cancel_event and cancel_event.is_set(): break @@ -449,6 +614,7 @@ class MLXInferenceBackend: enable_thinking = None, reasoning_effort = None, preserve_thinking = None, + presence_penalty = 0.0, ): from mlx_vlm import stream_generate as vlm_stream @@ -457,8 +623,7 @@ class MLXInferenceBackend: ) # Pick the chat-template-aware caller: processors with their own - # apply_chat_template + chat_template (e.g. Qwen2.5-VL) use it - # directly; else fall back to the nested tokenizer. + # apply_chat_template + chat_template (e.g. Qwen2.5-VL), else the nested tokenizer. chat_target = self._processor if ( getattr(self._processor, "apply_chat_template", None) is None @@ -479,16 +644,21 @@ class MLXInferenceBackend: # mlx_vlm's stream_generate handles pixel_values (None for text-only) images = [image] if image is not None else None - cumulative = "" + from core.inference.chat_template_helpers import detect_think_prefill + + # Re-emit an open prefill from the prompt (see _generate_text). + cumulative = detect_think_prefill(prompt, getattr(chat_target, "all_special_tokens", None)) + # Emit it before the first token so the block renders during prefill. + if cumulative: + yield cumulative logger.info( "VLM generating: prompt_len=%d, has_image=%s", len(prompt), image is not None, ) - # mlx_vlm.stream_generate forwards **kwargs into generate_step, which - # builds the sampler + logits_processors internally. - # GOTCHA: generate_step expects ``temperature=`` (long form); ``temp=`` - # silently falls into **kwargs and is ignored, stuck at greedy 0.0. + # stream_generate forwards **kwargs into generate_step (builds the + # sampler + logits_processors internally). GOTCHA: generate_step expects + # temperature= (long form); temp= is silently ignored, stuck at greedy 0.0. vlm_kwargs = dict( max_tokens = max_new_tokens, temperature = temperature, @@ -496,10 +666,23 @@ class MLXInferenceBackend: top_k = int(top_k or 0), min_p = float(min_p or 0.0), ) - if repetition_penalty is not None and float(repetition_penalty) not in ( + _rep_active = repetition_penalty is not None and float(repetition_penalty) not in ( 0.0, 1.0, - ): + ) + if presence_penalty: + # Presence needs a custom processor: pass the full list (repetition + + # presence) instead of the repetition_penalty shortcut so both apply. + from mlx_lm.sample_utils import make_logits_processors + + _vlm_processors = [] + if _rep_active: + _vlm_processors.extend( + make_logits_processors(repetition_penalty = float(repetition_penalty)) + ) + _vlm_processors.append(_make_mlx_presence_penalty_processor(float(presence_penalty))) + vlm_kwargs["logits_processors"] = _vlm_processors + elif _rep_active: vlm_kwargs["repetition_penalty"] = float(repetition_penalty) with self._generation_lock: @@ -534,7 +717,7 @@ class MLXInferenceBackend: cancel_event = None, **gen_kwargs, ) -> Generator[str, None, None]: - # MLX LoRA adapter toggling not yet supported — generate normally + # MLX LoRA adapter toggling not yet supported; generate normally yield from self.generate_chat_response(cancel_event = cancel_event, **gen_kwargs) def reset_generation_state(self): diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 5dbd5fb479..4fa0d3ed26 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -45,6 +45,22 @@ _DISPATCH_STOP_TIMEOUT = 5.0 _DISPATCH_IDLE_TIMEOUT = 30.0 _DISPATCH_DRAIN_TIMEOUT = 5.0 +# Max wait for a cancelled generation to release _gen_lock before unload_model +# tears the subprocess down. Only bounds a wedged worker. +_UNLOAD_GEN_LOCK_TIMEOUT = 15.0 + + +class GenStreamError(str): + """A stream chunk carrying a real backend/generation error, not model text. + + Subclasses str so existing display/logging consumers are unaffected, while + callers that must abort a distributed run on error (raise_on_streamed_error) + can distinguish a real error from model output whose visible text starts with + "Error:" by checking isinstance(chunk, GenStreamError). + """ + + __slots__ = () + class InferenceOrchestrator: """ @@ -60,7 +76,13 @@ class InferenceOrchestrator: self._cmd_queue: Any = None self._resp_queue: Any = None self._cancel_event: Any = None # mp.Event — set to cancel generation + # Set for the whole unload; the worker never clears it (unlike _cancel_event), + # so a generate queued behind the cancelled one is skipped, not run. + self._drain_event: Any = None self._gen_lock = threading.Lock() # Serializes generation + # Set during a switch so a generation winning the _gen_lock handoff bails + # instead of starting on the outgoing model. + self._unload_pending = False # Dispatcher state for compare mode (adapter-controlled requests): # bypass _gen_lock, send commands directly, read from per-request @@ -69,6 +91,12 @@ class InferenceOrchestrator: self._mailbox_lock = threading.Lock() self._dispatcher_thread: Optional[threading.Thread] = None self._dispatcher_stop = threading.Event() + # Serializes dispatcher start/stop. _generate_dispatched (compare mode) bypasses + # _gen_lock, so two concurrent compare requests can both reach _start_dispatcher; + # without this lock both could observe no live dispatcher and each spawn one, + # orphaning the extra thread (self._dispatcher_thread tracks only the last). The + # orphan later steals the "unloaded" reply off resp_queue and hangs unload_model. + self._dispatcher_lifecycle_lock = threading.Lock() # Local state mirrors (updated from subprocess responses) self.active_model_name: Optional[str] = None @@ -92,13 +120,11 @@ class InferenceOrchestrator: @property def default_models(self) -> list[str]: - # Wait up to 5s for background HF fetch - self._top_models_ready.wait(timeout = 5) top_gguf = self._top_gguf_cache or [] top_hub = self._top_hub_cache or [] - # Curated static defaults first, then HF download-ranked to backfill. - # Send extras so the frontend keeps 4 per category after removing - # downloaded ones. + # Never wait for the remote Hugging Face ranking during startup. Chat's + # first /api/models/list needs curated defaults immediately; the + # background fetch backfills extra choices on later calls. result: list[str] = [] seen: set[str] = set() for m in self._static_models + top_gguf + top_hub: @@ -159,6 +185,7 @@ class InferenceOrchestrator: self._cmd_queue = _CTX.Queue() self._resp_queue = _CTX.Queue() self._cancel_event = _CTX.Event() + self._drain_event = _CTX.Event() self._proc = _CTX.Process( target = run_without_native_path_secret, @@ -167,6 +194,7 @@ class InferenceOrchestrator: "cmd_queue": self._cmd_queue, "resp_queue": self._resp_queue, "cancel_event": self._cancel_event, + "drain_event": self._drain_event, "config": config, }, daemon = True, @@ -228,6 +256,7 @@ class InferenceOrchestrator: self._cmd_queue = None self._resp_queue = None self._cancel_event = None + self._drain_event = None logger.info("Inference subprocess shut down") def _cleanup(self): @@ -409,6 +438,7 @@ class InferenceOrchestrator: enable_thinking: Optional[bool] = None, reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, + presence_penalty: float = 0.0, ) -> dict: """Build the 'generate' command shared by the locked and dispatched paths.""" cmd = { @@ -423,6 +453,7 @@ class InferenceOrchestrator: "min_p": min_p, "max_new_tokens": max_new_tokens, "repetition_penalty": repetition_penalty, + "presence_penalty": presence_penalty, } # Only forward template kwargs the caller set, for older worker compat. if use_adapter is not None: @@ -456,12 +487,20 @@ class InferenceOrchestrator: cancel ack from that same source so stale events don't leak into the next request. """ + # Latch this stream's subprocess/queue: if a wedged worker is torn down and a + # later load spawns a fresh one, bail rather than re-block on the new queue + # under _gen_lock (deadlock). + initial_proc = self._proc + initial_resp_queue = self._resp_queue while True: + if self._proc is not initial_proc or self._resp_queue is not initial_resp_queue: + yield GenStreamError(f"Error: {self._subprocess_crash_message(crash_context)}") + return resp = read_one(read_timeout) if resp is None: # Check subprocess health if not self._ensure_subprocess_alive(): - yield f"Error: {self._subprocess_crash_message(crash_context)}" + yield GenStreamError(f"Error: {self._subprocess_crash_message(crash_context)}") return continue @@ -471,7 +510,7 @@ class InferenceOrchestrator: # Subprocess-level error (no request_id); request-scoped failures # arrive as gen_error below. if rtype == "error" and not resp.get("request_id"): - yield f"Error: {resp.get('error', 'Unknown error')}" + yield GenStreamError(f"Error: {resp.get('error', 'Unknown error')}") return if rtype == "token": @@ -486,40 +525,63 @@ class InferenceOrchestrator: stats_holder["stats"] = resp.get("stats") return elif rtype == "gen_error": - yield f"Error: {resp.get('error', 'Unknown error')}" + yield GenStreamError(f"Error: {resp.get('error', 'Unknown error')}") return # ------------------------------------------------------------------ # Dispatcher — per-request mailbox routing for compare mode # ------------------------------------------------------------------ - def _start_dispatcher(self) -> None: + def _start_dispatcher(self) -> bool: """Start the dispatcher thread if not already running. The dispatcher reads the shared resp_queue and routes responses to per-request mailbox queues, letting multiple adapter-controlled (compare) requests be in-flight without holding _gen_lock. - """ - if self._dispatcher_thread is not None and self._dispatcher_thread.is_alive(): - return - self._dispatcher_stop.clear() - self._dispatcher_thread = threading.Thread( - target = self._dispatcher_loop, - daemon = True, - name = "inference-dispatcher", - ) - self._dispatcher_thread.start() - logger.debug("Dispatcher thread started") + The whole check-then-spawn runs under _dispatcher_lifecycle_lock so + concurrent compare requests (which bypass _gen_lock) can't both observe + no live dispatcher and each spawn one. Returns True only for the caller + that actually started a new thread; False if one was already alive. + """ + with self._dispatcher_lifecycle_lock: + # Refuse to start while an unload is in progress. unload_model sets + # _unload_pending under this same lock before it stops the idle + # dispatcher, so a start queued behind that stop observes the unload + # here and bails. Without this a fresh dispatcher would be spawned + # after the stop, become the resp_queue reader, and consume the + # worker's "unloaded" reply (unroutable, so dropped) before + # unload_model's _wait_response sees it -- hanging the unload 300s. + if self._unload_pending: + return False + if self._dispatcher_thread is not None and self._dispatcher_thread.is_alive(): + return False + + self._dispatcher_stop.clear() + self._dispatcher_thread = threading.Thread( + target = self._dispatcher_loop, + daemon = True, + name = "inference-dispatcher", + ) + self._dispatcher_thread.start() + logger.debug("Dispatcher thread started") + return True def _stop_dispatcher(self) -> None: - """Signal the dispatcher to stop and wait for it.""" - if self._dispatcher_thread is None: - return - self._dispatcher_stop.set() - self._dispatcher_thread.join(timeout = _DISPATCH_STOP_TIMEOUT) - self._dispatcher_thread = None - logger.debug("Dispatcher thread stopped") + """Signal the dispatcher to stop and wait for it. + + Runs under _dispatcher_lifecycle_lock (paired with _start_dispatcher) so + a stop can't interleave with a concurrent start. Callers must NOT hold + _mailbox_lock here: this joins the dispatcher, and the dispatcher loop + takes _mailbox_lock, so holding it would deadlock the join. + """ + with self._dispatcher_lifecycle_lock: + if self._dispatcher_thread is None: + return + self._dispatcher_stop.set() + self._dispatcher_thread.join(timeout = _DISPATCH_STOP_TIMEOUT) + self._dispatcher_thread = None + logger.debug("Dispatcher thread stopped") def _dispatcher_loop(self) -> None: """Background loop: read resp_queue → route to mailboxes by request_id.""" @@ -581,6 +643,7 @@ class InferenceOrchestrator: reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, stats_holder: Optional[dict] = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Dispatched generation — sends command without holding _gen_lock. @@ -589,15 +652,32 @@ class InferenceOrchestrator: GPU work stays serialized; this only avoids orchestrator lock contention. """ if not self._ensure_subprocess_alive(): - yield "Error: Inference subprocess is not running" + yield GenStreamError("Error: Inference subprocess is not running") return if not self.active_model_name: - yield "Error: No active model" + yield GenStreamError("Error: No active model") + return + # Latch the target model so the recheck below can detect a switch that completed + # between _start_dispatcher and mailbox registration (mirrors the locked path's + # expected_model check). + expected_model = self.active_model_name + + # Switch in flight (unload waiting on _gen_lock). This path bypasses the lock, + # so without this early-out a compare request would enqueue a generate on the + # outgoing model and delay the switch. + if self._unload_pending: + yield GenStreamError("Error: model is being unloaded") return - # Ensure dispatcher is running - self._start_dispatcher() + # Ensure the dispatcher runs. _start_dispatcher serializes concurrent starters under + # _dispatcher_lifecycle_lock and returns True only for the caller that actually spawned + # the thread, so at most one dispatcher ever exists even when two compare requests race + # here. Derive dispatcher_preexisting from that atomic result (not a separate unlocked + # is_alive() read): if THIS call started the dispatcher and then bails on a racing + # unload, it must stop it again (see the unloading bail below). + started = self._start_dispatcher() + dispatcher_preexisting = not started request_id = str(uuid.uuid4()) @@ -617,6 +697,7 @@ class InferenceOrchestrator: min_p = min_p, max_new_tokens = max_new_tokens, repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, use_adapter = use_adapter, tools = tools, enable_thinking = enable_thinking, @@ -624,17 +705,49 @@ class InferenceOrchestrator: preserve_thinking = preserve_thinking, ) - # Create mailbox BEFORE sending command + # Create the mailbox BEFORE sending, rechecking _unload_pending under + # _mailbox_lock: an unload sets _unload_pending before _wait_dispatcher_idle + # reads _mailboxes under the same lock, so either the idle check sees this + # mailbox (and tears the dispatcher down) or we see the unload and bail. + # Registering after would orphan the mailbox and hang the compare stream forever. mailbox: queue.Queue = queue.Queue() with self._mailbox_lock: - self._mailboxes[request_id] = mailbox + # _unload_pending alone is not enough: an unload that ran fully since + # _start_dispatcher clears it in its finally and stops the dispatcher, so it + # reads False here though the dispatcher is gone and the model swapped. Also + # bail when the active model changed or the dispatcher died: a mailbox with no + # dispatcher to route gen_done/gen_error hangs the compare stream. + dispatcher_alive = ( + self._dispatcher_thread is not None and self._dispatcher_thread.is_alive() + ) + unloading = ( + self._unload_pending + or self.active_model_name != expected_model + or not dispatcher_alive + ) + if not unloading: + self._mailboxes[request_id] = mailbox + # When bailing without a mailbox, note whether any OTHER compare request still + # routes through the dispatcher; if none and this call started it, stop it below. + orphaned_dispatcher = unloading and not dispatcher_preexisting and not self._mailboxes + if unloading: + # A racing unload can pass its _wait_dispatcher_idle() while the dispatcher was + # stopped, then set _unload_pending. The one we just started would otherwise + # linger with no mailboxes, race unload_model's _wait_response for the "unloaded" + # reply off resp_queue, and drop it as unroutable -- hanging the unload 300s. Stop + # it here so the unload stays the sole resp_queue reader. Outside _mailbox_lock: + # _stop_dispatcher joins the dispatcher, which itself takes that lock. + if orphaned_dispatcher: + self._stop_dispatcher() + yield GenStreamError("Error: model is being unloaded") + return try: self._send_cmd(cmd) except RuntimeError as exc: with self._mailbox_lock: self._mailboxes.pop(request_id, None) - yield f"Error: {exc}" + yield GenStreamError(f"Error: {exc}") return def read_mailbox(timeout): @@ -676,14 +789,18 @@ class InferenceOrchestrator: return logger.warning("Timed out draining mailbox after cancel") - def _wait_dispatcher_idle(self) -> None: + def _wait_dispatcher_idle(self) -> bool: """Wait for all dispatched requests to complete, then stop dispatcher. - Called by _generate_inner before the _gen_lock path so the dispatcher - thread isn't competing for resp_queue reads. + Returns True if the dispatcher was stopped (all mailboxes drained, or no + dispatcher was running), and False if it was left running because compare + requests were still active after _DISPATCH_IDLE_TIMEOUT. + + Called before the _gen_lock path so the dispatcher thread isn't competing + for resp_queue reads. """ if self._dispatcher_thread is None or not self._dispatcher_thread.is_alive(): - return + return True # Wait for all mailboxes to be emptied (dispatched requests complete) deadline = time.monotonic() + _DISPATCH_IDLE_TIMEOUT @@ -704,8 +821,62 @@ class InferenceOrchestrator: "leaving dispatcher running for compare requests", len(self._mailboxes), ) - else: - self._stop_dispatcher() + return False + self._stop_dispatcher() + return True + + def share_distributed_object( + self, + obj, + timeout: Optional[float] = 300.0, + ): + """Share a small object through the worker's MLX distributed group.""" + if not self._ensure_subprocess_alive(): + raise RuntimeError("Inference subprocess is not running") + + self._wait_dispatcher_idle() + with self._mailbox_lock: + if self._mailboxes: + raise RuntimeError( + "Cannot share distributed objects while compare requests are active" + ) + request_id = str(uuid.uuid4()) + cmd = { + "type": "share_object", + "request_id": request_id, + "object": obj, + } + + with self._gen_lock: + self._send_cmd(cmd) + deadline = None if timeout is None else time.monotonic() + timeout + while deadline is None or time.monotonic() < deadline: + remaining = 1.0 if deadline is None else max(0.1, deadline - time.monotonic()) + resp = self._read_resp(timeout = min(remaining, 1.0)) + if resp is None: + if not self._ensure_subprocess_alive(): + raise RuntimeError(self._subprocess_crash_message("sharing chat turn")) + continue + + rtype = resp.get("type", "") + rid = resp.get("request_id") + if rid and rid != request_id: + logger.debug( + "Skipping response for request_id=%s while sharing request_id=%s", + rid, + request_id, + ) + continue + if rtype == "shared": + return resp.get("object") + if rtype == "share_error": + raise RuntimeError(resp.get("error", "Failed to share object")) + if rtype == "error": + raise RuntimeError(resp.get("error", "Subprocess error")) + if rtype == "status": + continue + + raise RuntimeError("Timeout waiting for distributed object share") # ------------------------------------------------------------------ # Public API — same interface as InferenceBackend @@ -722,6 +893,8 @@ class InferenceOrchestrator: approved_remote_code_fingerprint: Optional[str] = None, gpu_ids: Optional[list[int]] = None, subject: Optional[str] = None, + tensor_parallel: bool = False, + mlx_distributed: bool = False, ) -> bool: """Load a model for inference. @@ -747,6 +920,11 @@ class InferenceOrchestrator: "approved_remote_code_fingerprint": approved_remote_code_fingerprint, "subject": subject, "gpu_ids": gpu_ids, + "tensor_parallel": bool(tensor_parallel), + "mlx_distributed": bool(mlx_distributed), + "mlx_parallel_mode": ("tensor" if tensor_parallel else "pipeline") + if mlx_distributed + else None, } resolved_gpu_ids, gpu_selection = prepare_gpu_selection( gpu_ids, @@ -772,6 +950,19 @@ class InferenceOrchestrator: ) for attempt in range(2): + # Stop-loading (/unload -> cancel_load) aborts a load by discarding this + # model's loading marker. cancel_load only kills a live child; if the cancel + # lands before any child exists (GPU placement, or between retries) there is + # nothing to kill, and without this check the loop would spawn a worker and + # load the model after /unload reported it unloaded. Observe removal and stop. + if model_name not in self.loading_models: + logger.info( + "Load for '%s' was cancelled before spawn; not starting a worker", + model_name, + ) + self.active_model_name = None + self.models.clear() + return False logger.info( "Spawning fresh inference subprocess for '%s' " "(transformers %s.x, attempt %d/2%s)", @@ -783,6 +974,22 @@ class InferenceOrchestrator: sub_config["disable_xet"] = disable_xet self._spawn_subprocess(sub_config) + # A cancel can land after the pre-spawn recheck but while _spawn_subprocess + # is still creating the queues/process. cancel_load runs off the lifecycle + # gate, so its _shutdown_subprocess can see _proc still None and no-op, + # orphaning this fresh worker; the load would then wait for "loaded" and + # publish a model /unload reported unloaded, over a live subprocess nothing + # reaps. Recheck now the child exists and tear it down before publishing. + if model_name not in self.loading_models: + logger.info( + "Load for '%s' was cancelled during spawn; tearing the worker down", + model_name, + ) + self._shutdown_subprocess(timeout = 5) + self.active_model_name = None + self.models.clear() + return False + try: resp = self._wait_response("loaded") except DownloadStallError: @@ -803,8 +1010,31 @@ class InferenceOrchestrator: ) if resp.get("success"): + # A cancel can land while we were parked in _wait_response above. + # cancel_load (off the lifecycle gate) discards this model's loading + # marker BEFORE its teardown, so a Stop-loading that fired after the + # worker queued "loaded" (which we can still consume during cancel_load's + # shutdown window) shows up here only as the marker's removal. Without + # this recheck we would publish active_model_name/models for a model + # /unload reported cancelled, over a subprocess cancel_load just killed; + # its post-teardown re-clear cannot undo a publish that lands after it + # returns. Observe the removal and abort; cancel_load owns teardown. + if model_name not in self.loading_models: + logger.info( + "Load for '%s' was cancelled while waiting for 'loaded'; " + "not publishing the cancelled model", + model_name, + ) + self.active_model_name = None + self.models.clear() + return False model_info = resp.get("model_info", {}) self.active_model_name = model_info.get("identifier", model_name) + # A load always spawns a fresh subprocess holding only this model, so + # mirror that. A lingering stale name would pass unload_model's "not in + # self.models" guard, and the worker's absent-name fallback would unload + # its *active* model, not the already-gone one. + self.models = {} self.models[self.active_model_name] = { "is_vision": model_info.get("is_vision", False), "is_lora": model_info.get("is_lora", False), @@ -837,17 +1067,65 @@ class InferenceOrchestrator: self.models.clear() raise - def unload_model(self, model_name: str) -> bool: - """Unload a model from the subprocess.""" - if model_name in self.loading_models: - logger.info( - "Cancelling in-flight load for model '%s' by terminating subprocess", + def cancel_load(self, model_name: str) -> bool: + """Abort an in-flight load by terminating its subprocess. + + Returns True if a load for ``model_name`` (matched case-insensitively) was + cancelled, False if nothing was loading under that name. This only tears the + loading subprocess down -- it sends no command to a worker -- so, unlike the + rest of ``unload_model``, it is safe to run WITHOUT the inference lifecycle + gate. ``/unload`` calls it off-gate so the "stop loading" button can interrupt + a safetensors load that holds the gate for its whole (multi-minute) duration; + a gated cancel could never preempt that load. + """ + target = model_name + if target not in self.loading_models: + target = next( + (m for m in self.loading_models if m.lower() == model_name.lower()), model_name, ) - self._shutdown_subprocess(timeout = 0.5) - self.loading_models.discard(model_name) - self.active_model_name = None - self.models.clear() + if target not in self.loading_models: + return False + logger.info( + "Cancelling in-flight load for model '%s' by terminating subprocess", + target, + ) + # Discard the loading marker (and clear local state) BEFORE the teardown, not + # after. cancel_load runs off the lifecycle gate, alongside a load_model that + # rechecks this marker before each spawn. But _shutdown_subprocess can block (~1s + # tearing a live child down and joining the dispatcher), so clearing only after + # leaves a window where load_model reads the marker still set, passes its pre-spawn + # recheck, and loads the model after /unload reported it cancelled. Clear first. + self.loading_models.discard(target) + self.active_model_name = None + self.models.clear() + self._shutdown_subprocess(timeout = 0.5) + # Clear the local mirrors again AFTER the teardown. A racing off-gate load_model + # may still be parked in _wait_response("loaded"): its worker already queued a + # "loaded" reply, so during the shutdown window above (the 0.5s settle before the + # response queue is drained and nulled) that thread can consume it and repopulate + # active_model_name/models, undoing the pre-teardown clear. _shutdown_subprocess + # nulls the queue but not the mirrors, so without this second clear /unload reports + # success while the backend still advertises a killed model. The nulled queue lets + # no further "loaded" through, so re-clearing here wipes any repopulation. + self.active_model_name = None + self.models.clear() + return True + + def unload_model(self, model_name: str) -> bool: + """Unload a model from the subprocess.""" + # active_model_name can differ in case from the client's raw /unload name (the + # load path canonicalizes casing). Match case-insensitively and use the canonical + # spelling so the guard, unload command, and cleanup below hit the loaded model. + if ( + self.active_model_name is not None + and model_name != self.active_model_name + and model_name.lower() == self.active_model_name.lower() + ): + model_name = self.active_model_name + # In-flight load: tear its subprocess down (shared loading-cancel logic; no + # worker command sent). + if self.cancel_load(model_name): return True if not self._ensure_subprocess_alive(): @@ -857,30 +1135,93 @@ class InferenceOrchestrator: self.active_model_name = None return True - try: - self._send_cmd( - { - "type": "unload", - "model_name": model_name, - } - ) - resp = self._wait_response("unloaded") - - # Update local state + # Nothing loaded under this name: don't unload a stale model. The worker falls + # back to unloading its *active* model when the name is absent, so a stale unload + # (lost a race to a concurrent load) would hit the wrong one. + if model_name != self.active_model_name and model_name not in self.models: self.models.pop(model_name, None) - if self.active_model_name == model_name: - self.active_model_name = None - - logger.info("Model '%s' unloaded from subprocess", model_name) return True - except Exception as exc: - logger.error("Error unloading model '%s': %s", model_name, exc) - # Clear local state anyway - self.models.pop(model_name, None) - if self.active_model_name == model_name: - self.active_model_name = None - return False + # The subprocess runs commands sequentially, so a bare unload queues behind a + # running generate (a 2-3 min hang). Cancel first (via the mp.Event the worker + # polls each token), then take _gen_lock as sole resp_queue reader (like GGUF). + # + # Set _unload_pending under _dispatcher_lifecycle_lock so it is ordered ahead of + # the dispatcher stop that _wait_dispatcher_idle runs under the same lock: a + # compare request's _start_dispatcher queued behind that stop then observes the + # unload and refuses to spawn a fresh dispatcher that would eat the "unloaded" + # reply off resp_queue. This is a standalone acquisition (no _gen_lock held yet), + # so it keeps the _gen_lock -> _dispatcher_lifecycle_lock order and can't deadlock. + with self._dispatcher_lifecycle_lock: + self._unload_pending = True + # Cancelling only the running generation isn't enough: the worker clears + # cancel_event at each generate start, so a queued one would clear it and run the + # outgoing model to completion. drain_event, never cleared, makes any generate + # dequeued during the unload skip. + if self._drain_event is not None: + self._drain_event.set() + try: + self._cancel_generation() + acquired = self._gen_lock.acquire(timeout = _UNLOAD_GEN_LOCK_TIMEOUT) + if not acquired: + # Wedged worker: tear the subprocess down to free the GPU (next load respawns). + logger.warning( + "Unload: generation did not yield %.1fs after cancel; " + "shutting the inference subprocess down to free the model", + _UNLOAD_GEN_LOCK_TIMEOUT, + ) + self._shutdown_subprocess(timeout = 5) + self.models.pop(model_name, None) + if self.active_model_name == model_name: + self.active_model_name = None + return True + + try: + # Stop the compare-mode dispatcher so it can't consume the "unloaded" reply + # off resp_queue before we do. A dispatched generation bypasses _gen_lock, so + # a wedged one slips past the acquire above; if the dispatcher is still active + # it owns resp_queue and the queued unload hangs _wait_response behind the + # stuck generate. Mirror the wedged locked path: tear the subprocess down. + if not self._wait_dispatcher_idle(): + logger.warning( + "Unload: compare-mode dispatcher still active after idle " + "wait; shutting the inference subprocess down to free the model" + ) + self._shutdown_subprocess(timeout = 5) + self.models.pop(model_name, None) + if self.active_model_name == model_name: + self.active_model_name = None + return True + # Drop stale tokens so they can't be read as the unload reply. + self._drain_queue() + self._send_cmd( + { + "type": "unload", + "model_name": model_name, + } + ) + self._wait_response("unloaded") + + self.models.pop(model_name, None) + if self.active_model_name == model_name: + self.active_model_name = None + + logger.info("Model '%s' unloaded from subprocess", model_name) + return True + + except Exception as exc: + logger.error("Error unloading model '%s': %s", model_name, exc) + # Clear local state anyway + self.models.pop(model_name, None) + if self.active_model_name == model_name: + self.active_model_name = None + return False + finally: + self._gen_lock.release() + finally: + self._unload_pending = False + if self._drain_event is not None: + self._drain_event.clear() def generate_chat_response( self, @@ -899,6 +1240,7 @@ class InferenceOrchestrator: reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, stats_holder: Optional[dict] = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Generate response, streaming tokens from subprocess. @@ -908,6 +1250,8 @@ class InferenceOrchestrator: ``stats_holder``: caller-owned dict; on gen_done its "stats" key gets the worker's usage/timings. Request-scoped to avoid cross-stream reads. + + ``presence_penalty`` matches the GGUF sampling path (0 disables it). """ yield from self._generate_inner( messages = messages, @@ -926,6 +1270,7 @@ class InferenceOrchestrator: reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, stats_holder = stats_holder, + presence_penalty = presence_penalty, ) def generate_chat_completion_with_tools( @@ -945,6 +1290,7 @@ class InferenceOrchestrator: preserve_thinking: Optional[bool] = None, max_tool_iterations: int = 25, auto_heal_tool_calls: bool = True, + nudge_tool_calls: Optional[bool] = None, tool_call_timeout: int = 300, session_id: Optional[str] = None, rag_scope: Optional[dict] = None, @@ -952,6 +1298,7 @@ class InferenceOrchestrator: bypass_permissions: bool = False, use_adapter: Optional[Union[bool, str]] = None, stats_holder: Optional[dict] = None, + presence_penalty: float = 0.0, **_unused, ): """Run the safetensors agentic tool loop in the parent process, @@ -987,6 +1334,7 @@ class InferenceOrchestrator: preserve_thinking = preserve_thinking, # last turn wins, like the GGUF tool loop stats_holder = stats_holder, + presence_penalty = presence_penalty, ) if use_adapter is not None: yield from self.generate_with_adapter_control( @@ -1007,6 +1355,7 @@ class InferenceOrchestrator: execute_tool = execute_tool, cancel_event = cancel_event, auto_heal_tool_calls = auto_heal_tool_calls, + nudge_tool_calls = nudge_tool_calls, max_tool_iterations = max_tool_iterations, tool_call_timeout = tool_call_timeout, session_id = session_id, @@ -1053,6 +1402,7 @@ class InferenceOrchestrator: reasoning_effort: Optional[str] = None, preserve_thinking: Optional[bool] = None, stats_holder: Optional[dict] = None, + presence_penalty: float = 0.0, ) -> Generator[str, None, None]: """Inner generation logic — sends command to subprocess, yields tokens. @@ -1060,12 +1410,13 @@ class InferenceOrchestrator: readers don't consume each other's tokens off the shared resp_queue. """ if not self._ensure_subprocess_alive(): - yield "Error: Inference subprocess is not running" + yield GenStreamError("Error: Inference subprocess is not running") return if not self.active_model_name: - yield "Error: No active model" + yield GenStreamError("Error: No active model") return + expected_model = self.active_model_name # Drain any prior compare-mode dispatcher so we can read resp_queue. self._wait_dispatcher_idle() @@ -1074,6 +1425,14 @@ class InferenceOrchestrator: # consume and drop each other's token events. Hold _gen_lock across the # cmd build + send + whole stream so we stay the sole resp_queue reader. with self._gen_lock: + # Recheck under the lock: an unload we raced may have cleared/swapped the model. + # _unload_pending resets after the lock releases, so it can read False by now; + # the active-model check catches that handoff and a reload that swapped models, + # so we never generate on the wrong one. + if self._unload_pending or self.active_model_name != expected_model: + # Won the lock handoff during a switch; don't start on the outgoing model. + yield GenStreamError("Error: model is being unloaded") + return request_id = str(uuid.uuid4()) image_b64 = self._pil_to_base64(image) if image is not None else None cmd = self._build_generate_cmd( @@ -1087,6 +1446,7 @@ class InferenceOrchestrator: min_p = min_p, max_new_tokens = max_new_tokens, repetition_penalty = repetition_penalty, + presence_penalty = presence_penalty, use_adapter = use_adapter, tools = tools, enable_thinking = enable_thinking, @@ -1097,7 +1457,7 @@ class InferenceOrchestrator: try: self._send_cmd(cmd) except RuntimeError as exc: - yield f"Error: {exc}" + yield GenStreamError(f"Error: {exc}") return yield from self._consume_token_stream( @@ -1141,53 +1501,62 @@ class InferenceOrchestrator: raise RuntimeError("Inference subprocess is not running") if not self.active_model_name: raise RuntimeError("No active model") + expected_model = self.active_model_name - request_id = str(uuid.uuid4()) + # Serialize under _gen_lock (sole resp_queue reader) and refuse to start on the + # outgoing model once an unload is pending, like the text and audio-input paths. + # Without this a concurrent /audio/generate could run TTS on a model being switched. + with self._gen_lock: + # Recheck under the lock (see _generate_inner): a raced unload/switch may have + # cleared or swapped the model while we waited. + if self._unload_pending or self.active_model_name != expected_model: + raise RuntimeError("model is being unloaded") - cmd = { - "type": "generate_audio", - "request_id": request_id, - "text": text, - "temperature": temperature, - "top_p": top_p, - "top_k": top_k, - "min_p": min_p, - "max_new_tokens": max_new_tokens, - "repetition_penalty": repetition_penalty, - } - if use_adapter is not None: - cmd["use_adapter"] = use_adapter + request_id = str(uuid.uuid4()) - self._send_cmd(cmd) + cmd = { + "type": "generate_audio", + "request_id": request_id, + "text": text, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k, + "min_p": min_p, + "max_new_tokens": max_new_tokens, + "repetition_penalty": repetition_penalty, + } + if use_adapter is not None: + cmd["use_adapter"] = use_adapter - # Wait for audio_done or audio_error - deadline = time.monotonic() + 120.0 - while time.monotonic() < deadline: - remaining = max(0.1, deadline - time.monotonic()) - resp = self._read_resp(timeout = min(remaining, 1.0)) + self._send_cmd(cmd) - if resp is None: - if not self._ensure_subprocess_alive(): - raise RuntimeError(self._subprocess_crash_message("audio generation")) - continue + deadline = time.monotonic() + 120.0 + while time.monotonic() < deadline: + remaining = max(0.1, deadline - time.monotonic()) + resp = self._read_resp(timeout = min(remaining, 1.0)) - rtype = resp.get("type", "") + if resp is None: + if not self._ensure_subprocess_alive(): + raise RuntimeError(self._subprocess_crash_message("audio generation")) + continue - if rtype == "audio_done": - wav_bytes = base64.b64decode(resp["wav_base64"]) - sample_rate = resp["sample_rate"] - return wav_bytes, sample_rate + rtype = resp.get("type", "") - if rtype == "audio_error": - raise RuntimeError(resp.get("error", "Audio generation failed")) + if rtype == "audio_done": + wav_bytes = base64.b64decode(resp["wav_base64"]) + sample_rate = resp["sample_rate"] + return wav_bytes, sample_rate - if rtype == "error": - raise RuntimeError(resp.get("error", "Unknown error")) + if rtype == "audio_error": + raise RuntimeError(resp.get("error", "Audio generation failed")) - if rtype == "status": - continue + if rtype == "error": + raise RuntimeError(resp.get("error", "Unknown error")) - raise RuntimeError("Timeout waiting for audio generation (120s)") + if rtype == "status": + continue + + raise RuntimeError("Timeout waiting for audio generation (120s)") def generate_whisper_response( self, @@ -1247,13 +1616,20 @@ class InferenceOrchestrator: ) -> Generator[str, None, None]: """Shared inner logic for audio input generation (Whisper + ASR).""" if not self._ensure_subprocess_alive(): - yield "Error: Inference subprocess is not running" + yield GenStreamError("Error: Inference subprocess is not running") return if not self.active_model_name: - yield "Error: No active model" + yield GenStreamError("Error: No active model") return + expected_model = self.active_model_name with self._gen_lock: + # Recheck under the lock (see _generate_inner): a raced unload/switch may have + # cleared or swapped the model while we waited. + if self._unload_pending or self.active_model_name != expected_model: + # Won the lock handoff during a switch; don't start on the outgoing model. + yield GenStreamError("Error: model is being unloaded") + return request_id = str(uuid.uuid4()) # numpy array -> list for mp.Queue serialization @@ -1279,7 +1655,7 @@ class InferenceOrchestrator: try: self._send_cmd(cmd) except RuntimeError as exc: - yield f"Error: {exc}" + yield GenStreamError(f"Error: {exc}") return yield from self._consume_token_stream( diff --git a/studio/backend/core/inference/passthrough_healing.py b/studio/backend/core/inference/passthrough_healing.py index c73134b4a2..a444431f8d 100644 --- a/studio/backend/core/inference/passthrough_healing.py +++ b/studio/backend/core/inference/passthrough_healing.py @@ -29,10 +29,25 @@ 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 +# Only the formats this healer's parser can promote -- narrower than the loops' +# broader TOOL_XML_SIGNALS. A loop-only marker (Llama <|python_tag|>, bare +# [ARGS]) would buffer a streamed call as prose without promoting it, so keep a +# healer-aligned list. Mistral's [TOOL_CALLS] IS promotable, so it stays in. +_HEAL_SIGNALS = ( + "", + "<|tool_call>", + " bool: + return any(s in text for s in _HEAL_SIGNALS) + + # Read once at import (same convention as the other UNSLOTH_* switches). _HEALING_DISABLED = os.environ.get("UNSLOTH_DISABLE_TOOL_CALL_HEALING", "0") == "1" # Nudging is OPT-IN: per-request nudge_tool_calls=true, or flip the process @@ -44,7 +59,7 @@ 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) +_MAX_SIGNAL_LEN = max(len(s) for s in _HEAL_SIGNALS) # A suspected-but-unclosed tool block larger than this is declared a false # alarm and flushed, bounding memory on a model rambling XML-lookalike text. _MAX_HOLD_CHARS = 64 * 1024 @@ -198,7 +213,7 @@ def heal_openai_message_events( 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): + if not isinstance(content, str) or not _has_heal_signal(content): return None parsed, spans = parse_tool_calls_from_text(content, allow_incomplete = True, with_spans = True) tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None @@ -248,7 +263,7 @@ def heal_openai_message( def _earliest_signal(buffer: str) -> int: best = -1 - for signal in TOOL_XML_SIGNALS: + for signal in _HEAL_SIGNALS: index = buffer.find(signal) if index >= 0 and (best < 0 or index < best): best = index @@ -275,7 +290,7 @@ 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): + if any(signal.startswith(tail) for signal in _HEAL_SIGNALS): return length return 0 @@ -340,9 +355,10 @@ class StreamToolCallHealer: 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). + # HOLD: drain the first contiguous run per pass so events keep document + # order (a later declared call must not overtake an earlier undeclared one + # flushing as text). A run is one markup call OR a whole Mistral [TOOL_CALLS] + # array of contiguous spans, so later calls in it are not stranded as text. parsed, spans = parse_tool_calls_from_text( self._buffer, id_offset = self._id_offset, @@ -363,26 +379,32 @@ class StreamToolCallHealer: 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:] + pos = 0 + run_end = spans[0][1] + for order, (call, (start, end)) in enumerate(zip(parsed, spans)): + # Stop at the first gap or incomplete trailing block: leave it for the + # next pass to re-hold and stream incrementally, not flush as text early. + if order and start != run_end: + break + promoted = _promote( + [call], + self._allowed, + id_offset = self._id_offset, + tool_schemas = self._tool_schemas, + ) + if promoted: + # Flush any leading text, then drop the promoted markup span. + if self._buffer[pos:start]: + events.append(("text", self._buffer[pos:start])) + events.append(("tool_call", promoted[0])) + self._id_offset += 1 + else: + # Undeclared/unusable name: markup is DATA, flush it (and prior text) verbatim. + events.append(("text", self._buffer[pos:end])) + pos = end + run_end = end + # Everything past the drained run (later blocks) stays and is rescanned. + self._buffer = self._buffer[run_end:] self._holding = False def finalize(self) -> list: @@ -508,7 +530,7 @@ def nudge_should_retry( 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): + if not isinstance(text, str) or not _has_heal_signal(text): return False return not _heal_would_promote(text, allowed_tools, tools) diff --git a/studio/backend/core/inference/presence_penalty.py b/studio/backend/core/inference/presence_penalty.py new file mode 100644 index 0000000000..c73c513887 --- /dev/null +++ b/studio/backend/core/inference/presence_penalty.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Presence-penalty logits helpers for the safetensors/MLX inference paths. + +Kept in a dependency-light leaf module (torch + transformers only, no unsloth / +peft) so the pure logic can be imported and unit-tested without pulling in the +full inference backend. ``core.inference.inference`` re-exports these for the +runtime generate paths. +""" + +import torch + + +def apply_presence_penalty(input_ids, scores, penalty: float, prompt_len: int): + """OpenAI/llama.cpp presence penalty: subtract ``penalty`` once per distinct + completion token (positions >= prompt_len; prompt excluded, multiplicity + ignored, negatives raise). In place; zero is a no-op.""" + if not penalty: + return scores + vocab_size = scores.shape[-1] + for b in range(input_ids.shape[0]): + generated = input_ids[b, prompt_len:] + if generated.numel() == 0: + continue + seen = torch.unique(generated) + # Bound generated ids to the valid range [0, vocab_size). Real completion + # tokens are always in range, so this is a zero-regression safety net that + # drops any stray out-of-range or negative id before indexing (mirrors the + # MLX path's bound). Filtering both ends avoids indexing scores with a + # negative id (which would silently wrap to the wrong row). + seen = seen[(seen >= 0) & (seen < vocab_size)] + if seen.numel(): + scores[b, seen] = scores[b, seen] - penalty + return scores + + +def _make_presence_penalty_processor(penalty: float, prompt_len: int): + """``LogitsProcessorList`` for ``apply_presence_penalty``; ``None`` at zero penalty (generate call stays byte-identical).""" + if not penalty: + return None + from transformers import LogitsProcessor, LogitsProcessorList + + class _PresencePenaltyLogitsProcessor(LogitsProcessor): + @torch.no_grad() + def __call__(self, input_ids, scores): + return apply_presence_penalty(input_ids, scores, penalty, prompt_len) + + return LogitsProcessorList([_PresencePenaltyLogitsProcessor()]) diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 0c96378d6c..81c25b777e 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -14,6 +14,7 @@ parses tool calls from the cumulative text and dispatches via ``core.inference.tools``. """ +import bisect import re import threading from typing import Callable, Generator, Optional @@ -21,14 +22,38 @@ from typing import Callable, Generator, Optional from loggers import get_logger from core.inference.tool_call_parser import ( - _TOOL_ALL_PATS, + _GEMMA_BARE_TC_PREFIX_RE, + _GEMMA_BARE_TC_RE, + _TOOL_ALL_PATS as _PARSER_TOOL_ALL_PATS, + _TOOL_CLOSED_PATS as _PARSER_TOOL_CLOSED_PATS, + _balanced_brace_end, + _strip_function_xml_calls, + _strip_gemma_wrapperless_calls, + _strip_glm_calls, + _strip_mistral_closed_calls, + _strip_mistral_reasoning, BUDGET_EXHAUSTED_NUDGE, + MAX_ACT_REPROMPTS, RAG_MAX_SEARCHES_PER_TURN, RAG_SEARCH_CAP_NUDGE, TOOL_XML_SIGNALS, + is_short_intent_without_action, parse_tool_calls_from_text, + reprompt_to_act_message, + strip_leading_bare_json_call, + strip_llama3_leading_sentinels, strip_tool_markup, ) + +# The healer owns the bracket-tag + rehearsal strip helpers and their name-gated +# pattern lists, so the safetensors streaming strip stays aligned with the parser. +from core.tool_healing import ( + _REHEARSAL_TAIL_STRIP_RE, + _strip_bracket_tag_calls, + _think_spans_outside_tool_markup, + apply_tool_strip_patterns, + strip_outside_think, +) from core.inference.tool_loop_controller import ( ToolLoopController, coerce_tool_arguments, @@ -50,19 +75,213 @@ logger = get_logger(__name__) # Buffer cap while disambiguating a possible tool-call prefix. _MAX_BUFFER_CHARS = 32 +# Memory bound for holding a leading bare-JSON object whose top-level "{" never balances. +_MAX_BARE_JSON_BUFFER = 16384 + + +# No grammar constraint here (unlike llama-server's lazy grammar): collapse +# exact-duplicate calls and cap the count so a runaway turn cannot fan out. +_MAX_TOOL_CALLS_PER_TURN = 8 + + +def _active_tool_names(active_tools: list[dict]) -> list[str]: + names = [ + (tool.get("function") or {}).get("name") + for tool in active_tools + if isinstance(tool, dict) and isinstance(tool.get("function"), dict) + ] + return [name for name in names if name] + + +def _active_tool_names(active_tools: list[dict]) -> list[str]: + names = [ + (tool.get("function") or {}).get("name") + for tool in active_tools + if isinstance(tool, dict) and isinstance(tool.get("function"), dict) + ] + return [name for name in names if name] + + +# Unrestricted mode has no tool list, so any identifier may open a NAME[ARGS] rehearsal; +# ``[`` and each ARGS letter stay optional so a chunk split after ``NAME[`` is still held. +_UNRESTRICTED_REHEARSAL_RE = re.compile(r"[\w-]+(?:\[(?:A(?:R(?:G(?:S)?)?)?)?)?") + + +def _is_rehearsal_prefix( + stripped: str, + active_tools: list[dict], + *, + unrestricted: bool = False, +) -> bool: + """True if ``stripped`` is a (possibly partial) prefix of a ``NAME[ARGS]`` + rehearsal split across chunks (``web_search`` then ``[ARGS]{...}``). A space + means prose. Unrestricted mode accepts any identifier; else NAME must be active.""" + if not stripped or any(ch.isspace() for ch in stripped): + return False + if unrestricted: + return _UNRESTRICTED_REHEARSAL_RE.fullmatch(stripped) is not None + for name in _active_tool_names(active_tools): + if stripped == name or f"{name}[ARGS]".startswith(stripped): + return True + return False + + +def _held_rehearsal_tail_len( + text: str, + active_tools: list[dict], + *, + unrestricted: bool = False, +) -> int: + """Length of a trailing bare tool-name token that may be a split rehearsal call + (``...web_search`` with ``[ARGS]{...}`` still to arrive), so STREAMING can hold it + instead of leaking the name. Returns 0 for ordinary prose.""" + i = len(text) + while i > 0 and not text[i - 1].isspace(): + i -= 1 + tail = text[i:] + return ( + len(tail) + if tail and _is_rehearsal_prefix(tail, active_tools, unrestricted = unrestricted) + else 0 + ) + + +def _rehearsal_name_start( + candidate: str, + signal_pos: int, + active_tools: list[dict], + *, + unrestricted: bool = False, +) -> int: + """For an ``[ARGS]`` signal at ``signal_pos``, return the start of the preceding + bare tool-name token (``NAME[ARGS]``), else ``signal_pos`` unchanged when the + signal is not ``[ARGS]`` or NAME is not an active tool (restricted mode).""" + if not candidate.startswith("[ARGS]", signal_pos): + return signal_pos + j = signal_pos + while j > 0 and (candidate[j - 1].isalnum() or candidate[j - 1] in "_-"): + j -= 1 + if j < signal_pos and ( + unrestricted or candidate[j:signal_pos] in _active_tool_names(active_tools) + ): + return j + return signal_pos + + +def _earliest_tool_signal( + candidate: str, + signals, + active_tools: list[dict], + *, + unrestricted: bool = False, +) -> int: + """Index where the turn's first genuine tool-call boundary begins, or -1. + + Non-``[ARGS]`` markup wins on first occurrence. An ``[ARGS]`` hit is a rehearsal + only when an active tool name (any name in unrestricted mode) precedes it, so a + literal ``foo[ARGS]`` in prose is skipped rather than draining the turn; for a + real ``NAME[ARGS]`` the boundary is pulled back to NAME.""" + best = -1 + for sig in signals: + if sig != "[ARGS]": + p = candidate.find(sig) + if p >= 0 and (best < 0 or p < best): + best = p + continue + from_idx = 0 + while True: + p = candidate.find("[ARGS]", from_idx) + if p < 0: + break + name_start = _rehearsal_name_start( + candidate, p, active_tools, unrestricted = unrestricted + ) + if name_start < p: + # Genuine ``NAME[ARGS]``: the boundary is the start of NAME. + if best < 0 or name_start < best: + best = name_start + break + # Bare/prose [ARGS]: skip it so a later real call in the same chunk is still found. + from_idx = p + len("[ARGS]") + return best + + +def _has_genuine_tool_signal( + candidate: str, + signals, + active_tools: list[dict], + *, + unrestricted: bool = False, +) -> bool: + """True when ``candidate`` holds a genuine tool-call boundary for one of ``signals``. + + Non-``[ARGS]`` markers count on a substring hit; an ``[ARGS]`` hit is genuine only + when an active tool name (any in unrestricted mode) precedes it. Mirrors the + ``_earliest_tool_signal`` name-gating so BUFFERING / end-of-stream checks do not + drain inactive-name prose.""" + for sig in signals: + if sig == "[ARGS]": + if ( + _earliest_tool_signal( + candidate, ("[ARGS]",), active_tools, unrestricted = unrestricted + ) + >= 0 + ): + return True + continue + if sig in candidate: + return True + return False + def strip_tool_markup_streaming( text: str, *, auto_heal_tool_calls: bool = True, tool_protocol_active: bool = False, + enabled_tool_names: Optional[set] = None, ) -> str: - """Strip open-ended tool XML from display text without trimming whitespace.""" + """Strip open-ended tool XML from display text without trimming whitespace. + + Mirrors the parser-side ``strip_tool_markup`` segment scan (minus the final trim) so + streaming and final display agree: balanced strips first (nested JSON removed whole), + then the guarded function-XML / GLM scans that close at each call's REAL terminator so + literal markup inside argument values is data and trailing prose survives. Reasoning + ```` / ``[THINK]`` blocks are preserved verbatim (a rehearsed call inside one must + not be deleted, else the cumulative text shrinks then regrows). ``enabled_tool_names`` + keeps an inactive-name ``foo[ARGS]{..}`` / ``call:NAME{..}`` example visible (it is prose, + not a call), matching the parse / detection active-tool gate.""" if not (auto_heal_tool_calls or tool_protocol_active): return text - for pat in _TOOL_ALL_PATS: - text = pat.sub("", text) - return text + + # Drop a leading Magistral ``[THINK]...[/THINK]`` block (bracket reasoning form, not the + # ```` channel) so raw reasoning does not leak into streamed display; an unclosed + # leading block is held (dropped to EOF) until its closer streams in. + text = _strip_mistral_reasoning(text) + + def _seg(segment: str, is_last: bool) -> str: + # Same scan order as the parser's _strip_segment (seg_final -> is_last): balanced + # strips first, then the guarded function-XML / GLM scans, then the regex arms + # (DeepSeek / Kimi / closed forms). EOS-anchored tail arms run only on the last + # segment (a bare ``foo[ARGS]`` before is prose). Rehearsal strips are name-gated. + seg = _strip_mistral_closed_calls(segment) + seg = _strip_bracket_tag_calls(seg, enabled_tool_names = enabled_tool_names) + if is_last: + seg = _strip_gemma_wrapperless_calls(seg, enabled_tool_names) + seg = _strip_function_xml_calls(seg, final = is_last) + seg = _strip_glm_calls(seg, final = is_last) + pats = _PARSER_TOOL_ALL_PATS if is_last else _PARSER_TOOL_CLOSED_PATS + for pat in pats: + seg = pat.sub("", seg) + if is_last: + seg = apply_tool_strip_patterns( + seg, [_REHEARSAL_TAIL_STRIP_RE], enabled_tool_names = enabled_tool_names + ) + return seg + + # Preserve think blocks verbatim: stripping a rehearsed call inside one shrinks then + # regrows the cumulative text, corrupting append-by-length consumers. + return strip_outside_think(text, _seg) def _strip_tool_markup_final( @@ -70,10 +289,11 @@ def _strip_tool_markup_final( *, auto_heal_tool_calls: bool, tool_protocol_active: bool = False, + enabled_tool_names: Optional[set] = None, ) -> str: if not (auto_heal_tool_calls or tool_protocol_active): return text - return strip_tool_markup(text, final = True) + return strip_tool_markup(text, final = True, enabled_tool_names = enabled_tool_names) def _status_for_tool(tool_name: str, arguments: dict) -> str: @@ -81,25 +301,76 @@ def _status_for_tool(tool_name: str, arguments: dict) -> str: return status_for_tool(tool_name, arguments) +def _looks_like_enabled_bare_json(text: str, enabled_tool_names: Optional[set]) -> bool: + """True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False.""" + probe = strip_llama3_leading_sentinels(text.lstrip()) + if not (probe.startswith("{") and ('"name"' in probe or '"function"' in probe)): + return False + return strip_leading_bare_json_call(probe, enabled_tool_names) != probe + + _FUNCTION_SIGNAL_RE = re.compile(r"") _TOOL_CALL_NAME_RE = re.compile(r'"name"\s*:\s*"([\w-]+)"') +# Mistral name/v11 and rehearsal forms, aligned with the parser so the provisional +# render-html card fires for bracket-tag serializations too. +_MISTRAL_RENDER_NAME_RE = re.compile( + r"\[TOOL_CALLS\]\s*([\w-]+)(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*(?=\{)" +) +_REHEARSAL_RENDER_NAME_RE = re.compile(r"(? bool: - """Return True when the first drained tool call is clearly render_html.""" - function_match = _FUNCTION_SIGNAL_RE.search(content) - tool_call_index = content.find("") - if not function_match and tool_call_index < 0: + """Return True when the FIRST tool call in ``content`` is clearly render_html. + + Covers every serialization the loop executes (XML ```` / ````, + Mistral ``[TOOL_CALLS]``, rehearsal ``NAME[ARGS]``); the earliest marker wins so a + render_html marker inside another call's argument is treated as data. Markers inside + a ```` / ``[THINK]`` block are dropped since the parser skips them.""" + think_spans = _think_spans_outside_tool_markup(content) + _think_starts = [s for s, _e in think_spans] + + def _in_think(pos: int) -> bool: + if not think_spans: + return False + i = bisect.bisect_right(_think_starts, pos) - 1 + return i >= 0 and think_spans[i][0] <= pos < think_spans[i][1] + + def _first_outside(start: int, finder) -> int: + # First occurrence at/after ``start`` that is not inside a think span. + pos = finder(start) + while pos >= 0 and _in_think(pos): + pos = finder(pos + 1) + return pos + + candidates: list[tuple[int, str]] = [] + for fm in _FUNCTION_SIGNAL_RE.finditer(content): + if not _in_think(fm.start()): + candidates.append((fm.start(), fm.group(1))) + break + tc = _first_outside(0, lambda i: content.find("", i)) + if tc >= 0: + nm = _TOOL_CALL_NAME_RE.search(content[tc:]) + candidates.append((tc, nm.group(1) if nm else "")) + mt = _first_outside(0, lambda i: content.find("[TOOL_CALLS]", i)) + if mt >= 0: + mm = _MISTRAL_RENDER_NAME_RE.match(content, mt) + if mm: + candidates.append((mt, mm.group(1))) + else: + # Array shape: a bare ``"name"`` search can latch onto an argument key, so resolve the + # first call through the parser (it reads top-level names). + arr_calls = parse_tool_calls_from_text(content[mt:]) + if arr_calls: + candidates.append((mt, (arr_calls[0].get("function") or {}).get("name") or "")) + for rm in _REHEARSAL_RENDER_NAME_RE.finditer(content): + if not _in_think(rm.start(1)): + candidates.append((rm.start(1), rm.group(1))) + break + + if not candidates: return False - - if function_match and (tool_call_index < 0 or function_match.start() < tool_call_index): - return function_match.group(1) == "render_html" - - if tool_call_index >= 0: - name_match = _TOOL_CALL_NAME_RE.search(content[tool_call_index:]) - return bool(name_match and name_match.group(1) == "render_html") - - return False + _pos, name = min(candidates, key = lambda c: c[0]) + return name == "render_html" def _coerce_arguments_with_provenance( @@ -149,6 +420,7 @@ def run_safetensors_tool_loop( execute_tool: Callable[..., str], cancel_event: Optional[threading.Event] = None, auto_heal_tool_calls: bool = True, + nudge_tool_calls: Optional[bool] = None, max_tool_iterations: int = 25, tool_call_timeout: int = 300, session_id: Optional[str] = None, @@ -188,8 +460,17 @@ def run_safetensors_tool_loop( for _ev in _auto["events"]: yield _ev conversation.extend(_auto["messages"]) + # Autoinject ran a KB search outside the controller, so it counts as an + # executed tool for the plan-without-action gate. + rag_autoinjected = bool(_auto) unrestricted_tools = not tools + # Gate telling a genuine NAME[ARGS] rehearsal from inactive-name prose; built from the + # ORIGINAL tools list so a spent one-shot still reads as a tool name. None = unrestricted. + _enabled_names_gate = None if unrestricted_tools else set(_active_tool_names(tools)) + # Detection must see the same names as the strip gate (ORIGINAL list, incl. a spent + # one-shot), else its repeat is stripped but never drained and the turn ends blank. + _detect_tools = [] if unrestricted_tools else list(tools or []) tool_controller = ToolLoopController( tools = None if unrestricted_tools else tools, auto_heal_tool_calls = auto_heal_tool_calls, @@ -198,6 +479,14 @@ def run_safetensors_tool_loop( kb_search_count = 0 final_attempt_done = False next_call_id = 0 + reprompt_count = 0 + # A denied tool confirmation must not be answered with a plan-without-action + # re-prompt (which would raise the confirmation gate again). + tool_denied = False + # Real tool-call turns completed. Only turns that actually executed a tool count + # against ``max_tool_iterations``; a duplicate/disabled no-op correction turn (and a + # plan-without-action re-prompt) must not consume budget, matching the GGUF loop. + _executed_tool_iters = 0 def _tool_succeeded(tool_name: str) -> bool: key_prefix = f"{tool_name}:" @@ -215,9 +504,13 @@ def run_safetensors_tool_loop( _state_streaming = 1 _state_draining = 2 - for iteration in range(max_tool_iterations + 1): + # Reserve re-prompt slots so they don't eat the caller's tool budget. + _extra_iters = MAX_ACT_REPROMPTS if max_tool_iterations > 0 else 0 + for iteration in range(max_tool_iterations + _extra_iters + 1): if cancel_event is not None and cancel_event.is_set(): return + # Whether this turn ran a tool; a no-op-only turn stays False and doesn't consume budget. + _turn_executed_real_tool = False if final_attempt_done: active_tools: list[dict] = [] @@ -229,6 +522,8 @@ def run_safetensors_tool_loop( tool_protocol_active = not final_attempt_done and (unrestricted_tools or bool(active_tools)) tool_xml_signals = TOOL_XML_SIGNALS if tool_protocol_active else () + # Gate the markerless bare-JSON form on enabled names so an ordinary JSON answer isn't misread as a call. + _enabled_tool_names = None if unrestricted_tools else set(_active_tool_names(active_tools)) detect_state = _state_buffering content_buffer = "" @@ -304,17 +599,18 @@ def run_safetensors_tool_loop( if detect_state == _state_streaming: candidate = cumulative_display + delta - signal_pos = -1 - for sig in tool_xml_signals: - p = candidate.find(sig) - if p >= 0 and (signal_pos < 0 or p < signal_pos): - signal_pos = p + # Earliest genuine boundary: bare [ARGS] in prose is skipped; a real NAME[ARGS] is + # pulled back to NAME so the name is not flushed. + signal_pos = _earliest_tool_signal( + candidate, tool_xml_signals, _detect_tools, unrestricted = unrestricted_tools + ) if signal_pos >= 0: before_tool = candidate[:signal_pos] cleaned_before = strip_tool_markup_streaming( before_tool, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = tool_protocol_active, + enabled_tool_names = _enabled_names_gate, ) if len(cleaned_before) > len(last_emitted): last_emitted = cleaned_before @@ -345,10 +641,20 @@ def run_safetensors_tool_loop( cumulative_display, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = tool_protocol_active, + enabled_tool_names = _enabled_names_gate, ) - if len(cleaned) > len(last_emitted): - last_emitted = cleaned - yield {"type": "content", "text": cleaned} + # Hold a trailing bare active-tool-name (split rehearsal) until its [ARGS] arrives; + # released by later prose or the end-of-stream flush. + if tool_protocol_active: + _hold = _held_rehearsal_tail_len( + cleaned, _detect_tools, unrestricted = unrestricted_tools + ) + emit = cleaned[: len(cleaned) - _hold] if _hold else cleaned + else: + emit = cleaned + if len(emit) > len(last_emitted): + last_emitted = emit + yield {"type": "content", "text": emit} continue # BUFFERING: hold until we know it is not a tool call. @@ -366,6 +672,92 @@ def run_safetensors_tool_loop( if sig.startswith(stripped): is_prefix = True break + # Bracket-tag forms arrive mid-buffer, so substring-check too (mirrors GGUF); [ARGS] + # counts only with an active NAME so prose is not drained into a no-op. + if sig == "[ARGS]": + if ( + _earliest_tool_signal( + stripped, + ("[ARGS]",), + _detect_tools, + unrestricted = unrestricted_tools, + ) + >= 0 + ): + is_match = True + break + elif sig.startswith("[") and sig in stripped: + is_match = True + break + + # Split rehearsal: hold the bare name until its [ARGS] arrives and matches above. + is_rehearsal_prefix = False + if ( + not is_match + and not is_prefix + and tool_protocol_active + and _is_rehearsal_prefix(stripped, _detect_tools, unrestricted = unrestricted_tools) + ): + is_prefix = True + is_rehearsal_prefix = True + + # Llama-3.2 ``custom_tools`` emits a bare ``{"name":..,"parameters":..}`` with no XML + # signal. Hold a leading ``{`` (after any sentinel) until it closes: drain if it parses + # as a call, else stream as content. Non-call text is always recovered downstream. + bare_probe = strip_llama3_leading_sentinels(stripped) + if ( + not is_match + and not is_prefix + and tool_protocol_active + and bare_probe.startswith("{") + ): + if _balanced_brace_end(bare_probe, 0) is None: + if len(stripped) < _MAX_BARE_JSON_BUFFER: + continue # object still open -- keep buffering + elif _looks_like_enabled_bare_json(bare_probe, _enabled_tool_names): + # Oversized still-open ENABLED-tool call: stop holding (memory bound) but + # DRAIN instead of leaking the raw prefix; a giant ordinary JSON answer still streams. + detect_state = _state_draining + continue + elif parse_tool_calls_from_text( + content_buffer, + id_offset = next_call_id, + allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, + ): + # Closed object that parses as a bare-JSON call -- drain silently. + detect_state = _state_draining + continue + # Closed non-call object (or oversized non-call) -- stream as text. + + # Gemma wrapper-less ``call:NAME{...}`` has no tool_xml_signals entry: + # buffer it here or it streams raw until the end-of-turn safety net. + # ``(? len(last_emitted): last_emitted = cleaned @@ -398,7 +791,8 @@ def run_safetensors_tool_loop( "arguments": {}, "provenance": _tool_event_provenance(provisional = True), } - elif is_prefix and len(stripped) < _MAX_BUFFER_CHARS: + elif is_prefix and (is_rehearsal_prefix or len(stripped) < _MAX_BUFFER_CHARS): + # A rehearsal prefix is self-bounded; the buffer cap must not cut long MCP names short. continue else: detect_state = _state_streaming @@ -407,57 +801,121 @@ def run_safetensors_tool_loop( cumulative_display, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = tool_protocol_active, + enabled_tool_names = _enabled_names_gate, ) - if len(cleaned) > len(last_emitted): - last_emitted = cleaned - yield {"type": "content", "text": cleaned} + # Same trailing-name hold as STREAMING for this first flush out of BUFFERING. + if tool_protocol_active: + _hold = _held_rehearsal_tail_len( + cleaned, _detect_tools, unrestricted = unrestricted_tools + ) + emit = cleaned[: len(cleaned) - _hold] if _hold else cleaned + else: + emit = cleaned + if len(emit) > len(last_emitted): + last_emitted = emit + yield {"type": "content", "text": emit} # Stream finished -- resolve what we collected. if cancel_event is not None and cancel_event.is_set(): return if detect_state == _state_buffering: - # Buffer never resolved -- tool XML or plain content? + # Buffer never resolved: [ARGS] is name-gated so a prose answer with a literal + # ``foo[ARGS]{...}`` is not parsed. stripped = content_buffer.lstrip() + _bare_eos = strip_llama3_leading_sentinels(stripped) if ( stripped and tool_protocol_active - and any(sig in stripped for sig in tool_xml_signals) + and _has_genuine_tool_signal( + stripped, + tool_xml_signals, + _detect_tools, + unrestricted = unrestricted_tools, + ) ): detect_state = _state_draining + elif tool_protocol_active and _looks_like_enabled_bare_json( + _bare_eos, _enabled_tool_names + ): + # A held bare-JSON ENABLED-tool fragment has no XML signal; DRAIN it (an ordinary + # JSON answer falls through to the else and streams as content, GGUF parity). + detect_state = _state_draining else: + # Drain and fall through to STREAMING so the intent re-prompt + safety-net parser + # still fire on short emissions like "Let me search." that never exit BUFFERING. if content_buffer: cumulative_display += content_buffer - yield { - "type": "content", - "text": _strip_tool_markup_final( - cumulative_display, - auto_heal_tool_calls = auto_heal_tool_calls, - tool_protocol_active = False, - ), - } - yield {"type": "status", "text": ""} - return + cleaned = strip_tool_markup( + cumulative_display, final = True, enabled_tool_names = _enabled_tool_names + ) + if len(cleaned) > len(last_emitted): + last_emitted = cleaned + yield {"type": "content", "text": cleaned} + detect_state = _state_streaming if detect_state == _state_streaming: - # No tool detected mid-stream -- check for late tool XML. - safety_tc = None - saw_tool_signal = tool_protocol_active and any( - sig in content_accum for sig in tool_xml_signals + # Run the parser even with no XML signal (the Llama-3.2 bare-JSON form carries none); it's + # strict so plain answers stay untouched. Mirrors GGUF. + safety_tc = parse_tool_calls_from_text( + content_accum, + id_offset = next_call_id, + allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_tool_names, ) - if saw_tool_signal: - safety_tc = parse_tool_calls_from_text( - content_accum, - id_offset = next_call_id, - allow_incomplete = auto_heal_tool_calls, - ) if not safety_tc: - # Final answer: if a literal tool marker in prose was stripped - # during streaming but did not parse as a real call, restore the - # raw cumulative text for core callers. Route-level cleanup can - # still apply the Auto-Heal display policy. - if saw_tool_signal and content_accum: + # Re-prompt once on plan-without-action, before any tool runs + # (GGUF loop parity). The retry is gated on nudge_tool_calls so + # Studio callers (which send True) always nudge, while API callers + # who omit the flag keep today's no-reprompt behavior (opt-in). + stripped_answer = content_accum.strip() + if ( + auto_heal_tool_calls + and nudge_tool_calls + and active_tools + and reprompt_count < MAX_ACT_REPROMPTS + and not rag_autoinjected + and not tool_denied + and not any(record.executed for record in tool_controller.history) + and is_short_intent_without_action(stripped_answer) + ): + reprompt_count += 1 + logger.info( + "Safetensors re-prompt %d/%d: model responded without " + "calling tools (%d chars)", + reprompt_count, + MAX_ACT_REPROMPTS, + len(stripped_answer), + ) + conversation.append({"role": "assistant", "content": stripped_answer}) + tool_hint = " or ".join(_active_tool_names(active_tools)) or "an available tool" + conversation.append( + { + "role": "user", + "content": reprompt_to_act_message(tool_hint), + } + ) + # Empty status clears the badge and resets the route's + # per-turn text cursor before the re-prompted turn streams. + yield {"type": "status", "text": ""} + continue + + # Final answer. If a literal tool marker in prose was buffered but + # never parsed as a call, restore the raw text so the prose surfaces + # in full; route-level cleanup still applies the Auto-Heal policy. + if content_accum and any(sig in content_accum for sig in tool_xml_signals): yield {"type": "content", "text": content_accum} + else: + # Turn ended as a plain answer (no [ARGS] followed): the held rehearsal tail is real + # prose, release it. + final_clean = strip_tool_markup_streaming( + cumulative_display, + auto_heal_tool_calls = auto_heal_tool_calls, + tool_protocol_active = tool_protocol_active, + enabled_tool_names = _enabled_names_gate, + ) + if len(final_clean) > len(last_emitted): + yield {"type": "content", "text": final_clean} yield {"type": "status", "text": ""} return tool_calls = safety_tc @@ -465,31 +923,41 @@ def run_safetensors_tool_loop( content_accum, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = True, + enabled_tool_names = _enabled_names_gate, ) logger.info( "Safetensors safety net: parsed %d tool call(s) from streamed content", len(tool_calls), ) else: - # DRAINING: parse tool calls out of full content. + # DRAINING: parse tool calls out of full content. Gate the bare rehearsal on the + # ORIGINAL tool list (``_enabled_names_gate``), the same names detection/strip used to + # drain here: a spent one-shot (render_html) is off the active list but its re-emitted + # ``render_html[ARGS]{..}`` must still parse so it routes to the repeat no-op instead of + # being dropped into a blank continuation. tool_calls = parse_tool_calls_from_text( content_accum, id_offset = next_call_id, allow_incomplete = auto_heal_tool_calls, + enabled_tool_names = _enabled_names_gate, ) if not tool_calls: # Parser found nothing. Auto-Heal-enabled display cleanup # strips unparseable tool XML; disabled Auto-Heal preserves # the raw text so literal/malformed markup stays visible. if content_accum: - yield { - "type": "content", - "text": _strip_tool_markup_final( - content_accum, - auto_heal_tool_calls = auto_heal_tool_calls, - tool_protocol_active = False, - ), - } + _drain_text = _strip_tool_markup_final( + content_accum, + auto_heal_tool_calls = auto_heal_tool_calls, + tool_protocol_active = False, + enabled_tool_names = _enabled_tool_names, + ) + # Drained bare-JSON call that didn't parse: with Auto-Heal on, drop the fragment + # (plain JSON answers are left untouched); off keeps it visible per the strict contract. + if tool_protocol_active and auto_heal_tool_calls: + _drain_text = strip_leading_bare_json_call(_drain_text, _enabled_tool_names) + if _drain_text: + yield {"type": "content", "text": _drain_text} if provisional_render_html_started and not provisional_resolved: provisional_resolved = True yield { @@ -505,10 +973,14 @@ def run_safetensors_tool_loop( content_accum, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = True, + enabled_tool_names = _enabled_names_gate, ) if tool_calls: next_call_id += len(tool_calls) + # Strip a leading bare-JSON call from the kept content so it isn't replayed as text or + # next-turn history (``_strip_tool_markup_final`` only knows XML). No-op for plain JSON answers. + content_text = strip_leading_bare_json_call(content_text, _enabled_tool_names) if final_attempt_done: # Final-answer turn re-called a tool -- stop the loop. @@ -517,6 +989,27 @@ def run_safetensors_tool_loop( yield {"type": "status", "text": ""} return + # Collapse exact-duplicate calls and cap the count (runaway-turn guard). + if tool_calls: + seen_keys: set = set() + deduped: list = [] + for _tc in tool_calls: + _fn = _tc.get("function", {}) or {} + _key = (_fn.get("name", ""), str(_fn.get("arguments", ""))) + if _key in seen_keys: + continue + seen_keys.add(_key) + deduped.append(_tc) + if len(deduped) >= _MAX_TOOL_CALLS_PER_TURN: + break + if len(deduped) != len(tool_calls): + logger.info( + "Safetensors: collapsed %d repeated tool call(s) in one turn to %d", + len(tool_calls), + len(deduped), + ) + tool_calls = deduped + assistant_msg: dict = {"role": "assistant", "content": content_text} assistant_appended = False @@ -593,6 +1086,7 @@ def run_safetensors_tool_loop( "result": TOOL_REJECTED_MESSAGE, "provenance": decision.provenance, } + tool_denied = True denied_message = { "role": "tool", "name": decision.tool_name, @@ -634,6 +1128,8 @@ def run_safetensors_tool_loop( completion = tool_controller.record_result(decision, result) if provisional_match: provisional_resolved = True + # A tool ran this turn, so it counts against the caller's budget. + _turn_executed_real_tool = True yield completion.tool_end_event() conversation.append(completion.tool_message()) @@ -646,7 +1142,11 @@ def run_safetensors_tool_loop( if not unrestricted_tools and not tool_controller.active_tools(): final_attempt_done = True continue - if iteration + 1 >= max_tool_iterations and not final_attempt_done: + # Count only turns that executed a tool against the cap; a no-op correction turn doesn't + # consume budget so the model gets its nudge and another tool-enabled turn (GGUF parity). + if _turn_executed_real_tool: + _executed_tool_iters += 1 + if _executed_tool_iters >= max_tool_iterations and not final_attempt_done: # Budget exhausted; nudge a final plain answer. final_attempt_done = True conversation.append({"role": "user", "content": BUDGET_EXHAUSTED_NUDGE}) diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index ca3d1e4cbc..1ab1142eba 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -2,39 +2,123 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """ -Backend-neutral tool-call XML parser shared by GGUF and safetensors. -Tolerates missing closing tags in either ``{json}`` -or ``v...`` shape. +Backend-neutral tool-call parser shared by GGUF, safetensors, and MLX, so the +safetensors + MLX agentic loop sees the same call shape llama-server gives GGUF: + + - ``{json}`` (Qwen / Hermes) + - ``v`` (Qwen3.5 xml) + - ``<|python_tag|>NAME.call(k="v", ...)`` (Llama-3 built-in tools) + - ``<|python_tag|>{"name":..., "parameters":...}`` (Llama-3 custom) + - ``{"name":..., "parameters":...}`` (Llama-3.2 bare JSON) + - ``[TOOL_CALLS] [{...}, ...]`` (Mistral v0.3 / Nemo / Small) + - ``[TOOL_CALLS]name{json}`` (Mistral v11+ / Magistral) + - ``[TOOL_CALLS]name[ARGS]{json}`` (Ministral / Mistral Large 3) + - ``<|tool_call>call:NAME{k:<|"|>v<|"|>}`` (Gemma 4) + - ``<｜tool▁calls▁begin｜>...function<｜tool▁sep｜>NAME\\n``\\`\\`\\`json\\n{...}\\n\\`\\`\\`...`` (DeepSeek R1) + - ``<｜tool▁calls▁begin｜>...<｜tool▁call▁begin｜>NAME<｜tool▁sep｜>{json}<｜tool▁call▁end｜>...`` (DeepSeek V3 / V3.1) + - ``NAME\\nk\\nv...`` (GLM 4.5 / 4.6 / 4.7) + - ``<|tool_calls_section_begin|>...<|tool_call_begin|>functions.NAME:IDX<|tool_call_argument_begin|>{json}<|tool_call_end|>...`` (Kimi K2) + +Missing closing tags / brackets are tolerated: models often truncate mid-stream. """ +# Lazy annotations keep the standalone python 3.9 import working. +from __future__ import annotations + +import json +import re +from typing import Any, Optional + +# Qwen/Hermes, Qwen3.5 XML and Gemma 4 live in core.tool_healing; this module adds the rest. from core import tool_healing as _tool_healing -_TOOL_ALL_PATS = _tool_healing._TOOL_ALL_PATS +# Flip the streaming buffer STREAMING->DRAINING so partial markup never leaks. +TOOL_XML_SIGNALS = ( + "", + "", + "[TOOL_CALLS]", + "<|tool_call>", + # Bare reasoning-rehearsal marker (``name[ARGS]{...}``, no leading [TOOL_CALLS]); + # keeps a rehearsed call held in the stream so it is promoted, not leaked as prose. + "[ARGS]", + # DeepSeek R1 / V3 / V3.1 -- 5 opener variants llama.cpp keeps. + "<｜tool▁calls▁begin｜>", + "<｜tool▁call▁begin｜>", + "<｜tool_calls_begin｜>", + "<｜tool▁calls｜>", + "<｜tool calls begin｜>", + "<｜tool\\_calls\\_begin｜>", + # Kimi K2 / Moonshot. + "<|tool_calls_section_begin|>", + "<|tool_call_begin|>", +) -def parse_tool_calls_from_text( - content: str, - *, - id_offset: int = 0, - allow_incomplete: bool = True, -) -> list[dict]: - return _tool_healing.parse_tool_calls_from_text( - content, - id_offset = id_offset, - allow_incomplete = allow_incomplete, - ) +# DeepSeek opener variants; shared by parse and strip so a parsed signal is always stripped. +_DEEPSEEK_OPEN_ALT = ( + r"tool▁calls▁begin|tool_calls_begin|tool calls begin|tool\\_calls\\_begin|tool▁calls" +) +_DEEPSEEK_OPEN_RE_SRC = r"<｜(?:" + _DEEPSEEK_OPEN_ALT + r")｜>" + +# Closed pairs only (mid-stream); _TOOL_ALL_PATS also eats unclosed tails at +# end-of-turn. ``[\w-]+`` on ```` tracks OpenAI's +# ``^[a-zA-Z0-9_-]{1,64}$`` so hyphenated MCP names parse like built-ins. +_TOOL_CLOSED_PATS = [ + re.compile(r".*?", re.DOTALL), + # Span to the real ```` so a literal one inside a value can't truncate the strip. + re.compile( + r'' + r'(?:(?!).)*' + r"", + re.DOTALL, + ), + re.compile(r"<\|tool_call>.*?", re.DOTALL), + re.compile(r"\[TOOL_CALLS\]\s*\[.*?\](?:\s*)?", re.DOTALL), + # Mistral v11+ ``[TOOL_CALLS]name{json}`` (may chain), close at ``}``. + re.compile(r"\[TOOL_CALLS\]\s*[\w\.\-]+\s*(?:\[ARGS\])?\s*\{.*?\}", re.DOTALL), + # DeepSeek R1 / V3 / V3.1: full envelope (any opener variant) ... end. + re.compile(_DEEPSEEK_OPEN_RE_SRC + r".*?<｜tool▁calls▁end｜>", re.DOTALL), + # Kimi K2: ``<|tool_calls_section_begin|>...<|tool_calls_section_end|>``. + re.compile(r"<\|tool_calls_section_begin\|>.*?<\|tool_calls_section_end\|>", re.DOTALL), + # Kimi K2 section-less closed call; else the catch-all below eats trailing prose to EOS. + re.compile(r"<\|tool_call_begin\|>.*?<\|tool_call_end\|>", re.DOTALL), +] +_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ + re.compile(r".*$", re.DOTALL), + re.compile(r'.*$', re.DOTALL), + # Bare-word markers drop a trailing truncated call only when a call-shaped start + # follows; a prose mention (``See [TOOL_CALLS] docs...``) keeps its tail. Bare marker at EOF drops. + re.compile(r"<\|tool_call>(?=\s*call\s*:|\s*$).*$", re.DOTALL), + re.compile( + r"\[TOOL_CALLS\](?=\s*(?:[\[{]|[A-Za-z_][\w.\-]*(?:[\[{]|\s*$))|\s*$).*$", + re.DOTALL, + ), + re.compile( + r"<\|python_tag\|>(?=\s*(?:\{|[A-Za-z_][\w.]*\()|\s*$).*$", + re.DOTALL, + ), + # DeepSeek envelopes truncated mid-stream (any opener); same call-shaped lookahead as above. + re.compile( + _DEEPSEEK_OPEN_RE_SRC + r"(?=\s*(?:<｜tool▁call▁begin｜>|function)|\s*$).*$", + re.DOTALL, + ), + re.compile(r"<｜tool▁call▁begin｜>(?=\s*function|\s*$).*$", re.DOTALL), + # Kimi K2 envelope truncated. + re.compile( + r"<\|tool_calls_section_begin\|>(?=\s*<\|tool_call_begin\|>|\s*$).*$", + re.DOTALL, + ), + re.compile( + r"<\|tool_call_begin\|>(?=\s*[A-Za-z_][\w.\-]*:\d|\s*$).*$", + re.DOTALL, + ), + # Gemma wrapper-less ``call:NAME{...}`` is handled by ``_strip_gemma_wrapperless_calls`` (enabled-name gate). +] -def strip_tool_markup(text: str, *, final: bool = False) -> str: - return _tool_healing.strip_tool_call_markup(text, final = final) - - -# Prefixes the streaming buffer watches for to gate in-progress text. -TOOL_XML_SIGNALS = ("", "<|tool_call>", " bool: + stripped = text.strip() + return 0 < len(stripped) < REPROMPT_MAX_CHARS and INTENT_SIGNAL.search(stripped) is not None + + +def reprompt_to_act_message(tool_hint: str) -> str: + """The user message appended when re-prompting a plan-without-action turn.""" + return ( + "You have access to enabled tools. If a tool is needed to satisfy " + "the user's request or complete the action you described, call " + f"{tool_hint} now. If no tool is needed, provide the final answer " + "and follow the user's requested format." + ) + + +# Qwen / Hermes ``{json}``. +_TC_JSON_START_RE = re.compile(r"\s*\{") +# Qwen3.5 ```` and the attribute form ```` +# (MiniCPM-5, MiniMax-M2); name class ``[\w.\-]+`` lands in group(1) or group(2). +_TC_FUNC_START_RE = re.compile(r'\s*') +# Body ends at ```` (Hermes) or ```` (Qwen3.5 / MiniCPM-5) +# so it stops at the close even when prose follows (else prose leaked into args). +_TC_END_TAG_RE = re.compile(r"") +_TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") +# Horizontal whitespace only (``[^\S\n]*``, not ``\s*``) so the wrapping newline + +# first-line indentation survive; ``_trim_param_value`` trims one newline, preserving +# code indentation (SGLang qwen3_coder). +_TC_PARAM_START_RE = re.compile( + r'<(?:parameter|param)(?:=([\w\.\-]+)|\s+name="([\w\.\-]+)")>[^\S\n]*' +) +_TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") + +# Llama-3 ``<|python_tag|>NAME.call(...)``. +_LLAMA3_PYTHON_TAG = "<|python_tag|>" +_LLAMA3_PY_CALL_RE = re.compile( + r"<\|python_tag\|>\s*([\w\.\-]+)\s*\.\s*call\s*\(", +) +# Anchored at a fixed offset (char after ``<|python_tag|>``) plus the ``; NAME.call(`` +# chain separator; fixed-offset (not a free scan) ignores ``.call(`` inside JSON args. +_LLAMA3_PY_CALL_HEAD_RE = re.compile(r"\s*([\w\.\-]+)\s*\.\s*call\s*\(") +_LLAMA3_CALL_CHAIN_RE = re.compile(r"\s*;\s*([\w\.\-]+)\s*\.\s*call\s*\(") +# Llama-3 ``.call(k=v)`` kwarg tokens, hand-scanned below (not finditer) to stay +# linear on a truncated body; finditer retries every offset of a long run (ReDoS). +_LLAMA3_KEY_RE = re.compile(r"\w+") +_LLAMA3_WS_RE = re.compile(r"\s*") +# ints, decimals (1.5, 1., .5) and sci notation; trailing ``(?![\w.])`` stops a token +# like ``1.2.3`` being truncated to ``1.2`` (which would mis-parse the remainder). +_LLAMA3_NUM_RE = re.compile(r"-?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?(?![\w.])") +_LLAMA3_LIT_RE = re.compile(r"true|false|null") + +# Mistral ``[TOOL_CALLS]`` trigger. v11+ chains them, each followed by a bare name +# plus ``{json}`` (Magistral) or ``[ARGS]{json}`` (Ministral / Large 3). +_MISTRAL_TRIGGER = "[TOOL_CALLS]" +_MISTRAL_ARGS_MARKER = "[ARGS]" +# Mistral Small 3.2 emits ``name[CALL_ID][ARGS]{json}`` (absent on Ministral / +# Magistral); llama.cpp distinguishes the two on ``[CALL_ID]`` (common/chat.cpp). +_MISTRAL_CALL_ID_MARKER = "[CALL_ID]" +# Magistral wraps reasoning in ``[THINK]...[/THINK]``; a ``[TOOL_CALLS]`` inside +# that block is chain-of-thought, not a real call. +_MISTRAL_THINK_OPEN = "[THINK]" +_MISTRAL_THINK_CLOSE = "[/THINK]" +_MISTRAL_V11_NAME_RE = re.compile(r"\s*([\w\.\-]+)\s*") + +# DeepSeek markers (full-width pipe U+FF5C, block U+2581); five outer-open variants like llama.cpp. +_DEEPSEEK_BEGIN_RE = re.compile(_DEEPSEEK_OPEN_RE_SRC) +_DEEPSEEK_END = "<｜tool▁calls▁end｜>" +_DEEPSEEK_CALL_BEGIN = "<｜tool▁call▁begin｜>" +_DEEPSEEK_SEP = "<｜tool▁sep｜>" +_DEEPSEEK_CALL_END = "<｜tool▁call▁end｜>" +# R1 wraps args in a ```json fence with a ``function`` prefix; V3/V3.1 do not. +# Scanned with ``str.find`` -- the regex forms are O(N^2) on truncated bodies. +_DEEPSEEK_R1_FUNC_MARKER = "function" + _DEEPSEEK_SEP +_DEEPSEEK_R1_FENCE = "\n```json\n" +_DEEPSEEK_R1_CLOSE_RE = re.compile(r"```[\s\r\n]*" + re.escape(_DEEPSEEK_CALL_END)) + +# GLM 4.5-4.7: ``NAME[\n]K...``; the lookahead also allows a +# direct ````/```` (4.7 drops the newline, zero-arg calls close at once). +# Name class ``[\w.\-]+`` keeps prose like ``not a call`` unparsed; +# ``{`` stays with the Qwen JSON parser. +_GLM_TC_OPEN_RE = re.compile(r"\s*([\w.\-]+)\s*(?=\n||)") +_GLM_TC_CLOSE = "" +_GLM_ARG_KEY_OPEN = "" +_GLM_ARG_KEY_CLOSE = "" +_GLM_ARG_VAL_OPEN = "" +_GLM_ARG_VAL_CLOSE = "" +# Strings arrive raw, non-strings via tojson; only unambiguous JSON literals decode +# (bare ``42``/``true``/``null`` stay strings). +_GLM_JSON_NUMERIC_RE = re.compile(r"-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?") + +# Kimi K2 / Moonshot (ASCII pipes). Id ``functions.NAME:IDX`` -- strip ``functions.``/``:N`` for the name. +_KIMI_SECTION_BEGIN = "<|tool_calls_section_begin|>" +_KIMI_SECTION_END = "<|tool_calls_section_end|>" +_KIMI_CALL_BEGIN = "<|tool_call_begin|>" +_KIMI_ARG_BEGIN = "<|tool_call_argument_begin|>" +_KIMI_CALL_END = "<|tool_call_end|>" +_KIMI_ID_RE = re.compile(r"^(?:functions\.)?([\w\.\-]+)(?::(\d+))?$") + +# Gemma 4: ``<|tool_call>call:NAME{...}``, ``<|"|>`` wraps strings. +_GEMMA_TC_RE = re.compile(r"<\|tool_call>\s*call\s*:\s*([\w\.\-]+)\s*\{") +_GEMMA_STR_BEGIN = '<|"|>' +_GEMMA_STR_END = '<|"|>' +_GEMMA_TC_END = "" + +# skip_special_tokens strips the wrapper and ``<|"|>`` markers, so streamed Gemma calls +# arrive as bare ``call:NAME{k:v, ...}``; ``(? int | None: + """Index of the ``]`` matching ``[`` at ``text[start]`` (ignores brackets in JSON strings).""" + if start >= len(text) or text[start] != "[": + return None + depth = 0 + in_string = False + esc = False + i = start + while i < len(text): + ch = text[i] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_string = False + else: + if ch == '"': + in_string = True + elif ch == "[": + depth += 1 + elif ch == "]": + depth -= 1 + if depth == 0: + return i + i += 1 + return None + + +def _skip_mistral_call_id(text: str, pos: int) -> int: + """Skip an optional ``[CALL_ID]`` (Mistral Small 3.2); return the next token pos.""" + n = len(text) + i = pos + while i < n and text[i] in " \t\n\r": + i += 1 + if not text.startswith(_MISTRAL_CALL_ID_MARKER, i): + return pos + i += len(_MISTRAL_CALL_ID_MARKER) + while i < n and text[i] in " \t\n\r": + i += 1 + # The id is a short opaque token; stop at whitespace or the next marker. + while i < n and text[i] not in " \t\n\r[{": + i += 1 + while i < n and text[i] in " \t\n\r": + i += 1 + return i + + +def _strip_mistral_reasoning(content: str) -> str: + """Drop a leading Magistral ``[THINK]...[/THINK]`` so a ``[TOOL_CALLS]`` inside + reasoning is not taken as a real call; an unclosed ``[THINK]`` drops from it on.""" + i = 0 + n = len(content) + while i < n and content[i] in " \t\n\r": + i += 1 + if not content.startswith(_MISTRAL_THINK_OPEN, i): + return content + close = content.find(_MISTRAL_THINK_CLOSE, i + len(_MISTRAL_THINK_OPEN)) + if close == -1: + return content[:i] + return content[:i] + content[close + len(_MISTRAL_THINK_CLOSE) :] + + +def _strip_mistral_closed_calls(text: str) -> str: + """Strip cleanly-closed ``[TOOL_CALLS]`` blocks (array, ``name{json}``, + ``name[ARGS]{json}``) via balanced scanning -- a non-greedy ``\\{.*?\\}`` would + truncate at the first ``}`` and lose nested JSON. Unclosed runs are left for + ``final=True`` cleanup.""" + n = len(text) + out = [] + cursor = 0 + while cursor < n: + idx = text.find(_MISTRAL_TRIGGER, cursor) + if idx == -1: + out.append(text[cursor:]) + break + out.append(text[cursor:idx]) + body_start = idx + len(_MISTRAL_TRIGGER) + i = body_start + while i < n and text[i] in " \t\n\r": + i += 1 + # Array shape: ``[TOOL_CALLS] [...]``. + if i < n and text[i] == "[": + end = _balanced_bracket_end(text, i) + if end is None: + # Truncated; let caller buffer / final-strip. + out.append(text[idx:]) + break + cursor = end + 1 + if text.startswith("", cursor): + cursor += len("") + continue + # Single-object shape ``[TOOL_CALLS] { json }`` (no name/array): the parser + # accepts it, so the display strip must remove it too (else it leaks). + if i < n and text[i] == "{": + end = _balanced_brace_end(text, i) + if end is None: + out.append(text[idx:]) + break + cursor = end + 1 + if text.startswith("", cursor): + cursor += len("") + continue + # Named shape: ``[TOOL_CALLS] name [ARGS]? { json }``. + name_match = _MISTRAL_V11_NAME_RE.match(text, i) + if not name_match: + out.append(text[idx:body_start]) + cursor = body_start + continue + i = name_match.end() + while i < n and text[i] in " \t\n\r": + i += 1 + i = _skip_mistral_call_id(text, i) + if text.startswith(_MISTRAL_ARGS_MARKER, i): + i += len(_MISTRAL_ARGS_MARKER) + while i < n and text[i] in " \t\n\r": + i += 1 + if i >= n or text[i] != "{": + out.append(text[idx:i]) + cursor = i + continue + end = _balanced_brace_end(text, i) + if end is None: + out.append(text[idx:]) + break + cursor = end + 1 + # Consume the optional EOS marker too, mirroring the array shape, so a + # ``[TOOL_CALLS]name{json}`` tail doesn't leave ```` as content. + if text.startswith("", cursor): + cursor += len("") + return "".join(out) + + +def _strip_gemma_wrapperless_calls(text: str, enabled_tool_names: Optional[set] = None) -> str: + """Strip closed wrapper-less Gemma ``call:NAME{...}`` calls with balanced brace + scanning (nested arguments are removed whole). ``enabled_tool_names`` gates the + strip like the parser gate: a disabled/example name stays visible; ``None`` + strips every closed call.""" + if _whole_content_is_json_value(text): + return text + n = len(text) + out = [] + # Mirror the parse scan: a leading JSON answer's span is data, kept visible. + cursor = _leading_json_value_end(text) or 0 + if cursor: + out.append(text[:cursor]) + while cursor < n: + m = _GEMMA_BARE_TC_RE.search(text, cursor) + if not m: + out.append(text[cursor:]) + break + disabled = enabled_tool_names is not None and m.group(1) not in enabled_tool_names + brace = m.end() - 1 # _GEMMA_BARE_TC_RE consumes through the opening ``{`` + # Same boundary scanner as the parser: strip exactly what it consumed. + end = _gemma_body_brace_end(text, brace) + closed = end is not None + next_index = (end + 1) if closed else len(text) + if not closed: + # Unclosed call: drop an enabled call to EOS; keep a disabled/example name as prose. + out.append(text[cursor:] if disabled else text[cursor : m.start()]) + break + if disabled: + # Disabled/example name is prose: keep it whole. + out.append(text[cursor:next_index]) + else: + out.append(text[cursor : m.start()]) + cursor = next_index # already past the matching ``}`` + return "".join(out) + + +_FUNC_CLOSE_TAG_RE = re.compile(r"") + + +def _strip_function_xml_calls(text: str, *, final: bool) -> str: + """Strip ```` calls by mirroring the parser: an opener inside an open ```` is data and each call closes at its first ```` that is not parameter data; ``final`` drops a trailing unclosed call.""" + starts = [ + m for m in _TC_FUNC_START_RE.finditer(text) if not _inside_open_parameter(text, m.start()) + ] + if not starts: + return text + out: list[str] = [] + pos = 0 + for idx, m in enumerate(starts): + if m.start() < pos: + continue # opener already inside a previously consumed call span + out.append(text[pos : m.start()]) + next_start = starts[idx + 1].start() if idx + 1 < len(starts) else len(text) + close = None + for cm in _FUNC_CLOSE_TAG_RE.finditer(text, m.end(), next_start): + if not _inside_open_parameter(text, cm.start()): + close = cm # first close that is not parameter data = the real close + break + if close is not None: + pos = close.end() + elif final: + pos = len(text) # trailing unclosed call -- drop to EOF + else: + out.append(text[m.start() :]) # keep the unclosed call buffered mid-stream + pos = len(text) + break + out.append(text[pos:]) + return "".join(out) + + +def _glm_value_close( + text: str, + vs: int, + *, + strict: bool = False, +) -> int: + """Index of the ```` that really ends the GLM value at ``vs``: the + first one whose next non-space token is ````, ```` or + end-of-text AND that sits at balanced quote state (an embedded literal pair + like ``print("")`` lives inside a still-open string). + Quote openers are contextual (single quote only after punctuation, so + apostrophes are prose; double quote also at word start), mirroring the Gemma + scanners. If no candidate balances, the first token-valid one wins -- except + in ``strict`` mode (Auto-Heal off), which refuses the in-quote fallback rather + than execute truncated arguments. Returns -1 if unclosed.""" + n = len(text) + search = vs + first_candidate = -1 + quote = "" + prev = ":" + prev_raw = ":" + qpos = vs # quote-state cursor; advanced incrementally to each candidate + while True: + ve = text.find(_GLM_ARG_VAL_CLOSE, search) + if ve < 0: + return -1 if strict else first_candidate + j = ve + len(_GLM_ARG_VAL_CLOSE) + while j < n and text[j] in " \t\r\n": + j += 1 + if j >= n or text.startswith(_GLM_ARG_KEY_OPEN, j) or text.startswith(_GLM_TC_CLOSE, j): + while qpos < ve: + ch = text[qpos] + if quote: + if ch == "\\" and qpos + 1 < ve: + qpos += 2 + continue + if ch == quote: + quote = "" + elif ch in "\"'" and (prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())): + quote = ch + if not ch.isspace(): + prev = ch + prev_raw = ch + qpos += 1 + if not quote: + return ve + if first_candidate < 0: + first_candidate = ve + search = ve + len(_GLM_ARG_VAL_CLOSE) + + +def _strip_glm_calls(text: str, *, final: bool) -> str: + """Strip GLM 4.x calls by scanning to each call's REAL ```` (the one + after the last consumed ````, mirroring ``_parse_glm_tool_calls``), so + a literal ```` inside a value is data. Qwen ``{json}`` has + no NAME token and is left to the regex arms. ``final`` drops a truncated call to + EOS; otherwise it stays buffered.""" + out: list[str] = [] + cursor = 0 + n = len(text) + while True: + m = _GLM_TC_OPEN_RE.search(text, cursor) + if not m: + break + apos = m.end() + close = -1 + while True: + ks = text.find(_GLM_ARG_KEY_OPEN, apos) + tc = text.find(_GLM_TC_CLOSE, apos) + if tc >= 0 and (ks < 0 or tc < ks): + close = tc + break + if ks < 0: + break # no close and no more keys -- truncated body + ke = text.find(_GLM_ARG_KEY_CLOSE, ks + len(_GLM_ARG_KEY_OPEN)) + if ke < 0: + break + vstart = ke + len(_GLM_ARG_KEY_CLOSE) + while vstart < n and text[vstart] in " \t\r\n": + vstart += 1 + if not text.startswith(_GLM_ARG_VAL_OPEN, vstart): + apos = ke + len(_GLM_ARG_KEY_CLOSE) + continue + vs = vstart + len(_GLM_ARG_VAL_OPEN) + ve = _glm_value_close(text, vs) + if ve < 0: + break # unclosed -- truncated + apos = ve + len(_GLM_ARG_VAL_CLOSE) + if close >= 0: + out.append(text[cursor : m.start()]) + cursor = close + len(_GLM_TC_CLOSE) + continue + # Truncated GLM call (no real close yet). + if final: + out.append(text[cursor : m.start()]) + cursor = n + # Non-final: leave the unclosed call (and any tail) buffered as-is. + break + out.append(text[cursor:]) + return "".join(out) + + +def strip_tool_markup( + text: str, + *, + final: bool = False, + enabled_tool_names: Optional[set] = None, +) -> str: + """Strip tool-call markup. ``final=False`` keeps in-progress markup buffered; + ``final=True`` also drops trailing unclosed runs and trims. + + ``enabled_tool_names`` gates the name-conditioned forms so a disabled/example name in + prose is kept (mirrors the parser gate): the bare reasoning-rehearsal ``name[ARGS]{...}`` + and the markerless Gemma ``call:NAME{...}`` strip. ``None`` strips every closed call. + """ + if final: + # Drop a leading Magistral ``[THINK]...[/THINK]`` at end-of-turn; its bracket + # form is not the ```` the reasoning channel renders. + text = _strip_mistral_reasoning(text) + + def _strip_segment(segment: str, is_last: bool) -> str: + seg_final = final and is_last + seg = _strip_mistral_closed_calls(segment) + # Bare reasoning-rehearsal ``name[ARGS]{json}`` and the Mistral name form promote through + # the shared balanced scan, so strip them the same way (any nesting depth removed whole). + # The rehearsal arm is name-gated: an inactive ``foo[ARGS]{..}`` is prose and is kept. + seg = _tool_healing._strip_bracket_tag_calls(seg, enabled_tool_names = enabled_tool_names) + if seg_final: + # Markerless Gemma ``call:NAME{...}`` (name-gated, mirrors the parse gate); end-of-turn only. + seg = _strip_gemma_wrapperless_calls(seg, enabled_tool_names) + # Scan-strip the function-XML form (parser-accurate: a literal ```` in a + # value is data, not a call); the regex arms below cover the other formats. + seg = _strip_function_xml_calls(seg, final = seg_final) + # GLM 4.x: scan to the call's real so a literal one inside a value is data, + # not a leak. Qwen {json} is left to the regex arms. + seg = _strip_glm_calls(seg, final = seg_final) + pats = _TOOL_ALL_PATS if seg_final else _TOOL_CLOSED_PATS + for pat in pats: + seg = pat.sub("", seg) + if seg_final: + # Drop a trailing partial bare rehearsal (``name[ARGS]`` with a truncated or absent + # body) the balanced scan cannot close; gated so prose ``foo[ARGS] ...`` survives. + seg = _tool_healing.apply_tool_strip_patterns( + seg, + [_tool_healing._REHEARSAL_TAIL_STRIP_RE], + enabled_tool_names = enabled_tool_names, + ) + return seg + + # ```` / ``[THINK]`` reasoning is preserved verbatim (a rehearsed call inside it is + # not executed, so it must not be stripped from display either); a literal think marker + # inside a real call's arguments is that call's data and is stripped with the call. + result = _tool_healing.strip_outside_think(text, _strip_segment) + return result.strip() if final else result + + def has_tool_signal(text: str) -> bool: - """Return True if ``text`` contains any tool-call XML signal.""" return any(s in text for s in TOOL_XML_SIGNALS) + + +# A Qwen/Hermes ````/```` envelope whose arguments carry literal +# DeepSeek/Kimi markers must parse as the OUTER call. Detect it opening before the first +# marker so the pre-pass skips it. +_EMBEDDED_MARKER_RE = re.compile( + _DEEPSEEK_OPEN_RE_SRC + "|" + re.escape(_KIMI_SECTION_BEGIN) + "|" + re.escape(_KIMI_CALL_BEGIN) +) +# Covers ```` and the attribute form. ``<|python_tag|>`` is Llama-3's +# envelope too (built-in ``NAME.call(`` and custom ``{json}``), so a quoted DeepSeek/Kimi +# example is data; the call-shaped lookahead mirrors the ``_TOOL_ALL_PATS`` python_tag arm +# so a bare prose ``<|python_tag|>`` mention isn't treated as one. +_OUTER_ENVELOPE_OPEN_RE = re.compile( + r'|' + r"|<\|python_tag\|>(?=\s*(?:\{|[A-Za-z_][\w.]*\())" +) +# CLOSED outer envelopes, each spanning to its REAL final close so a literal +# ````/```` inside a value is data. Wrapped Gemma counts too. +_OUTER_ENVELOPE_CLOSED_PATS = ( + re.compile(r"(?:(?!).)*", re.DOTALL), + _TOOL_CLOSED_PATS[1], + re.compile(r"<\|tool_call>.*?", re.DOTALL), +) + + +def _marker_inside_leading_envelope(content: str, enabled_tool_names: Optional[set] = None) -> bool: + first_marker = _EMBEDDED_MARKER_RE.search(content) + if first_marker is None: + return False + # A leading bare-JSON or Mistral [TOOL_CALLS] call is an outer envelope too: + # a DS/Kimi marker in its argument strings is data. + i = 0 + n = len(content) + while i < n and content[i] in " \t\n\r": + i += 1 + if content.startswith("{", i): + end = _balanced_brace_end(content, i) + if end is not None and i < first_marker.start(): + name = _top_level_bare_json_name(content[i : end + 1]) + if name is not None and (enabled_tool_names is None or name in enabled_tool_names): + # The closed leading call owns the turn: a marker inside it is argument + # data, one after it a trailing example (same rule as the XML envelopes below). + return True + if name is not None and first_marker.start() <= end: + # A disabled-name leading object is prose (can't own the turn), but a marker + # inside its own strings stays data. A marker AFTER it falls through to the pre-pass. + return True + elif content.startswith(_MISTRAL_TRIGGER, i): + end = _mistral_region_end(content, i) + if end is not None and i < first_marker.start(): + return True + # A closed outer call PRECEDING the first marker owns the turn; the pre-pass must + # not steal a trailing example or argument data. + for _pat in _OUTER_ENVELOPE_CLOSED_PATS: + m = _pat.search(content) + if m is not None and m.start() < first_marker.start(): + return True + residue = content + for _pat in _OUTER_ENVELOPE_CLOSED_PATS: + residue = _pat.sub("", residue) + marker = _EMBEDDED_MARKER_RE.search(residue) + if marker is None: + return True + # A marker still stands; any opener left in the residue is UNCLOSED. One before the + # marker is a truncated outer call holding the marker as data: skip the pre-pass. + opener = _OUTER_ENVELOPE_OPEN_RE.search(residue) + return opener is not None and opener.start() < marker.start() + + +def _mistral_region_end(text: str, idx: int) -> int | None: + """Exclusive end of the balanced ``[TOOL_CALLS]`` call starting at ``idx``, + or ``None`` when truncated/unrecognised (same shapes as the strip scan: + array, single-object, and named ``name [CALL_ID]? [ARGS]? {json}``).""" + n = len(text) + i = idx + len(_MISTRAL_TRIGGER) + while i < n and text[i] in " \t\n\r": + i += 1 + if i < n and text[i] == "[": + end = _balanced_bracket_end(text, i) + return None if end is None else end + 1 + if i < n and text[i] == "{": + end = _balanced_brace_end(text, i) + return None if end is None else end + 1 + name_match = _MISTRAL_V11_NAME_RE.match(text, i) + if not name_match: + return None + i = name_match.end() + while i < n and text[i] in " \t\n\r": + i += 1 + i = _skip_mistral_call_id(text, i) + if text.startswith(_MISTRAL_ARGS_MARKER, i): + i += len(_MISTRAL_ARGS_MARKER) + while i < n and text[i] in " \t\n\r": + i += 1 + if i >= n or text[i] != "{": + return None + end = _balanced_brace_end(text, i) + return None if end is None else end + 1 + + +def _xml_signal_inside_leading_mistral(content: str) -> bool: + """True when a parseable Mistral call is the first tool emission in document order: it owns the turn, so later XML (quoted in its arguments or in trailing prose) is not promoted over it. A signal BEFORE the trigger keeps normal order.""" + trig = content.find(_MISTRAL_TRIGGER) + if trig < 0: + return False + first_xml = _first_foreign_tool_signal(content) + if first_xml is not None and first_xml < trig: + return False + # Only plain prose precedes the trigger: a visible preface must not hand + # the turn to a later XML literal (preamble-tolerant, like the + # wrapperless-Gemma guard). Prose that merely mentions the marker has no + # parseable region and keeps the normal order. + return _mistral_region_end(content, trig) is not None + + +def _parse_bare_rehearsals( + content: str, + *, + id_offset: int = 0, + enabled_tool_names: Optional[set] = None, +) -> list[dict]: + """Promote bare reasoning-rehearsal ``name[ARGS]{json}`` calls that a leading [TOOL_CALLS] + owns-the-turn parse would miss. Only the ``rehearsal`` kind is taken (a Mistral + ``[TOOL_CALLS]name[ARGS]{..}`` yields ``name`` and is not double-counted), and a rehearsal + inside a ```` / ``[THINK]`` block is reasoning, so it is skipped.""" + out: list[dict] = [] + think_spans = _tool_healing._think_spans_outside_tool_markup(content) + for start, end, kind, m in _tool_healing._iter_bracket_spans( + content, enabled_tool_names = enabled_tool_names + ): + if kind != "rehearsal": + continue + if any(s <= start < e for s, e in think_spans): + continue + try: + payload = json.loads(content[m.end() : end]) + except (json.JSONDecodeError, ValueError): + continue + if not isinstance(payload, dict): + continue + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": m.group(1), "arguments": json.dumps(payload)}, + } + ) + return out + + +_ATTR_FUNC_OPEN_RE = re.compile(r' int | None: + """Offset of the first tool signal a non-envelope parser would fire on + (XML forms plus ``<|python_tag|>``, which also runs before the Mistral parser).""" + first = None + for sig in ("", "<|tool_call>", ""): + p = content.find(sig) + if p >= 0 and (first is None or p < first): + first = p + attr = _ATTR_FUNC_OPEN_RE.search(content) + if attr is not None and (first is None or attr.start() < first): + first = attr.start() + # DeepSeek/Kimi markers are foreign to a JSON envelope too: a marker inside a leading + # object routes through the same guard (and, if disabled, the drop-and-parse-the-tail + # recursion, so a real call after the object is still reached). + marker = _EMBEDDED_MARKER_RE.search(content) + if marker is not None and (first is None or marker.start() < first): + first = marker.start() + return first + + +def _xml_signal_inside_leading_bare_json(content: str) -> bool: + """True when the first foreign tool signal is a quoted literal inside a + LEADING bare-JSON call object or JSON answer -- data, not a real call + (sibling of ``_xml_signal_inside_leading_mistral``).""" + i = 0 + n = len(content) + while i < n and content[i] in " \t\n\r": + i += 1 + if i >= n or content[i] not in "{[": + return False + if content[i] == "[": + # A leading array is only ever a structured answer; its literals are data. + end = _balanced_bracket_end(content, i) + if end is None: + return False + try: + json.loads(content[i : end + 1]) + except ValueError: + return False + first_xml = _first_foreign_tool_signal(content) + trig = content.find(_MISTRAL_TRIGGER) + if trig >= 0 and (first_xml is None or trig < first_xml): + first_xml = trig + return first_xml is not None and i < first_xml < end + end = _balanced_brace_end(content, i) + if end is None: + return False + if _top_level_bare_json_name(content[i : end + 1]) is None: + # A NAMELESS object that parses as real JSON is a structured answer / envelope too: + # quoted markup is data, and the decline path drops it and parses the tail. + # Non-JSON braced prose keeps the old behaviour. + try: + json.loads(content[i : end + 1]) + except ValueError: + return False + first_xml = _first_foreign_tool_signal(content) + # The Mistral trigger is foreign to a JSON envelope too (its parser runs first). + trig = content.find(_MISTRAL_TRIGGER) + if trig >= 0 and (first_xml is None or trig < first_xml): + first_xml = trig + # Inside the balanced body the signal is quoted data; after the closed object the + # leading call still owns the turn (mirrors the leading-Mistral rule). + return first_xml is not None and i < first_xml + + +def _signal_inside_leading_wrapperless_gemma( + content: str, enabled_tool_names: Optional[set] +) -> bool: + """True when the first foreign tool signal is a quoted literal inside (or + after) a LEADING enabled wrapper-less Gemma call (sibling of the + Mistral/bare-JSON leading guards). Markerless form, so gated on an enabled + name (``None`` keeps the name-agnostic behaviour).""" + first = _first_foreign_tool_signal(content) + # The Mistral trigger is foreign to a Gemma call too (its parser runs first). + trig = content.find(_MISTRAL_TRIGGER) + if trig >= 0 and (first is None or trig < first): + first = trig + if first is None: + return False + # A preamble before ``call:NAME{...}`` is normal; what matters is an ENABLED balanced + # call beginning before the first foreign signal. + cursor = 0 + while True: + m = _GEMMA_BARE_TC_RE.search(content, cursor) + if m is None or m.start() > first: + return False + if enabled_tool_names is not None and m.group(1) not in enabled_tool_names: + cursor = m.end() + continue + end = _gemma_body_brace_end(content, m.end() - 1) + if end is None: + return False + if m.end() - 1 < first <= end: + return True + # An enabled call that CLOSES before the signal still owns the turn (inside-or-after + # rule, as for closed bare-JSON/Mistral envelopes), gated on an enabled name. + return enabled_tool_names is not None and end < first + + +def _disabled_gemma_call_end_containing_signal( + content: str, enabled_tool_names: Optional[set] +) -> int | None: + """End offset (exclusive) of the earliest DISABLED wrapper-less Gemma call + whose balanced body contains the first foreign signal, else None. A disabled + name is prose, so the quoted literal is data: the caller drops the span and + recurses on the tail. An ENABLED call defers to the enabled-call guard.""" + if enabled_tool_names is None: + return None + first = _first_foreign_tool_signal(content) + # Mirror the enabled-call guard: the Mistral trigger is foreign here too. + trig = content.find(_MISTRAL_TRIGGER) + if trig >= 0 and (first is None or trig < first): + first = trig + if first is None: + return None + cursor = 0 + while True: + m = _GEMMA_BARE_TC_RE.search(content, cursor) + if m is None or m.start() > first: + return None + if m.group(1) in enabled_tool_names: + return None + end = _gemma_body_brace_end(content, m.end() - 1) + if end is None: + cursor = m.end() + continue + if m.end() - 1 < first <= end: + return end + 1 + cursor = end + 1 + + +def parse_tool_calls_from_text( + content: str, + *, + id_offset: int = 0, + allow_incomplete: bool = True, + enabled_tool_names: Optional[set] = None, +) -> list[dict]: + """Return OpenAI-format tool calls, first-match wins so calls are never double-counted. + + ``allow_incomplete=True`` (default) heals truncated calls (missing close tag / + unclosed parameter); ``False`` accepts only well-formed closed calls (trailing + prose tolerated), matching llama-server's strict path when Auto-Heal is off. + + ``enabled_tool_names`` gates only the markerless Llama-3.2 bare-JSON form (the + marker-based forms carry an explicit signal, so a disabled-tool name there is a + real call attempt). ``None`` keeps the name-agnostic behaviour.""" + # Drop Magistral [THINK]...[/THINK] BEFORE dispatch: a rehearsed call inside it must + # never be promoted, and the parse path must agree with the display strip. + content = _strip_mistral_reasoning(content) + + # A leading bare-JSON value is decided FIRST: a string argument quoting tool markup + # (XML or a Mistral trigger) must stay data, so the bare-JSON parser takes the outer + # call before any other pass. Precedes the Mistral guard, whose preamble tolerance + # would otherwise claim a trigger quoted inside the leading object. + if _xml_signal_inside_leading_bare_json(content): + calls = _parse_llama3_bare_json( + content, id_offset = id_offset, enabled_tool_names = enabled_tool_names + ) + if calls: + return calls + # Disabled/example name: the leading object is content. Drop it and parse the tail. + i = 0 + while i < len(content) and content[i] in " \t\n\r": + i += 1 + # The guard guarantees a balanced leading value (object or array). + end = (_balanced_brace_end if content[i] == "{" else _balanced_bracket_end)(content, i) + return parse_tool_calls_from_text( + content[end + 1 :], + id_offset = id_offset, + allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, + ) + + # A leading enabled wrapper-less Gemma call is decided BEFORE the Mistral guard: its + # body reads as prose to the preamble tolerance below, so a quoted [TOOL_CALLS] would + # otherwise steal the turn. + if _signal_inside_leading_wrapperless_gemma(content, enabled_tool_names): + calls = _parse_gemma_tool_calls( + content, + id_offset = id_offset, + allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, + ) + if calls: + return calls + + # A DISABLED wrapper-less Gemma call is prose: drop the span and parse the tail BEFORE + # the Mistral guard, whose preamble tolerance would otherwise parse a quoted trigger. + _prose_end = _disabled_gemma_call_end_containing_signal(content, enabled_tool_names) + if _prose_end is not None: + return parse_tool_calls_from_text( + content[_prose_end:], + id_offset = id_offset, + allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, + ) + + # A [TOOL_CALLS] call that is the first tool emission owns the turn: XML quoted in its + # arguments or trailing prose is not promoted over it, nor does a prose preface forfeit it. + if _xml_signal_inside_leading_mistral(content): + calls = _parse_mistral_tool_calls( + content, id_offset = id_offset, allow_incomplete = allow_incomplete + ) + if calls: + # A bare rehearsal ``name[ARGS]{..}`` after the Mistral call is a peer tool call, + # not foreign XML the owns-the-turn guard protects against: promote it too so a + # Mistral call and a rehearsal in one message both parse. + calls.extend( + _parse_bare_rehearsals( + content, + id_offset = id_offset + len(calls), + enabled_tool_names = enabled_tool_names, + ) + ) + return calls + + # DeepSeek/Kimi markers are unique, so try them first -- unless an outer envelope + # opens before the first marker (then the marker is argument data). + if not _marker_inside_leading_envelope(content, enabled_tool_names): + # Dispatch by earliest opener so a quoted DS example inside a Kimi call (or vice + # versa) can't hijack the turn via fixed parser order. + _ds = _DEEPSEEK_BEGIN_RE.search(content) + _ds_pos = _ds.start() if _ds else len(content) + _km_section = content.find(_KIMI_SECTION_BEGIN) + _km_bare = content.find(_KIMI_CALL_BEGIN) + _km_pos = min(p for p in (_km_section, _km_bare, len(content)) if p >= 0) + pre_pass = [ + (_ds_pos, _parse_deepseek_tool_calls), + (_km_pos, _parse_kimi_tool_calls), + ] + pre_pass.sort(key = lambda pair: pair[0]) + for _pos, parser in pre_pass: + calls = parser(content, id_offset = id_offset, allow_incomplete = allow_incomplete) + if calls: + return calls + + # A leading MiniCPM/MiniMax attribute-form call owns the turn: tool_healing doesn't know + # the wrapper, so a quoted in its parameter would beat + # the outer call. Any earlier signal keeps normal order. + attr = _ATTR_FUNC_OPEN_RE.search(content) + if attr is not None: + first_other = None + for sig in ( + "", + "<|tool_call>", + "", + _MISTRAL_TRIGGER, + ): + p = content.find(sig) + if p >= 0 and (first_other is None or p < first_other): + first_other = p + if first_other is None or attr.start() < first_other: + calls = _parse_function_xml( + content, id_offset = id_offset, allow_incomplete = allow_incomplete + ) + if calls: + return calls + + # A leading Llama-3 ``<|python_tag|>`` call owns the turn like the others: markup quoted + # in a ``.call(...)`` argument is not promoted. tool_healing does not know the tag, so + # gate it here. A foreign signal before the tag keeps normal order. + py_tag = content.find(_LLAMA3_PYTHON_TAG) + if py_tag >= 0: + first_other = None + for sig in ("", "<|tool_call>", "= 0 and (first_other is None or p < first_other): + first_other = p + attr = _ATTR_FUNC_OPEN_RE.search(content) + if attr is not None and (first_other is None or attr.start() < first_other): + first_other = attr.start() + if first_other is None or py_tag < first_other: + calls = _parse_llama3_python_tag( + content, id_offset = id_offset, allow_incomplete = allow_incomplete + ) + if calls: + return calls + + # Qwen/Hermes, Qwen3.5 XML, Gemma 4, plus Mistral [TOOL_CALLS] / bare rehearsal + # ``name[ARGS]{json}`` use the shared tool_healing parser (strict/Auto-Heal contract + + # nested-marker, trailing-prose, and ``<|"|>`` quoted-string handling the GGUF path + # relies on). ``enabled_tool_names`` gates the ambiguous bare-rehearsal form so an + # inactive ``foo[ARGS]{..}`` stays prose. + calls = _tool_healing.parse_tool_calls_from_text( + content, + id_offset = id_offset, + allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, + ) + if calls: + return calls + + # Formats tool_healing does not cover; these run only after it finds + # nothing, so a strict-rejected call is never re-healed here. Blank any + # JSON/Gemma marker coverage first: markup inside a marker's span (even one + # that failed to parse) is that call's data, not a sibling, so a nested + # ```` / ``<|python_tag|>`` / ``[TOOL_CALLS]`` must not be promoted. + fallback_content = content + coverage = _tool_healing.marker_coverage(content) + if coverage: + chars = list(content) + for cov_start, cov_end in coverage: + for i in range(cov_start, min(cov_end, len(chars))): + chars[i] = " " + fallback_content = "".join(chars) + for parser in ( + _parse_glm_tool_calls, # GLM 4.x name + _parse_function_xml, # attribute form + _parse_llama3_python_tag, # Llama-3 <|python_tag|> + _parse_mistral_tool_calls, # Mistral [TOOL_CALLS] + ): + calls = parser(fallback_content, id_offset = id_offset, allow_incomplete = allow_incomplete) + if calls: + return calls + + # Llama-3.2 bare ``{"name":..., "parameters":...}`` (strict shape). Only a LEADING call + # object matches and owns the turn, so an enabled ``call:NAME{...}`` in its arguments + # stays data (Gemma never starts ``{``). + calls = _parse_llama3_bare_json( + content, id_offset = id_offset, enabled_tool_names = enabled_tool_names + ) + if calls: + return calls + + # Gemma wrapper-less ``call:NAME{...}``: markerless, so the same enabled-name gate applies. + return _parse_gemma_tool_calls( + content, + id_offset = id_offset, + allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, + ) + + +def _parse_tool_call_json( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + out: list[dict] = [] + for m in _TC_JSON_START_RE.finditer(content): + brace_start = m.end() - 1 + end = _balanced_brace_end(content, brace_start) + if end is None: + continue + # Strict mode: a balanced JSON body that never closed its ```` + # is a truncated call, not a finished one. Trailing prose after the close + # is still tolerated (matches the GGUF strict path). + if not allow_incomplete and not content[end + 1 :].lstrip().startswith(""): + continue + try: + obj = json.loads(content[brace_start : end + 1]) + except (json.JSONDecodeError, ValueError): + continue + name = obj.get("name", "") + # Accept ``arguments`` (Hermes/Qwen) and ``parameters`` (Llama-3 drift). + args = obj.get("arguments") + if args is None: + args = obj.get("parameters", {}) + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + args_str = args + else: + args_str = json.dumps({"value": args}) + if not name: + continue + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": args_str}, + } + ) + return out + + +def _trim_param_value(val: str) -> str: + """Trim one wrapping newline the template adds around an XML parameter value + (``\nVALUE\n``), preserving inner indentation. + ``str.strip()`` destroyed code/diff indentation; SGLang's qwen3_coder trims only + the wrapping newline.""" + if val.startswith("\n"): + val = val[1:] + if val.endswith("\n"): + val = val[:-1] + return val + + +def _inside_open_parameter(text: str, pos: int) -> bool: + """True if ``pos`` sits inside an unclosed ````/```` block -- + i.e. a ```` / ```` opener at ``pos`` is a literal inside an + argument value (e.g. code that prints tool-call XML), not a real nested call. + Compares the last parameter opener before ``pos`` against the last + parameter/function close before it.""" + last_param_open = -1 + for m in _TC_PARAM_START_RE.finditer(text, 0, pos): + last_param_open = m.start() + if last_param_open < 0: + return False + # The parameter's OWN close tag decides: while it closes after ``pos`` the position is + # argument data, even across several literal function closes. Only an unclosed + # parameter (heal mode) falls back to the first function close. + own_closes = [ + c + for c in ( + text.find("", last_param_open), + text.find("", last_param_open), + ) + if c >= 0 + ] + if own_closes: + return min(own_closes) > pos + func_closes = [ + c + for c in ( + text.find("", last_param_open), + text.find("", last_param_open), + ) + if c >= 0 + ] + return not func_closes or pos < min(func_closes) + + +def _parse_function_xml( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + out: list[dict] = [] + # Skip ```` openers that are literals inside an open parameter value, + # else the nested marker is promoted to a second call and truncates the real argument. + func_starts = [ + fm + for fm in _TC_FUNC_START_RE.finditer(content) + if not _inside_open_parameter(content, fm.start()) + ] + for idx, fm in enumerate(func_starts): + # group(1) is ````, group(2) is ````. + func_name = fm.group(1) or fm.group(2) + body_start = fm.end() + next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content) + # The call ends at the FIRST / not inside an open + # parameter: a literal close in a code/search argument is skipped as data, and + # prose after the real close isn't folded into the last argument (mirrors + # _strip_function_xml_calls and tool_healing._func_close_index). + close_match = None + for cm in _TC_END_TAG_RE.finditer(content, body_start, next_func): + if not _inside_open_parameter(content, cm.start()): + close_match = cm + break + has_close = close_match is not None + if has_close: + body_end = close_match.start() + else: + body_end = min(len(content), next_func) + # Strict mode: an unclosed function call is truncated -- do not heal it. + if not allow_incomplete and not has_close: + continue + body = _TC_FUNC_CLOSE_RE.sub("", content[body_start:body_end]) + + args: dict = {} + param_unclosed = False + # A ```` opener inside an open parameter value is literal text. + param_starts = [ + pm + for pm in _TC_PARAM_START_RE.finditer(body) + if not _inside_open_parameter(body, pm.start()) + ] + if len(param_starts) == 1: + pm = param_starts[0] + raw_val = body[pm.end() :] + if not _TC_PARAM_CLOSE_RE.search(raw_val): + param_unclosed = True + val = _TC_PARAM_CLOSE_RE.sub("", raw_val) + args[pm.group(1) or pm.group(2)] = _trim_param_value(val) + else: + for pidx, pm in enumerate(param_starts): + val_start = pm.end() + next_param = ( + param_starts[pidx + 1].start() if pidx + 1 < len(param_starts) else len(body) + ) + raw_val = body[val_start:next_param] + if not _TC_PARAM_CLOSE_RE.search(raw_val): + param_unclosed = True + val = _TC_PARAM_CLOSE_RE.sub("", raw_val) + args[pm.group(1) or pm.group(2)] = _trim_param_value(val) + + # Strict mode: a dangling parameter means the call was cut off; a closed + # zero-parameter call stays valid. + if not allow_incomplete and param_unclosed: + continue + + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": func_name, "arguments": json.dumps(args)}, + } + ) + return out + + +def _llama3_kv_value(body: str, p: int, n: int) -> tuple[Any, int | None]: + """One ``.call`` value (string/number/true/false/null) at ``body[p:]``. + Returns ``(value, consumed_len)`` or ``(None, None)`` if none matches.""" + if p >= n: + return None, None + if body[p] == '"': + # ``"((?:\\.|[^"\\])*)"`` by hand so an unterminated quote is O(n), not O(n^2). + j = p + 1 + while j < n: + c = body[j] + if c == "\\": + # ``\\.`` needs a following non-newline char; else the body can't match. + if j + 1 >= n or body[j + 1] == "\n": + return None, None + j += 2 + continue + if c == '"': + raw = body[p + 1 : j] + # json.loads keeps \n/\uXXXX escapes and literal UTF-8 (emoji/CJK) intact. + try: + return json.loads('"' + raw + '"'), j + 1 - p + except (json.JSONDecodeError, ValueError): + return raw, j + 1 - p + j += 1 + return None, None # unterminated + nm = _LLAMA3_NUM_RE.match(body, p) + if nm: + v = nm.group(0) + # Scientific notation (1e-3, -2E+4, 0.5e2) and decimals decode as float; a bare + # integer stays int. ``"." in v`` alone missed the exponent forms (1e-3 -> 1). + return (float(v) if any(c in v for c in ".eE") else int(v)), nm.end() - p + lm = _LLAMA3_LIT_RE.match(body, p) + if lm: + return {"true": True, "false": False, "null": None}[lm.group(0)], lm.end() - p + return None, None + + +def _parse_llama3_kv_args(body: str) -> dict[str, Any]: + """``k=v, ...`` kwargs from a ``.call(...)`` body, left to right (later keys win). + Linear hand-scan replacing the quadratic ``_LLAMA3_KV_RE.finditer`` walk.""" + args: dict[str, Any] = {} + n = len(body) + i = 0 + while i < n: + km = _LLAMA3_KEY_RE.match(body, i) + if km is None: + i += 1 + continue + p = _LLAMA3_WS_RE.match(body, km.end()).end() + if p >= n or body[p] != "=": + i = km.end() + continue + p = _LLAMA3_WS_RE.match(body, p + 1).end() + val, length = _llama3_kv_value(body, p, n) + if length is None: + i = km.end() + continue + args[km.group(0)] = val + i = p + length + return args + + +def _parse_llama3_python_tag( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Parse the Llama-3 emissions: ``<|python_tag|>NAME.call(...)`` (built-in), + ``<|python_tag|>{"name":..., "parameters":...}`` (custom), multi-call via + ``; ``, ``parameters`` or ``arguments`` key.""" + out: list[dict] = [] + if _LLAMA3_PYTHON_TAG not in content: + return out + + # 1. ``NAME.call(...)`` built-in form, anchored to ``<|python_tag|>`` and optionally + # ``; ``-chained within one emission. Anchoring to the tag boundary (not a free scan) + # keeps a literal ``<|python_tag|>x.call(...)`` quoted in a custom-form JSON argument + # from being mistaken for a real built-in call. + pos = content.find(_LLAMA3_PYTHON_TAG) + truncated = False + while pos >= 0 and not truncated: + head = _LLAMA3_PY_CALL_HEAD_RE.match(content, pos + len(_LLAMA3_PYTHON_TAG)) + if head is None: + # Tag is the custom JSON form (``{...}``) or noise -- leave it to step 2. + break + name = head.group(1) + open_idx = head.end() + i = open_idx + while True: + i = open_idx + depth = 1 + in_string = False + esc = False + while i < len(content) and depth > 0: + ch = content[i] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_string = False + else: + if ch == '"': + in_string = True + elif ch == "(": + depth += 1 + elif ch == ")": + depth -= 1 + if depth == 0: + break + i += 1 + # Truncated ``.call(...)`` with no closing paren: reject in strict mode + # instead of executing a partial. + if not allow_incomplete and depth > 0: + truncated = True + break + body = content[open_idx:i] + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(_parse_llama3_kv_args(body)), + }, + } + ) + # ``)`` then optional ``; NAME.call(`` chains the next built-in call. + chain = _LLAMA3_CALL_CHAIN_RE.match(content, i + 1) + if chain is None: + break + name = chain.group(1) + open_idx = chain.end() + # Past the consumed region: a second ``<|python_tag|>`` may carry more calls. + pos = content.find(_LLAMA3_PYTHON_TAG, i + 1) + + # 2. ``<|python_tag|>{"name":..., "parameters":...}``. ``raw_decode`` peels multiple + # ``; ``-separated objects from one emission. + if not out: + decoder = json.JSONDecoder() + idx = content.find(_LLAMA3_PYTHON_TAG) + while idx >= 0: + search_from = idx + len(_LLAMA3_PYTHON_TAG) + cursor = search_from + while cursor < len(content): + brace = content.find("{", cursor) + if brace < 0: + break + # Stop at the next ``<|python_tag|>``. + next_tag = content.find(_LLAMA3_PYTHON_TAG, search_from, brace) + if next_tag >= 0: + break + try: + obj, end_offset = decoder.raw_decode(content[brace:]) + except (json.JSONDecodeError, ValueError): + cursor = brace + 1 + continue + if not isinstance(obj, dict): + cursor = brace + end_offset + continue + name = obj.get("name") or obj.get("function") or "" + args = obj.get("parameters") if "parameters" in obj else obj.get("arguments", {}) + # Skip rather than fabricate ``{"value": args}`` for a non-dict/non-string value. + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + args_str = args + else: + cursor = brace + end_offset + continue + if name: + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": args_str}, + } + ) + cursor = brace + end_offset + idx = content.find(_LLAMA3_PYTHON_TAG, cursor) + return out + + +# Llama-3 special-token sentinels (chainable, any order) plus the role label the +# template inserts between ``<|start_header_id|>`` and ``<|end_header_id|>``. +_LLAMA3_BARE_JSON_SENTINELS = ( + "<|begin_of_text|>", + "<|eot_id|>", + "<|start_header_id|>", + "<|end_header_id|>", + "<|eom_id|>", +) +_LLAMA3_HEADER_ROLES = ("assistant", "user", "system", "tool", "ipython") + + +def strip_llama3_leading_sentinels(content: str) -> str: + """Strip leading Llama-3 special-token sentinels (and the role label after + ``<|start_header_id|>``) that can leak from a prior turn before a bare-JSON tool + call. Shared by the parser and the streaming buffering guards so a + sentinel-prefixed ``{"name":...}`` is recognised the same everywhere.""" + stripped = content.lstrip() + while True: + stripped = stripped.lstrip() + matched = False + for sentinel in _LLAMA3_BARE_JSON_SENTINELS: + if stripped.startswith(sentinel): + stripped = stripped[len(sentinel) :] + if sentinel == "<|start_header_id|>": + for role in _LLAMA3_HEADER_ROLES: + if stripped.startswith(role): + stripped = stripped[len(role) :] + break + matched = True + break + if not matched: + return stripped + + +def _parse_llama3_bare_json( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, + enabled_tool_names: Optional[set] = None, +) -> list[dict]: + """Llama-3.2 ``custom_tools`` bare ``{"name":.., "parameters":{..}}`` (no ``<|python_tag|>``), + strict so prose/echoes don't fire. ``enabled_tool_names`` gates on the parsed name so an + ordinary JSON answer isn't misread as a call to a disabled tool; ``None`` is name-agnostic.""" + out: list[dict] = [] + stripped = strip_llama3_leading_sentinels(content) + if not stripped.startswith("{"): + return out + + decoder = json.JSONDecoder() + cursor = 0 + n = len(stripped) + while cursor < n: + # Skip whitespace and the Llama-3 ``;`` inter-call separator. + while cursor < n and stripped[cursor] in " \t\n\r;": + cursor += 1 + if cursor >= n or stripped[cursor] != "{": + break + try: + obj, end_offset = decoder.raw_decode(stripped[cursor:]) + except (json.JSONDecodeError, ValueError): + break + if not isinstance(obj, dict): + break + name = obj.get("name") or obj.get("function") or "" + if not isinstance(name, str) or not name: + break + # Markerless JSON is ambiguous: treat it as a call only when the name is an enabled + # tool, else it is an ordinary JSON answer. + if enabled_tool_names is not None and name not in enabled_tool_names: + break + # ``parameters`` must be a dict (Llama-3 spec); ``arguments`` may be a dict or + # JSON-string of one (OpenAI). Looser would fire on ``{"name":"x","parameters":"sentence"}``. + if "parameters" in obj: + args = obj.get("parameters") + if not isinstance(args, dict): + break + args_str = json.dumps(args) + elif "arguments" in obj: + args = obj.get("arguments") + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + try: + parsed = json.loads(args) + except (json.JSONDecodeError, ValueError): + break + if not isinstance(parsed, dict): + break + args_str = args + else: + break + else: + break + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": args_str}, + } + ) + cursor += end_offset + return out + + +def _parse_mistral_tool_calls( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Parse all Mistral emissions: pre-v11 ``[TOOL_CALLS][...]`` / ``[TOOL_CALLS]{...}`` + and v11+ ``[TOOL_CALLS]name{json}`` / ``[TOOL_CALLS]name[ARGS]{json}``.""" + out: list[dict] = [] + content = _strip_mistral_reasoning(content) + idx = content.find(_MISTRAL_TRIGGER) + if idx < 0: + return out + + # Disambiguate the first occurrence: array / single object (pre-v11), or bare-name (v11+). + j = idx + len(_MISTRAL_TRIGGER) + k = j + while k < len(content) and content[k] in " \t\n\r": + k += 1 + if k >= len(content): + return out + + if content[k] == "[": + return _parse_mistral_array(content, k, id_offset, allow_incomplete = allow_incomplete) + + if content[k] == "{": + # Pre-v11 single ``{"name":...}``; fall through without a ``name`` so v11+ still runs. + end = _balanced_brace_end(content, k) + if end is not None: + try: + obj = json.loads(content[k : end + 1]) + if isinstance(obj, dict) and obj.get("name"): + _consume_mistral_call(content[k : end + 1], out, id_offset) + return out + except (json.JSONDecodeError, ValueError): + pass + + # v11+: walk every ``[TOOL_CALLS]``, parsing ``name{json}`` or + # ``name[ARGS]{json}`` after each trigger. + pos = idx + while pos >= 0: + cur = pos + len(_MISTRAL_TRIGGER) + nm = _MISTRAL_V11_NAME_RE.match(content, cur) + if not nm: + pos = content.find(_MISTRAL_TRIGGER, cur) + continue + name = nm.group(1) + after_name = nm.end() + after_name = _skip_mistral_call_id(content, after_name) + if content.startswith(_MISTRAL_ARGS_MARKER, after_name): + after_name += len(_MISTRAL_ARGS_MARKER) + while after_name < len(content) and content[after_name] in " \t\n\r": + after_name += 1 + if after_name >= len(content) or content[after_name] != "{": + pos = content.find(_MISTRAL_TRIGGER, cur) + continue + end = _balanced_brace_end(content, after_name) + if end is None: + break + try: + args = json.loads(content[after_name : end + 1]) + except (json.JSONDecodeError, ValueError): + pos = content.find(_MISTRAL_TRIGGER, end + 1) + continue + if not isinstance(args, dict): + pos = content.find(_MISTRAL_TRIGGER, end + 1) + continue + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args), + }, + } + ) + pos = content.find(_MISTRAL_TRIGGER, end + 1) + return out + + +def _parse_mistral_array( + content: str, + start: int, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Pre-v11 ``[TOOL_CALLS] [{...}, ...]`` array form.""" + out: list[dict] = [] + j = start + depth = 0 + in_string = False + esc = False + while j < len(content): + ch = content[j] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_string = False + else: + if ch == '"': + in_string = True + elif ch == "[": + depth += 1 + elif ch == "]": + depth -= 1 + if depth == 0: + break + j += 1 + # An unclosed array (no matching ]) is a truncated call. In strict mode reject it + # instead of recovering objects by hand below. + if not allow_incomplete and depth != 0: + return out + body = content[start : j + 1] if depth == 0 else content[start:] + + try: + arr = json.loads(body) + if isinstance(arr, list): + for obj in arr: + if isinstance(obj, dict): + _consume_mistral_call(json.dumps(obj), out, id_offset) + return out + except (json.JSONDecodeError, ValueError): + if not allow_incomplete: + return out + + # Healing path for unclosed arrays: walk top-level objects, advancing past each balanced + # ``{...}`` instead of re-scanning from every ``{`` (quadratic ReDoS). + pos = 0 + blen = len(body) + while pos < blen: + brace = body.find("{", pos) + if brace < 0: + break + end = _balanced_brace_end(body, brace) + if end is None: + break # truncated mid-object: nothing after it can balance + _consume_mistral_call(body[brace : end + 1], out, id_offset) + pos = end + 1 + return out + + +def _consume_mistral_call(obj_text: str, out: list[dict], id_offset: int) -> None: + try: + obj = json.loads(obj_text) + except (json.JSONDecodeError, ValueError): + return + if not isinstance(obj, dict): + return + name = obj.get("name") or "" + # Mistral uses ``arguments``; accept the ``parameters`` alias too (sibling paths and + # SGLang's base detector alias it) so an array object keyed on it keeps args. + args = obj.get("arguments") + if args is None: + args = obj.get("parameters", {}) + if isinstance(args, dict): + args_str = json.dumps(args) + elif isinstance(args, str): + args_str = args + else: + args_str = json.dumps({"value": args}) + if name: + out.append( + { + "id": obj.get("id") or f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": args_str}, + } + ) + + +def _whole_content_is_json_value(text: str) -> bool: + """True when the entire content is one valid JSON value (a structured + answer, e.g. a response_format turn). Markerless scans must treat text + inside it as data: an answer documenting an enabled tool's syntax must + not execute that tool or have the example stripped from display.""" + t = text.strip() + if t[:1] not in "{[": + return False + try: + json.loads(t) + except ValueError: + return False + return True + + +def _leading_json_value_end(text: str) -> int | None: + """End index (exclusive) of a balanced LEADING JSON value that parses as + JSON: a structured answer possibly followed by prose. Markerless scans treat + its contents as data (extends ``_whole_content_is_json_value``); leading-keyed, + so a JSON blob mid-prose is not an answer span.""" + i = 0 + n = len(text) + while i < n and text[i].isspace(): + i += 1 + if i >= n or text[i] not in "{[": + return None + end = (_balanced_brace_end if text[i] == "{" else _balanced_bracket_end)(text, i) + if end is None: + return None + try: + json.loads(text[i : end + 1]) + except ValueError: + return None + return end + 1 + + +def _parse_gemma_tool_calls( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, + enabled_tool_names: Optional[set] = None, +) -> list[dict]: + """Gemma 4: ``<|tool_call>call:NAME{k:<|"|>v<|"|>, ...}``, plus the + ``skip_special_tokens`` stream where the wrapper and string markers were + stripped (bare ``call:NAME{k:v, ...}``). + + ``enabled_tool_names`` gates on the parsed name: the wrapper-less shape is + indistinguishable from prose documenting the syntax, so a disabled/example + name must not be stolen as a call. ``None`` keeps the name-agnostic behaviour.""" + out: list[dict] = [] + # The WRAPPED form (strict + nested-marker handling) is tool_healing's, which runs + # first: defer content with a wrapped opener. A marker literal alone is not enough -- + # a wrapper-less call mentioning ``<|tool_call>`` would be lost if deferred. + if _GEMMA_TC_RE.search(content): + return out + # A whole-content JSON value is a structured answer: quoted examples must not become calls. + if _whole_content_is_json_value(content): + return out + # Manual cursor: resume AFTER each consumed balanced body so a nested ``call:OTHER{...}`` + # in an argument is never re-matched. A leading JSON answer's span is data -- scan after it. + cursor = _leading_json_value_end(content) or 0 + while True: + m = _GEMMA_BARE_TC_RE.search(content, cursor) + if m is None: + break + name = m.group(1) + body_start = m.end() - 1 + end = _gemma_body_brace_end(content, body_start) + if end is None: + # Unclosed call: nothing parseable follows (mirrors the strip contract); + # scanning on would promote quoted argument text. + break + cursor = end + 1 + # Markerless: a disabled/example name is prose, not a call. + if enabled_tool_names is not None and name not in enabled_tool_names: + continue + body = content[body_start + 1 : end] + try: + args = _gemma_parse_stripped_body(body) + except Exception: + args = {} + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": name, "arguments": json.dumps(args)}, + } + ) + return out + + +def _balanced_brace_end(text: str, brace_pos: int) -> int | None: + """Index of the ``}`` matching ``{`` at ``brace_pos`` (ignores braces in JSON strings).""" + if brace_pos >= len(text) or text[brace_pos] != "{": + return None + depth = 0 + in_string = False + esc = False + i = brace_pos + while i < len(text): + ch = text[i] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_string = False + else: + if ch == '"': + in_string = True + elif ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return i + i += 1 + return None + + +def _gemma_body_brace_end(text: str, brace_pos: int) -> int | None: + """Index of the ``}`` closing the wrapper-less Gemma body at ``brace_pos``. + + Values are raw after ``skip_special_tokens``, so quoted strings (single or + double) hide braces; the quote rules mirror ``_gemma_parse_stripped_body`` so + the boundary always agrees with the body parser. Contextual openers: a single + quote opens only at value-start context (after ``:{[(,=`` -- apostrophes in + ``what's the weather`` are prose), a double quote also at word start (so + ``query:find "a, b"`` hides its delimiters).""" + if brace_pos >= len(text) or text[brace_pos] != "{": + return None + depth = 0 + quote = "" + prev = "" + prev_raw = "" + i = brace_pos + n = len(text) + while i < n: + ch = text[i] + if quote: + if ch == "\\" and i + 1 < n: + i += 2 + continue + if ch == quote: + quote = "" + elif ch in "\"'" and (prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())): + quote = ch + elif ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return i + if not ch.isspace(): + prev = ch + prev_raw = ch + i += 1 + return None + + +_BARE_JSON_NAME_RE = re.compile(r'"name"\s*:\s*"([^"]+)"') + + +def _top_level_bare_json_name(probe: str) -> Optional[str]: + """TOP-LEVEL ``"name"`` (or ``"function"`` alias, name wins) of a bare-JSON object, else None. + + Skips nested objects/arrays so a nested ``"name"`` isn't mistaken for the call name; a + truncated tail returns None so the caller keeps the text.""" + if not probe.startswith("{"): + return None + decoder = json.JSONDecoder() + function_value = None # the ``"function"`` alias, used only if no ``"name"`` key + i = 1 + n = len(probe) + while i < n: + while i < n and probe[i] in " \t\r\n,": + i += 1 + if i >= n or probe[i] == "}": + # End of the object with no top-level ``"name"``: fall back to a recorded ``"function"`` alias. + return function_value + if probe[i] != '"': + return None + try: + key, consumed = decoder.raw_decode(probe[i:]) + except (json.JSONDecodeError, ValueError): + return None + if not isinstance(key, str): + return None + i += consumed + while i < n and probe[i] in " \t\r\n": + i += 1 + if i >= n or probe[i] != ":": + return None + i += 1 + while i < n and probe[i] in " \t\r\n": + i += 1 + if key == "name": + if i < n and probe[i] == '"': + try: + value, _consumed = decoder.raw_decode(probe[i:]) + except (json.JSONDecodeError, ValueError): + return None + return value if isinstance(value, str) else None + return None + if key == "function" and function_value is None and i < n and probe[i] == '"': + # ``"function"`` aliases the call name. Record it but keep scanning: a top-level + # ``"name"`` still wins. + try: + value, consumed = decoder.raw_decode(probe[i:]) + except (json.JSONDecodeError, ValueError): + return None + if isinstance(value, str): + function_value = value + i += consumed + continue + # Skip a non-name top-level value; a truncated one can't prove a top-level name + # exists, so return None (keep the text). + if i < n and probe[i] == "{": + end = _balanced_brace_end(probe, i) + if end is None: + return None + i = end + 1 + elif i < n and probe[i] == "[": + end = _balanced_bracket_end(probe, i) + if end is None: + return None + i = end + 1 + else: + try: + _value, consumed = decoder.raw_decode(probe[i:]) + except (json.JSONDecodeError, ValueError): + return None + i += consumed + # No top-level ``"name"`` key: fall back to the ``"function"`` alias if seen. + return function_value + + +def strip_leading_bare_json_call(text: str, enabled_tool_names: Optional[set] = None) -> str: + """Remove leading Llama-3.2 bare-JSON calls (including a ``;``-chained run) + that ``strip_tool_markup`` misses; non-call text is unchanged and + ``enabled_tool_names`` gates like the parser. Consuming the whole chain + matters because the loops keep this text as next-turn assistant history: a + leftover executed call would be replayed alongside the structured + ``tool_calls``.""" + remainder = text + stripped_any = False + while True: + probe = strip_llama3_leading_sentinels(remainder.lstrip()) + # Skip the Llama-3 ``;`` inter-call separator between chained calls. + if stripped_any: + probe = probe.lstrip(" \t\n\r;") + if not (probe.startswith("{") and ('"name"' in probe or '"function"' in probe)): + return probe.lstrip() if stripped_any else text + if enabled_tool_names is not None: + # Only suppress when the leading object's TOP-LEVEL name is an enabled tool. A + # nested ``"name"`` (e.g. {"result":{"name":"web_search",...}}) is data, not the + # call name, so it must not gate the strip. An un-extractable name is kept. + name = _top_level_bare_json_name(probe) + if name not in enabled_tool_names: + return probe.lstrip() if stripped_any else text + end = _balanced_brace_end(probe, 0) + if end is None: + return "" # truncated bare-JSON call -- nothing recoverable + # A closed object must have the CALL SHAPE the parser accepts (dict ``parameters``, + # or dict / JSON-string ``arguments``). An ordinary JSON answer like + # {"name":"web_search","result":"no call"} is content, so the strip keeps it visible. + try: + obj = json.loads(probe[: end + 1]) + except (json.JSONDecodeError, ValueError): + return probe.lstrip() if stripped_any else text + if not _bare_json_call_shaped(obj): + return probe.lstrip() if stripped_any else text + remainder = probe[end + 1 :] + stripped_any = True + + +def _bare_json_call_shaped(obj) -> bool: + """The shape gate ``_parse_llama3_bare_json`` applies to a decoded object.""" + if not isinstance(obj, dict): + return False + # The parser requires a TOP-LEVEL name; a nested one (e.g. in a "result" value of an + # ordinary JSON answer) is data, and stripping it name-agnostically would delete content. + name = obj.get("name") or obj.get("function") or "" + if not isinstance(name, str) or not name: + return False + if "parameters" in obj: + return isinstance(obj.get("parameters"), dict) + args = obj.get("arguments") + if isinstance(args, dict): + return True + if isinstance(args, str): + try: + return isinstance(json.loads(args), dict) + except (json.JSONDecodeError, ValueError): + return False + return False + + +def _gemma_balanced_brace_end(text: str, brace_pos: int, hard_stop: int) -> int | None: + """Like ``_balanced_brace_end`` but skips ``<|"|>`` strings and matches {}/[] symmetrically.""" + if brace_pos >= len(text) or text[brace_pos] != "{": + return None + depth = 0 + i = brace_pos + while i < hard_stop: + if text.startswith(_GEMMA_STR_BEGIN, i): + close = text.find(_GEMMA_STR_END, i + len(_GEMMA_STR_BEGIN)) + if close < 0: + return None + i = close + len(_GEMMA_STR_END) + continue + ch = text[i] + if ch == "{" or ch == "[": + depth += 1 + elif ch == "}" or ch == "]": + depth -= 1 + if depth == 0: + return i + i += 1 + return None + + +def _gemma_parse_value( + text: str, + i: int, + *, + in_mapping: bool = False, +): + """Parse one Gemma arg value at ``i`` in a single O(n) forward pass; returns + ``(value, next_index, closed)``. ``closed`` is False when a string/object/array + runs off the end without its terminator, so the caller can fall back to raw. + ``in_mapping`` applies the top-level rule that a comma only ends the value + when a ``key:`` follows (array elements split on every top-level comma).""" + if text.startswith(_GEMMA_STR_BEGIN, i): + close = text.find(_GEMMA_STR_END, i + len(_GEMMA_STR_BEGIN)) + if close < 0: + return text[i + len(_GEMMA_STR_BEGIN) :], len(text), False + return text[i + len(_GEMMA_STR_BEGIN) : close], close + len(_GEMMA_STR_END), True + if text[i] == "{": + return _gemma_parse_mapping(text, i) + if text[i] == "[": + return _gemma_parse_array(text, i) + if text[i] in "\"'": + # Raw-quoted string: delimiters inside are data (``{city:"New, York"}`` is one + # value); returned unquoted like the top-level scalar coercion. + quote = text[i] + j = i + 1 + n = len(text) + while j < n: + if text[j] == "\\" and j + 1 < n: + j += 2 + continue + if text[j] == quote: + return text[i + 1 : j], j + 1, True + j += 1 + return text[i + 1 :], n, False + # Primitive / unquoted code: same delimiter rules as the top-level scan (bracket depth + # + contextual quote openers hide commas and closers). + end = i + n = len(text) + depth = 0 + quote = "" + prev = ":" + prev_raw = ":" + while end < n and not text.startswith(_GEMMA_STR_BEGIN, end): + ch = text[end] + if quote: + if ch == "\\" and end + 1 < n: + end += 2 + continue + if ch == quote: + quote = "" + elif ch in "\"'" and (prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())): + quote = ch + elif ch in "{[(": + depth += 1 + elif ch in "}])": + if depth == 0: + break + depth -= 1 + elif ch == "," and depth == 0: + if not in_mapping or _GEMMA_KEY_RE.match(text, end + 1): + break + if not ch.isspace(): + prev = ch + prev_raw = ch + end += 1 + if end == i: + # Stray delimiter where a value was expected: consume one char so callers always + # advance (no infinite loop on malformed input). + return "", i + 1, True + raw = text[i:end].strip() + if raw == "true": + return True, end, True + if raw == "false": + return False, end, True + if raw == "null": + return None, end, True + try: + return int(raw), end, True + except ValueError: + pass + try: + return float(raw), end, True + except ValueError: + pass + return raw, end, True + + +def _gemma_parse_array(text: str, start: int): + """Parse a Gemma ``[...]`` array at ``text[start] == '['`` in one forward + pass; returns ``(list, next_index, closed)``.""" + items: list[Any] = [] + i, n = start + 1, len(text) + while i < n: + while i < n and text[i] in " \t\n\r,": + i += 1 + if i < n and text[i] == "]": + return items, i + 1, True + if i >= n: + break + v, i, _closed = _gemma_parse_value(text, i) + items.append(v) + return items, i, False + + +def _gemma_coerce_scalar(raw: str) -> Any: + """Coerce an unquoted Gemma value to bool/int/float/None, else keep str + (quotes stripped first so quoted/unquoted variants compare identical).""" + raw = raw.strip() + if len(raw) >= 2 and raw[0] == raw[-1] and raw[0] in "\"'": + return raw[1:-1] + if raw == "true": + return True + if raw == "false": + return False + if raw == "null": + return None + try: + return int(raw) + except ValueError: + pass + try: + return float(raw) + except ValueError: + pass + return raw + + +def _gemma_strip_quoted_leaves(value: Any) -> Any: + """Recursively unquote quoted string leaves of a nested stripped-stream value, + so nested ``city:"New York"`` matches the top-level coercion (no stray quotes).""" + if isinstance(value, str): + v = value.strip() + if len(v) >= 2 and v[0] == v[-1] and v[0] in "\"'": + return v[1:-1] + return value + if isinstance(value, dict): + return {k: _gemma_strip_quoted_leaves(v) for k, v in value.items()} + if isinstance(value, list): + return [_gemma_strip_quoted_leaves(v) for v in value] + return value + + +def _gemma_parse_stripped_body(body: str) -> dict[str, Any]: + """Parse a quote-less Gemma arg body ``key:value, key2:value2`` (the + ``skip_special_tokens`` stream with ``<|"|>`` markers removed). Each value runs + to the next top-level ``, key:`` boundary, tracking ``{}``/``[]``/``()`` depth so + commas/braces inside a ``code`` / ``command`` value aren't truncated.""" + out: dict[str, Any] = {} + i, n = 0, len(body) + while i < n: + m = _GEMMA_KEY_RE.match(body, i) + if not m: + break + key = m.group(1) + i = m.end() + vstart = i + depth = 0 + quote = "" + # Contextual quote openers mirror _gemma_body_brace_end. + prev = ":" + prev_raw = ":" + while i < n: + ch = body[i] + if quote: + # A ``, key:`` shape inside the quoted string is not a boundary. + if ch == "\\" and i + 1 < n: + i += 2 + continue + if ch == quote: + quote = "" + elif ch in "\"'" and (prev in ":{[(,=" or (ch == '"' and prev_raw.isspace())): + quote = ch + elif ch in "{[(": + depth += 1 + elif ch in "}])": + if depth > 0: + depth -= 1 + elif ch == "," and depth == 0 and _GEMMA_KEY_RE.match(body, i + 1): + break + if not ch.isspace(): + prev = ch + prev_raw = ch + i += 1 + raw_val = body[vstart:i].strip() + if raw_val[:1] in "{[": + # Nested object/array: accept only a fully consumed, closed parse; a + # truncated/malformed value falls back to the raw string. + parsed, end, closed = _gemma_parse_value(raw_val, 0) + out[key] = ( + _gemma_strip_quoted_leaves(parsed) + if (closed and end == len(raw_val)) + else _gemma_coerce_scalar(raw_val) + ) + else: + out[key] = _gemma_coerce_scalar(raw_val) + if i < n and body[i] == ",": + i += 1 + return out + + +def _gemma_parse_mapping(text: str, start: int): + """Parse a Gemma ``{key:value, ...}`` mapping at ``text[start] == '{'`` in one + forward pass; returns ``(dict, next_index, closed)`` (``closed`` True iff the + matching ``}`` was reached).""" + out: dict[str, Any] = {} + i, n = start + 1, len(text) + while i < n: + while i < n and text[i] in " \t\n\r,": + i += 1 + if i < n and text[i] == "}": + return out, i + 1, True + if i >= n: + break + if text.startswith(_GEMMA_STR_BEGIN, i): + close = text.find(_GEMMA_STR_END, i + len(_GEMMA_STR_BEGIN)) + if close < 0: + break + key = text[i + len(_GEMMA_STR_BEGIN) : close] + i = close + len(_GEMMA_STR_END) + else: + kstart = i + while i < n and text[i] not in ":}": + i += 1 + key = text[kstart:i].strip() + while i < n and text[i] in " \t\n\r": + i += 1 + if i < n and text[i] == ":": + i += 1 + while i < n and text[i] in " \t\n\r": + i += 1 + if i >= n: + out[key] = None + break + if text[i] == "}": + out[key] = None + return out, i + 1, True + v, i, _closed = _gemma_parse_value(text, i, in_mapping = True) + out[key] = v + return out, i, False + + +# ── DeepSeek R1 / V3 / V3.1 ───────────────────────────────────────── + + +def _find_outside_json_strings(text: str, needle: str, start: int) -> int: + """Index of ``needle`` at/after ``start`` OUTSIDE any JSON string, or -1: a + marker inside an argument string must not be taken as the structural terminator.""" + i = start + n = len(text) + in_string = False + esc = False + while i < n: + ch = text[i] + if in_string: + if esc: + esc = False + elif ch == "\\": + esc = True + elif ch == '"': + in_string = False + i += 1 + continue + if ch == '"': + in_string = True + i += 1 + continue + if text.startswith(needle, i): + return i + i += 1 + return -1 + + +def _parse_deepseek_tool_calls( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """DeepSeek R1 / V3 / V3.1. + + R1: ``<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>NAME\\n``\\`\\`\\`json\\n{...}\\n\\`\\`\\`<｜tool▁call▁end｜>...`` + V3.x: ``<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>NAME<｜tool▁sep｜>{json}<｜tool▁call▁end｜>...`` + + Mirrors llama.cpp's pre-autoparser ``common_chat_parse_deepseek_r1`` / + ``_v3_1`` handling; tolerates the 5 opener variants llama.cpp keeps. + """ + out: list[dict] = [] + begin = _DEEPSEEK_BEGIN_RE.search(content) + if not begin: + return out + scan_start = begin.end() + # Envelope end OUTSIDE JSON strings: an argument may contain the literal end token, + # and a raw find would truncate the call. + end_pos = _find_outside_json_strings(content, _DEEPSEEK_END, scan_start) + # Strict mode: an unclosed envelope is truncated; reject, don't heal to EOF. + if not allow_incomplete and end_pos < 0: + return out + scan_end = end_pos if end_pos >= 0 else len(content) + body = content[scan_start:scan_end] + + # R1 path first: ``function<｜tool▁sep｜>NAME\n```json\n{...}\n```<｜tool▁call▁end｜>``. + pos = 0 + while pos < len(body): + fpos = body.find(_DEEPSEEK_R1_FUNC_MARKER, pos) + if fpos < 0: + break + name_start = fpos + len(_DEEPSEEK_R1_FUNC_MARKER) + nl = body.find("\n", name_start) + if nl < 0: + break + if not body.startswith(_DEEPSEEK_R1_FENCE, nl): + pos = name_start + continue + name = body[name_start:nl].strip() + json_start = nl + len(_DEEPSEEK_R1_FENCE) + # Walk a balanced ``{`` even if the trailing fence is truncated. + if json_start >= len(body) or body[json_start] != "{": + pos = json_start + continue + brace_end = _balanced_brace_end(body, json_start) + if brace_end is None: + break + try: + args = json.loads(body[json_start : brace_end + 1]) + except (json.JSONDecodeError, ValueError): + pos = brace_end + 1 + continue + if not isinstance(args, dict): + pos = brace_end + 1 + continue + # The closing fence + <｜tool▁call▁end｜> must IMMEDIATELY follow the JSON, else an + # unbounded search lands on a LATER call's terminator. Absent close: heal past the + # JSON (strict rejects); later well-formed calls are still kept. + after = brace_end + 1 + while after < len(body) and body[after] in " \t\r\n": + after += 1 + close_m = _DEEPSEEK_R1_CLOSE_RE.match(body, after) + if not allow_incomplete and close_m is None: + pos = brace_end + 1 + continue + if name: + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args), + }, + } + ) + pos = close_m.end() if close_m else brace_end + 1 + if out: + return out + + # V3 / V3.1: name then bare JSON. Use ``str.find`` for the sep marker and walk + # back for the name (a ``[^\n<]+`` regex search is O(N^2) on truncated bodies). + pos = 0 + while pos < len(body): + sep_pos = body.find(_DEEPSEEK_SEP, pos) + if sep_pos < 0: + break + # Walk left from sep_pos to the name start; stop at ``\n`` (turn boundary), ``<`` + # (tag start), or ``>`` (end of an optional ``<｜tool▁call▁begin｜>``). + name_start = sep_pos + while name_start > pos and body[name_start - 1] not in "\n<>": + name_start -= 1 + name = body[name_start:sep_pos].strip() + json_start = sep_pos + len(_DEEPSEEK_SEP) + while json_start < len(body) and body[json_start] in " \t\n\r": + json_start += 1 + if json_start >= len(body) or body[json_start] != "{": + pos = sep_pos + len(_DEEPSEEK_SEP) + continue + brace_end = _balanced_brace_end(body, json_start) + if brace_end is None: + break + # Strict mode: a real V3 call closes with the per-call <｜tool▁call▁end｜>; without + # it the call is truncated/merged, so skip it but keep scanning for a later + # well-formed call (matches Kimi strict). + if not allow_incomplete: + after = brace_end + 1 + while after < len(body) and body[after] in " \t\r\n": + after += 1 + if not body.startswith(_DEEPSEEK_CALL_END, after): + pos = brace_end + 1 + continue + try: + args = json.loads(body[json_start : brace_end + 1]) + except (json.JSONDecodeError, ValueError): + pos = brace_end + 1 + continue + if not isinstance(args, dict): + pos = brace_end + 1 + continue + if name: + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args), + }, + } + ) + # Advance just past the JSON; seeking the optional <｜tool▁call▁end｜> could land on + # a LATER call's end marker and skip the call between. + pos = brace_end + 1 + return out + + +# ── GLM 4.5 / 4.6 / 4.7 ───────────────────────────────────────────── + + +def _parse_glm_tool_calls( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """GLM 4.5 / 4.6 / 4.7. + + ``NAME[\\n]K[\\n]V + ...``. Multi-call is back-to-back blocks, no envelope. + Mirrors llama.cpp's GLM 4.x tool-call handling (``common_chat_params_init_glm_4_5`` + plus its generalized XML-style parser, llama.cpp PRs #15904 / #16932). + """ + out: list[dict] = [] + pos = 0 + while pos < len(content): + m = _GLM_TC_OPEN_RE.search(content, pos) + if not m: + break + name = m.group(1).strip() + apos = m.end() # absolute position in ``content``; advances past each pair + + args: dict[str, Any] = {} + valid = True + close = -1 + # Walk arg pairs directly against ``content``: a value may contain a literal + # , so the real close is the before the next . + # ``str.find`` keeps this linear. + while True: + ks = content.find(_GLM_ARG_KEY_OPEN, apos) + tc = content.find(_GLM_TC_CLOSE, apos) + if tc >= 0 and (ks < 0 or tc < ks): + close = tc + break + if ks < 0: + break # no close and no more keys -- truncated body + ke = content.find(_GLM_ARG_KEY_CLOSE, ks + len(_GLM_ARG_KEY_OPEN)) + if ke < 0: + break + vstart = ke + len(_GLM_ARG_KEY_CLOSE) + while vstart < len(content) and content[vstart] in " \t\r\n": + vstart += 1 + if not content.startswith(_GLM_ARG_VAL_OPEN, vstart): + # Key without : strict rejects the call; Auto-Heal skips it. + if not allow_incomplete: + valid = False + apos = ke + len(_GLM_ARG_KEY_CLOSE) + continue + vs = vstart + len(_GLM_ARG_VAL_OPEN) + # A first-match find on would truncate values containing literal + # close tags and execute corrupted arguments. + ve = _glm_value_close(content, vs, strict = not allow_incomplete) + key = content[ks + len(_GLM_ARG_KEY_OPEN) : ke].strip() + if ve < 0: + # Unclosed : strict rejects the whole call; Auto-Heal keeps the + # partial value (a truncated query is not a no-arg call). + if not allow_incomplete: + valid = False + break + # Bound the healed value at the next structural tag, not EOF, so a value + # missing only its can't swallow the markup after it. + nk = content.find(_GLM_ARG_KEY_OPEN, vs) + tc = content.find(_GLM_TC_CLOSE, vs) + bounds = [b for b in (nk, tc) if b >= 0] + if not bounds: + args[key] = content[vs:].rstrip() + break + bound = min(bounds) + args[key] = content[vs:bound].rstrip() + apos = bound + continue + raw_val = content[vs:ve] + apos = ve + len(_GLM_ARG_VAL_CLOSE) + # Decode only unambiguous JSON literals; else keep the value RAW so whitespace + # in string args survives (matches vLLM glm4_moe). ``"`` is left out of the + # probe: a verbatim string's quotes are meaningful. + probe = raw_val.strip() + if ( + probe[:1] in "{[" + or probe in ("true", "false", "null") + or _GLM_JSON_NUMERIC_RE.fullmatch(probe) + ): + try: + args[key] = json.loads(probe) + continue + except (json.JSONDecodeError, ValueError): + pass + args[key] = raw_val + + # Strict mode: a block with no is truncated; reject it. + if not allow_incomplete and close < 0: + valid = False + + if name and valid: + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args), + }, + } + ) + pos = close + len(_GLM_TC_CLOSE) if close >= 0 else len(content) + return out + + +# ── Kimi K2 / Moonshot ────────────────────────────────────────────── + + +def _parse_kimi_tool_calls( + content: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Kimi K2. + + ``<|tool_calls_section_begin|><|tool_call_begin|>functions.NAME:IDX + <|tool_call_argument_begin|>{json}<|tool_call_end|>... + <|tool_calls_section_end|>``. Full id is preserved on ``tool_calls + [i].id`` for round-trip through the chat template. Outer loop walks + every section in the stream (vLLM / SGLang parity); mirrors llama.cpp's + Kimi K2 handling via its generalized XML-style parser (llama.cpp PR #16932). + """ + out: list[dict] = [] + outer_pos = 0 + while True: + section_start = content.find(_KIMI_SECTION_BEGIN, outer_pos) + if section_start < 0: + break + scan_start = section_start + len(_KIMI_SECTION_BEGIN) + # Section end OUTSIDE JSON strings: an argument may contain the literal end token, + # and a raw find would drop the later valid call. + section_end = _find_outside_json_strings(content, _KIMI_SECTION_END, scan_start) + scan_end = section_end if section_end >= 0 else len(content) + body = content[scan_start:scan_end] + # Truncated tail: parse what we have, then exit. In strict mode a section with no + # <|tool_calls_section_end|> is truncated; reject it instead. + if section_end < 0: + if allow_incomplete: + out.extend( + _parse_kimi_section_body( + body, id_offset = id_offset + len(out), allow_incomplete = True + ) + ) + return out + outer_pos = section_end + len(_KIMI_SECTION_END) + out.extend( + _parse_kimi_section_body( + body, id_offset = id_offset + len(out), allow_incomplete = allow_incomplete + ) + ) + + # The section wrapper is optional (llama.cpp): a bare <|tool_call_begin|> call parses + # as one section when the loop matched nothing. + if not out and _KIMI_CALL_BEGIN in content: + out.extend( + _parse_kimi_section_body( + content, id_offset = id_offset, allow_incomplete = allow_incomplete + ) + ) + return out + + +def _parse_kimi_section_body( + body: str, + *, + id_offset: int, + allow_incomplete: bool = True, +) -> list[dict]: + """Parse one Kimi K2 section body (between begin / end markers).""" + out: list[dict] = [] + pos = 0 + while pos < len(body): + call_start = body.find(_KIMI_CALL_BEGIN, pos) + if call_start < 0: + break + id_start = call_start + len(_KIMI_CALL_BEGIN) + arg_begin = body.find(_KIMI_ARG_BEGIN, id_start) + if arg_begin < 0: + break + full_id = body[id_start:arg_begin].strip() + m = _KIMI_ID_RE.match(full_id) + if m: + # group(1) is the whole name; do NOT split on ``.`` -- a dotted MCP name stays intact. + name = m.group(1) + else: + base = full_id.split(":")[0] + name = base[len("functions.") :] if base.startswith("functions.") else base + # Drop bare-counter ids (``3``, ``42``) -- matches vLLM; SGLang infers the name + # from the tool schema, which we don't have here. + if name.isdigit(): + json_start = arg_begin + len(_KIMI_ARG_BEGIN) + brace_end = ( + _balanced_brace_end(body, json_start) + if (json_start < len(body) and body[json_start] == "{") + else None + ) + if brace_end is None: + pos = arg_begin + len(_KIMI_ARG_BEGIN) + else: + pos = brace_end + 1 + continue + json_start = arg_begin + len(_KIMI_ARG_BEGIN) + # Balanced brace lets a truncated trailing end marker still surface a call. + while json_start < len(body) and body[json_start] in " \t\n\r": + json_start += 1 + if json_start >= len(body) or body[json_start] != "{": + pos = arg_begin + len(_KIMI_ARG_BEGIN) + continue + brace_end = _balanced_brace_end(body, json_start) + if brace_end is None: + # Malformed / truncated JSON: skip this call but keep parsing later ones + # instead of dropping the rest of the section (vLLM recovers them). + nxt = body.find(_KIMI_CALL_BEGIN, json_start) + if nxt < 0: + break + pos = nxt + continue + try: + args = json.loads(body[json_start : brace_end + 1]) + except (json.JSONDecodeError, ValueError): + pos = brace_end + 1 + continue + if not isinstance(args, dict): + pos = brace_end + 1 + continue + if not allow_incomplete: + # Strict mode: this call must close with <|tool_call_end|> before the next + # <|tool_call_begin|>; otherwise it is truncated, so reject it. + end_marker = body.find(_KIMI_CALL_END, brace_end + 1) + next_call = body.find(_KIMI_CALL_BEGIN, brace_end + 1) + if end_marker < 0 or (next_call >= 0 and end_marker > next_call): + pos = brace_end + 1 + continue + if name: + out.append( + { + "id": full_id or f"call_{id_offset + len(out)}", + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(args), + }, + } + ) + # Advance past the JSON; seeking <|tool_call_end|> could skip a following call + # when this one's end marker is missing. + pos = brace_end + 1 + return out diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 4e27183d88..05dee39283 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -13,6 +13,7 @@ mp.Queue, and exits on shutdown or unload. Pattern follows core/training/worker. from __future__ import annotations import base64 +import json from loggers import get_logger import os import queue as _queue @@ -26,6 +27,9 @@ from typing import Any logger = get_logger(__name__) from utils.hardware import apply_gpu_ids +_SHARE_OBJECT_MAX_BYTES = 1 << 20 +_SHARE_OBJECT_ERROR_SIZE = -1 + # studio/backend root, prepended to sys.path so the spawned subprocess can # import the utils/core packages. _BACKEND_PATH = str(Path(__file__).resolve().parent.parent.parent) @@ -75,6 +79,17 @@ def _send_response(resp_queue: Any, response: dict) -> None: logger.error("Failed to send response: %s", exc) +def _encode_share_object(obj: Any) -> bytes: + data = json.dumps(obj, separators = (",", ":"), ensure_ascii = False).encode("utf-8") + if len(data) > _SHARE_OBJECT_MAX_BYTES: + raise ValueError("Distributed object share payload is too large") + return data + + +def _decode_share_object(data: Any) -> Any: + return json.loads(bytes(data.tolist()).decode("utf-8")) + + def _clean_token(value: str | None) -> str | None: """Normalize an HF token: blank or whitespace-only becomes None.""" return value if value and value.strip() else None @@ -329,14 +344,18 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1", ) try: - success = backend.load_model( - config = mc, - max_seq_length = config.get("max_seq_length", 2048), - load_in_4bit = load_in_4bit, - hf_token = hf_token, - trust_remote_code = trust_remote_code, - gpu_ids = config.get("resolved_gpu_ids"), - ) + load_kwargs = { + "config": mc, + "max_seq_length": config.get("max_seq_length", 2048), + "load_in_4bit": load_in_4bit, + "hf_token": hf_token, + "trust_remote_code": trust_remote_code, + "gpu_ids": config.get("resolved_gpu_ids"), + } + if getattr(backend, "device", None) == "mlx": + load_kwargs["parallel_mode"] = config.get("mlx_parallel_mode") + load_kwargs["distributed_group"] = config.get("_mlx_distributed_group") + success = backend.load_model(**load_kwargs) finally: heartbeat_stop.set() @@ -406,6 +425,32 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: ) +def _drain_skip_generate(cmd: dict, resp_queue: Any, drain_event) -> bool: + """Skip a generate queued behind a cancelled one during an unload. + + The parent sets ``drain_event`` for the whole unload. Because the parent's + per-token ``cancel_event`` is cleared at the start of every generate, a cancel + set while this generate was still queued would otherwise be lost when it is + dequeued. If the drain is in effect, emit an immediate (empty) ``gen_done`` so + the parent's stream/mailbox drains fast and the switch stays fast, and report + the generate was skipped so the caller does not clear the cancel or run it. + """ + if drain_event is None or not drain_event.is_set(): + return False + request_id = cmd.get("request_id", "") + logger.info("Skipping generate for request %s: unload draining", request_id) + _send_response( + resp_queue, + { + "type": "gen_done", + "request_id": request_id, + "cancelled": True, + "stats": None, + }, + ) + return True + + def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None: """Handle a generate command: stream tokens back via resp_queue. @@ -431,6 +476,7 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None: "min_p": cmd.get("min_p", 0.0), "max_new_tokens": cmd.get("max_new_tokens", 256), "repetition_penalty": cmd.get("repetition_penalty", 1.0), + "presence_penalty": cmd.get("presence_penalty", 0.0), "cancel_event": cancel_event, } @@ -494,6 +540,67 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None: ) +def _handle_share_object(backend, cmd: dict, resp_queue: Any) -> None: + """Share a small Python object across MLX distributed ranks.""" + request_id = cmd.get("request_id", "") + group = getattr(backend, "_distributed_group", None) + rank = int(getattr(backend, "_distributed_rank", 0) or 0) + world_size = int(getattr(backend, "_distributed_world_size", 1) or 1) + obj = cmd.get("object") + + try: + if group is None or world_size <= 1: + shared = obj + else: + import mlx.core as mx + if rank == 0: + if obj is None: + mx.eval(mx.distributed.all_sum(mx.array(0), group = group)) + shared = None + else: + try: + data = mx.array(_encode_share_object(obj), dtype = mx.uint8) + except Exception: + mx.eval( + mx.distributed.all_sum( + mx.array(_SHARE_OBJECT_ERROR_SIZE), + group = group, + ) + ) + raise + mx.eval(mx.distributed.all_sum(mx.array(data.size), group = group)) + mx.eval(mx.distributed.all_sum(data, group = group)) + shared = obj + else: + size = int(mx.distributed.all_sum(mx.array(0), group = group).item()) + if size == _SHARE_OBJECT_ERROR_SIZE: + raise RuntimeError("Failed to share distributed object") + if size == 0: + shared = None + else: + data = mx.zeros(size, dtype = mx.uint8) + data = mx.distributed.all_sum(data, group = group) + shared = _decode_share_object(data) + _send_response( + resp_queue, + { + "type": "shared", + "request_id": request_id, + "object": shared, + }, + ) + except Exception as exc: + _send_response( + resp_queue, + { + "type": "share_error", + "request_id": request_id, + "error": str(exc), + "stack": traceback.format_exc(limit = 20), + }, + ) + + def _handle_generate_audio(backend, cmd: dict, resp_queue: Any) -> None: """Handle TTS audio generation — returns WAV bytes + sample_rate.""" request_id = cmd.get("request_id", "") @@ -632,7 +739,14 @@ def _handle_unload(backend, cmd: dict, resp_queue: Any) -> None: ) -def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, config: dict) -> None: +def run_inference_process( + *, + cmd_queue: Any, + resp_queue: Any, + cancel_event, + config: dict, + drain_event = None, +) -> None: """Subprocess entrypoint. Persistent — runs the command loop until shutdown. Args: @@ -640,6 +754,10 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf resp_queue: mp.Queue for sending responses to parent. cancel_event: mp.Event the parent sets to cancel generation. config: Initial configuration dict with model info. + drain_event: mp.Event the parent sets for the duration of an unload. Unlike + cancel_event (cleared at the start of every generate), it is never cleared + here, so a generate still queued behind a cancelled one is skipped rather + than run — the cancel survives the queue handoff. """ os.environ["TOKENIZERS_PARALLELISM"] = "false" os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports @@ -682,9 +800,29 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf exc, ) try: - from core.inference.mlx_inference import MLXInferenceBackend + from core.inference.mlx_inference import MLXInferenceBackend, _init_mlx_distributed backend = MLXInferenceBackend() + if config.get("mlx_distributed"): + group, rank, size = _init_mlx_distributed() + config["_mlx_distributed_group"] = group + if size <= 1: + # A singleton group (MLX built without distributed support, + # or an invalid launch env/hostfile) would leave nonzero ranks + # looping forever on share_distributed_object. Fail the load + # instead of silently continuing without sharding. + raise RuntimeError( + "MLX distributed launch requested but initialized a singleton " + "group (size 1). Ensure the installed MLX has distributed " + "support and the launch environment/hostfile is valid, or run " + "without distributed." + ) + logger.info( + "MLX distributed initialized in worker: rank=%s size=%s mode=%s", + rank, + size, + config.get("mlx_parallel_mode"), + ) _send_response( resp_queue, {"type": "status", "message": "Loading model..."}, @@ -715,8 +853,19 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf cmd_type = cmd.get("type", "") try: if cmd_type == "generate": + if _drain_skip_generate(cmd, resp_queue, drain_event): + continue cancel_event.clear() + # Re-check the drain after clearing: the parent sets drain_event + # then cancel_event for an unload, so if that pair landed between + # the check above and this clear, the clear just erased the unload's + # cancel. Skip here so the outgoing model is not run to completion, + # which would stall the switch until the dispatcher idle-timeout. + if _drain_skip_generate(cmd, resp_queue, drain_event): + continue _handle_generate(backend, cmd, resp_queue, cancel_event) + elif cmd_type == "share_object": + _handle_share_object(backend, cmd, resp_queue) elif cmd_type == "load": if backend.active_model_name: backend.unload_model(backend.active_model_name) @@ -918,9 +1067,21 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf try: if cmd_type == "generate": + if _drain_skip_generate(cmd, resp_queue, drain_event): + continue cancel_event.clear() + # Re-check the drain after clearing: the parent sets drain_event then + # cancel_event for an unload, so if that pair landed between the check + # above and this clear, the clear just erased the unload's cancel. Skip + # here so the outgoing model is not run to completion, which would stall + # the switch until the dispatcher idle-timeout tears the subprocess down. + if _drain_skip_generate(cmd, resp_queue, drain_event): + continue _handle_generate(backend, cmd, resp_queue, cancel_event) + elif cmd_type == "share_object": + _handle_share_object(backend, cmd, resp_queue) + elif cmd_type == "load": if backend.active_model_name: backend.unload_model(backend.active_model_name) diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 345b4dd853..47d26209b4 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -63,6 +63,87 @@ def _install_torchao_stub_once() -> None: install_torchao_windows_rocm_stub() +class UnsafeEmbeddingModelError(RuntimeError): + """Raised when the embedding model repo is flagged unsafe. A distinct type so the + llama-server fallback paths re-raise it instead of masking a security block as a + routine ST failure.""" + + +def _ambient_hf_token() -> str | None: + """The HF token the loader itself would use (HF_TOKEN env or the cached login), so + the scan can reach a gated/private repo instead of failing open. None if unavailable.""" + try: + from huggingface_hub import get_token + return get_token() + except Exception: + return None + + +def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]: + """The module directories a SentenceTransformer load reads weights from, taken from + the repo's ``modules.json`` (each module's non-empty ``path``, e.g. ``0_Transformer``). + ST deserializes ``pytorch_model.bin`` from these dirs, so they are load roots for the + security scan: a flagged pickle directly under one must block. Returns () on any + failure (no modules.json, offline, malformed) so the guard never bricks the embedder. + """ + try: + import json + + from utils.paths import is_local_path + + if is_local_path(name): + from pathlib import Path + from utils.paths import normalize_path + + path = Path(normalize_path(name)).expanduser() / "modules.json" + if not path.is_file(): + return () + data = json.loads(path.read_text()) + else: + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import EntryNotFoundError + + try: + local = hf_hub_download(name, "modules.json", token = token or None) + except EntryNotFoundError: + return () + data = json.loads(open(local).read()) + subdirs = [] + for module in data or (): + sub = str((module or {}).get("path", "")).strip().strip("/") + if sub: + subdirs.append(sub) + return tuple(dict.fromkeys(subdirs)) + except Exception: + return () + + +def _guard_model_security(name: str) -> None: + """Refuse to load a repo HF flagged as unsafe: a poisoned pickle deserializes inside + SentenceTransformer regardless of trust_remote_code. Defense in depth behind the + /settings gate (a name can also arrive via env/default); local paths and unreachable + scans fail open inside evaluate_file_security. Never bricks the embedder on a gate error. + """ + try: + from utils.security import evaluate_file_security, security_load_subdirs + + token = _ambient_hf_token() + # Union the audio-model load roots with the ST module dirs so a flagged pickle + # directly under a Transformer module dir (0_Transformer/) blocks instead of + # passing as an unreferenced nested shard. + load_subdirs = tuple( + dict.fromkeys((*security_load_subdirs(name, token), *_st_module_subdirs(name, token))) + ) + blocked = evaluate_file_security(name, hf_token = token, load_subdirs = load_subdirs).blocked + except Exception: + return + if blocked: + raise UnsafeEmbeddingModelError( + f"Embedding model {name!r} is flagged as unsafe by Hugging Face's security " + "scan; refusing to load. Set a different RAG embedding model." + ) + + def _get(model_name: str | None = None): """Cached SentenceTransformer, (re)loading on a name change. Loaded in fp16 for a ~1.5x speedup at negligible accuracy loss.""" @@ -75,6 +156,7 @@ def _get(model_name: str | None = None): device = _device() logger.info("loading embedding model %s on %s", name, device) + _guard_model_security(name) _model = SentenceTransformer( name, device = device, model_kwargs = {"torch_dtype": "float16"} ) @@ -159,6 +241,8 @@ class _SentenceTransformersBackend: ): try: return _st_encode(texts, model_name = model_name, normalize = normalize) + except UnsafeEmbeddingModelError: + raise # a security block must hard-fail, not fall back to llama-server except Exception as st_err: # noqa: BLE001 - runtime ST/CUDA encode failure # ST loaded but this encode blew up; swap the process to the llama-server # embedder (so later encodes stay in one space) and retry. @@ -222,6 +306,8 @@ def _build_st_backend_or_fallback(): try: backend.warm(model_name = None) return backend + except UnsafeEmbeddingModelError: + raise # a security block must hard-fail, not fall back to llama-server except Exception as st_err: # noqa: BLE001 - any ST/torch import or load failure fallback = _try_make_llama_backend() if fallback is None: @@ -290,6 +376,37 @@ def _reset_backend() -> None: _backend_key = None +def active_backend_is_llama() -> bool: + """True when this process actually embeds via the llama-server (GGUF) backend. + + Reflects the ACTUAL built backend once one exists: an ``auto`` install that + resolves to sentence-transformers but then falls back to llama-server at + runtime (``_build_st_backend_or_fallback`` on a torch/CUDA load failure, or + ``_switch_to_llama_fallback`` on an encode failure) loads only inert GGUF, so + callers gating on the ST pickle must see llama here. Before any backend is + built, defers to the resolver (``auto`` -> ``_resolve_auto()``, else the raw + key) exactly as a fresh process would. Never raises: a backend probe must not + block saving a model.""" + try: + with _backend_lock: + backend = _backend + if backend is not None: + # A backend exists: report what it ACTUALLY is. A concrete + # sentence-transformers backend must return False even if the + # resolver would now pick llama, so its pickle stays gated. If the + # llama import fails we cannot be llama, so fall to the safe False. + try: + from .embed_llama_server import LlamaServerBackend + except Exception: # noqa: BLE001 - llama plumbing import must never block + return False + return isinstance(backend, LlamaServerBackend) + raw = (config.EMBED_BACKEND or "auto").strip().lower() + key = _resolve_auto() if raw in _AUTO_ALIASES else raw + return key in _LLAMA_ALIASES + except Exception: # noqa: BLE001 - a backend probe must never block saving + return False + + def warm(model_name: str | None = None) -> None: """Eagerly load the embedder so the first real request isn't slow.""" _get_backend().warm(model_name = model_name) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index e8367ad08c..1b6b05768a 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -1,38 +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 +# +# Bracket-tag, rehearsal, and thinking-block-strip logic adapted from forge +# (https://github.com/antoinezambelli/forge), Copyright (c) 2025-2026 +# Antoine Zambelli, used under the MIT License. -"""Lightweight tool-call XML parsing and stripping helpers. +"""Lightweight tool-call parsing and stripping helpers. External inference servers import this module without pulling in the inference -orchestrator, structlog, httpx, or the rest of the studio backend. +orchestrator, structlog, httpx, or the rest of the studio backend. Kept in +lockstep with ``core/inference/tool_call_parser.py`` so those servers +(llama-server wrappers, llama-swap, custom shims) reuse the same logic. Any +change here must also land there. + +Handles these serializations (see ``parse_tool_calls_from_text``): + +* ``{json}`` +* ``<|tool_call>call:name{...}`` (Gemma) +* ``v`` +* ``[TOOL_CALLS]name{json}`` (Mistral / Devstral fallback) +* ``name[ARGS]{json}`` (reasoning-model rehearsal) """ +# PEP 604 annotations must stay import-safe on Python 3.9 (requires-python >=3.9). +from __future__ import annotations + +import bisect import json import re -# Pre-compiled patterns for tool XML stripping. The hyphen in the name -# char-class lets dashed MCP tool/parameter names (mcp__srv__list-issues, -# issue-number) parse alongside the built-ins. +# One nesting level in the strip regexes; deeper may leak markup (still parsed). +_BRACKETED_JSON_ONE_LEVEL = r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}" + +# Rehearsal ``name[ARGS]{..}`` strips; group 1 = name for tool-list gating. Closed = +# complete body, tail = truncated; ``(?.*?`` rescans to EOF from every opener +# (quadratic on a stream of unclosed openers). Also reused by the quote-aware Gemma pre-pass. +_TC_JSON_CLOSED_PAT = re.compile(r".*?", re.DOTALL) +_TC_GEMMA_CLOSED_PAT = re.compile(r"<\|tool_call>.*?", re.DOTALL) +_TC_FUNC_CLOSED_PAT = re.compile(r".*?", re.DOTALL) _TOOL_CLOSED_PATS = [ - re.compile(r".*?", re.DOTALL), - re.compile(r"<\|tool_call>.*?", re.DOTALL), + _TC_JSON_CLOSED_PAT, + _TC_GEMMA_CLOSED_PAT, re.compile(r""), - re.compile(r".*?", re.DOTALL), + _TC_FUNC_CLOSED_PAT, + # Mirror the parser regexes: tolerate whitespace and v11 [CALL_ID]/[ARGS] metadata. + re.compile( + r"\[TOOL_CALLS\]\s*[\w-]+(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*" + + _BRACKETED_JSON_ONE_LEVEL, + re.DOTALL, + ), + _REHEARSAL_CLOSED_STRIP_RE, + # Drop the bare v11 [/TOOL_CALLS] closer the balanced scan leaves behind. + re.compile(r"\[/TOOL_CALLS\]"), ] -_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ +# Bare open markers strip a partial call mid-stream; the rehearsal tail needs `{` or EOF +# so prose ``foo[ARGS]`` survives. The XML open-tail forms reach EOF and are reused by +# _tool_call_markup_spans (a think tag in an unclosed call's args stays argument data). +_TOOL_OPEN_XML_TAIL_PATS = [ re.compile(r".*$", re.DOTALL), re.compile(r"<\|tool_call>.*$", re.DOTALL), re.compile(r".*$", re.DOTALL), ] +_TOOL_ALL_PATS = ( + _TOOL_CLOSED_PATS + + _TOOL_OPEN_XML_TAIL_PATS + + [ + re.compile(r"\[TOOL_CALLS\].*$", re.DOTALL), + _REHEARSAL_TAIL_STRIP_RE, + ] +) + +# Rehearsal strips (name in group 1); name-gated via ``enabled_tool_names``, strip-all when None. +_REHEARSAL_STRIP_PATS = frozenset({_REHEARSAL_CLOSED_STRIP_RE, _REHEARSAL_TAIL_STRIP_RE}) + +# Stripped before the quote-aware Gemma helper so a Gemma opener quoted in argument +# data cannot make the helper truncate the block and its tail. +_TOOL_CLOSED_BLOCK_PATS = [_TC_JSON_CLOSED_PAT, _TC_FUNC_CLOSED_PAT] +# A lazy closed-pair pattern whose close token is absent would rescan to EOF from every +# opener; skip that doomed (quadratic) pass. Shared by both strip helpers. +_PAT_REQUIRED_TOKEN = { + _TC_JSON_CLOSED_PAT: "", + _TC_GEMMA_CLOSED_PAT: "", + _TC_FUNC_CLOSED_PAT: "", +} + + +def strip_tool_patterns(text: str, patterns) -> str: + """Apply ``patterns`` in order, skipping closed-pair passes with no close token.""" + for pat in patterns: + token = _PAT_REQUIRED_TOKEN.get(pat) + if token is not None and token not in text: + continue + text = pat.sub("", text) + return text + + +def apply_tool_strip_patterns( + text: str, + patterns, + enabled_tool_names = None, +) -> str: + """Apply strip ``patterns`` to ``text``. A bare rehearsal ``name[ARGS]{..}`` pattern + strips only when ``name`` is an enabled tool (or when ``enabled_tool_names`` is + ``None``); every other pattern is removed unconditionally. A closed-pair pattern whose + close token is absent is skipped so an unclosed-marker stream stays linear.""" + for pat in patterns: + token = _PAT_REQUIRED_TOKEN.get(pat) + if token is not None and token not in text: + continue + if enabled_tool_names is not None and pat in _REHEARSAL_STRIP_PATS: + text = pat.sub(lambda m: "" if m.group(1) in enabled_tool_names else m.group(0), text) + else: + text = pat.sub("", text) + return text + # Pre-compiled patterns for tool-call XML parsing. _TC_JSON_START_RE = re.compile(r"\s*\{") -_TC_GEMMA_START_RE = re.compile(r"<\|tool_call>call:([\w-]+)\s*\{") +_TC_GEMMA_START_RE = re.compile(r"<\|tool_call>\s*call\s*:\s*([\w.\-]+)\s*\{") _TC_FUNC_START_RE = re.compile(r"\s*") _TC_END_TAG_RE = re.compile(r"") _TC_GEMMA_END_TAG_RE = re.compile(r"") _TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") -_TC_PARAM_START_RE = re.compile(r"\s*") +# Horizontal-whitespace trailing class keeps the wrapping newline; _trim_param_value trims it. +_TC_PARAM_START_RE = re.compile(r"[^\S\n]*") _TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") _GEMMA_QUOTE = '<|"|>' _PARAM_CLOSE_TAG = "" @@ -43,7 +142,62 @@ _FUNC_CLOSE_TAG = "" # must be identifier-shaped (start with a letter or underscore); a comma # followed by digits-then-colon is value text such as a timestamp or ratio # (`meet at 10:00, 11:00 tomorrow`), not a new key. -_GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w-]*\s*:") +_GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w.\-]*\s*:") + +# A candidate starting inside a think block is a rehearsal (block kept so literal tags in +# real args survive); ``$`` accepts an unclosed block mid-stream. +_THINK_TAG_RE = re.compile(r".*?(?:|$)|\[THINK\].*?(?:\[/THINK\]|$)", re.DOTALL) +# Bare open/close markers for prefilled-reasoning turns (template opens in the prompt). +_THINK_OPEN_RE = re.compile(r"|\[THINK\]") +_THINK_CLOSE_RE = re.compile(r"|\[/THINK\]") + +# Mistral canonical array: [TOOL_CALLS] + JSON list of {"name","arguments"} objects. +_MISTRAL_ARRAY_RE = re.compile(r"\[TOOL_CALLS\]\s*(?=\[)") + +# Mistral name form + v11 [ARGS]/[CALL_ID] shapes; [CALL_ID] is metadata, not the name, +# and hyphens keep dashed MCP names whole. +_MISTRAL_BRACKET_RE = re.compile( + r"\[TOOL_CALLS\]\s*([\w-]+)(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*(?=\{)" +) + +# Rehearsal ``name[ARGS]{json}`` (no [TOOL_CALLS]); the lookbehind keeps the v11 call-id +# from being taken as the function name. +_REHEARSAL_RE = re.compile(r"(? int | None: + """Return the end index of a balanced JSON object opening at ``start``, + or ``None`` if the braces don't balance. Honors escapes and strings. + """ + if start >= len(text) or text[start] != "{": + return None + depth = 0 + in_string = False + escape = False + for j in range(start, len(text)): + ch = text[j] + if escape: + escape = False + continue + if ch == "\\": + escape = True + continue + if in_string: + if ch == '"': + in_string = False + continue + if ch == '"': + in_string = True + elif ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return j + return None def _balanced_brace_end( @@ -109,6 +263,94 @@ def _balanced_bracket_end(src: str, start: int) -> int: return -1 +def _decode_array_items(text: str, body_start: int, body_end: int): + """Return ``(objs, ends)`` for each top-level element of the JSON array between + ``body_start`` (at or before its ``[``) and ``body_end`` (exclusive): the decoded + object and its absolute exclusive end offset. + + Decoding element-by-element with ``raw_decode`` tolerates the comma-less object + separators the repo's own Mistral/Ollama multi-call templates emit + (``[{...}{...}]``; see ollama_template_mappers.py). A single ``json.loads`` of the + whole body rejects that form and would drop every call. The ends also tile the + region across the calls' spans so a with_spans consumer strips each exactly once.""" + decoder = json.JSONDecoder() + objs: list = [] + ends: list[int] = [] + i = text.find("[", body_start) + if i < 0: + return objs, ends + i += 1 + while i < body_end: + while i < body_end and text[i] in " \t\r\n,": + i += 1 + if i >= body_end or text[i] == "]": + break + try: + obj, rel = decoder.raw_decode(text[i:body_end]) + except (json.JSONDecodeError, ValueError): + break + i += rel + objs.append(obj) + ends.append(i) + return objs, ends + + +def _iter_bracket_spans( + text: str, + start: int = 0, + enabled_tool_names = None, +): + """Yield ``(span_start, span_end, kind, match)`` for each balanced bracket-tag + call from ``start`` on, in document order; ``span_end`` exclusive. ``kind`` is + ``"array"`` ([TOOL_CALLS] [..]), ``"name"`` ([TOOL_CALLS]name{..}, incl. v11 + [CALL_ID]/[ARGS]) or ``"rehearsal"`` (name[ARGS]{..}). + + ``enabled_tool_names`` (set, or None = unrestricted) gates only the ambiguous + bare rehearsal form: name[ARGS]{..} is a call ONLY when ``name`` is enabled, so a + prose ``foo[ARGS]{..}`` (foo disabled) is neither parsed nor stripped. Explicit + [TOOL_CALLS] markers stay unconditional, keeping parse/strip/detection symmetric. + + Balance-only (no JSON validation) so strip and parse share one scan. The cursor + jumps past each consumed span, so a marker inside consumed JSON is never + re-matched and each regex re-searches only once its match falls behind: linear.""" + n = len(text) + specs = ( + ("array", _MISTRAL_ARRAY_RE), + ("name", _MISTRAL_BRACKET_RE), + ("rehearsal", _REHEARSAL_RE), + ) + nexts = {kind: rx.search(text, start) for kind, rx in specs} + cursor = start + while cursor < n: + for kind, rx in specs: + m = nexts[kind] + if m is not None and m.start() < cursor: + nexts[kind] = rx.search(text, cursor) + live = [(kind, m) for kind, m in nexts.items() if m is not None] + if not live: + return + kind, m = min(live, key = lambda km: km[1].start()) + if kind == "array": + end = _balanced_bracket_end(text, m.end()) + end = None if end < 0 else end + else: + end = _balanced_json_span(text, m.end()) + if end is None: + # Truncated body: skip and keep scanning; the caller's catch-all strips the tail. + cursor = m.end() + continue + if ( + kind == "rehearsal" + and enabled_tool_names is not None + and m.group(1) not in enabled_tool_names + ): + # Inactive-name rehearsal is prose: advance past its body without yielding. + cursor = end + 1 + continue + yield (m.start(), end + 1, kind, m) + cursor = end + 1 + + def _split_top_level_commas(src: str) -> list: """Split on commas that are not inside a nested ``[]``/``{}`` or a string.""" parts: list[str] = [] @@ -223,7 +465,7 @@ def _quote_gemma_object_keys(src: str) -> str: while i < len(src) and src[i].isspace(): i += 1 key_name_start = i - while i < len(src) and (src[i].isalnum() or src[i] in "_-"): + while i < len(src) and (src[i].isalnum() or src[i] in "_-."): i += 1 key_name = src[key_name_start:i] colon_pos = i @@ -267,7 +509,8 @@ def _quote_gemma_object_keys(src: str) -> str: json.loads(raw.strip()) parts.append(raw) except (json.JSONDecodeError, ValueError): - parts.append(json.dumps(raw.strip()) if raw.strip() else raw) + # Quote bare value; empty ({k:}) becomes "" so json.loads sees {"k":""} not invalid {"k":}. + parts.append(json.dumps(raw.strip())) else: parts.append(src[key_start:i]) return "".join(parts) @@ -291,9 +534,99 @@ def _inside_open_parameter(content: str, pos: int) -> bool: last_param_start = match.start() if last_param_start < 0: return False - last_param_close = content.rfind(_PARAM_CLOSE_TAG, 0, pos) - last_func_close = content.rfind(_FUNC_CLOSE_TAG, 0, pos) - return last_param_start > max(last_param_close, last_func_close) + # The parameter's OWN close tag decides: if it closes after ``pos`` the position is + # argument data (even across literal function closes); an unclosed one falls back to func close. + own_close = content.find(_PARAM_CLOSE_TAG, last_param_start) + if own_close >= 0: + return own_close > pos + func_close = content.find(_FUNC_CLOSE_TAG, last_param_start) + return func_close < 0 or pos < func_close + + +def _func_close_index(content: str, body_start: int, body: str) -> int: + """Index in ``body`` of the first ```` that is not argument + data (not inside an open parameter value); -1 when every close is data. + Taking the LAST close swallowed prose between the real close and a + literal ```` mentioned later in the answer.""" + idx = body.find(_FUNC_CLOSE_TAG) + while idx >= 0: + if not _inside_open_parameter(content, body_start + idx): + return idx + idx = body.find(_FUNC_CLOSE_TAG, idx + 1) + return -1 + + +def _trim_param_value(val: str) -> str: + """Trim the single wrapping newline the chat template adds around an XML + parameter value, preserving indentation inside VALUE (``str.strip()`` destroyed + code/diff argument indentation).""" + if val.startswith("\n"): + val = val[1:] + if val.endswith("\n"): + val = val[:-1] + return val + + +def _marker_coverage(content: str, markers) -> list[tuple[int, int]]: + """Coverage ``[start, end]`` per marker, used to skip markers that are another + call's data. Closes pair to markers via a per-format stack so an inner close + is not mistaken for the outer's. Unbalanced braces cover to EOF; balanced with + a paired close cover through it (markers before the close are data); balanced + without one cover only the braces, so a later sibling is still recovered.""" + n = len(content) + brace_regions = [(s, be) for (s, be, _k, _m) in markers if be >= 0] + events = [] # (position, order) with order 0 = braces-done, 1 = close marker + for idx, (_start, brace_end, _kind, _m) in enumerate(markers): + if brace_end >= 0: + events.append((brace_end, 0, _kind, idx)) + for kind, close_re in (("json", _TC_END_TAG_RE), ("gemma", _TC_GEMMA_END_TAG_RE)): + for cm in close_re.finditer(content): + # A close inside another call's balanced braces is quoted data; it + # must not pop an earlier close-less marker and swallow a sibling. + if any(s < cm.start() < be for s, be in brace_regions): + continue + events.append((cm.start(), 1, kind, cm.end())) + events.sort(key = lambda e: (e[0], e[1])) + waiting = {"json": [], "gemma": []} + close_end_for: dict[int, int] = {} + for _pos, order, kind, payload in events: + if order == 0: + waiting[kind].append(payload) # marker index, now awaiting its close + elif waiting[kind]: + close_end_for[waiting[kind].pop()] = payload # innermost open marker closes here + coverage = [] + for idx, (start, brace_end, _kind, _m) in enumerate(markers): + if brace_end < 0: + coverage.append((start, n)) + elif idx in close_end_for: + coverage.append((start, close_end_for[idx])) + else: + coverage.append((start, brace_end)) + return coverage + + +def _build_markers(content: str): + """JSON/Gemma tool markers as ``(start, brace_end, kind, match)`` in document + order; ``brace_end < 0`` marks an unbalanced (to-EOF) open.""" + markers = [] + for start_re, gemma, kind in ( + (_TC_JSON_START_RE, False, "json"), + (_TC_GEMMA_START_RE, True, "gemma"), + ): + for m in start_re.finditer(content): + if _inside_open_parameter(content, m.start()): + continue + brace_end = _balanced_brace_end(content, m.end() - 1, gemma_quotes = gemma) + markers.append((m.start(), brace_end, kind, m)) + markers.sort(key = lambda c: c[0]) + return markers + + +def marker_coverage(content: str) -> list[tuple[int, int]]: + """Coverage spans of JSON/Gemma tool markers so other parsers can treat markup + inside a marker's coverage (even a marker that failed to parse) as that call's + data rather than a sibling call.""" + return _marker_coverage(content, _build_markers(content)) def parse_tool_calls_from_text( @@ -301,6 +634,7 @@ def parse_tool_calls_from_text( *, id_offset: int = 0, allow_incomplete: bool = True, + enabled_tool_names = None, with_spans: bool = False, ): """Parse OpenAI-format tool calls from model text. @@ -309,55 +643,61 @@ def parse_tool_calls_from_text( {"name":"web_search","arguments":{"query":"..."}} <|tool_call>call:web_search{query:"..."} ... + [TOOL_CALLS]web_search{"query":"..."} (Mistral / Devstral fallback) + web_search[ARGS]{"query":"..."} (reasoning-model rehearsal) + + A call rehearsed inside a ```` / ``[THINK]`` block is skipped, not + executed; the block is kept so a literal tag in a real argument is preserved. With ``with_spans=True`` returns ``(tool_calls, spans)`` where ``spans[i]`` is the half-open ``(start, end)`` byte range of ``tool_calls[i]``'s markup in ``content`` (including its close tag when present), so a caller can remove exactly the parsed markup and keep every other byte intact. """ + # Candidates starting inside a think block are rehearsals, skipped; blocks are kept, and a + # think marker opening inside a call is argument data (excluded from spans). + _think_spans = _think_spans_outside_tool_markup(content) + _think_starts = [s for s, _e in _think_spans] + + def _in_think(pos: int) -> bool: + # Spans are ordered and non-overlapping; bisect gives O(log M) per candidate. + i = bisect.bisect_right(_think_starts, pos) - 1 + return i >= 0 and _think_spans[i][0] <= pos < _think_spans[i][1] + tool_calls: list[dict] = [] call_spans: list[tuple] = [] - # Collect 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): - if _inside_open_parameter(content, m.start()): - continue - end = _balanced_brace_end(content, m.end() - 1) - if end >= 0: - candidates.append((m.start(), end, "json", m)) - for m in _TC_GEMMA_START_RE.finditer(content): - if _inside_open_parameter(content, m.start()): - continue - end = _balanced_brace_end(content, m.end() - 1, gemma_quotes = True) - if end >= 0: - candidates.append((m.start(), end, "gemma", m)) - candidates.sort(key = lambda c: c[0]) - - candidate_spans = [(s, e) for s, e, _kind, _m in candidates] - for idx, (start, end, kind, m) in enumerate(candidates): - if any(s <= start and end <= e for j, (s, e) in enumerate(candidate_spans) if j != idx): + # Collect JSON/Gemma markers; _marker_coverage decides nesting so a marker inside + # another call's coverage (even one that failed to parse) is data, not executed. A + # marker opening inside a think block is a rehearsal and is skipped. + parsed_items = [] # (start, span_end, name, arguments) in document order + markers = [mk for mk in _build_markers(content) if not _in_think(mk[0])] + coverage = _marker_coverage(content, markers) + for idx, (start, brace_end, kind, m) in enumerate(markers): + if any(s <= start < e for j, (s, e) in enumerate(coverage) if j != idx): continue + if brace_end < 0: + continue # unclosed: not parseable; the fallback still excludes its XML if not allow_incomplete: - tail = content[end + 1 :].lstrip() + tail = content[brace_end + 1 :].lstrip() close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE if close_re.match(tail) is None: continue try: if kind == "json": - obj = json.loads(content[m.end() - 1 : end + 1]) + obj = json.loads(content[m.end() - 1 : brace_end + 1]) name = obj.get("name", "") - arguments = obj.get("arguments", {}) + # Accept ``parameters`` alias for ``arguments`` (Llama-3.2 drift inside Hermes). + arguments = obj.get("arguments") + if arguments is None: + arguments = obj.get("parameters", {}) if isinstance(arguments, dict): arguments = json.dumps(arguments) else: name = m.group(1) - arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : end])) + arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : brace_end])) except (json.JSONDecodeError, ValueError): continue - span_end = end + 1 + span_end = brace_end + 1 close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE ws = len(content[span_end:]) - len(content[span_end:].lstrip()) close_m = close_re.match(content, span_end + ws) @@ -369,7 +709,8 @@ def parse_tool_calls_from_text( fm for fm in _TC_FUNC_START_RE.finditer(content) if not _inside_open_parameter(content, fm.start()) - and not any(s <= fm.start() <= e for s, e in candidate_spans) + and not _in_think(fm.start()) + and not any(s <= fm.start() < e for s, e in coverage) ] for idx, fm in enumerate(func_starts): func_name = fm.group(1) @@ -382,7 +723,7 @@ def parse_tool_calls_from_text( body_end = len(content) body_end = min(body_end, next_func) body = content[body_start:body_end] - close_idx = body.rfind(_FUNC_CLOSE_TAG) + close_idx = _func_close_index(content, body_start, body) if close_idx >= 0: span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG) body = body[:close_idx] @@ -404,7 +745,7 @@ def parse_tool_calls_from_text( val = stripped_val[: -len(_PARAM_CLOSE_TAG)] else: val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[pm.group(1)] = val.strip() + arguments[pm.group(1)] = _trim_param_value(val) else: valid_params = True for pidx, pm in enumerate(param_starts): @@ -422,7 +763,7 @@ def parse_tool_calls_from_text( val = stripped_val[: -len(_PARAM_CLOSE_TAG)] else: val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[param_name] = val.strip() + arguments[param_name] = _trim_param_value(val) if not valid_params: continue @@ -444,19 +785,293 @@ def parse_tool_calls_from_text( } ) call_spans.append((start, span_end)) + + # Patterns 3+4: Mistral [TOOL_CALLS] and bare rehearsal via one balanced scan in document + # order, so a Mistral call and a rehearsal in one message both parse. + if not tool_calls: + for start, end, kind, m in _iter_bracket_spans( + content, enabled_tool_names = enabled_tool_names + ): + if _in_think(start): + continue + # Extend the region over an immediately-following v11 closer so with_spans consumers strip it too. + closer = re.match(r"\s*\[/TOOL_CALLS\]", content[end:]) + region_end = end + closer.end() if closer else end + if kind == "array": + # Decode elements individually (comma-tolerant): one json.loads of the whole + # body rejects the comma-less multi-call arrays Mistral/Ollama templates emit. + payload, item_ends = _decode_array_items(content, m.end(), end) + if not payload: + continue + # Tile the region so every byte belongs to exactly one span; a with_spans consumer + # keeps skipped bytes visible and strips promoted markup exactly once. + tile_start = start + last_span_idx = -1 + for item_idx, item in enumerate(payload): + if not isinstance(item, dict) or "name" not in item: + continue + args = item.get("arguments", {}) + if isinstance(args, str): + # ``arguments`` may itself be a JSON string (OpenAI spec). + try: + args = json.loads(args) + except (json.JSONDecodeError, ValueError): + pass + if not isinstance(args, (dict, str)): + # ``"arguments": null`` (or any non-object scalar) becomes {} like the + # path, not the string "null" auto-heal would mangle to + # a bogus {"query":"null"}. + args = {} + tool_calls.append( + { + "id": f"call_{id_offset + len(tool_calls)}", + "type": "function", + "function": { + "name": item.get("name", ""), + # A bare scalar string stays raw (like the path); + # json.dumps would double-encode it so the arg healer wraps + # "weather" with its literal quotes. + "arguments": args if isinstance(args, str) else json.dumps(args), + }, + } + ) + item_end = item_ends[item_idx] if item_idx < len(item_ends) else region_end + last_span_idx = len(call_spans) + call_spans.append((tile_start, item_end)) + tile_start = item_end + if last_span_idx >= 0: + tile_start, _tile_end = call_spans[last_span_idx] + call_spans[last_span_idx] = (tile_start, region_end) + else: + try: + payload = json.loads(content[m.end() : end]) + except (json.JSONDecodeError, ValueError): + continue + if not isinstance(payload, dict): + continue + tool_calls.append( + { + "id": f"call_{id_offset + len(tool_calls)}", + "type": "function", + "function": { + "name": m.group(1), + "arguments": json.dumps(payload), + }, + } + ) + call_spans.append((start, region_end)) + if with_spans: return tool_calls, call_spans return tool_calls -def strip_tool_call_markup(text: str, *, final: bool = False) -> str: +def _strip_bracket_tag_calls(text: str, enabled_tool_names = None) -> str: + """Strip complete [TOOL_CALLS] arrays / name / bare name[ARGS]{..} calls with one + balanced forward scan, so nested JSON args are removed whole (a fixed-depth regex + left two-level args behind). Truncated tails go to the caller's catch-all. Linear. + ``enabled_tool_names`` gates the rehearsal form (inactive-name prose kept; None + strips every span).""" + if len(text) > _MAX_BRACKET_SCAN_CHARS: + return text + out: list[str] = [] + cursor = 0 + for start, end, _kind, _m in _iter_bracket_spans(text, enabled_tool_names = enabled_tool_names): + out.append(text[cursor:start]) + cursor = end + out.append(text[cursor:]) + return "".join(out) + + +def _tool_call_markup_spans(text: str) -> list[tuple[int, int]]: + """Spans of tool-call markup, so a literal /[THINK] inside a call's args is + stripped WITH the call, not kept as a reasoning block. Covers closed XML/bracket + calls and an unclosed XML call (run via allow_incomplete); without the open-ended + span the unclosed call's markup would leak after execution.""" + # Skip a lazy closed-pair pattern whose close token is absent: its finditer would rescan + # to EOF from every opener (quadratic on a stream of unclosed openers). + spans = [ + m.span() + for pat in _TOOL_CLOSED_PATS + if (_PAT_REQUIRED_TOKEN.get(pat) is None or _PAT_REQUIRED_TOKEN[pat] in text) + for m in pat.finditer(text) + ] + spans.extend((start, end) for start, end, _kind, _m in _iter_bracket_spans(text)) + # An unclosed opener is a real incomplete call only outside closed/bracket spans. + for pat in _TOOL_OPEN_XML_TAIL_PATS: + for m in pat.finditer(text): + if not any(s <= m.start() < e for s, e in spans): + spans.append(m.span()) + return spans + + +def _think_spans_outside_tool_markup(text: str) -> list[tuple[int, int]]: + """/[THINK] block spans, minus any whose opening marker sits INSIDE a + tool-call span (that tag is argument data, not reasoning). Keeping it would drop a + real call after it as rehearsed and leak the call's markup. START tested only, so + a greedy unclosed past the call is still that call's argument data.""" + think_spans = [m.span() for m in _THINK_TAG_RE.finditer(text)] + call_spans = _tool_call_markup_spans(text) + # Prefilled reasoning: the template opens in the prompt, so add a leading span + # (0..close) to skip calls rehearsed there; guarded so a stray close in a normal answer is safe. + close = _THINK_CLOSE_RE.search(text) + if close is not None: + opener = _THINK_OPEN_RE.search(text) + if ( + (opener is None or close.start() < opener.start()) + and not any(cs <= close.start() < ce for cs, ce in call_spans) + and any(cs >= close.end() for cs, ce in call_spans) + ): + think_spans = [(0, close.end())] + think_spans + if not think_spans: + return think_spans + if not call_spans: + return think_spans + return [(s, e) for (s, e) in think_spans if not any(cs <= s < ce for cs, ce in call_spans)] + + +def strip_outside_think(text: str, strip_segment) -> str: + """Apply ``strip_segment(segment, is_last)`` to visible text around /[THINK] + blocks, preserving the blocks verbatim (tool-looking text inside is rehearsal). + ``is_last`` is True only after the final block, so trailing-tail patterns apply + only there. Shared by every strip path so they stay consistent.""" + # A think marker opening inside a complete call is argument text; excluding it lets the + # stripper see the whole call. START-tested, so an unclosed match stays argument data. + think_spans = _think_spans_outside_tool_markup(text) + if not think_spans: + return strip_segment(text, True) + pieces: list[str] = [] + prev = 0 + for s, e in think_spans: + pieces.append(strip_segment(text[prev:s], False)) + pieces.append(text[s:e]) + prev = e + pieces.append(strip_segment(text[prev:], True)) + return "".join(pieces) + + +def _strip_gemma_native_spans(text: str, *, final: bool) -> str: + """Remove complete Gemma-native spans, brace/quote-balanced so a literal + ```` in a quoted argument cannot truncate the span. An incomplete + span is dropped to EOF when ``final``, else kept (still streaming).""" + out: list[str] = [] + cursor = 0 + for match in _TC_GEMMA_START_RE.finditer(text): + start = match.start() + if start < cursor: + continue + brace_end = _balanced_brace_end(text, match.end() - 1, gemma_quotes = True) + if brace_end < 0: + # Unbalanced: nothing completes from here on. Drop the rest if final, + # else keep it; stop either way (rescanning would be quadratic). + if final: + out.append(text[cursor:start]) + cursor = len(text) + break + # Junk between } and is malformed-call markup: strip through + # the close, keep text after it. No close anywhere means stop (linear). + close = _TC_GEMMA_END_TAG_RE.search(text, brace_end + 1) + if close is None: + if final: + out.append(text[cursor:start]) + cursor = len(text) + break + out.append(text[cursor:start]) + cursor = close.end() + out.append(text[cursor:]) + return "".join(out) + + +def _gemma_span_ranges(text: str) -> list: + """``(start, end)`` of each complete Gemma-native span; same walk as + ``_strip_gemma_native_spans`` without stripping.""" + ranges: list[tuple] = [] + cursor = 0 + for match in _TC_GEMMA_START_RE.finditer(text): + start = match.start() + if start < cursor: + continue + brace_end = _balanced_brace_end(text, match.end() - 1, gemma_quotes = True) + if brace_end < 0: + break + close = _TC_GEMMA_END_TAG_RE.search(text, brace_end + 1) + if close is None: + break + ranges.append((start, close.end())) + cursor = close.end() + return ranges + + +def _strip_closed_blocks_outside_gemma(text: str) -> str: + """Closed JSON/function pre-pass that skips matches starting inside a complete + Gemma span: deleting across the span boundary would mangle the Gemma close and + truncate the tail. A skipped match resumes at the covering span's end, so a + real function-XML call after the span is still stripped.""" + ranges = _gemma_span_ranges(text) + if not ranges: + return strip_tool_patterns(text, _TOOL_CLOSED_BLOCK_PATS) + for pat in _TOOL_CLOSED_BLOCK_PATS: + token = _PAT_REQUIRED_TOKEN.get(pat) + if token is not None and token not in text: + continue + out: list[str] = [] + pos = 0 + while True: + m = pat.search(text, pos) + if m is None: + out.append(text[pos:]) + break + covering = next((r for r in ranges if r[0] <= m.start() < r[1]), None) + if covering is not None: + out.append(text[pos : covering[1]]) + pos = covering[1] + continue + out.append(text[pos : m.start()]) + pos = m.end() + new_text = "".join(out) + if new_text != text: + text = new_text + ranges = _gemma_span_ranges(text) + return text + + +def _strip_markup_segment( + text: str, + *, + final: bool, + enabled_tool_names = None, +) -> str: + # Bracket-tag calls (Mistral/rehearsal) first via balanced scan (any nesting depth, + # rehearsal name-gated); then the quote-aware Gemma-native passes so a literal + # in an argument cannot truncate a block; finally the regex XML/tail sweeps. + text = _strip_bracket_tag_calls(text, enabled_tool_names = enabled_tool_names) + text = _strip_closed_blocks_outside_gemma(text) + text = _strip_gemma_native_spans(text, final = final) + patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS + return apply_tool_strip_patterns(text, patterns, enabled_tool_names = enabled_tool_names) + + +def strip_tool_call_markup( + text: str, + *, + final: bool = False, + enabled_tool_names = None, +) -> str: """Strip tool-call XML markup from text. When ``final`` is False, only fully closed tool-call blocks are removed. When ``final`` is True, trailing incomplete tool-call blocks are removed too, and the result is stripped of surrounding whitespace. + + ```` / ``[THINK]`` reasoning is preserved verbatim (see + ``strip_outside_think``); the trailing-tail patterns apply only after the + last block. ``enabled_tool_names`` keeps an inactive-name ``foo[ARGS]{..}`` + example visible (it is prose, not a call) so display cleanup matches detection. """ - patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS - for pat in patterns: - text = pat.sub("", text) - return text.strip() if final else text + result = strip_outside_think( + text, + lambda seg, is_last: _strip_markup_segment( + seg, final = final and is_last, enabled_tool_names = enabled_tool_names + ), + ) + return result.strip() if final else result diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 20b2305a5a..96d1b90b16 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -62,7 +62,6 @@ from loggers import get_logger import time from pathlib import Path from typing import Any, Dict, List, Optional, Callable -from dataclasses import dataclass import pandas as pd from datasets import Dataset from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset @@ -71,7 +70,7 @@ from core.inference.llama_cpp import _hf_offline_if_dns_dead from utils.models import is_vision_model, detect_audio_type from utils.models.model_config import _env_offline from utils.datasets import format_and_template_dataset -from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, TEMPLATE_TO_RESPONSES_MAPPER +from utils.datasets.completion_masking import apply_completion_masking from utils.datasets.iterable import is_streaming_dataset as detect_streaming_dataset from utils.datasets.raw_text import prepare_raw_text_dataset, resolve_column_names from utils.paths import ( @@ -86,6 +85,11 @@ from utils.native_path_leases import child_env_without_native_path_secret from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, ) +from .training import ( + TrainingProgress, + create_mlx_trainer_adapter, + should_use_mlx_training_backend, +) logger = get_logger(__name__) @@ -104,31 +108,16 @@ def _build_report_targets(training_args) -> list[str] | str: return report_to or "none" -@dataclass -class TrainingProgress: - """Training progress tracking""" - - epoch: float = 0 - step: int = 0 - total_steps: int = 0 - loss: Optional[float] = None - learning_rate: Optional[float] = None - is_training: bool = False - is_completed: bool = False - error: Optional[str] = None - status_message: str = "Ready to train" # Current stage - elapsed_seconds: Optional[float] = None - eta_seconds: Optional[float] = None - grad_norm: Optional[float] = None - num_tokens: Optional[int] = None - eval_loss: Optional[float] = None - - class UnslothTrainer: """ Unsloth Training Backend """ + def __new__(cls, *args, **kwargs): + if cls is UnslothTrainer and should_use_mlx_training_backend(): + return create_mlx_trainer_adapter(*args, **kwargs) + return super().__new__(cls) + def __init__(self): self.model = None self.tokenizer = None @@ -3466,8 +3455,6 @@ class UnslothTrainer: # ========== TRAIN ON RESPONSES ONLY ========== # Raw-text datasets always train on all tokens. - instruction_part = None - response_part = None is_cpt = training_args.get("is_cpt", False) train_on_responses_enabled = ( False @@ -3484,113 +3471,93 @@ class UnslothTrainer: # DeepSeek OCR handles this internally in its collator, so skip # Audio VLM handles label masking in its collator, so skip + # Markers auto-detected from the chat template first, manual table + # as fallback; gpt-oss stays on its manual markers. See + # apply_completion_masking. if ( train_on_responses_enabled and not self.is_audio_vlm and not self.is_audio and not (is_deepseek_ocr or dataset_final_format == "alpaca") ): - try: - logger.info("Configuring train on responses only...\n") + from unsloth.chat_templates import train_on_responses_only - # Template mapping for this model - model_name_lower = self.model_name.lower() + logger.info("Configuring train on responses only...\n") - if model_name_lower in MODEL_TO_TEMPLATE_MAPPER: - template_name = MODEL_TO_TEMPLATE_MAPPER[model_name_lower] - logger.info(f"Detected template: {template_name}\n") + def _notify(level, message): + if level == "warning": + logger.warning(message) + else: + logger.info(f"{message}\n") - if template_name in TEMPLATE_TO_RESPONSES_MAPPER: - instruction_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name][ - "instruction" - ] - response_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name]["response"] + # No try/except: the helper handles detection failures and + # double misses itself, so an exception here is a real masking + # failure that must fail the run, not silently train on full + # sequences. + self.trainer, masking_applied = apply_completion_masking( + self.trainer, + self.model_name, + train_on_responses_only, + num_proc = config_args["dataset_num_proc"], + notify = _notify, + ) - logger.info(f"Instruction marker: {instruction_part[:50]}...\n") - logger.info(f"Response marker: {response_part[:50]}...\n") + if not masking_applied: + train_on_responses_enabled = False + + if masking_applied: + try: + # ── Safety net: check if all samples were filtered out ── + # train_on_responses_only masks non-response tokens with -100; a + # row becomes all -100 (Unsloth drops it) when the response + # template is not found in the formatted text. Usually a + # dataset/template mismatch (already-formatted data, or 'Train on + # completions' on data that doesn't match the model's chat + # template); only sometimes max_seq_length truncating the response + # away. Skip this len()-based check for streaming. + if detect_streaming_dataset(self.trainer.train_dataset): + logger.info("Skipping post-filter length check for streaming dataset\n") else: - logger.info( - f"No response mapping found for template: {template_name}\n" + filtered_len = len(self.trainer.train_dataset) + original_dataset_obj = ( + dataset["dataset"] if isinstance(dataset, dict) else dataset ) - train_on_responses_enabled = False - else: - logger.info(f"No template mapping found for model: {self.model_name}\n") - train_on_responses_enabled = False - - except Exception as e: - logger.warning(f"Could not configure train on responses: {e}") - train_on_responses_enabled = False - - # Apply train on responses only if we have valid parts - if ( - train_on_responses_enabled - and instruction_part - and response_part - and not self.is_audio_vlm - and not self.is_audio - and not (is_deepseek_ocr or dataset_final_format == "alpaca") - ): - try: - from unsloth.chat_templates import train_on_responses_only - - self.trainer = train_on_responses_only( - self.trainer, - instruction_part = instruction_part, - response_part = response_part, - num_proc = config_args["dataset_num_proc"], - ) - logger.info("Train on responses only configured successfully\n") - - # ── Safety net: check if all samples were filtered out ── - # train_on_responses_only masks non-response tokens with -100; - # a row becomes all -100 (and Unsloth drops it) when the response - # template is not found in the formatted text. That is usually a - # dataset/template mismatch (already-formatted data, or 'Train on - # completions' applied to data that doesn't match the model's chat - # template), and only sometimes max_seq_length truncating the - # response away. Skip this len()-based check for streaming. - if detect_streaming_dataset(self.trainer.train_dataset): - logger.info("Skipping post-filter length check for streaming dataset\n") - else: - filtered_len = len(self.trainer.train_dataset) - original_dataset_obj = ( - dataset["dataset"] if isinstance(dataset, dict) else dataset - ) - original_len = len(original_dataset_obj) - dropped = original_len - filtered_len - drop_pct = round(100 * dropped / original_len, 1) if original_len > 0 else 0 - - if filtered_len == 0 or drop_pct > 30: - max_seq = training_args.get("max_seq_length", 2048) - error_msg = ( - f"{dropped}/{original_len} samples ({drop_pct}%) were " - f"dropped after applying 'Train on completions': after " - f"masking, those rows had no trainable response tokens " - f"left. The usual cause is that this model's response " - f"template was not found in the formatted samples, so " - f"every token was masked out. That typically means the " - f"dataset is already formatted, or its structure does " - f"not match the model's chat template, so 'Train on " - f"completions' should be turned off for this dataset. " - f"Less commonly, a max_seq_length ({max_seq}) shorter " - f"than the prompt can truncate the response away; only " - f"raise it if your samples are actually longer than that." + original_len = len(original_dataset_obj) + dropped = original_len - filtered_len + drop_pct = ( + round(100 * dropped / original_len, 1) if original_len > 0 else 0 ) - logger.error(error_msg) - self._update_progress(error = error_msg, is_training = False) - return - if dropped > 0: - logger.info( - f"⚠️ {dropped}/{original_len} samples " - f"({drop_pct}%) were dropped (all labels " - f"masked). {filtered_len} samples remain.\n" - ) - logger.info(f"Post-filter dataset size: {filtered_len} samples\n") + if filtered_len == 0 or drop_pct > 30: + max_seq = training_args.get("max_seq_length", 2048) + error_msg = ( + f"{dropped}/{original_len} samples ({drop_pct}%) were " + f"dropped after applying 'Train on completions': after " + f"masking, those rows had no trainable response tokens " + f"left. The usual cause is that this model's response " + f"template was not found in the formatted samples, so " + f"every token was masked out. That typically means the " + f"dataset is already formatted, or its structure does " + f"not match the model's chat template, so 'Train on " + f"completions' should be turned off for this dataset. " + f"Less commonly, a max_seq_length ({max_seq}) shorter " + f"than the prompt can truncate the response away; only " + f"raise it if your samples are actually longer than that." + ) + logger.error(error_msg) + self._update_progress(error = error_msg, is_training = False) + return - except Exception as e: - logger.warning(f"Failed to apply train on responses only: {e}") - train_on_responses_enabled = False + if dropped > 0: + logger.info( + f"⚠️ {dropped}/{original_len} samples " + f"({drop_pct}%) were dropped (all labels " + f"masked). {filtered_len} samples remain.\n" + ) + logger.info(f"Post-filter dataset size: {filtered_len} samples\n") + + except Exception as e: + logger.warning(f"Post-masking dataset size check failed: {e}") else: if train_on_responses_enabled and is_deepseek_ocr: logger.info("Train on responses handled by DeepSeek OCR collator\n") diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index f4233fcf04..2ddda19951 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -14,17 +14,19 @@ import json as _json import math import multiprocessing as mp import os +import platform import queue import re import shutil import threading import time +import traceback import structlog from datetime import datetime, timezone from loggers import get_logger -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from pathlib import Path -from typing import Optional, Tuple, Any, TYPE_CHECKING +from typing import Optional, Tuple, Any, Callable, Union, TYPE_CHECKING if TYPE_CHECKING: import matplotlib.pyplot as plt @@ -98,6 +100,107 @@ def _coerce_optional_nonneg_float(name: str, value): return coerced +def is_apple_silicon_training_platform() -> bool: + return platform.system() == "Darwin" and platform.machine() == "arm64" + + +def is_mlx_training_device(device: Any) -> bool: + return ( + str(device).lower() == "mlx" + or str(device).lower().endswith(".mlx") + or getattr(device, "name", "").lower() == "mlx" + ) + + +def should_use_mlx_training_backend(*, device: Optional[Any] = None) -> bool: + if device is not None: + return is_mlx_training_device(device) + return is_apple_silicon_training_platform() + + +def _build_training_worker_config(values: dict[str, Any]) -> dict[str, Any]: + """Build the normalized worker config shared by Studio and the CLI adapter.""" + config = { + "model_name": values["model_name"], + "project_name": values.get("project_name"), + "training_type": values.get("training_type", "LoRA/QLoRA"), + "hf_token": values.get("hf_token", ""), + "load_in_4bit": values.get("load_in_4bit", True), + "max_seq_length": values.get("max_seq_length", 2048), + "vision_image_size": values.get("vision_image_size"), + "hf_dataset": values.get("hf_dataset", ""), + "local_datasets": values.get("local_datasets"), + "local_eval_datasets": values.get("local_eval_datasets"), + "format_type": values.get("format_type", ""), + "subset": values.get("subset"), + "train_split": values.get("train_split", "train"), + "eval_split": values.get("eval_split"), + "eval_steps": values.get("eval_steps", 0.00), + "dataset_streaming": values.get("dataset_streaming", False), + "dataset_slice_start": values.get("dataset_slice_start"), + "dataset_slice_end": values.get("dataset_slice_end"), + "custom_format_mapping": values.get("custom_format_mapping"), + "is_dataset_image": values.get("is_dataset_image", False), + "is_dataset_audio": values.get("is_dataset_audio", False), + "is_embedding": values.get("is_embedding", False), + "num_epochs": values.get("num_epochs", 3), + "learning_rate": values.get("learning_rate", "2e-4"), + "embedding_learning_rate": values.get("embedding_learning_rate"), + "batch_size": values.get("batch_size", 2), + "gradient_accumulation_steps": values.get("gradient_accumulation_steps", 4), + "warmup_steps": values.get("warmup_steps"), + "warmup_ratio": values.get("warmup_ratio"), + "max_steps": values.get("max_steps", 0), + "save_steps": values.get("save_steps", 0), + "weight_decay": values.get("weight_decay", 0.001), + "max_grad_norm": values.get("max_grad_norm", 0.0), + "max_grad_value": _coerce_optional_nonneg_float( + "max_grad_value", values.get("max_grad_value") + ), + "max_grad_leaf_norm": _coerce_optional_nonneg_float( + "max_grad_leaf_norm", values.get("max_grad_leaf_norm") + ), + "cast_norm_output_to_input_dtype": _coerce_optional_bool( + values.get("cast_norm_output_to_input_dtype"), True + ), + "random_seed": _coerce_seed(values.get("random_seed")), + "packing": values.get("packing", False), + "optim": values.get("optim", "adamw_8bit"), + "lr_scheduler_type": values.get("lr_scheduler_type", "linear"), + "use_lora": values.get("use_lora", True), + "lora_r": values.get("lora_r", 16), + "lora_alpha": values.get("lora_alpha", 16), + "lora_dropout": values.get("lora_dropout", 0.0), + "target_modules": values.get("target_modules"), + "gradient_checkpointing": values.get("gradient_checkpointing", "unsloth"), + "use_rslora": values.get("use_rslora", False), + "use_loftq": values.get("use_loftq", False), + "train_on_completions": values.get("train_on_completions", False), + "finetune_vision_layers": values.get("finetune_vision_layers", True), + "finetune_language_layers": values.get("finetune_language_layers", True), + "finetune_attention_modules": values.get("finetune_attention_modules", True), + "finetune_mlp_modules": values.get("finetune_mlp_modules", True), + "enable_wandb": values.get("enable_wandb", False), + "wandb_token": values.get("wandb_token"), + "wandb_project": values.get("wandb_project", "unsloth-training"), + "enable_tensorboard": values.get("enable_tensorboard", False), + "tensorboard_dir": values.get("tensorboard_dir", "runs"), + "resume_from_checkpoint": values.get("resume_from_checkpoint"), + "trust_remote_code": values.get("trust_remote_code", False), + "approved_remote_code_fingerprint": values.get("approved_remote_code_fingerprint"), + "subject": values.get("subject"), + "gpu_ids": values.get("gpu_ids"), + "s3_config": values.get("s3_config"), + "disable_xet": values.get("disable_xet", False), + } + for key in ("output_dir", "allow_external_output_dir"): + if key in values: + config[key] = values.get(key) + if config["training_type"] == "Full Finetuning": + config["load_in_4bit"] = False + return config + + _HF_TMP_CHECKPOINT_RE = re.compile(r"^tmp-checkpoint-\d+$") @@ -133,7 +236,7 @@ def _s3_dataset_name(s3_dataset: Any) -> Optional[str]: return f"s3://{bucket}/{prefix}" if prefix else f"s3://{bucket}" -def _cleanup_cancelled_checkpoints(output_dir: str | os.PathLike) -> None: +def _cleanup_cancelled_checkpoints(output_dir: Union[str, os.PathLike]) -> None: """Remove only HF Trainer ``tmp-checkpoint-/`` partials after a cancel. Completed ``checkpoint-/`` dirs survive. Symlinked output_dir / children @@ -183,7 +286,7 @@ PLOT_HEIGHT = 3.5 @dataclass class TrainingProgress: - """Mirror of trainer.TrainingProgress so the parent never imports heavy ML modules.""" + """Shared training progress payload for Studio and backend-aware trainers.""" epoch: float = 0 step: int = 0 @@ -200,6 +303,423 @@ class TrainingProgress: num_tokens: Optional[int] = None eval_loss: Optional[float] = None peak_memory_gb: Optional[float] = None + output_dir: Optional[str] = None + + +class _MLXTrainerAdapter: + """Adapts the legacy UnslothTrainer API to the shared Studio MLX worker path.""" + + def __init__(self): + self.model = None + self.tokenizer = None + self.trainer = None + self.training_thread = None + self.training_progress = TrainingProgress() + self.progress_callbacks: list[Callable[[TrainingProgress], None]] = [] + self.is_training = False + self.should_stop = False + self.save_on_stop = True + self.load_in_4bit = True + self.output_dir = None + + self.is_cpt = False + self.is_vlm = False + self.is_audio = False + self.is_audio_vlm = False + self.model_name = None + self.max_seq_length = None + + self._model_config: dict[str, Any] = {} + self._peft_config: dict[str, Any] = {} + self._dataset_config: dict[str, Any] = {} + self._event_queue: Optional[queue.Queue] = None + self._stop_queue: Optional[queue.Queue] = None + self._pump_thread: Optional[threading.Thread] = None + self._lock = threading.Lock() + + def _activate_transformers_for_model(self, model_name: str, hf_token: Optional[str]) -> None: + try: + from utils.transformers_version import activate_transformers_for_subprocess + activate_transformers_for_subprocess(model_name, hf_token) + except Exception as exc: + logger.warning("MLX trainer adapter Transformers activation failed", error = str(exc)) + + def add_progress_callback(self, callback: Callable[[TrainingProgress], None]): + self.progress_callbacks.append(callback) + + def _update_progress(self, **kwargs): + with self._lock: + for key, value in kwargs.items(): + if hasattr(self.training_progress, key): + setattr(self.training_progress, key, value) + progress = self.training_progress + for callback in self.progress_callbacks: + try: + callback(progress) + except Exception: + pass + + def load_model( + self, + model_name: str, + max_seq_length: int = 2048, + load_in_4bit: bool = True, + hf_token: Optional[str] = None, + is_dataset_image: bool = False, + is_dataset_audio: bool = False, + trust_remote_code: bool = False, + full_finetuning: bool = False, + gpu_ids: Optional[list[int]] = None, + ) -> bool: + self.model_name = model_name + self.max_seq_length = max_seq_length + self.load_in_4bit = load_in_4bit + self._audio_type = None + self._activate_transformers_for_model(model_name, hf_token) + try: + from utils.models import detect_audio_type, is_vision_model + + self._audio_type = detect_audio_type(model_name, hf_token) + if self._audio_type == "audio_vlm": + self.is_audio = False + self.is_audio_vlm = bool(is_dataset_audio) + self._audio_type = None + else: + self.is_audio = self._audio_type is not None + self.is_audio_vlm = False + vision = is_vision_model(model_name, hf_token = hf_token) if not self.is_audio else False + self.is_vlm = not self.is_audio_vlm and vision and bool(is_dataset_image) + except Exception as exc: + logger.warning("MLX trainer adapter model type detection failed", error = str(exc)) + self.is_vlm = False + self.is_audio = False + self.is_audio_vlm = False + self.model = object() + self.tokenizer = object() + self._model_config = { + "model_name": model_name, + "max_seq_length": max_seq_length, + "load_in_4bit": load_in_4bit, + "hf_token": hf_token or "", + "is_dataset_image": bool(is_dataset_image), + "is_dataset_audio": bool(is_dataset_audio), + "trust_remote_code": bool(trust_remote_code), + "gpu_ids": gpu_ids, + } + self._update_progress( + is_training = False, + is_completed = False, + error = None, + step = 0, + loss = 0.0, + epoch = 0, + status_message = f"Queued MLX model load: {model_name}", + ) + return True + + def prepare_model_for_training( + self, + use_lora: bool = True, + finetune_vision_layers: bool = True, + finetune_language_layers: bool = True, + finetune_attention_modules: bool = True, + finetune_mlp_modules: bool = True, + target_modules: Optional[Union[list, str]] = None, + lora_r: int = 16, + lora_alpha: int = 16, + lora_dropout: float = 0.0, + use_gradient_checkpointing: Union[str, bool] = "unsloth", + use_rslora: bool = False, + use_loftq: bool = False, + ) -> bool: + self._peft_config = { + "use_lora": bool(use_lora), + "lora_r": lora_r, + "lora_alpha": lora_alpha, + "lora_dropout": lora_dropout, + "target_modules": target_modules, + "gradient_checkpointing": use_gradient_checkpointing, + "use_rslora": bool(use_rslora), + "use_loftq": bool(use_loftq), + "finetune_vision_layers": bool(finetune_vision_layers), + "finetune_language_layers": bool(finetune_language_layers), + "finetune_attention_modules": bool(finetune_attention_modules), + "finetune_mlp_modules": bool(finetune_mlp_modules), + } + self._update_progress(status_message = "Queued MLX training setup") + return True + + def load_and_format_dataset( + self, + dataset_source: Optional[str], + format_type: str = "auto", + local_datasets: Optional[list[str]] = None, + local_eval_datasets: Optional[list[str]] = None, + custom_format_mapping: Optional[dict[str, Any]] = None, + subset: Optional[str] = None, + train_split: str = "train", + eval_split: Optional[str] = None, + dataset_streaming: bool = False, + eval_steps: float = 0.00, + dataset_slice_start: Optional[int] = None, + dataset_slice_end: Optional[int] = None, + is_cpt: bool = False, + s3_config: dict = None, + ) -> Optional[tuple]: + self._dataset_config = { + "hf_dataset": dataset_source or "", + "local_datasets": local_datasets, + "local_eval_datasets": local_eval_datasets, + "format_type": format_type or "", + "custom_format_mapping": custom_format_mapping, + "subset": subset, + "train_split": train_split or "train", + "eval_split": eval_split, + "dataset_streaming": bool(dataset_streaming), + "eval_steps": eval_steps or 0.0, + "dataset_slice_start": dataset_slice_start, + "dataset_slice_end": dataset_slice_end, + "s3_config": s3_config, + } + self.is_cpt = bool(is_cpt) + self._update_progress(status_message = "Queued MLX dataset load") + return ({"dataset": [], "final_format": "deferred_mlx_cli", "success": True}, None) + + def start_training( + self, + dataset = None, + eval_dataset = None, + **training_args, + ) -> bool: + if self.is_training and self.training_thread and self.training_thread.is_alive(): + return False + if self._pump_thread and self._pump_thread.is_alive(): + self._pump_thread.join(timeout = 2.0) + if self._pump_thread.is_alive(): + self._update_progress(error = "Previous training event pump is still finalizing") + return False + if not self._model_config: + self._update_progress(error = "Model not loaded") + return False + if not self._dataset_config: + self._update_progress(error = "Dataset not loaded") + return False + if self.is_cpt: + self._update_progress( + error = "Continued Pretraining is not supported for MLX training yet.", + is_training = False, + is_completed = False, + ) + return False + + config = self._build_worker_config(training_args) + event_queue = queue.Queue() + stop_queue = queue.Queue() + self._event_queue = event_queue + self._stop_queue = stop_queue + self.should_stop = False + self.is_training = True + self.training_progress = TrainingProgress( + is_training = True, + status_message = "Initializing MLX training...", + ) + + self.training_thread = threading.Thread( + target = self._run_training_thread, + args = (config, event_queue, stop_queue), + daemon = True, + ) + self._pump_thread = threading.Thread( + target = self._pump_events, + args = (event_queue, self.training_thread), + daemon = True, + ) + self.training_thread.start() + self._pump_thread.start() + return True + + def _build_worker_config(self, training_args: dict[str, Any]) -> dict[str, Any]: + peft = { + "use_lora": True, + "lora_r": 16, + "lora_alpha": 16, + "lora_dropout": 0.0, + "target_modules": None, + "gradient_checkpointing": "unsloth", + "use_rslora": False, + "use_loftq": False, + "finetune_vision_layers": True, + "finetune_language_layers": True, + "finetune_attention_modules": True, + "finetune_mlp_modules": True, + **self._peft_config, + } + output_dir = training_args.get("output_dir") + if output_dir: + output_dir = os.path.abspath(os.path.expanduser(str(output_dir))) + values = { + **self._model_config, + **self._dataset_config, + **training_args, + "training_type": ( + "Continued Pretraining" + if self.is_cpt + else "LoRA/QLoRA" + if peft["use_lora"] + else "Full Finetuning" + ), + **peft, + "output_dir": output_dir, + "allow_external_output_dir": bool(output_dir), + } + config = _build_training_worker_config(values) + config["resolved_gpu_ids"] = None + config["gpu_selection"] = None + return config + + def _run_training_thread( + self, config: dict[str, Any], event_queue: queue.Queue, stop_queue: queue.Queue + ): + try: + self._run_mlx_worker(config, event_queue, stop_queue) + except Exception as exc: + if event_queue is not None: + event_queue.put( + { + "type": "error", + "error": str(exc), + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + } + ) + + def _run_mlx_worker( + self, config: dict[str, Any], event_queue: queue.Queue, stop_queue: queue.Queue + ): + from .worker import run_mlx_training_process + run_mlx_training_process( + event_queue = event_queue, + stop_queue = stop_queue, + config = config, + ) + + def _pump_events(self, event_queue: queue.Queue, training_thread: threading.Thread): + while True: + event = None + try: + event = event_queue.get(timeout = 0.25) + except queue.Empty: + pass + if event is not None: + self._handle_event(event) + continue + if not training_thread.is_alive(): + self._drain_events(event_queue) + with self._lock: + if self.training_progress.is_training: + self.training_progress.is_training = False + if self.should_stop: + self.training_progress.status_message = "Training stopped." + elif ( + not self.training_progress.error + and not self.training_progress.is_completed + ): + self.training_progress.error = "Training process exited unexpectedly" + self.is_training = False + self._event_queue = None + self._stop_queue = None + return + + def _drain_events(self, event_queue: Optional[queue.Queue] = None): + event_queue = event_queue or self._event_queue + if event_queue is None: + return + while True: + try: + self._handle_event(event_queue.get_nowait()) + except queue.Empty: + return + + def _handle_event(self, event: dict[str, Any]): + etype = event.get("type") + if etype == "status": + self._update_progress( + status_message = event.get("status_message") or event.get("message") or "" + ) + return + if etype == "progress": + self._update_progress( + step = event.get("step", self.training_progress.step), + epoch = event.get("epoch", self.training_progress.epoch), + loss = event.get("loss", self.training_progress.loss), + learning_rate = event.get("learning_rate", self.training_progress.learning_rate), + total_steps = event.get("total_steps", self.training_progress.total_steps), + elapsed_seconds = event.get( + "elapsed_seconds", + self.training_progress.elapsed_seconds, + ), + eta_seconds = event.get("eta_seconds", self.training_progress.eta_seconds), + grad_norm = event.get("grad_norm", self.training_progress.grad_norm), + num_tokens = event.get("num_tokens", self.training_progress.num_tokens), + eval_loss = event.get("eval_loss", self.training_progress.eval_loss), + peak_memory_gb = event.get("peak_memory_gb", self.training_progress.peak_memory_gb), + ) + return + if etype == "complete": + status_message = event.get("status_message") or "Training completed" + output_dir = event.get("output_dir") + was_cancelled = self.should_stop or status_message.strip().lower() in { + "training cancelled", + "training stopped", + } + self.output_dir = output_dir + self._update_progress( + is_training = False, + is_completed = not was_cancelled, + error = None, + status_message = status_message, + output_dir = output_dir, + ) + self.is_training = False + return + if etype == "error": + self._update_progress( + is_training = False, + is_completed = False, + error = event.get("error") or event.get("message") or "Training failed", + ) + self.is_training = False + return + + def stop_training(self, save: bool = True): + self.should_stop = True + self.save_on_stop = bool(save) + if self._stop_queue is not None: + self._stop_queue.put({"type": "stop", "save": save}) + status_message = ( + "Stopping training and saving checkpoint..." if save else "Cancelling training..." + ) + self._update_progress(status_message = status_message) + return True + + def get_training_progress(self) -> TrainingProgress: + pump_thread = self._pump_thread + training_thread = self.training_thread + if ( + pump_thread is not None + and pump_thread.is_alive() + and (training_thread is None or not training_thread.is_alive()) + and threading.current_thread() is not pump_thread + ): + pump_thread.join(timeout = 5.0) + if pump_thread is None or not pump_thread.is_alive(): + self._drain_events() + with self._lock: + return replace(self.training_progress) + + +def create_mlx_trainer_adapter(*args, **kwargs): + return _MLXTrainerAdapter(*args, **kwargs) class TrainingBackend: @@ -296,86 +816,7 @@ class TrainingBackend: # treat this fresh setup as a recoverable death. self._pump_running = False - # Build config dict for the subprocess - config = { - "model_name": kwargs["model_name"], - "project_name": kwargs.get("project_name"), - "training_type": kwargs.get("training_type", "LoRA/QLoRA"), - "hf_token": kwargs.get("hf_token", ""), - "load_in_4bit": kwargs.get("load_in_4bit", True), - "max_seq_length": kwargs.get("max_seq_length", 2048), - "vision_image_size": kwargs.get("vision_image_size"), - "hf_dataset": kwargs.get("hf_dataset", ""), - "local_datasets": kwargs.get("local_datasets"), - "local_eval_datasets": kwargs.get("local_eval_datasets"), - "format_type": kwargs.get("format_type", ""), - "subset": kwargs.get("subset"), - "train_split": kwargs.get("train_split", "train"), - "eval_split": kwargs.get("eval_split"), - "eval_steps": kwargs.get("eval_steps", 0.00), - "dataset_streaming": kwargs.get("dataset_streaming", False), - "dataset_slice_start": kwargs.get("dataset_slice_start"), - "dataset_slice_end": kwargs.get("dataset_slice_end"), - "custom_format_mapping": kwargs.get("custom_format_mapping"), - "is_dataset_image": kwargs.get("is_dataset_image", False), - "is_dataset_audio": kwargs.get("is_dataset_audio", False), - "is_embedding": kwargs.get("is_embedding", False), - "num_epochs": kwargs.get("num_epochs", 3), - "learning_rate": kwargs.get("learning_rate", "2e-4"), - "embedding_learning_rate": kwargs.get("embedding_learning_rate"), - "batch_size": kwargs.get("batch_size", 2), - "gradient_accumulation_steps": kwargs.get("gradient_accumulation_steps", 4), - "warmup_steps": kwargs.get("warmup_steps"), - "warmup_ratio": kwargs.get("warmup_ratio"), - "max_steps": kwargs.get("max_steps", 0), - "save_steps": kwargs.get("save_steps", 0), - "weight_decay": kwargs.get("weight_decay", 0.001), - "max_grad_norm": kwargs.get("max_grad_norm", 0.0), - "max_grad_value": _coerce_optional_nonneg_float( - "max_grad_value", kwargs.get("max_grad_value") - ), - "max_grad_leaf_norm": _coerce_optional_nonneg_float( - "max_grad_leaf_norm", kwargs.get("max_grad_leaf_norm") - ), - "cast_norm_output_to_input_dtype": _coerce_optional_bool( - kwargs.get("cast_norm_output_to_input_dtype"), True - ), - # MLX/CUDA/embedding workers need an int (transformers.set_seed(None) raises). - "random_seed": _coerce_seed(kwargs.get("random_seed")), - "packing": kwargs.get("packing", False), - "optim": kwargs.get("optim", "adamw_8bit"), - "lr_scheduler_type": kwargs.get("lr_scheduler_type", "linear"), - "use_lora": kwargs.get("use_lora", True), - "lora_r": kwargs.get("lora_r", 16), - "lora_alpha": kwargs.get("lora_alpha", 16), - "lora_dropout": kwargs.get("lora_dropout", 0.0), - "target_modules": kwargs.get("target_modules"), - "gradient_checkpointing": kwargs.get("gradient_checkpointing", "unsloth"), - "use_rslora": kwargs.get("use_rslora", False), - "use_loftq": kwargs.get("use_loftq", False), - "train_on_completions": kwargs.get("train_on_completions", False), - "finetune_vision_layers": kwargs.get("finetune_vision_layers", True), - "finetune_language_layers": kwargs.get("finetune_language_layers", True), - "finetune_attention_modules": kwargs.get("finetune_attention_modules", True), - "finetune_mlp_modules": kwargs.get("finetune_mlp_modules", True), - "enable_wandb": kwargs.get("enable_wandb", False), - "wandb_token": kwargs.get("wandb_token"), - "wandb_project": kwargs.get("wandb_project", "unsloth-training"), - "enable_tensorboard": kwargs.get("enable_tensorboard", False), - "tensorboard_dir": kwargs.get("tensorboard_dir", "runs"), - "resume_from_checkpoint": kwargs.get("resume_from_checkpoint"), - "trust_remote_code": kwargs.get("trust_remote_code", False), - "approved_remote_code_fingerprint": kwargs.get("approved_remote_code_fingerprint"), - "subject": kwargs.get("subject"), - "gpu_ids": kwargs.get("gpu_ids"), - "s3_config": kwargs.get("s3_config"), - # Flipped to True only by the HTTP-fallback respawn after a stall. - "disable_xet": kwargs.get("disable_xet", False), - } - - # Full finetuning always runs in 16-bit; LoRA/QLoRA/CPT keep the request. - if config["training_type"] == "Full Finetuning": - config["load_in_4bit"] = False + config = _build_training_worker_config(kwargs) # Split GPU validation from placement around the VRAM hook: # * Explicit gpu_ids are validated here (raises -> the route returns 400 @@ -401,7 +842,7 @@ class TrainingBackend: ) defer_auto_selection = False - if _hw.DEVICE == _hw.DeviceType.MLX: + if should_use_mlx_training_backend(device = _hw.DEVICE): config["resolved_gpu_ids"] = None config["gpu_selection"] = None elif gpu_ids: @@ -1022,17 +1463,22 @@ class TrainingBackend: self._progress.is_training = True elif etype == "complete": - self._progress.is_training = False - self._progress.is_completed = True - self._output_dir = event.get("output_dir") msg = event.get("status_message", "Training completed") + stopped = self._should_stop or msg.strip().lower() in { + "training cancelled", + "training stopped", + } + self._progress.is_training = False + self._progress.is_completed = not stopped + self._output_dir = event.get("output_dir") + self._progress.output_dir = self._output_dir self._progress.status_message = msg if not self._db_run_created and self.current_job_id and self._db_config: db_action = "create_and_finalize" else: db_action = "finalize" db_action_kwargs = { - "status": "stopped" if self._should_stop else "completed", + "status": "stopped" if stopped else "completed", "output_dir": self._output_dir, } diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 17dc1299ca..5fb8fc2bb6 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1309,14 +1309,18 @@ def _normalize_mlx_studio_scheduler(value): def _resolve_mlx_local_dataset_files(file_paths: list) -> list[str]: - """Resolve Studio local dataset uploads without importing the GPU trainer.""" + """Resolve CLI paths and Studio local dataset uploads without importing the GPU trainer.""" from utils.paths import resolve_dataset_path all_files: list[str] = [] for dataset_file in file_paths or []: - file_path = ( - dataset_file if os.path.isabs(dataset_file) else str(resolve_dataset_path(dataset_file)) - ) + dataset_path = Path(os.path.expanduser(str(dataset_file))) + if dataset_path.is_absolute(): + file_path = str(dataset_path) + elif dataset_path.exists(): + file_path = str(dataset_path.resolve()) + else: + file_path = str(resolve_dataset_path(str(dataset_file))) file_path_obj = Path(file_path) if file_path_obj.is_dir(): @@ -1355,6 +1359,58 @@ def _mlx_local_dataset_loader_for_files(files: list[str]) -> str: raise ValueError(f"Unsupported dataset format: {files[0]}") +_MLX_WORKER_COMPLETE = "_mlx_worker_complete" + + +def _start_mlx_stop_poller(stop_queue): + import queue as _queue + import threading + + stop_save = [True] + stop_requested = [False] + trainer_ref = [None] + + def is_stop_requested(): + return stop_requested[0] + + def poll_stop(): + while True: + try: + msg = stop_queue.get(timeout = 0.25) + if msg and msg.get("type") == _MLX_WORKER_COMPLETE: + return + if msg and msg.get("type") == "stop": + stop_save[0] = msg.get("save", True) + stop_requested[0] = True + trainer = trainer_ref[0] + if trainer is not None: + trainer.stop_requested = True + return + except _queue.Empty: + continue + except (EOFError, OSError): + return + + stop_thread = threading.Thread(target = poll_stop, daemon = True) + stop_thread.start() + return stop_save, stop_requested, trainer_ref, is_stop_requested, stop_thread + + +def _resolve_mlx_output_dir(config, model_name): + from utils.paths import resolve_output_dir, default_run_dir_name + + output_dir = config.get("output_dir", "") + if not output_dir: + output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}" + return str(resolve_output_dir(output_dir)) + if config.get("allow_external_output_dir"): + output_path = Path(output_dir).expanduser() + if not output_path.is_absolute(): + output_path = Path.cwd() / output_path + return str(output_path.resolve()) + return str(resolve_output_dir(output_dir)) + + def _run_mlx_training(event_queue, stop_queue, config): """Self-contained MLX training path for Apple Silicon. @@ -1363,8 +1419,6 @@ def _run_mlx_training(event_queue, stop_queue, config): """ import time import math - import threading - import queue as _queue from pathlib import Path def _send(event_type, **kwargs): @@ -1374,31 +1428,9 @@ def _run_mlx_training(event_queue, stop_queue, config): kwargs["message"] = sm event_queue.put({"type": event_type, "ts": time.time(), **kwargs}) - _stop_save = [True] - _stop_requested = [False] - _trainer_ref = [None] - - def _is_stop_requested(): - return _stop_requested[0] - - def _poll_stop(): - while True: - try: - msg = stop_queue.get(timeout = 1.0) - if msg and msg.get("type") == "stop": - _stop_save[0] = msg.get("save", True) - _stop_requested[0] = True - trainer = _trainer_ref[0] - if trainer is not None: - trainer.stop_requested = True - return - except _queue.Empty: - continue - except (EOFError, OSError): - return - - stop_thread = threading.Thread(target = _poll_stop, daemon = True) - stop_thread.start() + _stop_save, _stop_requested, _trainer_ref, _is_stop_requested, _stop_thread = ( + _start_mlx_stop_poller(stop_queue) + ) _send("status", status_message = "Loading MLX libraries...") @@ -1699,6 +1731,7 @@ def _run_mlx_training(event_queue, stop_queue, config): # sharegpt+images) and text (alpaca/sharegpt/chatml → "text" column). format_type = config.get("format_type", "") custom_format_mapping = config.get("custom_format_mapping") + dataset_final_format = "" try: from utils.datasets import format_and_template_dataset def _fmt_progress(status_message = "", **_kw): @@ -1764,6 +1797,7 @@ def _run_mlx_training(event_queue, stop_queue, config): ) if info.get("success", True): dataset = info.get("dataset", dataset) + dataset_final_format = str(info.get("final_format", "") or "").lower() if eval_dataset is not None: ev = format_and_template_dataset( eval_dataset, @@ -1804,21 +1838,14 @@ def _run_mlx_training(event_queue, stop_queue, config): # ── 5. Build output dir ── # Resolve to ~/.unsloth/studio/outputs/ so the export page finds it - from utils.paths import resolve_output_dir, ensure_dir + from utils.paths import ensure_dir - output_dir = config.get("output_dir", "") - if not output_dir: - output_dir = build_default_output_dir_name( - model_name, - config.get("project_name"), - ) - output_dir = str(resolve_output_dir(output_dir)) + output_dir = _resolve_mlx_output_dir(config, model_name) ensure_dir(Path(output_dir)) # ── 6. Create trainer ── eval_steps_val = config.get("eval_steps", 0) or 0 if isinstance(eval_steps_val, float) and 0 < eval_steps_val < 1: - # Studio sometimes sends fraction-of-total-steps eval_steps_val = max(1, int(eval_steps_val * max_steps)) else: eval_steps_val = int(eval_steps_val) @@ -1869,6 +1896,9 @@ def _run_mlx_training(event_queue, stop_queue, config): eval_steps = eval_steps_val, ) + # Also gates the masking skip below, so defined outside the feature-detect block. + raw_text_mode = training_type == "Continued Pretraining" or format_type == "raw" + # Feature-detect optional fields so this PR works without the paired zoo bump. _supported_fields = getattr(MLXTrainingConfig, "__dataclass_fields__", {}) if "cast_norm_output_to_input_dtype" in _supported_fields: @@ -1882,7 +1912,6 @@ def _run_mlx_training(event_queue, stop_queue, config): if "max_grad_leaf_norm" in _supported_fields: mlx_config_kwargs["max_grad_leaf_norm"] = max_grad_leaf_norm if "append_eos" in _supported_fields: - raw_text_mode = training_type == "Continued Pretraining" or format_type == "raw" # Studio SFT formatting owns rendered examples; raw/CPT text still # needs MLX to append EOS like the CUDA raw-text path. mlx_config_kwargs["append_eos"] = bool(raw_text_mode) @@ -1903,29 +1932,27 @@ def _run_mlx_training(event_queue, stop_queue, config): _send("eval_configured") # ── 7. Apply train_on_responses_only if requested ── - if config.get("train_on_completions", False): + # Auto-detect markers from the chat template first, manual table as + # fallback. Mirror the CUDA skips: raw/CPT text has no chat turns and + # Alpaca-rendered text lacks the chat markers. Also check the resolved + # format, since format_type="auto" can land on alpaca or raw text. + if ( + config.get("train_on_completions", False) + and not raw_text_mode + and format_type != "alpaca" + and dataset_final_format not in ("alpaca", "raw_text") + ): _send("status", status_message = "Configuring response-only training...") - try: - from utils.datasets import ( - MODEL_TO_TEMPLATE_MAPPER, - TEMPLATE_TO_RESPONSES_MAPPER, - ) - - template_name = MODEL_TO_TEMPLATE_MAPPER.get(model_name.lower()) - markers = TEMPLATE_TO_RESPONSES_MAPPER.get(template_name) if template_name else None - if markers: - trainer = train_on_responses_only( - trainer, - instruction_part = markers["instruction"], - response_part = markers["response"], - ) - else: - _send( - "status", - status_message = f"train_on_completions skipped (no template for {model_name})", - ) - except Exception as e: - _send("status", status_message = f"train_on_completions failed: {e}") + # No catch: the helper handles detection failures and double misses, so + # an exception here is a real masking failure that must fail the run, + # not silently train on full sequences. + from utils.datasets.completion_masking import apply_completion_masking + trainer, _masking_applied = apply_completion_masking( + trainer, + model_name, + train_on_responses_only, + notify = lambda level, message: _send("status", status_message = message), + ) # ── 8. Setup wandb / tensorboard ── wandb_run = None @@ -2043,12 +2070,27 @@ def _run_mlx_training(event_queue, stop_queue, config): # ── 11. Run training ── gc.collect() mx.synchronize() - trainer.train(resume_from_checkpoint = resume_from_checkpoint) + _save_model = trainer.save_model + + def _skip_internal_final_save(*args, **kwargs): + raise ValueError("worker owns final save") + + trainer.save_model = _skip_internal_final_save + try: + trainer.train(resume_from_checkpoint = resume_from_checkpoint) + finally: + trainer.save_model = _save_model # ── 12. Save and finalize ── - if trainer.stop_requested and not _stop_save[0]: - # User clicked "Cancel" (save=False) — skip saving - _send("complete", output_dir = None, status_message = "Training cancelled") + if trainer.stop_requested: + if not _stop_save[0]: + # Cancel (save=False): skip saving. + _send("complete", output_dir = None, status_message = "Training cancelled") + else: + _send("status", status_message = "Saving stopped model...") + mx.synchronize() + trainer.save_model(output_dir) + _send("complete", output_dir = output_dir, status_message = "Training stopped") else: _send("status", status_message = "Saving model...") mx.synchronize() @@ -2067,6 +2109,79 @@ def _run_mlx_training(event_queue, stop_queue, config): pass +def _is_current_process_apple_silicon() -> bool: + import platform + return platform.system() == "Darwin" and platform.machine() == "arm64" + + +def run_mlx_training_process( + *, + event_queue: Any, + stop_queue: Any, + config: dict, + transformers_activated: bool = False, +) -> None: + """MLX worker entrypoint shared by Studio subprocesses and the CLI adapter.""" + model_name = config["model_name"] + + backend_path = str(Path(__file__).resolve().parent.parent.parent) + if backend_path not in sys.path: + sys.path.insert(0, backend_path) + + from utils.hf_xet_fallback import child_should_disable_xet + + if child_should_disable_xet(config): + os.environ["HF_HUB_DISABLE_XET"] = "1" + os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0" + + if not transformers_activated: + # Must precede detect_hardware(): its MLX stack check imports mlx_lm, hence transformers. + _activate_transformers_version_or_warn(model_name, config.get("hf_token") or None) + + from utils.hardware import hardware as _hw + + _hw.detect_hardware() + if _hw.DEVICE != _hw.DeviceType.MLX: + event_queue.put( + { + "type": "error", + "error": "MLX training requires Apple Silicon with the MLX backend available.", + "stack": "", + "ts": time.time(), + } + ) + return + + if config.get("is_dataset_audio"): + event_queue.put( + { + "type": "error", + "error": "Audio dataset training is not yet supported on Apple Silicon.", + "stack": "", + "ts": time.time(), + } + ) + return + + try: + try: + _run_mlx_training(event_queue, stop_queue, config) + finally: + try: + stop_queue.put({"type": _MLX_WORKER_COMPLETE}) + except (EOFError, OSError, ValueError): + pass + except Exception as exc: + event_queue.put( + { + "type": "error", + "error": str(exc), + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + } + ) + + def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> None: """Subprocess entrypoint. Fresh Python — no stale module state. @@ -2141,36 +2256,26 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> if backend_path not in sys.path: sys.path.insert(0, backend_path) + from .training import is_apple_silicon_training_platform, should_use_mlx_training_backend + + mlx_backend_requested = is_apple_silicon_training_platform() + + mlx_transformers_activated = False + if mlx_backend_requested and _is_current_process_apple_silicon(): + # Must precede detect_hardware(): its MLX stack check imports mlx_lm, hence transformers. + _activate_transformers_version_or_warn(model_name, config.get("hf_token") or None) + mlx_transformers_activated = True + from utils.hardware import hardware as _hw _hw.detect_hardware() - if _hw.DEVICE == _hw.DeviceType.MLX: - if config.get("is_dataset_audio"): - event_queue.put( - { - "type": "error", - "error": "Audio dataset training is not yet supported on Apple Silicon.", - "stack": "", - "ts": time.time(), - } - ) - return - # Activate correct transformers version (Gemma-4 needs a 5.x sidecar, etc.) - # Must happen before any transformers/mlx-lm imports in _run_mlx_training. - # Non-fatal: fall through with whatever version is installed, but log - # the failure instead of swallowing it (issue #6103). - _activate_transformers_version_or_warn(model_name, config.get("hf_token") or None) - try: - _run_mlx_training(event_queue, stop_queue, config) - except Exception as exc: - event_queue.put( - { - "type": "error", - "error": str(exc), - "stack": traceback.format_exc(limit = 20), - "ts": time.time(), - } - ) + if mlx_backend_requested or should_use_mlx_training_backend(device = _hw.DEVICE): + run_mlx_training_process( + event_queue = event_queue, + stop_queue = stop_queue, + config = config, + transformers_activated = mlx_transformers_activated, + ) return # ── 1. Activate correct transformers version BEFORE any ML imports ── @@ -2693,7 +2798,8 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> if backend_path not in sys.path: sys.path.insert(0, backend_path) - from core.training.trainer import UnslothTrainer, TrainingProgress + from core.training.training import TrainingProgress + from core.training.trainer import UnslothTrainer from utils.paths import ( ensure_dir, resolve_output_dir, diff --git a/studio/backend/hub/tests/test_model_services.py b/studio/backend/hub/tests/test_model_services.py index f05d8359ec..44701c0b64 100644 --- a/studio/backend/hub/tests/test_model_services.py +++ b/studio/backend/hub/tests/test_model_services.py @@ -105,6 +105,10 @@ def test_repo_id_validation_accepts_hf_repo_id_contract(repo_id): assert paths.is_valid_repo_id(repo_id) +def test_repo_id_validation_accepts_max_length_namespaced_repo(): + assert paths.is_valid_repo_id(f"{'a' * 96}/{'b' * 96}") + + @pytest.mark.parametrize( "repo_id", [ @@ -121,6 +125,48 @@ def test_repo_id_validation_rejects_unsafe_or_invalid_ids(repo_id): assert not paths.is_valid_repo_id(repo_id) +def test_download_state_preserves_readable_keys_when_safe(monkeypatch, tmp_path): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path) + + path = state_dir.marker_path("model", "Owner/Repo", "Q4_K_M") + + assert path is not None + assert path.name == "models--owner--repo--variant--q4_k_m.json" + + +@pytest.mark.parametrize("variant", ["bad variant with spaces", "q" * 64]) +def test_download_state_bounds_long_repo_variant_filenames(monkeypatch, tmp_path, variant): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path) + repo_id = f"{'a' * 96}/{'b' * 96}" + + assert paths.is_valid_repo_id(repo_id) + assert download_manifest.write_cancel_marker("model", repo_id, variant, "http") + assert download_manifest.write_manifest( + "model", + repo_id, + variant, + [download_manifest.ExpectedFile(path = "model.gguf", size = 1)], + "http", + ) + + marker_path = state_dir.marker_path("model", repo_id, variant) + manifest_path = state_dir.manifest_path("model", repo_id, variant) + + assert marker_path is not None + assert manifest_path is not None + assert "--sha256-" in marker_path.name + assert len(marker_path.name.encode("utf-8")) <= 255 + assert len(f".{marker_path.name}.tmp-00000000".encode("utf-8")) <= 255 + assert download_manifest.has_cancel_marker("model", repo_id, variant) + assert download_manifest.read_manifest("model", repo_id, variant) is not None + assert list(download_manifest.iter_variant_markers("model", repo_id)) == [ + (variant, marker_path) + ] + assert list(download_manifest.iter_variant_manifests("model", repo_id)) == [ + (variant, manifest_path) + ] + + class _RecordingLogger: def __init__(self): self.warnings = [] diff --git a/studio/backend/hub/utils/gguf_plan.py b/studio/backend/hub/utils/gguf_plan.py index 2abdb0fb79..18daa4f84e 100644 --- a/studio/backend/hub/utils/gguf_plan.py +++ b/studio/backend/hub/utils/gguf_plan.py @@ -102,16 +102,17 @@ def preferred_mmproj_sibling(siblings: Sequence) -> Optional[object]: def preferred_mtp_sibling(siblings: Sequence) -> Optional[object]: """The separate MTP drafter to fetch with every variant: the repo-root ``mtp-*.gguf`` copy unsloth ships for llama.cpp ``-hf`` auto-discovery - (Gemma 4). Same pick as the loader's drafter resolution (``mtp-`` basename - prefix, first in sort order) so download and load resolve the same file; - the higher-precision ``MTP/`` subdir copies are for explicit selection and - are not auto-fetched. None for repos with the head baked into the main - GGUF (Qwen).""" + (Gemma 4). Same pick as the loader's drafter resolution (root-level + ``mtp-`` prefix, first in sort order) so download and load resolve the same + file; the higher-precision ``MTP/`` subdir copies are for explicit + selection and are not auto-fetched. None for repos with the head baked into + the main GGUF (Qwen).""" + # Root-level only: the MTP/ subdir copies now share the mtp- prefix too. candidates = sorted( ( s for s in siblings - if (name := _gguf_rfilename(s)) and name.lower().rsplit("/", 1)[-1].startswith("mtp-") + if (name := _gguf_rfilename(s)) and "/" not in name and name.lower().startswith("mtp-") ), key = lambda s: getattr(s, "rfilename"), ) diff --git a/studio/backend/hub/utils/paths.py b/studio/backend/hub/utils/paths.py index afcb0b41dc..5435202565 100644 --- a/studio/backend/hub/utils/paths.py +++ b/studio/backend/hub/utils/paths.py @@ -181,15 +181,20 @@ def is_valid_repo_id(repo_id: str) -> bool: """Validate Hugging Face ``repo_name`` or ``namespace/repo_name`` IDs.""" if not repo_id or repo_id != repo_id.strip(): return False - if len(repo_id) > _MAX_REPO_ID_LENGTH or repo_id.endswith(".git"): + if repo_id.endswith(".git"): return False if "--" in repo_id or ".." in repo_id: return False segments = repo_id.split("/") if len(segments) not in (1, 2): return False + # Match huggingface_hub.validate_repo_id: the 96-char limit applies per + # segment (repo name / namespace), not to the whole "namespace/repo_name" + # string, so long-but-valid repo names are not falsely rejected. return all( - segment not in ("", ".", "..") and _VALID_REPO_ID_SEGMENT.fullmatch(segment) is not None + segment not in ("", ".", "..") + and len(segment) <= _MAX_REPO_ID_LENGTH + and _VALID_REPO_ID_SEGMENT.fullmatch(segment) is not None for segment in segments ) diff --git a/studio/backend/hub/utils/state_dir.py b/studio/backend/hub/utils/state_dir.py index a304477a3d..183e934724 100644 --- a/studio/backend/hub/utils/state_dir.py +++ b/studio/backend/hub/utils/state_dir.py @@ -11,8 +11,9 @@ cache lifecycle. Two subdirectories: manifests/ .json per-download expected-files manifest cancelled/ .json per-download cancel marker -The ```` mirrors HF's cache dir naming so a state file can be -eyeballed next to the on-disk repo it describes: +The ```` mirrors HF's cache dir naming while the resulting manifest, +cancel-marker, and atomic-write temp filenames fit common filesystem basename +limits. Very long repo IDs use a stable hash in the state key: models---- full snapshot models------variant-- GGUF variant @@ -49,6 +50,11 @@ _MANIFESTS_SUBDIR = "manifests" _CANCELLED_SUBDIR = "cancelled" _WORKERS_SUBDIR = "workers" _SAFE_VARIANT_FRAGMENT = re.compile(r"^[a-z0-9._-]{1,64}$") +_MAX_STATE_BASENAME_BYTES = 255 +_STATE_EXTENSION = ".json" +# _atomic_write_json writes "..tmp-<8hex>" beside the final file. +_ATOMIC_WRITE_TMP_OVERHEAD = len(".") + len(".tmp-") + 8 +_MAX_VARIANT_FRAGMENT_LENGTH = 64 def state_root() -> Optional[Path]: @@ -84,16 +90,35 @@ def repo_cache_basename(repo_type: RepoType, repo_id: str) -> str: return f"{repo_type}s--{repo_id.replace('/', '--')}".lower() +def _filename_bytes(name: str) -> int: + return len(name.encode("utf-8")) + + +def _state_filename_fits(entry_key: str) -> bool: + filename = f"{entry_key}{_STATE_EXTENSION}" + return _filename_bytes(filename) + _ATOMIC_WRITE_TMP_OVERHEAD <= _MAX_STATE_BASENAME_BYTES + + +def _state_repo_key(repo_type: RepoType, repo_id: str) -> str: + base = repo_cache_basename(repo_type, repo_id) + variant_prefix = f"{base}--variant--" + longest_variant_key = f"{variant_prefix}{'x' * _MAX_VARIANT_FRAGMENT_LENGTH}" + if _state_filename_fits(longest_variant_key): + return base + digest = hashlib.sha256(base.encode("utf-8")).hexdigest()[:32] + return f"{repo_type}s--sha256-{digest}" + + def variant_filename_prefix(repo_type: RepoType, repo_id: str) -> str: """Lowercased prefix every variant-keyed state file for this repo shares. The single source the download_manifest enumerators match against, so the scheme in :func:`_entry_key` cannot drift from them silently.""" - return f"{repo_cache_basename(repo_type, repo_id)}--variant--" + return f"{_state_repo_key(repo_type, repo_id)}--variant--" def _entry_key(repo_type: RepoType, repo_id: str, variant: Optional[str]) -> str: - base = repo_cache_basename(repo_type, repo_id) + base = _state_repo_key(repo_type, repo_id) if not variant: return base normalized_variant = variant.strip().lower() diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 31c100dbec..53b0f14b09 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -177,6 +177,7 @@ class GenerateRequest(BaseModel): temperature: float = Field(0.6, ge = 0.0, le = 2.0, description = "Sampling temperature") top_p: float = Field(0.95, ge = 0.0, le = 1.0, description = "Top-p sampling") top_k: int = Field(20, ge = -1, le = 100, description = "Top-k sampling") + min_p: float = Field(0.0, ge = 0.0, le = 1.0, description = "Min-p sampling") max_new_tokens: int = Field(2048, ge = 1, le = 4096, description = "Maximum tokens to generate") repetition_penalty: float = Field(1.0, ge = 1.0, le = 2.0, description = "Repetition penalty") presence_penalty: float = Field(0.0, ge = 0.0, le = 2.0, description = "Presence penalty") @@ -1146,7 +1147,8 @@ class CompletionMessage(BaseModel): """The assistant's complete response message.""" role: Literal["assistant"] = "assistant" - content: str + # ``None`` on a pure tool-call turn (OpenAI content=null); string otherwise. + content: Optional[str] = None refusal: Optional[str] = None reasoning_content: Optional[str] = None tool_calls: Optional[list[dict]] = None @@ -1531,12 +1533,41 @@ class AnthropicToolResultBlock(BaseModel): tool_use_id: str content: Union[str, list] = "" + @field_validator("content", mode = "before") + @classmethod + def _coerce_null_content(cls, v): + # Some clients send null content for an empty tool result; the str|list + # union would 400 on it, so treat null as "". + return "" if v is None else v + + +# Block types the converter translates explicitly. Anything else (thinking / +# redacted_thinking, a provider block a resumed session replays, or a future type) +# is accepted as an unknown block and dropped by the converter, rather than 400-ing +# the whole request on strict validation. +_KNOWN_ANTHROPIC_BLOCK_TYPES = frozenset({"text", "image", "tool_use", "tool_result"}) + + +class AnthropicUnknownBlock(BaseModel): + type: str + model_config = {"extra": "allow"} + + @field_validator("type") + @classmethod + def _only_unknown_types(cls, v): + # Known types parse as their typed models above (so a malformed known block + # still fails cleanly); this fallback only catches the rest. + if v in _KNOWN_ANTHROPIC_BLOCK_TYPES: + raise ValueError("known block type handled by its typed model") + return v + AnthropicContentBlock = Union[ AnthropicTextBlock, AnthropicImageBlock, AnthropicToolUseBlock, AnthropicToolResultBlock, + AnthropicUnknownBlock, ] @@ -1581,6 +1612,40 @@ class AnthropicMessage(BaseModel): role: Literal["user", "assistant"] content: Union[str, list[AnthropicContentBlock]] + @model_validator(mode = "before") + @classmethod + def _normalize_content(cls, data): + # Role-aware leniency that never silently drops real user input: + # - assistant: a resumed tool-only turn's null content -> "" (str|list would + # 400 on null; "" keeps the converter's `for block in content` safe). + # Unknown blocks (thinking / future types) validate via + # AnthropicUnknownBlock and are dropped by the converter. + # - user: keep strict. Null user content stays None so str|list rejects it + # (400) rather than forwarding an empty prompt; and reject block types the + # converter cannot translate, since it silently skips unknown user blocks + # -- a user turn made only of them would validate yet send no content + # (silent data loss). + if not isinstance(data, dict): + return data + content = data.get("content") + if data.get("role") == "assistant": + # Coerce only an explicit null (resumed tool-only turn). A missing + # content key stays malformed so the required-field check still 400s. + if "content" in data and content is None: + return {**data, "content": ""} + return data + if isinstance(content, list): + for block in content: + btype = ( + block.get("type") if isinstance(block, dict) else getattr(block, "type", None) + ) + # Guard the value: a non-string type is unsupported too, and a + # membership test on an unhashable value would raise TypeError + # (escaping as a 500 instead of a clean 400). + if not isinstance(btype, str) or btype not in _KNOWN_ANTHROPIC_BLOCK_TYPES: + raise ValueError(f"unsupported content block type {btype!r} in a user message") + return data + class AnthropicTool(BaseModel): # Client tools have input_schema; server tools may only have type/name. diff --git a/studio/backend/requirements/single-env/overrides-darwin-arm64.txt b/studio/backend/requirements/single-env/overrides-darwin-arm64.txt index 56b948644f..43f37b3183 100644 --- a/studio/backend/requirements/single-env/overrides-darwin-arm64.txt +++ b/studio/backend/requirements/single-env/overrides-darwin-arm64.txt @@ -10,3 +10,9 @@ transformers>=4.57.6 # anyio that also ImportErrors on TaskHandle and 500s the server. An override # wins the fight, so force one consistent <4.14 here too. anyio<4.14.0 + +# mlx-lm 0.31.3 regressed QK-norm archs (gemma4 / qwen3_5): strict load_weights +# rejects q_norm/k_norm, so those checkpoints fail to load. mlx-lm #1242. +# The override also drops it from transitive resolution; keep the >=0.22.0 floor +# (mirrors mlx_repair.py _MLX_MIN_VERSIONS) or the resolver could go below it. +mlx-lm>=0.22.0,!=0.31.3 diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py index 2c87ce8c6e..7a27a58a52 100644 --- a/studio/backend/routes/chat_history.py +++ b/studio/backend/routes/chat_history.py @@ -56,6 +56,7 @@ class ChatThread(BaseModel): projectId: Optional[str] = None archived: bool = False createdAt: int + updatedAt: Optional[int] = None openaiCodeExecContainerId: Optional[str] = None anthropicCodeExecContainerId: Optional[str] = None forkedFromThreadId: Optional[str] = None @@ -70,6 +71,7 @@ class ChatThreadPatch(BaseModel): projectId: Optional[str] = None archived: Optional[bool] = None createdAt: Optional[int] = None + updatedAt: Optional[int] = None openaiCodeExecContainerId: Optional[str] = None anthropicCodeExecContainerId: Optional[str] = None @@ -177,6 +179,7 @@ class ChatSettingsPayload(BaseModel): collapseHtmlArtifacts: Optional[bool] = None allowArtifactNetworkAccess: Optional[bool] = None autoHealToolCalls: Optional[bool] = None + nudgeToolCalls: Optional[bool] = None maxToolCallsPerMessage: Optional[int] = Field(default = None, ge = 1) toolCallTimeout: Optional[int] = Field(default = None, ge = 1) @@ -251,7 +254,7 @@ async def patch_thread( current_subject: str = Depends(get_current_subject), ): patch = payload.model_dump(exclude_unset = True) - for field in ("title", "modelType", "modelId", "archived", "createdAt"): + for field in ("title", "modelType", "modelId", "archived", "createdAt", "updatedAt"): if field in patch and patch[field] is None: raise HTTPException(status_code = 400, detail = f"{field} cannot be null") if patch.get("projectId") and get_chat_project(patch["projectId"]) is None: diff --git a/studio/backend/routes/data_recipe/seed.py b/studio/backend/routes/data_recipe/seed.py index 57a291291e..a5b75b7335 100644 --- a/studio/backend/routes/data_recipe/seed.py +++ b/studio/backend/routes/data_recipe/seed.py @@ -10,6 +10,7 @@ import binascii import json import os import re +import shutil from itertools import islice from pathlib import Path from typing import Any @@ -59,6 +60,9 @@ UNSTRUCTURED_ALLOWED_EXTS = {".pdf", ".docx", ".txt", ".md"} SEED_UPLOAD_DIR = seed_uploads_root() UNSTRUCTURED_UPLOAD_ROOT = unstructured_uploads_root() _SAFE_ID_RE = re.compile(r"^[a-zA-Z0-9_-]+$") +# Frontend-generated upload namespace (UUID4 hex). Legacy node ids (n1, ...) +# never match: those directories can be shared by several recipes. +_UPLOAD_UID_RE = re.compile(r"^[0-9a-f]{32}$") def _validate_safe_id(value: str, label: str) -> str: @@ -580,6 +584,39 @@ async def remove_unstructured_file(block_id: str, file_id: str): return {"status": "ok"} +@router.delete("/seed/unstructured-block/{block_id}") +async def remove_unstructured_block(block_id: str): + """Delete a block's upload directory; files on disk still count toward its quota. + + Only uid-namespaced directories may be bulk-deleted: they have exactly one + owning block. Legacy node-id directories (n1, ...) can be shared by other + recipes, so they are managed file-by-file instead. + """ + _validate_safe_id(block_id, "block_id") + if not _UPLOAD_UID_RE.match(block_id): + raise HTTPException(400, "Invalid block_id: only uid-namespaced blocks can be deleted") + + block_dir = (UNSTRUCTURED_UPLOAD_ROOT / block_id).resolve() + if not block_dir.is_relative_to(UNSTRUCTURED_UPLOAD_ROOT.resolve()): + raise HTTPException(400, "Invalid block_id: outside upload root") + if not block_dir.exists(): + return {"status": "ok", "deleted": False} + + try: + shutil.rmtree(block_dir) + except OSError as exc: + raise log_and_http_error( + exc, + 500, + "failed to delete uploaded files", + event = "data_recipe.seed.unstructured_block_delete_failed", + log = logger, + ) from exc + if block_dir.exists(): + raise HTTPException(500, "failed to delete uploaded files") + return {"status": "ok", "deleted": True} + + @router.post("/seed/inspect-upload", response_model = SeedInspectResponse) def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectResponse: if payload.file_ids is not None: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 17be222d93..ec5d309810 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -13,7 +13,7 @@ from pathlib import Path from fastapi import APIRouter, Depends, HTTPException, Request, status from fastapi.responses import StreamingResponse, JSONResponse, Response from starlette.requests import ClientDisconnect -from typing import Any, List, Optional, Union +from typing import Any, Callable, List, Optional, Union import json import httpx from loggers import get_logger @@ -28,6 +28,16 @@ import re as _re from utils.models import extract_model_size_b as _extract_model_size_b from utils.api_errors import openai_error_body, anthropic_error_body +from core.inference.llama_admission import ( + LlamaAdmissionCancelled, + LlamaAdmissionConfig, + LlamaAdmissionLease, + LlamaAdmissionQueueFull, + LlamaAdmissionReservation, + LlamaAdmissionTimeout, + get_llama_admission_queue, + llama_admission_config_from_env, +) def _positive_int_or_none(value: Any) -> Optional[int]: @@ -40,6 +50,43 @@ def _positive_int_or_none(value: Any) -> Optional[int]: return value_int if value_int > 0 else None +def _nonnegative_int_or_none(value: Any) -> Optional[int]: + if isinstance(value, bool): + return None + try: + value_int = int(value) + except (TypeError, ValueError): + return None + return value_int if value_int >= 0 else None + + +_MLX_MPI_DISTRIBUTED_ENV_PAIRS = ( + ("OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"), + ("PMI_RANK", "PMI_SIZE"), + ("PMIX_RANK", "PMIX_SIZE"), + ("MPI_RANK", "MPI_WORLD_SIZE"), + ("MV2_COMM_WORLD_RANK", "MV2_COMM_WORLD_SIZE"), +) + + +def _mlx_distributed_launch_detected() -> bool: + if _nonnegative_int_or_none(os.environ.get("MLX_RANK")) is not None: + world_size = _positive_int_or_none(os.environ.get("MLX_WORLD_SIZE")) + if world_size is not None and world_size > 1: + return True + return bool( + os.environ.get("MLX_HOSTFILE") + or os.environ.get("MLX_IBV_DEVICES") + or os.environ.get("MLX_JACCL_COORDINATOR") + or (os.environ.get("NCCL_HOST_IP") and os.environ.get("NCCL_PORT")) + ) + return any( + _nonnegative_int_or_none(os.environ.get(rank_env)) is not None + and (_positive_int_or_none(os.environ.get(size_env)) or 0) > 1 + for rank_env, size_env in _MLX_MPI_DISTRIBUTED_ENV_PAIRS + ) + + def _install_httpcore_asyncgen_silencer() -> None: """Silence benign httpx/httpcore asyncgen GC noise on Python 3.13. @@ -231,10 +278,87 @@ def _effective_max_tokens(payload): ) +_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV = "UNSLOTH_OPENAI_COMPAT_STREAM_STALL_TIMEOUT" + + +def _positive_float_env(env_name: str, default): + """Parse a positive float from an env var. A parseable non-positive value + returns ``None`` (0 disables the guarded feature); only unparseable or unset + values fall back to ``default``.""" + raw_value = os.environ.get(env_name) + if raw_value is None or not raw_value.strip(): + return default + try: + value = float(raw_value.strip()) + except ValueError: + return default + return value if value > 0 else None + + +def _effective_openai_max_tokens_from_values(max_tokens, max_completion_tokens = None): + """Resolve the OpenAI-compatible generation cap from raw request values. + + Prefers ``max_completion_tokens`` over the deprecated ``max_tokens``, and + returns ``None`` when both are omitted so callers keep their context-window + default (OpenAI treats an omitted cap as bounded only by the context + window). Explicit client caps pass through unchanged. + """ + + def _validate_explicit(value, param: str): + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int): + raise HTTPException( + status_code = 400, + detail = openai_error_body( + f"'{param}' must be an integer.", + status = 400, + code = "invalid_type", + param = param, + ), + ) + # The legacy completions spec declares ``minimum: 0`` for max_tokens, + # so 0 is a valid (if degenerate) cap and only negatives are rejected. + # The chat fields never reach here with 0 (pydantic enforces ge=1). + if value < 0: + raise HTTPException( + status_code = 400, + detail = openai_error_body( + f"'{param}' must be at least 0.", + status = 400, + code = "invalid_value", + param = param, + ), + ) + return value + + max_tokens = _validate_explicit(max_tokens, "max_tokens") + max_completion_tokens = _validate_explicit(max_completion_tokens, "max_completion_tokens") + return max_completion_tokens if max_completion_tokens is not None else max_tokens + + +def _effective_openai_max_tokens(payload): + return _effective_openai_max_tokens_from_values( + getattr(payload, "max_tokens", None), + getattr(payload, "max_completion_tokens", None), + ) + + def _wants_multiple_choices(payload) -> bool: return (payload.n or 1) > 1 +def _has_openai_tool_history(messages) -> bool: + for message in messages or []: + if isinstance(message, dict): + if message.get("role") == "tool" or message.get("tool_calls"): + return True + continue + if getattr(message, "role", None) == "tool" or getattr(message, "tool_calls", None): + return True + return False + + def _raise_unsupported_openai_parameter(param: str, message: str) -> None: raise HTTPException( status_code = 400, @@ -294,6 +418,14 @@ def _openai_stream_error_chunk(exc) -> dict: return openai_error_body(_friendly_error(exc), status = 500) +def _openai_stream_error_sse(error: dict) -> str: + return f"data: {json.dumps(error)}\n\ndata: [DONE]\n\n" + + +def _openai_stream_error_sse_bytes(error: dict) -> bytes: + return _openai_stream_error_sse(error).encode("utf-8") + + def _openai_passthrough_error(status_code, text) -> "HTTPException": """HTTPException for a non-200 upstream response on the OpenAI passthrough (tools / response_format). An over-context upstream error is mapped to a 400 @@ -508,20 +640,62 @@ def _drop_parallel_tool_call_deltas(chunk) -> bool: return changed -def _cap_parallel_tool_calls_sse_line(raw_line: str) -> str: - """Drop tool_call deltas whose index >= 1 from one streamed OpenAI SSE - ``data:`` line so only the first tool call survives (parallel_tool_calls=false, - best-effort). Non-tool / unparseable payloads are returned byte-for-byte.""" - payload = raw_line[len("data: ") :] +def _add_empty_content_to_reasoning_deltas(chunk: dict) -> bool: + """Make reasoning-only deltas palatable to strict OpenAI adapters. + + Some clients built on OpenAI-compatible streams ignore or reject chunks whose + delta only contains non-standard ``reasoning_content``. Preserve that field, + but add an empty standard ``content`` member so the chunk is still a valid + text-delta shape and downstream parsers keep the stream alive. + """ + changed = False + choices = chunk.get("choices") + if not isinstance(choices, list): + return False + for choice in choices: + if not isinstance(choice, dict): + continue + delta = choice.get("delta") + if not isinstance(delta, dict): + continue + if "reasoning_content" in delta and "content" not in delta: + delta["content"] = "" + changed = True + return changed + + +def _normalize_openai_passthrough_sse_line( + raw_line: str, *, cap_parallel_tool_calls: bool = False +) -> str: + """Normalize one passthrough OpenAI SSE ``data:`` line before relaying. + + The function is intentionally narrow: it leaves comments, blank events, + ``[DONE]``, and unparseable upstream bytes untouched; parsed chunks are + re-serialized only when a compatibility mutation is actually required. + """ + if not raw_line.startswith("data:"): + return raw_line + # Both mutations key off JSON object keys, so a line without either quoted + # key can never change; skip the parse on the per-token common case. + if '"reasoning_content"' not in raw_line and not ( + cap_parallel_tool_calls and '"tool_calls"' in raw_line + ): + return raw_line + payload = raw_line[len("data:") :].lstrip() if payload.strip() in ("", "[DONE]"): return raw_line try: obj = json.loads(payload) except Exception: return raw_line - if not _drop_parallel_tool_call_deltas(obj): + if not isinstance(obj, dict): return raw_line - return "data: " + json.dumps(obj, separators = (",", ":")) + changed = _add_empty_content_to_reasoning_deltas(obj) + if cap_parallel_tool_calls and _drop_parallel_tool_call_deltas(obj): + changed = True + if not changed: + return raw_line + return "data: " + json.dumps(obj, separators = (",", ":"), ensure_ascii = False) def _prompt_tokens_details(upstream): @@ -538,6 +712,49 @@ def _wants_stream_usage(payload) -> bool: return bool((payload.stream_options or {}).get("include_usage")) +_OPENAI_PASSTHROUGH_TERMINAL_GRACE_S = 2.0 +_SSE_DONE_LINE = "data: [DONE]" + + +def _openai_passthrough_sse_line_terminal_state(raw_line: str) -> Optional[str]: + """Classify OpenAI-compatible chat stream terminal markers. + + Some llama-server builds can emit the logical final chunk (``finish_reason``) + and optional usage chunk, then keep the HTTP stream open without sending the + OpenAI ``data: [DONE]`` sentinel. Classifying those chunks lets Studio close + the client stream promptly while preserving an optional trailing usage chunk. + """ + if not raw_line.startswith("data:"): + return None + data_str = raw_line[5:].lstrip() + if data_str == "[DONE]": + return "done" + try: + data = json.loads(data_str) + except json.JSONDecodeError: + return None + return _openai_passthrough_terminal_state_from_data(data) + + +def _openai_passthrough_terminal_state_from_data(data) -> Optional[str]: + """Dict-level core of ``_openai_passthrough_sse_line_terminal_state`` for + callers that already parsed the chunk (avoids a re-parse per relayed line).""" + if not isinstance(data, dict): + return None + if _monitor_openai_error_message(data): + return "error" + choices = data.get("choices") + if isinstance(choices, list): + if not choices and isinstance(data.get("usage"), dict): + return "usage" + for choice in choices: + if isinstance(choice, dict) and choice.get("finish_reason") is not None: + return "finish" + elif isinstance(data.get("usage"), dict): + return "usage" + return None + + def _openai_stream_usage_chunk( payload, completion_id, created, model_name, stream_usage, stream_timings ): @@ -603,6 +820,21 @@ def _chat_content_chunk(completion_id, created, model_name, text) -> str: ) +def _chat_reasoning_chunk(completion_id, created, model_name, text) -> str: + """Like ``_chat_content_chunk`` but on ``reasoning_content`` (renders the UI thinking block). + + Carries ``content: ""`` alongside, like the GGUF and passthrough paths, so + strict OpenAI adapters don't drop the reasoning-only delta. + """ + return _chat_chunk_sse( + completion_id, + created, + model_name, + delta = ChoiceDelta(content = "", reasoning_content = text), + finish_reason = None, + ) + + def _chat_final_chunk(completion_id, created, model_name, finish_reason) -> str: """Terminal stop chunk (empty delta) carrying the finish reason.""" return _chat_chunk_sse( @@ -614,6 +846,66 @@ def _chat_final_chunk(completion_id, created, model_name, finish_reason) -> str: ) +def _chat_tool_calls_chunk(completion_id, created, model_name, tool_calls) -> str: + """Delta chunk carrying OpenAI tool-call deltas (sibling of ``_chat_content_chunk``).""" + return _chat_chunk_sse( + completion_id, + created, + model_name, + delta = ChoiceDelta(tool_calls = tool_calls), + finish_reason = None, + ) + + +def _sf_heal_events_to_sse( + events, + completion_id, + created, + model_name, + state, + parallel_tool_calls, + monitor_id = None, +): + """Serialize ``StreamToolCallHealer`` events into chat SSE lines. + + ``state["idx"]`` tracks the call index across ``feed``/``finalize``; + ``parallel_tool_calls is False`` caps promotion to one call (GGUF parity). + The monitor is fed from the same events the client receives, never the + healed-away markup.""" + lines = [] + for kind, value in events: + if kind == "text": + if value: + lines.append(_chat_content_chunk(completion_id, created, model_name, value)) + api_monitor.append_reply(monitor_id, value) + continue + if parallel_tool_calls is False and state["idx"] >= 1: + continue + lines.append( + _chat_tool_calls_chunk( + completion_id, + created, + model_name, + [ + { + "index": state["idx"], + "id": value["id"], + "type": "function", + "function": value["function"], + } + ], + ) + ) + _fn = value.get("function") or {} + api_monitor.append_reply( + monitor_id, + ("[tool_calls] " if state["idx"] == 0 else "; ") + + f"{_fn.get('name', '')}({_fn.get('arguments', '')})", + ) + state["idx"] += 1 + return lines + + def _rewrite_cmpl_id(raw: bytes) -> bytes: """Rewrite llama-server's chat-style ``chatcmpl-`` ids to the ``cmpl-`` prefix OpenAI's legacy /v1/completions use. Anchored on the ``"id":`` key @@ -773,8 +1065,11 @@ def _llama_streaming_generation_timeout() -> httpx.Timeout: def _set_stream_response_read_timeout( - response: httpx.Response, read_timeout_s: float = _DEFAULT_STREAM_STALL_TIMEOUT_S + response: httpx.Response, read_timeout_s: Optional[float] = _DEFAULT_STREAM_STALL_TIMEOUT_S ) -> None: + # ``read_timeout_s = None`` clears httpx's read timeout (wait indefinitely), + # used when the stall guard is disabled so a stale first-token deadline + # can't keep timing out post-first-chunk gaps. try: timeout_ext = response.request.extensions.get("timeout") if isinstance(timeout_ext, dict): @@ -785,6 +1080,295 @@ def _set_stream_response_read_timeout( _STREAM_DISCONNECT_POLL_TIMEOUT_S = 0.25 _OPENAI_PASSTHROUGH_PREHEADER_STATUS_WINDOW_S = 0.1 +_OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S = 5.0 +_OPENAI_PASSTHROUGH_SSE_KEEPALIVE = ": keep-alive\n\n" +_OPENAI_LLAMA_ADMISSION_POLL_S = 0.25 + + +def _openai_llama_admission_capacity(request: Optional[Request], llama_backend = None) -> int: + """Serving slots available for one local llama-server backend. + + The loaded backend is the source of truth because it may have reduced + ``--parallel`` at load time to keep the model on GPU. The app state is a + launch-intent fallback for tests and for the short window before a backend + reports its committed runtime slots. + """ + slots = _positive_int_or_none(getattr(llama_backend, "effective_parallel_slots", None)) + if slots is not None: + return slots + try: + slots = getattr(request.app.state, "llama_parallel_slots", None) + except Exception: + slots = None + return _positive_int_or_none(slots) or 1 + + +def _openai_llama_admission_reserve( + *, request: Optional[Request], llama_backend +) -> tuple[LlamaAdmissionReservation, LlamaAdmissionConfig]: + config = llama_admission_config_from_env() + capacity = _openai_llama_admission_capacity(request, llama_backend) + key = str(getattr(llama_backend, "base_url", "llama-server")) + reservation = get_llama_admission_queue(key).reserve( + capacity = capacity, + config = config, + ) + return reservation, config + + +def _openai_admission_request_path(request: Optional[Request]) -> Optional[str]: + try: + return str(request.url.path) if request is not None else None + except Exception: + return None + + +def _openai_admission_log( + event: str, + reservation: Optional[LlamaAdmissionReservation] = None, + *, + snapshot = None, + request: Optional[Request], + mode: str, + wait_started_at: Optional[float] = None, + completion_id: Optional[str] = None, + level: str = "debug", +) -> None: + if snapshot is None and reservation is not None: + snapshot = reservation.snapshot_now() + wait_ms = None + if wait_started_at is not None: + wait_ms = int(max(0.0, time.monotonic() - wait_started_at) * 1000) + log = getattr(logger, level, logger.debug) + log( + "openai admission %s: mode=%s path=%s completion_id=%s capacity=%s active=%s queued=%s wait_ms=%s", + event, + mode, + _openai_admission_request_path(request), + completion_id, + getattr(snapshot, "capacity", None), + getattr(snapshot, "active", None), + getattr(snapshot, "queued", None), + wait_ms, + ) + + +def _openai_admission_error_body(exc: Exception, *, status_code: int) -> dict: + snapshot = getattr(exc, "snapshot", None) + message = str(exc) + if snapshot is not None: + message = ( + f"{message} " + f"(active={snapshot.active}, queued={snapshot.queued}, capacity={snapshot.capacity})" + ) + return openai_error_body(message, status = status_code) + + +def _openai_admission_http_exception(exc: Exception, *, status_code: int) -> HTTPException: + return HTTPException( + status_code = status_code, + detail = _openai_admission_error_body(exc, status_code = status_code), + ) + + +def _openai_admission_timeout_error( + reservation: LlamaAdmissionReservation, +) -> LlamaAdmissionTimeout: + return LlamaAdmissionTimeout( + "Timed out waiting for an available local llama-server generation slot", + snapshot = reservation.snapshot_now(), + ) + + +def _openai_admission_cancelled_error( + reservation: LlamaAdmissionReservation, +) -> LlamaAdmissionCancelled: + return LlamaAdmissionCancelled( + "Client disconnected before an upstream llama-server generation slot was available", + snapshot = reservation.snapshot_now(), + ) + + +async def _raise_if_openai_admission_cancelled( + reservation: LlamaAdmissionReservation, *, request: Optional[Request], cancel_event +) -> None: + if reservation.is_cancelled: + raise _openai_admission_cancelled_error(reservation) + if await _preheader_cancelled(cancel_event, request): + reservation.cancel() + raise _openai_admission_cancelled_error(reservation) + + +async def _wait_for_openai_admission_non_streaming( + reservation: LlamaAdmissionReservation, + config: LlamaAdmissionConfig, + *, + request: Optional[Request], + cancel_event, +) -> LlamaAdmissionLease: + lease = reservation.lease_nowait() + if lease is not None: + try: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + lease.release() + raise + except LlamaAdmissionCancelled: + lease.release() + raise + return lease + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + deadline = None if config.queue_timeout_s is None else time.monotonic() + config.queue_timeout_s + try: + while True: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + lease = reservation.lease_nowait() + if lease is not None: + try: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + lease.release() + raise + except LlamaAdmissionCancelled: + lease.release() + raise + return lease + wait_s = _OPENAI_LLAMA_ADMISSION_POLL_S + if deadline is not None: + remaining_s = deadline - time.monotonic() + if remaining_s <= 0: + reservation.cancel() + raise _openai_admission_timeout_error(reservation) + wait_s = min(wait_s, max(remaining_s, 0.001)) + try: + lease = await reservation.wait(wait_s) + except asyncio.TimeoutError: + continue + if lease is not None: + return lease + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + reservation.cancel() + raise + + +async def _openai_admission_wait_stream_chunks( + reservation: LlamaAdmissionReservation, + config: LlamaAdmissionConfig, + *, + request: Optional[Request], + cancel_event, +): + lease = reservation.lease_nowait() + if lease is not None: + yield lease + return + + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + deadline = None if config.queue_timeout_s is None else time.monotonic() + config.queue_timeout_s + keepalive_interval_s = max(0.001, config.keepalive_interval_s) + next_keepalive_at = time.monotonic() + keepalive_interval_s + try: + while True: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + lease = reservation.lease_nowait() + if lease is not None: + yield lease + return + + now = time.monotonic() + wait_s = min(_OPENAI_LLAMA_ADMISSION_POLL_S, max(next_keepalive_at - now, 0.001)) + if deadline is not None: + remaining_s = deadline - now + if remaining_s <= 0: + reservation.cancel() + raise _openai_admission_timeout_error(reservation) + wait_s = min(wait_s, max(remaining_s, 0.001)) + try: + lease = await reservation.wait(wait_s) + except asyncio.TimeoutError: + lease = None + if lease is not None: + yield lease + return + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + now = time.monotonic() + if now >= next_keepalive_at: + next_keepalive_at = now + keepalive_interval_s + yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE + except asyncio.CancelledError: + reservation.cancel() + raise + + +async def _close_openai_admitted_stream_iterator(iterator, *, cancelled: bool) -> None: + if iterator is None: + return + if cancelled: + athrow = getattr(iterator, "athrow", None) + if athrow is not None: + try: + await athrow(asyncio.CancelledError()) + except (asyncio.CancelledError, StopAsyncIteration, RuntimeError): + return + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + + +def _openai_compat_stream_stall_timeout(): + """Max silent gap after an OpenAI passthrough stream has produced data. + + If the socket goes silent after valid SSE data, this bounds how long the + client is kept open. Defaults to the backend-wide stall timeout so this + path stalls out like every sibling stream; set the env var to tighten it + for local serving, or to 0 to disable the guard. + """ + return _positive_float_env( + _OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, + _DEFAULT_STREAM_STALL_TIMEOUT_S, + ) + + +def _openai_passthrough_upstream_headers(*, llama_backend = None) -> dict: + headers = {} + auth_headers = getattr(llama_backend, "_auth_headers", None) + if isinstance(auth_headers, dict): + headers.update(auth_headers) + headers["Connection"] = "close" + return headers class _CompatSameTaskTimeout: @@ -841,17 +1425,14 @@ class _SameTaskStreamingResponse(StreamingResponse): **kwargs, ) -> None: super().__init__(*args, **kwargs) - # Async callable invoked when the client disconnects before the body - # iterator is ever advanced. A generator that never started cannot run - # its own try/finally, so a stream that acquires resources before its - # first yield (the passthrough opens an upstream httpx stream eagerly) - # passes this to release them. + # Released when the client disconnects before the body iterator starts: + # its try/finally never runs, so a stream that opens resources before the + # first yield (the passthrough's upstream httpx stream) passes this. self._unstarted_cleanup = unstarted_cleanup async def __call__(self, scope, receive, send) -> None: - # Track whether the body iterator was ever advanced: send() only emits a - # body message after the generator yields its first chunk, so a failure - # before then means it never entered its try/finally. + # send() emits a body message only after the first chunk, so no body + # message means the generator never entered its try/finally. body_started = False async def _tracking_send(message) -> None: @@ -862,15 +1443,11 @@ class _SameTaskStreamingResponse(StreamingResponse): try: await self.stream_response(_tracking_send) - except OSError: - # Client disconnected mid-send. + except OSError: # client disconnected mid-send if body_started: - # The generator produced at least one chunk and is suspended in - # its try/finally. Throw CancelledError into it (not aclose's - # GeneratorExit) so its `except asyncio.CancelledError` handler - # runs and finishes any api_monitor entry; GeneratorExit would - # skip it and only run `finally`. Fall back to aclose() without - # athrow. + # Generator is suspended in its try/finally: throw CancelledError + # (not aclose's GeneratorExit) so its handler finishes the + # api_monitor entry. Fall back to aclose() without athrow. athrow = getattr(self.body_iterator, "athrow", None) if athrow is not None: try: @@ -882,16 +1459,16 @@ class _SameTaskStreamingResponse(StreamingResponse): if aclose is not None: await aclose() else: - # http.response.start failed before the body iterator advanced, - # so its try/finally never armed and aclose()/athrow() are no-ops - # on an unstarted generator. Release any resources acquired - # before the first yield via the explicit cleanup hook. + # Generator never started; aclose()/athrow() are no-ops on it, so + # release eager resources via the hook. getattr guards a response + # built through __new__ without __init__ (tests, pickling). aclose = getattr(self.body_iterator, "aclose", None) if aclose is not None: await aclose() - if self._unstarted_cleanup is not None: + cleanup = getattr(self, "_unstarted_cleanup", None) + if cleanup is not None: try: - await self._unstarted_cleanup() + await cleanup() except Exception: pass raise ClientDisconnect() @@ -899,6 +1476,16 @@ class _SameTaskStreamingResponse(StreamingResponse): await self.background() +def _tracked_cancel_unstarted_cleanup(tracker): + """unstarted_cleanup that exits ``tracker`` on a pre-start disconnect, when + the generator's finally (which normally exits it) never runs.""" + + async def _cleanup() -> None: + tracker.__exit__(None, None, None) + + return _cleanup + + async def _aclose_stream_resources( *, watchers = (), @@ -909,7 +1496,8 @@ async def _aclose_stream_resources( """Tear down an httpx streaming generator's resources in the required order: cancel + await each watcher task, then aclose() the byte/line iterator, the response, and the client. Each step swallows its own exceptions so teardown - always completes. See _anthropic_passthrough_stream for the ordering rationale.""" + always completes; a close-time CancelledError is re-raised only after every + step has run. See _anthropic_passthrough_stream for the ordering rationale.""" for watcher in watchers: if watcher is not None: watcher.cancel() @@ -917,21 +1505,30 @@ async def _aclose_stream_resources( await watcher except (asyncio.CancelledError, Exception): pass + close_cancelled = False if iterator is not None: try: await iterator.aclose() + except asyncio.CancelledError: + close_cancelled = True except Exception: pass if resp is not None: try: await resp.aclose() + except asyncio.CancelledError: + close_cancelled = True except Exception: pass if client is not None: try: await client.aclose() + except asyncio.CancelledError: + close_cancelled = True except Exception: pass + if close_cancelled: + raise asyncio.CancelledError() async def _preheader_cancelled(cancel_event = None, request: Optional[Request] = None) -> bool: @@ -954,6 +1551,7 @@ async def _send_stream_with_preheader_cancel( req: httpx.Request, cancel_event = None, request: Optional[Request] = None, + mark_cancel_on_cancel: bool = True, ) -> Optional[httpx.Response]: if cancel_event is None and request is None: return await client.send(req, stream = True) @@ -985,7 +1583,7 @@ async def _send_stream_with_preheader_cancel( await _stop_send_task() return None except asyncio.CancelledError: - if cancel_event is not None: + if mark_cancel_on_cancel and cancel_event is not None: cancel_event.set() await _stop_send_task() raise @@ -1004,11 +1602,19 @@ async def _aiter_llama_stream_items( request: Optional[Request] = None, first_token_deadline: Optional[float] = None, response: Optional[httpx.Response] = None, - post_first_item_read_timeout_s: Optional[float] = _DEFAULT_STREAM_STALL_TIMEOUT_S, + post_first_item_read_timeout_s: Optional[ + Union[float, Callable[[], Optional[float]]] + ] = _DEFAULT_STREAM_STALL_TIMEOUT_S, ): if first_token_deadline is None: first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S last_item_at: Optional[float] = None + + def _post_first_timeout_s() -> Optional[float]: + if callable(post_first_item_read_timeout_s): + return post_first_item_read_timeout_s() + return post_first_item_read_timeout_s + while True: if cancel_event is not None and cancel_event.is_set(): return @@ -1029,15 +1635,14 @@ async def _aiter_llama_stream_items( async with _same_task_timeout(remaining_s): item = await async_iter.__anext__() else: + timeout_s = _post_first_timeout_s() if ( request is not None and response is not None - and post_first_item_read_timeout_s is not None + and timeout_s is not None and last_item_at is not None ): - stall_remaining_s = post_first_item_read_timeout_s - ( - time.monotonic() - last_item_at - ) + stall_remaining_s = timeout_s - (time.monotonic() - last_item_at) if stall_remaining_s <= 0: raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") _set_stream_response_read_timeout(response, stall_remaining_s) @@ -1054,19 +1659,16 @@ async def _aiter_llama_stream_items( if now >= first_token_deadline: raise continue - if ( - request is not None - and post_first_item_read_timeout_s is not None - and now - last_item_at < post_first_item_read_timeout_s - ): + timeout_s = _post_first_timeout_s() + if request is not None and timeout_s is not None and now - last_item_at < timeout_s: continue raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") - if ( - last_item_at is None - and response is not None - and post_first_item_read_timeout_s is not None - ): - _set_stream_response_read_timeout(response, post_first_item_read_timeout_s) + if last_item_at is None and response is not None: + # The first-token read deadline no longer applies once a chunk has + # arrived: switch to the stall timeout, or clear the read timeout + # entirely when the stall guard is disabled (callable returns None) + # so a long gap can't trip the stale first-token deadline. + _set_stream_response_read_timeout(response, _post_first_timeout_s()) last_item_at = time.monotonic() yield item @@ -1136,6 +1738,13 @@ from core.inference.key_exchange import decrypt_api_key from core.inference.model_ids import public_model_id from core.inference.api_monitor import api_monitor from core.inference.llama_http import nonstreaming_client +from core.inference.tool_call_parser import ( + _strip_function_xml_calls, + _strip_gemma_wrapperless_calls, + _strip_glm_calls, + _strip_mistral_closed_calls, +) +from core.inference.tool_call_parser import TOOL_XML_SIGNALS as _PARSER_TOOL_SIGNALS from core.inference.passthrough_healing import ( StreamToolCallHealer, heal_gate, @@ -1294,6 +1903,11 @@ async def artifact_preview_frame(allow_network: bool = False): ) +# Whitespace/escape-tolerant bare-JSON tool-template detector (matches pretty-printed and +# JSON-escaped ``{"name":`` plus the ``"function"`` alias), mirroring the parser's tolerance. +_BARE_JSON_NAME_MARKER_RE = _re.compile(r'\{\s*\\?"(?:name|function)\\?"\s*:') + + def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: """Classify reasoning/tool capabilities via the GGUF classifier so flags match across backends. gpt-oss is overridden: Harmony routes reasoning and @@ -1304,17 +1918,22 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: model_identifier = model_id, log_source = "safetensors", ) - # Our safetensors loop only parses {json}, - # ..., and Gemma native <|tool_call>.... - # Llama uses <|python_tag|>, Mistral uses [TOOL_CALLS]; advertising tools for - # those enables a pill the parser can't honour. GGUF is unaffected -- - # llama-server normalises every format into structured deltas. + # Markers any supported parser recognises (template advertises tools but + # uses none -> drop the pill). Reuse the parser's own signal list so this + # gate never drifts (a hand-maintained copy lost the DeepSeek variants); + # ```` is GLM's unique signal, absent from the shared set. The + # bare-JSON ``{"name":`` form is matched below with the whitespace/escape- + # tolerant ``_BARE_JSON_NAME_MARKER_RE`` so pretty-printed or escaped + # templates are not mis-classified as tool-less. + _PARSER_MARKERS = ( + *_PARSER_TOOL_SIGNALS, + "", + ) if ( flags.get("supports_tools") and chat_template - and "" not in chat_template - and "" not in chat_template + and not any(m in chat_template for m in _PARSER_MARKERS) + and not _BARE_JSON_NAME_MARKER_RE.search(chat_template) ): logger.info( "safetensors: template advertises tools but uses an " @@ -1335,6 +1954,87 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: return flags +def _generation_prompt_opens_think(template: Optional[str]) -> bool: + """True when rendering the template's generation prompt ends INSIDE an unclosed ````. + + Distinguishes templates that PREFILL an open ```` in the assistant generation + prompt (DeepSeek-R1, QwQ, Qwen3-Thinking) -- where the model emits only the closing + ```` and the extractor must start in reasoning mode -- from templates that merely + render PAST assistant ``...`` history while leaving the generation prompt + open with no ```` (e.g. Kimi-K2-Thinking), where the model self-emits its own block + and the extractor must start in normal mode. Renders a single-user-message probe with the + same sandbox transformers uses; on any failure returns True, preserving the historical + always-on prefill for templates that cannot be rendered here. + """ + if not template: + return False + try: + from jinja2.sandbox import ImmutableSandboxedEnvironment + + def _raise_exception(message: str): + raise RuntimeError(message) + + env = ImmutableSandboxedEnvironment( + trim_blocks = True, + lstrip_blocks = True, + extensions = ["jinja2.ext.loopcontrols"], + ) + env.filters["tojson"] = lambda value, **kwargs: json.dumps(value, ensure_ascii = False) + env.globals["raise_exception"] = _raise_exception + rendered = env.from_string(template).render( + messages = [{"role": "user", "content": "hi"}], + add_generation_prompt = True, + bos_token = "", + eos_token = "", + ) + except Exception: + return True + # ```` is not a substring of ```` (the ``/`` breaks it), so the last open + # tag sitting after the last close tag means the prompt ends inside an open block. + return rendered.rfind("") > rendered.rfind("") + + +def _sf_reasoning_prefill_mode( + features: dict, + enable_thinking: Optional[bool], + template: Optional[str] = None, + reasoning_effort: Optional[str] = None, +) -> bool: + """Whether a safetensors/MLX generation begins INSIDE an unclosed ````. + + ``enable_thinking`` templates (Qwen3/GLM) prefill an open ```` so the model + emits only the closing ````, and the extractor must start in reasoning mode. + Gated on the STANDARD ````/```` markers: bespoke channels (gemma's + ``<|think|>``) never emit ```` and would swallow the answer, so they and + gpt-oss and thinking-disabled requests return False. ``enable_thinking`` None + defaults thinking ON, so a plain request still prefills. + """ + if features.get("reasoning_style") not in ("enable_thinking", "enable_thinking_effort"): + return False + tpl = template or "" + if "" not in tpl and "" not in tpl: + return False + if features.get("reasoning_always_on"): + # enable_thinking_effort + always-on: the effort mechanism (not the prompt shape) keeps + # thinking on, so always-on wins over reasoning_effort and we prefill. + if features.get("reasoning_style") == "enable_thinking_effort": + return True + # ``reasoning_always_on`` fires on paired ``...`` anywhere in the + # template, including markup that only renders PAST assistant history (Kimi-K2-Thinking) + # while the generation prompt opens none. Prefill only when the generation prompt opens + # one, else the extractor captures a normal answer as reasoning_content and returns blank. + return _generation_prompt_opens_think(tpl) + if not features.get("supports_reasoning"): + return False + if enable_thinking is False: + return False + # Thinking-off arrives as reasoning_effort "none" on enable_thinking_effort models; honor it + # so we don't prefill and capture the answer. Plain enable_thinking models ignore effort. + if features.get("reasoning_style") == "enable_thinking_effort" and reasoning_effort == "none": + return False + return True + + def _effective_enable_tools(payload) -> Optional[bool]: """Resolve `payload.enable_tools` against the process-level tool policy. @@ -1347,6 +2047,20 @@ def _effective_enable_tools(payload) -> Optional[bool]: return policy if policy is not None else payload.enable_tools +def _explicit_studio_tool_loop_requested(payload) -> bool: + """True when the request itself asks Studio to execute local tools. + + Process-wide CLI policy can default Studio's tool loop on for ordinary chat, + but it must not steal OpenAI-compatible client tools or response_format + requests from the llama-server passthrough path. A policy of ``False`` + (--disable-tools) vetoes even an explicit ``enable_tools: true`` ask. + """ + from state.tool_policy import get_tool_policy + + policy = get_tool_policy() + return policy is not False and (payload.enable_tools is True or bool(payload.mcp_enabled)) + + # Cancel registry. Proxies (e.g. Colab) can swallow client fetch aborts so # is_disconnected() never fires. POST /inference/cancel looks up in-flight # cancel_events here by cancel_id (per-run) or session_id / completion_id @@ -1487,6 +2201,39 @@ async def _await_disconnect_then_cancel(request, cancel_event) -> None: return +def _cancelable_nonstreaming_client() -> httpx.AsyncClient: + return httpx.AsyncClient( + limits = httpx.Limits(max_connections = 1, max_keepalive_connections = 0), + trust_env = False, + ) + + +async def _await_cancel_or_disconnect_then_close_client( + *, cancel_event, request: Optional[Request], client: httpx.AsyncClient +) -> None: + """Close a dedicated non-streaming upstream client on cancel/disconnect. + + The shared ``nonstreaming_client()`` is pooled, so cancelable generation calls + use a per-request client. Closing it interrupts a blocked llama-server + request without affecting unrelated pooled non-streaming calls. + """ + try: + while True: + if cancel_event is not None and cancel_event.is_set(): + break + if request is not None and await request.is_disconnected(): + if cancel_event is not None: + cancel_event.set() + break + await asyncio.sleep(0.1) + try: + await client.aclose() + except Exception: + pass + except asyncio.CancelledError: + return + + async def _stop_local_disconnect_cancel_watcher(watcher) -> None: watcher.cancel() try: @@ -1605,30 +2352,155 @@ def _apply_rag_nudge(nudge: str, tools: list[dict], *, rag_scope) -> str: return nudge + " " + _RAG_GROUNDING_NUDGE -# Strip tool-call XML the speculative buffer in core/inference/llama_cpp.py -# split across the visible/DRAIN boundary. Four leak shapes: +# Strip leaked tool-call markup: every shared-parser format plus the leak shapes +# llama_cpp.py's speculative buffer splits across the visible/DRAIN boundary: # 1. well-formed `...` / `...` # 2. orphan opening to EOF (close was DRAINED) # 3. bare orphan close (open was DRAINED) # 4. tail-only `` (outer close truncated by EOS); anchored to # `\Z` so mid-text `` in user code samples survives. +# 5. Mistral `[TOOL_CALLS]name{json}` / rehearsal `name[ARGS]{json}`: the balanced +# scan removes the whole call (a non-greedy regex would truncate nested JSON). +# DeepSeek/GLM/Kimi envelopes are covered by the parser's own arms/scans, so a signal +# we parse is never left un-stripped; the DeepSeek opener alternation is the parser's own. +from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC + _TOOL_XML_RE = _re.compile( - # Hyphen in the name char-class matches MCP tool names with dashes - # (mcp__srv__list-issues) that would otherwise leak past this strip. - r"<(?:tool_call|function=[\w-]+)>.*?(?:|\Z)" + # Arm order/notes: the closed ```` arm runs first and extends + # to the call's REAL close so a literal ```` in a value does not + # leak the tail; the combined arm still catches ```` and orphan + # tails. The python_tag arm bounds only on REAL Llama control sentinels + # (stopping at any ``<|`` truncated on literal ``<|x|>`` tokens in values). + # The last arms cover DeepSeek envelopes (all opener variants), Kimi section + # blocks, and bare Kimi calls. Name class ``[\w.\-]`` mirrors the parser. + # Those three arms carry a call-shaped lookahead (matching the parser's + # ``_TOOL_ALL_PATS``): a prose answer that merely mentions a marker + # (``See <|tool_call_begin|> in the docs``) is only stripped when a real + # call actually follows the marker, or the marker is a bare fragment at EOF. + r'(?:(?!).)*' + r'|<(?:tool_call|function(?:=[\w.\-]+|\s+name="[\w.\-]+"))>.*?(?:|\Z)' r"|<\|tool_call>.*?(?:|\Z)" r"|" r"|" - r"|\s*\Z", + r"|<\|python_tag\|>(?:[^<]|<(?!\|(?:eot_id|eom_id|python_tag|start_header_id|end_header_id|begin_of_text|finetune_right_pad_id)\|))*" + r"|\[/TOOL_CALLS\]" + # Truncated canonical array (closing ``]`` lost to EOS): the balanced scan cannot remove + # it, so strip its tail here. + r"|\[TOOL_CALLS\]\s*\[.*\Z" + # Named / v11 forms and bare rehearsal; arms aligned with the parser regexes. + r"|\[TOOL_CALLS\]\s*[\w-]+(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*(?:\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}|.*?\Z)" + # Rehearsal: balanced/truncated body or bare marker at EOS only (prose ``foo[ARGS]`` + # survives); NAME captured as ``reh`` for the inactive-name display gate. + r"|(?[\w-]+)\[ARGS\]\s*(?:\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}|\{.*\Z|\Z)" + # DeepSeek envelopes (all opener variants), Kimi section blocks, and bare Kimi calls; + # each arm carries a call-shaped lookahead so prose merely mentioning a marker survives. + r"|" + + _DS_OPEN_SRC + + r"(?=\s*(?:<｜tool▁call▁begin｜>|function)|\s*$).*?(?:<｜tool▁calls▁end｜>|\Z)" + r"|<\|tool_calls_section_begin\|>(?=\s*<\|tool_call_begin\|>|\s*$).*?(?:<\|tool_calls_section_end\|>|\Z)" + r"|<\|tool_call_begin\|>(?=\s*[A-Za-z_][\w.\-]*:\d|\s*$).*?(?:<\|tool_call_end\|>|\Z)" + # ```` is the attribute-form alias of ```` (the parser accepts + # both); strip a tail-only orphan close of either spelling. + r"|\s*\Z", + _re.DOTALL, +) + +# Closed-only variant for segments before the last think block: the ``\Z``-anchored arms +# would treat a segment boundary as EOS and strip prose ``foo[ARGS]``. +_TOOL_XML_CLOSED_RE = _re.compile( + r"<(?:tool_call|function=[\w-]+)>.*?" + r"|<\|tool_call>.*?" + r"|" + r"|" + r"|\[/TOOL_CALLS\]", _re.DOTALL, ) -def _strip_tool_xml_for_display(text: str, *, auto_heal_tool_calls: bool) -> str: - """Apply route-level XML leak cleanup only when Auto-Heal is enabled.""" +def _gemma_strip_gate(tools) -> set: + """Enabled tool NAMES gating the wrapper-less Gemma strip (mirrors the + parser/loop gate: only an enabled ``call:foo{...}`` is a call). With NO tools + enabled this returns an EMPTY set, not ``None``: every ``call:NAME{...}`` is + then prose, and ``None`` would strip-all and delete a legitimate answer.""" + names = { + (t.get("function") or {}).get("name") + for t in (tools or []) + if isinstance(t, dict) and isinstance(t.get("function"), dict) + } + names.discard(None) + return names + + +def _display_tool_name_gate(active_tools): + """Active tool NAMES for gating the rehearsal display strip, or None when no tools + are enabled. ``None`` keeps the legacy strip-all behavior, mirroring the loop gate: + a bare ``NAME[ARGS]`` is a call only when NAME is active; without a tool list every + identifier stays ambiguous, so strip.""" + names = { + (t.get("function") or {}).get("name") + for t in (active_tools or []) + if isinstance(t, dict) and isinstance(t.get("function"), dict) + } + names.discard(None) + return names or None + + +def _strip_tool_xml_for_display( + text: str, + *, + auto_heal_tool_calls: bool, + enabled_tool_names: Optional[set] = None, +) -> str: + """Apply route-level XML leak cleanup only when Auto-Heal is enabled. + + Mirrors the parser-side segment scan: balanced strips first (Mistral, gated Gemma + wrapper-less, GLM real-close, guarded function-XML close at each call's REAL terminator + so literal markup inside a value is data), then the ``_TOOL_XML_RE`` arms cover the + DeepSeek / Kimi / orphan forms. ```` blocks are preserved verbatim and the + ``\\Z``-anchored tail arms run only on the last segment (prose ``foo[ARGS]`` before a + block survives). ``enabled_tool_names`` (when not None) gates the ambiguous bare-rehearsal + ``NAME[ARGS]{...}`` and wrapper-less Gemma ``call:NAME{...}`` strips on the active tool + list; an inactive NAME is prose and is kept. The ``[TOOL_CALLS]`` control-token arms strip + unconditionally regardless of NAME.""" if not auto_heal_tool_calls: return text - return _TOOL_XML_RE.sub("", text) + from core.tool_healing import _strip_bracket_tag_calls, strip_outside_think + + def _keep_inactive_rehearsal(m) -> str: + # Only the bare-rehearsal arm captures ``reh``; with a tool list an inactive + # NAME[ARGS]{...} is prose -- keep it. + if enabled_tool_names is not None: + name = m.groupdict().get("reh") + if name is not None and name not in enabled_tool_names: + return m.group(0) + return "" + + def _strip_segment(seg: str, is_last: bool) -> str: + # Scan strips close at each call's REAL terminator (a literal ```` or a + # nested marker quoted inside a value cannot truncate the strip); the regex arms below + # cover the attribute form and the DeepSeek / Kimi / orphan families. + seg = _strip_mistral_closed_calls(seg) + seg = _strip_bracket_tag_calls(seg, enabled_tool_names = enabled_tool_names) + if is_last: + seg = _strip_gemma_wrapperless_calls(seg, enabled_tool_names) + seg = _strip_glm_calls(seg, final = is_last) + seg = _strip_function_xml_calls(seg, final = is_last) + if is_last: + return _TOOL_XML_RE.sub(_keep_inactive_rehearsal, seg) + return _TOOL_XML_CLOSED_RE.sub("", seg) + + return strip_outside_think(text, _strip_segment) + + +def _strip_tool_xml(text: str, enabled_tool_names: Optional[set] = None) -> str: + # Mistral balanced-brace pre-strip (kept explicit so the regression guards see it), then + # the shared think-aware display strip -- the one raw _TOOL_XML_RE.sub lives inside + # _strip_tool_xml_for_display, so every route cleanup site shares it. ``enabled_tool_names`` + # gates the Gemma wrapper-less strip; ``None`` strips every closed call. + text = _strip_mistral_closed_calls(text) + return _strip_tool_xml_for_display( + text, auto_heal_tool_calls = True, enabled_tool_names = enabled_tool_names + ) logger = get_logger(__name__) @@ -2472,6 +3344,21 @@ def _automatic_model_load_may_run() -> bool: return get_openai_auto_switch_enabled() or get_auto_unload_idle_seconds() > 0 +def _no_model_loaded_detail(base: str) -> str: + """Append a pointer to the opt-in auto-switch toggle to a "no model loaded" + error, but only when it's off. Auto-switch (default off) cold-loads a + requested downloaded GGUF, so an off toggle is the usual reason a request + naming a listed model still 400/503s; surface the fix. With it on the name + simply didn't resolve to a local GGUF, so the hint would mislead and is omitted.""" + from utils.openai_auto_switch_settings import get_openai_auto_switch_enabled + + if get_openai_auto_switch_enabled(): + return base + return base + ( + " Or enable Model auto-switch (Settings > API) to load a requested model automatically." + ) + + async def _maybe_auto_switch_model( requested_model: Optional[str], fastapi_request: Request, @@ -2741,10 +3628,14 @@ def _remote_gguf_companion_bytes( info = model_info(repo, token = hf_token, files_metadata = True) total = 0 for sibling in info.siblings or []: - base = Path(sibling.rfilename or "").name.lower() + name = sibling.rfilename or "" + base = Path(name).name.lower() if not base.endswith(".gguf"): continue - if base.startswith("mtp-") or (include_mmproj and "mmproj" in base): + # Root-level mtp- only: -hf auto-fetches the repo-root drafter, not + # the MTP/ subdir copies (which now share the mtp- prefix too). + is_root_mtp = "/" not in name and base.startswith("mtp-") + if is_root_mtp or (include_mmproj and "mmproj" in base): total += getattr(sibling, "size", 0) or 0 return total except Exception as e: @@ -3129,6 +4020,15 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre status_code = 400, detail = "gpu_ids is not supported for GGUF models yet.", ) + if not config.is_gguf and _mlx_distributed_launch_detected(): + raise HTTPException( + status_code = 400, + detail = ( + "Studio does not support distributed MLX inference under " + "mlx.launch. Use `mlx.launch ... unsloth chat` or run Studio " + "without the distributed launcher." + ), + ) # Effective quantization (LoRA can flip 4-bit -> 16-bit); guard + load reuse it. effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit) @@ -3157,12 +4057,15 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre llama_backend = get_llama_cpp_backend() unsloth_backend = get_inference_backend() - # Unload any active Unsloth model to free VRAM + # Unload any active Unsloth model to free VRAM (off the event loop: + # unload takes _gen_lock and can wait on an in-flight stream). if unsloth_backend.active_model_name: logger.info( f"Unloading Unsloth model '{unsloth_backend.active_model_name}' before loading GGUF" ) - unsloth_backend.unload_model(unsloth_backend.active_model_name) + await asyncio.to_thread( + unsloth_backend.unload_model, unsloth_backend.active_model_name + ) # Inherit llama_extra_args from the previous load when the request # omits the field (the chat-settings Apply path doesn't round-trip @@ -3826,28 +4729,76 @@ async def unload_model(request: UnloadRequest, current_subject: str = Depends(ge # 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 + from core.inference.llama_keepwarm import inference_lifecycle_gate, note_model_unloaded try: - # Check if the GGUF backend has this model loaded or is loading it. - llama_backend = get_llama_cpp_backend() - if llama_backend.is_active and ( - llama_backend.model_identifier == request.model_path - or is_registered_native_path_label(llama_backend.model_identifier, request.model_path) - or not llama_backend.is_loaded + # "Stop loading" (frontend cancelLoading -> /unload) must abort a still-loading + # model promptly. /load holds the lifecycle gate for the whole (multi-minute) load, + # so gating first would make the cancel wait it out. cancel_load only tears the + # loading subprocess down (no unload command), so it is safe off-gate. + backend = get_inference_backend() + loading = getattr(backend, "get_loading_model", lambda: None)() + if ( + loading is not None + and hasattr(backend, "cancel_load") + and (request.model_path == loading or request.model_path.lower() == loading.lower()) ): - # 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() + if await asyncio.to_thread(backend.cancel_load, request.model_path): + note_model_unloaded() + logger.info(f"Cancelled in-flight load: {request.model_path}") + return UnloadResponse(status = "unloaded", model = request.model_path) + + # Same "stop loading" fast path for a still-loading GGUF (llama-server spawned, + # health check not yet passed). A gated unload would wait out the multi-minute + # load; unload_model() sets the cancel_event load_model polls off its own lock and + # kills the child, sending no worker command, so it is safe off-gate like + # cancel_load. The gated GGUF branch below handles the already-loaded case. Gate on + # the loading model (identifier or native label): the single llama-server loads one + # GGUF at a time, so an unload for a different model must not cancel this load. + llama_backend = get_llama_cpp_backend() + if ( + llama_backend.is_active + and not llama_backend.is_loaded + and ( + llama_backend.model_identifier == request.model_path + or is_registered_native_path_label( + llama_backend.model_identifier, request.model_path + ) + ) + ): + await asyncio.to_thread(llama_backend.unload_model) note_model_unloaded() - logger.info(f"Unloaded GGUF model: {request.model_path}") + logger.info(f"Cancelled in-flight GGUF load: {request.model_path}") return UnloadResponse(status = "unloaded", model = request.model_path) - # Otherwise, unload from Unsloth backend - backend = get_inference_backend() - backend.unload_model(request.model_path) - note_model_unloaded() - logger.info(f"Unloaded model: {request.model_path}") - return UnloadResponse(status = "unloaded", model = request.model_path) + # Serialize with /load under the same lifecycle gate: the Unsloth unload now runs + # off the event loop (asyncio.to_thread), so without this a concurrent /load could + # swap in a fresh subprocess mid-unload and the unload command would land on the + # new worker. The gate makes load and unload exclusive. + async with inference_lifecycle_gate(): + # Check if the GGUF backend has this model loaded or is loading it. + llama_backend = get_llama_cpp_backend() + if llama_backend.is_active and ( + llama_backend.model_identifier == request.model_path + or is_registered_native_path_label( + llama_backend.model_identifier, request.model_path + ) + or not llama_backend.is_loaded + ): + # A manual unload is a deliberate user action: tear down now even if a + # request is mid-stream (only the automatic idle loop defers to it). + llama_backend.unload_model() + note_model_unloaded() + logger.info(f"Unloaded GGUF model: {request.model_path}") + return UnloadResponse(status = "unloaded", model = request.model_path) + + # Unload from Unsloth backend off the event loop: unload takes _gen_lock, which + # a slow SSE stream paused between tokens still holds, so a sync call would block + # the loop that drives the stream's next token and the lock release. + backend = get_inference_backend() + await asyncio.to_thread(backend.unload_model, request.model_path) + note_model_unloaded() + logger.info(f"Unloaded model: {request.model_path}") + return UnloadResponse(status = "unloaded", model = request.model_path) except Exception as e: logger.error(f"Error unloading model: {e}", exc_info = True) @@ -4005,19 +4956,18 @@ async def generate_stream( temperature = request.temperature, top_p = request.top_p, top_k = request.top_k, + min_p = request.min_p, max_new_tokens = request.max_new_tokens, repetition_penalty = request.repetition_penalty, + presence_penalty = request.presence_penalty, cancel_event = cancel_event, ) _DONE = object() while True: if cancel_event.is_set(): - # The disconnect watcher set cancel_event between chunks. - # Reset the backend here: closing the Python generator does - # not signal a subprocess backend, so without this it keeps - # decoding after the client is gone. The finally's reset is - # guarded on cancel_event being unset, so it will not run - # again for this path. + # Watcher set cancel_event between chunks. Reset here: closing + # the generator does not signal a subprocess backend, so it would + # keep decoding. The finally's reset is guarded, so no double-run. backend.reset_generation_state() break chunk = await asyncio.to_thread(next, gen, _DONE) @@ -5516,7 +6466,7 @@ async def openai_chat_completions( if not backend.active_model_name: raise HTTPException( status_code = 400, - detail = "No model loaded. Call POST /inference/load first.", + detail = _no_model_loaded_detail("No model loaded. Call POST /inference/load first."), ) # Clean public id so the response never echoes a local path; the audio # branch below receives this sanitized label too. @@ -5627,6 +6577,7 @@ async def openai_chat_completions( return _SameTaskStreamingResponse( audio_input_stream(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -5692,11 +6643,11 @@ async def openai_chat_completions( # unaware of `role="tool"` messages and assistant messages that only # carry `tool_calls` (content=None) — both of which are valid in # multi-turn client-side tool loops. - effective_max_tokens = _effective_max_tokens(payload) + effective_max_tokens = _effective_openai_max_tokens(payload) normalized_stop = _normalize_stop_sequences(payload.stop) - _has_tool_messages = any(m.role == "tool" or m.tool_calls for m in payload.messages) + _has_tool_messages = _has_openai_tool_history(payload.messages) # Route guided-decoding requests through the verbatim passthrough so # ``response_format`` (JSON schema) reaches llama-server and the model's # GBNF-constrained output comes back unmodified. The non-passthrough GGUF @@ -5705,12 +6656,41 @@ async def openai_chat_completions( # free-form sampling. Guided decoding does not require ``supports_tools`` -- # the grammar machinery is independent of tool-call parsing. _has_response_format = _extract_response_format(payload) is not None - _tools_passthrough = llama_backend.supports_tools and ( - (payload.tools and len(payload.tools) > 0) or _has_tool_messages + _has_tool_catalog = bool(payload.tools and len(payload.tools) > 0) + _has_active_tool_catalog = _has_tool_catalog and payload.tool_choice != "none" + _has_client_tool_contract = _has_active_tool_catalog or _has_tool_messages + # The Studio tool loop needs a tool-capable backend, so a request that asks + # for it on a backend that can't run it (DiffusionGemma forces supports_tools + # off) must not steal client tools from the passthrough (#6851). + _studio_tool_loop_requested = ( + _explicit_studio_tool_loop_requested(payload) and llama_backend.supports_tools ) + _client_disabled_tool_calls = payload.tool_choice == "none" and not _studio_tool_loop_requested + _supports_tool_passthrough = getattr( + llama_backend, "supports_tool_passthrough", llama_backend.supports_tools + ) + _tools_passthrough = _supports_tool_passthrough and _has_client_tool_contract if ( using_gguf - and not _effective_enable_tools(payload) + and not _studio_tool_loop_requested + and _has_client_tool_contract + and not _supports_tool_passthrough + ): + raise _reject( + 400, + openai_error_body( + ( + "Client-supplied tools or tool-call history require a GGUF chat template " + "with tool-call support; the current model/template does not advertise tools." + ), + status = 400, + code = "unsupported_parameter", + param = "tools" if payload.tools else "messages", + ), + ) + if ( + using_gguf + and not _studio_tool_loop_requested and (_tools_passthrough or _has_response_format) ): if _wants_multiple_choices(payload): @@ -5756,12 +6736,20 @@ async def openai_chat_completions( completion_id, monitor_id = monitor_id, ) - return await _openai_passthrough_non_streaming( - llama_backend, - payload, - model_name, - monitor_id = monitor_id, - ) + _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) + _tracker = _TrackedCancel(cancel_event, *_cancel_keys) + _tracker.__enter__() + try: + return await _openai_passthrough_non_streaming( + llama_backend, + payload, + model_name, + monitor_id = monitor_id, + request = request, + cancel_event = cancel_event, + ) + finally: + _tracker.__exit__(None, None, None) # ── Parse messages (handles multimodal content parts) ───── # Reuse the pre-hook parse when auto-switch did it, else parse now. @@ -5821,6 +6809,8 @@ async def openai_chat_completions( ) def _gguf_chat_delta_line(delta: ChoiceDelta, finish_reason = None) -> str: + if delta.reasoning_content is not None and delta.content is None: + delta = delta.model_copy(update = {"content": ""}) chunk = ChatCompletionChunk( id = completion_id, created = created, @@ -5844,8 +6834,12 @@ async def openai_chat_completions( from state.tool_policy import get_tool_policy as _get_tool_policy_g _cli_policy = _get_tool_policy_g() - _tools_on = _effective_enable_tools(payload) - _mcp_allowed = bool(payload.mcp_enabled) and _cli_policy is not False + _tools_on = False if _client_disabled_tool_calls else _effective_enable_tools(payload) + _mcp_allowed = ( + not _client_disabled_tool_calls + and bool(payload.mcp_enabled) + and _cli_policy is not False + ) use_tools = (_tools_on or _mcp_allowed) and llama_backend.supports_tools if use_tools: @@ -5895,13 +6889,18 @@ async def openai_chat_completions( _gguf_auto_heal_tool_calls = ( payload.auto_heal_tool_calls if payload.auto_heal_tool_calls is not None else True ) + # Active tool names gating the bare-rehearsal strip, matching the loop gate. + _gguf_display_tool_names = _display_tool_name_gate(tools_to_use) # ── Strip stale tool-call XML from conversation history ─ for _msg in gguf_messages: if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str): + # Gate on enabled tool names, like the live strip, so a documented inactive + # ``foo[ARGS]{...}`` survives in the replayed prompt context. _msg["content"] = _strip_tool_xml_for_display( _msg["content"], auto_heal_tool_calls = _gguf_auto_heal_tool_calls, + enabled_tool_names = _gguf_display_tool_names, ).strip() def gguf_generate_with_tools(): @@ -5922,6 +6921,7 @@ async def openai_chat_completions( reasoning_effort = payload.reasoning_effort, preserve_thinking = payload.preserve_thinking, auto_heal_tool_calls = _gguf_auto_heal_tool_calls, + nudge_tool_calls = payload.nudge_tool_calls, max_tool_iterations = payload.max_tool_calls_per_message if payload.max_tool_calls_per_message is not None else 25, @@ -5938,6 +6938,24 @@ async def openai_chat_completions( bypass_permissions = bool(payload.bypass_permissions), ) + _tool_admission_mode = "chat_tool_stream" if payload.stream else "chat_tool_nonstream" + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = _tool_admission_mode, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + _tool_sentinel = object() _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) @@ -5946,6 +6964,8 @@ async def openai_chat_completions( async def gguf_tool_stream(): gen = None + next_task = None + stream_completed = False disconnect_watcher = asyncio.create_task( _await_disconnect_then_cancel(request, cancel_event) ) @@ -5983,7 +7003,14 @@ async def openai_chat_completions( api_monitor.finish(monitor_id, "cancelled") return - event = await asyncio.to_thread(next, gen, _tool_sentinel) + next_task = asyncio.create_task( + asyncio.to_thread(next, gen, _tool_sentinel) + ) + try: + event = await asyncio.shield(next_task) + finally: + if next_task.done(): + next_task = None if event is _tool_sentinel: break @@ -6035,6 +7062,7 @@ async def openai_chat_completions( clean_cumulative = _strip_tool_xml_for_display( raw_cumulative, auto_heal_tool_calls = _gguf_auto_heal_tool_calls, + enabled_tool_names = _gguf_display_tool_names, ) new_text = clean_cumulative[len(prev_text) :] prev_text = clean_cumulative @@ -6081,6 +7109,7 @@ async def openai_chat_completions( api_monitor.finish( monitor_id, "cancelled" if cancel_event.is_set() else "completed" ) + stream_completed = True yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -6093,19 +7122,157 @@ async def openai_chat_completions( # Recover if an MTP+tensor crash killed the server mid-stream. get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) error_chunk = _openai_stream_error_chunk(e) - yield f"data: {json.dumps(error_chunk)}\n\n" + yield _openai_stream_error_sse(error_chunk) finally: - await _stop_local_disconnect_cancel_watcher(disconnect_watcher) - if gen is not None: - try: - gen.close() - except (RuntimeError, ValueError): - pass - _tracker.__exit__(None, None, None) + try: + if not stream_completed: + cancel_event.set() + task_to_drain = next_task + next_task = None + while task_to_drain is not None and not task_to_drain.done(): + try: + await asyncio.shield(task_to_drain) + except asyncio.CancelledError: + cancel_event.set() + continue + except Exception: + break + if task_to_drain is not None and task_to_drain.done(): + try: + task_to_drain.exception() + except (asyncio.CancelledError, Exception): + pass + if gen is not None and not stream_completed: + try: + await asyncio.to_thread(gen.close) + except (RuntimeError, ValueError): + pass + except Exception: + logger.debug( + "Error closing GGUF tool stream generator during cleanup", + exc_info = True, + ) + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) + finally: + _tracker.__exit__(None, None, None) if payload.stream: + stream_lease = reservation.lease_nowait() + admission_wait_started_at = None + if stream_lease is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = _tool_admission_mode, + completion_id = completion_id, + level = "debug", + ) + + async def admitted_gguf_tool_stream(): + lease = stream_lease + stream_started = False + stream_cancelled = False + try: + if lease is None: + async for wait_item in _openai_admission_wait_stream_chunks( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ): + if isinstance(wait_item, str): + yield wait_item + continue + lease = wait_item + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + break + if lease is None: + return + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + iterator = gguf_tool_stream() + stream_started = True + try: + async for chunk in iterator: + yield chunk + except asyncio.CancelledError: + stream_cancelled = True + raise + finally: + await _close_openai_admitted_stream_iterator( + iterator, + cancelled = stream_cancelled, + ) + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + yield _openai_stream_error_sse( + _openai_admission_error_body(exc, status_code = 503) + ) + except LlamaAdmissionCancelled: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except HTTPException as exc: + status_code = getattr(exc, "status_code", 500) or 500 + detail = exc.detail + error = ( + detail + if isinstance(detail, dict) and "error" in detail + else openai_error_body(str(detail), status = status_code) + ) + api_monitor.fail(monitor_id, str(detail)) + yield _openai_stream_error_sse(error) + finally: + if lease is not None: + lease.release() + if not stream_started: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + _tracker.__exit__(None, None, None) + + async def _gguf_tool_admission_unstarted_cleanup() -> None: + api_monitor.finish(monitor_id, "cancelled") + if stream_lease is not None: + stream_lease.release() + reservation.cancel() + _tracker.__exit__(None, None, None) + return _SameTaskStreamingResponse( - gguf_tool_stream(), + admitted_gguf_tool_stream(), + unstarted_cleanup = _gguf_tool_admission_unstarted_cleanup, media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -6140,6 +7307,7 @@ async def openai_chat_completions( full_text = _strip_tool_xml_for_display( event.get("text", ""), auto_heal_tool_calls = _gguf_auto_heal_tool_calls, + enabled_tool_names = _gguf_display_tool_names, ) return full_text, usage, finish finally: @@ -6150,10 +7318,61 @@ async def openai_chat_completions( except (RuntimeError, ValueError): pass + drain_task = None + + async def _drain_cancelled_gguf_tool_task(): + if drain_task is None: + return + while not drain_task.done(): + try: + await asyncio.shield(drain_task) + except asyncio.CancelledError: + cancel_event.set() + continue + except Exception: + break + if drain_task.done(): + try: + drain_task.exception() + except (asyncio.CancelledError, Exception): + pass + + admission_lease = None + admission_wait_started_at = None try: - full_text, completion_usage, completion_finish = await asyncio.to_thread( - _drain_gguf_tool_loop + if reservation.lease_nowait() is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = _tool_admission_mode, + completion_id = completion_id, + level = "debug", + ) + admission_lease = await _wait_for_openai_admission_non_streaming( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, ) + if admission_wait_started_at is not None: + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + drain_task = asyncio.create_task(asyncio.to_thread(_drain_gguf_tool_loop)) + full_text, completion_usage, completion_finish = await asyncio.shield(drain_task) reasoning_text, visible_text = _extract_responses_reasoning( full_text, parse_think_markers = _responses_should_parse_think_markers( @@ -6199,6 +7418,48 @@ async def openai_chat_completions( monitor_id, "cancelled" if cancel_event.is_set() else "completed" ) return _model_json_response(response) + except asyncio.CancelledError: + cancel_event.set() + await _drain_cancelled_gguf_tool_task() + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise _openai_admission_http_exception(exc, status_code = 503) + except LlamaAdmissionCancelled as exc: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = _tool_admission_mode, + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise HTTPException( + status_code = 499, + detail = _openai_admission_error_body(exc, status_code = 499), + ) except Exception as e: logger.error(f"Error during GGUF tool completion: {e}", exc_info = True) api_monitor.fail(monitor_id, _friendly_error(e)) @@ -6219,6 +7480,8 @@ async def openai_chat_completions( ) raise HTTPException(status_code = 500, detail = safe_error_detail(e)) finally: + if admission_lease is not None: + admission_lease.release() _tracker.__exit__(None, None, None) # ── Standard GGUF path (no tools) ───────────────────── @@ -6253,11 +7516,31 @@ async def openai_chat_completions( _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) _tracker = _TrackedCancel(cancel_event, *_cancel_keys) _tracker.__enter__() + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _tracker.__exit__(None, None, None) + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "chat_standard_stream", + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) async def gguf_stream_chunks(): disconnect_watcher = asyncio.create_task( _await_disconnect_then_cancel(request, cancel_event) ) + gen = None + next_task = None + stream_completed = False try: yield _chat_role_chunk(completion_id, created, model_name) @@ -6276,7 +7559,14 @@ async def openai_chat_completions( cancel_event.set() api_monitor.finish(monitor_id, "cancelled") return - cumulative = await asyncio.to_thread(next, gen, _gguf_sentinel) + next_task = asyncio.create_task( + asyncio.to_thread(next, gen, _gguf_sentinel) + ) + try: + cumulative = await asyncio.shield(next_task) + finally: + if next_task.done(): + next_task = None if cumulative is _gguf_sentinel: break # Capture server metadata for the final usage chunk @@ -6345,6 +7635,7 @@ async def openai_chat_completions( api_monitor.finish( monitor_id, "cancelled" if cancel_event.is_set() else "completed" ) + stream_completed = True yield "data: [DONE]\n\n" except asyncio.CancelledError: @@ -6355,13 +7646,156 @@ async def openai_chat_completions( logger.error(f"Error during GGUF streaming: {e}", exc_info = True) api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = _openai_stream_error_chunk(e) - yield f"data: {json.dumps(error_chunk)}\n\n" + yield _openai_stream_error_sse(error_chunk) finally: - await _stop_local_disconnect_cancel_watcher(disconnect_watcher) - _tracker.__exit__(None, None, None) + try: + if not stream_completed: + cancel_event.set() + task_to_drain = next_task + next_task = None + while task_to_drain is not None and not task_to_drain.done(): + try: + await asyncio.shield(task_to_drain) + except asyncio.CancelledError: + cancel_event.set() + continue + except Exception: + break + if task_to_drain is not None and task_to_drain.done(): + try: + task_to_drain.exception() + except (asyncio.CancelledError, Exception): + pass + if gen is not None and not stream_completed: + try: + await asyncio.to_thread(gen.close) + except (RuntimeError, ValueError): + pass + except Exception: + logger.debug( + "Error closing GGUF stream generator during cleanup", + exc_info = True, + ) + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) + finally: + _tracker.__exit__(None, None, None) + + stream_lease = reservation.lease_nowait() + admission_wait_started_at = None + if stream_lease is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "chat_standard_stream", + completion_id = completion_id, + level = "debug", + ) + + async def admitted_gguf_stream_chunks(): + lease = stream_lease + stream_started = False + stream_cancelled = False + try: + if lease is None: + async for wait_item in _openai_admission_wait_stream_chunks( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ): + if isinstance(wait_item, str): + yield wait_item + continue + lease = wait_item + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "chat_standard_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + break + if lease is None: + return + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + iterator = gguf_stream_chunks() + stream_started = True + try: + async for chunk in iterator: + yield chunk + except asyncio.CancelledError: + stream_cancelled = True + raise + finally: + await _close_openai_admitted_stream_iterator( + iterator, + cancelled = stream_cancelled, + ) + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "chat_standard_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + yield _openai_stream_error_sse( + _openai_admission_error_body(exc, status_code = 503) + ) + except LlamaAdmissionCancelled: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "chat_standard_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except HTTPException as exc: + status_code = getattr(exc, "status_code", 500) or 500 + detail = exc.detail + error = ( + detail + if isinstance(detail, dict) and "error" in detail + else openai_error_body(str(detail), status = status_code) + ) + api_monitor.fail(monitor_id, str(detail)) + yield _openai_stream_error_sse(error) + finally: + if lease is not None: + lease.release() + if not stream_started: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + _tracker.__exit__(None, None, None) + + async def _gguf_admission_unstarted_cleanup() -> None: + api_monitor.finish(monitor_id, "cancelled") + if stream_lease is not None: + stream_lease.release() + reservation.cancel() + _tracker.__exit__(None, None, None) return _SameTaskStreamingResponse( - gguf_stream_chunks(), + admitted_gguf_stream_chunks(), + unstarted_cleanup = _gguf_admission_unstarted_cleanup, media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -6370,57 +7804,189 @@ async def openai_chat_completions( }, ) else: + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "chat_standard_nonstream", + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + + _cancel_keys = (payload.cancel_id, payload.session_id, completion_id) + _tracker = _TrackedCancel(cancel_event, *_cancel_keys) + _tracker.__enter__() + admission_lease = None + admission_wait_started_at = None + try: + if reservation.lease_nowait() is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "chat_standard_nonstream", + completion_id = completion_id, + level = "debug", + ) + admission_lease = await _wait_for_openai_admission_non_streaming( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ) + if admission_wait_started_at is not None: + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "chat_standard_nonstream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "chat_standard_nonstream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise _openai_admission_http_exception(exc, status_code = 503) + except LlamaAdmissionCancelled as exc: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "chat_standard_nonstream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) + raise HTTPException( + status_code = 499, + detail = _openai_admission_error_body(exc, status_code = 499), + ) + try: # ``n`` requests several independent completions; the single # decode slot yields one at a time, so loop sequentially. - _n = payload.n or 1 + drain_task = None - _choices = [] - _monitor_replies = [] - _prompt_tokens = 0 - _sum_completion = 0 - _prompt_details = None - for _idx in range(_n): - # Stop spawning the remaining choices once cancelled. - if cancel_event.is_set(): - break - full_text = "" - completion_usage = None - completion_finish = None - for token in gguf_generate(_idx): - if isinstance(token, dict): - if token.get("type") == "metadata": - completion_usage = token.get("usage") - completion_finish = token.get("finish_reason") + async def _drain_cancelled_gguf_task(): + if drain_task is None: + return + while not drain_task.done(): + try: + await asyncio.shield(drain_task) + except asyncio.CancelledError: + cancel_event.set() continue - full_text = token + except Exception: + break + if drain_task.done(): + try: + drain_task.exception() + except (asyncio.CancelledError, Exception): + pass - reasoning_text, visible_text = _extract_responses_reasoning( - full_text, - parse_think_markers = _responses_should_parse_think_markers( - payload, - llama_backend, - ), - ) - message_kwargs = {"content": visible_text} - if reasoning_text: - message_kwargs["reasoning_content"] = reasoning_text - _choices.append( - CompletionChoice( - index = _idx, - message = CompletionMessage(**message_kwargs), - finish_reason = _clamp_finish_reason(completion_finish), + def _drain_gguf_choices(): + _n = payload.n or 1 + _choices = [] + _monitor_replies = [] + _prompt_tokens = 0 + _sum_completion = 0 + _prompt_details = None + for _idx in range(_n): + # Stop spawning the remaining choices once cancelled. + if cancel_event.is_set(): + break + full_text = "" + completion_usage = None + completion_finish = None + for token in gguf_generate(_idx): + if isinstance(token, dict): + if token.get("type") == "metadata": + completion_usage = token.get("usage") + completion_finish = token.get("finish_reason") + continue + full_text = token + + reasoning_text, visible_text = _extract_responses_reasoning( + full_text, + parse_think_markers = _responses_should_parse_think_markers( + payload, + llama_backend, + ), ) + message_kwargs = {"content": visible_text} + if reasoning_text: + message_kwargs["reasoning_content"] = reasoning_text + _choices.append( + CompletionChoice( + index = _idx, + message = CompletionMessage(**message_kwargs), + finish_reason = _clamp_finish_reason(completion_finish), + ) + ) + _monitor_replies.append(visible_text) + if completion_usage: + # The prompt is shared across all n choices, so count its + # tokens ONCE (OpenAI bills only generated tokens for each + # extra choice). Only completion_tokens accumulates. + _prompt_tokens = completion_usage.get("prompt_tokens") or _prompt_tokens + _sum_completion += completion_usage.get("completion_tokens") or 0 + if _prompt_details is None: + _prompt_details = completion_usage.get("prompt_tokens_details") + return ( + _n, + _choices, + _monitor_replies, + _prompt_tokens, + _sum_completion, + _prompt_details, ) - _monitor_replies.append(visible_text) - if completion_usage: - # The prompt is shared across all n choices, so count its - # tokens ONCE (OpenAI bills only generated tokens for each - # extra choice). Only completion_tokens accumulates. - _prompt_tokens = completion_usage.get("prompt_tokens") or _prompt_tokens - _sum_completion += completion_usage.get("completion_tokens") or 0 - if _prompt_details is None: - _prompt_details = completion_usage.get("prompt_tokens_details") + + drain_task = asyncio.create_task(asyncio.to_thread(_drain_gguf_choices)) + ( + _n, + _choices, + _monitor_replies, + _prompt_tokens, + _sum_completion, + _prompt_details, + ) = await asyncio.shield(drain_task) response = ChatCompletion( id = completion_id, @@ -6452,6 +8018,11 @@ async def openai_chat_completions( api_monitor.finish(monitor_id) return _model_json_response(response) + except asyncio.CancelledError: + cancel_event.set() + await _drain_cancelled_gguf_task() + api_monitor.finish(monitor_id, "cancelled") + raise except Exception as e: logger.error(f"Error during GGUF completion: {e}", exc_info = True) api_monitor.fail(monitor_id, _friendly_error(e)) @@ -6471,7 +8042,10 @@ async def openai_chat_completions( ), ) raise HTTPException(status_code = 500, detail = safe_error_detail(e)) - + finally: + if admission_lease is not None: + admission_lease.release() + _tracker.__exit__(None, None, None) # ── Standard Unsloth path ───────────────────────────────── # Decode image (from content parts OR legacy field) @@ -6511,6 +8085,25 @@ async def openai_chat_completions( _sf_tpl = (_sf_model_info.get("chat_template_info") or {}).get("template") _sf_features = _detect_safetensors_features(backend, _sf_tpl) + # GGUF parity: enable_thinking templates prefill an unclosed ; split into + # reasoning_content deltas so the UI renders the block for safetensors and MLX. + _sf_parse_think = bool( + _sf_features.get("supports_reasoning") or _sf_features.get("reasoning_always_on") + ) + # Prefilled-open only for prefill styles with thinking on; gpt-oss uses the normal mode. + _sf_reasoning_prefilled = _sf_reasoning_prefill_mode( + _sf_features, + payload.enable_thinking, + _sf_tpl, + reasoning_effort = payload.reasoning_effort, + ) + + def _new_sf_reasoning_extractor(): + return _ResponsesReasoningExtractor( + parse_think_markers = _sf_parse_think, + reasoning_prefilled = _sf_reasoning_prefilled, + ) + cancel_event = threading.Event() completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" created = int(time.time()) @@ -6586,6 +8179,8 @@ async def openai_chat_completions( _sf_auto_heal_tool_calls = ( payload.auto_heal_tool_calls if payload.auto_heal_tool_calls is not None else True ) + # Active tool names gating the bare-rehearsal strip, matching the loop gate. + _sf_display_tool_names = _display_tool_name_gate(_sf_tools_to_use) # Strip stale tool-call XML from prior assistant turns. _sf_chat_messages = [] @@ -6597,6 +8192,7 @@ async def openai_chat_completions( "content": _strip_tool_xml_for_display( _msg["content"], auto_heal_tool_calls = _sf_auto_heal_tool_calls, + enabled_tool_names = _sf_display_tool_names, ).strip(), } ) @@ -6617,11 +8213,13 @@ async def openai_chat_completions( min_p = payload.min_p, max_tokens = effective_max_tokens, repetition_penalty = payload.repetition_penalty, + presence_penalty = payload.presence_penalty, cancel_event = cancel_event, enable_thinking = payload.enable_thinking, reasoning_effort = payload.reasoning_effort, preserve_thinking = payload.preserve_thinking, auto_heal_tool_calls = _sf_auto_heal_tool_calls, + nudge_tool_calls = payload.nudge_tool_calls, max_tool_iterations = _sf_tool_budget, tool_call_timeout = payload.tool_call_timeout if payload.tool_call_timeout is not None @@ -6652,6 +8250,19 @@ async def openai_chat_completions( gen = sf_generate_with_tools() prev_text = "" + reasoning_extractor = _new_sf_reasoning_extractor() + + def _sf_flush_reasoning(): + # Drain the extractor at turn/stream end (mirrors GGUF); only visible text hits the monitor. + fr, fv = reasoning_extractor.finish() + out = [] + if fr: + out.append(_chat_reasoning_chunk(completion_id, created, model_name, fr)) + if fv: + api_monitor.append_reply(monitor_id, fv) + out.append(_chat_content_chunk(completion_id, created, model_name, fv)) + return out + while True: if cancel_event.is_set(): backend.reset_generation_state() @@ -6668,7 +8279,11 @@ async def openai_chat_completions( if event["type"] == "status": if not event["text"]: + # Iteration boundary: flush reasoning, then a fresh prefilled extractor for the next turn. + for _c in _sf_flush_reasoning(): + yield _c prev_text = "" + reasoning_extractor = _new_sf_reasoning_extractor() status_data = json.dumps( { "type": "tool_status", @@ -6680,7 +8295,11 @@ async def openai_chat_completions( if event["type"] in ("tool_start", "tool_end"): if event["type"] == "tool_start": + # Flush reasoning before tool_start so the thinking block closes ahead of the card. + for _c in _sf_flush_reasoning(): + yield _c prev_text = "" + reasoning_extractor = _new_sf_reasoning_extractor() yield f"data: {json.dumps(event)}\n\n" continue @@ -6689,14 +8308,24 @@ async def openai_chat_completions( clean_cumulative = _strip_tool_xml_for_display( raw_cumulative, auto_heal_tool_calls = _sf_auto_heal_tool_calls, + enabled_tool_names = _sf_display_tool_names, ) new_text = clean_cumulative[len(prev_text) :] prev_text = clean_cumulative if not new_text: continue - api_monitor.append_reply(monitor_id, new_text) - yield _chat_content_chunk(completion_id, created, model_name, new_text) + # Split reasoning vs visible; only visible reaches the monitor. + reasoning_delta, visible_delta = reasoning_extractor.feed(new_text) + if reasoning_delta: + yield _chat_reasoning_chunk( + completion_id, created, model_name, reasoning_delta + ) + if visible_delta: + api_monitor.append_reply(monitor_id, visible_delta) + yield _chat_content_chunk(completion_id, created, model_name, visible_delta) + for _c in _sf_flush_reasoning(): + yield _c yield _chat_final_chunk(completion_id, created, model_name, "stop") # Usage chunk from the last turn, same shape as the # GGUF tool loop's metadata. Request-scoped holder, so @@ -6736,7 +8365,7 @@ async def openai_chat_completions( "type": "server_error", }, } - yield f"data: {json.dumps(error_chunk)}\n\n" + yield _openai_stream_error_sse(error_chunk) finally: await _stop_local_disconnect_cancel_watcher(disconnect_watcher) if gen is not None: @@ -6749,6 +8378,7 @@ async def openai_chat_completions( if payload.stream: return _SameTaskStreamingResponse( sf_tool_stream(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_sf_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -6770,22 +8400,32 @@ async def openai_chat_completions( full_text = _strip_tool_xml_for_display( event.get("text", ""), auto_heal_tool_calls = _sf_auto_heal_tool_calls, + enabled_tool_names = _sf_display_tool_names, ) return full_text content_text = await asyncio.to_thread(_drain_to_text) - api_monitor.set_reply(monitor_id, content_text) + # Split prefilled out of the visible answer (GGUF parity); the monitor gets visible text only. + _reasoning_text, _visible_text = _extract_responses_reasoning( + content_text, + parse_think_markers = _sf_parse_think, + reasoning_prefilled = _sf_reasoning_prefilled, + ) + api_monitor.set_reply(monitor_id, _visible_text) _stats = _sf_stats_holder.get("stats") if _stats: _monitor_usage(monitor_id, _stats.get("usage")) api_monitor.finish(monitor_id, "cancelled" if cancel_event.is_set() else "completed") + _sf_msg_kwargs = {"content": _visible_text} + if _reasoning_text: + _sf_msg_kwargs["reasoning_content"] = _reasoning_text response = ChatCompletion( id = completion_id, created = created, model = model_name, choices = [ CompletionChoice( - message = CompletionMessage(content = content_text), + message = CompletionMessage(**_sf_msg_kwargs), finish_reason = "stop", ) ], @@ -6819,6 +8459,7 @@ async def openai_chat_completions( min_p = payload.min_p, max_new_tokens = effective_max_tokens or 2048, repetition_penalty = payload.repetition_penalty, + presence_penalty = payload.presence_penalty, ) # Forward reasoning kwargs; the worker/template wrapper peels off any the # template doesn't accept. @@ -6829,25 +8470,87 @@ async def openai_chat_completions( if payload.preserve_thinking is not None: gen_kwargs["preserve_thinking"] = payload.preserve_thinking + # ── Client-tool passthrough (safetensors + MLX) ────────────── + # Client tools (or tool-result history) without server-side tools: render + # tools into the template, generate one turn, heal text-form calls (#6801). + # supports_tools=False falls through to plain relay (GGUF gate parity). + _sf_has_tool_msgs = any(m.role == "tool" or m.tool_calls for m in payload.messages) + # Gate on _sf_use_tools (did the server-side path claim the request?), not + # raw mcp_enabled: an empty MCP registry must not silently drop client tools. + _sf_client_tools = ( + not _effective_enable_tools(payload) + and not _sf_use_tools + and image is None + and not _sf_is_gptoss + and _sf_features.get("supports_tools", False) + and ((payload.tools and len(payload.tools) > 0) or _sf_has_tool_msgs) + ) + _sf_heal = ( + heal_gate(payload.auto_heal_tool_calls, payload.tools, payload.tool_choice) + if _sf_client_tools + else None + ) + if _sf_client_tools: + # Re-derive from payload.messages so tool_calls / role="tool" history + # survives templating; fold system/developer into one leading system + # message (templates reject "developer") and clear prompt to avoid a dup. + gen_kwargs["messages"] = _set_or_prepend_system_message( + _structured_tool_history_for_local_template( + _flatten_content_parts_for_local_template(_openai_messages_for_passthrough(payload)) + ), + system_prompt, + ) + gen_kwargs["system_prompt"] = "" + # tool_choice="none": keep history templating but advertise no tools + # (heal_gate is off, markup would relay as prose). A forced function + # narrows templating to that one schema. Both mirror the GGUF path, + # where llama-server honors tool_choice itself. + _sf_tc = payload.tool_choice + _sf_forced = None + if isinstance(_sf_tc, dict) and isinstance(_sf_tc.get("function"), dict): + _sf_forced = _sf_tc["function"].get("name") + if _sf_tc == "none": + gen_kwargs["tools"] = None + elif isinstance(_sf_forced, str): + gen_kwargs["tools"] = [ + t + for t in payload.tools or [] + if isinstance(t, dict) + and isinstance(t.get("function"), dict) + and t["function"].get("name") == _sf_forced + ] or None + else: + gen_kwargs["tools"] = payload.tools + # Request-scoped usage/timings receptacle (filled at gen_done). stats_holder: dict = {} if payload.use_adapter is not None: - def generate(): + def generate(messages_override = None): + kw = ( + gen_kwargs + if messages_override is None + else {**gen_kwargs, "messages": messages_override} + ) return backend.generate_with_adapter_control( use_adapter = payload.use_adapter, cancel_event = cancel_event, stats_holder = stats_holder, - **gen_kwargs, + **kw, ) else: - def generate(): + def generate(messages_override = None): + kw = ( + gen_kwargs + if messages_override is None + else {**gen_kwargs, "messages": messages_override} + ) return backend.generate_chat_response( cancel_event = cancel_event, stats_holder = stats_holder, - **gen_kwargs, + **kw, ) # ── Streaming response ──────────────────────────────────────── @@ -6863,7 +8566,14 @@ async def openai_chat_completions( try: yield _chat_role_chunk(completion_id, created, model_name) + # Client-tool passthrough: heal text-form calls on the fly + # (None => relay verbatim). + healer = StreamToolCallHealer(_sf_heal, payload.tools) if _sf_heal else None + heal_state = {"idx": 0} + prev_text = "" + # Split prefilled into reasoning_content deltas (GGUF parity); single turn, serves MLX. + reasoning_extractor = _new_sf_reasoning_extractor() # Run the sync generator in a thread pool to avoid blocking the # event loop. Critical for compare mode: two SSE requests arrive # concurrently but the orchestrator serializes them via @@ -6892,10 +8602,76 @@ async def openai_chat_completions( prev_text = cumulative if not new_text: continue - api_monitor.append_reply(monitor_id, new_text) - yield _chat_content_chunk(completion_id, created, model_name, new_text) + # Split prefilled reasoning first (GGUF/MLX parity), + # then route only the visible text through the client-tool + # healer so tool markup inside a reasoning block is not promoted. + reasoning_delta, visible_delta = reasoning_extractor.feed(new_text) + if reasoning_delta: + yield _chat_reasoning_chunk( + completion_id, created, model_name, reasoning_delta + ) + if visible_delta: + if healer is None: + # Monitor mirrors the verbatim relay; with healing on, + # _sf_heal_events_to_sse records the healed events instead. + api_monitor.append_reply(monitor_id, visible_delta) + yield _chat_content_chunk( + completion_id, created, model_name, visible_delta + ) + else: + for line in _sf_heal_events_to_sse( + healer.feed(visible_delta), + completion_id, + created, + model_name, + heal_state, + payload.parallel_tool_calls, + monitor_id, + ): + yield line - yield _chat_final_chunk(completion_id, created, model_name, "stop") + final_reasoning, final_visible = reasoning_extractor.finish() + if final_reasoning: + yield _chat_reasoning_chunk(completion_id, created, model_name, final_reasoning) + if final_visible: + if healer is None: + api_monitor.append_reply(monitor_id, final_visible) + yield _chat_content_chunk(completion_id, created, model_name, final_visible) + else: + for line in _sf_heal_events_to_sse( + healer.feed(final_visible), + completion_id, + created, + model_name, + heal_state, + payload.parallel_tool_calls, + monitor_id, + ): + yield line + + # A cancelled stream must not promote buffered-but-incomplete + # markup: finalize()'s allow_incomplete heal would execute a tool + # the user just cancelled. Disconnect returns earlier; "Stop" only + # sets cancel_event, so guard on it here too. + _cancelled = cancel_event.is_set() + if healer is not None and not _cancelled: + for line in _sf_heal_events_to_sse( + healer.finalize(), + completion_id, + created, + model_name, + heal_state, + payload.parallel_tool_calls, + monitor_id, + ): + yield line + + _finish = ( + "tool_calls" + if (healer is not None and not _cancelled and healer.healed) + else "stop" + ) + yield _chat_final_chunk(completion_id, created, model_name, _finish) # Usage chunk (choices=[], usage set), same shape as the # GGUF path so the speed popover works for MLX too. # Request-scoped holder, so concurrent streams cannot @@ -6934,13 +8710,14 @@ async def openai_chat_completions( "type": "server_error", }, } - yield f"data: {json.dumps(error_chunk)}\n\n" + yield _openai_stream_error_sse(error_chunk) finally: await _stop_local_disconnect_cancel_watcher(disconnect_watcher) _tracker.__exit__(None, None, None) return _SameTaskStreamingResponse( stream_chunks(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -6956,18 +8733,96 @@ async def openai_chat_completions( for token in generate(): full_text = token + # Split prefilled reasoning (GGUF parity); also covers MLX via + # the shared generate(). Client-tool healing then runs on the visible + # text so tool markup inside a reasoning block is never promoted. + _reasoning_text, _visible_text = _extract_responses_reasoning( + full_text, + parse_think_markers = _sf_parse_think, + reasoning_prefilled = _sf_reasoning_prefilled, + ) + # Client-tool passthrough: promote text-form calls; opt-in single + # nudge retry on unparseable tool markup. + _msg = {"role": "assistant", "content": _visible_text} + if _reasoning_text: + _msg["reasoning_content"] = _reasoning_text + _finish = "stop" + if _sf_heal: + if heal_openai_message(_msg, _sf_heal, payload.tools): + _finish = "tool_calls" + elif nudge_enabled(payload.nudge_tool_calls): + _data = { + "choices": [{"message": {"role": "assistant", "content": _visible_text}}] + } + if nudge_should_retry(_data, _sf_heal, payload.tools): + # A failed retry must not 500 the request; keep the first + # response (GGUF nudge parity). The retry's generate() + # overwrites stats_holder, so save the first attempt's stats + # and restore them if the retry is discarded. + _first_stats = stats_holder.get("stats") + try: + retry_text = "" + for token in generate( + [*gen_kwargs["messages"], *nudge_messages(_data, _sf_heal)] + ): + retry_text = token + # Re-split reasoning on the retry so its visible text is + # what heals into a call (and reaches the monitor). + _retry_reasoning, _retry_visible = _extract_responses_reasoning( + retry_text, + parse_think_markers = _sf_parse_think, + reasoning_prefilled = _sf_reasoning_prefilled, + ) + retry_msg = {"role": "assistant", "content": _retry_visible} + if _retry_reasoning: + retry_msg["reasoning_content"] = _retry_reasoning + if heal_openai_message(retry_msg, _sf_heal, payload.tools): + _visible_text, _msg, _finish = ( + _retry_visible, + retry_msg, + "tool_calls", + ) + else: + # Retry produced no healable call -> first response wins. + stats_holder["stats"] = _first_stats + except Exception as retry_exc: + logger.debug( + "Nudge retry failed; keeping first response: %s", retry_exc + ) + stats_holder["stats"] = _first_stats + # parallel_tool_calls=false: cap to one call (GGUF parity). + if payload.parallel_tool_calls is False: + _tcs = _msg.get("tool_calls") + if isinstance(_tcs, list) and len(_tcs) > 1: + _msg["tool_calls"] = _tcs[:1] + response = ChatCompletion( id = completion_id, created = created, model = model_name, choices = [ CompletionChoice( - message = CompletionMessage(content = full_text), - finish_reason = "stop", + message = CompletionMessage( + content = _msg["content"], + reasoning_content = _msg.get("reasoning_content"), + tool_calls = _msg.get("tool_calls"), + ), + finish_reason = _finish, ) ], ) - api_monitor.set_reply(monitor_id, full_text) + _monitor_reply = _msg.get("content") or "" + if _finish == "tool_calls": + _tcs = _msg.get("tool_calls") or [] + _calls_text = "; ".join( + f"{(tc.get('function') or {}).get('name', '')}" + f"({(tc.get('function') or {}).get('arguments', '')})" + for tc in _tcs + ) + _monitor_reply = (_msg.get("content") or "") + ( + f"[tool_calls] {_calls_text}" if _calls_text else "" + ) + api_monitor.set_reply(monitor_id, _monitor_reply) _stats = stats_holder.get("stats") if _stats: _monitor_usage(monitor_id, _stats.get("usage")) @@ -7344,7 +9199,7 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge if not llama_backend.is_loaded: raise HTTPException( status_code = 503, - detail = "No GGUF model loaded. Load a GGUF model first.", + detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), ) if not isinstance(body, dict): # Re-read to re-raise a malformed-body error (post-503, pre-feature behavior); @@ -7353,8 +9208,12 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge if not isinstance(body, dict): raise HTTPException(status_code = 400, detail = "Request body must be a JSON object") - if body.get("max_tokens") is None: - body["max_tokens"] = llama_backend.context_length or _DEFAULT_MAX_TOKENS_FLOOR + _resolved_max_tokens = _effective_openai_max_tokens_from_values(body.get("max_tokens")) + body["max_tokens"] = ( + _resolved_max_tokens + if _resolved_max_tokens is not None + else (llama_backend.context_length or _DEFAULT_MAX_TOKENS_FLOOR) + ) target_url = f"{llama_backend.base_url}/v1/completions" is_stream = body.get("stream", False) prompt_text = _flatten_monitor_prompt(body.get("prompt", "")) @@ -7447,7 +9306,7 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge logger.error("openai_completions stream error: %s", e) api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = _openai_stream_error_chunk(e) - yield f"data: {json.dumps(error_chunk)}\n\n".encode("utf-8") + yield _openai_stream_error_sse_bytes(error_chunk) return api_monitor.finish(monitor_id, "cancelled") return @@ -7462,7 +9321,7 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge logger.error("openai_completions stream error: %s", e) api_monitor.fail(monitor_id, _friendly_error(e)) error_chunk = _openai_stream_error_chunk(e) - yield f"data: {json.dumps(error_chunk)}\n\n".encode("utf-8") + yield _openai_stream_error_sse_bytes(error_chunk) return finally: await _aclose_stream_resources( @@ -7556,7 +9415,7 @@ async def openai_embeddings(request: Request, current_subject: str = Depends(get if not llama_backend.is_loaded: raise HTTPException( status_code = 503, - detail = "No GGUF model loaded. Load a GGUF model first.", + detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), ) if not isinstance(body, dict): # Re-read to re-raise a malformed-body error (post-503, pre-feature behavior); @@ -7790,10 +9649,18 @@ def _responses_marker_holdback(text: str, markers: tuple[str, ...]) -> int: class _ResponsesReasoningExtractor: """Split local markup into Responses reasoning and visible text.""" - def __init__(self, *, parse_think_markers: bool = False) -> None: + def __init__( + self, + *, + parse_think_markers: bool = False, + reasoning_prefilled: bool = False, + ) -> None: self._buffer = "" - self._in_reasoning = False - self._parse_think_markers = parse_think_markers + # reasoning_prefilled: the template inserts an unclosed , so output begins inside + # the block; start in reasoning until the first close tag. Existing callers pass False. + self._in_reasoning = reasoning_prefilled + # Splitting requires marker parsing; a prefilled open implies it. + self._parse_think_markers = parse_think_markers or reasoning_prefilled def feed( self, @@ -7816,14 +9683,21 @@ class _ResponsesReasoningExtractor: if self._in_reasoning: close_idx = self._buffer.find(_RESPONSES_THINK_CLOSE) if close_idx != -1: - reasoning_parts.append(self._buffer[:close_idx]) + reasoning_parts.append( + self._buffer[:close_idx].replace(_RESPONSES_THINK_OPEN, "") + ) self._buffer = self._buffer[close_idx + len(_RESPONSES_THINK_CLOSE) :] self._in_reasoning = False continue - keep = _responses_marker_holdback(self._buffer, (_RESPONSES_THINK_CLOSE,)) + # Hold back a trailing partial of either marker: the close (clean split across chunks) + # and a stray open (a re-emitted is suppressed, not leaked). + keep = _responses_marker_holdback( + self._buffer, (_RESPONSES_THINK_CLOSE, _RESPONSES_THINK_OPEN) + ) if keep == len(self._buffer): break - reasoning_parts.append(self._buffer[:-keep] if keep else self._buffer) + emit = self._buffer[:-keep] if keep else self._buffer + reasoning_parts.append(emit.replace(_RESPONSES_THINK_OPEN, "")) self._buffer = self._buffer[-keep:] if keep else "" break @@ -7860,7 +9734,7 @@ class _ResponsesReasoningExtractor: return "", remaining if self._in_reasoning: self._in_reasoning = False - return remaining, "" + return remaining.replace(_RESPONSES_THINK_OPEN, ""), "" return "", remaining.replace(_RESPONSES_THINK_CLOSE, "") @@ -7869,8 +9743,12 @@ def _extract_responses_reasoning( reasoning_content: Any = None, *, parse_think_markers: bool = False, + reasoning_prefilled: bool = False, ) -> tuple[str, str]: - extractor = _ResponsesReasoningExtractor(parse_think_markers = parse_think_markers) + extractor = _ResponsesReasoningExtractor( + parse_think_markers = parse_think_markers, + reasoning_prefilled = reasoning_prefilled, + ) reasoning, visible = extractor.feed(text, reasoning_content) final_reasoning, final_visible = extractor.finish() return reasoning + final_reasoning, visible + final_visible @@ -8277,7 +10155,7 @@ async def _responses_stream( # so the client sees a useful error instead of a dangling stream. raise HTTPException( status_code = 400, - detail = ( + detail = _no_model_loaded_detail( "Streaming /v1/responses requires a GGUF model loaded via " "llama-server. Use non-streaming /v1/responses, " "/v1/chat/completions, or load a GGUF model." @@ -8299,6 +10177,52 @@ async def _responses_stream( ) body["stream_options"] = {"include_usage": True} target_url = f"{llama_backend.base_url}/v1/chat/completions" + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "responses_stream", + completion_id = resp_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + + def _responses_admission_failed_sse(exc: Exception, *, status_code: int) -> str: + return ( + "event: response.failed\n" + "data: " + + json.dumps( + { + "type": "response.failed", + "response": { + "id": resp_id, + "object": "response", + "created_at": created_at, + "status": "failed", + "model": _llama_public_model_id(llama_backend, payload.model) + or payload.model, + "output": [], + "usage": { + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0, + }, + "error": { + "code": status_code, + "message": str(exc), + }, + }, + } + ) + + "\n\n" + ) async def event_generator(): # Clean public id for every response envelope. Prefer the loaded model's @@ -9108,14 +11032,111 @@ async def _responses_stream( api_monitor.finish(monitor_id) yield _sse("response.completed", completed_response) + async def admitted_event_generator(): + lease = reservation.lease_nowait() + admission_wait_started_at = None + stream_started = False + stream_cancelled = False + iterator = None + try: + if lease is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "responses_stream", + completion_id = resp_id, + level = "debug", + ) + async for wait_item in _openai_admission_wait_stream_chunks( + reservation, + admission_config, + request = request, + cancel_event = None, + ): + if isinstance(wait_item, str): + yield wait_item + continue + lease = wait_item + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "responses_stream", + wait_started_at = admission_wait_started_at, + completion_id = resp_id, + level = "debug", + ) + break + if lease is None: + return + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = None, + ) + iterator = event_generator() + stream_started = True + try: + async for chunk in iterator: + yield chunk + except asyncio.CancelledError: + stream_cancelled = True + api_monitor.finish(monitor_id, "cancelled") + raise + finally: + await _close_openai_admitted_stream_iterator( + iterator, + cancelled = stream_cancelled, + ) + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "responses_stream", + wait_started_at = admission_wait_started_at, + completion_id = resp_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + yield _responses_admission_failed_sse(exc, status_code = 503) + except LlamaAdmissionCancelled: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "responses_stream", + wait_started_at = admission_wait_started_at, + completion_id = resp_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + finally: + if lease is not None: + lease.release() + if not stream_started: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + + async def _responses_admission_unstarted_cleanup() -> None: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + return _SameTaskStreamingResponse( - event_generator(), + admitted_event_generator(), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", "Connection": "close", "X-Accel-Buffering": "no", }, + unstarted_cleanup = _responses_admission_unstarted_cleanup, ) @@ -9373,7 +11394,7 @@ async def anthropic_count_tokens( if not llama_backend.is_loaded: raise HTTPException( status_code = 503, - detail = "No GGUF model loaded. Load a GGUF model first.", + detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), ) # Same Anthropic → OpenAI translation as anthropic_messages: system is @@ -9385,8 +11406,11 @@ async def anthropic_count_tokens( # Apply the same sanitization /messages does before generation, so the count # matches the prompt the real request would build (otherwise empty-assistant # sentinels / synthetic tool history inflate the count or hit the fallback). - openai_messages = _strip_provider_synthetic_tool_history( - _drop_empty_assistant_sentinels(openai_messages) + # Coalesce adjacent user turns left behind by dropping an empty / null assistant + # turn, so a strict GGUF chat template does not 400 on non-alternating roles + # (mirrors the GGUF chat path); a no-op for already-alternating histories. + openai_messages = _coalesce_consecutive_user_turns( + _strip_provider_synthetic_tool_history(_drop_empty_assistant_sentinels(openai_messages)) ) openai_tools = anthropic_tools_to_openai(payload.tools or []) or None @@ -9443,7 +11467,7 @@ async def anthropic_messages( if not llama_backend.is_loaded and not _automatic_model_load_may_run(): raise HTTPException( status_code = 503, - detail = "No GGUF model loaded. Load a GGUF model first.", + detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), ) # max_tokens is a required field on the Anthropic Messages API; real Anthropic @@ -9494,7 +11518,7 @@ async def anthropic_messages( if not llama_backend.is_loaded: raise HTTPException( status_code = 503, - detail = "No GGUF model loaded. Load a GGUF model first.", + detail = _no_model_loaded_detail("No GGUF model loaded. Load a GGUF model first."), ) # Advertised repo id after an auto-switch load, else a clean public id, never @@ -9514,8 +11538,11 @@ async def anthropic_messages( # builders apply the same strip; without it an Anthropic /v1/messages caller # replaying a prior provider-side tool_use forwards fake builtin tool # history to a backend with no matching function declarations. - openai_messages = _strip_provider_synthetic_tool_history( - _drop_empty_assistant_sentinels(openai_messages) + # Coalesce adjacent user turns left behind by dropping an empty / null assistant + # turn, so a strict GGUF chat template does not 400 on non-alternating roles + # (mirrors the GGUF chat path); a no-op for already-alternating histories. + openai_messages = _coalesce_consecutive_user_turns( + _strip_provider_synthetic_tool_history(_drop_empty_assistant_sentinels(openai_messages)) ) # Enforce vision guard + re-encode embedded images to PNG so the Anthropic @@ -9564,7 +11591,9 @@ async def anthropic_messages( and not _has_image ) client_tools = ( - not server_tools and len(openai_client_tools) > 0 and llama_backend.supports_tools + not server_tools + and len(openai_client_tools) > 0 + and getattr(llama_backend, "supports_tool_passthrough", llama_backend.supports_tools) ) # Anthropic tool_choice.disable_parallel_tool_use caps the response to a @@ -9697,10 +11726,16 @@ async def anthropic_messages( else: openai_messages.insert(0, {"role": "system", "content": _nudge}) - # Strip stale tool-call XML from conversation + # Strip stale tool-call XML via the protected display helper (think rehearsal and [TOOL_CALLS] + # prose survive), gated on enabled tool names so documented inactive examples are kept. + _anthropic_history_gate = _display_tool_name_gate(openai_tools) for _msg in openai_messages: if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str): - _msg["content"] = _TOOL_XML_RE.sub("", _msg["content"]).strip() + _msg["content"] = _strip_tool_xml_for_display( + _msg["content"], + auto_heal_tool_calls = True, + enabled_tool_names = _anthropic_history_gate, + ).strip() def _run_tool_gen(): return llama_backend.generate_chat_completion_with_tools( @@ -9717,6 +11752,7 @@ async def anthropic_messages( cancel_event = cancel_event, max_tool_iterations = 25, auto_heal_tool_calls = True, + nudge_tool_calls = payload.nudge_tool_calls, tool_call_timeout = 300, session_id = payload.session_id, # Anthropic passthrough has no rag_scope field (RAG is local-only). @@ -9745,6 +11781,7 @@ async def anthropic_messages( message_id, model_name, disable_parallel_tool_use = _disable_parallel, + openai_tools = openai_tools, ) ) @@ -9798,6 +11835,10 @@ async def _anthropic_tool_stream( """Streaming response for the tool-calling path.""" _sentinel = object() + # Gate the display strip on the declared tools: an inactive NAME[ARGS]{...} in a final + # answer is prose and must survive in the delivered text. + _display_names = _display_tool_name_gate(openai_tools) + # Prompt-token count for message_start.usage.input_tokens. count_chat_tokens # makes blocking HTTP calls to llama-server, so run it off the event loop. # Pass the tools so tool-schema tokens are counted (the generator renders @@ -9823,11 +11864,8 @@ async def _anthropic_tool_stream( drop_until_tool_end = False gen = run_gen() - # Concurrent disconnect watcher: the loop only polls is_disconnected() - # between events, so a client disconnect during a long prefill or - # generation step would otherwise hold the decode slot until the next - # event or a failed send. The watcher sets cancel_event so the backend - # stops promptly. + # Watcher to cancel on disconnect: the in-loop poll fires only between + # events, so a mid-prefill disconnect would otherwise hold the decode slot. disconnect_watcher = asyncio.create_task( _await_disconnect_then_cancel(request, cancel_event) ) @@ -9852,9 +11890,15 @@ async def _anthropic_tool_stream( captured_finish_reason = _fr # Strip leaked tool-call XML from content events first, so a # content event that was purely tool XML doesn't count as text. + # Protected helper preserves rehearsal and balanced + # [TOOL_CALLS] trailing prose (raw _TOOL_XML_RE.sub corrupts both). if etype == "content": event = dict(event) - event["text"] = _TOOL_XML_RE.sub("", event["text"]) + event["text"] = _strip_tool_xml_for_display( + event["text"], + auto_heal_tool_calls = True, + enabled_tool_names = _display_names, + ) # disable_parallel_tool_use: keep only the first tool_use block, # dropping every later tool_start and its paired tool_end (robust # to empty tool-call ids — tracked by state, not id matching). @@ -9919,11 +11963,8 @@ async def _anthropic_plain_stream( captured_finish_reason = None gen = run_gen() - # Concurrent disconnect watcher: the loop only polls is_disconnected() - # between chunks, so a client disconnect during a long prefill or - # generation step would otherwise hold the decode slot until the next - # chunk or a failed send. The watcher sets cancel_event so the backend - # stops promptly. + # Watcher to cancel on disconnect: the in-loop poll fires only between + # chunks, so a mid-prefill disconnect would otherwise hold the decode slot. disconnect_watcher = asyncio.create_task( _await_disconnect_then_cancel(request, cancel_event) ) @@ -10012,6 +12053,7 @@ async def _anthropic_tool_non_streaming( message_id, model_name, disable_parallel_tool_use = False, + openai_tools = None, ): """Non-streaming response for the tool-calling path. @@ -10030,6 +12072,9 @@ async def _anthropic_tool_non_streaming( usage = {} prev_text = "" captured_finish_reason = None + # Gate the display strip on the declared tools: an inactive NAME[ARGS]{...} in a final + # answer is prose and must survive in the delivered text. + _display_names = _display_tool_name_gate(openai_tools) # Pending client tool_use; cleared by tool_end (server execution) or # trailing text. See the stop_reason mapping below. ends_on_tool_use = False @@ -10039,8 +12084,10 @@ async def _anthropic_tool_non_streaming( for event in events: etype = event.get("type", "") if etype == "content": - # Strip leaked tool-call XML - clean = _TOOL_XML_RE.sub("", event["text"]) + # Strip leaked tool XML (protected helper keeps think rehearsal and trailing prose). + clean = _strip_tool_xml_for_display( + event["text"], auto_heal_tool_calls = True, enabled_tool_names = _display_names + ) new = clean[len(prev_text) :] prev_text = clean if new: @@ -10164,13 +12211,15 @@ def _build_passthrough_payload( ): body = { "messages": openai_messages, - "tools": openai_tools, - "tool_choice": tool_choice, "temperature": temperature, "top_p": top_p, "top_k": top_k, "stream": stream, } + if openai_tools: + body["tools"] = openai_tools + if tool_choice is not None: + body["tool_choice"] = tool_choice if seed is not None: body["seed"] = seed if stream and stream_options is not None: @@ -10377,13 +12426,17 @@ async def _anthropic_passthrough_stream( yield event return finally: - await _aclose_stream_resources( - watchers = (cancel_watcher, disconnect_watcher), - iterator = lines_iter, - resp = resp, - client = client, - ) - _tracker.__exit__(None, None, None) + # Same shape as the OpenAI passthrough: a close-time CancelledError + # re-raised by _aclose_stream_resources must not skip the tracker exit. + try: + await _aclose_stream_resources( + watchers = (cancel_watcher, disconnect_watcher), + iterator = lines_iter, + resp = resp, + client = client, + ) + finally: + _tracker.__exit__(None, None, None) for line in emitter.finish(): yield line @@ -10510,9 +12563,16 @@ async def _anthropic_passthrough_non_streaming( 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. + # only for opted-out or no-client-tool requests. Protected helper (not + # raw _TOOL_XML_RE.sub): preserves rehearsal and balanced + # [TOOL_CALLS] trailing prose, gated on the declared tools so an + # inactive NAME[ARGS]{...} example in the final text is kept. if not healing_active: - text = _TOOL_XML_RE.sub("", text) + text = _strip_tool_xml_for_display( + text, + auto_heal_tool_calls = True, + enabled_tool_names = _display_tool_name_gate(openai_tools), + ) text = text.strip() if text: content_blocks.append(AnthropicResponseTextBlock(text = text)) @@ -10760,6 +12820,55 @@ def _openai_messages_for_passthrough(payload) -> list[dict]: return messages +def _flatten_content_parts_for_local_template(messages: list[dict]) -> list[dict]: + """Flatten OpenAI content-part lists to plain strings. + + Local text templates take string content and raise on part lists (e.g. a + remote ``image_url`` that leaves ``image is None``): keep the text parts, + drop the rest, like the plain non-GGUF path. GGUF keeps the parts.""" + out = [] + for msg in messages: + content = msg.get("content") + if isinstance(content, list): + text_parts = [ + part.get("text", "") + for part in content + if isinstance(part, dict) and part.get("type") == "text" + ] + msg = {**msg, "content": "\n".join(text_parts) if text_parts else ""} + out.append(msg) + return out + + +def _structured_tool_history_for_local_template(messages: list[dict]) -> list[dict]: + """Deserialize assistant ``tool_calls[].function.arguments`` JSON strings to + mappings for local templating. + + Clients send prior-turn arguments as JSON strings, but local templates take + mappings (some raise on strings). Only the internal messages copy is + rewritten; the HTTP response stays OpenAI-shaped and unparseable strings + are left untouched.""" + out = [] + for msg in messages: + tool_calls = msg.get("tool_calls") + if isinstance(tool_calls, list) and tool_calls: + new_calls = [] + for tc in tool_calls: + fn = tc.get("function") if isinstance(tc, dict) else None + args = fn.get("arguments") if isinstance(fn, dict) else None + if isinstance(args, str): + try: + parsed = json.loads(args) + except ValueError: + parsed = None + if isinstance(parsed, dict): + tc = {**tc, "function": {**fn, "arguments": parsed}} + new_calls.append(tc) + msg = {**msg, "tool_calls": new_calls} + out.append(msg) + return out + + def _openai_messages_for_gguf_chat(payload, is_vision: bool) -> tuple[list[dict], bool]: """Build llama-server messages for the standard GGUF chat path. @@ -10824,6 +12933,9 @@ def _build_openai_passthrough_body( system_prompt, _, _ = _extract_content_parts(payload.messages) messages = _set_or_prepend_system_message(messages, system_prompt) tool_choice = payload.tool_choice if payload.tool_choice is not None else "auto" + tools = payload.tools + if payload.tool_choice == "none" and not _has_openai_tool_history(payload.messages): + tools = None # Forward per-request reasoning fields (enable_thinking / reasoning_effort / # preserve_thinking) via chat_template_kwargs so the Jinja template renders # in the caller's mode, gated on the active template's capabilities exactly @@ -10839,12 +12951,12 @@ def _build_openai_passthrough_body( ) return _build_passthrough_payload( messages, - payload.tools, + tools, payload.temperature, payload.top_p, payload.top_k, # Honor max_completion_tokens on the tools/response_format passthrough too. - _effective_max_tokens(payload), + _effective_openai_max_tokens(payload), payload.stream, stop = payload.stop, min_p = payload.min_p, @@ -10868,29 +12980,219 @@ async def _openai_passthrough_stream( completion_id, monitor_id: Optional[str] = None, ): - """Streaming client-side pass-through for /v1/chat/completions. - - Forwards the client's OpenAI function-calling request to llama-server and - relays the SSE stream back verbatim, preserving llama-server's native - response ``id``, ``finish_reason`` (including ``"tool_calls"``), - ``delta.tool_calls``, and any client-requested trailing ``usage`` chunk so - the client sees a standard OpenAI response. - """ - target_url = f"{llama_backend.base_url}/v1/chat/completions" - body = _build_openai_passthrough_body( - payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend - ) - # 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__() + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _tracker.__exit__(None, None, None) + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "chat_passthrough_stream", + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + + lease = reservation.lease_nowait() + if lease is not None: + try: + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + lease.release() + _tracker.__exit__(None, None, None) + raise + except LlamaAdmissionCancelled as exc: + lease.release() + _tracker.__exit__(None, None, None) + api_monitor.finish(monitor_id, "cancelled") + raise HTTPException( + status_code = 499, + detail = _openai_admission_error_body(exc, status_code = 499), + ) + return await _openai_passthrough_stream_admitted( + request, + cancel_event, + llama_backend, + payload, + model_name, + completion_id, + monitor_id = monitor_id, + admission_lease = lease, + tracker = _tracker, + ) + + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "chat_passthrough_stream", + completion_id = completion_id, + level = "debug", + ) + + async def _queued_stream(): + admitted_started = False + admitted_body_owns_cleanup = False + admitted_response = None + admitted_body_cancelled = False + try: + async for wait_item in _openai_admission_wait_stream_chunks( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ): + if isinstance(wait_item, str): + yield wait_item + continue + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "chat_passthrough_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + admitted_response = await _openai_passthrough_stream_admitted( + request, + cancel_event, + llama_backend, + payload, + model_name, + completion_id, + monitor_id = monitor_id, + admission_lease = wait_item, + tracker = _tracker, + ) + admitted_started = True + iterator = admitted_response.body_iterator + admitted_body_owns_cleanup = True + try: + async for chunk in iterator: + yield chunk + except asyncio.CancelledError: + admitted_body_cancelled = True + raise + finally: + await _close_openai_admitted_stream_iterator( + iterator, + cancelled = admitted_body_cancelled, + ) + if not admitted_body_owns_cleanup: + cleanup = getattr(admitted_response, "_unstarted_cleanup", None) + if cleanup is not None: + await cleanup() + return + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "chat_passthrough_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + yield _openai_stream_error_sse(_openai_admission_error_body(exc, status_code = 503)) + except LlamaAdmissionCancelled: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "chat_passthrough_stream", + wait_started_at = admission_wait_started_at, + completion_id = completion_id, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + return + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise + except HTTPException as exc: + status_code = getattr(exc, "status_code", 500) or 500 + detail = exc.detail + error = ( + detail + if isinstance(detail, dict) and "error" in detail + else openai_error_body(str(detail), status = status_code) + ) + api_monitor.fail(monitor_id, str(detail)) + yield _openai_stream_error_sse(error) + finally: + if not admitted_started: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + _tracker.__exit__(None, None, None) + + async def _queued_unstarted_cleanup() -> None: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + _tracker.__exit__(None, None, None) + + return _SameTaskStreamingResponse( + _queued_stream(), + media_type = "text/event-stream", + headers = { + "Cache-Control": "no-cache", + "Connection": "close", + "X-Accel-Buffering": "no", + }, + unstarted_cleanup = _queued_unstarted_cleanup, + ) + + +async def _openai_passthrough_stream_admitted( + request, + cancel_event, + llama_backend, + payload, + model_name, + completion_id, + monitor_id: Optional[str] = None, + *, + admission_lease: LlamaAdmissionLease, + tracker, +): + """Streaming client-side pass-through after Studio granted an upstream slot. + + Forwards the client's OpenAI function-calling request to llama-server and + relays the SSE stream back with minimal normalization (reasoning-only + deltas gain ``content: ""``; errors and missing terminal markers get a + closing ``[DONE]``), preserving llama-server's native response ``id``, + ``finish_reason`` (including ``"tool_calls"``), ``delta.tool_calls``, and + any client-requested trailing ``usage`` chunk so the client sees a + standard OpenAI response. + + Reasoning/tool-call splitting is delegated to llama-server (``--jinja + --reasoning-format auto``), so ``delta.content`` carries no raw markup and is + deliberately not re-parsed locally, unlike the ``/completion`` paths. + """ + _tracker = tracker + target_url = f"{llama_backend.base_url}/v1/chat/completions" + upstream_headers = _openai_passthrough_upstream_headers(llama_backend = llama_backend) + client = None resp = None send_task: Optional[asyncio.Task[Optional[httpx.Response]]] = None @@ -10912,6 +13214,17 @@ async def _openai_passthrough_stream( # Keep tracker cleanup paired if pre-header dispatch is cancelled. try: + body = _build_openai_passthrough_body( + payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend + ) + # Text-form tool calls from small models get promoted to structured calls on + # the way back (declared client tools only); requests without tools or with + # auto_heal_tool_calls=false keep the unhealed relay. tool_choice constrains + # the allowlist ("none" disables, a forced function narrows to it). + _allowed_tools = heal_gate( + payload.auto_heal_tool_calls, body.get("tools"), body.get("tool_choice") + ) + # Keep the pre-header window short so accepted SSE clients receive # immediate headers in the common timeout-reduced stall. client = httpx.AsyncClient( @@ -10925,12 +13238,16 @@ async def _openai_passthrough_stream( while True: try: - req = client.build_request( - "POST", target_url, json = body, headers = {"Connection": "close"} - ) + req = client.build_request("POST", target_url, json = body, headers = upstream_headers) first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S send_task = asyncio.create_task( - _send_stream_with_preheader_cancel(client, req, cancel_event, request = request) + _send_stream_with_preheader_cancel( + client, + req, + cancel_event, + request = request, + mark_cancel_on_cancel = False, + ) ) done, _ = await asyncio.wait( {send_task}, @@ -10956,13 +13273,17 @@ async def _openai_passthrough_stream( if resp is None and send_task is not None and not send_task.done(): break if resp is None: + if cancel_event is not None: + cancel_event.set() api_monitor.finish(monitor_id, "cancelled") - await _aclose_send_task(send_task) try: - await client.aclose() - except Exception: - pass - _tracker.__exit__(None, None, None) + await _aclose_send_task(send_task) + await _aclose_stream_resources(client = client) + finally: + try: + admission_lease.release() + finally: + _tracker.__exit__(None, None, None) return _SameTaskStreamingResponse( iter(()), media_type = "text/event-stream", @@ -10987,6 +13308,7 @@ async def _openai_passthrough_stream( await resp.aclose() except Exception: pass + resp = None # Opt-in overflow policy: shrink and retry instead of a fatal 400. if ( _truncate_budget > 0 @@ -11015,11 +13337,14 @@ async def _openai_passthrough_stream( disconnect_watcher = None nonlocal resp, send_task, first_token_deadline, _truncate_budget + nonlocal client monitor_done = False saw_finish_reason = False saw_done = False saw_stream_error = False + saw_stream_item = False saw_tool_call_delta = False + terminal_seen = False last_chunk_id = completion_id last_chunk_model = model_name last_chunk_created = int(time.time()) @@ -11080,6 +13405,13 @@ async def _openai_passthrough_stream( lines.append("data: " + json.dumps(chunk, ensure_ascii = False)) return lines + stall_timeout_s = _openai_compat_stream_stall_timeout() + + def _terminal_read_timeout_s() -> Optional[float]: + if terminal_seen: + return _OPENAI_PASSTHROUGH_TERMINAL_GRACE_S + return stall_timeout_s + def _heal_transform(chunk_data: dict, raw_line: str) -> list: """SSE lines to emit in place of one upstream line (healing on).""" choices = chunk_data.get("choices") @@ -11148,24 +13480,47 @@ async def _openai_passthrough_stream( 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 send_task is not None: + last_keepalive_at = time.monotonic() + while not send_task.done(): + # Wake often enough that _preheader_cancelled keeps + # cancel/disconnect latency sub-second during prefill; + # keepalives still pace off last_keepalive_at. + wait_timeout = min( + _STREAM_DISCONNECT_POLL_TIMEOUT_S, + _OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S, + ) + done, _ = await asyncio.wait( + {send_task}, + timeout = wait_timeout, + return_when = asyncio.FIRST_COMPLETED, + ) + if send_task in done: + break + if await _preheader_cancelled(cancel_event, request): + api_monitor.finish(monitor_id, "cancelled") + return + # The downstream SSE response is already committed; + # keep strict clients and proxies from treating a long + # llama-server prefill/header wait as a dead stream. + now = time.monotonic() + if ( + now - last_keepalive_at + >= _OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S + ): + last_keepalive_at = now + yield _OPENAI_PASSTHROUGH_SSE_KEEPALIVE + if resp is None: + try: + resp = send_task.result() + except httpx.RequestError as e: + logger.error( + "openai passthrough stream: upstream unreachable: %s", e + ) + api_monitor.fail(monitor_id, _friendly_error(e)) + yield _openai_stream_error_sse(_openai_stream_error_chunk(e)) + return + send_task = None if resp is None: api_monitor.finish(monitor_id, "cancelled") @@ -11193,12 +13548,16 @@ async def _openai_passthrough_stream( ): _truncate_budget -= 1 req = client.build_request( - "POST", target_url, json = body, headers = {"Connection": "close"} + "POST", target_url, json = body, headers = upstream_headers ) first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S send_task = asyncio.create_task( _send_stream_with_preheader_cancel( - client, req, cancel_event, request = request + client, + req, + cancel_event, + request = request, + mark_cancel_on_cancel = False, ) ) continue @@ -11213,7 +13572,7 @@ async def _openai_passthrough_stream( ) ) api_monitor.fail(monitor_id, err_text[:500]) - yield f"data: {json.dumps(error_payload)}\n\n" + yield _openai_stream_error_sse(error_payload) return cancel_watcher = asyncio.create_task(_await_cancel_then_close(cancel_event, resp)) @@ -11227,12 +13586,14 @@ async def _openai_passthrough_stream( request = request, first_token_deadline = first_token_deadline, response = resp, + post_first_item_read_timeout_s = _terminal_read_timeout_s, ): if not raw_line: continue - if not raw_line.startswith("data: "): + if not raw_line.startswith("data:"): continue - data_text = raw_line[6:].strip() + saw_stream_item = True + data_text = raw_line[5:].strip() if data_text == "[DONE]": saw_done = True # Upstream ended without a finish chunk: heal the residue @@ -11264,13 +13625,11 @@ async def _openai_passthrough_stream( yield raw_line + "\n\n" monitor_done = True break - # Honor parallel_tool_calls=false (best-effort): drop tool_call - # deltas with index>=1 so only the first call streams. Only - # lines carrying tool_calls are reparsed; everything else is - # relayed byte-for-byte. - if payload.parallel_tool_calls is False and '"tool_calls"' in raw_line: - raw_line = _cap_parallel_tool_calls_sse_line(raw_line) - data_text = raw_line[6:].strip() + raw_line = _normalize_openai_passthrough_sse_line( + raw_line, + cap_parallel_tool_calls = payload.parallel_tool_calls is False, + ) + data_text = raw_line[5:].strip() try: chunk_data = json.loads(data_text) except json.JSONDecodeError: @@ -11291,16 +13650,15 @@ async def _openai_passthrough_stream( delta = choice.get("delta") if isinstance(delta, dict) and delta.get("tool_calls"): saw_tool_call_delta = True - # Detect an upstream error chunk independently of API - # monitoring: when monitor_id is None (skip_api_monitor), - # _monitor_openai_sse_line returns before inspecting the - # error, so without this the synthetic-finish guard would - # emit a successful finish_reason after a failed stream. + # Detect an error chunk independently of API monitoring + # (skip_api_monitor returns early), else the synthetic + # finish would fire after a failed stream. if _monitor_openai_error_message(chunk_data): saw_stream_error = True # With healing active, a content-bearing line may be replaced by - # held/promoted chunks; otherwise the single upstream line - # relays verbatim (monitored exactly as emitted either way). + # held/promoted chunks; otherwise the single (already + # normalized) line relays unchanged (monitored exactly as + # emitted either way). if ( healer is not None and not healer.dormant @@ -11350,6 +13708,27 @@ async def _openai_passthrough_stream( yield out_line + "\n\n" if monitor_event == "done": monitor_done = True + break + terminal_state = ( + _openai_passthrough_terminal_state_from_data(chunk_data) + if out_line is raw_line + else _openai_passthrough_sse_line_terminal_state(out_line) + ) + if terminal_state == "usage" or ( + terminal_state == "finish" and not _wants_stream_usage(payload) + ): + done_line = _SSE_DONE_LINE + _monitor_openai_sse_line( + monitor_id, + done_line, + llama_backend.context_length, + ) + yield done_line + "\n\n" + saw_done = True + monitor_done = True + break + if terminal_state == "finish": + terminal_seen = True if monitor_done: break if not saw_done and not saw_stream_error and not cancel_event.is_set(): @@ -11371,7 +13750,7 @@ async def _openai_passthrough_stream( llama_backend.context_length, ) yield finish_line + "\n\n" - done_line = "data: [DONE]" + done_line = _SSE_DONE_LINE _monitor_openai_sse_line( monitor_id, done_line, @@ -11387,6 +13766,29 @@ async def _openai_passthrough_stream( except asyncio.CancelledError: api_monitor.finish(monitor_id, "cancelled") raise + except httpx.ReadTimeout as e: + if terminal_seen and not saw_stream_error and not cancel_event.is_set(): + done_line = _SSE_DONE_LINE + _monitor_openai_sse_line( + monitor_id, + done_line, + llama_backend.context_length, + ) + yield done_line + "\n\n" + api_monitor.finish(monitor_id) + return + if cancel_event.is_set(): + api_monitor.finish(monitor_id, "cancelled") + return + logger.error( + "openai passthrough stream %s: %s", + "stalled mid-response" if saw_stream_item else "timeout", + e, + ) + api_monitor.fail(monitor_id, _friendly_error(e)) + get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) + err = _openai_stream_error_chunk(e) + yield _openai_stream_error_sse(err) except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e: # Watcher closed resp on cancel. Emit nothing extra; the client # initiated the cancel or already disconnected. @@ -11395,6 +13797,16 @@ async def _openai_passthrough_stream( get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) raise api_monitor.finish(monitor_id, "cancelled") + except HTTPException as exc: + status_code = getattr(exc, "status_code", 500) or 500 + detail = exc.detail + error_payload = ( + detail + if isinstance(detail, dict) and "error" in detail + else openai_error_body(str(detail), status = status_code) + ) + api_monitor.fail(monitor_id, str(detail)) + yield _openai_stream_error_sse(error_payload) except Exception as e: if cancel_event.is_set(): api_monitor.finish(monitor_id, "cancelled") @@ -11404,25 +13816,37 @@ async def _openai_passthrough_stream( api_monitor.fail(monitor_id, _friendly_error(e)) get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) err = _openai_stream_error_chunk(e) - yield f"data: {json.dumps(err)}\n\n" + yield _openai_stream_error_sse(err) finally: - await _aclose_send_task(send_task) - await _aclose_stream_resources( - watchers = (cancel_watcher, disconnect_watcher), - iterator = lines_iter, - resp = resp, - client = client, - ) - _tracker.__exit__(None, None, None) + # _aclose_stream_resources re-raises a close-time CancelledError + # only after finishing teardown, and the tracker exits either way. + try: + await _aclose_send_task(send_task) + await _aclose_stream_resources( + watchers = (cancel_watcher, disconnect_watcher), + iterator = lines_iter, + resp = resp, + client = client, + ) + finally: + try: + admission_lease.release() + finally: + _tracker.__exit__(None, None, None) async def _unstarted_cleanup() -> None: # Client disconnected before the body stream started, so _stream()'s # finally never ran. Release the eagerly-opened upstream resp/client # and the cancel-registry entry here; the watchers and line iterator # are created inside _stream(), so there is nothing else to close. - await _aclose_send_task(send_task) - await _aclose_stream_resources(resp = resp, client = client) - _tracker.__exit__(None, None, None) + try: + await _aclose_send_task(send_task) + await _aclose_stream_resources(resp = resp, client = client) + finally: + try: + admission_lease.release() + finally: + _tracker.__exit__(None, None, None) return _SameTaskStreamingResponse( _stream(), @@ -11434,10 +13858,22 @@ 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) + except BaseException as exc: + if isinstance(exc, asyncio.CancelledError): + if cancel_event is not None: + cancel_event.set() + api_monitor.finish(monitor_id, "cancelled") + else: + detail = exc.detail if isinstance(exc, HTTPException) else _friendly_error(exc) + api_monitor.fail(monitor_id, str(detail)) + try: + await _aclose_send_task(send_task) + await _aclose_stream_resources(resp = resp, client = client) + finally: + try: + admission_lease.release() + finally: + _tracker.__exit__(None, None, None) raise @@ -11446,6 +13882,109 @@ async def _openai_passthrough_non_streaming( payload, model_name, monitor_id: Optional[str] = None, + *, + request: Optional[Request] = None, + cancel_event = None, +): + """Non-streaming pass-through guarded by local llama-server admission.""" + try: + reservation, admission_config = _openai_llama_admission_reserve( + request = request, + llama_backend = llama_backend, + ) + except LlamaAdmissionQueueFull as exc: + _openai_admission_log( + "queue-full", + snapshot = exc.snapshot, + request = request, + mode = "chat_passthrough_nonstream", + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 429) + + lease = None + admission_wait_started_at = None + try: + if reservation.lease_nowait() is None: + admission_wait_started_at = time.monotonic() + _openai_admission_log( + "queued", + reservation, + request = request, + mode = "chat_passthrough_nonstream", + level = "debug", + ) + lease = await _wait_for_openai_admission_non_streaming( + reservation, + admission_config, + request = request, + cancel_event = cancel_event, + ) + if admission_wait_started_at is not None: + _openai_admission_log( + "granted-after-wait", + reservation, + request = request, + mode = "chat_passthrough_nonstream", + wait_started_at = admission_wait_started_at, + level = "debug", + ) + await _raise_if_openai_admission_cancelled( + reservation, + request = request, + cancel_event = cancel_event, + ) + return await _openai_passthrough_non_streaming_upstream( + llama_backend, + payload, + model_name, + monitor_id = monitor_id, + request = request, + cancel_event = cancel_event, + ) + except LlamaAdmissionTimeout as exc: + _openai_admission_log( + "timeout", + reservation, + request = request, + mode = "chat_passthrough_nonstream", + wait_started_at = admission_wait_started_at, + level = "warning", + ) + api_monitor.fail(monitor_id, str(exc)) + raise _openai_admission_http_exception(exc, status_code = 503) + except LlamaAdmissionCancelled as exc: + _openai_admission_log( + "cancelled-before-upstream", + reservation, + request = request, + mode = "chat_passthrough_nonstream", + wait_started_at = admission_wait_started_at, + level = "debug", + ) + api_monitor.finish(monitor_id, "cancelled") + raise HTTPException( + status_code = 499, + detail = _openai_admission_error_body(exc, status_code = 499), + ) + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + reservation.cancel() + raise + finally: + if lease is not None: + lease.release() + + +async def _openai_passthrough_non_streaming_upstream( + llama_backend, + payload, + model_name, + monitor_id: Optional[str] = None, + *, + request: Optional[Request] = None, + cancel_event = None, ): """Non-streaming client-side pass-through for /v1/chat/completions. @@ -11454,20 +13993,67 @@ async def _openai_passthrough_non_streaming( ``tool_calls``, and accurate ``usage`` token counts. """ target_url = f"{llama_backend.base_url}/v1/chat/completions" + upstream_headers = _openai_passthrough_upstream_headers(llama_backend = llama_backend) body = _build_openai_passthrough_body( payload, backend_ctx = llama_backend.context_length, llama_backend = llama_backend ) + body["stream"] = False + body.pop("stream_options", None) _truncate_budget = ( _OVERFLOW_TRUNCATE_MAX_RETRIES if _overflow_truncation_requested(payload) else 0 ) - while True: - try: - resp = await nonstreaming_client().post( + + async def _post(body_to_send): + if cancel_event is None and request is None: + return await nonstreaming_client().post( target_url, - json = body, + json = body_to_send, + headers = upstream_headers, timeout = _llama_non_streaming_generation_timeout(), ) + + if cancel_event is None: + cancel = threading.Event() + else: + cancel = cancel_event + client = _cancelable_nonstreaming_client() + watcher = asyncio.create_task( + _await_cancel_or_disconnect_then_close_client( + cancel_event = cancel, + request = request, + client = client, + ) + ) + try: + try: + response = await client.post( + target_url, + json = body_to_send, + headers = upstream_headers, + timeout = _llama_non_streaming_generation_timeout(), + ) + except httpx.RequestError: + if cancel.is_set(): + raise asyncio.CancelledError() + raise + if cancel.is_set(): + raise asyncio.CancelledError() + return response + finally: + watcher.cancel() + try: + await watcher + except (asyncio.CancelledError, Exception): + pass + try: + await client.aclose() + except Exception: + pass + + while True: + try: + resp = await _post(body) except asyncio.CancelledError: api_monitor.finish(monitor_id, "cancelled") raise @@ -11535,15 +14121,14 @@ async def _openai_passthrough_non_streaming( "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(), - ) + retry_resp = await _post(retry_body) if retry_resp.status_code == 200: retry_data = retry_resp.json() if response_has_promotable_calls(retry_data, _allowed_tools, body.get("tools")): resp, data = retry_resp, retry_data + except asyncio.CancelledError: + api_monitor.finish(monitor_id, "cancelled") + raise except (httpx.RequestError, ValueError) as exc: logger.warning("tool-call nudge retry failed; keeping original: %s", exc) diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 862bce8be8..bbee374334 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -32,6 +32,7 @@ from utils.helper_precache_settings import ( helper_model_disabled_by_env, set_helper_precache_enabled, ) +from utils.coding_agents import CODING_AGENTS, detect_installed_coding_agents from utils.openai_auto_switch_settings import ( DEFAULT_AUTO_UNLOAD_IDLE_SECONDS, DEFAULT_OPENAI_AUTO_SWITCH_ENABLED, @@ -174,6 +175,19 @@ def update_helper_precache( return _helper_precache_response(enabled) +class CodingAgentsResponse(BaseModel): + # All agents `unsloth start` supports, in the CLI's declared order. + agents: tuple[str, ...] = CODING_AGENTS + # Subset of `agents` whose CLI binary was found on PATH; the frontend uses + # this to default the API-keys panel to a command the user can run as-is. + detected: list[str] + + +@router.get("/coding-agents", response_model = CodingAgentsResponse) +def get_coding_agents(current_subject: str = Depends(get_current_subject)) -> CodingAgentsResponse: + return CodingAgentsResponse(detected = detect_installed_coding_agents()) + + @router.get("/openai-auto-switch", response_model = OpenAIAutoSwitchResponse) def get_openai_auto_switch( current_subject: str = Depends(get_current_subject), @@ -260,17 +274,29 @@ def _embedding_model_response() -> EmbeddingModelResponse: ) -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 - +def _ambient_hf_token() -> Optional[str]: + """The HF token the loader would use (HF_TOKEN env or the cached login), so a gated + repo is scanned rather than failing open. None if unavailable.""" try: - raw = (rag_config.EMBED_BACKEND or "auto").strip().lower() - key = embeddings._resolve_auto() if raw in embeddings._AUTO_ALIASES else raw + from huggingface_hub import get_token + return get_token() + except Exception: + return None + + +def _llama_backend_active() -> bool: + """True when this install actually embeds via the llama-server (GGUF) backend. + + Delegates to the embeddings module so a runtime fallback from + sentence-transformers to llama-server (after a torch/CUDA load or encode + failure) is honored: in that state the process loads only inert GGUF, so the + ST pickle gate below must not hard-block a repo whose GGUF companion is clean. + Before any backend is built this still reflects the resolver.""" + from core.rag import embeddings + try: + return embeddings.active_backend_is_llama() except Exception: # noqa: BLE001 - backend probe must never block saving return False - return key in embeddings._LLAMA_ALIASES def _resolves_as_local_gguf(model: str) -> bool: @@ -357,6 +383,8 @@ def update_embedding_model( """Set the RAG embedding model. Unless ``force`` is set, the repo is verified to be an embedding model via HF metadata; an unverifiable model (wrong type, typo, gated repo, or no network) returns 409 so the UI can offer "save anyway". + A repo flagged unsafe by HF's security scan returns 403 instead: a hard block + that ``force`` cannot bypass, so the UI must not offer "save anyway". Documents indexed under the previous model must be re-uploaded.""" from utils.models import is_embedding_model @@ -370,15 +398,51 @@ def update_embedding_model( event = "settings.update_embedding_model_failed", log = logger, ) from exc + hf_token = (payload.hf_token or "").strip() or None # The env/default model needs no verification; saving it is a no-op override. # A local GGUF on the llama-server backend is accepted as-is: it is exactly # what the backend loads, and HF metadata cannot verify a local path. - 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 + is_local_gguf = _llama_backend_active() and _resolves_as_local_gguf(model) + # The pickle gate only matters for the sentence-transformers backend, which is what + # deserializes pickles. On the llama-server backend the embedder loads GGUF files + # (inert) from effective_gguf_repo(), so scanning the ST repo's pickle here would + # wrongly reject a custom repo whose GGUF companion is clean; the GGUF availability + # checks below cover that path instead. + scan_st_pickle = ( + model != default_embedding_model() and not is_local_gguf and not _llama_backend_active() + ) + if scan_st_pickle: + # Malware/pickle gate before we persist a repo the embedder later loads with + # SentenceTransformer. Runs even under force (force only skips the is-embedding + # type check for offline/local repos HF cannot verify); local paths and + # unreachable scans fail open inside evaluate_file_security. + from utils.security import evaluate_file_security, security_load_subdirs + from core.rag.embeddings import _st_module_subdirs + + # Fall back to the loader's own token so a gated/private repo is actually scanned + # (a token-less scan fails open for exactly the repo that would still load). + scan_token = hf_token or _ambient_hf_token() + # Include the ST module dirs (0_Transformer/) so a flagged pickle directly under + # one blocks instead of passing as an unreferenced nested shard. + load_subdirs = tuple( + dict.fromkeys( + ( + *security_load_subdirs(model, scan_token), + *_st_module_subdirs(model, scan_token), + ) + ) + ) + if evaluate_file_security(model, hf_token = scan_token, load_subdirs = load_subdirs).blocked: + # 403, not 409: the client routes every 409 into the forceable "save anyway" + # flow, but this block is a hard, non-forceable security refusal. + raise HTTPException( + status_code = 403, + detail = ( + f"{model!r} is flagged as unsafe by Hugging Face's security scan and " + "cannot be used as the embedding model." + ), + ) + if model != default_embedding_model() and not payload.force and not is_local_gguf: from core.rag import config as rag_config # A GGUF-named repo on the llama-server backend is loaded from its .gguf diff --git a/studio/backend/run.py b/studio/backend/run.py index ccd0113972..2cc6c4a93e 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -12,6 +12,79 @@ import time from pathlib import Path from typing import Optional + +def _fix_torch_cuda_ld_path(): + """Prepend torch's bundled CUDA libs to LD_LIBRARY_PATH. + + PyTorch wheels ship their own CUDA runtime (libcudart, libcublas, ...) in + ``site-packages/nvidia/*/lib``. On Linux the dynamic linker reads + LD_LIBRARY_PATH before the RUNPATH baked into torch's .so files, so a + pre-existing LD_LIBRARY_PATH pointing at a different system CUDA (e.g. + /usr/local/cuda-13/lib64 from conda or a Docker base image) shadows torch's + libs and triggers "undefined symbol" errors when torch is imported. Detect + torch's lib dirs (without importing torch) and prepend them. Returns True if + LD_LIBRARY_PATH was changed. + """ + if sys.platform != "linux": + return False + ld_path = os.environ.get("LD_LIBRARY_PATH", "") + if not ld_path: + return False + try: + import importlib.util + + spec = importlib.util.find_spec("torch") + if not spec or not spec.origin: + return False + torch_dir = os.path.dirname(spec.origin) + site_pkgs = os.path.dirname(torch_dir) + nvidia_dir = os.path.join(site_pkgs, "nvidia") + + lib_dirs = [] + torch_lib = os.path.join(torch_dir, "lib") + if os.path.isdir(torch_lib): + lib_dirs.append(torch_lib) + if os.path.isdir(nvidia_dir): + for sub in sorted(os.listdir(nvidia_dir)): + lib = os.path.join(nvidia_dir, sub, "lib") + if os.path.isdir(lib): + lib_dirs.append(lib) + if not lib_dirs: + return False + + existing = ld_path.split(":") + if existing[: len(lib_dirs)] == lib_dirs: + return False # already at the front, nothing to do + + torch_set = set(lib_dirs) + cleaned = [p for p in existing if p not in torch_set] + os.environ["LD_LIBRARY_PATH"] = ":".join(lib_dirs + cleaned) + return True + except Exception: + return False + + +_LD_FIXED_SENTINEL = "_UNSLOTH_STUDIO_LD_FIXED" + + +def _maybe_reexec_for_cuda_ld_path(): + """Re-exec once so the dynamic linker sees the corrected LD_LIBRARY_PATH. + + LD_LIBRARY_PATH is read at process start, so editing os.environ in-process + cannot fix the running interpreter; a single re-exec is required. Call only + from a true entry point (the ``if __name__ == "__main__"`` block), never at + import time, because os.execv replaces the whole process (an embedder such + as Colab that does ``from run import run_server`` must not be re-exec'd). + """ + if _LD_FIXED_SENTINEL in os.environ: + return + if not _fix_torch_cuda_ld_path(): + return + os.environ[_LD_FIXED_SENTINEL] = "1" + argv = getattr(sys, "orig_argv", None) or [sys.executable, *sys.argv] + os.execv(sys.executable, argv) + + # Suppress C-level dependency warnings globally (e.g. SwigPyPacked). os.environ["PYTHONWARNINGS"] = "ignore" @@ -1457,6 +1530,12 @@ def _build_arg_parser(): # For direct execution (also invoked by CLI via os.execvp / subprocess). if __name__ == "__main__": + # Correct a conflicting system CUDA on LD_LIBRARY_PATH before torch is + # imported (below, via run_server). Re-execs once on Linux so the dynamic + # linker uses torch's bundled CUDA libs; no-op on other platforms, when + # LD_LIBRARY_PATH is unset or already correct, or after the single re-exec. + _maybe_reexec_for_cuda_ld_path() + import signal import traceback diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index 41a9adcc29..87aa50ee26 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -240,6 +240,7 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: project_id TEXT, archived INTEGER NOT NULL DEFAULT 0, created_at INTEGER NOT NULL, + updated_at INTEGER, openai_code_exec_container_id TEXT, anthropic_code_exec_container_id TEXT, forked_from_thread_id TEXT, @@ -261,6 +262,24 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_thread_id TEXT") if "forked_from_message_id" not in chat_thread_cols: conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_message_id TEXT") + if "updated_at" not in chat_thread_cols: + conn.execute("ALTER TABLE chat_threads ADD COLUMN updated_at INTEGER") + # Floor at created_at: forked threads copy older ancestor messages, + # so the fork's creation time must win over the branch message times. + conn.execute( + """ + UPDATE chat_threads SET updated_at = MAX( + COALESCE( + ( + SELECT MAX(m.created_at) FROM chat_messages m + WHERE m.thread_id = chat_threads.id + ), + created_at + ), + created_at + ) + """ + ) conn.execute( """ CREATE TABLE IF NOT EXISTS chat_messages ( @@ -992,6 +1011,9 @@ def _chat_thread_from_row(row: sqlite3.Row) -> dict: "projectId": data.get("project_id") or None, "archived": bool(data["archived"]), "createdAt": data["created_at"], + "updatedAt": data.get("updated_at") + if data.get("updated_at") is not None + else data["created_at"], "openaiCodeExecContainerId": data.get("openai_code_exec_container_id"), "anthropicCodeExecContainerId": data.get("anthropic_code_exec_container_id"), "forkedFromThreadId": data.get("forked_from_thread_id"), @@ -1039,8 +1061,8 @@ def upsert_chat_thread(thread: dict) -> dict: conn.execute( """ INSERT INTO chat_threads - (id, title, model_type, model_id, pair_id, project_id, archived, created_at, openai_code_exec_container_id, anthropic_code_exec_container_id, forked_from_thread_id, forked_from_message_id) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + (id, title, model_type, model_id, pair_id, project_id, archived, created_at, updated_at, openai_code_exec_container_id, anthropic_code_exec_container_id, forked_from_thread_id, forked_from_message_id) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET title = excluded.title, model_type = excluded.model_type, @@ -1049,6 +1071,7 @@ def upsert_chat_thread(thread: dict) -> dict: project_id = excluded.project_id, archived = excluded.archived, created_at = excluded.created_at, + updated_at = COALESCE(excluded.updated_at, chat_threads.updated_at), openai_code_exec_container_id = excluded.openai_code_exec_container_id, anthropic_code_exec_container_id = excluded.anthropic_code_exec_container_id, forked_from_thread_id = excluded.forked_from_thread_id, @@ -1063,6 +1086,7 @@ def upsert_chat_thread(thread: dict) -> dict: thread.get("projectId"), 1 if thread.get("archived") else 0, int(thread["createdAt"]), + int(thread["updatedAt"]) if thread.get("updatedAt") is not None else None, thread.get("openaiCodeExecContainerId"), thread.get("anthropicCodeExecContainerId"), thread.get("forkedFromThreadId"), @@ -1084,6 +1108,7 @@ def update_chat_thread(id: str, patch: dict) -> Optional[dict]: "projectId": ("project_id", patch.get("projectId")), "archived": ("archived", 1 if patch.get("archived") else 0), "createdAt": ("created_at", patch.get("createdAt")), + "updatedAt": ("updated_at", patch.get("updatedAt")), "openaiCodeExecContainerId": ( "openai_code_exec_container_id", patch.get("openaiCodeExecContainerId"), @@ -1155,7 +1180,8 @@ def list_chat_threads( conn = get_connection() try: rows = conn.execute( - f"SELECT * FROM chat_threads {where} ORDER BY created_at DESC", + f"SELECT * FROM chat_threads {where} " + "ORDER BY COALESCE(updated_at, created_at) DESC, created_at DESC", values, ).fetchall() return [_chat_thread_from_row(row) for row in rows] @@ -1394,6 +1420,44 @@ def _raise_if_chat_message_thread_conflicts( ) +def _bump_chat_thread_updated_at( + conn: sqlite3.Connection, thread_id: str, message_created_at: int +) -> None: + conn.execute( + """ + UPDATE chat_threads + SET updated_at = MAX(COALESCE(updated_at, created_at), ?) + WHERE id = ? + """, + (message_created_at, thread_id), + ) + + +def _recompute_chat_thread_updated_at(conn: sqlite3.Connection, thread_id: str) -> None: + """Set updated_at from the remaining messages, floored at created_at. + + Unlike the ratchet-only bump, this can lower updated_at -- needed after + pruning, which may delete the thread's newest message. + """ + conn.execute( + """ + UPDATE chat_threads + SET updated_at = MAX( + COALESCE( + ( + SELECT MAX(m.created_at) FROM chat_messages m + WHERE m.thread_id = chat_threads.id + ), + created_at + ), + created_at + ) + WHERE id = ? + """, + (thread_id,), + ) + + def upsert_chat_message(message: dict) -> dict: conn = get_connection() try: @@ -1432,6 +1496,7 @@ def upsert_chat_message(message: dict) -> dict: int(message["createdAt"]), ), ) + _bump_chat_thread_updated_at(conn, message["threadId"], int(message["createdAt"])) conn.commit() return message except Exception: @@ -1484,6 +1549,12 @@ def sync_chat_messages( for m in messages ], ) + if prune_missing: + _recompute_chat_thread_updated_at(conn, thread_id) + elif messages: + _bump_chat_thread_updated_at( + conn, thread_id, max(int(m["createdAt"]) for m in messages) + ) conn.commit() return list_chat_messages(thread_id) except ChatMessageConflictError: diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index a6c1fcda9c..54454f6563 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -52,6 +52,49 @@ from io import BytesIO as _BytesIO from types import SimpleNamespace +def _emitter_client_text(events: list[str]) -> str: + """Concatenate the text_delta payloads an SSE event list carries.""" + text = "" + for line in events: + for raw in line.split("\n"): + raw = raw.strip() + if not raw.startswith("data: "): + continue + data = json.loads(raw[len("data: ") :]) + delta = data.get("delta", {}) + if delta.get("type") == "text_delta": + text += delta.get("text", "") + return text + + +def test_anthropic_emitter_closes_reasoning_only_think_block(): + # A reasoning-only reply streams X live then shrinks to bare X at EOF. + # This emitter diffs cumulative snapshots and drops the shrink, so without a + # closing pass the client text would end on an unclosed . finish() + # must balance it. + emitter = AnthropicStreamEmitter() + events = emitter.start("msg_1", "m") + events += emitter.feed({"type": "content", "text": "The capital"}) + events += emitter.feed({"type": "content", "text": "The capital of France is Paris."}) + # The generator's final bare-text shrink (dropped by the cumulative diff). + events += emitter.feed({"type": "content", "text": "The capital of France is Paris."}) + events += emitter.finish() + + assert _emitter_client_text(events) == "The capital of France is Paris." + + +def test_anthropic_emitter_does_not_double_close_balanced_think(): + # A reasoning-then-answer reply already closes its own ; the balancer + # must not append a second one. + emitter = AnthropicStreamEmitter() + events = emitter.start("msg_1", "m") + events += emitter.feed({"type": "content", "text": "Thinking."}) + events += emitter.feed({"type": "content", "text": "Thinking.Answer."}) + events += emitter.finish() + + assert _emitter_client_text(events) == "Thinking.Answer." + + def test_streamed_anthropic_tool_use_records_api_monitor_reply(monkeypatch): import routes.inference as inf_mod @@ -889,6 +932,24 @@ class TestAnthropicToolNonStreaming: assert tool_blocks[0]["name"] == "render_html" assert tool_blocks[0]["input"] == {"code": ""} + def test_display_strip_gates_on_declared_tools(self): + # A final answer containing NAME[ARGS]{json} is gated on the declared tools: undeclared + # ``foo`` markup is prose and survives, the declared web_search rehearsal strips. + def _run_gen(): + yield { + "type": "content", + "text": 'Try foo[ARGS]{"x": 1} but not web_search[ARGS]{"q": "hi"} here.', + } + + tools = [{"type": "function", "function": {"name": "web_search", "parameters": {}}}] + response = asyncio.run( + _anthropic_tool_non_streaming(_run_gen, "msg_1", "m", openai_tools = tools) + ) + body = json.loads(response.body) + text = "".join(b["text"] for b in body["content"] if b["type"] == "text") + assert 'foo[ARGS]{"x": 1}' in text # inactive name preserved as prose + assert "web_search[ARGS]" not in text # active name stripped from display + # ===================================================================== # Pass-through emitter tests (client-side tool execution path) @@ -1709,3 +1770,187 @@ class TestAnthropicMessagesToolRouting: _drive(anthropic_messages(payload, request = None, current_subject = "t")) assert backend.calls[0][0] == "plain" + + +def test_resumed_session_thinking_and_null_content_do_not_400(): + # A resumed session replays assistant turns with `thinking` (and sometimes null) + # content. Those must be accepted (thinking dropped by the converter), not 400ed. + from pydantic import ValidationError + + req = AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "secret reasoning", "signature": "s"}, + {"type": "text", "text": "the answer"}, + {"type": "tool_use", "id": "t1", "name": "f", "input": {}}, + ], + }, + {"role": "assistant", "content": None}, # tool-only turn serialized as null + ], + ) + # Known blocks still parse as their typed models; only the unknown one is loose. + assert type(req.messages[1].content[0]).__name__ == "AnthropicUnknownBlock" + assert type(req.messages[1].content[1]).__name__ == "AnthropicTextBlock" + assert req.messages[2].content == "" # null coerced + + openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages]) + assistant = next(m for m in openai if m["role"] == "assistant" and m.get("content")) + assert assistant["content"] == "the answer" + assert "secret reasoning" not in json.dumps(openai) # thinking never forwarded + + # A malformed KNOWN block still fails cleanly instead of being swallowed. + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [{"role": "assistant", "content": [{"type": "tool_use", "name": "f"}]}], + ) + + +def test_user_null_content_rejected(): + # The null->"" leniency is assistant-only; a null user content must be rejected + # at the boundary, not coerced into an empty prompt and forwarded to the model. + from pydantic import ValidationError + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [{"role": "user", "content": None}], + ) + + +def test_user_unknown_block_rejected_not_silently_dropped(): + # The converter skips user blocks it cannot translate, so a user turn whose only + # block is unknown would validate yet forward no content. Reject at the boundary + # to avoid that silent data loss (the assistant fallback is unaffected). + from pydantic import ValidationError + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": [{"type": "document", "source": {}}]}, + ], + ) + + +def test_user_translatable_blocks_still_accepted(): + # text / image / tool_result are translatable, so a real user message built from + # them must still pass; the unknown-block guard only trips on other types. + req = AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image", + "source": {"type": "base64", "media_type": "image/png", "data": "AA"}, + }, + {"type": "tool_result", "tool_use_id": "t1", "content": "ok"}, + ], + } + ], + ) + assert [type(b).__name__ for b in req.messages[0].content] == [ + "AnthropicTextBlock", + "AnthropicImageBlock", + "AnthropicToolResultBlock", + ] + + openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages]) + assert any(m["role"] == "tool" and m["tool_call_id"] == "t1" for m in openai) + + +def test_user_malformed_known_block_still_rejected(): + # The guard only allow-lists a user block's *type*; the union still validates its + # shape, so a known-but-malformed block (tool_result without tool_use_id) fails. + from pydantic import ValidationError + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": [{"type": "tool_result", "content": "x"}]}, + ], + ) + + +def test_user_content_block_non_string_type_rejected_cleanly(): + # A user block whose `type` is a non-string (unhashable list / dict, or a stray + # int) must fail as a clean validation error, not raise TypeError from the + # frozenset membership test and escape as a 500. + from pydantic import ValidationError + for bad_type in ([], {}, 5): + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [{"role": "user", "content": [{"type": bad_type}]}], + ) + + +def test_assistant_missing_content_key_still_rejected(): + # The null -> "" leniency is only for an EXPLICIT null. An assistant message that + # omits content entirely stays malformed and must fail required-field validation. + from pydantic import ValidationError + + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [{"role": "assistant"}], + ) + # An explicit null is still accepted and coerced (regression guard). + req = AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": None}, + ], + ) + assert req.messages[1].content == "" + + +def test_resumed_null_assistant_between_users_coalesced_on_messages_route(monkeypatch): + # user -> assistant(null) -> user is now accepted: the null assistant turn coerces + # to "" and is dropped. The route must then coalesce the two remaining user turns + # so a strict GGUF chat template does not 400 on non-alternating roles. + backend = _mock_backend(monkeypatch, context_length = 2048) + + class _Req: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/messages") + method = "POST" + + async def is_disconnected(self): + return False + + payload = AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": "first question"}, + {"role": "assistant", "content": None}, + {"role": "user", "content": "please continue"}, + ], + ) + + response = _drive(anthropic_messages(payload, request = _Req(), current_subject = "t")) + assert response.status_code == 200 + + [(_path, kwargs)] = backend.calls + user_turns = [m for m in kwargs["messages"] if m.get("role") == "user"] + assert len(user_turns) == 1 # the two user turns were merged, not left adjacent + merged = user_turns[0]["content"] + if isinstance(merged, list): + merged = " ".join(p.get("text", "") for p in merged if isinstance(p, dict)) + assert "first question" in merged and "please continue" in merged diff --git a/studio/backend/tests/test_chat_eos_template_refresh.py b/studio/backend/tests/test_chat_eos_template_refresh.py new file mode 100644 index 0000000000..75d0117015 --- /dev/null +++ b/studio/backend/tests/test_chat_eos_template_refresh.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Mapper models whose own tokenizer ships no chat_template have their turn-end +eos resolved at LOAD from an empty template (document eos only). The effective +template is installed later, at generate time, via get_chat_template, so the +turn-end-eos cache must be refreshed then; otherwise generate_stream runs past +the ChatML <|im_end|> boundary and loops (the exact bug this PR fixes). +""" + +import sys +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +# These tests construct InferenceBackend, pulling the full stack. CI may lack +# unsloth/unsloth_zoo (ImportError) or have a broken CUDA/bitsandbytes setup +# (RuntimeError); skip at module level so collection is not aborted (exit 2). +try: + from core.inference import inference as inf_mod # noqa: E402 + from core.inference.inference import InferenceBackend # noqa: E402 +except (ImportError, RuntimeError) as exc: # pragma: no cover - env-dependent + pytest.skip( + f"full inference backend unavailable ({type(exc).__name__}: {exc})", + allow_module_level = True, + ) + +_CHATML = "{% for m in messages %}<|im_start|>{{m.role}}\n{{m.content}}<|im_end|>{% endfor %}" +_GEMMA = "{% for m in messages %}{{m.role}}\n{{m.content}}{% endfor %}" + + +class _FakeTokenizer: + def __init__( + self, + eos_id, + chat_template = "", + token_ids = None, + ): + self.eos_token_id = eos_id + self.chat_template = chat_template + self.pad_token_id = eos_id + self.unk_token_id = None + self._ids = dict(token_ids or {}) + + def convert_tokens_to_ids(self, tok): + return self._ids.get(tok) + + +def test_turn_end_eos_refreshed_after_generate_time_template(monkeypatch): + import utils.datasets as ds + + backend = InferenceBackend.__new__(InferenceBackend) + backend.active_model_name = "unsloth/qwen2.5-0.5b" + + # No chat_template at load, so the cache stored only the document eos, though + # <|im_end|> is atomic in the vocab (unused until the mapper installs a template). + bare_tok = _FakeTokenizer(151643, chat_template = "", token_ids = {"<|im_end|>": 151645}) + model_info = { + "tokenizer": bare_tok, + "is_vision": False, + "chat_turn_end_eos_ids": [151643], + } + backend.models = {backend.active_model_name: model_info} + + # The mapper installs a ChatML template (turns end with <|im_end|>) at generate time. + templated_tok = _FakeTokenizer(151643, chat_template = _CHATML, token_ids = {"<|im_end|>": 151645}) + monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: templated_tok) + monkeypatch.setattr( + ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "qwen-2.5"}, raising = False + ) + + # Stub the tail so the generator runs through the refresh without a real model. + monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False) + monkeypatch.setattr( + backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False + ) + monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False) + + list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}])) + + # After the template is applied the cache must include the ChatML turn-end id. + assert model_info["chat_turn_end_eos_ids"] == [151643, 151645] + + +def test_turn_end_eos_refresh_preserves_load_time_ids_on_destructive_swap(monkeypatch): + # Regression: get_chat_template can return a remapped tokenizer (Gemma: + # folded onto the eos id) while generate_stream re-reads the original. Resolving on + # the swap yields a narrower set, so the refresh must UNION, never overwrite. + import utils.datasets as ds + + backend = InferenceBackend.__new__(InferenceBackend) + backend.active_model_name = "unsloth/gemma-2b-it" + + # Original tokenizer (used by generate_stream): =107 distinct from + # eos=1, so the load-time cache resolved to [1, 107]. + orig_tok = _FakeTokenizer(1, chat_template = _GEMMA, token_ids = {"": 107}) + model_info = { + "tokenizer": orig_tok, + "is_vision": False, + "chat_turn_end_eos_ids": [1, 107], + } + backend.models = {backend.active_model_name: model_info} + + # Destructively-swapped tokenizer: now maps onto eos id 1, so + # resolving on it yields only [1] (drops 107). + swapped_tok = _FakeTokenizer(1, chat_template = _GEMMA, token_ids = {"": 1}) + monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: swapped_tok) + monkeypatch.setattr( + ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "gemma-3"}, raising = False + ) + + monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False) + monkeypatch.setattr( + backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False + ) + monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False) + + list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}])) + + # The load-time =107 must survive: overwriting with the swapped + # [1] would regress and loop past the turn. + assert model_info["chat_turn_end_eos_ids"] == [1, 107] + + +def test_turn_end_eos_refresh_resolves_marker_id_on_original_not_remapped(monkeypatch): + # Yi-style map_eos_token=True: the original carries <|im_end|> at its own id, but + # get_chat_template folds it onto the doc-eos id. generate_stream uses the original, + # so read marker strings from the mapped template but ids from the original. + import utils.datasets as ds + + backend = InferenceBackend.__new__(InferenceBackend) + backend.active_model_name = "01-ai/yi-6b" + + # Original: no template of its own, doc eos = 2, <|im_end|> atomic = 7. + orig_tok = _FakeTokenizer(2, chat_template = "", token_ids = {"<|im_end|>": 7}) + model_info = { + "tokenizer": orig_tok, + "is_vision": False, + "chat_turn_end_eos_ids": [2], + } + backend.models = {backend.active_model_name: model_info} + + # Remapped tokenizer: ChatML template, but <|im_end|> folded onto doc-eos id 2. + remapped_tok = _FakeTokenizer(2, chat_template = _CHATML, token_ids = {"<|im_end|>": 2}) + monkeypatch.setattr(inf_mod, "get_chat_template", lambda tok, chat_template = None: remapped_tok) + monkeypatch.setattr( + ds, "MODEL_TO_TEMPLATE_MAPPER", {backend.active_model_name: "chatml"}, raising = False + ) + + monkeypatch.setattr(backend, "_normalize_top_k", lambda k: k, raising = False) + monkeypatch.setattr( + backend, "_apply_chat_template_for_generation", lambda *a, **k: "PROMPT", raising = False + ) + monkeypatch.setattr(backend, "generate_stream", lambda *a, **k: iter(()), raising = False) + + list(backend._generate_chat_response_inner(messages = [{"role": "user", "content": "hi"}])) + + # The real <|im_end|>=7 (original vocab) must be recovered, not the remapped 2. + assert model_info["chat_turn_end_eos_ids"] == [2, 7] + + +class _FakeProcessor: + """A ProcessorMixin-like container: carries the chat_template itself and + wraps the real text tokenizer as ``.tokenizer`` (the vision layout).""" + + def __init__(self, chat_template, tokenizer): + self.chat_template = chat_template + self.tokenizer = tokenizer + + +def test_resolve_chat_eos_reads_vision_processor_template(): + # Vision model: the chat_template lives on the processor while the inner tokenizer + # ships none. _resolve_chat_eos must read the marker from the processor but resolve + # its id on the inner tokenizer, and repair generation_config. + from types import SimpleNamespace + + inner_tok = _FakeTokenizer(1, chat_template = "", token_ids = {"": 107}) + processor = _FakeProcessor(_GEMMA, inner_tok) + model = SimpleNamespace(generation_config = SimpleNamespace(eos_token_id = 1)) + + backend = InferenceBackend.__new__(InferenceBackend) + backend.active_model_name = "unsloth/gemma-3-4b-it" + model_info = {"model": model, "tokenizer": processor, "processor": processor, "is_vision": True} + backend.models = {backend.active_model_name: model_info} + + backend._resolve_chat_eos(backend.active_model_name) + + assert model_info["chat_turn_end_eos_ids"] == [1, 107] + # generation_config repaired so the vision .generate() path stops at the turn. + assert model.generation_config.eos_token_id == [1, 107] diff --git a/studio/backend/tests/test_chat_history_routes.py b/studio/backend/tests/test_chat_history_routes.py index 2a6ebe244f..a60ac700bf 100644 --- a/studio/backend/tests/test_chat_history_routes.py +++ b/studio/backend/tests/test_chat_history_routes.py @@ -91,6 +91,17 @@ def test_chat_settings_payload_accepts_fast_mode_presets(): assert dumped["customPresets"][0]["params"]["fastMode"] is True +def test_chat_settings_payload_accepts_nudge_tool_calls(): + # extra="forbid" 400s PUT /api/chat/settings on unknown keys, so the + # frontend's persisted nudgeToolCalls needs a payload field (like + # autoHealToolCalls). + payload = chat_history.ChatSettingsPayload.model_validate( + {"autoHealToolCalls": True, "nudgeToolCalls": False} + ) + dumped = payload.model_dump(exclude_unset = True) + assert dumped == {"autoHealToolCalls": True, "nudgeToolCalls": False} + + def test_chat_inference_settings_covers_frontend_persisted_fields(): # Drift guard: every InferenceParams field the UI persists (all but # checkpoint) must exist on ChatInferenceSettings, else extra="forbid" diff --git a/studio/backend/tests/test_chat_history_storage.py b/studio/backend/tests/test_chat_history_storage.py index aa19df15fe..0239410734 100644 --- a/studio/backend/tests/test_chat_history_storage.py +++ b/studio/backend/tests/test_chat_history_storage.py @@ -4,6 +4,7 @@ import os import platform import shutil +import sqlite3 import threading import uuid from pathlib import Path @@ -11,6 +12,7 @@ from pathlib import Path import pytest from storage import studio_db +from utils.paths import studio_db_path def _reset_studio_db( @@ -108,6 +110,138 @@ def test_sync_chat_messages_upserts_without_pruning(tmp_path, monkeypatch): assert by_id["msg-2"]["content"] == [{"type": "text", "text": "updated text"}] +def test_chat_thread_updated_at_bumps_on_message_writes(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + thread = studio_db.upsert_chat_thread(_thread()) + assert thread["updatedAt"] == thread["createdAt"] + + studio_db.upsert_chat_message(_message("msg-1", 1_700_000_000_500, "hi")) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500 + + studio_db.upsert_chat_message(_message("msg-0", 1_600_000_000_000, "old")) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500 + + studio_db.sync_chat_messages( + "thread-1", + [_message("msg-2", 1_700_000_001_000, "newer")], + ) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_001_000 + + +def test_chat_thread_updated_at_recomputed_when_pruning(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + thread = studio_db.upsert_chat_thread(_thread()) + studio_db.sync_chat_messages( + "thread-1", + [ + _message("msg-1", 1_700_000_000_500, "older"), + _message("msg-2", 1_700_000_001_000, "newest"), + ], + prune_missing = True, + ) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_001_000 + + # Pruning the newest message must lower updated_at to the remaining one. + studio_db.sync_chat_messages( + "thread-1", + [_message("msg-1", 1_700_000_000_500, "older")], + prune_missing = True, + ) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500 + + # Pruning every message falls back to created_at. + studio_db.sync_chat_messages("thread-1", [], prune_missing = True) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == thread["createdAt"] + + +def test_chat_thread_updated_at_survives_thread_resave(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + studio_db.upsert_chat_message(_message("msg-1", 1_700_000_000_500, "hi")) + + studio_db.upsert_chat_thread(_thread()) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500 + + +def test_list_chat_threads_orders_by_last_activity(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + older = _thread("thread-old") + older["createdAt"] = 1_700_000_000_000 + newer = _thread("thread-new") + newer["createdAt"] = 1_700_000_100_000 + studio_db.upsert_chat_thread(older) + studio_db.upsert_chat_thread(newer) + assert [t["id"] for t in studio_db.list_chat_threads()] == ["thread-new", "thread-old"] + + studio_db.upsert_chat_message( + _message("msg-1", 1_700_000_200_000, "hi", thread_id = "thread-old") + ) + assert [t["id"] for t in studio_db.list_chat_threads()] == ["thread-old", "thread-new"] + + +def test_chat_threads_updated_at_migration_backfills_from_messages(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + db_path = studio_db_path() + db_path.parent.mkdir(parents = True, exist_ok = True) + conn = sqlite3.connect(str(db_path)) + try: + conn.execute( + """ + CREATE TABLE chat_threads ( + id TEXT NOT NULL PRIMARY KEY, + title TEXT NOT NULL, + model_type TEXT NOT NULL, + model_id TEXT, + pair_id TEXT, + archived INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL + ) + """ + ) + conn.execute( + """ + CREATE TABLE chat_messages ( + id TEXT NOT NULL PRIMARY KEY, + thread_id TEXT NOT NULL, + parent_id TEXT, + role TEXT NOT NULL, + content_json TEXT NOT NULL, + attachments_json TEXT, + metadata_json TEXT, + created_at INTEGER NOT NULL + ) + """ + ) + conn.execute( + "INSERT INTO chat_threads (id, title, model_type, created_at) VALUES (?, ?, ?, ?)", + ("thread-with-msgs", "Old", "base", 1_700_000_000_000), + ) + conn.execute( + "INSERT INTO chat_threads (id, title, model_type, created_at) VALUES (?, ?, ?, ?)", + ("thread-empty", "Empty", "base", 1_700_000_050_000), + ) + # Fork-like thread: copied ancestor messages predate the thread itself. + conn.execute( + "INSERT INTO chat_threads (id, title, model_type, created_at) VALUES (?, ?, ?, ?)", + ("thread-fork", "Fork", "base", 1_700_000_100_000), + ) + conn.executemany( + "INSERT INTO chat_messages (id, thread_id, role, content_json, created_at) VALUES (?, ?, ?, ?, ?)", + [ + ("m1", "thread-with-msgs", "user", "[]", 1_700_000_001_000), + ("m2", "thread-with-msgs", "assistant", "[]", 1_700_000_002_000), + ("m3", "thread-fork", "user", "[]", 1_700_000_001_000), + ], + ) + conn.commit() + finally: + conn.close() + + assert studio_db.get_chat_thread("thread-with-msgs")["updatedAt"] == 1_700_000_002_000 + assert studio_db.get_chat_thread("thread-empty")["updatedAt"] == 1_700_000_050_000 + assert studio_db.get_chat_thread("thread-fork")["updatedAt"] == 1_700_000_100_000 + + def test_chat_projects_delete_cascades_threads_and_messages(tmp_path, monkeypatch): _reset_studio_db(tmp_path, monkeypatch) project = studio_db.upsert_chat_project(_project()) diff --git a/studio/backend/tests/test_chat_template_tool_arguments.py b/studio/backend/tests/test_chat_template_tool_arguments.py new file mode 100644 index 0000000000..13d1ecabaa --- /dev/null +++ b/studio/backend/tests/test_chat_template_tool_arguments.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""apply_chat_template_for_generation must coerce assistant tool_call arguments +from the OpenAI JSON-string form to a dict before rendering. Strict tool +templates (e.g. mlx-community Qwen3.5 checkpoints) iterate arguments.items() and +raise "Can only get item pairs from a mapping." on the string form when a prior +tool call is re-rendered on the next turn (MLX + transformers paths). +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +from core.inference.chat_template_helpers import ( # noqa: E402 + _normalize_tool_call_arguments, + apply_chat_template_for_generation, +) + + +def _conv(arguments): + return [ + {"role": "user", "content": "weather?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "type": "function", + "id": "c1", + "function": {"name": "web_search", "arguments": arguments}, + } + ], + }, + {"role": "tool", "name": "web_search", "content": "21C sunny"}, + ] + + +class _StrictTemplateTokenizer: + """Mimics a strict Qwen tool template: rejects string tool_call arguments.""" + + def apply_chat_template( + self, + messages, + *, + tokenize = False, + add_generation_prompt = True, + **kw, + ): + for msg in messages: + for call in msg.get("tool_calls", []) or []: + args = call.get("function", {}).get("arguments") + if isinstance(args, str): + raise TypeError("Can only get item pairs from a mapping.") + return "RENDERED" + + +def test_string_arguments_are_parsed_to_dict(): + out = _normalize_tool_call_arguments(_conv('{"query": "sweden"}')) + args = out[1]["tool_calls"][0]["function"]["arguments"] + assert args == {"query": "sweden"} + + +def test_dict_arguments_untouched_and_no_copy(): + conv = _conv({"query": "sweden"}) + assert _normalize_tool_call_arguments(conv) is conv + + +def test_non_json_string_left_as_is(): + out = _normalize_tool_call_arguments(_conv("not json")) + assert out[1]["tool_calls"][0]["function"]["arguments"] == "not json" + + +def test_render_succeeds_on_strict_template_with_string_arguments(): + # Regression: strict template + string args used to raise. + result = apply_chat_template_for_generation(_StrictTemplateTokenizer(), _conv('{"query": "x"}')) + assert result == "RENDERED" + + +class _RecordingTokenizer: + """Lenient template: renders whatever arguments it is given (string or dict).""" + + def __init__(self): + self.seen_arguments = None + + def apply_chat_template( + self, + messages, + *, + tokenize = False, + add_generation_prompt = True, + **kw, + ): + for msg in messages: + for call in msg.get("tool_calls", []) or []: + self.seen_arguments = call.get("function", {}).get("arguments") + return "RENDERED" + + +def test_lenient_template_receives_original_string_untouched(): + # Lenient template must see the exact original string, not a coerced dict. + tok = _RecordingTokenizer() + apply_chat_template_for_generation(tok, _conv('{"query": "x"}')) + assert tok.seen_arguments == '{"query": "x"}' + + +def test_messages_without_tool_calls_pass_through_unchanged(): + conv = [{"role": "user", "content": "hi"}] + assert _normalize_tool_call_arguments(conv) is conv + + +class _RaiseExceptionTemplateTokenizer: + """Mimics the bundled gemma-4.jinja: rejects string tool_call arguments via + ``raise_exception(...)``, which surfaces as a Jinja error, NOT a TypeError.""" + + def apply_chat_template( + self, + messages, + *, + tokenize = False, + add_generation_prompt = True, + **kw, + ): + for msg in messages: + for call in msg.get("tool_calls", []) or []: + args = call.get("function", {}).get("arguments") + if isinstance(args, str): + raise ValueError( + "chat_template: tool_calls[].function.arguments must be a " + "JSON object (mapping), not a string." + ) + return "RENDERED" + + +def test_render_succeeds_on_raise_exception_template_with_string_arguments(): + # Regression: gemma-4.jinja rejects string args via a non-TypeError; retry must still coerce. + result = apply_chat_template_for_generation( + _RaiseExceptionTemplateTokenizer(), _conv('{"query": "x"}') + ) + assert result == "RENDERED" + + +def test_unrelated_template_error_still_propagates_with_dict_args(): + # Failure unrelated to string args (dict args, nothing to coerce) must propagate. + class _AlwaysRaises: + def apply_chat_template(self, messages, **kw): + raise ValueError("template is broken") + + with pytest.raises(ValueError, match = "broken"): + apply_chat_template_for_generation(_AlwaysRaises(), _conv({"query": "x"})) diff --git a/studio/backend/tests/test_chat_turn_end_eos.py b/studio/backend/tests/test_chat_turn_end_eos.py new file mode 100644 index 0000000000..c49e39f8fe --- /dev/null +++ b/studio/backend/tests/test_chat_turn_end_eos.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""chat_eos: resolve assistant-turn-end stop tokens from the chat_template and +repair generation_config so a chat model whose eos is a bare document terminator +(Qwen3.5: config eos <|endoftext|>, turns end with <|im_end|>) stops at the turn +boundary instead of running past it and looping. Dependency-light: imported here +without the full inference stack. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +from core.inference.chat_eos import ( # noqa: E402 + chat_eos_repair, + resolve_chat_turn_end_eos_ids, + resolve_chat_turn_end_eos_ids_using, +) + + +class _FakeTokenizer: + def __init__( + self, + eos_id, + chat_template = "", + token_ids = None, + unk_token_id = None, + ): + self.eos_token_id = eos_id + self.chat_template = chat_template + self.unk_token_id = unk_token_id + self._ids = dict(token_ids or {}) + + def convert_tokens_to_ids(self, tok): + return self._ids.get(tok, self.unk_token_id) + + +# ---- resolve_chat_turn_end_eos_ids --------------------------------------- + +_CHATML = "{% for m in messages %}<|im_start|>{{m.role}}\n{{m.content}}<|im_end|>{% endfor %}" + + +def test_qwen35_adds_im_end_from_template(): + # eos synced to <|endoftext|> (248044); template uses <|im_end|> (248046). + tok = _FakeTokenizer(248044, chat_template = _CHATML, token_ids = {"<|im_end|>": 248046}) + assert resolve_chat_turn_end_eos_ids(tok) == [248044, 248046] + + +def test_marker_in_vocab_but_not_in_template_is_ignored(): + # Base/coder model: <|im_end|> is in the vocab but the template does not use + # it, so it must not become a stop token. + tok = _FakeTokenizer(248044, chat_template = "{{ messages }}", token_ids = {"<|im_end|>": 248046}) + assert resolve_chat_turn_end_eos_ids(tok) == [248044] + + +def test_harmony_template_is_left_untouched(): + # gpt-oss/harmony: <|end|> is a channel delimiter, not the turn end. + harmony = "<|start|>assistant<|channel|>analysis<|message|>...<|end|>" + tok = _FakeTokenizer(200002, chat_template = harmony, token_ids = {"<|end|>": 200007}) + assert resolve_chat_turn_end_eos_ids(tok) == [200002] + + +def test_llama3_eot_id_from_template(): + tok = _FakeTokenizer(128001, chat_template = "...<|eot_id|>...", token_ids = {"<|eot_id|>": 128009}) + assert resolve_chat_turn_end_eos_ids(tok) == [128001, 128009] + + +def test_gemma4_turn_marker_from_template(): + # Gemma-4 ends turns with while keeping a document eos, so must + # be added as a stop token. + tok = _FakeTokenizer( + 1, chat_template = ".........", token_ids = {"": 106} + ) + assert resolve_chat_turn_end_eos_ids(tok) == [1, 106] + + +def test_resolve_using_reads_markers_from_template_but_ids_from_generation_tokenizer(): + # map_eos_token=True: the mapped template remaps <|im_end|> onto the doc-eos id, + # but the original keeps it atomic. Reading marker STRINGS from the template but + # IDS on the original recovers the real turn-end id (7), not the doc-eos id (2). + template_tok = _FakeTokenizer(2, chat_template = _CHATML, token_ids = {"<|im_end|>": 2}) + id_tok = _FakeTokenizer(2, chat_template = "", token_ids = {"<|im_end|>": 7}) + assert resolve_chat_turn_end_eos_ids_using(template_tok, id_tok) == [2, 7] + # Same tokenizer for both reproduces the plain resolve (load-time behaviour). + assert resolve_chat_turn_end_eos_ids_using(template_tok, template_tok) == [2] + + +def test_list_eos_preserved(): + tok = _FakeTokenizer([1, 2], chat_template = _CHATML, token_ids = {"<|im_end|>": 2}) + assert resolve_chat_turn_end_eos_ids(tok) == [1, 2] + + +def test_missing_marker_maps_to_unk_and_is_skipped(): + tok = _FakeTokenizer(7, chat_template = _CHATML, token_ids = {}, unk_token_id = 0) + assert resolve_chat_turn_end_eos_ids(tok) == [7] + + +def test_starling_barred_end_of_turn_from_template(): + # OpenChat/Starling end turns with the BARRED <|end_of_turn|> (distinct from + # Gemma's ). eos synced to =2, turn marker at 32000. + starling = "GPT4 Correct Assistant: hi<|end_of_turn|>" + tok = _FakeTokenizer(2, chat_template = starling, token_ids = {"<|end_of_turn|>": 32000}) + assert resolve_chat_turn_end_eos_ids(tok) == [2, 32000] + + +def test_dict_chat_template_scans_all_variants(): + # Hermes-3 style: chat_template is a {name: template} dict. Detection must scan + # every variant, not bail because the container is not a plain str. + tmpl = {"default": "{{ messages }}", "tool_use": _CHATML} + tok = _FakeTokenizer(2, chat_template = tmpl, token_ids = {"<|im_end|>": 5}) + assert resolve_chat_turn_end_eos_ids(tok) == [2, 5] + + +def test_list_of_dicts_chat_template_scans_all_variants(): + # tokenizer_config.json stores multi-templates as a list of {name, template}. + tmpl = [{"name": "default", "template": _CHATML}] + tok = _FakeTokenizer(2, chat_template = tmpl, token_ids = {"<|im_end|>": 5}) + assert resolve_chat_turn_end_eos_ids(tok) == [2, 5] + + +def test_dict_harmony_template_left_untouched(): + # A multi-variant container whose variant is harmony must still be left alone. + tmpl = {"default": "<|start|>assistant<|channel|>analysis<|message|>...<|end|>"} + tok = _FakeTokenizer(200002, chat_template = tmpl, token_ids = {"<|end|>": 200007}) + assert resolve_chat_turn_end_eos_ids(tok) == [200002] + + +# ---- chat_eos_repair ------------------------------------------------------ + + +def test_repair_adds_missing_turn_end(): + assert chat_eos_repair(248044, [248044, 248046]) == [248044, 248046] + + +def test_repair_from_missing_generation_config_eos(): + assert chat_eos_repair(None, [248046]) == [248046] + + +def test_repair_noop_when_already_covered(): + assert chat_eos_repair([248046, 248044], [248046]) is None + + +def test_repair_noop_when_no_turn_end_ids(): + assert chat_eos_repair(248044, []) is None diff --git a/studio/backend/tests/test_coding_agents.py b/studio/backend/tests/test_coding_agents.py new file mode 100644 index 0000000000..b19da1dded --- /dev/null +++ b/studio/backend/tests/test_coding_agents.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for coding-agent CLI detection used by the API-keys settings panel.""" + +from unittest.mock import patch + +from utils.coding_agents import CODING_AGENTS, detect_installed_coding_agents + + +def test_matches_unsloth_start_subcommands(): + # Each entry must be an actual `unsloth start ` subcommand name + # (unsloth_cli/commands/start.py). Spelled out here rather than imported + # from that module, which pulls in the CLI's heavier dependencies. + assert CODING_AGENTS == ("claude", "codex", "openclaw", "opencode", "hermes", "pi") + + +def test_detects_only_agents_present_on_path(): + installed = {"claude", "opencode"} + with patch( + "utils.coding_agents.shutil.which", + side_effect = lambda name: f"/usr/bin/{name}" if name in installed else None, + ): + assert detect_installed_coding_agents() == ["claude", "opencode"] + + +def test_returns_empty_list_when_nothing_is_installed(): + with patch("utils.coding_agents.shutil.which", return_value = None): + assert detect_installed_coding_agents() == [] + + +def test_preserves_declared_order_regardless_of_path_lookup_order(): + with patch( + "utils.coding_agents.shutil.which", + side_effect = lambda name: name if name in ("pi", "claude", "hermes") else None, + ): + assert detect_installed_coding_agents() == ["claude", "hermes", "pi"] + + +def test_treats_a_path_lookup_error_as_not_installed(): + # An advisory check: shutil.which raising for one entry (e.g. a permission + # error walking a PATH directory) should not take down the whole endpoint, + # and should not stop the remaining agents from being checked. + def flaky_which(name: str): + if name == "codex": + raise OSError("permission denied") + return name if name == "claude" else None + + with patch("utils.coding_agents.shutil.which", side_effect = flaky_which): + assert detect_installed_coding_agents() == ["claude"] diff --git a/studio/backend/tests/test_completion_masking.py b/studio/backend/tests/test_completion_masking.py new file mode 100644 index 0000000000..be0d8a69bd --- /dev/null +++ b/studio/backend/tests/test_completion_masking.py @@ -0,0 +1,314 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Completion-only masking policy: auto-detect first, manual table fallback. + +Covers utils.datasets.completion_masking.apply_completion_masking, shared by +the CUDA trainer (core/training/trainer.py) and the MLX worker +(core/training/worker.py): + - unmapped models use chat template auto-detection (previously masking was + silently disabled), + - gpt-oss goes auto-first too (its quantized checkpoints ship a template + the manual markers cannot match), + - an auto-detection failure falls back to the template table markers, + - a table miss after an auto failure warns and leaves the trainer unchanged. +""" + +from __future__ import annotations + +import pytest + +from utils.datasets.completion_masking import apply_completion_masking, lookup_manual_markers +from utils.datasets.model_mappings import TEMPLATE_TO_RESPONSES_MAPPER + + +class _Trainer: + """Sentinel trainer; train_fn wraps it in a new object when applied.""" + + +class _Recorder: + """Fake train_on_responses_only that records calls.""" + + def __init__(self): + self.calls = [] + + def __call__(self, trainer, **kwargs): + self.calls.append(kwargs) + wrapped = _Trainer() + wrapped.wrapped_from = trainer + return wrapped + + +def _detect_ok(processor): + return "", "" + + +def _detect_fail(processor): + raise ValueError( + "Unsloth: Could not reliably auto-detect response_part - " + "pass instruction_part and response_part." + ) + + +_AUTO = {"instruction_part": "", "response_part": ""} + + +class _Notes: + def __init__(self): + self.messages = [] + + def __call__(self, level, message): + self.messages.append((level, message)) + + def warnings(self): + return [m for level, m in self.messages if level == "warning"] + + +def test_unmapped_model_uses_auto_detection(): + # Unmapped model: the auto path applies masking (was silently disabled). + trainer = _Trainer() + train_fn = _Recorder() + notes = _Notes() + + result, applied = apply_completion_masking( + trainer, "LiquidAI/LFM2-8B-A1B", train_fn, notify = notes, detect_fn = _detect_ok + ) + + assert applied is True + assert result.wrapped_from is trainer + assert train_fn.calls == [dict(_AUTO)] # applied with the detected markers + assert notes.warnings() == [] + + +def test_mapped_model_prefers_auto_detection(): + trainer = _Trainer() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "unsloth/Qwen3-0.6B", train_fn, detect_fn = _detect_ok + ) + + assert applied is True + assert train_fn.calls == [dict(_AUTO)] + + +def test_gpt_oss_uses_auto_detection_first(): + # The quantized gpt-oss checkpoints ship a template without the + # <|channel|>final header, where the manual markers match nothing; auto + # derives markers from the template the checkpoint actually ships. + trainer = _Trainer() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "unsloth/gpt-oss-20b", train_fn, detect_fn = _detect_ok + ) + + assert applied is True + assert train_fn.calls == [dict(_AUTO)] + + +def test_gpt_oss_detection_failure_falls_back_to_manual_markers(): + trainer = _Trainer() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "unsloth/gpt-oss-20b", train_fn, detect_fn = _detect_fail + ) + + assert applied is True + expected = TEMPLATE_TO_RESPONSES_MAPPER["gpt-oss"] + assert train_fn.calls == [ + { + "instruction_part": expected["instruction"], + "response_part": expected["response"], + } + ] + + +def test_auto_failure_falls_back_to_template_table(): + trainer = _Trainer() + train_fn = _Recorder() + notes = _Notes() + + result, applied = apply_completion_masking( + trainer, "unsloth/Qwen3-0.6B", train_fn, notify = notes, detect_fn = _detect_fail + ) + + assert applied is True + assert result.wrapped_from is trainer + expected = TEMPLATE_TO_RESPONSES_MAPPER["qwen3"] + assert train_fn.calls == [ + { + "instruction_part": expected["instruction"], + "response_part": expected["response"], + }, + ] + assert any("falling back to the template table" in m for m in notes.warnings()) + + +def test_application_failure_propagates_not_fallback(): + # Detection succeeds; a failure while APPLYING the masking must propagate, + # never silently fall back to full-sequence training. + def train_fn(trainer, **kwargs): + raise RuntimeError("dataset map worker crashed") + + with pytest.raises(RuntimeError, match = "dataset map worker crashed"): + apply_completion_masking(_Trainer(), "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_ok) + + +def test_preset_tokenizer_markers_used_directly(): + # Preset unsloth marker attrs skip detection; zoo reuses them on a bare call. + class _Tok: + _unsloth_input_part = "" + _unsloth_output_part = "" + + trainer = _Trainer() + trainer.processing_class = _Tok() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_fail + ) + assert applied is True + assert train_fn.calls == [{}] # bare call, stored parts + + +def test_table_miss_warns_and_disables_without_crashing(): + trainer = _Trainer() + train_fn = _Recorder() + notes = _Notes() + + result, applied = apply_completion_masking( + trainer, "some-org/not-in-any-mapper", train_fn, notify = notes, detect_fn = _detect_fail + ) + + assert applied is False + assert result is trainer # unchanged: full sequence training + assert train_fn.calls == [] # detection failed; nothing applied + assert any("could not be applied" in m for m in notes.warnings()) + assert any("full sequences" in m for m in notes.warnings()) + + +def test_num_proc_forwarded_only_when_given(): + # CUDA path passes num_proc; the MLX path omits it. + train_fn = _Recorder() + apply_completion_masking( + _Trainer(), "unsloth/Qwen3-0.6B", train_fn, num_proc = 4, detect_fn = _detect_ok + ) + assert train_fn.calls == [dict(_AUTO, num_proc = 4)] + + train_fn = _Recorder() + apply_completion_masking( + _Trainer(), "unsloth/Qwen3-0.6B", train_fn, num_proc = 4, detect_fn = _detect_fail + ) + assert train_fn.calls[0]["num_proc"] == 4 + + train_fn = _Recorder() + apply_completion_masking(_Trainer(), "unsloth/Qwen3-0.6B", train_fn, detect_fn = _detect_ok) + assert train_fn.calls == [dict(_AUTO)] + + +def test_manual_fallback_failure_propagates_to_caller(): + # Errors while applying the manual fallback must propagate to the caller. + def train_fn(trainer, **kwargs): + raise RuntimeError("boom") + + with pytest.raises(RuntimeError, match = "boom"): + apply_completion_masking(_Trainer(), "unsloth/gpt-oss-20b", train_fn) + + +def test_notify_is_optional(): + train_fn = _Recorder() + _, applied = apply_completion_masking( + _Trainer(), "some-org/not-in-any-mapper", train_fn, detect_fn = _detect_fail + ) + assert applied is False + + +def test_lookup_manual_markers(): + template, instruction, response = lookup_manual_markers("unsloth/Qwen3-0.6B") + assert template == "qwen3" + assert instruction == TEMPLATE_TO_RESPONSES_MAPPER["qwen3"]["instruction"] + assert response == TEMPLATE_TO_RESPONSES_MAPPER["qwen3"]["response"] + + template, instruction, response = lookup_manual_markers("some-org/unknown") + assert (template, instruction, response) == (None, None, None) + + template, instruction, response = lookup_manual_markers(None) + assert (template, instruction, response) == (None, None, None) + + +def test_renamed_gpt_oss_gets_template_markers(): + # Name-detected as gpt-oss but not in the exact-name table: must use the + # gpt-oss markers, not fall through to full-sequence training. + trainer = _Trainer() + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "some-org/gpt-oss-20b-sft", train_fn, detect_fn = _detect_fail + ) + assert applied is True + expected = TEMPLATE_TO_RESPONSES_MAPPER["gpt-oss"] + assert train_fn.calls == [ + { + "instruction_part": expected["instruction"], + "response_part": expected["response"], + } + ] + + +class _FakeTokenizerWrapper: + """mlx-lm TokenizerWrapper semantics: plain reads delegate to the wrapped + tokenizer, underscore attrs do not (so preset markers are hidden).""" + + def __init__(self, tokenizer): + object.__setattr__(self, "_tokenizer", tokenizer) + + def __getattr__(self, attr): + if attr.startswith("_"): + return object.__getattribute__(self, attr) + return getattr(object.__getattribute__(self, "_tokenizer"), attr) + + +_FakeTokenizerWrapper.__name__ = "TokenizerWrapper" + + +def test_mlx_tokenizer_wrapper_unwrapped_for_preset_markers(): + # Markers live on the inner HF tokenizer that the wrapper hides; the helper + # must unwrap so the preset bare-call path still fires on MLX. + class _Tok: + _unsloth_input_part = "" + _unsloth_output_part = "" + + trainer = _Trainer() + trainer.tokenizer = _FakeTokenizerWrapper(_Tok()) + train_fn = _Recorder() + + _, applied = apply_completion_masking( + trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_fail + ) + assert applied is True + assert train_fn.calls == [{}] # bare call, stored parts + + +def test_mlx_tokenizer_wrapper_unwrapped_for_detection(): + # Detection must see the real tokenizer, not the wrapper, so it does not + # depend on the loader's __call__ patch. + class _Tok: + pass + + inner = _Tok() + trainer = _Trainer() + trainer.tokenizer = _FakeTokenizerWrapper(inner) + train_fn = _Recorder() + seen = [] + + def detect(processor): + seen.append(processor) + return "", "" + + _, applied = apply_completion_masking( + trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = detect + ) + assert applied is True + assert seen == [inner] diff --git a/studio/backend/tests/test_compute_buffer.py b/studio/backend/tests/test_compute_buffer.py index 5d14c5c5bd..8408f8203d 100644 --- a/studio/backend/tests/test_compute_buffer.py +++ b/studio/backend/tests/test_compute_buffer.py @@ -65,12 +65,14 @@ def _backend( vocab = 248320, embd = 5120, mla = None, + arch = None, ): """Backend with just the dims the compute-buffer estimate reads.""" b = LlamaCppBackend.__new__(LlamaCppBackend) b._vocab_size = vocab b._embedding_length = embd b._key_length_mla = mla # non-None -> MLA (compressed attention) + b._architecture = arch # GGUF general.architecture (e.g. 'deepseek4') return b @@ -290,3 +292,62 @@ class TestContextBufferMLA: b = _backend(embd = 6144, mla = 256) est = b._compute_buffer_ctx_bytes(754688, cache_type_kv = "q8_0") / MIB assert est <= 4141 * 1.7 + + +class TestContextBufferDSV4: + """DeepSeek-V4 (deepseek4) reserves a large lightning-indexer / sparse-attention + compute buffer the KQ-mask and MLA rates miss (present even with an f16 cache). + Measured on UD-Q4_K_XL (ub=512): ~2 GiB at 16k ctx, ~65.5 GiB at 1M. The auto-fit + must see this so it does not commit the full 1M train context and OOM (spilling + to CPU at ~4 tok/s).""" + + _MEASURED_1M_GIB = 65.5 # 70353790464 B compute-graph reserve that OOM'd at 1M ctx + GIB = 1024**3 + + def test_covers_measured_1m_buffer(self): + b = _backend(embd = 4096, arch = "deepseek4") + gib = b._compute_buffer_ctx_bytes(1048576, cache_type_kv = "f16") / self.GIB + assert gib >= self._MEASURED_1M_GIB, f"under-reserved {gib:.1f} < {self._MEASURED_1M_GIB}" + + def test_not_wildly_over_at_1m(self): + # Within ~1.3x of measured so the fit still grants a large (~256k) context. + b = _backend(embd = 4096, arch = "deepseek4") + gib = b._compute_buffer_ctx_bytes(1048576, cache_type_kv = "f16") / self.GIB + assert gib <= self._MEASURED_1M_GIB * 1.3 + + def test_fires_for_f16_cache(self): + # The bug: an f16 (default) cache took the tiny mask-only path. DSV4 must + # reserve GiB, not the ~MiB a non-DSV4 model reserves at the same ctx. + dsv4 = _backend(embd = 4096, arch = "deepseek4")._compute_buffer_ctx_bytes( + 262144, cache_type_kv = "f16" + ) + other = _backend(embd = 4096, arch = "qwen3")._compute_buffer_ctx_bytes( + 262144, cache_type_kv = "f16" + ) + assert dsv4 > 40 * other + + def test_cache_type_independent(self): + # Indexer scratch is present for an f16 and a quantized cache alike. + b = _backend(embd = 4096, arch = "deepseek4") + assert b._compute_buffer_ctx_bytes( + 262144, cache_type_kv = "f16" + ) == b._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0") + + def test_flat_floor_at_small_ctx(self): + # ~2 GiB indexer scratch present even at tiny ctx (covers the measured 16k ~2 GiB). + b = _backend(embd = 4096, arch = "deepseek4") + assert b._compute_buffer_ctx_bytes(16384, cache_type_kv = "f16") / self.GIB >= 2.0 + + def test_scales_with_context_and_ubatch(self): + b = _backend(embd = 4096, arch = "deepseek4") + assert b._compute_buffer_ctx_bytes(131072) > b._compute_buffer_ctx_bytes(65536) + assert b._compute_buffer_ctx_bytes(131072, n_ubatch = 1024) > b._compute_buffer_ctx_bytes( + 131072, n_ubatch = 256 + ) + + def test_non_dsv4_unchanged(self): + # Regression guard: a non-deepseek4 model keeps the mask-only f16 rate. + b = _backend(embd = 4096, arch = "llama") + per_tok = b._compute_buffer_ctx_bytes(100000, cache_type_kv = "f16") / 100000 + expected = 512 * 2 * LlamaCppBackend._CTX_COMPUTE_F16_MASK_SAFETY + assert per_tok == pytest.approx(expected, rel = 1e-6) diff --git a/studio/backend/tests/test_data_recipe_seed.py b/studio/backend/tests/test_data_recipe_seed.py index 09e22116ed..58bbd24061 100644 --- a/studio/backend/tests/test_data_recipe_seed.py +++ b/studio/backend/tests/test_data_recipe_seed.py @@ -124,3 +124,100 @@ def test_unstructured_upload_import_errors_stay_generic(monkeypatch, tmp_path, e assert result.status == "error" assert result.error == "Text extraction failed." assert _block_files(seed_route) == [] + + +_TEST_UPLOAD_UID = "0f" * 16 + + +def test_remove_unstructured_block_deletes_directory(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + _run_upload(seed_route, "notes.txt", b"hello", block_id = _TEST_UPLOAD_UID) + assert _block_files(seed_route, _TEST_UPLOAD_UID) != [] + + result = asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID)) + + assert result == {"status": "ok", "deleted": True} + assert not (seed_route.UNSTRUCTURED_UPLOAD_ROOT / _TEST_UPLOAD_UID).exists() + + +def test_remove_unstructured_block_missing_directory_is_ok(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + + result = asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID)) + + assert result == {"status": "ok", "deleted": False} + + +def test_remove_unstructured_block_rejects_unsafe_ids(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + + with pytest.raises(seed_route.HTTPException) as exc: + asyncio.run(seed_route.remove_unstructured_block("../escape")) + + assert exc.value.status_code == 400 + + +def test_remove_unstructured_block_rejects_legacy_node_ids(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + _run_upload(seed_route, "notes.txt", b"hello", block_id = "n1") + assert _block_files(seed_route, "n1") != [] + + with pytest.raises(seed_route.HTTPException) as exc: + asyncio.run(seed_route.remove_unstructured_block("n1")) + + assert exc.value.status_code == 400 + assert _block_files(seed_route, "n1") != [] + + +def test_remove_unstructured_block_rejects_symlink_escape(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + outside = tmp_path / "outside" + outside.mkdir() + (outside / "victim.txt").write_text("keep me") + root = seed_route.UNSTRUCTURED_UPLOAD_ROOT + root.mkdir(parents = True) + (root / _TEST_UPLOAD_UID).symlink_to(outside) + + with pytest.raises(seed_route.HTTPException) as exc: + asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID)) + + assert exc.value.status_code == 400 + assert (outside / "victim.txt").exists() + + +def test_remove_unstructured_block_fails_if_directory_remains(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + root = seed_route.UNSTRUCTURED_UPLOAD_ROOT + block_dir = root / _TEST_UPLOAD_UID + block_dir.mkdir(parents = True) + (block_dir / "victim.txt").write_text("keep me") + + calls = [] + + def noop_rmtree(path, *args, **kwargs): + calls.append((path, args, kwargs)) + + monkeypatch.setattr(seed_route.shutil, "rmtree", noop_rmtree) + + with pytest.raises(seed_route.HTTPException) as exc: + asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID)) + + assert calls + assert exc.value.status_code == 500 + assert block_dir.exists() + + +def test_total_upload_quota_is_scoped_per_block(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + monkeypatch.setattr(seed_route, "UNSTRUCTURED_RECIPE_UPLOAD_TOTAL_MAX_BYTES", 10) + + first = _run_upload(seed_route, "a.txt", b"123456789") + assert first.status == "ok" + + with pytest.raises(seed_route.HTTPException) as exc: + _run_upload(seed_route, "b.txt", b"123") + assert exc.value.status_code == 413 + + # Another block starts with its own untouched budget. + other = _run_upload(seed_route, "c.txt", b"123", block_id = "other") + assert other.status == "ok" diff --git a/studio/backend/tests/test_deepseek_v4_thinking_effort.py b/studio/backend/tests/test_deepseek_v4_thinking_effort.py new file mode 100644 index 0000000000..19808ad0d7 --- /dev/null +++ b/studio/backend/tests/test_deepseek_v4_thinking_effort.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""DeepSeek-V4-Flash reasoning toggle: None / High / Max. + +The GGUF template gates thinking with ``enable_thinking`` and only branches +``reasoning_effort`` on ``'max'`` (an escalation layered over plain thinking). +Detection used to return the single level ``['max']``, so the UI collapsed to +None / Max and the plain-thinking tier was unreachable. Detection now surfaces +``'high'`` as that plain tier, giving None / High / Max. These tests pin the +classifier, the GLM-style parity case, and the full request-kwargs -> rendered +prompt path for each state (the model itself is too large to load here). +""" + +from __future__ import annotations + +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +_backend_root = Path(__file__).resolve().parent.parent +if str(_backend_root) not in sys.path: + sys.path.insert(0, str(_backend_root)) + + +# Faithful slice of the DeepSeek-V4-Flash GGUF template: the enable_thinking +# gate, the sole ``reasoning_effort == 'max'`` escalation, and the plain-think +# fallback. Any non-'max' effort renders as ordinary thinking. +DEEPSEEK_V4_TEMPLATE = """ +{%- if not thinking is defined -%} + {%- if enable_thinking is defined -%} + {%- set thinking = enable_thinking -%} + {%- else -%} + {%- set thinking = false -%} + {%- endif -%} +{%- endif -%} +{%- if not reasoning_effort is defined -%} + {%- set reasoning_effort = none -%} +{%- endif -%} +{{- bos_token -}} +{%- if thinking and reasoning_effort == 'max' -%} + {{- 'Reasoning Effort: Absolute maximum with no shortcuts permitted.\\n\\n' -}} +{%- endif -%} +{%- for message in messages -%} + {{- '<|User|>' + (message['content'] or '') -}} +{%- endfor -%} +{%- if add_generation_prompt -%} + {{- '<|Assistant|>' -}} + {%- if thinking -%}{{- '' -}}{%- else -%}{{- '' -}}{%- endif -%} +{%- endif -%} +""" + + +# GLM-5.2-style: branches on two effort literals, so 'high' already exists as +# the sub-'max' tier and detection must leave the pair untouched. +GLM_STYLE_TEMPLATE = """ +{%- if enable_thinking -%} + {%- if reasoning_effort == 'high' -%}{{- 'H' -}} + {%- elif reasoning_effort == 'max' -%}{{- 'M' -}} + {%- endif -%} +{%- endif -%} +""" + + +# A ['max']-only template under a non-deepseek id: the synthetic 'high' is scoped +# to deepseek-v4, so this must stay ['max'] (no phantom 'high'). +NON_DEEPSEEK_MAX_ONLY_TEMPLATE = DEEPSEEK_V4_TEMPLATE + + +# A template whose sole effort literal is a sub-'max' level: the guard targets +# only the ['max']-alone case, so a lone 'high' stays a singleton. +HIGH_ONLY_TEMPLATE = """ +{%- if enable_thinking and reasoning_effort == 'high' -%}{{- 'H' -}}{%- endif -%} +""" + + +def _render(template: str, **kwargs) -> str: + jinja2 = pytest.importorskip("jinja2") + env = jinja2.Environment() + tmpl = env.from_string(template) + return tmpl.render(bos_token = "", add_generation_prompt = True, **kwargs) + + +# -- Classifier ------------------------------------------------------- + + +def test_deepseek_v4_surfaces_high_as_plain_tier(): + """Sole 'max' escalation expands to ['high', 'max'] so None/High/Max show.""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash") + assert flags["supports_reasoning"] is True + assert flags["reasoning_style"] == "enable_thinking_effort" + assert flags["reasoning_effort_levels"] == ["high", "max"] + + +def test_glm_style_two_level_template_unchanged(): + """A template that already names a sub-'max' tier is left as-is.""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(GLM_STYLE_TEMPLATE, "unsloth/GLM-5.2") + assert flags["reasoning_style"] == "enable_thinking_effort" + assert flags["reasoning_effort_levels"] == ["high", "max"] + + +def test_synthetic_high_scoped_to_deepseek_v4(): + """The same ['max']-only template under a non-deepseek id keeps ['max'].""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(NON_DEEPSEEK_MAX_ONLY_TEMPLATE, "vendor/OtherHybrid-GGUF") + assert flags["reasoning_effort_levels"] == ["max"] + + +def test_guard_does_not_fire_for_sub_max_singleton(): + """The expansion targets only ['max']; a lone 'high' stays a singleton.""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(HIGH_ONLY_TEMPLATE, "custom/high-only") + assert flags["reasoning_effort_levels"] == ["high"] + + +# -- Request kwargs -> rendered prompt, for each state ---------------- + + +def _kwargs_for(flags: dict, enable_thinking, reasoning_effort): + """Drive the real backend method with a shim carrying the detected flags.""" + from core.inference.llama_cpp import LlamaCppBackend + + shim = SimpleNamespace( + _supports_reasoning = flags["supports_reasoning"], + _reasoning_always_on = flags["reasoning_always_on"], + _reasoning_style = flags["reasoning_style"], + _reasoning_effort_levels = flags["reasoning_effort_levels"], + _supports_preserve_thinking = flags["supports_preserve_thinking"], + ) + build = LlamaCppBackend._request_reasoning_kwargs.__get__(shim) + return build(enable_thinking, reasoning_effort, None) or {} + + +def _flags(): + from core.inference.llama_cpp import detect_reasoning_flags + return detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash") + + +def test_none_state_renders_non_thinking(): + """UI 'None' -> enable_thinking=false -> closed , no preamble.""" + kwargs = _kwargs_for(_flags(), enable_thinking = False, reasoning_effort = None) + assert kwargs == {"enable_thinking": False} + out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + assert out.endswith("") + assert "Absolute maximum" not in out + + +def test_high_state_renders_plain_thinking(): + """UI 'High' -> et=true, effort=high -> open , no max preamble.""" + kwargs = _kwargs_for(_flags(), enable_thinking = True, reasoning_effort = "high") + assert kwargs == {"enable_thinking": True, "reasoning_effort": "high"} + out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + assert out.endswith("") + assert "Absolute maximum" not in out + + +def test_max_state_injects_max_preamble(): + """UI 'Max' -> et=true, effort=max -> open plus the max preamble.""" + kwargs = _kwargs_for(_flags(), enable_thinking = True, reasoning_effort = "max") + assert kwargs == {"enable_thinking": True, "reasoning_effort": "max"} + out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + assert out.endswith("") + assert "Absolute maximum" in out + + +def test_high_effort_alone_enables_thinking(): + """API caller sending only reasoning_effort='high' (no enable_thinking) still + gets thinking on, so the newly exposed High mode renders correctly.""" + kwargs = _kwargs_for(_flags(), enable_thinking = None, reasoning_effort = "high") + assert kwargs == {"enable_thinking": True, "reasoning_effort": "high"} + out = _render(DEEPSEEK_V4_TEMPLATE, messages = [{"role": "user", "content": "hi"}], **kwargs) + assert out.endswith("") + assert "Absolute maximum" not in out diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py new file mode 100644 index 0000000000..940b35d7ba --- /dev/null +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -0,0 +1,365 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""The RAG embedding model must pass the malware/pickle gate before it is persisted or +loaded. A flagged repo (or any repo saved with force) previously reached +SentenceTransformer unscanned, bypassing the normal model-load protections.""" + +from pathlib import Path +import sys +import types as _types + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +import routes.settings as settings + + +class _Decision: + def __init__(self, blocked): + self.blocked = blocked + + +def _security_stub(blocked): + mod = _types.ModuleType("utils.security") + mod.evaluate_file_security = lambda *a, **k: _Decision(blocked) + mod.security_load_subdirs = lambda *a, **k: () + return mod + + +@pytest.fixture +def client(monkeypatch): + # The settings scan unions in the ST module dirs read from modules.json; keep it + # offline and deterministic for the endpoint tests that use this fixture. + import core.rag.embeddings as embeddings + + monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ()) + saved: dict = {} + monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) + monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + monkeypatch.setattr(settings, "_llama_backend_active", lambda: False) + monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) + monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) + monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + return TestClient(app, raise_server_exceptions = False), saved + + +def test_flagged_repo_is_blocked_even_with_force(client, monkeypatch): + c, saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) + r = c.put( + "/embedding-model", json = {"embedding_model": "attacker/malicious-embed", "force": True} + ) + # 403, not the forceable 409, so the client does not offer "save anyway". + assert r.status_code == 403 + assert "model" not in saved # force must not persist a flagged repo + + +def test_flagged_repo_is_blocked_without_force(client, monkeypatch): + c, saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) + r = c.put("/embedding-model", json = {"embedding_model": "attacker/malicious-embed"}) + assert r.status_code == 403 + assert "model" not in saved + + +def test_hard_block_uses_non_forceable_status(client, monkeypatch): + # The forceable verification path uses 409; the hard security block must be distinct + # (403) so the frontend never routes it into the "save anyway" force flow. + c, _saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) + blocked = c.put("/embedding-model", json = {"embedding_model": "attacker/malicious-embed"}) + assert blocked.status_code == 403 + + # A verification failure (not-an-embedding-model) stays forceable at 409. + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) + monkeypatch.setattr(settings, "is_embedding_model", lambda *a, **k: False, raising = False) + import utils.models as _models + + monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: False) + unverified = c.put("/embedding-model", json = {"embedding_model": "acme/not-an-embedder"}) + assert unverified.status_code == 409 + + +def test_llama_backend_skips_the_st_pickle_scan(monkeypatch): + # On the llama-server backend the embedder loads GGUF (inert), not the ST repo's + # pickle, so a flagged ST repo with a clean GGUF companion must not be rejected here. + saved: dict = {} + monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) + monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + monkeypatch.setattr(settings, "_llama_backend_active", lambda: True) + monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) + monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) + monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + # force skips the GGUF availability checks; the ST pickle gate is what we assert is skipped. + called = {"scanned": False} + mod = _types.ModuleType("utils.security") + + def _fail(*a, **k): + called["scanned"] = True + return _Decision(True) + + mod.evaluate_file_security = _fail + mod.security_load_subdirs = lambda *a, **k: () + monkeypatch.setitem(sys.modules, "utils.security", mod) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + c = TestClient(app, raise_server_exceptions = False) + r = c.put( + "/embedding-model", + json = {"embedding_model": "attacker/flagged-st-clean-gguf", "force": True}, + ) + assert r.status_code == 200 + assert called["scanned"] is False # the ST pickle scan never ran on the llama path + assert saved.get("model") == "attacker/flagged-st-clean-gguf" + + +def test_runtime_llama_fallback_skips_the_st_pickle_scan(monkeypatch): + # auto resolves to sentence-transformers (GPU present) but the embedder fell back to + # llama-server at runtime (torch/CUDA load or encode failure), so the process now loads + # only inert GGUF. The real _llama_backend_active() must reflect that cached fallback, + # so a flagged ST repo with a clean GGUF companion must not be hard-blocked here. + import core.rag.embeddings as embeddings + from core.rag.embed_llama_server import LlamaServerBackend + + # Simulate the runtime fallback: the process-wide backend is a LlamaServerBackend even + # though the auto resolver would still say sentence-transformers. + monkeypatch.setattr(embeddings, "_backend", LlamaServerBackend()) + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers") + monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ()) + + saved: dict = {} + monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) + monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + # Deliberately do NOT monkeypatch settings._llama_backend_active: this test exercises the + # real delegation to embeddings.active_backend_is_llama() so the cached fallback is honored. + monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) + monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) + monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + + called = {"scanned": False} + mod = _types.ModuleType("utils.security") + + def _fail(*a, **k): + called["scanned"] = True + return _Decision(True) + + mod.evaluate_file_security = _fail + mod.security_load_subdirs = lambda *a, **k: () + monkeypatch.setitem(sys.modules, "utils.security", mod) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + c = TestClient(app, raise_server_exceptions = False) + r = c.put( + "/embedding-model", + json = {"embedding_model": "attacker/flagged-st-clean-gguf", "force": True}, + ) + assert r.status_code == 200 + assert called["scanned"] is False # the ST pickle scan never ran on the llama fallback + assert saved.get("model") == "attacker/flagged-st-clean-gguf" + + +def test_active_backend_is_llama_reflects_cache_and_resolver(monkeypatch): + # active_backend_is_llama() reports the ACTUAL built backend when one exists, and defers + # to the resolver (fresh-process behavior) when none has been built yet. + import core.rag.embeddings as embeddings + import core.rag.config as rag_config + from core.rag.embed_llama_server import LlamaServerBackend + + # A cached llama backend wins even when auto would resolve to sentence-transformers. + monkeypatch.setattr(rag_config, "EMBED_BACKEND", "auto") + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers") + monkeypatch.setattr(embeddings, "_backend", LlamaServerBackend()) + assert embeddings.active_backend_is_llama() is True + + # A cached ST backend reports False even when the resolver now picks llama, so its + # pickle stays gated (the cached backend, not the resolver, is what actually embeds). + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "llama-server") + monkeypatch.setattr(embeddings, "_backend", embeddings._SentenceTransformersBackend()) + assert embeddings.active_backend_is_llama() is False + + # No cached backend -> the resolver decides, unchanged from before. + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers") + monkeypatch.setattr(embeddings, "_backend", None) + assert embeddings.active_backend_is_llama() is False # auto -> sentence-transformers + + monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "llama-server") + assert embeddings.active_backend_is_llama() is True # auto -> llama-server + + # An explicit (non-auto) key is honored verbatim without a cached backend. + monkeypatch.setattr(rag_config, "EMBED_BACKEND", "llama-server") + assert embeddings.active_backend_is_llama() is True + + +def test_settings_scan_scopes_module_subdirs(monkeypatch): + # The settings scan must pass the ST module dirs (0_Transformer/) as load roots so a + # pickle directly under one blocks; assert those subdirs reach evaluate_file_security. + saved: dict = {} + monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) + monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + monkeypatch.setattr(settings, "_llama_backend_active", lambda: False) + monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) + monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) + monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + + import core.rag.embeddings as embeddings + + monkeypatch.setattr( + embeddings, "_st_module_subdirs", lambda name, token = None: ("0_Transformer",) + ) + seen = {} + + def _capture(*a, **k): + seen["subdirs"] = tuple(k.get("load_subdirs") or ()) + return _Decision(False) + + mod = _types.ModuleType("utils.security") + mod.security_load_subdirs = lambda *a, **k: () + mod.evaluate_file_security = _capture + monkeypatch.setitem(sys.modules, "utils.security", mod) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + c = TestClient(app, raise_server_exceptions = False) + r = c.put( + "/embedding-model", json = {"embedding_model": "acme/embed-with-module-dir", "force": True} + ) + assert r.status_code == 200 + assert "0_Transformer" in seen["subdirs"] + + +def test_clean_repo_saves_under_force(client, monkeypatch): + c, saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) + r = c.put("/embedding-model", json = {"embedding_model": "acme/clean-embed", "force": True}) + assert r.status_code == 200 + assert saved.get("model") == "acme/clean-embed" + + +def test_load_sink_refuses_flagged_model(monkeypatch): + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) + import core.rag.embeddings as embeddings + with pytest.raises(embeddings.UnsafeEmbeddingModelError): + embeddings._guard_model_security("attacker/malicious-embed") + + +def test_load_sink_allows_clean_model(monkeypatch): + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) + import core.rag.embeddings as embeddings + embeddings._guard_model_security("acme/clean-embed") # no raise + + +def test_sink_threads_ambient_token_into_scan(monkeypatch): + # A gated repo set via env/default has no request token; the guard must feed the + # loader's own token to the scan, or it fails open for the repo that still loads. + seen = {} + mod = _types.ModuleType("utils.security") + mod.security_load_subdirs = ( + lambda name, token = None: seen.setdefault("subdirs_token", token) or () + ) + mod.evaluate_file_security = lambda *a, **k: seen.setdefault( + "scan_token", k.get("hf_token") + ) or _Decision(False) + monkeypatch.setitem(sys.modules, "utils.security", mod) + import core.rag.embeddings as embeddings + + monkeypatch.setattr(embeddings, "_ambient_hf_token", lambda: "hf_ambient") + embeddings._guard_model_security("acme/gated-embed") + assert seen["scan_token"] == "hf_ambient" + assert seen["subdirs_token"] == "hf_ambient" + + +def test_sink_scopes_st_module_subdirs_into_scan(monkeypatch): + # A flagged pickle directly under a Transformer module dir (0_Transformer/) must + # reach the scan as a load root; assert the guard unions the module dirs into + # load_subdirs so evaluate_file_security treats such a pickle as root-level. + seen = {} + + def _capture(*a, **k): + seen["subdirs"] = tuple(k.get("load_subdirs") or ()) + return _Decision(False) + + mod = _types.ModuleType("utils.security") + mod.security_load_subdirs = lambda name, token = None: () + mod.evaluate_file_security = _capture + monkeypatch.setitem(sys.modules, "utils.security", mod) + import core.rag.embeddings as embeddings + + monkeypatch.setattr(embeddings, "_ambient_hf_token", lambda: None) + monkeypatch.setattr( + embeddings, "_st_module_subdirs", lambda name, token = None: ("0_Transformer",) + ) + embeddings._guard_model_security("acme/embed-with-module-dir") + assert "0_Transformer" in seen["subdirs"] + + +def test_st_module_subdirs_reads_local_modules_json(tmp_path, monkeypatch): + # The helper must parse each module's non-empty "path" from a local repo's + # modules.json and drop the root-level ("") Transformer entry. + import json + import core.rag.embeddings as embeddings + + (tmp_path / "modules.json").write_text( + json.dumps( + [ + {"idx": 0, "name": "0", "path": "0_Transformer", "type": "..."}, + {"idx": 1, "name": "1", "path": "1_Pooling", "type": "..."}, + {"idx": 2, "name": "2", "path": "", "type": "..."}, + ] + ) + ) + subdirs = embeddings._st_module_subdirs(str(tmp_path), None) + assert subdirs == ("0_Transformer", "1_Pooling") + + +def test_st_module_subdirs_swallows_errors(monkeypatch): + # Any failure (no modules.json, offline, malformed) returns () so the guard never + # bricks the embedder. + import huggingface_hub + import core.rag.embeddings as embeddings + + def _boom(*a, **k): + raise RuntimeError("offline") + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", _boom) + assert embeddings._st_module_subdirs("acme/no-such-repo-xyz", None) == () + + +def test_security_block_is_not_swallowed_by_llama_fallback(monkeypatch): + # The ST encode fallback must re-raise a security block, not swap to llama-server. + import core.rag.embeddings as embeddings + + def _boom(*a, **k): + raise embeddings.UnsafeEmbeddingModelError("flagged") + + monkeypatch.setattr(embeddings, "_st_encode", _boom) + monkeypatch.setattr( + embeddings, + "_switch_to_llama_fallback", + lambda err: pytest.fail("security block must not fall back to llama-server"), + ) + with pytest.raises(embeddings.UnsafeEmbeddingModelError): + embeddings._SentenceTransformersBackend().encode(["hi"]) diff --git a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py index 8df8d37a52..e3055d2127 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -1,15 +1,8 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Edge cases in Gemma-native tool-call parsing. - -Covers two failure modes: - 1. A bare (unquoted) string argument that contains a comma, e.g. - ``location:New York, NY`` -- the comma must not be treated as the next - key boundary, or the whole call is dropped. - 2. A tool-call marker that appears INSIDE another call's argument string is - data, not a real call, so it must not be promoted to a second tool call. -""" +"""Gemma-native tool-call parsing edge cases: commas inside bare string values, +and markers inside another call's argument data staying data.""" from __future__ import annotations @@ -21,7 +14,11 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) -from core.inference.tool_call_parser import parse_tool_calls_from_text +from core.inference.tool_call_parser import ( + _gemma_parse_value, + parse_tool_calls_from_text, +) +from core.tool_healing import strip_tool_call_markup def _args(call: dict) -> dict: @@ -40,14 +37,22 @@ def test_bare_string_argument_with_comma_is_kept(): def test_normal_multi_key_arguments_still_split(): calls = parse_tool_calls_from_text('<|tool_call>call:f{a:1,b:hello,c:"x,y"}') assert len(calls) == 1, calls - # Numbers stay numeric, bare strings get quoted, an explicit quoted comma - # stays inside its value. assert _args(calls[0]) == {"a": 1, "b": "hello", "c": "x,y"} +def test_empty_bare_value_becomes_empty_string_not_dropped(): + # An empty bare value (``{query:}``) must serialise as ``""`` (``{"query":}`` is invalid JSON and dropped the call). + calls = parse_tool_calls_from_text("<|tool_call>call:search{query:,unit:celsius}") + assert len(calls) == 1, calls + assert _args(calls[0]) == {"query": "", "unit": "celsius"} + + only = parse_tool_calls_from_text("<|tool_call>call:get{q:}") + assert len(only) == 1, only + assert _args(only[0]) == {"q": ""} + + def test_bare_value_with_timestamps_after_comma_is_kept(): - # A comma followed by digits-then-colon (a timestamp/ratio) is value text, - # not a new key, so the whole query must be preserved as one argument. + # A comma before digits-then-colon (timestamp/ratio) is value text, not a key. calls = parse_tool_calls_from_text( "<|tool_call>call:remind{query:meet at 10:00, 11:00 tomorrow,priority:high}" ) @@ -55,9 +60,16 @@ def test_bare_value_with_timestamps_after_comma_is_kept(): assert _args(calls[0]) == {"query": "meet at 10:00, 11:00 tomorrow", "priority": "high"} +def test_wrapperless_bare_value_with_timestamps_after_comma_is_kept(): + # The wrapper-less Gemma form (no <|tool_call> markers) goes through the + # _gemma_parse_stripped_body scanner and its _GEMMA_KEY_RE. + calls = parse_tool_calls_from_text("call:web_search{query:meet at 10:00, 11:00 tomorrow}") + assert len(calls) == 1, calls + assert calls[0]["function"]["name"] == "web_search" + assert _args(calls[0]) == {"query": "meet at 10:00, 11:00 tomorrow"} + + def test_marker_inside_json_argument_is_not_a_second_call(): - # A python call whose `code` argument contains a Gemma marker string. The - # marker is data and must not execute as a second `terminal` call. content = ( '{"name":"python","arguments":{"code":' '"x = 1 # <|tool_call>call:terminal{command:ls}"}}' @@ -75,8 +87,6 @@ def test_two_separate_gemma_calls_both_parse(): def test_mixed_format_calls_preserve_document_order(): - # A Gemma-native call precedes a JSON-format call in the text; tools execute - # in returned order, so `create` must come before `read`. content = ( "<|tool_call>call:create{path:a} then " '{"name":"read","arguments":{"path":"a"}}' @@ -86,8 +96,6 @@ def test_mixed_format_calls_preserve_document_order(): def test_json_marker_inside_gemma_argument_is_not_a_second_call(): - # The reverse of the JSON-outer case: a JSON-style marker inside a Gemma - # call's quoted argument is code text, not a second `terminal` call. content = ( '<|tool_call>call:python{code:<|"|>' 'print({"name":"terminal","arguments":{"command":"ls"}})' @@ -98,18 +106,14 @@ def test_json_marker_inside_gemma_argument_is_not_a_second_call(): def test_nested_gemma_marker_in_unquoted_arg_does_not_run_inner_call(): - # An UNQUOTED Gemma value containing a literal marker: the outer object fails - # to normalize (the inner braces/marker break the JSON), but the inner marker - # is nested in the outer candidate span, so it must not be promoted to a - # standalone `terminal` call. The safe outcome is no executed tool call. + # An UNQUOTED Gemma value containing a literal marker: the marker is nested in the outer + # candidate span, so it must not be promoted to a standalone `terminal` call (no tool call). content = "<|tool_call>call:python{code:<|tool_call>call:terminal{command:ls}}" calls = parse_tool_calls_from_text(content) assert "terminal" not in [c["function"]["name"] for c in calls], calls def test_bare_string_array_argument_is_quoted(): - # Gemma may emit an array of bare strings without per-element quotes; they - # must be quoted so the call is not dropped. calls = parse_tool_calls_from_text("<|tool_call>call:label{labels:[bug,ui]}") assert len(calls) == 1, calls assert _args(calls[0]) == {"labels": ["bug", "ui"]} @@ -123,8 +127,6 @@ def test_array_keeps_numbers_and_quoted_elements(): def test_array_of_objects_is_normalised(): - # Arrays of objects are a common tool-schema shape; their (unquoted) keys and - # bare values must be normalised too, not left verbatim, or the call drops. calls = parse_tool_calls_from_text( "<|tool_call>call:batch{items:[{path:a,mode:r},{path:b,mode:w}]}" ) @@ -138,9 +140,6 @@ def test_nested_array_elements_are_normalised(): def test_gemma_marker_inside_xml_parameter_is_not_a_second_call(): - # An XML-style call whose value contains a - # Gemma marker: the marker is the parameter's data, not a separate terminal - # call, so only the python call must be returned. content = ( "" "x = 1 # <|tool_call>call:terminal{command:ls}" @@ -159,3 +158,254 @@ def test_json_marker_inside_xml_parameter_is_not_a_second_call(): ) calls = parse_tool_calls_from_text(content) assert [c["function"]["name"] for c in calls] == ["python"], calls + + +def test_unclosed_think_literal_inside_tool_argument_does_not_hide_later_call(): + # A literal inside a completed call's arguments is argument data; both calls must parse. + text = '[TOOL_CALLS]a{"x":"literal marker"} b[ARGS]{"y":2}' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["a", "b"], calls + + +def test_real_think_block_with_rehearsal_inside_still_skips_only_the_rehearsal(): + # A genuine reasoning block still hides its rehearsal while a real call after it parses. + text = 'web_search[ARGS]{"q":"draft"}real[ARGS]{"q":"go"}' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["real"], calls + + +def test_wrapperless_nested_object_argument_is_parsed(): + # skip_special_tokens stream: wrapper and <|"|> markers stripped, so a nested object arrives bare. + calls = parse_tool_calls_from_text("call:f{loc:{city:NYC},n:3}") + assert len(calls) == 1 + assert _args(calls[0]) == {"loc": {"city": "NYC"}, "n": 3} + + +def test_wrapperless_array_argument_is_parsed(): + calls = parse_tool_calls_from_text("call:label{labels:[bug,ui],n:2}") + assert len(calls) == 1 + assert _args(calls[0]) == {"labels": ["bug", "ui"], "n": 2} + + +def test_wrapperless_deeply_nested_object_and_array_are_preserved(): + # The single-pass parser must keep multi-level nesting (objects inside + # objects, arrays inside arrays) intact, not flatten or drop it. + calls = parse_tool_calls_from_text( + "call:f{loc:{city:NYC,geo:{lat:1,lng:2}},tags:[a,b,[c,d]],n:3}" + ) + assert len(calls) == 1 + assert _args(calls[0]) == { + "loc": {"city": "NYC", "geo": {"lat": 1, "lng": 2}}, + "tags": ["a", "b", ["c", "d"]], + "n": 3, + } + + +def test_gemma_parse_array_advances_on_stray_brace(): + # Regression: a stray '}' / ']' / ',' where an array element is expected must + # not stall _gemma_parse_value at the same index (it looped forever before). + from core.inference.tool_call_parser import _gemma_parse_array + + items, end, closed = _gemma_parse_array("[a,}]", 0) + assert end == 5 and closed is True # consumed through the closing ']' + assert items[0] == "a" + + +def test_gemma_close_marker_inside_quoted_arg_is_not_leaked_when_stripping(): + # Parse keeps the quoted close marker as data; strip removes the whole span. + text = '<|tool_call>call:python{code:<|"|>print("")<|"|>}' + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1, calls + assert _args(calls[0]) == {"code": 'print("")'} + assert strip_tool_call_markup("before " + text + " after") == "before after" + assert strip_tool_call_markup("before " + text + " after", final = True) == "before after" + + +def test_nested_xml_in_malformed_gemma_call_does_not_execute(): + # The failed Gemma candidate's span still covers its nested . + text = ( + "<|tool_call>call:outer{code:id" + ", broken:{x}}" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_unbalanced_gemma_call_with_xml_does_not_execute(): + # Unclosed braces cover to EOF, so the trailing is excluded. + text = ( + "<|tool_call>call:outer{code:" + "id" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_standalone_function_xml_still_parses(): + text = "id" + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["terminal"], calls + + +def test_xml_between_braces_and_close_marker_does_not_execute(): + # Coverage runs to the close marker, so in the gap is data. + text = ( + "<|tool_call>call:outer{broken:{x}}" + "id" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_balanced_inner_call_inside_unclosed_outer_does_not_execute(): + text = "<|tool_call>call:outer{code:<|tool_call>call:terminal{command:id}" + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_strip_preserves_text_after_malformed_gemma_close(): + # Junk before the close is a malformed span: strip through it, keep the tail. + text = "pre <|tool_call>call:t{a:1} note post" + assert strip_tool_call_markup(text) == "pre post" + assert strip_tool_call_markup(text, final = True) == "pre post" + + +def test_malformed_closed_gemma_span_is_stripped(): + assert ( + strip_tool_call_markup('before <|tool_call>{"name":"x"} after') + == "before after" + ) + + +def test_valid_call_after_missing_close_is_recovered(): + # A close-less call covers only its braces, so the later call is recovered. + text = "<|tool_call>call:a{x:1} <|tool_call>call:b{y:2}" + names_inc = [ + c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = True) + ] + assert "b" in names_inc, names_inc + names_strict = [ + c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = False) + ] + assert names_strict == ["b"], names_strict + + +def test_strip_non_final_keeps_incomplete_gemma_block(): + text = "before <|tool_call>call:t{" + assert strip_tool_call_markup(text) == text + assert strip_tool_call_markup(text, final = True) == "before" + + +def test_json_call_between_gemma_braces_and_close_does_not_execute(): + # A JSON call between the outer's braces and its close is covered data. + text = ( + "<|tool_call>call:outer{broken:{x}}" + '{"name":"terminal","arguments":{"command":"id"}}' + "" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_gemma_call_between_gemma_braces_and_close_does_not_execute(): + # Same escape with a Gemma-native inner marker. + text = "<|tool_call>call:outer{broken:{x}}<|tool_call>call:terminal{command:id}" + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_strip_final_keeps_text_after_closed_xml_with_inner_gemma_opener(): + # The to-EOF Gemma sweep must not eat visible text after . + text = ( + 'before print("<|tool_call>") after' + ) + assert strip_tool_call_markup(text, final = True) == "before after" + assert strip_tool_call_markup(text) == "before after" + + +def test_strip_final_keeps_text_after_closed_block_with_call_form_gemma_opener(): + # A call-form Gemma opener quoted in a closed block must not truncate it. + xml = "<|tool_call>call:t{" + json_block = ( + '{"name":"python","arguments":{"code":"<|tool_call>call:t{"}}' + ) + for block in (xml, json_block): + text = "before " + block + " after" + assert strip_tool_call_markup(text, final = True) == "before after", block + assert strip_tool_call_markup(text) == "before after", block + + +def test_function_sibling_after_close_less_gemma_marker_is_recovered(): + # The close-less marker covers only its braces; the XML sibling is recovered. + text = ( + "<|tool_call>call:bad{broken:{x}} " + "id" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert [c["function"]["name"] for c in calls] == ["terminal"], calls + + +def test_valid_call_after_close_less_marker_with_quoted_close_token_is_recovered(): + # A close token quoted in the later call must not extend the earlier + # close-less marker's coverage over that call. + gemma = '<|tool_call>call:a{x:1} <|tool_call>call:b{note:<|"|><|"|>}' + names = [ + c["function"]["name"] for c in parse_tool_calls_from_text(gemma, allow_incomplete = False) + ] + assert names == ["b"], names + json_text = ( + '{"name":"a","arguments":{}} ' + '{"name":"b","arguments":{"x":""}}' + ) + names_j = [ + c["function"]["name"] for c in parse_tool_calls_from_text(json_text, allow_incomplete = False) + ] + assert "b" in names_j, names_j + + +def test_gemma_parse_value_always_advances_on_stray_delimiter(): + # A stray delimiter (`,`, `}`, `]`) at the primitive position must still advance the + # index by at least one, or a caller looping on it spins forever at 100% CPU (DoS). + for delim in (",", "}", "]"): + text = delim + "rest" + value, nxt, _explicit = _gemma_parse_value(text, 0) + assert nxt > 0, (delim, value, nxt) + + +def test_malformed_gemma_array_does_not_hang(): + # ``[},]`` puts a stray ``}`` at the primitive position inside a list body. + # On the buggy parser this hangs the server; guard with a wall-clock timeout + # so the regression fails loudly instead of blocking CI forever. + import threading + + result: dict = {} + + def _run(): + result["calls"] = parse_tool_calls_from_text("<|tool_call>call:f{a:[},]}") + + t = threading.Thread(target = _run, daemon = True) + t.start() + t.join(timeout = 10.0) + assert not t.is_alive(), "parse_tool_calls_from_text hung on malformed array input" + + +def test_malformed_gemma_mapping_value_does_not_hang(): + # A stray ``}`` where a mapping value is expected must also terminate. + import threading + + result: dict = {} + + def _run(): + result["calls"] = parse_tool_calls_from_text("<|tool_call>call:f{a:}},b:1}") + + t = threading.Thread(target = _run, daemon = True) + t.start() + t.join(timeout = 10.0) + assert not t.is_alive(), "parse_tool_calls_from_text hung on malformed mapping input" diff --git a/studio/backend/tests/test_hf_xet_fallback.py b/studio/backend/tests/test_hf_xet_fallback.py index 9e40fbf508..2fff744b64 100644 --- a/studio/backend/tests/test_hf_xet_fallback.py +++ b/studio/backend/tests/test_hf_xet_fallback.py @@ -1,18 +1,16 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Unit tests for utils.hf_xet_fallback: the no-progress watchdog, the Xet->HTTP -transport policy, and the HF_HUB_DISABLE_XET precondition the fallback rests on. -CPU-only, no network, no real subprocess (the per-attempt download seam is -monkeypatched). +"""Tests for the Studio shim over the shared unsloth_zoo Xet -> HTTP fallback. + +The transport-policy matrix is tested once in unsloth_zoo; here we assert only the +Studio seam: re-exporting the shared API and injecting the marker-aware +prepare_cache_for_transport on the HTTP retry. CPU-only, no network, no real subprocess. """ from __future__ import annotations -import subprocess import sys -import threading -import time import types as _types from pathlib import Path @@ -22,9 +20,8 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) -# Stub heavy/unavailable deps before importing the module under test. Use the -# real structlog when present; a bare stub left in sys.modules would break later -# modules that log at import time. +# Stub heavy/unavailable deps before importing the module under test. Use real structlog when present; +# a bare stub would break later modules that log at import time. _loggers_stub = _types.ModuleType("loggers") _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) sys.modules.setdefault("loggers", _loggers_stub) @@ -34,171 +31,59 @@ except ImportError: sys.modules["structlog"] = _types.ModuleType("structlog") import huggingface_hub -from huggingface_hub import constants as hf_constants + +try: + import unsloth_zoo.hf_xet_fallback as _shared_mod + shared = _shared_mod +except Exception: # noqa: BLE001 - still collect degraded-path tests when unsloth_zoo is unavailable + shared = None import utils.hf_xet_fallback as xf -# --------------------------------------------------------------------------- # -# Watchdog: fires only on a constant-size .incomplete, sparse-aware byte total. -# --------------------------------------------------------------------------- # -REPO = "ztest/xet-watchdog" - - -@pytest.fixture -def hf_cache(tmp_path, monkeypatch): - monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) - return tmp_path - - -def _blobs_dir(root: Path, repo_id: str = REPO) -> Path: - d = root / f"models--{repo_id.replace('/', '--')}" / "blobs" - d.mkdir(parents = True, exist_ok = True) - return d - - -def _wait( - predicate, - timeout: float = 2.0, - step: float = 0.02, -) -> bool: - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - if predicate(): - return True - time.sleep(step) - return predicate() - - -def test_constant_incomplete_fires_stall(hf_cache): - blobs = _blobs_dir(hf_cache) - (blobs / "deadbeef.incomplete").write_bytes(b"\0" * 1024) # never grows - - calls: list[str] = [] - stop = xf.start_watchdog( - repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3 - ) - try: - assert _wait( - lambda: len(calls) >= 1, timeout = 3.0 - ), "watchdog never fired on a constant-size .incomplete" - finally: - stop.set() - assert "stalled" in calls[0].lower() - - -def test_growing_incomplete_never_stalls(hf_cache): - blobs = _blobs_dir(hf_cache) - part = blobs / "growing.incomplete" - part.write_bytes(b"\0" * 1024) - - grow_stop = threading.Event() - - def _grow(): - size = 1024 - while not grow_stop.wait(0.05): - size += 4096 - part.write_bytes(b"\0" * size) - - grower = threading.Thread(target = _grow, daemon = True) - grower.start() - - calls: list[str] = [] - stop = xf.start_watchdog( - repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3 - ) - try: - time.sleep(1.0) # well past stall_timeout, but bytes keep growing - assert calls == [], "watchdog fired despite continuous progress" - finally: - stop.set() - grow_stop.set() - - -def test_no_incomplete_never_stalls(hf_cache): - blobs = _blobs_dir(hf_cache) - (blobs / "finalized_blob").write_bytes(b"\0" * 4096) # no .incomplete - - calls: list[str] = [] - stop = xf.start_watchdog( - repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3 - ) - try: - time.sleep(0.8) - assert calls == [], "watchdog fired with no active .incomplete" - finally: - stop.set() - - -def test_stall_fires_at_most_once(hf_cache): - blobs = _blobs_dir(hf_cache) - (blobs / "frozen.incomplete").write_bytes(b"\0" * 2048) - - calls: list[str] = [] - stop = xf.start_watchdog( - repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.2 - ) - try: - assert _wait(lambda: len(calls) >= 1, timeout = 3.0) - time.sleep(0.6) # keep ticking; must not fire again - assert len(calls) == 1, f"on_stall fired {len(calls)} times, expected exactly 1" - finally: - stop.set() - - -def test_get_state_empty_cache(hf_cache): - assert xf.get_hf_download_state([REPO]) == (0, False) - - -def test_get_state_absent_cache_root(tmp_path, monkeypatch): - monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path / "no-such-cache")) - assert xf.get_hf_download_state([REPO]) == (0, False) - - -def test_get_state_skips_local_paths(hf_cache): - # Filesystem paths are not HF repo IDs and must be ignored without error. - assert xf.get_hf_download_state(["/abs/path", "./rel", "~user", "c:\\x"]) == (0, False) - - -def test_get_state_sparse_aware(hf_cache): - blobs = _blobs_dir(hf_cache) - sparse = blobs / "sparse.incomplete" - with open(sparse, "wb") as f: - f.truncate(64 * 1024 * 1024) # large apparent size, few allocated blocks - st = sparse.stat() - if getattr(st, "st_blocks", 0) == 0: - pytest.skip("filesystem does not report st_blocks; sparse accounting unavailable") - total, has_incomplete = xf.get_hf_download_state([REPO]) - assert has_incomplete is True - assert total < st.st_size, "sparse partial counted at apparent size, not allocated blocks" - - -# --------------------------------------------------------------------------- # -# Transport policy: cached short-circuit, cancel, error propagation, and the -# single Xet->HTTP fallback. _run_download_attempt is faked, so no real spawn. -# --------------------------------------------------------------------------- # DL_REPO, FILE = "ztest/xet-dl", "model-Q4_K_XL.gguf" -@pytest.fixture(autouse = True) -def _no_real_cache_hit(monkeypatch): - """Default: the cached probe misses; tests override it to force a hit.""" +def _requires_shared(): + if shared is None: + pytest.skip("unsloth_zoo.hf_xet_fallback is not installed in this environment") + + +def test_shim_reexports_shared_api(): + _requires_shared() + assert xf.DownloadStallError is shared.DownloadStallError + for name in ( + "start_watchdog", + "get_hf_download_state", + "child_should_disable_xet", + "hf_hub_download_with_xet_fallback", + "snapshot_download_with_xet_fallback", + ): + assert hasattr(xf, name), f"shim missing {name}" + + +def test_child_should_disable_xet_truth_table(): + assert xf.child_should_disable_xet({"disable_xet": True}) is True + assert xf.child_should_disable_xet({"disable_xet": False}) is False + assert xf.child_should_disable_xet({}) is False + + +def test_shim_injects_studio_prepare_on_http_retry(monkeypatch): + """A Xet stall retries over HTTP and the shim runs Studio's marker-aware + ``prepare_cache_for_transport(..., 'http')`` before the retry.""" + _requires_shared() + for var in ("UNSLOTH_DISABLE_XET", "UNSLOTH_STABLE_DOWNLOADS", "HF_HUB_DISABLE_XET"): + monkeypatch.delenv(var, raising = False) monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: None) + seen_disable_xet = [] -class _FakeAttempt: - """Records calls to the download seam and returns scripted results.""" - - def __init__(self, results): - self._results = list(results) - self.calls = [] - - def __call__( - self, + def fake_attempt( repo_id, - filename, - token, *, + kind, + params, + token, repo_type, disable_xet, cancel_event, @@ -208,146 +93,277 @@ class _FakeAttempt: on_status, force_download = False, ): - self.calls.append( - _types.SimpleNamespace( - repo_id = repo_id, - filename = filename, - disable_xet = disable_xet, - repo_type = repo_type, - ) - ) - return self._results[len(self.calls) - 1] + seen_disable_xet.append(disable_xet) + return ("ok", "/cache/model.gguf") if disable_xet else ("stall", None) + monkeypatch.setattr(shared, "_run_download_attempt", fake_attempt) -def _install(monkeypatch, results): - fake = _FakeAttempt(results) - monkeypatch.setattr(xf, "_run_download_attempt", fake) - return fake - - -def test_cached_file_short_circuits(monkeypatch, tmp_path): - cached = tmp_path / "cached.gguf" - cached.write_bytes(b"\0" * 8) - monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: str(cached)) - fake = _install(monkeypatch, []) # must not be called - - out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert out == str(cached) - assert fake.calls == [], "spawned a download for an already-cached file" - - -def test_cancel_before_start_raises_no_attempt(monkeypatch): - fake = _install(monkeypatch, []) - ev = threading.Event() - ev.set() - with pytest.raises(RuntimeError, match = "Cancelled"): - xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None, cancel_event = ev) - assert fake.calls == [] - - -def test_nonstall_error_propagates_without_fallback(monkeypatch): - fake = _install(monkeypatch, [("error", "RepositoryNotFoundError: 404 not found")]) - with pytest.raises(RuntimeError, match = "RepositoryNotFoundError"): - xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert len(fake.calls) == 1, "deterministic error must not trigger an HTTP fallback" - assert fake.calls[0].disable_xet is False - - -def test_immediate_success_uses_xet_only(monkeypatch): - prepared = [] - monkeypatch.setattr( - "hub.utils.download_registry.prepare_cache_for_transport", - lambda *a, **k: prepared.append(a), - ) - fake = _install(monkeypatch, [("ok", "/cache/model.gguf")]) - out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert out == "/cache/model.gguf" - assert len(fake.calls) == 1 and fake.calls[0].disable_xet is False - assert prepared == [], "no cache prep should run when Xet succeeds first try" - - -def test_stall_then_http_fallback_succeeds(monkeypatch): prepared = [] monkeypatch.setattr( "hub.utils.download_registry.prepare_cache_for_transport", lambda repo_type, repo_id, mode, *a, **k: prepared.append((repo_type, repo_id, mode)), ) - fake = _install(monkeypatch, [("stall", None), ("ok", "/cache/model.gguf")]) out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) assert out == "/cache/model.gguf" - assert len(fake.calls) == 2 - assert fake.calls[0].disable_xet is False # Xet first - assert fake.calls[1].disable_xet is True # HTTP fallback - assert prepared == [("model", DL_REPO, "http")], "must prep cache for HTTP before the retry" + assert seen_disable_xet == [False, True] # Xet first, then HTTP + assert prepared == [("model", DL_REPO, "http")], "shim must run Studio's marker-aware prep" -def test_second_stall_raises_download_stall_error(monkeypatch): - monkeypatch.setattr( - "hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None - ) - fake = _install(monkeypatch, [("stall", None), ("stall", None)]) - with pytest.raises(xf.DownloadStallError): - xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert len(fake.calls) == 2 +def test_shim_snapshot_injects_studio_prepare(monkeypatch): + """The snapshot wrapper forwards Studio's marker-aware prep, like the file wrapper.""" + captured = {} + + def fake_snapshot(repo_id, **kwargs): + captured["repo_id"] = repo_id + captured["prepare_for_http_fn"] = kwargs.get("prepare_for_http_fn") + return "/tmp/snap-dir" + + monkeypatch.setattr(xf, "_shared_snapshot_download_with_xet_fallback", fake_snapshot) + out = xf.snapshot_download_with_xet_fallback("org/model") + assert out == "/tmp/snap-dir" + assert captured["repo_id"] == "org/model" + assert captured["prepare_for_http_fn"] is xf._studio_prepare_for_http -def test_cancelled_midattempt_raises_no_fallback(monkeypatch): - fake = _install(monkeypatch, [("cancelled", None)]) - with pytest.raises(RuntimeError, match = "Cancelled"): - xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert len(fake.calls) == 1 +def test_degrades_gracefully_without_shared_helper(monkeypatch): + """On an older unsloth_zoo lacking the shared helper, the shim still imports (Studio + boots) and exposes stub API doing plain HF downloads with the watchdog disabled.""" + import importlib + + class _BlockShared: + def find_spec( + self, + name, + path = None, + target = None, + ): + if name == "unsloth_zoo.hf_xet_fallback": + raise ModuleNotFoundError(f"No module named '{name}'", name = name) + return None + + finder = _BlockShared() + saved_shared = sys.modules.pop("unsloth_zoo.hf_xet_fallback", None) + saved_shim = sys.modules.pop("utils.hf_xet_fallback", None) + sys.meta_path.insert(0, finder) + try: + degraded = importlib.import_module("utils.hf_xet_fallback") + + # Boots without raising and mirrors the shared API surface. + assert issubclass(degraded.DownloadStallError, RuntimeError) + assert degraded.child_should_disable_xet({"disable_xet": True}) is True + assert degraded.get_hf_download_state(["x"]) is None # unmeasurable + event = degraded.start_watchdog(repo_ids = ["x"], on_stall = lambda m: None) + assert hasattr(event, "set") and not event.is_set() # never fires + + # Degraded mode still emits heartbeats so the inactivity deadline is not tripped. + import time as _time + + beats = [] + hb_stop = degraded.start_watchdog( + repo_ids = ["x"], + on_stall = lambda m: None, + on_heartbeat = beats.append, + interval = 0.02, + ) + try: + deadline = _time.monotonic() + 2.0 + while not beats and _time.monotonic() < deadline: + _time.sleep(0.02) + assert beats, "degraded watchdog emitted no heartbeat" + finally: + hb_stop.set() + + # Downloads fall back to plain huggingface_hub (no watchdog, no crash). + called = {} + + def _fake_snapshot(repo_id, **kwargs): + called["repo_id"] = repo_id + return "/snap-dir" + + monkeypatch.setattr(huggingface_hub, "snapshot_download", _fake_snapshot) + assert degraded.snapshot_download_with_xet_fallback("org/model") == "/snap-dir" + assert called["repo_id"] == "org/model" + + # Cancellation still holds: an already-set cancel_event aborts before the HF download. + import threading as _threading + + cancelled = _threading.Event() + cancelled.set() + called.clear() + with pytest.raises(RuntimeError, match = "Cancelled"): + degraded.snapshot_download_with_xet_fallback("org/model", cancel_event = cancelled) + assert "repo_id" not in called, "degraded download ran despite cancellation" + finally: + sys.meta_path.remove(finder) + sys.modules.pop("utils.hf_xet_fallback", None) + if saved_shared is not None: + sys.modules["unsloth_zoo.hf_xet_fallback"] = saved_shared + if saved_shim is not None: + sys.modules["utils.hf_xet_fallback"] = saved_shim -def test_per_file_independent_fallback(monkeypatch): - """A stalled shard falls back; a sibling shard that succeeds does not.""" - monkeypatch.setattr( - "hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None - ) - fake = _install(monkeypatch, [("ok", "/a"), ("stall", None), ("ok", "/b")]) - assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardA.gguf", None) == "/a" - assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardB.gguf", None) == "/b" - assert [c.disable_xet for c in fake.calls] == [False, False, True] +def test_degrades_when_unsloth_zoo_entirely_absent(): + """When unsloth_zoo is absent entirely, the import raises + ModuleNotFoundError(name='unsloth_zoo') (top-level package). Guard that the shim still + degrades and does not re-raise, breaking every Studio import that pulls it in.""" + import importlib + + class _BlockZoo: + def find_spec( + self, + name, + path = None, + target = None, + ): + # Whole package absent, so ModuleNotFoundError.name is the top-level 'unsloth_zoo'. + if name == "unsloth_zoo" or name.startswith("unsloth_zoo."): + raise ModuleNotFoundError("No module named 'unsloth_zoo'", name = "unsloth_zoo") + return None + + finder = _BlockZoo() + saved = { + k: v + for k, v in list(sys.modules.items()) + if k == "unsloth_zoo" or k.startswith("unsloth_zoo.") + } + for k in saved: + del sys.modules[k] + saved_shim = sys.modules.pop("utils.hf_xet_fallback", None) + sys.meta_path.insert(0, finder) + try: + degraded = importlib.import_module("utils.hf_xet_fallback") + # Boots without raising and exposes the stub API. + assert issubclass(degraded.DownloadStallError, RuntimeError) + assert degraded.get_hf_download_state(["x"]) is None + event = degraded.start_watchdog(repo_ids = ["x"], on_stall = lambda m: None) + assert hasattr(event, "set") and not event.is_set() + finally: + sys.meta_path.remove(finder) + sys.modules.pop("utils.hf_xet_fallback", None) + sys.modules.update(saved) + if saved_shim is not None: + sys.modules["utils.hf_xet_fallback"] = saved_shim -# --------------------------------------------------------------------------- # -# Precondition: HF_HUB_DISABLE_XET is read at import time, so assert its effect -# in a FRESH interpreter (huggingface/huggingface_hub#3266 once ignored it). -# --------------------------------------------------------------------------- # -def _safe_path() -> str: +def test_degrades_when_shared_helper_import_raises_importerror(): + """unsloth_zoo can be installed yet fail to import when torch is missing (llama.cpp/GGUF-only + Studio), raising ImportError not ModuleNotFoundError. The shim must degrade for that too.""" + import importlib + + class _BlockWithImportError: + def find_spec( + self, + name, + path = None, + target = None, + ): + if name == "unsloth_zoo.hf_xet_fallback": + # Mirror a torch-less install: a plain ImportError with no .name. + raise ImportError("Unsloth: Pytorch is not installed.") + return None + + finder = _BlockWithImportError() + saved_shared = sys.modules.pop("unsloth_zoo.hf_xet_fallback", None) + saved_zoo = sys.modules.pop("unsloth_zoo", None) + saved_shim = sys.modules.pop("utils.hf_xet_fallback", None) + sys.meta_path.insert(0, finder) + try: + degraded = importlib.import_module("utils.hf_xet_fallback") + assert issubclass(degraded.DownloadStallError, RuntimeError) + assert degraded.get_hf_download_state(["x"]) is None + event = degraded.start_watchdog(repo_ids = ["x"], on_stall = lambda m: None) + assert hasattr(event, "set") and not event.is_set() + finally: + sys.meta_path.remove(finder) + sys.modules.pop("utils.hf_xet_fallback", None) + if saved_shared is not None: + sys.modules["unsloth_zoo.hf_xet_fallback"] = saved_shared + if saved_zoo is not None: + sys.modules["unsloth_zoo"] = saved_zoo + if saved_shim is not None: + sys.modules["utils.hf_xet_fallback"] = saved_shim + + +def test_retries_under_light_gpu_init_when_import_fails(monkeypatch): + """GPU detection in unsloth_zoo's __init__ raises NotImplementedError on a GPU-less host. The shim + retries under UNSLOTH_ZOO_DISABLE_GPU_INIT=1, restores the env, and degrades if the retry fails. + The backend loads lazily (first use of a heavy helper), so this triggers the load explicitly + before asserting the retry/degrade behavior.""" + import importlib import os - return os.environ.get("PATH", "") + + monkeypatch.delenv("UNSLOTH_ZOO_DISABLE_GPU_INIT", raising = False) + seen_env = [] + + class _GpuGatedBlocker: + def find_spec( + self, + name, + path = None, + target = None, + ): + # Crash is in unsloth_zoo's __init__, so intercept "unsloth_zoo" itself (the parent). + if name == "unsloth_zoo": + # Record the env each attempt sees; raise the no-GPU error both times so the shim + # degrades. + seen_env.append(os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT")) + raise NotImplementedError("Unsloth cannot find any torch accelerator") + return None + + finder = _GpuGatedBlocker() + saved = { + k: v + for k, v in list(sys.modules.items()) + if k == "unsloth_zoo" or k.startswith("unsloth_zoo.") + } + for k in saved: + del sys.modules[k] + saved_shim = sys.modules.pop("utils.hf_xet_fallback", None) + sys.meta_path.insert(0, finder) + try: + degraded = importlib.import_module("utils.hf_xet_fallback") + # Import is light (lazy backend); unsloth_zoo not loaded yet. + assert seen_env == [], seen_env + # First use of a heavy helper triggers the load (attempt without the light env, then a retry + # with it set); accessing DownloadStallError drives it via __getattr__. + stall_error = degraded.DownloadStallError + assert seen_env == [None, "1"], seen_env + # Both attempts raised -> Studio still boots in degraded mode. + assert issubclass(stall_error, RuntimeError) + # The env override must not leak past the load. + assert os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") is None + finally: + sys.meta_path.remove(finder) + sys.modules.pop("utils.hf_xet_fallback", None) + sys.modules.update(saved) + if saved_shim is not None: + sys.modules["utils.hf_xet_fallback"] = saved_shim -def test_disable_xet_constant_set_in_fresh_interpreter(): - code = ( - "from huggingface_hub import constants as c; " - "import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is True else 17)" - ) - proc = subprocess.run( - [sys.executable, "-c", code], - env = {"HF_HUB_DISABLE_XET": "1", "PATH": _safe_path()}, - capture_output = True, - text = True, - ) - assert proc.returncode == 0, ( - f"HF_HUB_DISABLE_XET=1 did not set constants.HF_HUB_DISABLE_XET=True " - f"(rc={proc.returncode}): {proc.stderr}" - ) +def test_importing_child_should_disable_xet_stays_light(monkeypatch): + """Regression guard for the stale-transformers-sidecar bug: importing the shim (and + ``child_should_disable_xet``) must NOT pull in ``transformers``/``unsloth_zoo``. The worker calls + this at startup to decide the Xet env flip BEFORE activating the sidecar; an eager import here + would cache the default transformers 4.57.x in sys.modules, defeating the sidecar sys.path prepend + and breaking 5.x models (Qwen3.5/GLM/gemma-4).""" + import importlib + for name in [ + m + for m in list(sys.modules) + if m == "transformers" + or m.startswith("transformers.") + or m == "unsloth_zoo" + or m.startswith("unsloth_zoo.") + or m == "utils.hf_xet_fallback" + ]: + monkeypatch.delitem(sys.modules, name, raising = False) -def test_default_leaves_xet_enabled(): - code = ( - "from huggingface_hub import constants as c; " - "import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is False else 17)" - ) - proc = subprocess.run( - [sys.executable, "-c", code], - env = {"PATH": _safe_path()}, # no HF_HUB_DISABLE_XET - capture_output = True, - text = True, - ) - assert proc.returncode == 0, ( - f"without the env var, constants.HF_HUB_DISABLE_XET was not False " - f"(rc={proc.returncode}): {proc.stderr}" - ) + mod = importlib.import_module("utils.hf_xet_fallback") + # The lightweight decision works without the heavy backend. + assert mod.child_should_disable_xet({"disable_xet": True}) is True + assert mod.child_should_disable_xet({}) is False + # And nothing heavy was imported as a side effect. + assert "transformers" not in sys.modules, "importing the shim must not import transformers" + assert "unsloth_zoo" not in sys.modules, "importing the shim must not import unsloth_zoo" diff --git a/studio/backend/tests/test_inference_default_models_non_blocking.py b/studio/backend/tests/test_inference_default_models_non_blocking.py new file mode 100644 index 0000000000..83a8e7bbfb --- /dev/null +++ b/studio/backend/tests/test_inference_default_models_non_blocking.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Default Chat model metadata must not block on remote Hugging Face discovery.""" + +from __future__ import annotations + +import sys +import time +from pathlib import Path + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +from core.inference.orchestrator import InferenceOrchestrator # noqa: E402 + + +def test_default_models_returns_static_defaults_before_top_fetch(monkeypatch): + sleep_seconds = 2.0 + + def _slow_fetch(self: InferenceOrchestrator) -> None: + time.sleep(sleep_seconds) + self._top_gguf_cache = ["unsloth/slow-GGUF"] + self._top_models_ready.set() + + monkeypatch.setattr(InferenceOrchestrator, "_fetch_top_models", _slow_fetch) + + orchestrator = InferenceOrchestrator() + started = time.monotonic() + defaults = orchestrator.default_models + elapsed = time.monotonic() - started + + assert elapsed < 0.5, f"default_models blocked for {elapsed:.2f}s" + assert defaults == orchestrator._static_models + assert "unsloth/slow-GGUF" not in defaults + + deadline = time.monotonic() + sleep_seconds + 5 + while not orchestrator._top_models_ready.is_set() and time.monotonic() < deadline: + time.sleep(0.05) + + assert "unsloth/slow-GGUF" in orchestrator.default_models diff --git a/studio/backend/tests/test_inference_dispatcher_resilience.py b/studio/backend/tests/test_inference_dispatcher_resilience.py index 60d7503485..6184496d78 100644 --- a/studio/backend/tests/test_inference_dispatcher_resilience.py +++ b/studio/backend/tests/test_inference_dispatcher_resilience.py @@ -112,7 +112,7 @@ def test_route_llama_streaming_async_clients_disable_proxy_env(): continue calls.append(node) - assert len(calls) == 4 + assert len(calls) == 5 for call in calls: assert any( kw.arg == "trust_env" and isinstance(kw.value, ast.Constant) and kw.value.value is False diff --git a/studio/backend/tests/test_install_resolve_prebuilt.py b/studio/backend/tests/test_install_resolve_prebuilt.py index e9941d9e62..e97ca47717 100644 --- a/studio/backend/tests/test_install_resolve_prebuilt.py +++ b/studio/backend/tests/test_install_resolve_prebuilt.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""install_llama_prebuilt.py: host->repo mapping and the --resolve-prebuilt mode. +"""install_llama_prebuilt.py: the --resolve-prebuilt probe (plans against the fork +by default; --published-repo overrides). These back the in-app update for source-build (markerless) installs: the backend asks the installer whether an official prebuilt exists for this host without @@ -24,9 +25,7 @@ if str(_studio) not in sys.path: ilp = importlib.import_module("install_llama_prebuilt") -if not hasattr(ilp, "published_repo_for_host") or not hasattr( - ilp, "resolve_simple_install_release_plans" -): +if not hasattr(ilp, "resolve_simple_install_release_plans"): pytest.skip("PR symbols not present - check branch", allow_module_level = True) FORK = ilp.DEFAULT_PUBLISHED_REPO # unslothai/llama.cpp @@ -56,71 +55,25 @@ def _host(**kw): return ilp.HostInfo(**base) -def test_published_repo_for_host(): - # CPU-only Linux (x64 and arm64) -> ggml-org upstream. - assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True)) == UPSTREAM - assert ( - ilp.published_repo_for_host(_host(is_linux = True, is_arm64 = True, machine = "aarch64")) - == UPSTREAM - ) - # GPU Linux -> fork. - assert ( - ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_usable_nvidia = True)) - == FORK - ) - assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_rocm = True)) == FORK - # CPU-only Windows -> ggml-org (setup.ps1: the fork ships no win-cpu bundle). - assert ( - ilp.published_repo_for_host(_host(system = "Windows", is_windows = True, is_x86_64 = True)) - == UPSTREAM - ) - # GPU Windows -> fork. - assert ( - ilp.published_repo_for_host( - _host(system = "Windows", is_windows = True, is_x86_64 = True, has_usable_nvidia = True) - ) - == FORK - ) - # macOS -> fork regardless of GPU (ggml-org macOS bundles need too-new macOS). - assert ( - ilp.published_repo_for_host( - _host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64") - ) - == FORK - ) - # Linux with AMD tooling but no probed GPU -> fork (setup.sh routes on tooling). - assert ( - ilp.published_repo_for_host( - _host(is_linux = True, is_x86_64 = True), linux_amd_tooling_present = True - ) - == FORK - ) - # The tooling hint is Linux-only: Windows CPU stays on ggml-org. - assert ( - ilp.published_repo_for_host( - _host(system = "Windows", is_windows = True, is_x86_64 = True), - linux_amd_tooling_present = True, - ) - == UPSTREAM - ) - - -def test_macos_intel_and_arm_both_route_to_fork(): - # macOS uses the unslothai fork's own Mac prebuilts for BOTH arm64 and Intel; - # there is no longer any upstream-on-macOS default path, so the obsolete - # pre-macOS-26 pin (b9415) is gone. - assert ( - ilp.published_repo_for_host( - _host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64") - ) - == FORK - ) - assert ( - ilp.published_repo_for_host( - _host(system = "Darwin", is_macos = True, is_x86_64 = True, machine = "x86_64") - ) - == FORK +def test_force_cpu_clears_all_gpu_attributes_including_intel(): + # --cpu-fallback is the "select the CPU prebuilt even when a GPU is present" + # escape hatch. It must drop EVERY GPU attribute, including has_intel_gpu, or + # the planner still prepends the Vulkan asset on an Intel-GPU host. + host = _host( + is_linux = True, + is_x86_64 = True, + has_usable_nvidia = True, + has_physical_nvidia = True, + has_rocm = True, + rocm_gfx_target = "gfx1100", + has_intel_gpu = True, ) + forced = ilp._apply_host_overrides(host, force_cpu = True) + assert forced.has_usable_nvidia is False + assert forced.has_physical_nvidia is False + assert forced.has_rocm is False + assert forced.rocm_gfx_target is None + assert forced.has_intel_gpu is False def test_macos_upstream_pin_only_for_explicit_pre26_upstream(): @@ -188,15 +141,13 @@ def test_resolve_prebuilt_unavailable(monkeypatch, capsys): assert out["repo"] == FORK -def test_resolve_prebuilt_linux_amd_tooling_routes_to_fork(monkeypatch, capsys): - # CPU-probed Linux host but rocminfo on PATH: the dispatch must route to the - # fork so a HIP source build is not offered an upstream CPU prebuilt. - monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True)) - monkeypatch.setattr(ilp.shutil, "which", lambda tool: tool == "rocminfo") +def _run_resolve_capture_host(monkeypatch, capsys): + """Drive --resolve-prebuilt and return the host the resolver was handed.""" seen = {} def _resolver(tag, host, repo, published_release_tag): seen["repo"] = repo + seen["host"] = host raise ilp.PrebuiltFallback("no asset") monkeypatch.setattr(ilp, "resolve_simple_install_release_plans", _resolver) @@ -207,10 +158,33 @@ def test_resolve_prebuilt_linux_amd_tooling_routes_to_fork(monkeypatch, capsys): ) assert ilp.main() == ilp.EXIT_SUCCESS out = json.loads(capsys.readouterr().out.strip().splitlines()[-1]) + return seen, out + + +def test_resolve_prebuilt_cpu_linux_routes_to_fork(monkeypatch, capsys): + # CPU-only Linux host (no GPU): the dispatch routes to the fork, which now + # ships the CPU prebuilt -- it no longer falls back to ggml-org upstream. + monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True)) + seen, out = _run_resolve_capture_host(monkeypatch, capsys) assert seen["repo"] == FORK assert out["repo"] == FORK +def test_resolve_prebuilt_rocm_sdk_only_host_still_offered_cpu(monkeypatch, capsys): + # A CPU-only host that merely has ROCm/HIP SDK tools on PATH (no AMD GPU, so + # detect_host leaves has_rocm False) is a valid CPU-prebuilt target. The probe + # must NOT reclassify it as ROCm from tool presence alone and suppress the CPU + # bundle -- that would deny the fork CPU prebuilt to a legitimate CPU source + # build. The host is left CPU-only and resolves against the fork. + monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True)) + monkeypatch.setattr( + ilp.shutil, "which", lambda tool: "/opt/rocm/bin/hipconfig" if tool == "hipconfig" else None + ) + seen, out = _run_resolve_capture_host(monkeypatch, capsys) + assert seen["repo"] == FORK + assert seen["host"].has_rocm is False + + # Blackwell floor is sm_100 (data-center B100/B200, B300/GB300), below consumer # sm_120 -- 120 wrongly excluded data-center hosts from the prebuilt selection. @@ -360,3 +334,386 @@ def test_sm103_host_drops_cuda128_windows_build(): ) kept_b200 = ilp._drop_blackwell_incapable_windows_cuda(b200, [cuda128, cuda129]) assert [a.name for a in kept_b200] == [cuda128.name, cuda129.name] + + +def _upstream_release(tag, asset_names): + return { + "tag_name": tag, + "assets": [ + {"name": n, "browser_download_url": f"https://example/{n}"} for n in asset_names + ], + } + + +def test_direct_upstream_arm64_intel_prefers_vulkan(): + # Auto-detected Intel GPU on Linux arm64 -> Vulkan prebuilt first, CPU + # second (mirrors the x86_64 branch; ggml-org ships the arm64 Vulkan asset). + host = _host(is_linux = True, is_arm64 = True, machine = "aarch64", has_intel_gpu = True) + rel = _upstream_release( + "b9925", + ["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"], + ) + plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") + kinds = [a.install_kind for a in plan.attempts] + assert kinds[0] == "linux-vulkan", kinds + assert "linux-arm64" in kinds + assert plan.attempts[0].name == "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz" + + +def test_direct_upstream_intel_with_hidden_nvidia_is_cpu_only(): + # A host with a physical NVIDIA hidden via CUDA_VISIBLE_DEVICES (physical + # True, usable False) + an Intel iGPU must NOT get the Vulkan archive even + # when planning directly against upstream: Vulkan ignores CUDA_VISIBLE_DEVICES + # and could grab the reserved card. It falls through to the CPU asset. + host = _host( + is_linux = True, + is_x86_64 = True, + has_intel_gpu = True, + has_physical_nvidia = True, + has_usable_nvidia = False, + ) + rel = _upstream_release( + "b9925", + ["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"], + ) + plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") + assert [a.install_kind for a in plan.attempts] == ["linux-cpu"] + + +def test_direct_upstream_arm64_without_intel_is_cpu_only(): + host = _host(is_linux = True, is_arm64 = True, machine = "aarch64") + rel = _upstream_release( + "b9925", + ["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"], + ) + plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") + assert [a.install_kind for a in plan.attempts] == ["linux-arm64"] + + +def test_direct_upstream_x86_intel_prefers_vulkan(): + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + rel = _upstream_release( + "b9925", + ["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"], + ) + plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") + kinds = [a.install_kind for a in plan.attempts] + assert kinds[0] == "linux-vulkan", kinds + assert "linux-cpu" in kinds + + +def test_linux_vulkan_health_glob_matches_bare_cpu_lib(): + # The widened glob must cover both arch-suffixed (x64) and bare (arm64) CPU + # libs so a valid Vulkan install is not re-flagged unhealthy every check. + choice = ilp.AssetChoice( + repo = UPSTREAM, + tag = "b9925", + name = "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", + url = "https://example/x", + source_label = "upstream", + install_kind = "linux-vulkan", + ) + groups = ilp.runtime_payload_health_groups(choice) + assert ["libggml-cpu*.so*"] in groups + assert ["libggml-cpu-*.so*"] not in groups + + +def test_route_to_vulkan_prebuilt_auto_intel_goes_upstream_and_drops_fork_pin(): + # Routing fork -> upstream also drops the fork release pin, which is in a + # different tag namespace and would make the upstream resolver miss. + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = False) + assert repo == UPSTREAM + assert tag == "" + assert routed.has_intel_gpu is True + + +def test_route_to_vulkan_prebuilt_preserves_explicit_upstream_pin(): + # A pin set WITH an explicit upstream repo is already on upstream -> kept. + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + _routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, UPSTREAM, "b9596", force_cpu = False) + assert repo == UPSTREAM + assert tag == "b9596" + + +def test_route_to_vulkan_prebuilt_cpu_fallback_wins(): + # --cpu-fallback suppresses Vulkan routing even for an Intel host. + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = True) + assert repo == FORK + assert tag == "b9596-mix-abc" + assert routed is host + + +def test_route_to_vulkan_prebuilt_hidden_nvidia_not_rerouted(): + # A mixed NVIDIA+Intel host that hid NVIDIA (CUDA_VISIBLE_DEVICES=""/-1): + # physical NVIDIA present but not usable. Must NOT auto-route to Vulkan, or + # Vulkan (which ignores CUDA_VISIBLE_DEVICES) could grab the reserved GPU. + host = _host( + is_linux = True, + is_x86_64 = True, + has_intel_gpu = True, + has_physical_nvidia = True, + has_usable_nvidia = False, + ) + _routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False) + assert repo == FORK + + +def test_route_to_vulkan_prebuilt_rocm_host_not_rerouted(): + # An Intel iGPU alongside a usable ROCm GPU stays on its ROCm/fork path. + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True, has_rocm = True) + _routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False) + assert repo == FORK + + +def test_route_to_vulkan_prebuilt_non_intel_unchanged(): + host = _host(is_linux = True, is_x86_64 = True) + routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False) + assert repo == FORK + assert routed is host + + +def test_resolve_prebuilt_intel_host_routes_to_upstream(monkeypatch, capsys): + # The --resolve-prebuilt probe must agree with the install path: an + # auto-detected Intel host resolves against upstream (Vulkan), not the fork. + monkeypatch.setattr( + ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + ) + seen, out = _run_resolve_capture_host(monkeypatch, capsys) + assert seen["repo"] == UPSTREAM + assert out["repo"] == UPSTREAM + + +# --------------------------------------------------------------------------- +# windows_intel_gpu_in_registry: the in-process Windows Intel probe. A fake +# winreg module stands in for the real registry so the walk runs anywhere. +# --------------------------------------------------------------------------- + + +class _FakeRegKey: + def __init__( + self, + subkeys = None, + values = None, + denied = False, + ): + self.subkeys = subkeys or {} + self.values = values or {} + self.denied = denied + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + +class _FakeWinreg: + HKEY_LOCAL_MACHINE = object() + + def __init__(self, root_key): + self._root_key = root_key + + def OpenKey(self, parent, name): + if parent is self.HKEY_LOCAL_MACHINE: + # Pin the production constant: a typo'd class GUID must fail here, + # not silently return the fake tree. + if name != ilp._WINDOWS_DISPLAY_CLASS_KEY: + raise FileNotFoundError(name) + if self._root_key is None: + raise FileNotFoundError(name) + return self._root_key + key = parent.subkeys.get(name) + if key is None: + # Real winreg raises OSError, never KeyError, for a missing key. + raise FileNotFoundError(name) + if key.denied: + raise PermissionError(name) + return key + + def QueryInfoKey(self, key): + return (len(key.subkeys), len(key.values), 0) + + def EnumKey(self, key, index): + return list(key.subkeys)[index] + + def QueryValueEx(self, key, value_name): + if value_name not in key.values: + raise FileNotFoundError(value_name) + return (key.values[value_name], 1) + + +def _probe_with_display_class(monkeypatch, adapters): + # The helper lazily does `import winreg`; plant the fake in sys.modules the + # same way unsloth_cli/tests/test_start.py fakes it for _refresh_windows_path. + monkeypatch.setitem(sys.modules, "winreg", _FakeWinreg(_FakeRegKey(subkeys = adapters))) + return ilp.windows_intel_gpu_in_registry() + + +def test_windows_intel_registry_matches_vendor_id(monkeypatch): + assert ( + _probe_with_display_class( + monkeypatch, + { + "0000": _FakeRegKey( + values = { + "MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0&SUBSYS_12345678", + "DriverDesc": "Intel(R) Arc(TM) A770 Graphics", + } + ), + }, + ) + is True + ) + + +def test_windows_intel_registry_matches_driver_desc_without_device_id(monkeypatch): + assert ( + _probe_with_display_class( + monkeypatch, + { + "0000": _FakeRegKey(values = {"DriverDesc": "Intel(R) UHD Graphics 630"}), + }, + ) + is True + ) + + +def test_windows_intel_registry_ignores_non_intel_adapters(monkeypatch): + assert ( + _probe_with_display_class( + monkeypatch, + { + "0000": _FakeRegKey( + values = { + "MatchingDeviceId": r"PCI\VEN_10DE&DEV_2684", + "DriverDesc": "NVIDIA GeForce RTX 4090", + } + ), + "0001": _FakeRegKey( + values = { + "MatchingDeviceId": r"PCI\VEN_1002&DEV_744C", + "DriverDesc": "AMD Radeon RX 7900 XTX", + } + ), + }, + ) + is False + ) + + +def test_windows_intel_registry_skips_restricted_properties_subkey(monkeypatch): + # The real class key carries an ACL-restricted "Properties" subkey and can + # deny access to individual adapter keys; neither may abort the walk. + assert ( + _probe_with_display_class( + monkeypatch, + { + "Properties": _FakeRegKey(denied = True), + "0000": _FakeRegKey(denied = True), + "0001": _FakeRegKey( + values = { + "MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0", + } + ), + }, + ) + is True + ) + + +def test_windows_intel_registry_missing_class_key_is_false(monkeypatch): + monkeypatch.setitem(sys.modules, "winreg", _FakeWinreg(None)) + assert ilp.windows_intel_gpu_in_registry() is False + + +def _detect_windows_host( + monkeypatch, + winreg_fake, + powershell_stdout = "", +): + """Drive the real detect_host() as a GPU-less Windows host with a fake + registry, recording every run_capture invocation. Pins the wiring the + unit tests above cannot see: registry-first, CIM only on a registry miss.""" + monkeypatch.setitem(sys.modules, "winreg", winreg_fake) + monkeypatch.setattr(ilp.platform, "system", lambda: "Windows") + monkeypatch.setattr(ilp.platform, "machine", lambda: "AMD64") + for _env in ( + "CUDA_VISIBLE_DEVICES", + "HIP_VISIBLE_DEVICES", + "ROCR_VISIBLE_DEVICES", + "HIP_PATH", + "ROCM_PATH", + ): + monkeypatch.delenv(_env, raising = False) + monkeypatch.setattr( + ilp.shutil, + "which", + lambda name: "powershell" if name in ("powershell", "pwsh") else None, + ) + captured = [] + + def _fake_run_capture(command, **kwargs): + captured.append(command[0]) + if command[0] == "powershell": + return SimpleNamespace(returncode = 0, stdout = powershell_stdout, stderr = "") + return SimpleNamespace(returncode = 1, stdout = "", stderr = "") + + monkeypatch.setattr(ilp, "run_capture", _fake_run_capture) + return ilp.detect_host(), captured + + +def test_detect_host_registry_intel_skips_cim_probe(monkeypatch): + winreg = _FakeWinreg( + _FakeRegKey( + subkeys = { + "0000": _FakeRegKey(values = {"MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0"}), + } + ) + ) + host, captured = _detect_windows_host(monkeypatch, winreg) + assert host.has_intel_gpu is True + assert "powershell" not in captured + + +def test_detect_host_cim_fallback_fires_on_registry_miss(monkeypatch): + winreg = _FakeWinreg( + _FakeRegKey( + subkeys = { + "0000": _FakeRegKey(values = {"MatchingDeviceId": r"PCI\VEN_10DE&DEV_2684"}), + } + ) + ) + host, captured = _detect_windows_host( + monkeypatch, winreg, powershell_stdout = "Intel(R) Arc(TM) A770 Graphics" + ) + assert host.has_intel_gpu is True + assert "powershell" in captured + + +def test_windows_intel_registry_unexpected_error_is_false(monkeypatch): + # The probe is advisory: even a non-OSError bug in the walk must return + # False (deferring to the CIM fallback), never crash detect_host. + class _ExplodingWinreg: + HKEY_LOCAL_MACHINE = object() + + def OpenKey(self, parent, name): + raise TypeError(name) + + monkeypatch.setitem(sys.modules, "winreg", _ExplodingWinreg()) + assert ilp.windows_intel_gpu_in_registry() is False + + +def test_detect_host_cim_rescues_exploding_registry(monkeypatch): + class _ExplodingWinreg: + HKEY_LOCAL_MACHINE = object() + + def OpenKey(self, parent, name): + raise TypeError(name) + + host, captured = _detect_windows_host( + monkeypatch, _ExplodingWinreg(), powershell_stdout = "Intel(R) Arc(TM) A770 Graphics" + ) + assert host.has_intel_gpu is True + assert "powershell" in captured diff --git a/studio/backend/tests/test_llama_admission.py b/studio/backend/tests/test_llama_admission.py new file mode 100644 index 0000000000..2f04e81926 --- /dev/null +++ b/studio/backend/tests/test_llama_admission.py @@ -0,0 +1,320 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +import asyncio +import os +import sys +import threading + +import pytest + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from core.inference import llama_admission +from core.inference.llama_admission import ( + ADMISSION_CONTROL_ENV, + ADMISSION_KEEPALIVE_INTERVAL_ENV, + ADMISSION_MAX_QUEUE_ENV, + ADMISSION_QUEUE_TIMEOUT_ENV, + DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S, + DEFAULT_ADMISSION_MAX_QUEUE, + DEFAULT_ADMISSION_QUEUE_TIMEOUT_S, + LlamaAdmissionConfig, + LlamaAdmissionQueueFull, + get_llama_admission_queue, + llama_admission_config_from_env, + reset_llama_admission_queues, +) + + +@pytest.fixture(autouse = True) +def _reset_queues(): + reset_llama_admission_queues() + yield + reset_llama_admission_queues() + + +def test_admission_config_defaults(monkeypatch): + for name in ( + ADMISSION_CONTROL_ENV, + ADMISSION_QUEUE_TIMEOUT_ENV, + ADMISSION_KEEPALIVE_INTERVAL_ENV, + ADMISSION_MAX_QUEUE_ENV, + ): + monkeypatch.delenv(name, raising = False) + + config = llama_admission_config_from_env() + + assert config.enabled is True + assert config.queue_timeout_s == DEFAULT_ADMISSION_QUEUE_TIMEOUT_S + assert config.keepalive_interval_s == DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S + assert config.max_queue == DEFAULT_ADMISSION_MAX_QUEUE + + +def test_admission_config_env_overrides(monkeypatch): + monkeypatch.setenv(ADMISSION_CONTROL_ENV, "off") + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0") + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.25") + monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "0") + + config = llama_admission_config_from_env() + + assert config.enabled is False + assert config.queue_timeout_s is None + assert config.keepalive_interval_s == 0.25 + assert config.max_queue is None + + +def test_admission_config_positive_queue_timeout_env(monkeypatch): + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "600") + + config = llama_admission_config_from_env() + + assert config.queue_timeout_s == 600.0 + + +def test_fifo_capacity_one_grants_next_waiter_on_release(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + third = queue.reserve(capacity = 1, config = config) + + first_lease = first.lease_nowait() + assert first_lease is not None + assert second.lease_nowait() is None + assert third.lease_nowait() is None + assert queue.snapshot().queued == 2 + + first_lease.release() + second_lease = await second.wait(0.1) + assert second_lease is not None + assert third.lease_nowait() is None + + second_lease.release() + third_lease = await third.wait(0.1) + assert third_lease is not None + third_lease.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_queue_full_rejects_excess_waiter(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig(max_queue = 1) + + first = queue.reserve(capacity = 1, config = config) + queued = queue.reserve(capacity = 1, config = config) + + assert first.lease_nowait() is not None + assert queued.lease_nowait() is None + with pytest.raises(LlamaAdmissionQueueFull): + queue.reserve(capacity = 1, config = config) + + asyncio.run(_run()) + + +def test_disabled_admission_bypasses_active_slot_limit(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig(enabled = False) + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + + assert first.lease_nowait() is not None + assert second.lease_nowait() is not None + assert queue.snapshot().active == 0 + assert queue.snapshot().queued == 0 + + asyncio.run(_run()) + + +def test_cancelling_promoted_waiter_releases_slot(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + first_lease = first.lease_nowait() + + first_lease.release() + await asyncio.sleep(0) + second.cancel() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_cancelling_promoted_waiter_before_delivery_releases_slot(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + first_lease = first.lease_nowait() + + first_lease.release() + second.cancel() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_external_waiter_future_cancel_invalidates_reservation(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + first_lease = first.lease_nowait() + assert first_lease is not None + assert second._waiter is not None + + second._waiter.future.cancel() + + assert second.lease_nowait() is None + assert second.is_cancelled is True + assert await second.wait(0.01) is None + + first_lease.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_wait_returns_none_when_waiter_future_cancelled_during_wait(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + first_lease = first.lease_nowait() + assert first_lease is not None + assert second._waiter is not None + + wait_task = asyncio.create_task(second.wait(1.0)) + await asyncio.sleep(0) + second._waiter.future.cancel() + + assert await asyncio.wait_for(wait_task, timeout = 0.1) is None + assert second.is_cancelled is True + + first_lease.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_capacity_increase_promotes_existing_waiter_fifo(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + first = queue.reserve(capacity = 1, config = config) + second = queue.reserve(capacity = 1, config = config) + + first_lease = first.lease_nowait() + assert first_lease is not None + assert second.lease_nowait() is None + assert queue.snapshot().active == 1 + assert queue.snapshot().queued == 1 + + third = queue.reserve(capacity = 2, config = config) + + second_lease = await second.wait(0.1) + assert second_lease is not None + assert third.lease_nowait() is None + + snapshot = queue.snapshot() + assert snapshot.capacity == 2 + assert snapshot.active == 2 + assert snapshot.queued == 1 + + first_lease.release() + third_lease = await third.wait(0.1) + assert third_lease is not None + + second_lease.release() + third_lease.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_lease_release_is_idempotent_under_concurrent_calls(): + async def _run(): + queue = get_llama_admission_queue("http://llama.test") + config = LlamaAdmissionConfig() + + reservation = queue.reserve(capacity = 1, config = config) + lease = reservation.lease_nowait() + assert lease is not None + + threads = [threading.Thread(target = lease.release) for _ in range(16)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + +def test_new_key_evicts_idle_prior_load_queues(): + # Each model load carries a fresh ephemeral port, so a new base_url key must + # not leave the drained queues from earlier loads accumulating forever. + get_llama_admission_queue("http://127.0.0.1:1001") + get_llama_admission_queue("http://127.0.0.1:1002") + assert set(llama_admission._QUEUES) == {"http://127.0.0.1:1002"} + + get_llama_admission_queue("http://127.0.0.1:1003") + assert set(llama_admission._QUEUES) == {"http://127.0.0.1:1003"} + + +def test_new_key_retains_in_flight_prior_load_queue(): + config = LlamaAdmissionConfig() + busy = get_llama_admission_queue("http://127.0.0.1:2001") + + async def _run(): + reservation = busy.reserve(capacity = 1, config = config) + lease = reservation.lease_nowait() + assert lease is not None + + # A new load must not drop a queue that still has an in-flight request. + get_llama_admission_queue("http://127.0.0.1:2002") + assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2001", "http://127.0.0.1:2002"} + + # Once it drains, the next load reclaims it. + lease.release() + get_llama_admission_queue("http://127.0.0.1:2003") + assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2003"} + + asyncio.run(_run()) diff --git a/studio/backend/tests/test_llama_cpp_effective_parallel_slots.py b/studio/backend/tests/test_llama_cpp_effective_parallel_slots.py new file mode 100644 index 0000000000..5525bc3ea9 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_effective_parallel_slots.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +import os +import sys + +import pytest + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from core.inference import llama_cpp as llama_cpp_module +from core.inference.llama_cpp import LlamaCppBackend + + +@pytest.fixture +def backend(monkeypatch): + monkeypatch.setattr(LlamaCppBackend, "_kill_orphaned_servers", lambda self: 0) + monkeypatch.setattr(llama_cpp_module.atexit, "register", lambda *_args, **_kwargs: None) + return LlamaCppBackend() + + +def test_effective_parallel_slots_initial_value_is_one(backend): + assert backend.effective_parallel_slots == 1 + + +def test_effective_parallel_slots_commit_uses_final_positive_parallel(backend): + backend._commit_effective_parallel_slots(3) + + assert backend.effective_parallel_slots == 3 + + +@pytest.mark.parametrize("value", [None, 0, -2, "not-an-int"]) +def test_effective_parallel_slots_commit_invalid_value_falls_back_to_one(backend, value): + backend._commit_effective_parallel_slots(value) + + assert backend.effective_parallel_slots == 1 + + +def test_effective_parallel_slots_reset_returns_to_one(backend): + backend._commit_effective_parallel_slots(4) + + backend._reset_effective_parallel_slots() + + assert backend.effective_parallel_slots == 1 + + +def test_effective_parallel_slots_unload_resets_to_one(backend): + backend._commit_effective_parallel_slots(4) + + backend.unload_model() + + assert backend.effective_parallel_slots == 1 diff --git a/studio/backend/tests/test_llama_cpp_freshness.py b/studio/backend/tests/test_llama_cpp_freshness.py index 2a2e113585..08e1334ac9 100644 --- a/studio/backend/tests/test_llama_cpp_freshness.py +++ b/studio/backend/tests/test_llama_cpp_freshness.py @@ -137,9 +137,9 @@ def test_read_install_marker_finds_windows_cmake_layout(tmp_path): @pytest.mark.parametrize("repo", ["unslothai/llama.cpp", "ggml-org/llama.cpp"]) def test_read_install_marker_carries_published_repo_dynamically(tmp_path, repo): - # The freshness check queries whichever release repo the marker records, - # so CUDA (unslothai), CPU/macOS (ggml-org), and ROCm all get the right - # "latest" tag. + # The freshness check queries whichever release repo the marker records: + # new installs record the fork, legacy CPU/macOS markers still say ggml-org, + # and both must get the right "latest" tag. install_dir = tmp_path / "llama.cpp" _write_marker(install_dir, tag = "b9000", published_repo = repo) bin_path = _fake_binary(install_dir, layout = "cmake") diff --git a/studio/backend/tests/test_llama_cpp_stream_cancel.py b/studio/backend/tests/test_llama_cpp_stream_cancel.py new file mode 100644 index 0000000000..1c73f4d17c --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_stream_cancel.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +import contextlib +import os +import sys +import threading + +import httpx +import pytest + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from core.inference.llama_cpp import LlamaCppBackend, _LlamaStreamCancelled + + +def _backend_stub() -> LlamaCppBackend: + backend = LlamaCppBackend.__new__(LlamaCppBackend) + backend._process = object() + backend._healthy = True + backend._port = 48848 + backend._effective_context_length = 4096 + backend._supports_reasoning = False + backend._reasoning_always_on = False + backend._reasoning_style = "enable_thinking" + backend._supports_preserve_thinking = False + return backend + + +def test_stream_cancel_uses_internal_exception_not_generator_exit(): + class FakeResponse: + status_code = 200 + + def close(self): + pass + + class FakeStream: + def __enter__(self): + return FakeResponse() + + def __exit__(self, *_args): + return False + + class FakeClient: + def stream(self, *_args, **_kwargs): + return FakeStream() + + cancel_event = threading.Event() + + with pytest.raises(Exception) as exc_info: + with LlamaCppBackend._stream_with_retry( + FakeClient(), + "http://llama.test/v1/chat/completions", + {}, + cancel_event, + ): + cancel_event.set() + raise httpx.ReadError("client closed") + + assert exc_info.type is _LlamaStreamCancelled + assert not issubclass(exc_info.type, GeneratorExit) + + +def test_generate_chat_completion_swallows_internal_stream_cancel(monkeypatch): + backend = _backend_stub() + + @contextlib.contextmanager + def fake_open_stream(*_args, **_kwargs): + raise _LlamaStreamCancelled + + monkeypatch.setattr(backend, "_open_stream", fake_open_stream) + + chunks = list( + backend.generate_chat_completion( + [{"role": "user", "content": "hi"}], + cancel_event = threading.Event(), + ) + ) + + assert chunks == [] diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index 05d2a0b80a..afac1f5249 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -20,7 +20,11 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) -from core.inference.llama_cpp import _PROVISIONAL_ARGS_MIN_CHARS, LlamaCppBackend +from core.inference.llama_cpp import ( + _MAX_REPROMPTS, + _PROVISIONAL_ARGS_MIN_CHARS, + LlamaCppBackend, +) from state import tool_approvals from state.tool_approvals import TOOL_REJECTED_MESSAGE, resolve_tool_decision @@ -217,7 +221,7 @@ def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch): assert assistant_messages[-1]["tool_calls"][0]["function"]["name"] == "render_html" -def test_buffered_reasoning_answer_emits_backend_summary(monkeypatch): +def test_streamed_reasoning_answer_emits_backend_summary(monkeypatch): stream = [ _sse({"reasoning_content": "I am thinking."}), _sse({"reasoning_content": " Still thinking."}), @@ -236,17 +240,236 @@ def test_buffered_reasoning_answer_emits_backend_summary(monkeypatch): ) ) + content_texts = [e["text"] for e in events if e["type"] == "content"] + # Reasoning streams live during BUFFERING instead of arriving as one block: + # each reasoning delta is emitted immediately, wrapped in . + assert content_texts[0] == "I am thinking." + assert content_texts[1] == "I am thinking. Still thinking." + # The final event closes the block and appends the answer. + assert content_texts[-1] == "I am thinking. Still thinking.Final answer." + summary_index = next( i for i, event in enumerate(events) if event["type"] == "reasoning_summary" ) - content_index = next(i for i, event in enumerate(events) if event["type"] == "content") - assert summary_index < content_index + final_content_index = max(i for i, event in enumerate(events) if event["type"] == "content") + assert summary_index < final_content_index assert events[summary_index]["duration_ms"] == 62000 - assert ( - events[content_index]["text"] - == "I am thinking. Still thinking.Final answer." + + +def test_reasoning_streams_incrementally_with_tools(monkeypatch): + # Regression (DeepSeek "thinking doesn't stream"): with a tool/pill active the + # tool-loop generator must stream reasoning token-by-token like the no-tool + # path, not accumulate it and dump one buffered block. + stream = [ + _sse({"reasoning_content": "Step one."}), + _sse({"reasoning_content": " Step two."}), + _sse({"reasoning_content": " Step three."}), + _sse({"content": "Done."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 2.0, 3.0, 4.0, 4.0]) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "think then answer"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) ) + reasoning_stage = [ + e["text"] + for e in events + if e["type"] == "content" + and e["text"].startswith("") + and "" not in e["text"] + ] + # One live emission per reasoning delta -- not a single dump. + assert reasoning_stage == [ + "Step one.", + "Step one. Step two.", + "Step one. Step two. Step three.", + ] + final = [e["text"] for e in events if e["type"] == "content"][-1] + assert final == "Step one. Step two. Step three.Done." + + +def test_reasoning_only_reply_matches_no_tool_path_with_tools(monkeypatch): + # A reasoning-only turn (whole answer in reasoning_content, no content, no + # tool) with a tool active streams the reasoning live, then resolves to the + # bare reasoning text -- identical to the no-tool generate_chat_completion + # path -- so the non-streaming drain still returns it as `content`, not an + # empty answer. + stream = [ + _sse({"reasoning_content": "The capital of France is Paris."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 5.0, 5.0]) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "just think"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + content_texts = [e["text"] for e in events if e["type"] == "content"] + # Reasoning streamed live during BUFFERING (the fix). + assert content_texts[0] == "The capital of France is Paris." + # Resolves to bare reasoning, matching the no-tool sibling. + assert content_texts[-1] == "The capital of France is Paris." + + +def test_reasoning_before_structured_tool_closes_think_block(monkeypatch): + # Regression: reasoning streamed live during BUFFERING must be closed with + # before a structured tool_call drains, so consumers without a + # reasoning extractor (Anthropic /v1/messages) never receive an unclosed + # . Mirrors the is_match (XML tool signal) path. + tool_stream = [ + _sse({"reasoning_content": "Let me search."}), + *_structured_tool_call("web_search", {"query": "weather"}, "call_1"), + ] + final_stream = [ + _sse({"content": "It is sunny."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [tool_stream, final_stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 2.0, 3.0, 4.0, 4.0]) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", lambda name, arguments, **_kwargs: "sunny" + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + tool_start_index = next(i for i, e in enumerate(events) if e["type"] == "tool_start") + content_before_tool = [e["text"] for e in events[:tool_start_index] if e["type"] == "content"] + # Reasoning streamed live, then closed before the tool -- balanced block. + assert content_before_tool[0] == "Let me search." + assert content_before_tool[-1] == "Let me search." + + +def _replay_route_reasoning_extractor(cumulatives: list[str]) -> tuple[str, str]: + """Replay the route's cumulative suffix-diff + reasoning extractor (the + shared core of routes/inference.py gguf_stream_chunks and the tool-loop + consumer) over content snapshots. Returns (visible, reasoning).""" + from routes.inference import _ResponsesReasoningExtractor + + extractor = _ResponsesReasoningExtractor(parse_think_markers = True) + prev_text = "" + visible: list[str] = [] + reasoning: list[str] = [] + for cumulative in cumulatives: + new_text = cumulative[len(prev_text) :] + prev_text = cumulative + if not new_text: + continue + reasoning_delta, visible_delta = extractor.feed(new_text) + if reasoning_delta: + reasoning.append(reasoning_delta) + if visible_delta: + visible.append(visible_delta) + final_reasoning, final_visible = extractor.finish() + if final_reasoning: + reasoning.append(final_reasoning) + if final_visible: + visible.append(final_visible) + return "".join(visible), "".join(reasoning) + + +def test_reasoning_only_route_output_matches_no_tool_path(monkeypatch): + # Parity contract: a reasoning-only reply must reach the client identically + # whether tools are on or off. Both generators stream live then + # resolve to the bare reasoning text; the route's suffix-diff + extractor + # must therefore produce the same (visible, reasoning) split for both. + stream = [ + _sse({"reasoning_content": "The capital"}), + _sse({"reasoning_content": " of France is Paris."}), + _done(), + ] + + tool_backend = _make_backend(monkeypatch, [list(stream)], []) + _patch_monotonic(monkeypatch, [1.0, 2.0, 2.0]) + tool_cumulatives = [ + e["text"] + for e in tool_backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "capital of France?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + if e.get("type") == "content" + ] + + no_tool_backend = _make_backend(monkeypatch, [list(stream)], []) + no_tool_cumulatives = [ + y + for y in no_tool_backend.generate_chat_completion( + messages = [{"role": "user", "content": "capital of France?"}], + ) + if isinstance(y, str) + ] + + # Both paths stream the reasoning live with the same leading shape. (Raw + # yield lists aren't compared verbatim: the tool path emits a pre-existing + # duplicate trailing event that the route's suffix-diff dedupes.) + assert tool_cumulatives[:3] == no_tool_cumulatives[:3] + # The contract that matters: identical route-level output. + tool_out = _replay_route_reasoning_extractor(tool_cumulatives) + no_tool_out = _replay_route_reasoning_extractor(no_tool_cumulatives) + assert tool_out == no_tool_out + # Pin the shared contract so a change to either path shows up here. + _visible, reasoning = tool_out + assert reasoning == "The capital of France is Paris." + + +def test_reasoning_before_bare_json_tool_closes_think_block(monkeypatch): + # _drain_silently sibling of the structured-tool close: a bare-JSON tool call + # with a live reasoning prefix must also close before draining, and + # must never leak the drained call text as content. + tool_stream = [ + _sse({"reasoning_content": "Searching now."}), + _sse({"content": '{"name":"web_search","arguments":{"query":"weather"}}'}), + _done(), + ] + final_stream = [ + _sse({"content": "It is sunny."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [tool_stream, final_stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 2.0, 3.0, 4.0, 4.0]) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", lambda name, arguments, **_kwargs: "sunny" + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + tool_start_index = next(i for i, e in enumerate(events) if e["type"] == "tool_start") + content_before_tool = [e["text"] for e in events[:tool_start_index] if e["type"] == "content"] + assert content_before_tool[0] == "Searching now." + assert content_before_tool[-1] == "Searching now." + # The bare-JSON call text was drained, never surfaced as content. + assert not any('"name"' in t for t in content_before_tool) + def test_consumed_tool_final_pass_emits_latest_reasoning_summary(monkeypatch): tool_stream = [ @@ -1036,9 +1259,11 @@ def test_render_html_success_does_not_reprompt_render_html_intent(monkeypatch): def test_internal_reprompt_attempts_do_not_duplicate_visible_text(monkeypatch): """No-tool re-prompt attempts should not concatenate into the UI.""" - streams = [ - [_sse({"content": "I will use render_html now."}), _done()], - [_sse({"content": "Understood. I will use render_html now."}), _done()], + # One initial response plus one stream per re-prompt; derive the count from the shared cap. + streams = [[_sse({"content": "I will use render_html now."}), _done()]] + streams += [ + [_sse({"content": "Understood. I will use render_html now."}), _done()] + for _ in range(_MAX_REPROMPTS) ] payloads: list[dict] = [] backend = _make_backend(monkeypatch, streams, payloads) @@ -1073,7 +1298,7 @@ def test_internal_reprompt_attempts_do_not_duplicate_visible_text(monkeypatch): content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] assert content_texts == ["I will use render_html now."] - assert len(payloads) == 2 + assert len(payloads) == _MAX_REPROMPTS + 1 def test_forced_reprompt_plain_final_answer_is_visible(monkeypatch): @@ -1162,6 +1387,48 @@ def test_internal_reprompt_disabled_when_auto_heal_disabled(monkeypatch): assert len(payloads) == 1 +def test_internal_reprompt_disabled_when_nudge_tool_calls_false(monkeypatch): + # Explicit nudge_tool_calls=False disables the plan-without-action + # re-prompt even with Auto-Heal on (None keeps the default-on behavior). + streams = [[_sse({"content": "I will use render_html now."}), _done()]] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + def fake_execute_tool(name, arguments, **_kwargs): + raise AssertionError(f"unexpected tool execution: {name} {arguments}") + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + tools = [ + { + "type": "function", + "function": { + "name": "render_html", + "description": "Render HTML.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + } + ] + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "Make a red square."}], + tools = tools, + max_tool_iterations = 1, + auto_heal_tool_calls = True, + nudge_tool_calls = False, + ) + ) + + content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] + assert content_texts == ["I will use render_html now."] + assert len(payloads) == 1 + + def test_auto_heal_disabled_parses_well_formed_xml_when_tools_enabled(monkeypatch): streams = [ [ @@ -1200,6 +1467,66 @@ def test_auto_heal_disabled_parses_well_formed_xml_when_tools_enabled(monkeypatc ) +def test_textual_mistral_marker_not_leaked_when_inline_with_preface(monkeypatch): + # Textual Mistral ``[TOOL_CALLS]`` inline with visible preface: the DRAINING flush must use the + # shared parser patterns (which know ``[TOOL_CALLS]``); the legacy set leaked the marker to clients. + streams = [ + [_sse({"content": 'Let me search. [TOOL_CALLS]web_search{"query":"cats"}'}), _done()], + [_sse({"content": "done"}), _done()], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "result" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "cats"})] + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("[TOOL_CALLS]" not in t for t in content_texts), content_texts + assert any("Let me search." in t for t in content_texts) + + +def test_textual_llama_python_tag_marker_not_leaked(monkeypatch): + # Same leak class for the Llama-3 built-in ``<|python_tag|>NAME.call(...)`` form. + streams = [ + [_sse({"content": '<|python_tag|>web_search.call(query="cats")'}), _done()], + [_sse({"content": "done"}), _done()], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "result" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "cats"})] + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("<|python_tag|>" not in t for t in content_texts), content_texts + + def test_reprompted_tool_call_still_streams_final_answer(monkeypatch): """Suppression ends once a forced re-prompt actually calls a tool.""" @@ -1738,6 +2065,255 @@ def test_empty_tool_call_id_does_not_emit_provisional_card(monkeypatch): assert calls == [("python", {"code": big_code})] +def _streamed_content(text: str, frag: int = 4) -> list[str]: + """Stream content token-by-token like llama-server; ``frag`` sets the chunk size.""" + chunks = [_sse({"content": text[i : i + frag]}) for i in range(0, len(text), frag)] + chunks.append(_done()) + return chunks + + +def test_bare_json_tool_call_streamed_is_not_leaked_and_executes(monkeypatch): + """A wrapper-less bare-JSON call must be held while incomplete, drained silently, and executed with nothing leaking.""" + + bare_call = '{"name": "web_search", "parameters": {"query": "weather in Sydney"}}' + first_stream = _streamed_content(bare_call) + final_stream = [_sse({"content": "It is sunny in Sydney."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "Weather: sunny, 22C." + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather in Sydney?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + # The tool ran with the parsed arguments. + assert calls == [("web_search", {"query": "weather in Sydney"})] + assert any( + event.get("type") == "tool_end" and event.get("tool_name") == "web_search" + for event in events + ) + + # The bare JSON never leaked to the user-visible stream. + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all('"name"' not in t for t in content_texts), content_texts + assert all("web_search" not in t for t in content_texts), content_texts + # The post-tool synthesis is still streamed. + assert any("sunny in Sydney" in t for t in content_texts), content_texts + + +def test_ordinary_json_with_name_key_is_shown_not_treated_as_tool_call(monkeypatch): + """Markerless JSON with a non-enabled name is the answer, not a phantom call.""" + + answer = '{"name": "Alice", "parameters": {"age": 30}}' + first_stream = _streamed_content(answer) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda n, a, **_k: (calls.append((n, a)) or "x"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "give me a person record"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert any("Alice" in t for t in content_texts), content_texts + + +def test_incomplete_bare_json_truncation_is_not_leaked(monkeypatch): + """If generation is cut off mid bare-JSON object (no closing brace), the held + fragment must be stripped at stream end rather than dumped to the user.""" + + truncated = '{"name": "web_search", "parameters": {"query": "weather in S' + stream = _streamed_content(truncated) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("no complete call")), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all('{"name"' not in t for t in content_texts), content_texts + + +def test_gguf_truncated_ordinary_json_with_name_key_is_shown_not_suppressed(monkeypatch): + """A truncated markerless object whose "name" is NOT an enabled tool (a person + record cut off mid-stream, ``{"name":"Alice","age":``) must still be shown. The + end-of-stream ``_is_bare_tc`` heuristic routed any ``{...,"name",...}`` fragment + to DRAINING (dropped); it is now gated on the enabled tool names so only a real + truncated tool call is suppressed, ordinary JSON streams through.""" + + truncated = '{"name": "Alice", "age": 30, "bio": "loves ' + stream = _streamed_content(truncated) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda n, a, **_k: (calls.append((n, a)) or "x"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "start a person record"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert any("Alice" in t for t in content_texts), content_texts + + +def test_gguf_truncated_disabled_name_json_is_preserved_when_tools_active(monkeypatch): + """A truncated JSON answer with a non-enabled name must still be shown (resolvers are gated on enabled names).""" + + truncated = '{"name": "Alice", "parameters": {"age": 30' + stream = _streamed_content(truncated) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda n, a, **_k: (calls.append((n, a)) or "x"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "give json"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert any("Alice" in t for t in content_texts), content_texts + + +def test_gguf_truncated_enabled_name_json_is_still_suppressed(monkeypatch): + """Counterpart guard: a truncated ENABLED-tool bare call (``web_search``) cut off + mid-JSON still must NOT leak -- the gate only spares disabled / non-tool names.""" + + truncated = '{"name": "web_search", "parameters": {"query": "weather in S' + stream = _streamed_content(truncated) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("no complete call")), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("web_search" not in t for t in content_texts), content_texts + assert all('{"name"' not in t for t in content_texts), content_texts + + +def test_gguf_oversized_disabled_name_json_is_preserved(monkeypatch): + """An oversized still-open JSON answer with a non-enabled name streams as content, not a phantom drain.""" + + cap = 16384 + big = "A" * (cap + 5000) + answer = '{"name":"Alice","parameters":{"bio":"' + big # never closes + first_stream = [_sse({"content": answer[i : i + 2000]}) for i in range(0, len(answer), 2000)] + first_stream.append(_done()) + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda n, a, **_k: (calls.append((n, a)) or "x"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "long json"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert any("Alice" in t for t in content_texts), content_texts[:1] + + +def test_gemma_wrapperless_call_streamed_is_not_leaked_and_executes(monkeypatch): + """Gemma 4 GGUF (skip_special_tokens) streams a wrapper-less ``call:NAME{..}`` + with no XML signal. Like bare JSON, the BUFFERING scan must recognise it via + _GEMMA_BARE_TC_RE, drain it silently, and execute the tool -- never leaking + the ``call:`` markup to the user-visible stream.""" + + gemma_call = 'call:web_search{query:"weather in Sydney"}' + first_stream = _streamed_content(gemma_call) + final_stream = [_sse({"content": "It is sunny in Sydney."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "Weather: sunny, 22C." + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather in Sydney?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "weather in Sydney"})] + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("call:" not in t for t in content_texts), content_texts + assert any("sunny in Sydney" in t for t in content_texts), content_texts + + def _usage_done(usage: dict, finish_reason: str = "stop") -> str: """A terminal SSE chunk carrying llama-server's ``usage`` block, the way the real server reports it on the final chunk of a completion.""" @@ -1813,3 +2389,497 @@ def test_metadata_event_omits_prompt_tokens_details_when_absent(monkeypatch): metadata = [e for e in events if e.get("type") == "metadata"] assert metadata, "expected a metadata event" assert "prompt_tokens_details" not in metadata[-1]["usage"] + + +def test_gguf_rehearsal_name_split_before_args_is_not_leaked(monkeypatch): + """Finding 6: a rehearsal call whose name (``web_search``) and ``[ARGS]{...}`` + arrive in separate content deltas must hold the bare name in the buffer until + ``[ARGS]`` flips it to a drain. Without _is_rehearsal_prefix the GGUF path + streams the tool name as visible content before the call executes.""" + + first_stream = [ + _sse({"content": "web_search"}), + _sse({"content": '[ARGS]{"query":"cats"}'}), + _done(), + ] + final_stream = [_sse({"content": "Found cats."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "result" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "cats"})], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("web_search" not in t for t in content_texts), content_texts + assert all("[ARGS]" not in t for t in content_texts), content_texts + + +def test_gguf_initial_buffer_flush_holds_split_rehearsal_name(monkeypatch): + """The first flush out of BUFFERING (prose plus a trailing active-tool-name in + the first delta, ``[ARGS]{...}`` in the next) must apply the same trailing-name + hold the STREAMING branch uses. The first delta has spaces so it is not a + rehearsal prefix and falls to the initial flush, which previously emitted the + bare name before the call drained.""" + + first_stream = [ + _sse({"content": "I will use web_search"}), + _sse({"content": '[ARGS]{"query":"cats"}'}), + _done(), + ] + final_stream = [_sse({"content": "Found cats."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "cats"})], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("web_search" not in t for t in content_texts), content_texts + assert all("[ARGS]" not in t for t in content_texts), content_texts + + +def test_gguf_rehearsal_name_after_prose_in_streaming_is_not_leaked(monkeypatch): + """Finding 9: the BUFFERING guard only covers a rehearsal at the turn start. + When prose has already streamed (STREAMING state) and the model then emits the + tool name and ``[ARGS]{...}`` in later deltas, the bare name must still be held, + not flushed as visible content before the call drains.""" + + first_stream = [ + _sse({"content": "Let me think. "}), + _sse({"content": "I will search "}), + _sse({"content": "web_search"}), + _sse({"content": '[ARGS]{"query":"cats"}'}), + _done(), + ] + final_stream = [_sse({"content": "Found cats."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "cats"})], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("web_search" not in t for t in content_texts), content_texts + + +def test_gguf_plain_answer_ending_with_tool_name_word_is_preserved(monkeypatch): + """End-of-stream flush: a plain answer that ENDS on a tool-name word with no + ``[ARGS]`` following is real prose and must not be dropped by the streaming + rehearsal hold.""" + + first_stream = [ + _sse({"content": "I think "}), + _sse({"content": "you should "}), + _sse({"content": "web_search"}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "advise"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert any(t.rstrip().endswith("web_search") for t in content_texts), content_texts + + +def test_gguf_long_tool_name_split_rehearsal_is_not_capped_and_executes(monkeypatch): + """Finding 11: a realistic MCP name longer than the 32-char buffer cap split as + NAME then [ARGS]{...} must still be held (a rehearsal prefix is self-bounding), + so the name does not leak and the call executes.""" + name = "mcp__github__create_pull_request" + assert len(name) >= 32, len(name) + + first_stream = [ + _sse({"content": name}), + _sse({"content": '[ARGS]{"x":1}'}), + _done(), + ] + final_stream = [_sse({"content": "done"}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda n, a, **_k: (calls.append((n, a)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "go"}], + tools = [{"type": "function", "function": {"name": name}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [(name, {"x": 1})], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert not any(name in t for t in content_texts), content_texts + + +def test_gguf_streaming_keeps_bare_args_before_think_block(monkeypatch): + """F4: the GGUF streaming strip must run its open-ended ``[ARGS]`` tail cleanup + only on the LAST segment. A bare ``foo[ARGS]`` (no JSON body, ``foo`` not a tool) + before a block is prose, not a truncated call, so the final visible text + must keep it verbatim instead of dropping ``foo[ARGS]`` and corrupting the + sentence.""" + + first_stream = [ + _sse({"content": "Please pass foo[ARGS] "}), + _sse({"content": "pause "}), + _sse({"content": "to the template."}), + _done(), + ] + backend = _make_backend(monkeypatch, [first_stream], []) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "x"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert content_texts, events + assert content_texts[-1] == "Please pass foo[ARGS] pause to the template." + + +def test_gguf_inactive_name_args_in_prose_is_not_drained(monkeypatch): + """BUG A: an inactive-name ``foo[ARGS]{...}`` in a prose answer must not be treated + as a tool call. The BUFFERING and end-of-stream safety-net ``[ARGS]`` checks gate on + active tool names (like the safetensors loop and the mid-stream path), so ``foo`` + (``web_search`` is the only enabled tool) is neither drained/parsed into a disabled + no-op nor forced into another generation turn.""" + first_stream = [ + _sse({"content": 'foo[ARGS]{"x":1} is just syntax.'}), + _done(), + ] + backend = _make_backend(monkeypatch, [first_stream], []) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "x"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 2, + ) + ) + + # No tool executed for the inactive name; a spurious no-op re-prompt would exhaust the + # single supplied stream and error. + assert calls == [], calls + assert not any(e.get("type") in ("tool_start", "tool_end") for e in events), events + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + # The inactive ``foo[ARGS]{...}`` is prose: the name-gated strip keeps the whole sentence. + assert any('foo[ARGS]{"x":1} is just syntax.' in t for t in content_texts), content_texts + + +def test_gguf_inactive_rehearsal_before_active_call_executes_and_keeps_prose(monkeypatch): + """BUG X (#5704): an inactive ``foo[ARGS]{...}`` before a real ``web_search[ARGS]{...}`` + in one delta must NOT swallow the real call; web_search executes while the inactive + rehearsal stays visible as prose.""" + first_stream = [ + _sse({"content": 'foo[ARGS]{"a":1} web_search[ARGS]{"query":"cats"}'}), + _done(), + ] + final_stream = [_sse({"content": "Found cats."}), _done()] + backend = _make_backend(monkeypatch, [first_stream, final_stream], []) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + # The real call runs; ``foo`` is not executed as a phantom disabled call. + assert calls == [("web_search", {"query": "cats"})], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + # The inactive rehearsal is preserved as prose; the active one is stripped. + assert any('foo[ARGS]{"a":1}' in t for t in content_texts), content_texts + assert all("web_search[ARGS]" not in t for t in content_texts), content_texts + + +def test_gguf_rehearsal_detection_recognises_spent_one_shot_with_original_tools(): + # Rehearsal detection is fed the ORIGINAL tool list, so a spent one-shot's re-emitted + # repeat is still detected (matching the strip gate) instead of blanking the turn. + from core.inference.llama_cpp import _gguf_has_genuine_tool_signal + from core.inference.tool_call_parser import TOOL_XML_SIGNALS + + repeat = 'render_html[ARGS]{"code":"x"}' + active_only = [{"type": "function", "function": {"name": "web_search"}}] + original = active_only + [{"type": "function", "function": {"name": "render_html"}}] + assert not _gguf_has_genuine_tool_signal(repeat, TOOL_XML_SIGNALS, active_only) + assert _gguf_has_genuine_tool_signal(repeat, TOOL_XML_SIGNALS, original) + + +def test_gguf_rehearsal_prefix_and_tail_hold_recognise_spent_one_shot(): + # The BUFFERING prefix check and STREAMING/flush tail-holds use the ORIGINAL tool list, + # so a spent one-shot's split repeat is held rather than leaked as visible text. + from core.inference.llama_cpp import _held_rehearsal_tail_len, _is_rehearsal_prefix + + active_only = [{"type": "function", "function": {"name": "web_search"}}] + original = active_only + [{"type": "function", "function": {"name": "render_html"}}] + assert not _is_rehearsal_prefix("render_html", active_only) + assert _is_rehearsal_prefix("render_html", original) + assert _held_rehearsal_tail_len("answer render_html", active_only) == 0 + assert _held_rehearsal_tail_len("answer render_html", original) == len("render_html") + + +def test_gguf_oversized_bare_json_not_leaked_and_executes(monkeypatch): + """An oversized bare-JSON call drains rather than streams, and still executes via the safety net.""" + + cap = 16384 + big = "A" * (cap + 5000) + full = '{"name":"python","parameters":{"code":"' + big + '"}}' + first_stream = [_sse({"content": full[i : i + 2000]}) for i in range(0, len(full), 2000)] + first_stream.append(_done()) + final_stream = [_sse({"content": "done"}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "OK"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "run"}], + tools = [{"type": "function", "function": {"name": "python"}}], + max_tool_iterations = 1, + ) + ) + + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert not any(t.lstrip().startswith('{"name') for t in content_texts), content_texts[:1] + assert calls and calls[0][0] == "python" + assert len(calls[0][1].get("code", "")) > cap + + +def test_gguf_bare_json_call_not_replayed_in_next_turn_content(monkeypatch): + """After a bare-JSON call executes, the kept assistant message must not carry the raw call as content.""" + + import copy + + first_stream = [ + _sse({"content": '{"name":"web_search","parameters":{"query":"cats"}}'}), + _done(), + ] + final_stream = [_sse({"content": "Found."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + monkeypatch.setattr("core.inference.tools.execute_tool", lambda *_a, **_k: "RESULT") + + list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 2, + ) + ) + + assert len(payloads) >= 2 + asst = [m for m in payloads[1]["messages"] if m.get("role") == "assistant"] + assert asst and not any('"name"' in (m.get("content") or "") for m in asst), asst + + +def test_gguf_textual_fallback_caps_distinct_tool_calls_per_turn(monkeypatch): + """A single textual-fallback turn that parses many DISTINCT tool calls must be + capped at _MAX_TOOL_CALLS_PER_TURN (structured delta.tool_calls are grammar + bounded by llama-server; text parsed from content is not). Mirrors the + safetensors loop so one runaway turn cannot fan out into dozens of executions.""" + from core.inference.llama_cpp import _MAX_TOOL_CALLS_PER_TURN + + n = _MAX_TOOL_CALLS_PER_TURN + 4 + blocks = "".join( + '{"name":"t%d","arguments":{"i":%d}}' % (i, i) for i in range(n) + ) + first_stream = [_sse({"content": blocks}), _done()] + final_stream = [_sse({"content": "done"}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "OK"), + ) + + list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "go"}], + tools = [{"type": "function", "function": {"name": f"t{i}"}} for i in range(n)], + max_tool_iterations = 1, + ) + ) + + assert len(calls) == _MAX_TOOL_CALLS_PER_TURN, [c[0] for c in calls] + # The cap keeps the first calls in order (no reordering / drop of leading ones). + assert [c[0] for c in calls] == [f"t{i}" for i in range(_MAX_TOOL_CALLS_PER_TURN)] + + +def test_gguf_textual_fallback_collapses_duplicate_tool_calls(monkeypatch): + """Exact-duplicate textual calls in one turn collapse to a single execution.""" + blocks = '{"name":"web_search","arguments":{"query":"cats"}}' * 5 + first_stream = [_sse({"content": blocks}), _done()] + final_stream = [_sse({"content": "done"}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "OK"), + ) + + list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert len(calls) == 1, [c[0] for c in calls] + + +def test_gguf_drain_truncated_enabled_name_json_preserved_when_auto_heal_disabled(monkeypatch): + """Auto-Heal OFF keeps a truncated enabled-name fragment visible; ON suppresses it (strip gated on auto_heal_tool_calls).""" + + trunc = '{"name":"web_search","parameters":{"query":"weather' + + def _run(auto_heal): + stream = [_sse({"content": trunc}), _done()] + backend = _make_backend(monkeypatch, [stream], []) + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "x"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + auto_heal_tool_calls = auto_heal, + ) + ) + contents = "".join(e.get("text", "") for e in events if e.get("type") == "content") + return calls, contents + + calls_off, contents_off = _run(False) + assert calls_off == [], calls_off + assert "web_search" in contents_off, contents_off + + calls_on, contents_on = _run(True) + assert calls_on == [], calls_on + assert "web_search" not in contents_on, contents_on + + +def test_gguf_valid_tool_calls_respect_max_tool_iterations(monkeypatch): + """Re-prompt slots must not extend the tool budget: stop after ``max_tool_iterations`` executed rounds.""" + # More tool-call streams than the budget: if re-prompt slots leaked into the budget (the bug) the + # loop would run 2+3=5 rounds; honouring it stops after 2, then a tool-less final-answer pass. + streams = [ + _structured_tool_call("web_search", {"query": f"q{i}"}, f"call_{i}") for i in range(6) + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search repeatedly"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 2, + ) + ) + + # Exactly two executed tool rounds, then one final-answer pass. + assert len(calls) == 2, calls + assert len(payloads) == 3, len(payloads) + # The final pass is the budget-exhausted nudge and carries no tools. + assert _tool_names(payloads[2]) == [], _tool_names(payloads[2]) + assert any( + m.get("role") == "user" and "used all available tool calls" in m.get("content", "") + for m in payloads[2]["messages"] + ), payloads[2]["messages"] diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index 4ceffbf75b..f405ebcbd1 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -448,6 +448,48 @@ def test_start_update_happy_path(monkeypatch, tmp_path): assert popen_kwargs["env"]["UNSLOTH_PROGRESS_PERCENT_STEP"] == "5" +def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path): + # A Vulkan install (marker asset carries 'vulkan') must re-assert + # UNSLOTH_FORCE_VULKAN on update, or detect_host on a GPU box re-routes to + # CUDA/ROCm and silently replaces the Vulkan build. + install_dir = tmp_path / "llama.cpp" + binary = _write_install( + install_dir, + "b9493", + repo = "ggml-org/llama.cpp", + asset = "llama-b9493-bin-ubuntu-vulkan-x64.tar.gz", + ) + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + + def _on_start(cmd): + _write_install( + install_dir, + "b9518", + repo = "ggml-org/llama.cpp", + asset = "llama-b9518-bin-ubuntu-vulkan-x64.tar.gz", + ) + + popen_kwargs: dict = {} + _patch_installer_popen( + monkeypatch, + lines = ["installed\n"], + on_start = _on_start, + captured_kwargs = popen_kwargs, + ) + + assert upd.start_update()["started"] is True + deadline = time.time() + 10 + while time.time() < deadline: + job = upd.get_update_status()["job"] + if job["state"] in ("success", "error"): + break + time.sleep(0.05) + assert job["state"] == "success", job + assert popen_kwargs["env"]["UNSLOTH_FORCE_VULKAN"] == "1" + + def test_start_update_reports_full_release_tag(monkeypatch, tmp_path): install_dir = tmp_path / "llama.cpp" binary = _write_install(install_dir, "b9595") @@ -594,9 +636,10 @@ def test_install_cmd_fork_rocm_marker_forwards_has_rocm(monkeypatch, tmp_path): def test_install_cmd_ggml_cpu_marker_has_no_cpu_fallback(monkeypatch, tmp_path): - # CPU installs come from ggml-org. Re-running into the same install-dir/repo - # reproduces the same CPU bundle; --cpu-fallback (which force-drops GPU - # detection) is reserved for setup.sh's arm64 rescue and must not appear here. + # Legacy CPU installs recorded a ggml-org marker (new installs use the fork). + # Re-running into the same install-dir/repo reproduces the same CPU bundle; + # --cpu-fallback (which force-drops GPU detection) is reserved for setup.sh's + # arm64 rescue and must not appear here. cmd = _capture_install_cmd( monkeypatch, tmp_path, diff --git a/studio/backend/tests/test_llama_cpp_vulkan_probe.py b/studio/backend/tests/test_llama_cpp_vulkan_probe.py new file mode 100644 index 0000000000..92aaab4873 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_vulkan_probe.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Vulkan free-VRAM reader regression tests on a synthetic probe output. + +Covers the post-probe handling in +``LlamaCppBackend._get_gpu_free_memory_vulkan``: + + * integrated GPUs (probe reports is_igpu=1) leave a flat per-device host + margin matching llama.cpp's --fit-target, so context auto-sizing can't + over-commit shared RAM, and report total 0 (shared RAM is not a budget), + * discrete GPUs (is_igpu=0) keep their free untouched and pass their real + total through so the fit can reserve absolute headroom, + * an inherited ``GGML_VK_VISIBLE_DEVICES`` is passed through to ggml unchanged + (ggml applies it), not stripped or filtered in Python -- the probe reports + ggml's compact ordinal, which load_model pins with ``--device Vulkan``. + +The ggml Vulkan library is never loaded: subprocess.run is mocked to emit +the tab-separated lines the real ``_vulkan_probe.py`` would print. +""" + +from __future__ import annotations + +import subprocess +import sys +import types as _types +from pathlib import Path +from unittest import mock + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +import importlib as _importlib # noqa: E402 + + +def _maybe_stub(name: str, builder): + try: + _importlib.import_module(name) + except ImportError: + sys.modules[name] = builder() + + +def _build_loggers_stub(): + m = _types.ModuleType("loggers") + m.get_logger = lambda name: __import__("logging").getLogger(name) + return m + + +_maybe_stub("loggers", _build_loggers_stub) +_maybe_stub("structlog", lambda: _types.ModuleType("structlog")) + +from core.inference import llama_cpp as _llama_mod # noqa: E402 +from core.inference.llama_cpp import ( # noqa: E402 + LlamaCppBackend, + _llama_lib_dir, + _vulkan_lib_filename, +) + +MIB = 1024 * 1024 +GIB = 1024 * MIB + + +def _make_vulkan_install(tmp_path: Path) -> str: + """A binary whose sibling dir holds the Vulkan ggml lib, so the + reader's ``is_vulkan_backend`` sibling-file check passes.""" + bindir = tmp_path / "build" / "bin" + bindir.mkdir(parents = True) + binary = bindir / ("llama-server.exe" if sys.platform == "win32" else "llama-server") + binary.write_bytes(b"stub") + (bindir / _vulkan_lib_filename()).write_bytes(b"stub") + return str(binary) + + +def _mock_probe(rows: list[str], captured_env: dict | None = None): + """Patch subprocess.run so the _vulkan_probe.py call returns ``rows`` + (already tab-formatted), recording the env it was launched with.""" + real_run = subprocess.run + + def fake_run(cmd, *args, **kwargs): + if isinstance(cmd, list) and any("_vulkan_probe" in str(c) for c in cmd): + if captured_env is not None: + captured_env.clear() + captured_env.update(kwargs.get("env") or {}) + return subprocess.CompletedProcess( + args = cmd, returncode = 0, stdout = "\n".join(rows), stderr = "" + ) + return real_run(cmd, *args, **kwargs) + + return mock.patch("subprocess.run", side_effect = fake_run) + + +def _row( + idx: int, + free_bytes: int, + is_igpu: int, + total_bytes: int = 0, +) -> str: + return f"{idx}\t{free_bytes}\t{is_igpu}\t{total_bytes}" + + +def test_integrated_gpu_leaves_host_margin(tmp_path): + binary = _make_vulkan_install(tmp_path) + # iGPU with 30 GiB free; reserve a flat 1024 MiB (llama.cpp --fit-target). + # total stays 0: shared system RAM is not a VRAM budget for the fit. + rows = [_row(0, 30 * GIB, is_igpu = 1, total_bytes = 32 * GIB)] + with _mock_probe(rows): + gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary) + assert gpus == [(0, 30 * 1024 - 1024, 0)], gpus + + +def test_discrete_gpu_free_is_untouched_and_total_passed_through(tmp_path): + binary = _make_vulkan_install(tmp_path) + # 6 GiB free on a partially occupied 24 GiB card: free is untouched and the + # real total flows through so the fit reserves absolute headroom (CUDA/ROCm + # parity) instead of the looser free*frac budget. + rows = [_row(0, 6 * GIB, is_igpu = 0, total_bytes = 24 * GIB)] + with _mock_probe(rows): + gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary) + assert gpus == [(0, 6 * 1024, 24 * 1024)], gpus + + +def test_large_discrete_gpu_is_untouched(tmp_path): + binary = _make_vulkan_install(tmp_path) + # A 48 GiB discrete card stays untouched regardless of size; only the + # iGPU flag triggers the host margin, never a VRAM/RAM ratio. + rows = [_row(0, 47 * GIB, is_igpu = 0, total_bytes = 48 * GIB)] + with _mock_probe(rows): + gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary) + assert gpus == [(0, 47 * 1024, 48 * 1024)], gpus + + +def test_inherited_visible_devices_mask_is_passed_through_to_probe(tmp_path, monkeypatch): + # The mask is NOT stripped or filtered in Python: ggml parses it in raw + # physical-device space while this probe reports the compact post-filter + # ordinal, so mixing spaces would be wrong. It is passed through unchanged + # so ggml applies it to the same device list the launch will enumerate. + binary = _make_vulkan_install(tmp_path) + monkeypatch.setenv("GGML_VK_VISIBLE_DEVICES", "1") + captured: dict = {} + rows = [_row(0, 23 * GIB, is_igpu = 0, total_bytes = 24 * GIB)] + with _mock_probe(rows, captured_env = captured): + LlamaCppBackend._get_gpu_free_memory_vulkan(binary) + assert captured.get("GGML_VK_VISIBLE_DEVICES") == "1", captured + + +def test_vulkan_pin_args_uses_device_names_not_env_mask(): + # Pin by compact device name via --device (the space the probe reports and + # the registry names), never by writing a compact ordinal into the raw + # GGML_VK_VISIBLE_DEVICES index space. + assert LlamaCppBackend._vulkan_pin_args([0]) == ["--device", "Vulkan0"] + assert LlamaCppBackend._vulkan_pin_args([1, 2]) == ["--device", "Vulkan1,Vulkan2"] + assert LlamaCppBackend._vulkan_pin_args(None) == [] + assert LlamaCppBackend._vulkan_pin_args([]) == [] + + +def test_vulkan_only_build_is_detected(tmp_path): + binary = _make_vulkan_install(tmp_path) + assert LlamaCppBackend._is_vulkan_backend(binary) is True + + +def test_multi_backend_build_is_not_vulkan_only(tmp_path): + # A custom build that ships CUDA (or HIP) alongside Vulkan must NOT be + # treated as Vulkan-only, or its CUDA GPU would be probed/pinned as a Vulkan + # device; defer to the CUDA/HIP path instead. + binary = _make_vulkan_install(tmp_path) + cuda = "ggml-cuda.dll" if sys.platform == "win32" else "libggml-cuda.so" + (_llama_lib_dir(binary) / cuda).write_bytes(b"stub") + assert LlamaCppBackend._is_vulkan_backend(binary) is False + + +@pytest.mark.skipif(sys.platform == "win32", reason = "shell wrapper fallback is POSIX") +def test_shell_wrapper_entrypoint_resolves_to_real_lib_dir(tmp_path): + # create_exec_entrypoint falls back to a #!/bin/sh wrapper at the install root + # when it cannot symlink; _find_llama_server_binary returns that root entrypoint, + # so _llama_lib_dir must follow the wrapper's exec target to build/bin -- else + # _is_vulkan_backend misses libggml-vulkan.so and the Vulkan probe/pin silently + # never engage on a valid Vulkan install. + import os + + binary = _make_vulkan_install(tmp_path) # tmp_path/build/bin/llama-server + vulkan lib + bindir = Path(binary).parent + wrapper = tmp_path / "llama-server" + wrapper.write_text('#!/bin/sh\nexec "$(dirname "$0")/build/bin/llama-server" "$@"\n') + os.chmod(wrapper, 0o755) + assert _llama_lib_dir(str(wrapper)) == bindir + assert LlamaCppBackend._is_vulkan_backend(str(wrapper)) is True + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/studio/backend/tests/test_llama_route_timeouts.py b/studio/backend/tests/test_llama_route_timeouts.py index 24866bd03e..b4666e3d18 100644 --- a/studio/backend/tests/test_llama_route_timeouts.py +++ b/studio/backend/tests/test_llama_route_timeouts.py @@ -203,3 +203,87 @@ def test_preheader_send_cleanup_on_disconnect_and_cancel(): asyncio.run(_run(False)) asyncio.run(_run(True)) + + +def test_stream_stall_timeout_callable_re_resolved_each_read(): + # The OpenAI passthrough passes a callable so the stall bound can switch to + # the short post-terminal grace mid-stream; it must be re-resolved per read, + # not captured once at generator start. + async def _run(): + response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}})) + values = iter([100.0, 2.0]) + seen = [] + + class _Request: + async def is_disconnected(self): + return False + + class _Items: + def __init__(self): + self.count = 0 + + async def __anext__(self): + self.count += 1 + if self.count > 3: + raise StopAsyncIteration + return "data: {}" + + async for _ in inf_mod._aiter_llama_stream_items( + _Items(), + cancel_event = threading.Event(), + request = _Request(), + response = response, + first_token_deadline = time.monotonic() + 1, + post_first_item_read_timeout_s = lambda: next(values, 5.0), + ): + seen.append(response.request.extensions["timeout"].get("read")) + + assert len(seen) == 3 + # The callable is resolved right after the first item (arming the + # post-first window) and again before each later read, consuming + # successive values. + assert seen[0] == 100.0 + assert 1.0 <= seen[1] <= 2.0 + assert 4.0 <= seen[2] <= 5.0 + + asyncio.run(_run()) + + +def test_stream_stall_timeout_disabled_clears_read_timeout(): + # UNSLOTH_OPENAI_COMPAT_STREAM_STALL_TIMEOUT=0 disables the stall guard, so + # the callable returns None. Once a chunk has arrived the leftover + # first-token read timeout must be cleared, else a long post-first-chunk gap + # trips a stale deadline the operator asked to turn off. + async def _run(): + response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}})) + seen = [] + + class _Request: + async def is_disconnected(self): + return False + + class _Items: + def __init__(self): + self.count = 0 + + async def __anext__(self): + self.count += 1 + if self.count > 2: + raise StopAsyncIteration + return "data: {}" + + async for _ in inf_mod._aiter_llama_stream_items( + _Items(), + cancel_event = threading.Event(), + request = _Request(), + response = response, + first_token_deadline = time.monotonic() + 5, + post_first_item_read_timeout_s = lambda: None, + ): + seen.append(response.request.extensions["timeout"].get("read")) + + # The first-token path armed a finite read timeout; after the first chunk + # with the guard disabled, it is cleared to None on every subsequent read. + assert seen == [None, None], seen + + asyncio.run(_run()) diff --git a/studio/backend/tests/test_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py index 12239e7113..6d26d075cf 100644 --- a/studio/backend/tests/test_mcp_servers.py +++ b/studio/backend/tests/test_mcp_servers.py @@ -587,10 +587,12 @@ def test_tool_xml_strip_handles_hyphenated_function_names(): import re as _re from pathlib import Path + from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC + src = (Path(__file__).resolve().parent.parent / "routes/inference.py").read_text() m = _re.search(r"_TOOL_XML_RE = _re\.compile\((.*?)\n\)", src, _re.DOTALL) assert m, "could not extract _TOOL_XML_RE" - ns: dict = {"_re": _re} + ns: dict = {"_re": _re, "_DS_OPEN_SRC": _DS_OPEN_SRC} exec(f"_TOOL_XML_RE = _re.compile({m.group(1)})", ns) rx = ns["_TOOL_XML_RE"] stripped = rx.sub( diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index 9871965ce8..55a3198a6b 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -4,6 +4,8 @@ import sys import types from types import SimpleNamespace +import pytest + class _DummyMetal: @staticmethod @@ -100,6 +102,32 @@ def test_mlx_inference_text_load_forwards_studio_settings(monkeypatch): ] assert backend._is_vlm is False assert isinstance(backend._tokenizer, _DummyTokenizer) + # Non-LoRA text model: no base_model on the record. + assert backend.models["fake/text"]["base_model"] is None + + +def test_mlx_text_lora_record_keeps_base_model_for_native_template(monkeypatch): + # A LoRA adapter's own tokenizer often ships no chat template; the native tool-calling template + # lives on the base model. + _install_fake_mlx(monkeypatch) + calls = [] + _install_fake_fast_mlx(monkeypatch, calls) + + from core.inference.mlx_inference import MLXInferenceBackend + + backend = MLXInferenceBackend() + config = SimpleNamespace( + identifier = "fake/text-adapter", + is_vision = False, + is_lora = True, + base_model = "fake/text-base", + ) + + assert backend.load_model(config, max_seq_length = 4096, hf_token = "hf-token") + + record = backend.models["fake/text-adapter"] + assert record["is_lora"] is True + assert record["base_model"] == "fake/text-base" def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewrite( @@ -159,6 +187,129 @@ def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewri assert isinstance(backend._tokenizer, _DummyTokenizer) +def test_mlx_inference_distributed_vlm_forwards_group_to_fast_mlx(monkeypatch): + _install_fake_mlx(monkeypatch) + calls = [] + _install_fake_fast_mlx(monkeypatch, calls) + from core.inference.mlx_inference import MLXInferenceBackend + + group = SimpleNamespace(size = lambda: 2, rank = lambda: 0) + config = SimpleNamespace(identifier = "fake/vlm", is_vision = True, is_lora = False) + for mode, group_key in (("tensor", "tensor_group"), ("pipeline", "pipeline_group")): + calls.clear() + assert MLXInferenceBackend().load_model(config, parallel_mode = mode, distributed_group = group) + _, kwargs = calls.pop() + assert kwargs["text_only"] is False and kwargs[group_key] is group + + calls.clear() + singleton = SimpleNamespace(size = lambda: 1, rank = lambda: 0) + assert MLXInferenceBackend().load_model( + config, parallel_mode = "tensor", distributed_group = singleton + ) + assert not {"tensor_group", "pipeline_group"} & set(calls.pop()[1]) + + config = SimpleNamespace(identifier = "fake/adapter", is_vision = False, is_lora = True) + with pytest.raises(ValueError, match = "LoRA adapter repos"): + MLXInferenceBackend().load_model(config, parallel_mode = "tensor", distributed_group = group) + + +@pytest.mark.parametrize("accepts_backend", (True, False)) +def test_mlx_distributed_init_selects_jaccl_backend(monkeypatch, accepts_backend): + _install_fake_mlx(monkeypatch) + from core.inference.mlx_inference import _init_mlx_distributed + + group = SimpleNamespace(rank = lambda: 1, size = lambda: 2) + calls = [] + + def _init(**kwargs): + calls.append(kwargs) + if kwargs and not accepts_backend: + raise TypeError("backend keyword unsupported") + return group + + sys.modules["mlx.core"].distributed = SimpleNamespace(init = _init) + monkeypatch.setenv("MLX_JACCL_COORDINATOR", "127.0.0.1:12345") + monkeypatch.setenv("MLX_IBV_DEVICES", "/tmp/devices.json") + + assert _init_mlx_distributed() == (group, 1, 2) + assert calls == ([{"backend": "jaccl"}] if accepts_backend else [{"backend": "jaccl"}, {}]) + + +def test_worker_share_object_receives_distributed_payload(monkeypatch): + from core.inference import worker + + shared_obj = {"type": "turn", "text": "hi"} + payload = worker._encode_share_object(shared_obj) + + def _array(value): + val = value.item() if hasattr(value, "item") else value + return SimpleNamespace( + item = lambda: val, + tolist = lambda: list(val) if hasattr(val, "__iter__") else [val], + ) + + mlx_pkg = types.ModuleType("mlx") + mlx_core = types.ModuleType("mlx.core") + mlx_core.uint8 = "uint8" + mlx_core.array = _array + mlx_core.zeros = lambda *_a, **_k: _array([]) + + def _all_sum(value, group = None): + value = value.item() if hasattr(value, "item") else value + return _array(len(payload)) if value == 0 else _array(payload) + + mlx_core.distributed = SimpleNamespace(all_sum = _all_sum) + mlx_pkg.core = mlx_core + monkeypatch.setitem(sys.modules, "mlx", mlx_pkg) + monkeypatch.setitem(sys.modules, "mlx.core", mlx_core) + + responses = [] + worker._handle_share_object( + SimpleNamespace( + _distributed_group = object(), + _distributed_rank = 1, + _distributed_world_size = 2, + ), + {"type": "share_object", "request_id": "rid", "object": None}, + SimpleNamespace(put = responses.append), + ) + + response = responses[0] + assert response["object"] == shared_obj + + +def test_worker_share_object_oversize_notifies_peers(monkeypatch): + from core.inference import worker + + calls = [] + + mlx_pkg = types.ModuleType("mlx") + mlx_core = types.ModuleType("mlx.core") + mlx_core.array = lambda value, **_kwargs: SimpleNamespace(item = lambda: value) + mlx_core.eval = lambda value: value + mlx_core.distributed = SimpleNamespace( + all_sum = lambda value, group = None: calls.append(value.item()) or value + ) + mlx_pkg.core = mlx_core + monkeypatch.setitem(sys.modules, "mlx", mlx_pkg) + monkeypatch.setitem(sys.modules, "mlx.core", mlx_core) + monkeypatch.setattr(worker, "_SHARE_OBJECT_MAX_BYTES", 8) + + responses = [] + worker._handle_share_object( + SimpleNamespace( + _distributed_group = object(), + _distributed_rank = 0, + _distributed_world_size = 2, + ), + {"type": "share_object", "request_id": "rid", "object": {"text": "too long"}}, + SimpleNamespace(put = responses.append), + ) + + assert calls == [worker._SHARE_OBJECT_ERROR_SIZE] + assert responses[0]["type"] == "share_error" + + # Regression: generate_chat_response must accept the four template kwargs # (tools / enable_thinking / reasoning_effort / preserve_thinking) so the route # layer can forward UI toggles. The old signature raised @@ -188,12 +339,12 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): _install_fake_mlx(monkeypatch) from core.inference.mlx_inference import MLXInferenceBackend - captured = {} + # The text path renders once with tools, then the native-template fallback makes a second no- + # tools probe call (tools=None) to detect whether the template dropped the schema. + captured_calls = [] def _fake_apply(tokenizer, messages, **kwargs): - captured["tokenizer"] = tokenizer - captured["messages"] = messages - captured["kwargs"] = kwargs + captured_calls.append({"tokenizer": tokenizer, "messages": messages, "kwargs": kwargs}) return "" monkeypatch.setattr( @@ -248,8 +399,15 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): ) ) assert out == ["hi"] - # The toggled kwargs must reach the chat-template helper. - assert captured["kwargs"]["tools"] == [{"function": {"name": "web_search"}}] - assert captured["kwargs"]["enable_thinking"] is True - assert captured["kwargs"]["reasoning_effort"] == "medium" - assert captured["kwargs"]["preserve_thinking"] is True + # The toggled kwargs must reach the chat-template helper on the real render + # (one of the calls carries the tools; the fallback probe passes tools=None). + tool_renders = [ + c + for c in captured_calls + if c["kwargs"].get("tools") == [{"function": {"name": "web_search"}}] + ] + assert tool_renders, captured_calls + render = tool_renders[0] + assert render["kwargs"]["enable_thinking"] is True + assert render["kwargs"]["reasoning_effort"] == "medium" + assert render["kwargs"]["preserve_thinking"] is True diff --git a/studio/backend/tests/test_mlx_repair.py b/studio/backend/tests/test_mlx_repair.py index 1b0cbf9df1..365cc46410 100644 --- a/studio/backend/tests/test_mlx_repair.py +++ b/studio/backend/tests/test_mlx_repair.py @@ -271,6 +271,29 @@ def test_stack_available_requires_runtime_imports_and_versions(monkeypatch): assert imported == list(mr._MLX_RUNTIME_IMPORTS) +def test_mlx_packages_exclude_known_bad_mlx_lm(): + # mlx-lm 0.31.3 regressed QK-norm archs (gemma4 / qwen3_5); the install spec + # must exclude it so the resolver picks 0.31.2 or >=0.31.4. See mlx-lm #1242. + (mlx_lm_spec,) = [p for p in mr.MLX_PACKAGES if p.startswith("mlx-lm")] + assert mlx_lm_spec == "mlx-lm>=0.22.0,!=0.31.3" + + +@pytest.mark.parametrize("bad_form", ["0.31.3", "0.31.3.0"]) +def test_known_bad_installed_mlx_lm_triggers_repair(monkeypatch, bad_form): + # An installed 0.31.3 counts as unsatisfied so the self-heal replaces it; + # parsed-Version compare also catches the trailing-zero form 0.31.3.0. + import importlib.metadata as metadata + + def _version(name): + return bad_form if name == "mlx-lm" else mr._MLX_MIN_VERSIONS[name] + + monkeypatch.setattr(metadata, "version", _version) + monkeypatch.setattr( + mr.importlib, "import_module", lambda _n: pytest.fail("versions must gate imports") + ) + assert mr.mlx_stack_available() is False + + def test_no_op_off_apple_silicon(monkeypatch): monkeypatch.setattr(mr, "is_apple_silicon", lambda: False) called = {"n": 0} diff --git a/studio/backend/tests/test_model_update_robustness.py b/studio/backend/tests/test_model_update_robustness.py index 9cf2a62c39..300eb587b3 100644 --- a/studio/backend/tests/test_model_update_robustness.py +++ b/studio/backend/tests/test_model_update_robustness.py @@ -5,8 +5,8 @@ Covers: * GGUF variant listing computes update_available from the already-fetched sibling metadata instead of a second Hub call. - * hf_hub_download_with_xet_fallback(force_download=True) bypasses the - try_to_load_from_cache cache-first early-return. + * hf_hub_download_with_xet_fallback forwards force_download through the shim to the + shared unsloth_zoo helper (which owns the cache-first early-return and its bypass). The cache "Update" action now runs through the download manager as a normal managed download (so it shows in the Downloads panel with progress + cancel), @@ -341,44 +341,26 @@ def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path): # ── hf_hub_download_with_xet_fallback force_download bypass (X2/F2) ─── -def test_force_download_bypasses_cache_first_early_return(monkeypatch): - """force_download=True skips the try_to_load_from_cache early-return and - proceeds to the real download path; force_download=False returns the cached - path without ever attempting a download (X2/F2).""" - import huggingface_hub as hf +def test_force_download_is_forwarded_through_the_shim(monkeypatch): + """The shim's contract is to forward force_download unchanged to the shared helper (which owns the + cache-first early-return and bypass). Verify both False and True reach it (X2/F2).""" import utils.hf_xet_fallback as X - cached_path = "/cache/blob/cached.gguf" + seen = [] - # Pretend the blob IS cached on disk (try_to_load_from_cache is imported - # inside the function from huggingface_hub, and os.path.exists must agree). - monkeypatch.setattr(hf, "try_to_load_from_cache", lambda *a, **k: cached_path, raising = False) - monkeypatch.setattr(X.os.path, "exists", lambda p: True, raising = False) + def fake_shared(repo_id, filename, token, **kwargs): + seen.append(kwargs.get("force_download")) + return "/downloaded/path" - attempts = [] + monkeypatch.setattr(X, "_shared_hf_hub_download_with_xet_fallback", fake_shared, raising = True) - def fake_attempt(repo_id, filename, token, **kwargs): - attempts.append( - {"repo_id": repo_id, "filename": filename, "force": kwargs.get("force_download")} - ) - return ("ok", "/freshly/downloaded/path") - - monkeypatch.setattr(X, "_run_download_attempt", fake_attempt, raising = True) - - # force_download=False: cache-first early-return, no download attempt. - out = X.hf_hub_download_with_xet_fallback( + X.hf_hub_download_with_xet_fallback( "unsloth/repo", "model.gguf", token = None, force_download = False ) - assert out == cached_path - assert attempts == [] # never reached the real download - - # force_download=True: bypass the early-return, run the real download. - out2 = X.hf_hub_download_with_xet_fallback( + X.hf_hub_download_with_xet_fallback( "unsloth/repo", "model.gguf", token = None, force_download = True ) - assert out2 == "/freshly/downloaded/path" - assert len(attempts) == 1 - assert attempts[0]["force"] is True + assert seen == [False, True] # the shim forwards force_download to the shared helper unchanged # ── multi-revision GGUF blob comparison and update reclaim ── diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index d5e8d13652..86e528ae67 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -23,7 +23,11 @@ if _BACKEND_DIR not in sys.path: from hub.utils.download_manifest import ExpectedFile from hub.utils.gguf import is_mtp_drafter_path -from hub.utils.gguf_plan import build_gguf_variant_plans, plan_from_expected_files +from hub.utils.gguf_plan import ( + build_gguf_variant_plans, + plan_from_expected_files, + preferred_mtp_sibling, +) from utils.models.model_config import ( _is_mtp_drafter, detect_gguf_model, @@ -37,6 +41,8 @@ from utils.models.model_config import ( DRAFTER_CASES = [ ("mtp-gemma-4-12b-it.gguf", True), ("MTP/gemma-4-12b-it-Q8_0-MTP.gguf", True), + # New-scheme MTP/ copies carry the mtp- basename prefix too. + ("MTP/mtp-gemma-4-E4B-it-BF16.gguf", True), ("foo/MTP/bar.gguf", True), ("gemma-4-12b-it-Q8_0.gguf", False), # Baked-in Qwen MTP repos: the head is inside the main GGUF, the file @@ -274,3 +280,178 @@ def test_detect_gguf_model_rejects_mtp_subdir_copy(tmp_path): assert detect_gguf_model(str(copy)) is None # Selecting the MTP dir itself must not surface the copies as models. assert detect_gguf_model(str(sub)) is None + + +# ── Root drafter wins over new-scheme MTP/ copies ──────────────────── +# The MTP/ copies were renamed to share the mtp- basename prefix (e.g. +# MTP/mtp-gemma-4-E4B-it-BF16.gguf). Auto-fetch/load must still resolve the +# small repo-root drafter, not a sort-first MTP/ copy (uppercase precedes +# lowercase, so the subdir path would otherwise win). + +NEW_SCHEME_SIBLINGS = [ + _sib("gemma-4-12b-it-Q4_K_M.gguf", 4_000, "main-q4"), + _sib("gemma-4-12b-it-Q8_0.gguf", 8_000, "main-q8"), + _sib("mtp-gemma-4-12b-it.gguf", 100, "drafter"), + _sib("MTP/mtp-gemma-4-12b-it-Q8_0.gguf", 100, "mtp-sub-q8"), + _sib("MTP/mtp-gemma-4-12b-it-BF16.gguf", 200, "mtp-sub-bf16"), + _sib("mmproj-F16.gguf", 500, "mmproj"), +] + + +def test_preferred_mtp_sibling_prefers_root_over_new_scheme_copies(): + picked = preferred_mtp_sibling(NEW_SCHEME_SIBLINGS) + assert picked is not None and picked.rfilename == "mtp-gemma-4-12b-it.gguf" + + +def test_variant_plans_new_scheme_uses_root_drafter(): + plans = build_gguf_variant_plans(NEW_SCHEME_SIBLINGS) + assert set(plans) == {"q4_k_m", "q8_0"} + for plan in plans.values(): + assert "mtp-gemma-4-12b-it.gguf" in plan.target_filenames + assert not any("MTP/" in name for name in plan.target_filenames) + assert "drafter" in plan.companion_hashes + # Download size = main + mmproj + root drafter (not the 200-byte BF16 copy). + assert plans["q4_k_m"].download_size_bytes == 4_600 + + +def test_download_mtp_prefers_root_over_new_scheme_copies(monkeypatch): + # _pick_mtp is nested; capture it via the companion-download seam. + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) # online: skip reuse probe + captured = {} + + def _fake_companion( + *, + hf_repo, + hf_token, + pick, + label, + cancel_event = None, + ): + captured["pick"] = pick + return None + + b = LlamaCppBackend() + b._download_companion_gguf = _fake_companion + b._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + + repo_files = [ + "MTP/mtp-gemma-4-E4B-it-BF16.gguf", + "MTP/mtp-gemma-4-E4B-it-Q4_0.gguf", + "MTP/mtp-gemma-4-E4B-it-Q8_0.gguf", + "gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf", + "mmproj-F16.gguf", + "mtp-gemma-4-E4B-it.gguf", + ] + assert captured["pick"](repo_files) == "mtp-gemma-4-E4B-it.gguf" + + +# ── Reuse an on-disk drafter offline; fetch fresh online ───────────── + + +def _seed_snapshot(tmp_path, names): + snap = tmp_path / "snap" + for rel in names: + f = snap / rel + f.parent.mkdir(parents = True, exist_ok = True) + f.write_bytes(b"x") + return snap + + +def test_download_mtp_reuses_cached_root_drafter_offline(tmp_path, monkeypatch): + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + snap = _seed_snapshot( + tmp_path, + [ + "gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf", + "mtp-gemma-4-E4B-it.gguf", + "MTP/mtp-gemma-4-E4B-it-BF16.gguf", + "mmproj-F16.gguf", + ], + ) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap]) + + got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it.gguf" + + +def test_download_mtp_reuses_cached_subdir_copy_when_no_root_offline(tmp_path, monkeypatch): + # Pre-fix build may have fetched only the MTP/ copy; reuse it offline. + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + snap = _seed_snapshot( + tmp_path, + [ + "gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf", + "MTP/mtp-gemma-4-E4B-it-BF16.gguf", + ], + ) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap]) + + got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it-BF16.gguf" + + +def test_download_mtp_prefers_root_across_snapshots_offline(tmp_path, monkeypatch): + # A newer partial snapshot holds only the MTP/ copy; an older one has the + # root. Must still return the small root, not the large subdir copy. + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + snap_partial = _seed_snapshot(tmp_path / "new", ["MTP/mtp-gemma-4-E4B-it-BF16.gguf"]) + snap_full = _seed_snapshot(tmp_path / "old", ["mtp-gemma-4-E4B-it.gguf"]) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap_partial, snap_full]) + + got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it.gguf" + + +def test_download_mtp_reuse_follows_snapshot_order_offline(tmp_path, monkeypatch): + # Two snapshots both hold a root drafter; newest-first order must win so a + # fresh main GGUF is not paired with a stale drafter revision. + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + newest = _seed_snapshot(tmp_path / "newest", ["mtp-gemma-4-E4B-it.gguf"]) + oldest = _seed_snapshot(tmp_path / "oldest", ["mtp-gemma-4-E4B-it.gguf"]) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [newest, oldest]) + + got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") + assert got is not None and Path(got).parent.parent.name == "newest" + + +def test_download_mtp_online_skips_cache_reuse(tmp_path, monkeypatch): + # Online, do not reuse a cached copy: go to the download path so a changed + # drafter is refetched (hf_hub_download checks the current revision). + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + snap = _seed_snapshot(tmp_path, ["mtp-gemma-4-E4B-it.gguf"]) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap]) + + reached = {} + + def _fake_companion( + *, + hf_repo, + hf_token, + pick, + label, + cancel_event = None, + ): + reached["hit"] = True + return None + + b = LlamaCppBackend() + b._download_companion_gguf = _fake_companion + assert b._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") is None + assert reached.get("hit") is True diff --git a/studio/backend/tests/test_native_template_trust_remote_code.py b/studio/backend/tests/test_native_template_trust_remote_code.py new file mode 100644 index 0000000000..60dc80f64c --- /dev/null +++ b/studio/backend/tests/test_native_template_trust_remote_code.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for trust_remote_code in the native-template fallback. + +``render_native_template`` re-fetches a model's native chat template from its +repo when an Unsloth override template (mistral, gemma-4) dropped the tools +schema. For a model loaded with ``trust_remote_code=True`` whose tokenizer repo +carries custom code, the secondary ``AutoTokenizer.from_pretrained`` must re-use +that same consent or transformers raises (it requires ``trust_remote_code`` to +instantiate a custom tokenizer class), the ``except`` swallows it, and the +request silently keeps the tool-dropping prompt even though the user already +consented to remote code for the model load. + +These tests pin that the stored ``trust_remote_code`` is threaded to the reload, +that the reload is skipped (returns ``None`` without executing code) when no +consent is stored, and that both backend ``model_info`` dicts persist the flag at +load time so the read lands on a value ``load_model`` actually set. +""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +import pytest + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# ``chat_template_helpers`` is dependency-light (copy / logging / typing, with the +# transformers import deferred inside the function). Load it directly so the test +# runs without importing the heavy ``core.inference`` package (unsloth / torch). +_HELPERS_PATH = Path(_BACKEND_DIR) / "core" / "inference" / "chat_template_helpers.py" +_spec = importlib.util.spec_from_file_location("_native_tpl_trc_test", _HELPERS_PATH) +chat_template_helpers = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(chat_template_helpers) + +render_native_template = chat_template_helpers.render_native_template + + +# A native template that emits a tools section only when tools are provided, so the +# with-tools vs no-tools render differs and ``render_native_template`` accepts it. +_NATIVE_TEMPLATE = ( + "{% for m in messages %}{{ m['role'] }}: {{ m['content'] }}\n{% endfor %}" + "{% if tools %}[AVAILABLE_TOOLS]{{ tools }}[/AVAILABLE_TOOLS]\n{% endif %}" + "{% if add_generation_prompt %}assistant:{% endif %}" +) + +_MESSAGES = [{"role": "user", "content": "what is the weather"}] +_TOOLS = [{"type": "function", "function": {"name": "get_weather"}}] + + +class _JinjaTokenizer: + """Minimal tokenizer whose ``apply_chat_template`` renders ``self.chat_template``. + + Stands in for the live model tokenizer that ``render_native_template`` shallow- + copies and re-points at the native template before rendering. + """ + + def __init__(self, chat_template): + self.chat_template = chat_template + + def apply_chat_template( + self, + messages, + tokenize = False, + add_generation_prompt = True, + tools = None, + **kwargs, + ): + from jinja2 import BaseLoader, Environment + env = Environment(loader = BaseLoader()) + return env.from_string(self.chat_template).render( + messages = messages, + tools = tools, + add_generation_prompt = add_generation_prompt, + ) + + +def _install_custom_code_tokenizer(monkeypatch): + """Patch ``AutoTokenizer.from_pretrained`` to mimic a custom-code repo: raise + unless ``trust_remote_code`` is truthy, else return a tokenizer carrying the + native template. Records the ``trust_remote_code`` it was called with.""" + pytest.importorskip("jinja2") + from transformers import AutoTokenizer + + calls = {} + + def fake_from_pretrained( + model_id, + *args, + trust_remote_code = False, + token = None, + **kwargs, + ): + calls["trust_remote_code"] = trust_remote_code + calls["model_id"] = model_id + calls["token"] = token + if not trust_remote_code: + # Mirrors transformers.dynamic_module_utils.resolve_trust_remote_code: + # has_remote_code and not has_local_code and not trust_remote_code -> ValueError. + raise ValueError( + f"The repository {model_id} contains custom code which must be executed " + "to correctly load the model. Please pass the argument " + "`trust_remote_code=True` to allow custom code to be run." + ) + return _JinjaTokenizer(_NATIVE_TEMPLATE) + + monkeypatch.setattr(AutoTokenizer, "from_pretrained", staticmethod(fake_from_pretrained)) + return calls + + +def _model_info(trust_remote_code): + return { + "native_chat_template": None, # force the repo reload path + "base_model": None, # non-LoRA: template_source == active_model_name + "trust_remote_code": trust_remote_code, + # Live tokenizer that gets shallow-copied + re-pointed at the native template. + "tokenizer": _JinjaTokenizer("OVERRIDE-THAT-DROPS-TOOLS"), + } + + +def test_native_reload_passes_stored_trust_remote_code(monkeypatch): + """With ``trust_remote_code`` stored on ``model_info`` the custom-code reload + succeeds and the tools-advertising native prompt is returned. This FAILS before + the fix (reload omits the flag, raises, is swallowed, returns None).""" + calls = _install_custom_code_tokenizer(monkeypatch) + model_info = _model_info(trust_remote_code = True) + + out = render_native_template( + model_info = model_info, + active_model_name = "acme/custom-tokenizer-model", + messages = _MESSAGES, + tools = _TOOLS, + ) + + assert out is not None, "native fallback should render the tools prompt with consent" + assert "[AVAILABLE_TOOLS]" in out + assert "get_weather" in out + assert calls["trust_remote_code"] is True # the stored consent was threaded through + # A successful fetch is cached so the next tool turn skips the reload. + assert model_info["native_chat_template"] == _NATIVE_TEMPLATE + + +def test_native_reload_without_consent_returns_none(monkeypatch): + """Without stored consent the custom-code reload raises, is swallowed, and + ``render_native_template`` returns None (no unconsented code execution). Proves + the stored flag -- not a hard-coded True -- drives the reload.""" + calls = _install_custom_code_tokenizer(monkeypatch) + model_info = _model_info(trust_remote_code = False) + + out = render_native_template( + model_info = model_info, + active_model_name = "acme/custom-tokenizer-model", + messages = _MESSAGES, + tools = _TOOLS, + ) + + assert out is None + assert calls["trust_remote_code"] is False + # A failed fetch must not be cached as "no template" (would pin the tool drop). + assert model_info["native_chat_template"] is None + + +def test_backend_model_info_persists_trust_remote_code(): + """Both backends must store ``trust_remote_code`` on their per-model info dict so + ``render_native_template`` can source the consent value. Guards against the read + landing on a key ``load_model`` never sets (which would silently no-op the fix).""" + inf = (Path(_BACKEND_DIR) / "core" / "inference" / "inference.py").read_text() + mlx = (Path(_BACKEND_DIR) / "core" / "inference" / "mlx_inference.py").read_text() + assert '"trust_remote_code": trust_remote_code,' in inf + assert '"trust_remote_code": trust_remote_code,' in mlx diff --git a/studio/backend/tests/test_nudge_tool_calls_wiring.py b/studio/backend/tests/test_nudge_tool_calls_wiring.py new file mode 100644 index 0000000000..e03fd0c7d7 --- /dev/null +++ b/studio/backend/tests/test_nudge_tool_calls_wiring.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Wiring guard for the plan-without-action ``nudge_tool_calls`` policy. + +Decided policy: the re-prompt is ALWAYS ON for the Studio inference paths +(safetensors, GGUF/llama_cpp, MLX) and OPT-IN for the API (/v1 OpenAI-compat + +Anthropic-compat, controlled by the request's ``nudge_tool_calls``, default off). + +Mechanism (verified here without loading a model): + + * every backend tool-loop entry point accepts and forwards ``nudge_tool_calls`` + (safetensors -> ``InferenceBackend``; MLX -> ``InferenceOrchestrator``; both + call the shared ``run_safetensors_tool_loop``; GGUF -> ``LlamaCppBackend``); + * the safetensors/MLX loop gates the retry on a truthy flag (new retry -> + opt-in), while the GGUF loop keeps its pre-existing default-on behaviour + (``None`` keeps nudging) so an omitted flag never disables GGUF; + * the API request models default the flag to ``None`` (opt-in / off); + * the Studio-facing routes forward the request's flag, and the Studio frontend + sends ``nudge_tool_calls: true`` -- exercised behaviourally in + ``test_safetensors_tool_loop.py`` and ``test_llama_cpp_tool_loop.py``. +""" + +import inspect + +from core.inference.llama_cpp import LlamaCppBackend +from core.inference.orchestrator import InferenceOrchestrator +from core.inference.safetensors_agentic import run_safetensors_tool_loop + +try: + # core.inference.inference imports unsloth at module scope, which requires + # unsloth_zoo. The dependency-light backend CI matrix job does not install + # it, so the safetensors InferenceBackend is folded into the checks below + # only when the unsloth stack is importable (local runs / full CI); the + # other entry points are always checked. + from core.inference.inference import InferenceBackend +except ImportError: + InferenceBackend = None + + +def _params(fn): + return inspect.signature(fn).parameters + + +def test_shared_loop_accepts_nudge_flag(): + assert "nudge_tool_calls" in _params(run_safetensors_tool_loop) + + +def test_backends_accept_the_flag(): + methods = [ + InferenceOrchestrator.generate_chat_completion_with_tools, + LlamaCppBackend.generate_chat_completion_with_tools, + ] + if InferenceBackend is not None: # safetensors path; needs the unsloth stack + methods.append(InferenceBackend.generate_chat_completion_with_tools) + for method in methods: + assert "nudge_tool_calls" in _params(method), method.__qualname__ + + +def test_delegating_backends_forward_the_flag_to_the_shared_loop(): + # safetensors (in-process transformers) and MLX (parent-process orchestrator) + # both delegate to run_safetensors_tool_loop; GGUF runs its own in-file loop + # and consumes the flag directly (asserted separately by the gate test). + methods = [InferenceOrchestrator.generate_chat_completion_with_tools] + if InferenceBackend is not None: # safetensors path; needs the unsloth stack + methods.append(InferenceBackend.generate_chat_completion_with_tools) + for method in methods: + src = inspect.getsource(method) + assert "nudge_tool_calls = nudge_tool_calls" in src, method.__qualname__ + + +def test_safetensors_loop_is_opt_in_while_gguf_stays_default_on(): + # Safetensors/MLX: the retry is new here, so it requires a truthy flag. + sf_src = inspect.getsource(run_safetensors_tool_loop) + assert "and nudge_tool_calls" in sf_src + # GGUF: pre-existing nudge must not be accidentally disabled -- an omitted + # (None) flag keeps nudging; only an explicit False turns it off. + gguf_src = inspect.getsource(LlamaCppBackend.generate_chat_completion_with_tools) + assert "nudge_tool_calls is None or nudge_tool_calls" in gguf_src + + +def test_api_request_models_default_the_flag_off(): + from models.inference import AnthropicMessagesRequest, ChatCompletionRequest + for model in (ChatCompletionRequest, AnthropicMessagesRequest): + field = model.model_fields["nudge_tool_calls"] + assert field.default is None, model.__name__ + + +def test_studio_routes_forward_the_request_flag(): + # The Studio chat frontend posts to /v1/chat/completions and /v1/messages + # with nudge_tool_calls=true; the route handlers forward the request value + # (external API clients that omit it fall back to the opt-in default). + from routes import inference as routes_inference + for handler in ( + routes_inference.openai_chat_completions, + routes_inference.anthropic_messages, + ): + src = inspect.getsource(handler) + assert "nudge_tool_calls = payload.nudge_tool_calls" in src, handler.__name__ diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index 4499881c4d..e24e2ca451 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -79,9 +79,11 @@ from huggingface_hub import constants as hf_constants from core.inference.llama_cpp import ( LlamaCppBackend, + _cached_colocated_split_main, _gguf_files_for_variant, _hf_offline_if_dns_dead, _probe_dns_dead, + _resolve_repo_id_casing, ) from utils.models.model_config import ( _detect_gguf_from_hf_cache, @@ -217,7 +219,7 @@ class TestGgufVariantFileResolution: downloaded.append(filename) return f"/fake/{repo_id}/{filename}" - monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) with ( patch( "huggingface_hub.list_repo_files", @@ -239,6 +241,214 @@ class TestGgufVariantFileResolution: assert downloaded == ["tinyllamas/stories260K.gguf"] assert out == "/fake/ggml-org/models/tinyllamas/stories260K.gguf" + def test_download_reuses_older_snapshot_when_current_ref_snapshot_is_partial( + self, monkeypatch, hf_cache + ): + # Cross-snapshot reuse is an offline-resilience path: online, hf_hub_download + # resumes the partial current-ref download and revalidates the revision instead + # of serving an older snapshot's same-name blob. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + backend = LlamaCppBackend() + repo = "unsloth/vision-GGUF" + old = _build_cache( + hf_cache, + repo, + {"model-UD-Q4_K_XL.gguf": 4}, + snapshot_sha = "a" * 40, + ) + _build_cache( + hf_cache, + repo, + {"mtp-model.gguf": 1}, + snapshot_sha = "b" * 40, + ) + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = path, size = 4) for path in paths if path] + + def fail_download(*_args, **_kwargs): + raise AssertionError("should reuse the cached GGUF instead of downloading") + + with ( + patch( + "huggingface_hub.list_repo_files", + lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf", "mtp-model.gguf"], + ), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_gguf( + hf_repo = repo, + hf_variant = "UD-Q4_K_XL", + ) + + assert out == str(old / "model-UD-Q4_K_XL.gguf") + + def test_download_reuses_cached_gguf_when_lowercase_partial_cache_shadows_it( + self, monkeypatch, hf_cache + ): + # Case-variant cross-dir reuse is offline-only; online the canonical repo id + # resolves up front and hf_hub_download fetches the current revision. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + backend = LlamaCppBackend() + canonical_repo = "unsloth/gemma-4-E2B-it-GGUF" + requested_repo = "unsloth/gemma-4-e2b-it-gguf" + gguf_file = "gemma-4-E2B-it-UD-Q4_K_XL.gguf" + snap = _build_cache( + hf_cache, + canonical_repo, + {gguf_file: 4}, + snapshot_sha = "a" * 40, + ) + lower_snap = _build_cache( + hf_cache, + requested_repo, + {"mtp-gemma-4-E2B-it.gguf": 1}, + snapshot_sha = "b" * 40, + ) + os.utime(lower_snap, (2000, 2000)) + os.utime(snap, (1000, 1000)) + seen_repos: list[str] = [] + + def fake_list_repo_files(repo_id, token = None): + seen_repos.append(repo_id) + return [gguf_file] + + def fake_get_paths_info( + repo_id, + paths, + token = None, + ): + seen_repos.append(repo_id) + return [_types.SimpleNamespace(path = path, size = 4) for path in paths if path] + + def fake_cache(repo_id, filename, *args, **kwargs): + seen_repos.append(repo_id) + return str(snap / filename) if repo_id == canonical_repo else None + + def fail_download(*_args, **_kwargs): + raise AssertionError("should reuse the cached GGUF instead of downloading") + + with ( + patch("huggingface_hub.list_repo_files", fake_list_repo_files), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", fake_cache), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_gguf( + hf_repo = requested_repo, + hf_variant = "UD-Q4_K_XL", + ) + + assert out == str(snap / gguf_file) + assert seen_repos + + def test_download_online_does_not_reuse_old_snapshot(self, monkeypatch, hf_cache): + # Online, an older same-name snapshot must not be served (it may be a stale + # revision); hf_hub_download is called so the current revision is fetched and + # its etag revalidated. + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + backend = LlamaCppBackend() + repo = "unsloth/vision-GGUF" + _build_cache(hf_cache, repo, {"model-UD-Q4_K_XL.gguf": 4}, snapshot_sha = "a" * 40) + downloaded: list[str] = [] + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p] + + def fake_download( + repo_id, + filename, + token = None, + **kwargs, + ): + downloaded.append(filename) + return f"/fresh/{filename}" + + with ( + patch( + "huggingface_hub.list_repo_files", + lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf"], + ), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + ): + out = backend._download_gguf(hf_repo = repo, hf_variant = "UD-Q4_K_XL") + + assert downloaded == ["model-UD-Q4_K_XL.gguf"] + assert out == "/fresh/model-UD-Q4_K_XL.gguf" + + def test_download_reuses_older_snapshot_when_offline_env_is_true(self, monkeypatch, hf_cache): + # HF_HUB_OFFLINE accepts truthy spellings beyond "1" (true/yes/on); the offline + # cache reuse must trigger for those too, otherwise the earlier Hub calls run + # offline while this branch still attempts hf_hub_download and the cached GGUF + # cannot load. + monkeypatch.setenv("HF_HUB_OFFLINE", "true") + backend = LlamaCppBackend() + repo = "unsloth/vision-GGUF" + old = _build_cache(hf_cache, repo, {"model-UD-Q4_K_XL.gguf": 4}, snapshot_sha = "a" * 40) + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p] + + def fail_download(*_args, **_kwargs): + raise AssertionError("should reuse the cached GGUF instead of downloading") + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf"]), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_gguf(hf_repo = repo, hf_variant = "UD-Q4_K_XL") + + assert out == str(old / "model-UD-Q4_K_XL.gguf") + + def test_download_companion_resolves_from_case_variant_snapshot_offline( + self, monkeypatch, hf_cache + ): + # Offline, resolve_cached_repo_id_case can keep a partial lower-case spelling, + # so the companion (mmproj) must resolve from whichever case-variant snapshot + # actually holds it rather than being dropped by an hf_hub_download on the + # wrong casing. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + backend = LlamaCppBackend() + canonical_repo = "unsloth/gemma-4-E2B-it-GGUF" + requested_repo = "unsloth/gemma-4-e2b-it-gguf" + snap = _build_cache(hf_cache, canonical_repo, {"mmproj-F16.gguf": 4}, snapshot_sha = "a" * 40) + # A partial lower-case dir exists so casing resolution keeps the requested spelling. + _build_cache(hf_cache, requested_repo, {"config.json": 1}, snapshot_sha = "b" * 40) + + _offline_exc = type("OfflineModeIsEnabled", (Exception,), {}) + + def fake_list_repo_files(repo_id, token = None): + raise _offline_exc("offline") + + def fail_download(*_args, **_kwargs): + raise AssertionError("should resolve the companion from cache, not download") + + with ( + patch("huggingface_hub.list_repo_files", fake_list_repo_files), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_mmproj(hf_repo = requested_repo) + + assert out == str(snap / "mmproj-F16.gguf") + def test_download_includes_uppercase_split_gguf_shards(self, monkeypatch, tmp_path): backend = LlamaCppBackend() downloaded: list[str] = [] @@ -264,7 +474,7 @@ class TestGgufVariantFileResolution: downloaded.append(filename) return f"/fake/{repo_id}/{filename}" - monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) with ( patch("huggingface_hub.list_repo_files", lambda *_a, **_k: files), patch("huggingface_hub.get_paths_info", fake_get_paths_info), @@ -279,6 +489,48 @@ class TestGgufVariantFileResolution: assert downloaded == files assert out == "/fake/org/repo/model-Q4_K_M-00001-of-00002.GGUF" + def test_download_refetches_split_gguf_when_shards_span_snapshots(self, monkeypatch, hf_cache): + # The cached main shard lives in an older snapshot; its sibling shard is only + # in a newer, separate snapshot. Reusing the main shard alone would leave + # llama.cpp unable to resolve the sibling, so the whole set must be re-fetched + # together (co-located) rather than served split across snapshot dirs. + backend = LlamaCppBackend() + repo = "org/split" + files = [ + "model-Q4_K_M-00001-of-00002.gguf", + "model-Q4_K_M-00002-of-00002.gguf", + ] + _build_cache(hf_cache, repo, {files[0]: 4}, snapshot_sha = "a" * 40) + _build_cache(hf_cache, repo, {files[1]: 4}, snapshot_sha = "b" * 40) + downloaded: list[str] = [] + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p] + + def fake_download( + repo_id, + filename, + token = None, + **_kwargs, + ): + downloaded.append(filename) + return f"/fake/{repo_id}/{filename}" + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: files), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + ): + out = backend._download_gguf(hf_repo = repo, hf_variant = "Q4_K_M") + + assert downloaded == files + assert out == f"/fake/{repo}/{files[0]}" + def _siblings(items: dict[str, int]): """Mock ``hf_model_info(...).siblings`` payload.""" @@ -315,6 +567,21 @@ class TestIterHfCacheSnapshots: out = list(_iter_hf_cache_snapshots("unsloth/multi")) assert [p.name for p in out] == ["b" * 40, "a" * 40] + def test_skips_snapshot_when_mtime_is_unavailable(self, hf_cache, monkeypatch): + stale = _build_cache(hf_cache, "unsloth/multi", {"x.gguf": 1}, snapshot_sha = "a" * 40) + good = _build_cache(hf_cache, "unsloth/multi", {"y.gguf": 1}, snapshot_sha = "b" * 40) + original_stat = Path.stat + + def flaky_stat(self, *args, **kwargs): + if self == stale: + raise FileNotFoundError(str(self)) + return original_stat(self, *args, **kwargs) + + monkeypatch.setattr(Path, "stat", flaky_stat) + + out = list(_iter_hf_cache_snapshots("unsloth/multi")) + assert out == [good] + def test_repo_id_match_is_case_insensitive(self, hf_cache): _build_cache(hf_cache, "unsloth/Foo-GGUF", {"Foo-Q4_K_M.gguf": 1}) # Lookup with different org/name casing still resolves @@ -347,6 +614,87 @@ class TestListGgufVariantsFromCache: assert _list_gguf_variants_from_hf_cache("unsloth/absent") is None +class TestCachedColocatedSplitMain: + def test_prefers_older_complete_snapshot_over_newer_partial(self, hf_cache): + # Newer snapshot has only shard 1; older snapshot has the complete set. The + # complete older snapshot must win so the split GGUF can load co-located. + shard1 = "m-00001-of-00002.gguf" + shard2 = "m-00002-of-00002.gguf" + old = _build_cache( + hf_cache, "unsloth/split-GGUF", {shard1: 100, shard2: 100}, snapshot_sha = "a" * 40 + ) + new = _build_cache(hf_cache, "unsloth/split-GGUF", {shard1: 100}, snapshot_sha = "b" * 40) + os.utime(old, (1000, 1000)) + os.utime(new, (2000, 2000)) + + main = _cached_colocated_split_main("unsloth/split-GGUF", shard1, [shard2], {}) + assert main is not None + assert main.startswith(str(old)) + + def test_returns_none_when_shards_span_snapshots(self, hf_cache): + shard1 = "m-00001-of-00002.gguf" + shard2 = "m-00002-of-00002.gguf" + a = _build_cache(hf_cache, "unsloth/split-GGUF", {shard1: 100}, snapshot_sha = "a" * 40) + b = _build_cache(hf_cache, "unsloth/split-GGUF", {shard2: 100}, snapshot_sha = "b" * 40) + os.utime(a, (1000, 1000)) + os.utime(b, (2000, 2000)) + + assert _cached_colocated_split_main("unsloth/split-GGUF", shard1, [shard2], {}) is None + + +class TestResolveRepoIdCasing: + def test_maps_to_canonical_casing(self, monkeypatch): + monkeypatch.setattr( + "utils.paths.resolve_cached_repo_id_case", + lambda repo: "unsloth/Gemma-4-GGUF" if repo.lower() == "unsloth/gemma-4-gguf" else repo, + ) + # A companion download passed the resolved id reads the same cache entry + # as the main GGUF instead of missing it under the requested casing. + assert _resolve_repo_id_casing("unsloth/gemma-4-gguf") == "unsloth/Gemma-4-GGUF" + + def test_passthrough_on_resolver_error(self, monkeypatch): + def boom(_repo): + raise RuntimeError("resolver unavailable") + + monkeypatch.setattr("utils.paths.resolve_cached_repo_id_case", boom) + assert _resolve_repo_id_casing("unsloth/gemma-4-gguf") == "unsloth/gemma-4-gguf" + + def test_companion_only_newer_snapshot_does_not_shadow_real_variants(self, hf_cache): + # A newer snapshot holds only a vision projector fetched on demand, + # while the quant files live in an older snapshot. The newer snapshot + # must not shadow the real variants; the vision flag carries over. + old = _build_cache( + hf_cache, + "unsloth/vision-GGUF", + {"vision-Q4_K_M.gguf": 100}, + snapshot_sha = "a" * 40, + ) + new = _build_cache( + hf_cache, + "unsloth/vision-GGUF", + {"mmproj-vision-F16.gguf": 10}, + snapshot_sha = "b" * 40, + ) + os.utime(old, (1000, 1000)) + os.utime(new, (2000, 2000)) + + out = _list_gguf_variants_from_hf_cache("unsloth/vision-GGUF") + assert out is not None + variants, has_vision = out + assert [v.quant for v in variants] == ["Q4_K_M"] + assert has_vision is True + + def test_companion_only_cache_returns_empty_variants_with_vision(self, hf_cache): + # Only a vision projector is cached anywhere: report the vision flag + # with an empty variant list rather than None. + _build_cache(hf_cache, "unsloth/vision-GGUF", {"mmproj-vision-F16.gguf": 10}) + out = _list_gguf_variants_from_hf_cache("unsloth/vision-GGUF") + assert out is not None + variants, has_vision = out + assert variants == [] + assert has_vision is True + + class TestListGgufVariantsOffline: def test_offline_env_short_circuits_api(self, hf_cache, clean_offline_env, monkeypatch): _build_cache(hf_cache, "unsloth/a", {"a-UD-Q4_K_XL.gguf": 1}) diff --git a/studio/backend/tests/test_openai_auto_switch.py b/studio/backend/tests/test_openai_auto_switch.py index 7d2e2213b3..d02a2a4f7e 100644 --- a/studio/backend/tests/test_openai_auto_switch.py +++ b/studio/backend/tests/test_openai_auto_switch.py @@ -3037,3 +3037,60 @@ def test_acquire_swap_gate_is_cancellation_safe(): inference_route._auto_switch_process_lock.release() asyncio.run(asyncio.wait_for(main(), timeout = 5)) + + +def test_no_model_loaded_detail_appends_hint_only_when_off(monkeypatch): + # The "no model loaded" errors point at the opt-in auto-switch toggle so a + # request naming a listed-but-unloaded model is self-explanatory -- but only + # when it's off. With it on the name simply didn't resolve, so no hint. + base = "No GGUF model loaded. Load a GGUF model first." + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False) + off = inference_route._no_model_loaded_detail(base) + assert off.startswith(base) + assert "Model auto-switch" in off and "Settings > API" in off + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True) + assert inference_route._no_model_loaded_detail(base) == base + + +def _run_responses_stream_no_model(monkeypatch, *, enabled, active_model_name): + # Drive _responses_stream's GGUF-not-loaded guard: llama backend unloaded, + # inference backend maybe holding a non-GGUF model. Returns the 400 detail. + from fastapi import HTTPException + from models.inference import ResponsesRequest, ChatMessage + + monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: enabled) + monkeypatch.setattr( + inference_route, "get_llama_cpp_backend", lambda: _FakeBackend(loaded_id = None) + ) + monkeypatch.setattr( + inference_route, + "get_inference_backend", + lambda: type("_B", (), {"active_model_name": active_model_name})(), + ) + payload = ResponsesRequest(model = "unsloth/Qwen3.5-4B-GGUF", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + with pytest.raises(HTTPException) as exc: + asyncio.run(inference_route._responses_stream(payload, messages, None)) + assert exc.value.status_code == 400 + return exc.value.detail + + +def test_responses_stream_hint_matches_toggle_regardless_of_active_model(monkeypatch): + # Streaming /v1/responses shares the GGUF-only 400 with the other "no model + # loaded" sites, so the auto-switch hint attaches whenever the toggle is + # off -- including while a non-GGUF model is active, since auto-switch + # evicts it to load a resolved GGUF (_maybe_auto_switch_model's resolver + # branch has no active-model guard, unlike its reload-stash branch). Only + # the toggle being on suppresses it. + hinted = _run_responses_stream_no_model(monkeypatch, enabled = False, active_model_name = None) + assert "Model auto-switch" in hinted + + on = _run_responses_stream_no_model(monkeypatch, enabled = True, active_model_name = None) + assert "Model auto-switch" not in on + + non_gguf_loaded = _run_responses_stream_no_model( + monkeypatch, enabled = False, active_model_name = "unsloth/Llama-3.2-1B-Instruct" + ) + assert "Model auto-switch" in non_gguf_loaded diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index ccbd78e2b1..910818d7d8 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -8,6 +8,7 @@ import sys import asyncio import json import threading +import time from types import SimpleNamespace _backend = os.path.join(os.path.dirname(__file__), "..") @@ -29,7 +30,17 @@ from core.inference.anthropic_compat import ( anthropic_tool_choice_to_openai, ) from core.inference.api_monitor import ApiMonitor +from core.inference.llama_admission import ( + ADMISSION_KEEPALIVE_INTERVAL_ENV, + ADMISSION_MAX_QUEUE_ENV, + ADMISSION_QUEUE_TIMEOUT_ENV, + LlamaAdmissionCancelled, + LlamaAdmissionConfig, + get_llama_admission_queue, + reset_llama_admission_queues, +) from routes.inference import ( + _aclose_stream_resources, _build_chat_request, _build_openai_passthrough_body, _build_passthrough_payload, @@ -38,24 +49,69 @@ from routes.inference import ( _coalesce_consecutive_user_turns, _drop_empty_assistant_sentinels, _effective_max_tokens, + _effective_openai_max_tokens, + _effective_openai_max_tokens_from_values, _extract_content_parts, _friendly_error, _friendly_upstream_error, _merge_user_content, _monitor_openai_chunk, _monitor_openai_sse_event, + _normalize_openai_passthrough_sse_line, + _openai_compat_stream_stall_timeout, + _openai_llama_admission_capacity, _openai_messages_for_gguf_chat, + _openai_passthrough_sse_line_terminal_state, + _openai_passthrough_upstream_headers, _openai_passthrough_non_streaming, _openai_passthrough_stream, + _responses_stream, + _openai_stream_error_sse, _openai_stream_usage_chunk, + _openai_admission_wait_stream_chunks, + _wait_for_openai_admission_non_streaming, _proxy_to_external_provider, _SameTaskStreamingResponse, + _OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, _set_or_prepend_system_message, openai_completions, openai_embeddings, openai_chat_completions, ) -from state.tool_policy import reset_tool_policy +from state.tool_policy import reset_tool_policy, set_tool_policy + + +@pytest.fixture(autouse = True) +def _reset_admission_queues(): + reset_llama_admission_queues() + yield + reset_llama_admission_queues() + + +def test_aclose_stream_resources_attempts_remaining_closes_after_cancel(): + class Closeable: + def __init__(self, *, cancel = False): + self.cancel = cancel + self.closed = False + + async def aclose(self): + self.closed = True + if self.cancel: + raise asyncio.CancelledError() + + async def _run(): + iterator = Closeable(cancel = True) + resp = Closeable() + client = Closeable() + + with pytest.raises(asyncio.CancelledError): + await _aclose_stream_resources(iterator = iterator, resp = resp, client = client) + + assert iterator.closed + assert resp.closed + assert client.closed + + asyncio.run(_run()) class TestFriendlyUpstreamError: @@ -548,6 +604,263 @@ class TestChatCompletionRequestToolFields: assert "n > 1 is not supported" in entry["error"] assert monitor.active_count() == 0 + def test_client_tools_rejected_when_gguf_template_has_no_tool_support(self, monkeypatch): + import routes.inference as inference_route + + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + is_vision = False + _is_audio = False + context_length = 4096 + + def generate_chat_completion(self, **_kwargs): + raise AssertionError("client tools must not fall through to the standard GGUF path") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object"}, + }, + } + ], + }, + ) + + self._assert_unsupported_param(resp, "tools") + assert "does not advertise tools" in resp.json()["error"]["message"] + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "does not advertise tools" in entry["error"] + assert monitor.active_count() == 0 + + def test_client_tools_use_passthrough_capability_when_tool_loop_is_disabled(self, monkeypatch): + import routes.inference as inference_route + + captured = {} + + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + supports_tool_passthrough = True + is_vision = False + _is_audio = False + context_length = 4096 + base_url = "http://llama.passthrough-capability.test" + _request_reasoning_kwargs = lambda *_args, **_kwargs: None + + def generate_chat_completion(self, **_kwargs): + raise AssertionError("client tools must use passthrough") + + def generate_chat_completion_with_tools(self, **_kwargs): + raise AssertionError("Studio tool loop must stay disabled") + + async def fake_passthrough(llama_backend, payload, model_name, **kwargs): + captured["body"] = inference_route._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + inference_route.api_monitor.finish(kwargs.get("monitor_id")) + return inference_route.JSONResponse({"ok": True, "model": model_name}) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) + monkeypatch.setattr( + inference_route, + "_openai_passthrough_non_streaming", + fake_passthrough, + ) + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "use client tool"}], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object"}, + }, + } + ], + }, + ) + + assert resp.status_code == 200 + assert resp.json()["ok"] is True + assert captured["body"]["tools"][0]["function"]["name"] == "lookup" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert monitor.active_count() == 0 + + def test_enable_tools_on_non_tool_backend_keeps_client_tools_on_passthrough(self, monkeypatch): + # DiffusionGemma forces supports_tools off while passthrough stays + # available (#6851): enable_tools=True must not steal client tools + # from the passthrough into a Studio tool loop that cannot run. + import routes.inference as inference_route + + captured = {} + + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + supports_tool_passthrough = True + is_vision = False + _is_audio = False + context_length = 4096 + base_url = "http://llama.passthrough-capability.test" + _request_reasoning_kwargs = lambda *_args, **_kwargs: None + + def generate_chat_completion(self, **_kwargs): + raise AssertionError("client tools must use passthrough") + + def generate_chat_completion_with_tools(self, **_kwargs): + raise AssertionError("Studio tool loop cannot run on a non-tool backend") + + async def fake_passthrough(llama_backend, payload, model_name, **kwargs): + captured["body"] = inference_route._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + inference_route.api_monitor.finish(kwargs.get("monitor_id")) + return inference_route.JSONResponse({"ok": True, "model": model_name}) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) + monkeypatch.setattr( + inference_route, + "_openai_passthrough_non_streaming", + fake_passthrough, + ) + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "use client tool"}], + "enable_tools": True, + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object"}, + }, + } + ], + }, + ) + + assert resp.status_code == 200 + assert resp.json()["ok"] is True + assert captured["body"]["tools"][0]["function"]["name"] == "lookup" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert monitor.active_count() == 0 + + def test_tool_choice_none_allows_tool_catalog_without_tool_template(self, monkeypatch): + import routes.inference as inference_route + + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + is_vision = False + _is_audio = False + context_length = 4096 + + def generate_chat_completion(self, **kwargs): + assert kwargs["max_tokens"] is None + yield "plain response" + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object"}, + }, + } + ], + "tool_choice": "none", + }, + ) + + assert resp.status_code == 200 + assert resp.json()["choices"][0]["message"]["content"] == "plain response" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "plain response" + assert monitor.active_count() == 0 + + def test_tool_call_history_rejected_when_gguf_template_has_no_tool_support(self, monkeypatch): + import routes.inference as inference_route + + class _GGUFBackend: + is_loaded = True + model_identifier = "test-gguf" + supports_tools = False + is_vision = False + _is_audio = False + context_length = 4096 + + def generate_chat_completion(self, **_kwargs): + raise AssertionError( + "tool-call history must not fall through to the standard GGUF path" + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inference_route, "api_monitor", monitor) + client = self._v1_client(monkeypatch, _GGUFBackend()) + resp = client.post( + "/v1/chat/completions", + json = { + "messages": [ + {"role": "user", "content": "use a tool"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "{}"}, + ], + }, + ) + + self._assert_unsupported_param(resp, "messages") + assert "does not advertise tools" in resp.json()["error"]["message"] + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "does not advertise tools" in entry["error"] + assert monitor.active_count() == 0 + def test_n_rejected_for_non_gguf_path(self, monkeypatch): class _NoGGUFBackend: is_loaded = False @@ -748,11 +1061,32 @@ class TestBuildPassthroughPayloadToolChoice: ) assert body.get("stream_options") == {"include_usage": False} + def test_response_format_without_tools_omits_tool_fields(self): + args = self._args() + args["openai_tools"] = None + + body = _build_passthrough_payload( + **args, + response_format = {"type": "json_object"}, + ) + + assert body["response_format"] == {"type": "json_object"} + assert "tools" not in body + assert "tool_choice" not in body + def test_repetition_penalty_renamed(self): body = _build_passthrough_payload(**self._args(), repetition_penalty = 1.1) assert body.get("repeat_penalty") == 1.1 assert "repetition_penalty" not in body + def test_omitted_passthrough_max_tokens_uses_backend_context(self): + args = self._args() + args["max_tokens"] = None + + body = _build_passthrough_payload(**args, backend_ctx = 4096) + + assert body["max_tokens"] == 4096 + def test_passthrough_body_merges_system_and_developer_messages(self): payload = ChatCompletionRequest( model = "default", @@ -772,6 +1106,74 @@ class TestBuildPassthroughPayloadToolChoice: ] +class TestOpenAIPassthroughSSETerminalState: + def test_done_sentinel(self): + assert _openai_passthrough_sse_line_terminal_state("data: [DONE]") == "done" + + def test_finish_reason_with_space(self): + line = 'data: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}' + assert _openai_passthrough_sse_line_terminal_state(line) == "finish" + + def test_finish_reason_without_space(self): + line = 'data:{"choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}' + assert _openai_passthrough_sse_line_terminal_state(line) == "finish" + + def test_usage_chunk(self): + line = 'data: {"choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2}}' + assert _openai_passthrough_sse_line_terminal_state(line) == "usage" + + def test_error_chunk(self): + line = 'data: {"error":{"message":"boom"}}' + assert _openai_passthrough_sse_line_terminal_state(line) == "error" + + def test_cap_parallel_tool_calls_accepts_no_space_after_data_colon(self): + line = ( + 'data:{"choices":[{"delta":{"tool_calls":[' + '{"index":0,"function":{"name":"a"}},' + '{"index":1,"function":{"name":"b"}}]}}]}' + ) + + capped = _normalize_openai_passthrough_sse_line(line, cap_parallel_tool_calls = True) + + data = json.loads(capped[len("data:") :].lstrip()) + assert data["choices"][0]["delta"]["tool_calls"] == [ + {"index": 0, "function": {"name": "a"}} + ] + + def test_plain_content_line_is_returned_identically(self): + # The relay dispatches terminal classification on `out_line is raw_line`, + # so the no-mutation path must return the identical string object. + line = 'data: {"choices":[{"index":0,"delta":{"content":"hello"},"finish_reason":null}]}' + assert _normalize_openai_passthrough_sse_line(line) is line + assert _normalize_openai_passthrough_sse_line(line, cap_parallel_tool_calls = True) is line + + def test_reasoning_key_inside_content_text_keeps_line_identical(self): + # Fast-path substring gate fires, but the parse finds nothing to change: + # the original object must come back so the relay stays byte-identical. + line = ( + 'data: {"choices":[{"index":0,"delta":{"content":' + '"mentions \\"reasoning_content\\" in text"},"finish_reason":null}]}' + ) + assert _normalize_openai_passthrough_sse_line(line) is line + + def test_reasoning_only_delta_gets_empty_content(self): + line = ( + 'data: {"choices":[{"index":0,' + '"delta":{"reasoning_content":"thinking"},' + '"finish_reason":null}]}' + ) + + normalized = _normalize_openai_passthrough_sse_line(line) + + data = json.loads(normalized[len("data:") :].lstrip()) + delta = data["choices"][0]["delta"] + assert delta["reasoning_content"] == "thinking" + assert delta["content"] == "" + + def test_reasoning_normalization_preserves_done_sentinel(self): + assert _normalize_openai_passthrough_sse_line("data: [DONE]") == "data: [DONE]" + + # ===================================================================== # Passthrough reasoning kwargs — enable_thinking / reasoning_effort / # preserve_thinking must reach llama-server via chat_template_kwargs, @@ -892,6 +1294,172 @@ class TestOpenAICompatibilityHelpers: payload = SimpleNamespace(max_tokens = 128, max_completion_tokens = 64) assert _effective_max_tokens(payload) == 64 + def test_openai_compat_max_tokens_returns_none_when_omitted(self): + payload = SimpleNamespace(max_tokens = None, max_completion_tokens = None) + assert _effective_openai_max_tokens(payload) is None + + @pytest.mark.parametrize( + ("payload", "expected"), + [ + (SimpleNamespace(max_tokens = 8192, max_completion_tokens = None), 8192), + (SimpleNamespace(max_tokens = 8192, max_completion_tokens = 256), 256), + ], + ) + def test_openai_compat_explicit_values_pass_through(self, payload, expected): + assert _effective_openai_max_tokens(payload) == expected + + @pytest.mark.parametrize( + ("payload", "param"), + [ + (SimpleNamespace(max_tokens = "128", max_completion_tokens = None), "max_tokens"), + (SimpleNamespace(max_tokens = True, max_completion_tokens = None), "max_tokens"), + (SimpleNamespace(max_tokens = 12.5, max_completion_tokens = None), "max_tokens"), + ( + SimpleNamespace(max_tokens = None, max_completion_tokens = "128"), + "max_completion_tokens", + ), + ], + ) + def test_openai_compat_max_tokens_rejects_non_integer_explicit_values(self, payload, param): + with pytest.raises(HTTPException) as exc: + _effective_openai_max_tokens(payload) + + assert exc.value.status_code == 400 + assert exc.value.detail["error"]["param"] == param + assert exc.value.detail["error"]["code"] == "invalid_type" + + def test_openai_compat_max_tokens_zero_is_valid_and_negative_rejected(self): + # Legacy completions spec: max_tokens has minimum 0, so 0 must pass + # through; only negatives are invalid_value. + assert _effective_openai_max_tokens_from_values(0) == 0 + + with pytest.raises(HTTPException) as exc: + _effective_openai_max_tokens_from_values(-1) + + assert exc.value.status_code == 400 + assert exc.value.detail["error"]["code"] == "invalid_value" + assert exc.value.detail["error"]["param"] == "max_tokens" + + def test_chat_reasoning_chunk_carries_empty_content(self): + from routes.inference import _chat_reasoning_chunk + + line = _chat_reasoning_chunk("chatcmpl-test", 123, "gguf", "thinking...") + chunk = json.loads(line[len("data: ") :]) + delta = chunk["choices"][0]["delta"] + + assert delta["reasoning_content"] == "thinking..." + assert delta["content"] == "" + + def test_passthrough_upstream_headers_include_backend_auth(self): + headers = _openai_passthrough_upstream_headers( + llama_backend = SimpleNamespace(_auth_headers = {"Authorization": "Bearer secret"}), + ) + + assert headers["Authorization"] == "Bearer secret" + assert headers["Connection"] == "close" + + def test_openai_admission_capacity_prefers_backend_effective_slots(self): + request = SimpleNamespace( + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + ) + backend = SimpleNamespace(effective_parallel_slots = 3) + + assert _openai_llama_admission_capacity(request, backend) == 3 + + @pytest.mark.parametrize("backend_value", [None, 0, -1, "not-an-int"]) + def test_openai_admission_capacity_falls_back_to_app_state(self, backend_value): + request = SimpleNamespace( + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 2)) + ) + backend = SimpleNamespace(effective_parallel_slots = backend_value) + + assert _openai_llama_admission_capacity(request, backend) == 2 + + def test_openai_admission_capacity_falls_back_to_one_without_request(self): + assert _openai_llama_admission_capacity(None, SimpleNamespace()) == 1 + + def test_openai_admission_non_streaming_exits_invalidated_waiter(self): + async def _run(): + queue = get_llama_admission_queue("http://llama.invalidated.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + reservation = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()) + assert reservation._waiter is not None + + reservation._waiter.future.cancel() + + with pytest.raises(LlamaAdmissionCancelled): + await asyncio.wait_for( + _wait_for_openai_admission_non_streaming( + reservation, + LlamaAdmissionConfig(), + request = None, + cancel_event = None, + ), + timeout = 0.1, + ) + + blocker.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_openai_admission_stream_exits_invalidated_waiter(self): + async def _run(): + queue = get_llama_admission_queue("http://llama.invalidated.stream.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + reservation = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()) + assert reservation._waiter is not None + + reservation._waiter.future.cancel() + + chunks = _openai_admission_wait_stream_chunks( + reservation, + LlamaAdmissionConfig(), + request = None, + cancel_event = None, + ) + with pytest.raises(LlamaAdmissionCancelled): + await asyncio.wait_for(chunks.__anext__(), timeout = 0.1) + + blocker.release() + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_openai_compat_stream_stall_timeout_uses_default(self, monkeypatch): + monkeypatch.delenv(_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, raising = False) + assert _openai_compat_stream_stall_timeout() == 120.0 + + def test_openai_compat_stream_stall_timeout_uses_env_override(self, monkeypatch): + monkeypatch.setenv(_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, "4.5") + assert _openai_compat_stream_stall_timeout() == 4.5 + + @pytest.mark.parametrize("raw_value", ["", "not-a-float"]) + def test_openai_compat_stream_stall_timeout_invalid_env_uses_default( + self, monkeypatch, raw_value + ): + monkeypatch.setenv(_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, raw_value) + assert _openai_compat_stream_stall_timeout() == 120.0 + + @pytest.mark.parametrize("raw_value", ["0", "-1"]) + def test_openai_compat_stream_stall_timeout_non_positive_env_disables( + self, monkeypatch, raw_value + ): + monkeypatch.setenv(_OPENAI_COMPAT_STREAM_STALL_TIMEOUT_ENV, raw_value) + assert _openai_compat_stream_stall_timeout() is None + + def test_openai_stream_error_sse_closes_with_done(self): + error = {"error": {"message": "boom"}} + assert _openai_stream_error_sse(error) == ( + 'data: {"error": {"message": "boom"}}\n\n' "data: [DONE]\n\n" + ) + @pytest.mark.parametrize( "finish_reason", ["stop", "length", "tool_calls", "content_filter", "function_call"], @@ -1515,9 +2083,659 @@ class TestGgufVisionToolRouting: assert "".join(d.get("reasoning_content", "") for d in deltas) == "plan" assert "".join(d.get("content", "") for d in deltas) == "visible" assert all("" not in d.get("content", "") for d in deltas) + assert all("content" in d for d in deltas if "reasoning_content" in d) [entry] = result.monitor.snapshot() assert entry["reply"] == "visible" + def test_standard_gguf_stream_queued_request_sends_keepalive_before_generation( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + def _generate(**_kwargs): + raise AssertionError("standard GGUF generation must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + queue = get_llama_admission_queue("http://llama.standard.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + stream = True, + ) + response = await openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + snapshot = queue.snapshot() + assert snapshot.active == 1 + assert snapshot.queued == 1 + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_standard_gguf_stream_close_after_first_chunk_cleans_tracker(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + cancel_id = "standard-stream-close-cleanup" + + def _generate(**_kwargs): + yield "visible" + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + stream = True, + cancel_id = cancel_id, + ) + response = await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + iterator = response.body_iterator + assert cancel_id in inf_mod._CANCEL_REGISTRY + await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + aclose = getattr(iterator, "aclose", None) + assert aclose is not None + await aclose() + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert get_llama_admission_queue("http://llama.standard.test").snapshot().active == 0 + + asyncio.run(_run()) + + def test_standard_gguf_stream_task_cancel_after_first_chunk_finalizes_monitor( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + started = threading.Event() + released = threading.Event() + + def _generate(**kwargs): + cancel_event = kwargs["cancel_event"] + started.set() + while not cancel_event.is_set(): + time.sleep(0.005) + released.set() + yield from () + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + stream = True, + ) + response = await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + iterator = response.body_iterator + assert await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + pending = asyncio.create_task(iterator.__anext__()) + assert await asyncio.to_thread(started.wait, 1.0) + + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, timeout = 1.0) + + assert released.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + assert get_llama_admission_queue("http://llama.standard.test").snapshot().active == 0 + + asyncio.run(_run()) + + def test_gguf_tool_stream_queued_request_sends_keepalive_before_generation(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + async def fake_select_tools(*_args, **_kwargs): + return [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _generate(**_kwargs): + raise AssertionError("GGUF tool loop must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.tool.test", + effective_parallel_slots = 1, + generate_chat_completion = lambda **_kwargs: "unused", + generate_chat_completion_with_tools = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) + + queue = get_llama_admission_queue("http://llama.tool.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + enable_tools = True, + stream = True, + ) + response = await openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + snapshot = queue.snapshot() + assert snapshot.active == 1 + assert snapshot.queued == 1 + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_gguf_tool_stream_task_cancel_after_first_chunk_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_select_tools(*_args, **_kwargs): + return [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + started = threading.Event() + released = threading.Event() + + def _tools(**kwargs): + cancel_event = kwargs["cancel_event"] + started.set() + while not cancel_event.is_set(): + time.sleep(0.005) + released.set() + yield from () + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.tool.test", + effective_parallel_slots = 1, + generate_chat_completion = lambda **_kwargs: "unused", + generate_chat_completion_with_tools = _tools, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + enable_tools = True, + stream = True, + ) + response = await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + iterator = response.body_iterator + assert await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + pending = asyncio.create_task(iterator.__anext__()) + assert await asyncio.to_thread(started.wait, 1.0) + + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, timeout = 1.0) + + assert released.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + assert get_llama_admission_queue("http://llama.tool.test").snapshot().active == 0 + + asyncio.run(_run()) + + def test_global_enable_tools_does_not_preempt_response_format_passthrough(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + set_tool_policy(True) + captured = {} + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**_kwargs): + raise AssertionError("Studio tool loop should not steal response_format") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.policy.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + + async def fake_passthrough(llama_backend, payload, model_name, **_kwargs): + captured["body"] = inf_mod._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + return inf_mod.JSONResponse({"ok": True, "model": model_name}) + + try: + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming", + fake_passthrough, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "json"}], + response_format = {"type": "json_object"}, + ) + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + + assert json.loads(response.body)["ok"] is True + assert captured["body"]["response_format"] == {"type": "json_object"} + assert "tools" not in captured["body"] + assert "tool_choice" not in captured["body"] + finally: + reset_tool_policy() + + def test_global_enable_tools_does_not_replace_client_tools_passthrough(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + set_tool_policy(True) + captured = {} + client_tools = [ + { + "type": "function", + "function": { + "name": "client_lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**_kwargs): + raise AssertionError("Studio tool loop should not replace client tools") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.policy.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + + async def fake_passthrough(llama_backend, payload, model_name, **_kwargs): + captured["body"] = inf_mod._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + return inf_mod.JSONResponse({"ok": True, "model": model_name}) + + try: + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming", + fake_passthrough, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "use client tool"}], + tools = client_tools, + ) + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + + assert json.loads(response.body)["ok"] is True + assert captured["body"]["tools"] == client_tools + assert captured["body"]["tool_choice"] == "auto" + finally: + reset_tool_policy() + + def test_global_enable_tools_honors_client_tool_choice_none(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + set_tool_policy(True) + client_tools = [ + { + "type": "function", + "function": { + "name": "client_lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _plain(**kwargs): + assert kwargs["max_tokens"] is None + yield "plain response" + + def _tools(**_kwargs): + raise AssertionError("tool_choice='none' must not start Studio's tool loop") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.policy.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + + try: + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "do not use tools"}], + tools = client_tools, + tool_choice = "none", + ) + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + + assert json.loads(response.body)["choices"][0]["message"]["content"] == "plain response" + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert entry["reply"] == "plain response" + assert monitor.active_count() == 0 + finally: + reset_tool_policy() + + def test_enabled_tools_without_enable_tools_keeps_response_format_passthrough( + self, monkeypatch + ): + import routes.inference as inf_mod + + reset_tool_policy() + captured = {} + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**_kwargs): + raise AssertionError("enabled_tools alone must not start Studio's tool loop") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.enabled-tools.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + + async def fake_passthrough(llama_backend, payload, model_name, **_kwargs): + captured["body"] = inf_mod._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + return inf_mod.JSONResponse({"ok": True, "model": model_name}) + + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_openai_passthrough_non_streaming", fake_passthrough) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "json"}], + enabled_tools = ["web_search"], + response_format = {"type": "json_object"}, + ) + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + + assert json.loads(response.body)["ok"] is True + assert captured["body"]["response_format"] == {"type": "json_object"} + + def test_enabled_tools_without_enable_tools_keeps_client_tools_passthrough(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + captured = {} + client_tools = [ + { + "type": "function", + "function": { + "name": "client_lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**_kwargs): + raise AssertionError("enabled_tools alone must not start Studio's tool loop") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.enabled-tools.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + + async def fake_passthrough(llama_backend, payload, model_name, **_kwargs): + captured["body"] = inf_mod._build_openai_passthrough_body( + payload, + backend_ctx = llama_backend.context_length, + llama_backend = llama_backend, + ) + return inf_mod.JSONResponse({"ok": True, "model": model_name}) + + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_openai_passthrough_non_streaming", fake_passthrough) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "use client tool"}], + enabled_tools = ["web_search"], + tools = client_tools, + ) + response = self._drive( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + + assert json.loads(response.body)["ok"] is True + assert captured["body"]["tools"] == client_tools + assert captured["body"]["tool_choice"] == "auto" + def test_reasoning_capable_gguf_stream_splits_reasoning_by_default(self, monkeypatch): def _generate(**_kwargs): yield "planvisible" @@ -1639,6 +2857,313 @@ class TestGgufVisionToolRouting: [entry] = result.monitor.snapshot() assert entry["reply"] == "visible" + def test_standard_gguf_non_streaming_admission_timeout_before_generation(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + def _generate(**_kwargs): + raise AssertionError("standard GGUF generation must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + queue = get_llama_admission_queue("http://llama.standard.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + ) + try: + with pytest.raises(HTTPException) as exc: + await openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + assert exc.value.status_code == 503 + finally: + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_standard_gguf_non_streaming_cancel_id_stops_queued_request_before_generation( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + def _generate(**_kwargs): + raise AssertionError("standard GGUF generation must not start after cancel_id") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + queue = get_llama_admission_queue("http://llama.standard.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + cancel_id = "standard-nonstream-admission-cancel" + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + cancel_id = cancel_id, + ) + task = asyncio.create_task( + openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + ) + try: + for _ in range(50): + if cancel_id in inf_mod._CANCEL_REGISTRY: + break + await asyncio.sleep(0.01) + assert cancel_id in inf_mod._CANCEL_REGISTRY + assert inf_mod._cancel_by_cancel_id_or_stash(cancel_id) == 1 + with pytest.raises(HTTPException) as exc: + await asyncio.wait_for(task, timeout = 0.5) + assert exc.value.status_code == 499 + finally: + if not task.done(): + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + blocker.release() + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_standard_gguf_non_streaming_admission_task_cancel_cleans_tracker_and_slot( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + cancel_id = "standard-nonstream-task-cancel" + + async def fake_wait(*_args, **_kwargs): + raise asyncio.CancelledError() + + def _generate(**_kwargs): + raise AssertionError("standard GGUF generation must not start after task cancel") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.standard.test", + effective_parallel_slots = 1, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr( + inf_mod, + "_wait_for_openai_admission_non_streaming", + fake_wait, + ) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + cancel_id = cancel_id, + ) + with pytest.raises(asyncio.CancelledError): + await openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert get_llama_admission_queue("http://llama.standard.test").snapshot().active == 0 + + asyncio.run(_run()) + + def test_gguf_tool_non_streaming_admission_timeout_before_generation(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + async def fake_select_tools(*_args, **_kwargs): + return [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + def _generate(**_kwargs): + raise AssertionError("GGUF tool loop must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.tool.test", + effective_parallel_slots = 1, + generate_chat_completion = lambda **_kwargs: "unused", + generate_chat_completion_with_tools = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) + + queue = get_llama_admission_queue("http://llama.tool.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + enable_tools = True, + ) + try: + with pytest.raises(HTTPException) as exc: + await openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + assert exc.value.status_code == 503 + finally: + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_gguf_tool_non_streaming_cancel_drains_worker_before_releasing_slot(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_select_tools(*_args, **_kwargs): + return [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + + started = threading.Event() + released = threading.Event() + + def _tools(**kwargs): + cancel_event = kwargs["cancel_event"] + started.set() + while not cancel_event.is_set(): + time.sleep(0.005) + released.set() + yield from () + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.tool.test", + effective_parallel_slots = 1, + generate_chat_completion = lambda **_kwargs: "unused", + generate_chat_completion_with_tools = _tools, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_select_request_tools", fake_select_tools) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + enable_tools = True, + ) + task = asyncio.create_task( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + assert await asyncio.to_thread(started.wait, 1.0) + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout = 1.0) + + assert released.is_set() + assert get_llama_admission_queue("http://llama.tool.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + def test_non_streaming_gguf_n_records_all_monitor_replies(self, monkeypatch): import routes.inference as inf_mod @@ -1691,6 +3216,58 @@ class TestGgufVisionToolRouting: assert entry["completion_tokens"] == 3 assert monitor.active_count() == 0 + def test_non_streaming_gguf_cancel_drains_worker(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + started = threading.Event() + released = threading.Event() + + def _generate(**kwargs): + cancel_event = kwargs["cancel_event"] + started.set() + while not cancel_event.is_set(): + time.sleep(0.005) + released.set() + yield from () + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + ) + task = asyncio.create_task( + openai_chat_completions( + payload, + request = self._Request(), + current_subject = "test", + ) + ) + assert await asyncio.to_thread(started.wait, 1.0) + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout = 1.0) + + assert released.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + def test_standard_gguf_merges_system_and_developer_messages(self, monkeypatch): import routes.inference as inf_mod @@ -1801,7 +3378,12 @@ class TestApiMonitorProviderAndCompletionStreams: async def is_disconnected(self): return False - async def _run_passthrough_stream(self, monkeypatch, lines): + async def _run_passthrough_stream( + self, + monkeypatch, + lines, + stream_options = None, + ): import routes.inference as inf_mod class Request: @@ -1829,6 +3411,7 @@ class TestApiMonitorProviderAndCompletionStreams: model = "default", messages = [ChatMessage(role = "user", content = "hi")], stream = True, + stream_options = stream_options, tools = [ { "type": "function", @@ -1900,7 +3483,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) @@ -1913,6 +3496,134 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_stream_forwards_backend_auth_headers(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + captured_headers = {} + + async def fake_send(_client, req, *_args, **_kwargs): + captured_headers.update(dict(req.headers)) + return httpx.Response(200, content = b"") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _auth_headers = {"Authorization": "Bearer secret"}, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ) + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + + assert "data: [DONE]\n\n" in "".join(chunks) + assert captured_headers["authorization"] == "Bearer secret" + assert captured_headers["connection"] == "close" + + asyncio.run(_run()) + + def test_passthrough_stream_keepalive_while_upstream_headers_are_pending(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + + async def fake_send(*_args, **_kwargs): + await gate.wait() + return httpx.Response(200, content = b"") + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr( + inf_mod, + "_OPENAI_PASSTHROUGH_PENDING_RESPONSE_KEEPALIVE_S", + 0.01, + ) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + + first = await asyncio.wait_for(response.body_iterator.__anext__(), timeout = 0.2) + assert first == ": keep-alive\n\n" + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + body = "".join(chunks) + assert "data: [DONE]\n\n" in body + + asyncio.run(_run()) + def test_passthrough_stream_preheader_non_200_in_window(self, monkeypatch): async def _run(): import routes.inference as inf_mod @@ -2044,7 +3755,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) gate.set() @@ -2055,6 +3766,7 @@ class TestApiMonitorProviderAndCompletionStreams: body = "".join(chunks) assert "data:" in body assert '"error"' in body + assert "data: [DONE]" in body [entry] = monitor.snapshot() assert entry["status"] == "error" assert "bad" in entry["error"] @@ -2107,7 +3819,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) @@ -2117,7 +3829,13 @@ class TestApiMonitorProviderAndCompletionStreams: async for chunk in response.body_iterator ] body = "".join(chunks) - payload = json.loads(body.removeprefix("data: ").strip()) + events = [ + line.removeprefix("data: ") + for line in body.splitlines() + if line.startswith("data: ") + ] + assert events[-1] == "[DONE]" + payload = json.loads(events[0]) assert payload["error"]["code"] == "context_length_exceeded" assert payload["error"]["param"] == "messages" assert isinstance(payload["error"], dict) @@ -2190,7 +3908,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) @@ -2207,6 +3925,105 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_stream_preheader_immediate_context_retry_adopts_delayed_response( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + gate = asyncio.Event() + calls = [] + err_body = json.dumps( + { + "error": { + "message": "request (10000 tokens) exceeds the available context size (2048 tokens)", + "n_prompt_tokens": 10000, + "n_ctx": 2048, + } + } + ).encode("utf-8") + ok_lines = [ + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,' + '"model":"gguf","choices":[{"index":0,"delta":{"content":"OK"},' + '"finish_reason":null}]}', + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,' + '"model":"gguf","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + "data: [DONE]", + ] + + async def fake_send(_client, req, *_args, **_kwargs): + calls.append(json.loads(req.content.decode("utf-8"))) + if len(calls) == 1: + return httpx.Response(400, content = err_body) + await gate.wait() + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + for line in ok_lines: + yield line + + class Request: + async def is_disconnected(self): + return False + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + + messages = [ + ChatMessage(role = "system", content = "system"), + *[ + ChatMessage(role = "user", content = f"turn {idx} " + ("x" * 1000)) + for idx in range(8) + ], + ] + payload = ChatCompletionRequest( + model = "default", + messages = messages, + stream = True, + context_overflow = "truncate_middle", + ) + response = await asyncio.wait_for( + _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 2048, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "chatcmpl-test", + "chatcmpl-test", + monitor_id = monitor_id, + ), + timeout = 0.2, + ) + assert isinstance(response, _SameTaskStreamingResponse) + + gate.set() + chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in response.body_iterator + ] + body = "".join(chunks) + + assert "OK" in body + assert "context_length_exceeded" not in body + assert len(calls) == 2 + assert len(calls[1]["messages"]) < len(calls[0]["messages"]) + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + + asyncio.run(_run()) + def test_passthrough_stream_preheader_delayed_request_error_cleans_up(self, monkeypatch): async def _run(): import routes.inference as inf_mod @@ -2252,7 +4069,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) assert isinstance(response, _SameTaskStreamingResponse) assert cancel_id in inf_mod._CANCEL_REGISTRY @@ -2323,13 +4140,13 @@ class TestApiMonitorProviderAndCompletionStreams: monitor_id = monitor_id, ) ) - await asyncio.wait_for(entered.wait(), timeout = 0.2) + await asyncio.wait_for(entered.wait(), timeout = 5.0) assert cancel_id in inf_mod._CANCEL_REGISTRY task.cancel() with pytest.raises(asyncio.CancelledError): await task - await asyncio.wait_for(cancelled.wait(), timeout = 0.2) + await asyncio.wait_for(cancelled.wait(), timeout = 5.0) assert cancel_id not in inf_mod._CANCEL_REGISTRY asyncio.run(_run()) @@ -2389,13 +4206,13 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ), - timeout = 0.2, + timeout = 5.0, ) 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.wait_for(returned.wait(), timeout = 5.0) await asyncio.sleep(0) await response._unstarted_cleanup() assert upstream_response.is_closed @@ -2647,6 +4464,150 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_completions_omitted_max_tokens_falls_back_to_context(self, monkeypatch): + # With no env knobs set, an omitted max_tokens must forward the + # backend's context length, exactly as on main. + async def _run(): + import routes.inference as inf_mod + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/completions") + method = "POST" + + async def json(self): + return {"prompt": "hi", "stream": False} + + captured = [] + + class CapturingClient: + async def post(self, _url, *, json, **_kwargs): + captured.append(dict(json)) + return httpx.Response( + 200, + json = { + "id": "cmpl-test", + "choices": [{"text": "ok"}], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + }, + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: CapturingClient()) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + base_url = "http://llama.test", + context_length = 4096, + model_identifier = "gguf", + ), + ) + + await openai_completions(Request(), current_subject = "test") + + assert captured[0]["max_tokens"] == 4096 + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_completions_forwards_spec_valid_zero_max_tokens(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/completions") + method = "POST" + + async def json(self): + return {"prompt": "hi", "stream": False, "max_tokens": 0} + + captured = [] + + class CapturingClient: + async def post(self, _url, *, json, **_kwargs): + captured.append(dict(json)) + return httpx.Response( + 200, + json = { + "id": "cmpl-test", + "choices": [{"text": "", "finish_reason": "length"}], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 0, + "total_tokens": 1, + }, + }, + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: CapturingClient()) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + base_url = "http://llama.test", + context_length = 4096, + model_identifier = "gguf", + ), + ) + + await openai_completions(Request(), current_subject = "test") + + assert captured[0]["max_tokens"] == 0 + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_completions_rejects_non_integer_max_tokens_before_forwarding(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/completions") + method = "POST" + + async def json(self): + return {"prompt": "hi", "stream": False, "max_tokens": "128"} + + class UnusedClient: + async def post(self, *_args, **_kwargs): + raise AssertionError("invalid max_tokens must not reach llama-server") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "nonstreaming_client", lambda: UnusedClient()) + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + base_url = "http://llama.test", + context_length = 4096, + model_identifier = "gguf", + ), + ) + + with pytest.raises(HTTPException) as exc: + await openai_completions(Request(), current_subject = "test") + + assert exc.value.status_code == 400 + assert exc.value.detail["error"]["param"] == "max_tokens" + assert exc.value.detail["error"]["code"] == "invalid_type" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + def test_monitor_openai_chunk_records_all_choice_replies(self, monkeypatch): import routes.inference as inf_mod @@ -2793,6 +4754,8 @@ class TestApiMonitorProviderAndCompletionStreams: yield 'data: {"choices":[{"delta":{"content":"hello"}}]}' await asyncio.sleep(3600) + cancel_id = "passthrough-stream-delete-cancel" + monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inf_mod, "api_monitor", monitor) monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) @@ -2807,6 +4770,7 @@ class TestApiMonitorProviderAndCompletionStreams: model = "default", messages = [ChatMessage(role = "user", content = "hi")], stream = True, + cancel_id = cancel_id, tools = [ { "type": "function", @@ -2824,6 +4788,7 @@ class TestApiMonitorProviderAndCompletionStreams: SimpleNamespace( base_url = "http://llama.test", context_length = 4096, + _auth_headers = {"Authorization": "Bearer secret"}, _request_reasoning_kwargs = lambda *_args, **_kwargs: None, ), payload, @@ -2835,6 +4800,7 @@ class TestApiMonitorProviderAndCompletionStreams: iterator = response.body_iterator first = await anext(iterator) assert "hello" in first + assert cancel_id in inf_mod._CANCEL_REGISTRY pending = asyncio.create_task(anext(iterator)) await asyncio.sleep(0) @@ -2846,6 +4812,271 @@ class TestApiMonitorProviderAndCompletionStreams: assert entry["status"] == "cancelled" assert entry["reply"] == "hello" assert monitor.active_count() == 0 + assert cancel_id not in inf_mod._CANCEL_REGISTRY + + asyncio.run(_run()) + + def test_passthrough_stream_immediate_task_cancel_releases_admission_and_tracker( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + async def fake_cancel_check(*_args, **_kwargs): + raise asyncio.CancelledError() + + cancel_id = "passthrough-stream-immediate-task-cancel" + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_raise_if_openai_admission_cancelled", + fake_cancel_check, + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + backend = SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + + with pytest.raises(asyncio.CancelledError): + await _openai_passthrough_stream( + self._Request(), + threading.Event(), + backend, + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert get_llama_admission_queue("http://llama.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_queued_cancel_before_inner_first_chunk_runs_cleanup( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + body_holder = {} + cleanup_called = threading.Event() + + async def fake_admitted(*_args, admission_lease, tracker, **_kwargs): + async def cleanup(): + admission_lease.release() + tracker.__exit__(None, None, None) + cleanup_called.set() + + class BlockingBody: + def __init__(self): + self.started = threading.Event() + self.closed = False + + def __aiter__(self): + return self + + async def __anext__(self): + self.started.set() + await asyncio.sleep(3600) + raise StopAsyncIteration + + async def aclose(self): + self.closed = True + await cleanup() + + body = BlockingBody() + body_holder["body"] = body + return _SameTaskStreamingResponse( + body, + media_type = "text/event-stream", + unstarted_cleanup = cleanup, + ) + + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_stream_admitted", + fake_admitted, + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + cancel_id = "queued-inner-unstarted-cleanup" + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + assert cancel_id in inf_mod._CANCEL_REGISTRY + + blocker.release() + pending = asyncio.create_task(iterator.__anext__()) + for _ in range(100): + if "body" in body_holder: + break + await asyncio.sleep(0.01) + body = body_holder["body"] + assert await asyncio.to_thread(body.started.wait, 1.0) + + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, timeout = 1.0) + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + assert body_holder["body"].closed + assert cleanup_called.is_set() + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert queue.snapshot().active == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_queued_cancel_after_inner_first_chunk_finalizes_monitor( + self, monkeypatch + ): + async def _run(): + import routes.inference as inf_mod + + class Request(self._Request): + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + + async def fake_admitted( + *_args, + monitor_id = None, + admission_lease, + tracker, + **_kwargs, + ): + async def cleanup(): + admission_lease.release() + tracker.__exit__(None, None, None) + + async def body(): + try: + yield 'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n' + await asyncio.sleep(3600) + except asyncio.CancelledError: + inf_mod.api_monitor.finish(monitor_id, "cancelled") + raise + finally: + await cleanup() + + return _SameTaskStreamingResponse( + body(), + media_type = "text/event-stream", + unstarted_cleanup = cleanup, + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_stream_admitted", + fake_admitted, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + cancel_id = "queued-inner-cancel-monitor" + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + cancel_id = cancel_id, + ) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + + blocker.release() + first = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert "hello" in first + + pending = asyncio.create_task(iterator.__anext__()) + await asyncio.sleep(0) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, timeout = 1.0) + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + assert cancel_id not in inf_mod._CANCEL_REGISTRY + assert queue.snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 asyncio.run(_run()) @@ -2931,6 +5162,313 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_usage_done_are_separate_sse_events(self, monkeypatch): + async def _run(): + result = await self._run_passthrough_stream( + monkeypatch, + [ + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"m","choices":[{"index":0,"delta":{"role":"assistant"},"finish_reason":null}]}', + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"m","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}', + 'data: {"id":"chatcmpl-test","object":"chat.completion.chunk","created":1,"model":"m","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2}}', + ], + stream_options = {"include_usage": True}, + ) + + assert ( + '"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2' in result.body + ) + assert "data: [DONE]" in result.body + assert "}\n\ndata: [DONE]\n\n" in result.body + assert "}\ndata: [DONE]\n\n" not in result.body + + asyncio.run(_run()) + + def test_passthrough_stream_queued_request_sends_keepalive_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + url = SimpleNamespace(path = "/v1/chat/completions") + + async def is_disconnected(self): + return False + + async def fail_admitted(*_args, **_kwargs): + raise AssertionError("upstream must not start while request is queued") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_openai_passthrough_stream_admitted", fail_admitted) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + ) + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + snapshot = queue.snapshot() + assert snapshot.active == 1 + assert snapshot.queued == 1 + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_admission_timeout_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + url = SimpleNamespace(path = "/v1/chat/completions") + + async def is_disconnected(self): + return False + + async def fail_upstream(*_args, **_kwargs): + raise AssertionError("upstream must not start while request is queued") + + monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0.01") + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming_upstream", + fail_upstream, + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + try: + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + request = Request(), + cancel_event = threading.Event(), + ) + assert exc.value.status_code == 503 + finally: + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_admission_queue_full_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + url = SimpleNamespace(path = "/v1/chat/completions") + + async def is_disconnected(self): + return False + + async def fail_upstream(*_args, **_kwargs): + raise AssertionError("upstream must not start when admission queue is full") + + monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "1") + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming_upstream", + fail_upstream, + ) + + queue = get_llama_admission_queue("http://llama.test") + blocker = queue.reserve( + capacity = 1, + config = LlamaAdmissionConfig(max_queue = 1), + ).lease_nowait() + queued = queue.reserve(capacity = 1, config = LlamaAdmissionConfig(max_queue = 1)) + assert blocker is not None + assert queued.lease_nowait() is None + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + try: + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + request = Request(), + cancel_event = threading.Event(), + ) + assert exc.value.status_code == 429 + finally: + queued.cancel() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_immediate_cancel_stops_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fail_upstream(*_args, **_kwargs): + raise AssertionError("upstream must not start after client cancellation") + + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming_upstream", + fail_upstream, + ) + + cancel_event = threading.Event() + cancel_event.set() + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + + with pytest.raises(HTTPException) as exc: + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + cancel_event = cancel_event, + ) + + assert exc.value.status_code == 499 + assert get_llama_admission_queue("http://llama.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_admission_task_cancel_finalizes_monitor(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fake_wait(*_args, **_kwargs): + raise asyncio.CancelledError() + + async def fail_upstream(*_args, **_kwargs): + raise AssertionError("upstream must not start after admission task cancel") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_wait_for_openai_admission_non_streaming", + fake_wait, + ) + monkeypatch.setattr( + inf_mod, + "_openai_passthrough_non_streaming_upstream", + fail_upstream, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + ) + + with pytest.raises(asyncio.CancelledError): + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + effective_parallel_slots = 1, + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + cancel_event = threading.Event(), + ) + + assert get_llama_admission_queue("http://llama.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + def test_passthrough_non_streaming_cancel_finalizes_monitor(self, monkeypatch): async def _run(): import routes.inference as inf_mod @@ -2990,6 +5528,407 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_non_streaming_cancel_closes_blocked_upstream_post(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class HangingCancelableClient: + def __init__(self): + self.started = asyncio.Event() + self.closed = asyncio.Event() + + async def post(self, *_args, **_kwargs): + self.started.set() + await self.closed.wait() + raise httpx.ReadError("client closed") + + async def aclose(self): + self.closed.set() + + class Request: + async def is_disconnected(self): + return False + + client = HangingCancelableClient() + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_cancelable_nonstreaming_client", + lambda: client, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + cancel_event = threading.Event() + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + task = asyncio.create_task( + _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + request = Request(), + cancel_event = cancel_event, + ) + ) + await asyncio.wait_for(client.started.wait(), 0.2) + cancel_event.set() + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 0.5) + + assert client.closed.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_route_registers_cancel_id(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class HangingCancelableClient: + def __init__(self): + self.started = asyncio.Event() + self.closed = asyncio.Event() + + async def post(self, *_args, **_kwargs): + self.started.set() + await self.closed.wait() + raise httpx.ReadError("client closed") + + async def aclose(self): + self.closed.set() + + class Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/chat/completions") + method = "POST" + + async def is_disconnected(self): + return False + + cancel_id = "passthrough-nonstream-cancel-id" + client = HangingCancelableClient() + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_cancelable_nonstreaming_client", lambda: client) + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + monkeypatch.setattr( + inf_mod, + "get_llama_cpp_backend", + lambda: SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + base_url = "http://llama.test", + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + generate_chat_completion = _plain, + ), + ) + + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + cancel_id = cancel_id, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + task = asyncio.create_task( + openai_chat_completions( + payload, + request = Request(), + current_subject = "test", + ) + ) + await asyncio.wait_for(client.started.wait(), 0.2) + assert cancel_id in inf_mod._CANCEL_REGISTRY + assert inf_mod._cancel_by_cancel_id_or_stash(cancel_id) == 1 + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 0.5) + + assert client.closed.is_set() + assert cancel_id not in inf_mod._CANCEL_REGISTRY + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_disconnect_closes_blocked_upstream_post(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class HangingCancelableClient: + def __init__(self): + self.started = asyncio.Event() + self.closed = asyncio.Event() + + async def post(self, *_args, **_kwargs): + self.started.set() + await self.closed.wait() + raise httpx.ReadError("client closed") + + async def aclose(self): + self.closed.set() + + class Request: + def __init__(self): + self.disconnected = False + + async def is_disconnected(self): + return self.disconnected + + client = HangingCancelableClient() + request = Request() + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "_cancelable_nonstreaming_client", + lambda: client, + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + cancel_event = threading.Event() + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + task = asyncio.create_task( + _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + request = request, + cancel_event = cancel_event, + ) + ) + await asyncio.wait_for(client.started.wait(), 0.2) + request.disconnected = True + + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 0.5) + + assert client.closed.is_set() + assert cancel_event.is_set() + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_non_streaming_forwards_backend_auth_headers(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + captured = {} + + class FakeNonStreamingClient: + async def post(self, *_args, **kwargs): + captured["headers"] = kwargs.get("headers") + return httpx.Response( + 200, + json = { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 123, + "model": "gguf", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "OK"}, + "finish_reason": "stop", + } + ], + }, + ) + + monitor = ApiMonitor(max_entries = 3) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "nonstreaming_client", + lambda: FakeNonStreamingClient(), + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + response = await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _auth_headers = {"Authorization": "Bearer secret"}, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + ) + + assert json.loads(response.body)["choices"][0]["message"]["content"] == "OK" + assert captured["headers"]["Authorization"] == "Bearer secret" + assert captured["headers"]["Connection"] == "close" + + asyncio.run(_run()) + + def test_passthrough_non_streaming_forces_upstream_stream_false(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + captured = {} + + class FakeNonStreamingClient: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + async def post(self, *_args, **kwargs): + captured["json"] = kwargs.get("json") + return httpx.Response( + 200, + json = { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 123, + "model": "gguf", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "OK"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 1, + "completion_tokens": 1, + "total_tokens": 2, + }, + }, + ) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr( + inf_mod, + "nonstreaming_client", + lambda: FakeNonStreamingClient(), + ) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + stream_options = {"include_usage": True}, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + await _openai_passthrough_non_streaming( + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + monitor_id = monitor_id, + ) + + assert captured["json"]["stream"] is False + assert "stream_options" not in captured["json"] + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + + asyncio.run(_run()) + def test_passthrough_clean_eof_finalizes_monitor(self, monkeypatch): async def _run(): result = await self._run_passthrough_stream( @@ -3009,6 +5948,216 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_finish_without_done_closes_stream_early(self, monkeypatch): + # Some llama-server builds emit the finish chunk and then hold the HTTP + # stream open without sending [DONE]; the terminal classifier must end + # the client stream promptly instead of hanging on the open socket. + async def _run(): + import routes.inference as inf_mod + + class Request: + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + yield 'data: {"choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":null}]}' + yield 'data: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}' + await asyncio.Event().wait() # upstream never closes + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + + async def _consume(): + return [chunk async for chunk in response.body_iterator] + + chunks = await asyncio.wait_for(_consume(), timeout = 2) + body = "".join(chunks) + + assert '"finish_reason":"stop"' in body.replace(" ", "") + assert body.endswith("data: [DONE]\n\n") + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stall_after_finish_closes_cleanly(self, monkeypatch): + # include_usage keeps the stream open past the finish chunk waiting for + # the usage chunk; if that never arrives, the post-terminal grace path + # must close with a clean [DONE], not an in-band error. + async def _run(): + import routes.inference as inf_mod + + class Request: + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + yield 'data: {"choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":null}]}' + yield 'data: {"choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}' + raise httpx.ReadTimeout("usage chunk never arrived") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + stream_options = {"include_usage": True}, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + chunks = [chunk async for chunk in response.body_iterator] + body = "".join(chunks) + + assert '"type":"api_error"' not in body.replace(" ", "") + assert body.endswith("data: [DONE]\n\n") + [entry] = monitor.snapshot() + assert entry["status"] == "completed" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_stall_after_data_emits_error(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + yield 'data: {"choices":[{"delta":{"content":"hello"}}]}' + raise httpx.ReadTimeout("upstream went silent") + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + chunks = [chunk async for chunk in response.body_iterator] + body = "".join(chunks) + + assert 'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n' in body + assert '"finish_reason"' not in body.replace(" ", "") + assert '"type":"api_error"' in body.replace(" ", "") + assert "still processing the prompt" in body + assert body.endswith("data: [DONE]\n\n") + [entry] = monitor.snapshot() + assert entry["status"] == "error" + assert "still processing the prompt" in entry["error"] + assert entry["reply"] == "hello" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + class TestApiMonitorSafetensorsUsage: class _Request: @@ -3512,6 +6661,15 @@ class TestApiMonitorAudioInput: class TestResponsesChatTemplateKwargs: _messages = [ChatMessage(role = "user", content = "What is 100 - 67?")] + class _Request: + app = SimpleNamespace(state = SimpleNamespace(llama_parallel_slots = 1)) + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/responses") + method = "POST" + + async def is_disconnected(self): + return False + def test_enable_thinking_lifted_from_extra_body(self): payload = ResponsesRequest( model = "qwen-local", @@ -3544,6 +6702,113 @@ class TestResponsesChatTemplateKwargs: chat_req = _build_chat_request(payload, self._messages, stream = False) assert chat_req.enable_thinking is None + def test_responses_stream_queued_request_sends_keepalive_before_upstream(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fail_send(*_args, **_kwargs): + raise AssertionError("responses upstream must not start while queued") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + base_url = "http://llama.responses.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.01") + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fail_send) + + queue = get_llama_admission_queue("http://llama.responses.test") + blocker = queue.reserve(capacity = 1, config = LlamaAdmissionConfig()).lease_nowait() + assert blocker is not None + monitor_id = monitor.start( + endpoint = "/v1/responses", + method = "POST", + model = "qwen-local", + prompt = "hi", + ) + payload = ResponsesRequest(model = "qwen-local", input = "hi", stream = True) + + response = await _responses_stream( + payload, + [ChatMessage(role = "user", content = "hi")], + self._Request(), + monitor_id, + ) + iterator = response.body_iterator + try: + chunk = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert chunk == ": keep-alive\n\n" + snapshot = queue.snapshot() + assert snapshot.active == 1 + assert snapshot.queued == 1 + finally: + aclose = getattr(iterator, "aclose", None) + if aclose is not None: + await aclose() + blocker.release() + + snapshot = queue.snapshot() + assert snapshot.active == 0 + assert snapshot.queued == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_responses_stream_cancel_after_created_finalizes_monitor_and_slot(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + async def fail_send(*_args, **_kwargs): + raise AssertionError("responses upstream must not start after created cancel") + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + base_url = "http://llama.responses.test", + context_length = 4096, + effective_parallel_slots = 1, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fail_send) + monitor_id = monitor.start( + endpoint = "/v1/responses", + method = "POST", + model = "qwen-local", + prompt = "hi", + ) + payload = ResponsesRequest(model = "qwen-local", input = "hi", stream = True) + + response = await _responses_stream( + payload, + [ChatMessage(role = "user", content = "hi")], + self._Request(), + monitor_id, + ) + iterator = response.body_iterator + first = await asyncio.wait_for(iterator.__anext__(), timeout = 0.2) + assert "event: response.created" in first + + with pytest.raises(asyncio.CancelledError): + await iterator.athrow(asyncio.CancelledError()) + + assert get_llama_admission_queue("http://llama.responses.test").snapshot().active == 0 + [entry] = monitor.snapshot() + assert entry["status"] == "cancelled" + assert monitor.active_count() == 0 + + asyncio.run(_run()) + # ===================================================================== # GGUF chat-template role alternation: coalesce orphaned user turns left diff --git a/studio/backend/tests/test_orchestrator_unload_cancel.py b/studio/backend/tests/test_orchestrator_unload_cancel.py new file mode 100644 index 0000000000..fe3c6d5a0d --- /dev/null +++ b/studio/backend/tests/test_orchestrator_unload_cancel.py @@ -0,0 +1,1622 @@ +# SPDX-License-Identifier: AGPL-3.0-only +"""unload_model cancels an in-flight generation instead of waiting it out. + +The sequential subprocess used to queue ``unload`` behind a running ``generate``, +hanging the UI. ``unload_model`` now cancels first (the mp.Event the worker checks +each token) and takes ``_gen_lock`` before the unload round-trip. +""" + +import threading +import time + +import pytest + +from core.inference import orchestrator as orch_mod +from core.inference.orchestrator import InferenceOrchestrator + + +def _bare_orchestrator(): + """An orchestrator without the real __init__ subprocess/network.""" + o = InferenceOrchestrator.__new__(InferenceOrchestrator) + o._gen_lock = threading.Lock() + o._cancel_event = threading.Event() # stands in for the mp.Event + o._drain_event = threading.Event() # stands in for the unload-drain mp.Event + o._proc = object() # truthy so _ensure_subprocess_alive reports alive + o._cmd_queue = object() + o._resp_queue = object() + o._dispatcher_thread = None + o._dispatcher_stop = threading.Event() + o._dispatcher_lifecycle_lock = threading.Lock() + o._unload_pending = False + o.active_model_name = "m" + o.models = {"m": {}} + o.loading_models = set() + return o + + +def test_unload_cancels_inflight_generation_then_unloads(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + sent = [] + monkeypatch.setattr(o, "_send_cmd", lambda cmd: sent.append(cmd)) + monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + + # A generation holds _gen_lock and releases it only once cancelled. + o._gen_lock.acquire() + + def releaser(): + o._cancel_event.wait(timeout = 5) # released only after the cancel fires + o._gen_lock.release() + + t = threading.Thread(target = releaser) + t.start() + + start = time.monotonic() + ok = o.unload_model("m") + elapsed = time.monotonic() - start + t.join(timeout = 5) + + assert ok is True + assert o._cancel_event.is_set(), "generation must be cancelled before the unload" + assert {"type": "unload", "model_name": "m"} in sent + assert o.active_model_name is None + assert "m" not in o.models + # Waited on the released-after-cancel lock, not a full generation. + assert elapsed < 2.0 + + +def test_unload_no_active_generation_unloads_normally(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + sent = [] + monkeypatch.setattr(o, "_send_cmd", lambda cmd: sent.append(cmd)) + monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + + ok = o.unload_model("m") + + assert ok is True + assert {"type": "unload", "model_name": "m"} in sent + assert o.active_model_name is None + # Lock released for the next caller. + assert o._gen_lock.acquire(blocking = False) + o._gen_lock.release() + + +def test_unload_falls_back_to_shutdown_when_generation_wont_yield(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(orch_mod, "_UNLOAD_GEN_LOCK_TIMEOUT", 0.2) + shutdown = [] + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr(o, "_send_cmd", lambda cmd: pytest.fail("must not send unload when wedged")) + + # A wedged worker never releases _gen_lock, even after the cancel. + o._gen_lock.acquire() + + ok = o.unload_model("m") + + assert ok is True + assert shutdown, "should tear the subprocess down to free the GPU" + assert o.active_model_name is None + + +def test_unload_tears_down_when_compare_dispatcher_wedged(monkeypatch): + # A wedged compare-mode generation bypasses _gen_lock, so the acquire guard + # misses it and _send_cmd/_wait_response would hang on resp_queue. Unload must + # instead tear the subprocess down, like the wedged locked-generation path. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(orch_mod, "_DISPATCH_IDLE_TIMEOUT", 0.2) + + # A live dispatcher whose mailbox never drains == a wedged compare-mode gen. + o._mailbox_lock = threading.Lock() + o._mailboxes = {"req-1": object()} + + class _AliveThread: + def is_alive(self): + return True + + o._dispatcher_thread = _AliveThread() + + shutdown = [] + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send unload with a wedged dispatcher") + ) + monkeypatch.setattr( + o, + "_wait_response", + lambda t, timeout = 300.0: pytest.fail( + "must not wait on resp_queue with a wedged dispatcher" + ), + ) + + # _gen_lock is free (compare mode never took it), so the acquire guard passes. + ok = o.unload_model("m") + + assert ok is True + assert shutdown, "should tear the subprocess down to free the GPU" + assert o.active_model_name is None + assert "m" not in o.models + + +def test_consume_token_stream_bails_when_subprocess_swapped(monkeypatch): + # After a wedged-worker teardown a fresh load swaps _proc/_resp_queue; the + # still-live generation thread must detect the swap and bail, not re-block on + # the new queue while holding _gen_lock. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_subprocess_crash_message", lambda ctx: "inference subprocess restarted" + ) + + def read_one(timeout): + o._proc = object() # simulate the reload swapping the subprocess + return None + + gen = o._consume_token_stream(read_one, lambda: None, crash_context = "generation") + msg = next(gen) + + assert "restarted" in msg + with pytest.raises(StopIteration): + next(gen) + + +def test_unload_pending_clears_after_unload(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_send_cmd", lambda cmd: None) + monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + + o.unload_model("m") + + # The flag must not leak past the unload, else every later generation bails. + assert o._unload_pending is False + + +def test_generation_bails_when_unload_pending(monkeypatch): + # Winning the _gen_lock handoff mid-switch must not start on the outgoing model. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + o._unload_pending = True + + out = list(o._generate_inner(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + # It released (or never held) the lock, so the pending unload can proceed. + assert o._gen_lock.acquire(blocking = False) + o._gen_lock.release() + + +def test_dispatched_generation_bails_when_unload_pending(monkeypatch): + # Compare-mode bypasses _gen_lock, so it must early-out on a pending switch or + # it enqueues a generate on the outgoing model and delays the unload. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_start_dispatcher", lambda: pytest.fail("must not start a generation mid-switch") + ) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate mid-switch") + ) + o._unload_pending = True + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + + +def test_audio_input_generation_bails_when_unload_pending(monkeypatch): + # The audio path takes _gen_lock but must also skip the outgoing model mid-switch. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate mid-switch") + ) + o._unload_pending = True + + out = list(o._generate_audio_input_inner(audio_array = [0.0, 0.1])) + + assert any("unloaded" in chunk.lower() for chunk in out) + # Lock released so the pending unload can proceed. + assert o._gen_lock.acquire(blocking = False) + o._gen_lock.release() + + +def test_audio_response_bails_when_unload_pending(monkeypatch): + # TTS (generate_audio_response) is blocking, so it RAISES rather than starting on the + # outgoing model mid-switch; it takes _gen_lock and must release it either way. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send audio generate mid-switch") + ) + o._unload_pending = True + + with pytest.raises(RuntimeError, match = "unload"): + o.generate_audio_response("hello") + + # Lock released so the pending unload can proceed. + assert o._gen_lock.acquire(blocking = False) + o._gen_lock.release() + + +# ---------------------------------------------------------------------------- +# Preserve unload cancels across the queue handoff (drain_event) — items #1/#4. +# ---------------------------------------------------------------------------- + + +def test_worker_drain_skip_emits_cancelled_gen_done_when_draining(): + # The worker clears cancel_event at the start of every generate, so a cancel set + # while a generate is still queued would be lost when it is dequeued. drain_event + # is the durable signal: while it is set the worker skips the generate (emitting an + # immediate gen_done so the stream/mailbox drains) instead of running it. + import queue as _queue + + from core.inference.worker import _drain_skip_generate + + drain = threading.Event() + rq: _queue.Queue = _queue.Queue() + cmd = {"type": "generate", "request_id": "r1"} + + # Not draining -> run normally (do not skip, emit nothing). + assert _drain_skip_generate(cmd, rq, drain) is False + assert rq.empty() + # Missing event (older worker) -> also runs normally. + assert _drain_skip_generate(cmd, rq, None) is False + assert rq.empty() + + # Draining -> skip and emit a cancelled gen_done for this request_id. + drain.set() + assert _drain_skip_generate(cmd, rq, drain) is True + resp = rq.get_nowait() + assert resp["type"] == "gen_done" + assert resp["request_id"] == "r1" + assert resp["cancelled"] is True + + +def test_worker_generate_branches_check_drain_before_clearing_cancel(): + # Both worker command loops (MLX fast-path + GPU) must consult the drain skip + # before clearing cancel_event and running, so a queued generate can't clear an + # unload-initiated cancel and run the outgoing model to completion. Each loop + # checks the drain twice -- once before the clear and once after -- so a + # drain+cancel pair that lands in the window between them is still caught. + import inspect + + from core.inference import worker + + src = inspect.getsource(worker.run_inference_process) + assert src.count("_drain_skip_generate(cmd, resp_queue, drain_event)") == 4 + + +def test_worker_generate_rechecks_drain_after_clearing_cancel(): + # The exact interleaving item #3 describes: the drain check reads unset, then the + # parent sets drain+cancel for an unload, then the worker clears cancel_event + # (erasing that cancel). A second drain check *after* the clear catches it and + # skips the generate instead of running the outgoing model to completion. + import queue as _queue + + from core.inference.worker import _drain_skip_generate + + drain = threading.Event() + cancel = threading.Event() + rq: _queue.Queue = _queue.Queue() + cmd = {"type": "generate", "request_id": "r1"} + + # 1. Pre-clear drain check: not draining yet -> run (no skip, no emit). + assert _drain_skip_generate(cmd, rq, drain) is False + assert rq.empty() + + # 2. Parent starts an unload: sets drain, then cancel (orchestrator order). + drain.set() + cancel.set() + + # 3. Worker clears cancel at the start of the generate -- erasing the cancel. + cancel.clear() + assert not cancel.is_set() + + # 4. Post-clear drain re-check catches the erased cancel and skips. + assert _drain_skip_generate(cmd, rq, drain) is True + resp = rq.get_nowait() + assert resp["type"] == "gen_done" and resp["cancelled"] is True + + +def test_unload_sets_drain_event_during_switch_and_clears_after(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + + seen = {} + + def record_send(cmd): + # drain_event must be set for the whole unload round-trip so any generate the + # worker dequeues in this window is skipped, not run. + seen["drain_set"] = o._drain_event.is_set() + + monkeypatch.setattr(o, "_send_cmd", record_send) + + assert o.unload_model("m") is True + assert seen.get("drain_set") is True + # Cleared on exit so a later generation (e.g. unloading a non-active model, or a + # reused subprocess) is not wrongly skipped. + assert o._drain_event.is_set() is False + + +def test_unload_clears_drain_event_even_on_wedged_teardown(monkeypatch): + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(orch_mod, "_UNLOAD_GEN_LOCK_TIMEOUT", 0.2) + monkeypatch.setattr(o, "_send_cmd", lambda cmd: pytest.fail("must not send when wedged")) + + # A wedged worker never releases _gen_lock; unload tears the subprocess down. The + # real teardown nulls _drain_event, so emulate that so the finally exercises its guard. + def fake_shutdown(timeout = 5): + o._drain_event = None + + monkeypatch.setattr(o, "_shutdown_subprocess", fake_shutdown) + o._gen_lock.acquire() + + assert o.unload_model("m") is True # must not raise in the drain_event clear + + +# ---------------------------------------------------------------------------- +# Recheck the active model after the lock wait — items #2/#3. +# ---------------------------------------------------------------------------- + + +def test_generation_rechecks_model_after_lock_wait(monkeypatch): + # A request passes the pre-lock active-model check, then blocks on _gen_lock while + # an unload clears/swaps the model. Even if _unload_pending was already reset (the + # unload's finally runs after the lock release), the under-lock active-model recheck + # must make it bail instead of sending a generate to the wrong/unloaded backend. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not generate on a swapped/unloaded model") + ) + + reached_lock = threading.Event() + # _wait_dispatcher_idle runs after the pre-lock check and before acquiring the lock; + # signalling here means the generator captured the model and is about to block. + monkeypatch.setattr(o, "_wait_dispatcher_idle", lambda: (reached_lock.set(), True)[1]) + + o.active_model_name = "m" + o._unload_pending = False + o._gen_lock.acquire() # stand in for an in-flight unload holding the lock + + out: list = [] + + def run(): + out.extend(o._generate_inner(messages = [{"role": "user", "content": "hi"}])) + + t = threading.Thread(target = run) + t.start() + assert reached_lock.wait(timeout = 5) + # Unload finished: model swapped, pending already cleared. Release the lock. + o.active_model_name = "other" + o._gen_lock.release() + t.join(timeout = 5) + + assert out and any("unloaded" in chunk.lower() for chunk in out) + + +def test_generation_rechecks_model_when_unloaded_to_none(monkeypatch): + # Same race, but the unload left no active model (a plain unload, not a switch). + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not generate after the model was unloaded") + ) + reached_lock = threading.Event() + monkeypatch.setattr(o, "_wait_dispatcher_idle", lambda: (reached_lock.set(), True)[1]) + + o.active_model_name = "m" + o._unload_pending = False + o._gen_lock.acquire() + + out: list = [] + t = threading.Thread( + target = lambda: out.extend(o._generate_inner(messages = [{"role": "user", "content": "hi"}])) + ) + t.start() + assert reached_lock.wait(timeout = 5) + o.active_model_name = None + o._gen_lock.release() + t.join(timeout = 5) + + assert out and any("unloaded" in chunk.lower() for chunk in out) + + +# ---------------------------------------------------------------------------- +# Don't unload a stale model name (worker's active-model fallback) — item #5. +# ---------------------------------------------------------------------------- + + +def test_unload_of_stale_name_does_not_touch_active_model(monkeypatch): + # If the named model isn't loaded (e.g. a concurrent load already swapped in a + # different one), unload must not send a command the worker would satisfy by + # unloading its *active* model. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send an unload for a stale model name") + ) + o.active_model_name = "current" + o.models = {"current": {}} + + assert o.unload_model("stale") is True + # The active model is left intact. + assert o.active_model_name == "current" + assert "current" in o.models + + +def test_unload_matches_active_model_case_insensitively(monkeypatch): + # active_model_name can differ in case from the raw model_path a client sends + # to /unload (the load path canonicalizes casing). The stale-name guard must + # match case-insensitively too; otherwise it no-ops the unload and leaves the + # model resident while reporting success. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + sent = [] + monkeypatch.setattr(o, "_send_cmd", lambda cmd: sent.append(cmd)) + monkeypatch.setattr(o, "_wait_response", lambda t, timeout = 300.0: {"type": "unloaded"}) + monkeypatch.setattr(o, "_drain_queue", lambda: []) + + o.active_model_name = "unsloth/Qwen3-4B" + o.models = {"unsloth/Qwen3-4B": {}} + + # Client unloads with the casing it originally typed, before canonicalization. + assert o.unload_model("unsloth/qwen3-4b") is True + # The guard did not no-op: an unload for the canonical active model reached + # the worker (not the raw lowercase name, so the worker matches it directly). + assert {"type": "unload", "model_name": "unsloth/Qwen3-4B"} in sent + # Local state is cleared for the canonical name, not left stale. + assert o.active_model_name is None + assert o.models == {} + + +def test_unload_of_stale_name_still_no_ops_after_case_insensitive_match(monkeypatch): + # The case-insensitive match must only rescue the active model; a genuinely + # different model name (case-insensitively too) must still no-op so the + # worker's absent-name fallback can't tear down the active model. + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send an unload for a stale model name") + ) + o.active_model_name = "unsloth/Qwen3-4B" + o.models = {"unsloth/Qwen3-4B": {}} + + assert o.unload_model("unsloth/Llama-3.1-8B") is True + assert o.active_model_name == "unsloth/Qwen3-4B" + assert "unsloth/Qwen3-4B" in o.models + + +def test_load_does_not_accumulate_stale_models_defeating_the_unload_guard(monkeypatch): + # A load always spawns a fresh subprocess holding only the new model, so + # self.models must mirror that instead of accumulating the previous model's name. + # Otherwise switching A -> B leaves 'A' in self.models, so a later unload('A') + # passes the "not in self.models" guard and the worker's absent-name fallback + # unloads the *active* model B. + import types + + from utils import transformers_version as _tv + + o = _bare_orchestrator() + o.active_model_name = None + o.models = {} + + monkeypatch.setattr(_tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda *a, **k: ([], {})) + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_shutdown_subprocess", lambda *a, **k: None) + monkeypatch.setattr(o, "_spawn_subprocess", lambda cfg: None) + monkeypatch.setattr(orch_mod.time, "sleep", lambda *_a, **_k: None) + + def _load(name): + monkeypatch.setattr( + o, + "_wait_response", + lambda expected, timeout = 300.0: { + "type": "loaded", + "success": True, + "model_info": {"identifier": name, "display_name": name}, + }, + ) + assert o.load_model(types.SimpleNamespace(identifier = name, gguf_variant = None)) is True + + _load("modelA") + _load("modelB") # switch to B without unloading A first + + # self.models mirrors the single live model; the swapped-out name is gone. + assert o.active_model_name == "modelB" + assert set(o.models) == {"modelB"} + + # A stale unload of the swapped-out model must not reach the worker (whose + # absent-name fallback would unload the active model B). + monkeypatch.setattr(o, "_send_cmd", lambda cmd: pytest.fail("stale unload reached the worker")) + assert o.unload_model("modelA") is True + assert o.active_model_name == "modelB" + assert "modelB" in o.models + + +def test_unload_route_serializes_with_loads_via_lifecycle_gate(monkeypatch): + # Item #5: /unload must hold the same lifecycle gate as /load so a concurrent load + # can't swap the backend subprocess/queues mid-unload. + import asyncio + + import routes.inference as inference_route + from core.inference import llama_keepwarm as kw + from models.inference import UnloadRequest + + class _Llama: + is_active = False + is_loaded = False + model_identifier = None + + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: _Llama()) + monkeypatch.setattr(inference_route, "is_registered_native_path_label", lambda *a: False) + + unloaded: list = [] + + class _Backend: + active_model_name = "m" + models = {"m": {}} + + def unload_model(self, name): + unloaded.append(name) + return True + + monkeypatch.setattr(inference_route, "get_inference_backend", lambda: _Backend()) + + async def scenario(): + # Hold the real gate, exactly as an in-flight /load would. + assert kw._lifecycle_lock.acquire(blocking = False) + try: + task = asyncio.ensure_future( + inference_route.unload_model(UnloadRequest(model_path = "m"), "tester") + ) + # Yield to the loop repeatedly: the route must stay blocked on the gate. + for _ in range(10): + await asyncio.sleep(0.01) + assert unloaded == [], "unload ran while the lifecycle gate was held" + assert not task.done() + finally: + kw._lifecycle_lock.release() + resp = await task + assert resp.status == "unloaded" + assert unloaded == ["m"] + + asyncio.run(scenario()) + + +# ---------------------------------------------------------------------------- +# Cancel an in-flight load OFF the lifecycle gate (Stop-loading regression). +# /load holds the gate for the whole load, so a gated /unload could never +# interrupt it; cancel_load only tears the loading subprocess down. +# ---------------------------------------------------------------------------- + + +def test_cancel_load_terminates_loading_subprocess_and_sends_no_command(monkeypatch): + o = _bare_orchestrator() + o.loading_models = {"m"} + o.active_model_name = None + o.models = {} + shutdown = [] + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("cancel_load must not send a worker command") + ) + + assert o.cancel_load("m") is True + assert shutdown, "must tear the loading subprocess down" + assert "m" not in o.loading_models + assert o.active_model_name is None + # A name that is not loading -> no-op, returns False so the caller takes the gate. + assert o.cancel_load("other") is False + + +def test_cancel_load_matches_loading_model_case_insensitively(monkeypatch): + o = _bare_orchestrator() + o.loading_models = {"unsloth/Qwen3-4B"} + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: None) + + assert o.cancel_load("unsloth/qwen3-4b") is True + assert o.loading_models == set() + + +def test_unload_model_cancels_a_loading_model_via_cancel_load(monkeypatch): + # unload_model still cancels an in-flight load (shared logic with cancel_load). + o = _bare_orchestrator() + o.loading_models = {"m"} + o.active_model_name = None + shutdown = [] + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send a command to cancel a load") + ) + + assert o.unload_model("m") is True + assert shutdown + assert "m" not in o.loading_models + + +def test_unload_route_cancels_in_flight_load_without_waiting_on_gate(monkeypatch): + # The regression: /unload wrapped its whole body in the lifecycle gate, so the + # Stop-loading button (cancelLoading -> /unload) could not interrupt a safetensors + # load that holds the gate for its full duration. The cancel must run off-gate. + import asyncio + + import routes.inference as inference_route + from core.inference import llama_keepwarm as kw + from models.inference import UnloadRequest + + class _Llama: + is_active = False + is_loaded = False + model_identifier = None + + monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: _Llama()) + monkeypatch.setattr(inference_route, "is_registered_native_path_label", lambda *a: False) + + cancelled: list = [] + + class _Backend: + active_model_name = None + models: dict = {} + + def get_loading_model(self): + return "m" + + def cancel_load(self, name): + cancelled.append(name) + return True + + def unload_model(self, name): + pytest.fail("must not take the gated unload path for a still-loading model") + + monkeypatch.setattr(inference_route, "get_inference_backend", lambda: _Backend()) + + async def scenario(): + # Hold the real gate, exactly as an in-flight /load would. + assert kw._lifecycle_lock.acquire(blocking = False) + try: + # Even with the gate held, the loading-cancel must go through. + resp = await inference_route.unload_model(UnloadRequest(model_path = "m"), "tester") + assert resp.status == "unloaded" + assert cancelled == ["m"] + finally: + kw._lifecycle_lock.release() + + asyncio.run(scenario()) + + +# ---------------------------------------------------------------------------- +# A dispatched (compare-mode) request that races an unload must not orphan its +# mailbox after _wait_dispatcher_idle stops the dispatcher. +# ---------------------------------------------------------------------------- + + +def test_dispatched_bails_when_unload_flips_before_mailbox_registration(monkeypatch): + # The request passes the pre-work _unload_pending check, then an unload sets + # _unload_pending and _wait_dispatcher_idle stops the dispatcher (mailboxes empty) + # before this request registers its mailbox. The recheck under _mailbox_lock must + # make it bail, or the worker's skipped-generate reply has nothing to route it and + # the compare stream hangs on an orphaned mailbox. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_start_dispatcher", lambda: None) + + # Flip the unload flag after the pre-work check (626) but before mailbox + # registration -- exactly the window _wait_dispatcher_idle exploits. + def flip(*a, **k): + o._unload_pending = True + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", flip) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate after the unload flipped") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + assert o._mailboxes == {}, "must not leave an orphaned mailbox" + + +# ---------------------------------------------------------------------------- +# Dispatched path: bail when a cleared-pending unload swapped the model or +# tore the dispatcher down during the pre-registration window -- item #2. +# ---------------------------------------------------------------------------- + + +class _AliveDispatcher: + """Stand-in dispatcher thread that reports itself alive.""" + + def is_alive(self): + return True + + +def test_dispatched_bails_when_model_swapped_before_mailbox_registration(monkeypatch): + # The request passes the pre-work checks, then a full unload+reload completes + # (clearing _unload_pending) before this request registers its mailbox. The + # under-lock recheck must notice active_model_name changed and bail, instead of + # sending a generate that lands on the swapped-in model. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = _AliveDispatcher() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_start_dispatcher", lambda: None) + + # Swap the active model after the pre-work check but before registration, + # with _unload_pending already back to False (the unload finally ran). + def swap(*a, **k): + o.active_model_name = "other" + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", swap) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not generate on the swapped-in model") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + assert o._mailboxes == {}, "must not leave an orphaned mailbox" + + +def test_dispatched_bails_when_dispatcher_stopped_before_mailbox_registration(monkeypatch): + # Same window, but the unload was a same-model reload so active_model_name is + # unchanged; the give-away is that the dispatcher was stopped. Registering a + # mailbox with no dispatcher to route the reply would hang the compare stream. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = _AliveDispatcher() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_start_dispatcher", lambda: None) + + def stop_dispatcher(*a, **k): + o._dispatcher_thread = None # unload's _stop_dispatcher cleared it + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", stop_dispatcher) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not generate with the dispatcher stopped") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + assert o._mailboxes == {}, "must not leave an orphaned mailbox" + + +def test_dispatched_happy_path_registers_and_sends(monkeypatch): + # Guard against a false bail: with the model unchanged and the dispatcher alive, + # the recheck must let the generate through (register a mailbox and send). + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = _AliveDispatcher() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_start_dispatcher", lambda: None) + monkeypatch.setattr( + o, "_build_generate_cmd", lambda *a, **k: {"type": "generate", "request_id": "r1"} + ) + sent = [] + monkeypatch.setattr(o, "_send_cmd", lambda cmd: sent.append(cmd)) + + # Feed one gen_done so the consumer returns promptly. + def fake_consume(read_mailbox, drainer, **k): + mbox = o._mailboxes.get("r1") + if mbox is not None: + mbox.put({"type": "gen_done", "request_id": "r1"}) + yield "" + + monkeypatch.setattr(o, "_consume_token_stream", fake_consume) + + list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert sent, "happy path must send the generate command" + assert o._mailboxes == {}, "mailbox popped in finally" + + +# ---------------------------------------------------------------------------- +# load_model observes a cancel that discarded its loading marker -- item #4. +# ---------------------------------------------------------------------------- + + +def test_load_model_aborts_when_cancelled_before_spawn(monkeypatch): + # Stop-loading during GPU placement discards the loading marker (cancel_load) with + # no child yet to kill. load_model must observe the removal and not spawn a worker + # that loads the model after /unload already reported it unloaded. + o = _bare_orchestrator() + o.active_model_name = None + o.models = {} + o.loading_models = set() + o._proc = None + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) + monkeypatch.setattr(o, "_shutdown_subprocess", lambda *a, **k: None) + monkeypatch.setattr( + o, "_spawn_subprocess", lambda cfg: pytest.fail("must not spawn a worker after a cancel") + ) + + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "needs_transformers_5", lambda name: False) + + # cancel_load discards the marker while we resolve GPU placement. + def cancel_during_gpu(gpu_ids, **k): + o.loading_models.discard("m") + return ([0], "sel") + + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", cancel_during_gpu) + + class _Cfg: + identifier = "m" + + ok = o.load_model(_Cfg()) + + assert ok is False + assert o.active_model_name is None + assert o.models == {} + + +def test_load_model_proceeds_when_not_cancelled(monkeypatch): + # Guard against a false abort: an uncancelled load keeps its marker and spawns. + o = _bare_orchestrator() + o.active_model_name = None + o.models = {} + o.loading_models = set() + o._proc = None + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) + monkeypatch.setattr(o, "_shutdown_subprocess", lambda *a, **k: None) + + spawned = [] + monkeypatch.setattr(o, "_spawn_subprocess", lambda cfg: spawned.append(cfg)) + monkeypatch.setattr( + o, + "_wait_response", + lambda t, timeout = 300.0: {"success": True, "model_info": {"identifier": "m"}}, + ) + + import utils.transformers_version as tv + + monkeypatch.setattr(tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda gpu_ids, **k: ([0], "sel")) + + class _Cfg: + identifier = "m" + + ok = o.load_model(_Cfg()) + + assert ok is True + assert spawned, "uncancelled load must spawn a worker" + assert o.active_model_name == "m" + + +def test_load_model_aborts_when_cancelled_during_spawn(monkeypatch): + # Stop-loading can land AFTER the pre-spawn marker recheck but while + # _spawn_subprocess is still creating the queues/process, so cancel_load's + # _shutdown_subprocess finds _proc not yet alive and no-ops. load_model must + # recheck the marker once the child exists and tear the orphaned worker down, + # instead of waiting for "loaded" and publishing a model /unload already + # reported as unloaded (a live subprocess nothing later reaps). + import types + + from utils import transformers_version as tv + + o = _bare_orchestrator() + o.active_model_name = None + o.models = {} + o.loading_models = {"m"} + o._proc = None + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) + monkeypatch.setattr(tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda gpu_ids, **k: ([0], "sel")) + + # The cancel lands during the spawn window: cancel_load already discarded the + # marker, but its teardown no-oped because _proc was not alive yet. + def spawn_then_cancel(cfg): + o.loading_models.discard("m") + + monkeypatch.setattr(o, "_spawn_subprocess", spawn_then_cancel) + + shutdown = [] + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: shutdown.append(timeout)) + monkeypatch.setattr( + o, + "_wait_response", + lambda t, timeout = 300.0: pytest.fail( + "must not wait for 'loaded' after a cancel during spawn" + ), + ) + + ok = o.load_model(types.SimpleNamespace(identifier = "m", gguf_variant = None)) + + assert ok is False + assert shutdown, "must tear the orphaned worker down" + assert o.active_model_name is None + assert o.models == {} + assert "m" not in o.loading_models + + +# ---------------------------------------------------------------------------- +# /unload cancels a still-loading GGUF off the lifecycle gate -- item #1. +# ---------------------------------------------------------------------------- + + +def test_unload_cancels_loading_gguf_off_gate(monkeypatch): + # A still-loading GGUF (is_active, not is_loaded) must be cancelled off the gate: + # /load holds the lifecycle gate for the whole load, so a gated unload would wait + # it out. Assert the gate is never entered and unload_model() runs. + import asyncio as _asyncio + + import routes.inference as ri + from core.inference import llama_keepwarm + + gate_entered = {"v": False} + + class _Gate: + async def __aenter__(self): + gate_entered["v"] = True + return self + + async def __aexit__(self, *a): + return False + + class _LlamaBackend: + is_active = True + is_loaded = False + model_identifier = "gguf-model" + + def __init__(self): + self.unloaded = False + + def unload_model(self): + self.unloaded = True + + llama = _LlamaBackend() + + class _Unsloth: + def get_loading_model(self): + return None # no Unsloth load in flight -> Unsloth fast path skipped + + monkeypatch.setattr(ri, "get_llama_cpp_backend", lambda: llama) + monkeypatch.setattr(ri, "get_inference_backend", lambda: _Unsloth()) + monkeypatch.setattr(llama_keepwarm, "inference_lifecycle_gate", lambda: _Gate()) + monkeypatch.setattr(llama_keepwarm, "note_model_unloaded", lambda: None) + + req = ri.UnloadRequest(model_path = "gguf-model") + resp = _asyncio.run(ri.unload_model(req, current_subject = "s")) + + assert getattr(resp, "status", None) == "unloaded" + assert llama.unloaded is True, "must cancel the loading GGUF via unload_model()" + assert gate_entered["v"] is False, "must handle the loading GGUF off the lifecycle gate" + + +def test_unload_loaded_gguf_still_uses_gate(monkeypatch): + # Guard: an already-loaded GGUF (is_loaded True) is NOT caught by the off-gate + # fast path; it goes through the gate as before. + import asyncio as _asyncio + + import routes.inference as ri + from core.inference import llama_keepwarm + + gate_entered = {"v": False} + + class _Gate: + async def __aenter__(self): + gate_entered["v"] = True + return self + + async def __aexit__(self, *a): + return False + + class _LlamaBackend: + is_active = True + is_loaded = True + model_identifier = "gguf-model" + + def __init__(self): + self.unloaded = False + + def unload_model(self): + self.unloaded = True + + llama = _LlamaBackend() + + class _Unsloth: + def get_loading_model(self): + return None + + monkeypatch.setattr(ri, "get_llama_cpp_backend", lambda: llama) + monkeypatch.setattr(ri, "get_inference_backend", lambda: _Unsloth()) + monkeypatch.setattr(ri, "is_registered_native_path_label", lambda a, b: False) + monkeypatch.setattr(llama_keepwarm, "inference_lifecycle_gate", lambda: _Gate()) + monkeypatch.setattr(llama_keepwarm, "note_model_unloaded", lambda: None) + + req = ri.UnloadRequest(model_path = "gguf-model") + resp = _asyncio.run(ri.unload_model(req, current_subject = "s")) + + assert getattr(resp, "status", None) == "unloaded" + assert llama.unloaded is True + assert gate_entered["v"] is True, "loaded GGUF unload must still take the gate" + + +def test_unload_of_mismatched_loading_gguf_skips_off_gate_fast_path(monkeypatch): + # A still-loading GGUF X (is_active, not is_loaded) must NOT be torn down by the + # off-gate fast path when /unload names a DIFFERENT model Y. The single llama-server + # can only load one GGUF at a time, so this fast path is "stop loading THIS model"; + # without a target check it fires for any in-flight GGUF and would abort an unrelated + # load (e.g. a second tab unloading Y kills the load of X). A mismatched target must + # fall through to the lifecycle gate (where, in production, it waits out X's /load and + # then no-ops) instead of taking the off-gate teardown. + import asyncio as _asyncio + + import routes.inference as ri + from core.inference import llama_keepwarm + + gate_entered = {"v": False} + + class _Gate: + async def __aenter__(self): + gate_entered["v"] = True + return self + + async def __aexit__(self, *a): + return False + + class _LlamaBackend: + is_active = True + is_loaded = False + model_identifier = "gguf-X" + + def __init__(self): + self.unloaded = False + + def unload_model(self): + self.unloaded = True + + llama = _LlamaBackend() + + class _Unsloth: + def get_loading_model(self): + return None # no Unsloth load in flight -> Unsloth fast path skipped + + monkeypatch.setattr(ri, "get_llama_cpp_backend", lambda: llama) + monkeypatch.setattr(ri, "get_inference_backend", lambda: _Unsloth()) + monkeypatch.setattr(ri, "is_registered_native_path_label", lambda a, b: False) + monkeypatch.setattr(llama_keepwarm, "inference_lifecycle_gate", lambda: _Gate()) + monkeypatch.setattr(llama_keepwarm, "note_model_unloaded", lambda: None) + + req = ri.UnloadRequest(model_path = "gguf-Y") # different from the loading model X + _asyncio.run(ri.unload_model(req, current_subject = "s")) + + assert gate_entered["v"] is True, ( + "a mismatched-target unload must not use the off-gate GGUF fast path; " + "it would cancel the wrong in-flight load" + ) + + +# ---------------------------------------------------------------------------- +# cancel_load clears its loading marker BEFORE tearing the subprocess down, so a +# racing off-gate load_model observes the cancel during the shutdown window. +# ---------------------------------------------------------------------------- + + +def test_cancel_load_clears_marker_before_shutdown(monkeypatch): + # cancel_load runs off the lifecycle gate, concurrently with a load_model that + # rechecks the loading marker before each spawn to observe the cancel. + # _shutdown_subprocess can block (tearing a live child down / joining the compare + # dispatcher), so discarding the marker only AFTER it leaves a long window in which + # that load_model reads the marker still set, passes its pre-spawn recheck, and + # spawns + loads the model after /unload already reported it cancelled. The marker + # (and local state) must be cleared before the teardown. + o = _bare_orchestrator() + o.loading_models = {"m"} + o.active_model_name = "m" + o.models = {"m": {}} + + at_shutdown = {} + + def record_shutdown(timeout = 5): + at_shutdown["marker_present"] = "m" in o.loading_models + at_shutdown["active"] = o.active_model_name + at_shutdown["models"] = dict(o.models) + + monkeypatch.setattr(o, "_shutdown_subprocess", record_shutdown) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("cancel_load must not send a worker command") + ) + + assert o.cancel_load("m") is True + assert at_shutdown.get("marker_present") is False, ( + "the loading marker must be cleared before _shutdown_subprocess so a concurrent " + "load_model pre-spawn recheck observes the cancel during the shutdown window" + ) + assert at_shutdown.get("active") is None + assert at_shutdown.get("models") == {} + assert "m" not in o.loading_models + assert o.active_model_name is None + assert o.models == {} + + +def test_cancel_load_reclears_state_when_racing_load_repopulates_during_teardown(monkeypatch): + # cancel_load (off the lifecycle gate) can race a load_model whose worker already + # queued its successful "loaded" reply. cancel_load discards the loading marker and + # clears the local mirrors, then tears the subprocess down; but the still-running + # load_model thread can consume that "loaded" DURING the teardown window and repopulate + # active_model_name/models. _shutdown_subprocess nulls the queues but never touches those + # mirrors, so without a second clear /unload reports success while the backend keeps + # advertising a model whose worker was just killed. cancel_load must re-clear after the + # teardown so no phantom loaded model survives. + import types + + from utils import transformers_version as _tv + + o = _bare_orchestrator() + o.loading_models = {"m"} + o.active_model_name = None + o.models = {} + o._proc = None # no prior subprocess -> load_model goes straight to the spawn loop + + monkeypatch.setattr(_tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda *a, **k: ([], {})) + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) + monkeypatch.setattr(o, "_spawn_subprocess", lambda cfg: None) + + parked = threading.Event() # load_model is parked in _wait_response("loaded") + release_loaded = threading.Event() # cancel_load lets the load consume "loaded" + load_done = threading.Event() + + def blocking_wait_response(expected, timeout = 300.0): + parked.set() + assert release_loaded.wait(timeout = 5) + return { + "type": "loaded", + "success": True, + "model_info": {"identifier": "m", "display_name": "m"}, + } + + monkeypatch.setattr(o, "_wait_response", blocking_wait_response) + + load_result: dict = {} + + def run_load(): + try: + load_result["ok"] = o.load_model( + types.SimpleNamespace(identifier = "m", gguf_variant = None) + ) + except Exception as exc: # noqa: BLE001 + load_result["exc"] = exc + finally: + load_done.set() + + loader = threading.Thread(target = run_load) + loader.start() + assert parked.wait(timeout = 5), "load_model must reach _wait_response" + + # The teardown IS the window in which the racing load repopulates the mirrors: the + # marker is already discarded here, so release the load and wait for it to finish + # repopulating, mirroring the 0.5s cancel-settle inside the real _shutdown_subprocess. + def racing_shutdown(timeout = 0.5): + release_loaded.set() + assert load_done.wait(timeout = 5), "the racing load must repopulate during teardown" + + monkeypatch.setattr(o, "_shutdown_subprocess", racing_shutdown) + + assert o.cancel_load("m") is True + loader.join(timeout = 5) + + # Fail-without: load_model set active_model_name/models during racing_shutdown and + # cancel_load left them set, so the backend advertises a model whose worker was killed. + assert o.active_model_name is None, "cancel_load must not leave a repopulated active model" + assert o.models == {}, "cancel_load must not leave a repopulated models mirror" + assert "m" not in o.loading_models + + +# ---------------------------------------------------------------------------- +# A dispatched (compare-mode) request that starts the dispatcher and then bails on +# a racing unload must stop the dispatcher it started, or that orphaned dispatcher +# steals the worker's "unloaded" reply and hangs unload_model on its 300s timeout. +# ---------------------------------------------------------------------------- + + +def test_dispatched_bail_stops_orphan_dispatcher_it_started(monkeypatch): + # The request passes the pre-work _unload_pending check and starts the dispatcher + # (none was running), then an unload sets _unload_pending so the under-lock recheck + # bails. The just-started dispatcher, left running with no mailboxes, competes with + # unload_model()'s _wait_response for the worker's "unloaded" reply off the shared + # resp_queue and drops it as unroutable, hanging the unload until its 300s timeout. + # The bail must stop the dispatcher it started. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = None # none running -> this call starts it + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + + started = {"v": False} + stopped = {"v": False} + + def fake_start(): + started["v"] = True + o._dispatcher_thread = _AliveDispatcher() + return True # _start_dispatcher returns True for the caller that spawned it + + def fake_stop(): + stopped["v"] = True + o._dispatcher_thread = None + + monkeypatch.setattr(o, "_start_dispatcher", fake_start) + monkeypatch.setattr(o, "_stop_dispatcher", fake_stop) + + # An unload flips _unload_pending after the pre-work check but before registration. + def flip(*a, **k): + o._unload_pending = True + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", flip) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate after the unload flipped") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + assert started["v"], "this call started the dispatcher" + assert stopped["v"], "the bail must stop the dispatcher it started (no other mailboxes)" + assert o._mailboxes == {} + + +def test_dispatched_bail_keeps_dispatcher_with_other_active_mailbox(monkeypatch): + # Guard against over-stopping: if another compare request registered a mailbox on the + # dispatcher this call started, the bail must NOT stop it, or that request's token + # routing dies mid-stream. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = None + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_start_dispatcher", lambda: setattr(o, "_dispatcher_thread", _AliveDispatcher()) + ) + monkeypatch.setattr( + o, + "_stop_dispatcher", + lambda: pytest.fail("must not stop a dispatcher another compare request is using"), + ) + + # A concurrent compare request registers its mailbox, then an unload flips the flag. + def flip(*a, **k): + o._mailboxes["other"] = object() + o._unload_pending = True + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", flip) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate after the unload flipped") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + assert set(o._mailboxes) == {"other"}, "the other request's mailbox is untouched" + + +def test_dispatched_bail_keeps_preexisting_dispatcher(monkeypatch): + # Guard: if the dispatcher was already running before this request (an earlier compare + # request started it), a bail must not stop it even with no mailboxes now -- this + # request did not start it and another may re-use it. Only the call that starts an + # otherwise-idle dispatcher during the race is responsible for stopping it. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._unload_pending = False + o._dispatcher_thread = _AliveDispatcher() # already running + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_start_dispatcher", lambda: None) + monkeypatch.setattr( + o, "_stop_dispatcher", lambda: pytest.fail("must not stop a pre-existing dispatcher") + ) + + def flip(*a, **k): + o._unload_pending = True + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", flip) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate after the unload flipped") + ) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert any("unloaded" in chunk.lower() for chunk in out) + + +# ---------------------------------------------------------------------------- +# load_model rechecks the loading marker AFTER _wait_response("loaded") and +# BEFORE publishing -- item #6. cancel_load's post-teardown re-clear only wipes a +# repopulation that lands during its shutdown; a publish that lands after +# cancel_load returns survives it, so the recheck must abort the publish itself. +# ---------------------------------------------------------------------------- + + +def test_load_model_aborts_publish_when_cancelled_after_wait_response(monkeypatch): + # cancel_load (off the lifecycle gate) discards the loading marker BEFORE its teardown + # and re-clears the mirrors AFTER it. A racing load_model can consume its worker's + # already-queued "loaded" reply and reach the publish block only AFTER cancel_load has + # fully returned -- so cancel_load's post-teardown re-clear cannot undo that publish. + # Without a marker recheck between _wait_response("loaded") and the publish, load_model + # advertises active_model_name/models for a model /unload already reported cancelled, + # over a subprocess cancel_load just killed. The recheck must observe the discarded + # marker and abort the publish. + import types + + from utils import transformers_version as _tv + + o = _bare_orchestrator() + o.loading_models = {"m"} + o.active_model_name = None + o.models = {} + o._proc = None # no prior subprocess -> load_model goes straight to the spawn loop + + monkeypatch.setattr(_tv, "needs_transformers_5", lambda name: False) + monkeypatch.setattr(orch_mod, "prepare_gpu_selection", lambda *a, **k: ([], {})) + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: False) + monkeypatch.setattr(o, "_spawn_subprocess", lambda cfg: None) + # cancel_load tears the worker down; a no-op keeps the test off real subprocesses. + monkeypatch.setattr(o, "_shutdown_subprocess", lambda timeout = 5: None) + + parked = threading.Event() # load_model reached _wait_response("loaded") + cancel_done = threading.Event() # cancel_load fully returned (marker discarded + re-clear) + load_done = threading.Event() + + def blocking_wait_response(expected, timeout = 300.0): + parked.set() + # Do not consume "loaded" until cancel_load has fully returned, so the publish + # would land AFTER cancel_load's post-teardown re-clear -- the window the + # re-clear alone cannot cover. + assert cancel_done.wait(timeout = 5) + return { + "type": "loaded", + "success": True, + "model_info": {"identifier": "m", "display_name": "m"}, + } + + monkeypatch.setattr(o, "_wait_response", blocking_wait_response) + + load_result: dict = {} + + def run_load(): + try: + load_result["ok"] = o.load_model( + types.SimpleNamespace(identifier = "m", gguf_variant = None) + ) + except Exception as exc: # noqa: BLE001 + load_result["exc"] = exc + finally: + load_done.set() + + loader = threading.Thread(target = run_load) + loader.start() + assert parked.wait(timeout = 5), "load_model must reach _wait_response" + + # cancel_load runs to completion while the load is parked: it discards the marker and + # re-clears the mirrors (post-teardown), then returns. Only then let the load consume + # "loaded" and attempt to publish. + assert o.cancel_load("m") is True + cancel_done.set() + + loader.join(timeout = 5) + assert load_done.is_set() + + # Fail-without: load_model published active_model_name/models for 'm' AFTER cancel_load + # returned, advertising a cancelled model over a killed subprocess. + assert load_result.get("ok") is False, "the cancelled load must not report success" + assert o.active_model_name is None, "must not publish a cancelled model's active name" + assert o.models == {}, "must not publish a cancelled model's mirror" + assert "m" not in o.loading_models + + +# ---------------------------------------------------------------------------- +# Concurrent compare-mode requests must not each spawn a dispatcher. Compare mode +# (_generate_dispatched) deliberately bypasses _gen_lock, so two requests can reach +# _start_dispatcher at once. Without _dispatcher_lifecycle_lock the check-then-spawn +# races: both observe no live dispatcher and each start one. The extra dispatcher is +# orphaned (self._dispatcher_thread tracks only the last) and later consumes the +# "unloaded" reply off the shared resp_queue before unload_model's _wait_response, +# hanging the unload on its 300s timeout. The lifecycle lock must serialize the +# check-then-spawn so exactly one dispatcher thread is ever created. +# ---------------------------------------------------------------------------- + + +def test_concurrent_start_dispatcher_spawns_exactly_one(): + import queue as _queue + + o = _bare_orchestrator() + o._resp_queue = _queue.Queue() # real queue so the dispatcher loop blocks and stays alive + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._dispatcher_thread = None + o._dispatcher_stop = threading.Event() + o._dispatcher_lifecycle_lock = threading.Lock() + + n = 32 + # A barrier aligns every thread on the check-then-spawn window: without the lifecycle + # lock several would clear the "is a dispatcher alive?" check together and each spawn one. + barrier = threading.Barrier(n) + results: list = [] + results_lock = threading.Lock() + + def racer(): + barrier.wait() + started = o._start_dispatcher() + with results_lock: + results.append(started) + + threads = [threading.Thread(target = racer, name = f"racer-{i}") for i in range(n)] + for t in threads: + t.start() + for t in threads: + t.join(timeout = 5) + + try: + # _start_dispatcher returns True only for the caller that actually spawned a thread. + # Exactly one caller may win; every other must observe the dispatcher alive and bail. + assert results.count(True) == 1, f"expected exactly one spawn, got {results.count(True)}" + assert results.count(False) == n - 1 + # And exactly one live dispatcher thread exists -- no orphan racing resp_queue. + live = [ + t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive() + ] + assert len(live) == 1, f"expected one live dispatcher, found {len(live)}" + assert o._dispatcher_thread is live[0] + finally: + o._stop_dispatcher() + + # Stop joins and clears it; no dispatcher thread must survive. + assert o._dispatcher_thread is None + remaining = [ + t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive() + ] + assert remaining == [], "dispatcher must be stopped and joined" + + +# ---------------------------------------------------------------------------- +# A compare request whose _start_dispatcher is queued behind an unload's +# _stop_dispatcher must NOT spawn a fresh dispatcher. The idle-dispatcher stop +# and the queued start both serialize on _dispatcher_lifecycle_lock; if the +# queued start spawned a new dispatcher after the stop, it would become the +# resp_queue reader and consume unload_model's "unloaded" reply (unroutable, so +# dropped) before _wait_response saw it -- hanging the unload on its 300s +# timeout. unload_model sets _unload_pending under the SAME lifecycle lock ahead +# of the stop, so _start_dispatcher observes it and refuses. +# ---------------------------------------------------------------------------- + + +def test_start_dispatcher_refuses_while_unload_pending(): + # Direct unit guard: with an unload in progress (_unload_pending set under the + # lifecycle lock by unload_model), _start_dispatcher must refuse and spawn nothing, + # even though no dispatcher is currently running. + import queue as _queue + + o = _bare_orchestrator() + o._resp_queue = _queue.Queue() # a spawned dispatcher would block-read here and stay alive + o._dispatcher_thread = None + o._dispatcher_stop = threading.Event() + o._dispatcher_lifecycle_lock = threading.Lock() + o._unload_pending = True + + started = o._start_dispatcher() + + assert started is False, "must not start a dispatcher while an unload is pending" + assert o._dispatcher_thread is None, "no dispatcher thread may be created" + live = [t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive()] + assert live == [], "no dispatcher may exist to consume the unloaded reply" + + +def test_start_dispatcher_resumes_after_unload_clears(): + # Guard the other direction: once the unload finishes and clears _unload_pending, a + # later compare request must be able to start the dispatcher again (the gate must not + # wedge). Proves the refusal above is scoped to the unload, not permanent. + import queue as _queue + + o = _bare_orchestrator() + o._resp_queue = _queue.Queue() + o._dispatcher_thread = None + o._dispatcher_stop = threading.Event() + o._dispatcher_lifecycle_lock = threading.Lock() + o._unload_pending = False + + try: + assert ( + o._start_dispatcher() is True + ), "a fresh dispatcher must start once no unload is pending" + assert o._dispatcher_thread is not None and o._dispatcher_thread.is_alive() + finally: + o._stop_dispatcher() + + assert o._dispatcher_thread is None + + +def test_queued_start_behind_unload_stop_spawns_no_dispatcher(): + # Codex's exact ordering, forced deterministically: an unload holds + # _dispatcher_lifecycle_lock across its _stop_dispatcher (the idle dispatcher's join + # is gated by an event), while a compare request's _start_dispatcher is queued behind + # it on the same lock. When the stop releases the lock the queued start must observe + # _unload_pending (set under the lock ahead of the stop) and refuse: no fresh + # dispatcher may be left running to steal the "unloaded" reply. + import queue as _queue + + o = _bare_orchestrator() + o._resp_queue = _queue.Queue() # a spawned dispatcher would block-read here and stay alive + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._dispatcher_stop = threading.Event() + o._dispatcher_lifecycle_lock = threading.Lock() + o._unload_pending = False + + start_queued = threading.Event() # release the stop's join once the start is queued behind it + join_may_finish = threading.Event() + + class _IdleDispatcher: + # Stand-in for the idle compare-mode dispatcher the unload stops. Its join blocks + # until we confirm the compare _start_dispatcher is queued behind the stop, so the + # stop provably holds _dispatcher_lifecycle_lock across that window. + def is_alive(self): + return True + + def join(self, timeout = None): + assert start_queued.wait(timeout = 5), "compare start must queue behind the stop" + assert join_may_finish.wait(timeout = 5) + + o._dispatcher_thread = _IdleDispatcher() + + def unload_side(): + # unload_model's sequence: set _unload_pending under the lifecycle lock, then stop + # the idle dispatcher (also under the lock, via _wait_dispatcher_idle). + with o._dispatcher_lifecycle_lock: + o._unload_pending = True + o._stop_dispatcher() + + started_result = {} + + def compare_side(): + started_result["v"] = o._start_dispatcher() + + u = threading.Thread(target = unload_side, name = "unload-side") + u.start() + # Let the unload set _unload_pending, enter _stop_dispatcher, and block in the gated join + # while holding the lifecycle lock. + time.sleep(0.2) + + c = threading.Thread(target = compare_side, name = "compare-side") + c.start() + # Let the compare _start_dispatcher block on the lifecycle lock (queued behind the stop). + time.sleep(0.2) + + start_queued.set() # the start is now queued behind the stop + join_may_finish.set() # let the stop's join complete and release the lock + + u.join(timeout = 5) + c.join(timeout = 5) + + assert started_result.get("v") is False, "the queued start must refuse while unloading" + assert o._dispatcher_thread is None, "the stop cleared it and the queued start spawned nothing" + live = [t for t in threading.enumerate() if t.name == "inference-dispatcher" and t.is_alive()] + assert live == [], "no fresh dispatcher may be left to consume the unloaded reply" diff --git a/studio/backend/tests/test_passthrough_healing.py b/studio/backend/tests/test_passthrough_healing.py index 06316e2243..9b22ab8b05 100644 --- a/studio/backend/tests/test_passthrough_healing.py +++ b/studio/backend/tests/test_passthrough_healing.py @@ -280,6 +280,56 @@ class TestStreamHealer: assert [c["id"] for c in calls] == ["call_0", "call_1"] assert _events_text(events).strip() == "then" + def test_mistral_array_multiple_calls_all_promoted_in_stream(self): + # A canonical Mistral [TOOL_CALLS] array carries several calls under a + # SINGLE signal. Draining only the first call would leave the residue + # starting at ",{...}]" (no signal), so later calls in the same array + # must be promoted in the same pass, not flushed as raw text. + healer = StreamToolCallHealer({"get_weather", "get_time"}) + array = ( + '[TOOL_CALLS][{"name":"get_weather","arguments":{"city":"Paris"}},' + '{"name":"get_time","arguments":{"tz":"UTC"}}]' + ) + events = healer.feed(array) + healer.finalize() + calls = _events_calls(events) + assert [c["function"]["name"] for c in calls] == ["get_weather", "get_time"] + assert [c["id"] for c in calls] == ["call_0", "call_1"] + assert _events_text(events) == "" + + def test_mistral_array_multiple_calls_promoted_char_by_char(self): + healer = StreamToolCallHealer({"get_weather", "get_time"}) + array = ( + '[TOOL_CALLS][{"name":"get_weather","arguments":{"city":"Paris"}},' + '{"name":"get_time","arguments":{"tz":"UTC"}}]' + ) + events = [] + for ch in array: + events += healer.feed(ch) + events += healer.finalize() + calls = _events_calls(events) + assert [c["function"]["name"] for c in calls] == ["get_weather", "get_time"] + assert _events_text(events) == "" + + def test_mistral_array_undeclared_middle_kept_as_text_others_promoted(self): + # A mid-array element for a tool that is not declared must survive as + # text while the declared neighbours on either side still promote in + # document order. + healer = StreamToolCallHealer({"a", "c"}) + array = ( + '[TOOL_CALLS][{"name":"a","arguments":{}},' + '{"name":"b","arguments":{}},{"name":"c","arguments":{}}]' + ) + events = healer.feed(array) + healer.finalize() + assert [c["function"]["name"] for c in _events_calls(events)] == ["a", "c"] + assert '"b"' in _events_text(events) + + def test_mistral_array_then_trailing_prose(self): + healer = StreamToolCallHealer({"a", "b"}) + array = '[TOOL_CALLS][{"name":"a","arguments":{}},{"name":"b","arguments":{}}]' + events = healer.feed(f"{array} all done") + healer.finalize() + assert [c["function"]["name"] for c in _events_calls(events)] == ["a", "b"] + assert "all done" in _events_text(events) + def test_incomplete_call_healed_at_finalize(self): healer = StreamToolCallHealer({"Bash"}) events = healer.feed('{"name":"Bash","arguments":{"cmd":"ls"}}') @@ -465,6 +515,7 @@ class ScriptedClient: _url, json = None, timeout = None, + headers = None, ): self.posts.append(json) return httpx.Response(200, json = self.bodies[min(len(self.posts) - 1, len(self.bodies) - 1)]) @@ -1356,3 +1407,42 @@ class TestOpenaiStreamingRoute: assert chunks[0] == line + "\n\n" # byte-for-byte relay asyncio.run(_run()) + + +class TestHealerSignalAlignment: + """The passthrough healer buffers only formats its parser can promote. + The loops' bare [ARGS] rehearsal signal is gated on active tool names + there; ungated in the healer it would stall legitimate prose until + finalization without ever producing a promotable call.""" + + def test_heal_signals_are_promotable_formats_only(self): + from core.inference.passthrough_healing import _HEAL_SIGNALS + assert set(_HEAL_SIGNALS) == { + "", + "<|tool_call>", + "`` containing a literal ``<`` (e.g. ``if x < 10``). +* Kimi K2 dotted name ``functions.my.tool:0`` keeps its full name + (``my.tool``) after stripping only the ``functions.`` prefix and + ``:idx`` suffix, while the full id is preserved on the call. +* Kimi K2 bare-counter id (no ``functions.`` prefix, no ``:IDX``) is + dropped rather than surfaced under a numeric name. +* DeepSeek V3.1 truncated mid-stream produces an empty result without + raising. +* ``routes.inference._strip_tool_xml`` strips the DeepSeek envelope and + the Kimi section markers added by this PR. +""" + +import json + +import pytest + +from core.inference.tool_call_parser import ( + parse_tool_calls_from_text, + strip_tool_markup, +) + + +# GLM string-vs-JSON-encoded value coercion (finding B in plan) + + +@pytest.mark.parametrize( + "raw_val, expected_python", + [ + # Bare numeric / bool / null shapes are still treated as JSON + # literals (ambiguous with strings; the template doesn't tell us). + ("42", 42), + ("true", True), + ("false", False), + ("null", None), + ("3.14", 3.14), + ("-7", -7), + ("1e3", 1000.0), + ], +) +def test_glm_numeric_and_bool_literals_are_json_decoded(raw_val, expected_python): + text = ( + "n\n" + f"v\n" + f"{raw_val}\n" + "" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["v"] == expected_python + + +@pytest.mark.parametrize( + "raw_val", + [ + "hello world", # plain prose + "True", # Python literal, NOT JSON -- no longer eaten by ast.literal_eval + "None", # Python literal, NOT JSON -- no longer eaten by ast.literal_eval + "if x < 10: pass", # code with literal < (well, < not in arg_value here) + "{not valid json", # looks like an object but is malformed -- must stay raw + "[oops", # looks like an array but is malformed + ], +) +def test_glm_non_json_shapes_stay_raw(raw_val): + text = ( + "n\n" + f"v\n" + f"{raw_val}\n" + "" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["v"] == raw_val + assert isinstance(args["v"], str) + + +def test_glm_json_object_arg_decoded(): + text = ( + "nest\n" + "opts\n" + '{"limit": 10}\n' + "" + ) + calls = parse_tool_calls_from_text(text) + args = json.loads(calls[0]["function"]["arguments"]) + assert args["opts"] == {"limit": 10} + + +def test_glm_json_array_arg_decoded(): + text = ( + "nest\n" + "ids\n" + "[1, 2, 3]\n" + "" + ) + calls = parse_tool_calls_from_text(text) + args = json.loads(calls[0]["function"]["arguments"]) + assert args["ids"] == [1, 2, 3] + + +def test_glm_arg_value_with_literal_less_than(): + text = ( + "run\n" + "code\n" + "if x < 10: pass\n" + "" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "if x < 10: pass" + + +# GLM 4.7 no-newline emission shape + + +def test_glm_4_7_no_newlines_between_name_and_arg_key(): + """GLM 4.7 strips the ``\\n`` after the name (``{{- ... -}}`` in the + template) so ```` follows directly. Parser must accept both.""" + text = ( + "get_weather" + "cityLondon" + "unitscelsius" + "" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_weather" + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"city": "London", "units": "celsius"} + + +def test_glm_4_7_no_newlines_multi_call(): + """Back-to-back GLM 4.7 calls without intervening newlines.""" + text = ( + "ax1" + "by2" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 2 + assert calls[0]["function"]["name"] == "a" + assert calls[1]["function"]["name"] == "b" + + +def test_glm_4_7_does_not_break_qwen_path(): + """Qwen ``{json}`` still dispatches to Qwen; GLM's + first-char ``[^\\n<{]`` excludes ``{``.""" + text = '{"name":"web_search","arguments":{"q":"x"}}' + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "web_search" + + +# Kimi K2 dotted name + bare counter (finding C in plan) + + +def test_kimi_dotted_namespace_keeps_full_dotted_name(): + # A dotted Kimi id keeps its FULL name; only the ``functions.`` prefix and ``:idx`` suffix drop (vLLM parity). + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.my.tool:0" + "<|tool_call_argument_begin|>{}" + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "my.tool" + assert calls[0]["id"] == "functions.my.tool:0" + + +def test_kimi_two_sections_in_one_stream_both_parse(): + """Outer loop walks every ``<|tool_calls_section_begin|>...end|>`` + so vLLM / SGLang parity holds even on multi-section streams.""" + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.a:0" + '<|tool_call_argument_begin|>{"x":1}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + " some prose between sections " + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.b:0" + '<|tool_call_argument_begin|>{"y":2}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 2 + assert calls[0]["function"]["name"] == "a" + assert calls[1]["function"]["name"] == "b" + assert calls[0]["id"] == "functions.a:0" + assert calls[1]["id"] == "functions.b:0" + + +def test_kimi_bare_counter_id_is_dropped(): + """Bare-digit id (``3``) is dropped (matches vLLM); SGLang infers + name from schema, which we don't have at parse time.""" + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>3" + "<|tool_call_argument_begin|>{}" + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + calls = parse_tool_calls_from_text(text) + assert calls == [] + + +# DeepSeek truncated mid-stream + + +def test_deepseek_v3_1_huge_truncated_body_is_linear(): + """Adversarial input: DeepSeek envelope with no JSON brace and a + 50k-char body. A regex-based ``[^\\n<]+?`` name capture is O(N^2) + here; the parser uses ``str.find`` on the sep marker so it stays + linear. Budget 1s to flag any future regression.""" + import time as _time + + text = "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>fn<｜tool▁sep｜>" + "x" * 50_000 + start = _time.time() + calls = parse_tool_calls_from_text(text) + elapsed = _time.time() - start + assert elapsed < 1.0, f"V3 path is non-linear: {elapsed:.2f}s" + assert calls == [] + + +def test_deepseek_r1_huge_fenceless_body_is_linear(): + """R1 detection used a greedy ``([^\\n]+)\\n```json`` regex that is O(N^2) on a + fence-less body of repeated ``function`` tokens. The parser now scans with + ``str.find``; budget 1s to flag any regression.""" + import time as _time + + text = "<｜tool▁calls▁begin｜>" + "function<｜tool▁sep｜>a" * 40_000 + start = _time.time() + calls = parse_tool_calls_from_text(text) + elapsed = _time.time() - start + assert elapsed < 1.0, f"R1 path is non-linear: {elapsed:.2f}s" + assert calls == [] + + +def test_glm_unclosed_body_many_arg_keys_is_linear(): + """An unclosed GLM ```` body runs to EOF; a lazy-group ``finditer`` + over many bare ```` tokens was O(N^2). The parser now walks pairs with + ``str.find``; budget 1s.""" + import time as _time + + text = "foo\n" + "k" * 40_000 + start = _time.time() + parse_tool_calls_from_text(text) + elapsed = _time.time() - start + assert elapsed < 1.0, f"GLM path is non-linear: {elapsed:.2f}s" + + +def test_deepseek_r1_fenced_json_parses(): + """R1 wraps args in a ```json fence after ``functionNAME``.""" + import json as _json + + text = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n" + "```json\n" + '{"city":"NYC","unit":"c"}\n' + "```<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_weather" + assert _json.loads(calls[0]["function"]["arguments"]) == {"city": "NYC", "unit": "c"} + + +def test_deepseek_v3_1_truncated_arguments_drops_call_without_crash(): + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city":"Tokyo"' # no closing brace, no end markers + ) + calls = parse_tool_calls_from_text(text) + assert calls == [] + + +def test_deepseek_v3_1_truncated_after_end_marker_still_yields_call(): + text = ( + "<｜tool▁calls▁begin｜>" "<｜tool▁call▁begin｜>get_time" "<｜tool▁sep｜>" '{"city":"Tokyo"}' + # neither <｜tool▁call▁end｜> nor <｜tool▁calls▁end｜> + ) + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_time" + assert json.loads(calls[0]["function"]["arguments"]) == {"city": "Tokyo"} + + +# Routes-layer strip across the three new families + + +def test_routes_layer_strip_removes_deepseek_envelope(): + from routes.inference import _strip_tool_xml as _routes_strip + + text = ( + "before " + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + '<｜tool▁sep｜>{"city":"Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + " after" + ) + stripped = _routes_strip(text) + assert stripped == "before after" + + +def test_routes_layer_strip_removes_kimi_section(): + from routes.inference import _strip_tool_xml as _routes_strip + + text = ( + "before " + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"q":"x"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + " after" + ) + stripped = _routes_strip(text) + assert stripped == "before after" + + +def test_routes_layer_strip_removes_glm_block(): + """``.*?`` covers GLM via the Qwen pattern.""" + from routes.inference import _strip_tool_xml as _routes_strip + + text = ( + "before " + "web_search\n" + "q\nx\n" + "" + " after" + ) + stripped = _routes_strip(text) + assert stripped == "before after" + + +# strip_tool_markup (parser-level finalise path) over the new families + + +def test_strip_tool_markup_handles_deepseek_envelope(): + text = ( + "before " + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + '<｜tool▁sep｜>{"city":"Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + " after" + ) + stripped = strip_tool_markup(text, final = True) + assert "before" in stripped and "after" in stripped + assert "｜tool▁" not in stripped + assert "get_time" not in stripped and "Tokyo" not in stripped + + +def test_strip_tool_markup_handles_kimi_section(): + text = ( + "before " + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"q":"x"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + " after" + ) + stripped = strip_tool_markup(text, final = True) + assert "before" in stripped and "after" in stripped + assert "tool_calls_section_begin" not in stripped + + +# Round-2 review findings: GLM quoted-string / unclosed-arg, DeepSeek +# strict terminator, nested wrapper-less Gemma strip + + +def test_glm_quoted_string_arg_keeps_its_quotes(): + # A GLM string value emitted verbatim that itself begins with a quote. + text = ( + "web_search\n" + "query\n" + '"exact phrase"\n' + "" + ) + calls = parse_tool_calls_from_text(text) + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == '"exact phrase"' + + +def test_glm_unclosed_arg_value_is_rejected_in_strict_mode(): + # Closing present but a value never closes: strict mode must reject + # the whole call rather than execute it with the argument silently dropped. + text = ( + "web_search\n" + "query\n" + "Tokyo weather" # no + "" + ) + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + # With Auto-Heal the partial value is kept, not dropped to a no-arg call. + healed = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(healed) == 1 + args = json.loads(healed[0]["function"]["arguments"]) + assert "Tokyo weather" in args.get("query", "") + + +def test_deepseek_v3_missing_call_terminator_rejected_in_strict_mode(): + # Envelope closes but the per-call <｜tool▁call▁end｜> is absent. Strict mode + # must reject (it is truncated/merged); Auto-Heal still parses it. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + '<｜tool▁sep｜>{"city":"Tokyo"}' + "<｜tool▁calls▁end｜>" # envelope end only, no per-call end + ) + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + healed = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(healed) == 1 + assert healed[0]["function"]["name"] == "get_time" + + +def test_deepseek_v3_with_call_terminator_parses_in_strict_mode(): + # Sanity: a well-formed V3 call (with the per-call end marker) still parses + # under strict mode after the terminator check. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + '<｜tool▁sep｜>{"city":"Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_time" + + +def test_strip_tool_markup_removes_nested_wrapperless_gemma_call(): + # Wrapper-less Gemma call with a NESTED object arg: the balanced helper must strip the whole call, not leave a trailing ``}``. + text = "answer: call:f{loc:{city:NYC},n:3} done" + stripped = strip_tool_markup(text, final = True) + assert "call:f" not in stripped + assert "}" not in stripped + assert "answer:" in stripped and "done" in stripped + + +# Pass-3 review findings: bare-Kimi streaming (non-final) strip symmetry +# and the wrapper-less Gemma route-display strip + + +def test_strip_tool_markup_non_final_removes_bare_kimi_call(): + # A bare ``<|tool_call_begin|>...<|tool_call_end|>`` (no section wrapper): the CLOSED (final=False) strip must remove it too. + text = ( + "before " + "<|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"q":"x"}' + "<|tool_call_end|>" + " after" + ) + stripped = strip_tool_markup(text, final = False) + assert "tool_call_begin" not in stripped + assert "tool_call_end" not in stripped + assert "before" in stripped and "after" in stripped + + +def test_routes_layer_strip_removes_wrapperless_gemma_call(): + # Gemma 4 (skip_special_tokens) emits a wrapper-less ``call:NAME{..}`` with no XML markers. + from routes.inference import _strip_tool_xml as _routes_strip + + text = 'before call:web_search{query:"weather in Sydney"} after' + stripped = _routes_strip(text) + assert "call:web_search" not in stripped + assert "before" in stripped and "after" in stripped + + +def test_deepseek_envelope_end_inside_arg_string_is_not_a_truncation(): + # A DeepSeek V3.1 call whose argument string contains the literal envelope-end token must not be dropped. + content = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>web_search<｜tool▁sep｜>" + '{"query":"what does <｜tool▁calls▁end｜> mean"}' + "<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + calls = parse_tool_calls_from_text(content) + assert len(calls) == 1, calls + assert calls[0]["function"]["name"] == "web_search" + assert json.loads(calls[0]["function"]["arguments"]) == { + "query": "what does <｜tool▁calls▁end｜> mean" + } + + +def test_glm_value_containing_literal_arg_value_close_is_preserved(): + # A GLM string argument may legitimately contain . + content = ( + "runcode" + 'print("")' + ) + calls = parse_tool_calls_from_text(content) + assert len(calls) == 1, calls + assert json.loads(calls[0]["function"]["arguments"]) == {"code": 'print("")'} + + +def test_attribute_form_function_with_embedded_marker_runs_outer_call(): + # is a supported envelope; a DeepSeek/Kimi marker inside one of its + # parameter values is data, not a second call. + content = ( + '' + "The Kimi format is <|tool_call_begin|>functions.delete_all:0" + "<|tool_call_argument_begin|>{}<|tool_call_end|>" + "" + ) + calls = parse_tool_calls_from_text(content) + assert [c["function"]["name"] for c in calls] == ["respond"], calls + + +def test_wrapperless_gemma_call_gated_by_enabled_tools(): + # Once skip_special_tokens removes the <|tool_call> wrapper, call:NAME{...} is + # indistinguishable from prose documenting the Gemma syntax. + prose = "Here is an example of the syntax: call:foo{x:1}. That shows how tools work." + assert parse_tool_calls_from_text(prose, enabled_tool_names = {"web_search"}) == [] + # The display strip is gated the same way, so the example survives in the answer. + assert "call:foo{x:1}" in strip_tool_markup( + prose, final = True, enabled_tool_names = {"web_search"} + ) + # An enabled name is still a real call (parsed, and stripped from display). + real = "Answer. call:web_search{query:hi}" + calls = parse_tool_calls_from_text(real, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert "call:web_search" not in strip_tool_markup( + real, final = True, enabled_tool_names = {"web_search"} + ) + + +def test_kimi_section_end_inside_arg_string_is_not_a_truncation(): + # In a multi-call Kimi section, a later call whose argument holds the literal section-end token must not truncate the section. + content = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.search:0<|tool_call_argument_begin|>" + '{"q":"cats"}<|tool_call_end|>' + "<|tool_call_begin|>functions.explain:1<|tool_call_argument_begin|>" + '{"text":"the token <|tool_calls_section_end|> means end"}<|tool_call_end|>' + "<|tool_calls_section_end|>" + ) + calls = parse_tool_calls_from_text(content) + assert [c["function"]["name"] for c in calls] == ["search", "explain"], calls + assert json.loads(calls[1]["function"]["arguments"]) == { + "text": "the token <|tool_calls_section_end|> means end" + } + + +def test_closed_envelope_before_deepseek_block_owns_turn(): + # Document order is the contract: a CLOSED / call that precedes a + # DeepSeek/Kimi block owns the turn, even when prose frames it as an example. + deepseek = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>search_web\n" + "```json\n" + '{"query":"weather in Paris"}\n' + "```" + "<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + prose = ( + 'A Qwen call looks like {"name":"example_tool","arguments":{}}.\n' + ) + calls = parse_tool_calls_from_text(prose + deepseek) + assert [c["function"]["name"] for c in calls] == ["example_tool"], calls + + kimi = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.lookup:0" + '<|tool_call_argument_begin|>{"id":7}<|tool_call_end|><|tool_calls_section_end|>' + ) + calls_k = parse_tool_calls_from_text("Example: {} and now:\n" + kimi) + assert [c["function"]["name"] for c in calls_k] == ["demo"], calls_k + + +def test_marker_inside_closed_outer_envelope_still_runs_outer_call(): + # The guard must fire when the marker sits INSIDE a closed outer / envelope's arguments: the OUTER call wins. + outer = ( + "what does <｜tool▁calls▁begin｜> mean" + ) + calls = parse_tool_calls_from_text(outer) + # The outer envelope is the real call; the embedded DeepSeek marker must not + # hijack the parse into a spurious tool. + assert [c["function"]["name"] for c in calls] == ["lookup"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "q": "what does <｜tool▁calls▁begin｜> mean" + } + + +def test_truncated_outer_envelope_with_embedded_marker_heals_outer_call(): + # A TRUNCATED outer call embedding a DeepSeek/Kimi marker in its argument still Auto-Heals as the outer call. + trunc = 'x = "<｜tool▁calls▁begin｜>sample"' + calls = parse_tool_calls_from_text(trunc) + assert [c["function"]["name"] for c in calls] == ["python"], calls + + +def test_python_tag_call_with_embedded_marker_runs_outer_call(): + # ``<|python_tag|>`` is Llama-3's tool-call envelope, so a DeepSeek/Kimi example quoted + # in its argument is data: the OUTER python_tag call (``web_search``) must run, not the + # embedded marker (``delete_all``). + kimi = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.delete_all:0" + "<|tool_call_argument_begin|>{}<|tool_call_end|><|tool_calls_section_end|>" + ) + deepseek = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>delete_all<｜tool▁sep｜>{}" + "<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + for embedded in (kimi, deepseek): + builtin = '<|python_tag|>web_search.call(query="explain ' + embedded + '")' + calls = parse_tool_calls_from_text(builtin, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + custom = ( + '<|python_tag|>{"name":"web_search","parameters":' + '{"query":"explain ' + embedded + '"}}' + ) + calls = parse_tool_calls_from_text(custom, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + + # A bare ``<|python_tag|>`` prose mention (no call shape) must NOT be treated as an + # envelope: a real Kimi call after it still parses (the call-shaped lookahead guard). + prose = "The token <|python_tag|> is used. " + kimi + calls = parse_tool_calls_from_text(prose) + assert [c["function"]["name"] for c in calls] == ["delete_all"], calls + + +def test_gemma_wrapperless_quoted_value_with_comma_not_split(): + # A wrapper-less Gemma call whose quoted value contains ``, key:``. + text = 'call:web_search{query:"weather, location: Boston", limit:3}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "query": "weather, location: Boston", + "limit": 3, + } + + +def test_literal_close_tag_in_xml_arg_before_marker_runs_outer_call(): + # A literal ```` inside an outer XML argument (before a marker) is not the envelope close: the span reaches the REAL final close. + text = ( + 'x = " ' + "<|tool_call_begin|>functions.delete_all:0<|tool_call_argument_begin|>{}" + '<|tool_call_end|>"' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["python"], calls + + +def test_literal_tool_call_close_in_qwen_json_before_marker_runs_outer_call(): + # A Qwen/Hermes whose JSON argument holds a literal then a marker must run the OUTER call. + text = ( + '{"name":"search","arguments":{"query":"explain then ' + "<|tool_call_begin|>functions.delete_all:0<|tool_call_argument_begin|>{}" + '<|tool_call_end|>"}}' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["search"], calls + # Back-to-back Qwen calls still parse independently (real-close span must keep the + # negative-lookahead that separates adjacent calls). + bb = ( + '{"name":"a","arguments":{}}' + '{"name":"b","arguments":{}}' + ) + assert [c["function"]["name"] for c in parse_tool_calls_from_text(bb)] == ["a", "b"] + + +def test_r1_heal_keeps_later_call_when_first_omits_close_fence(): + # DeepSeek R1 multi-call where the FIRST call has balanced JSON but omits its close + # fence/terminator, followed by a well-formed second call. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n```json\n" + '{"city":"SF"}\n```' # no <｜tool▁call▁end｜> + "<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_time\n```json\n" + '{"tz":"UTC"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + heal = [c["function"]["name"] for c in parse_tool_calls_from_text(text)] + assert "get_time" in heal, heal + # Strict keeps the later well-formed call; heal must be a superset. + strict = [ + c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = False) + ] + assert set(strict) <= set(heal), (strict, heal) + + +def test_wrapperless_gemma_nested_call_in_arg_is_not_a_second_call(): + # A wrapper-less Gemma call whose quoted argument mentions another enabled tool must not execute that nested name. + text = 'call:web_search{query:"explain call:delete_all{target:files}"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "query": "explain call:delete_all{target:files}" + } + # Two genuinely separate calls still both parse. + two = "call:web_search{query:hi}call:get_time{tz:UTC}" + assert [ + c["function"]["name"] + for c in parse_tool_calls_from_text(two, enabled_tool_names = {"web_search", "get_time"}) + ] == ["web_search", "get_time"] + + +def test_leading_bare_json_call_owns_quoted_gemma_snippet(): + # Document order: a leading Llama-3.2 bare-JSON call with trailing prose owns the turn. + text = ( + '{"name":"lookup","parameters":{"note":"use call:web_search{query:cats} for this"}}\n' + "That is the call I would make." + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"lookup", "web_search"}) + assert [c["function"]["name"] for c in calls] == ["lookup"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "note": "use call:web_search{query:cats} for this" + } + + # Same with the ``;`` inter-call separator: both real calls parse, the + # quoted snippet still does not. + two = ( + '{"name":"lookup","parameters":{"note":"see call:web_search{query:cats}"}};' + '{"name":"lookup","parameters":{"q":"second"}}' + ) + calls_two = parse_tool_calls_from_text(two, enabled_tool_names = {"lookup", "web_search"}) + assert [c["function"]["name"] for c in calls_two] == ["lookup", "lookup"], calls_two + + +def test_leading_gemma_call_still_wins_over_trailing_json_example(): + # Reverse control: a real leading Gemma call followed by a bare-JSON example keeps the Gemma call (bare JSON matches only a LEADING object). + text = 'call:web_search{query:cats} Example JSON: {"name":"demo_tool","parameters":{}}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "demo_tool"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + + # And prose-only enabled Gemma syntax (no leading JSON) still promotes: the + # markerless by-design behaviour is unchanged. + prose = "You can run call:web_search{query:cats} to search." + calls_p = parse_tool_calls_from_text(prose, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls_p] == ["web_search"], calls_p + + +def test_leading_gemma_call_owns_quoted_mistral_trigger(): + # A leading wrapper-less Gemma call whose argument quotes a Mistral trigger must win: the [TOOL_CALLS] literal is data. + text = 'call:web_search{query:"docs say [TOOL_CALLS]delete_all{}"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "query": "docs say [TOOL_CALLS]delete_all{}" + } + + # Reverse control: a real leading Mistral call still parses normally. + real = '[TOOL_CALLS]delete_all{"x":1}' + calls_m = parse_tool_calls_from_text(real, enabled_tool_names = {"web_search", "delete_all"}) + assert [c["function"]["name"] for c in calls_m] == ["delete_all"], calls_m + + # A DISABLED Gemma example quoting the trigger is dropped as prose and a + # real call after it still parses (drop-the-span recursion). + mixed = ( + 'Example: call:demo{note:"see [TOOL_CALLS]delete_all{}"}\n' + '[TOOL_CALLS]web_search{"q":"real"}' + ) + calls_d = parse_tool_calls_from_text(mixed, enabled_tool_names = {"web_search", "delete_all"}) + assert [c["function"]["name"] for c in calls_d] == ["web_search"], calls_d + + +def test_chained_bare_json_owns_kimi_marker_in_later_call(): + # Document order: two ;-chained bare-JSON calls own the turn even when the second's argument quotes a complete Kimi snippet. + kimi = ( + "<|tool_call_begin|>functions.delete_all:0" + "<|tool_call_argument_begin|>{}<|tool_call_end|>" + ) + two = ( + '{"name":"lookup","parameters":{"q":"first"}};' + '{"name":"lookup","parameters":{"note":"' + kimi + '"}}' + ) + calls = parse_tool_calls_from_text(two, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["lookup", "lookup"], calls + + # Reverse control: prose followed by a real Kimi block still parses. + real = "Let me check.\n<|tool_calls_section_begin|>" + kimi + "<|tool_calls_section_end|>" + calls_k = parse_tool_calls_from_text(real, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls_k] == ["delete_all"], calls_k + + # A closed leading Mistral call preceding a trailing Kimi example owns the + # turn too (same closed-call-precedes-marker rule). + mistral = '[TOOL_CALLS]lookup{"q":"first"} then example ' + kimi + calls_m = parse_tool_calls_from_text(mistral, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls_m] == ["lookup"], calls_m + + +def test_nested_gemma_values_keep_commas_and_parens(): + # Nested wrapper-less Gemma mappings/arrays use the top-level delimiter rules, so nested arguments are not split. + calls = parse_tool_calls_from_text( + "call:python{opts:{code:print(1,2),lang:py}}", enabled_tool_names = {"python"} + ) + assert [c["function"]["name"] for c in calls] == ["python"], calls + assert json.loads(calls[0]["function"]["arguments"]) == { + "opts": {"code": "print(1,2)", "lang": "py"} + } + + arr = parse_tool_calls_from_text( + "call:python{opts:[1,2,{a:f(1,2)}]}", enabled_tool_names = {"python"} + ) + assert json.loads(arr[0]["function"]["arguments"]) == {"opts": [1, 2, {"a": "f(1,2)"}]} + + prose_comma = parse_tool_calls_from_text( + "call:python{opts:{note:hello, world}}", enabled_tool_names = {"python"} + ) + assert json.loads(prose_comma[0]["function"]["arguments"]) == {"opts": {"note": "hello, world"}} + + quoted = parse_tool_calls_from_text( + 'call:python{opts:{q:say "a, b" now,n:3}}', enabled_tool_names = {"python"} + ) + assert json.loads(quoted[0]["function"]["arguments"]) == { + "opts": {"q": 'say "a, b" now', "n": 3} + } + + # Controls: nested quoted values and multi-key mappings are unchanged, and + # a truncated nested value still falls back to the raw string. + nested_q = parse_tool_calls_from_text( + 'call:python{loc:{city:"New York"}}', enabled_tool_names = {"python"} + ) + assert json.loads(nested_q[0]["function"]["arguments"]) == {"loc": {"city": "New York"}} + multi = parse_tool_calls_from_text( + "call:python{opts:{a:1,b:2},n:3}", enabled_tool_names = {"python"} + ) + assert json.loads(multi[0]["function"]["arguments"]) == {"opts": {"a": 1, "b": 2}, "n": 3} + trunc = parse_tool_calls_from_text( + "call:python{opts:{code:print(1,2}}", enabled_tool_names = {"python"} + ) + assert json.loads(trunc[0]["function"]["arguments"]) == {"opts": "{code:print(1,2}"} + + +def test_multi_gemma_calls_own_turn_over_signal_in_later_call(): + # Document order: when the first enabled Gemma call closes before the first foreign signal, the leading call still owns the turn. + en = {"get_time", "web_search", "delete_all"} + both = parse_tool_calls_from_text( + 'call:get_time{} call:web_search{query:"docs say [TOOL_CALLS]delete_all{}"}', + enabled_tool_names = en, + ) + assert [c["function"]["name"] for c in both] == ["get_time", "web_search"], both + assert json.loads(both[1]["function"]["arguments"]) == { + "query": "docs say [TOOL_CALLS]delete_all{}" + } + + # XML and Kimi markers in the later call's strings stay data too. + xml = parse_tool_calls_from_text( + 'call:get_time{} call:web_search{query:"see delete_all"}', + enabled_tool_names = en, + ) + assert [c["function"]["name"] for c in xml] == ["get_time", "web_search"], xml + kimi = parse_tool_calls_from_text( + 'call:get_time{} call:web_search{query:"see <|tool_call_begin|>' + 'functions.delete_all:0<|tool_call_argument_begin|>{}<|tool_call_end|>"}', + enabled_tool_names = en, + ) + assert [c["function"]["name"] for c in kimi] == ["get_time", "web_search"], kimi + + # A trailing prose example after the closed leading call defers the same way. + prose = parse_tool_calls_from_text( + "call:get_time{} Example: [TOOL_CALLS]delete_all{}", enabled_tool_names = en + ) + assert [c["function"]["name"] for c in prose] == ["get_time"], prose + + +def test_multi_gemma_ownership_reverse_controls(): + # A real leading Mistral/XML call with a trailing Gemma example keeps the leading call; a signal before every Gemma call keeps normal order. + en = {"get_time", "web_search", "delete_all"} + mistral = parse_tool_calls_from_text( + '[TOOL_CALLS][{"name":"delete_all","arguments":{}}] Example: call:web_search{query:cats}', + enabled_tool_names = en, + ) + assert [c["function"]["name"] for c in mistral] == ["delete_all"], mistral + xml_first = parse_tool_calls_from_text( + '{"name":"delete_all","arguments":{}} call:web_search{query:cats}', + enabled_tool_names = en, + ) + assert [c["function"]["name"] for c in xml_first] == ["delete_all"], xml_first + agnostic = parse_tool_calls_from_text( + 'call:foo{} {"name":"delete_all","arguments":{}}' + ) + assert [c["function"]["name"] for c in agnostic] == ["delete_all"], agnostic + + +def test_disabled_leading_bare_json_does_not_hide_later_marker_call(): + # A leading bare-JSON object with a NOT-enabled name is prose: the real DeepSeek/Kimi call after it still parses. + kimi = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"q":"cats"}<|tool_call_end|><|tool_calls_section_end|>' + ) + calls = parse_tool_calls_from_text( + '{"name":"draft","parameters":{}} ' + kimi, enabled_tool_names = {"web_search"} + ) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"q": "cats"} + + deepseek = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>web_search\n" + '```json\n{"q":"cats"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + calls_ds = parse_tool_calls_from_text( + '{"name":"draft","parameters":{}} ' + deepseek, enabled_tool_names = {"web_search"} + ) + assert [c["function"]["name"] for c in calls_ds] == ["web_search"], calls_ds + + +def test_disabled_leading_bare_json_ownership_controls(): + kimi_delete = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.delete_all:0" + "<|tool_call_argument_begin|>{}<|tool_call_end|><|tool_calls_section_end|>" + ) + # ENABLED leading name still owns the turn (document order, the shipped + # inside-or-after rule). + owns = parse_tool_calls_from_text( + '{"name":"web_search","parameters":{"q":"first"}} ' + kimi_delete, + enabled_tool_names = {"web_search", "delete_all"}, + ) + assert [c["function"]["name"] for c in owns] == ["web_search"], owns + # A marker INSIDE the disabled object's own strings stays data: the span + # is prose, the tail holds no call, so nothing parses. + inside = parse_tool_calls_from_text( + '{"name":"draft","parameters":{"note":"see <|tool_call_begin|>functions.delete_all:0' + '<|tool_call_argument_begin|>{}<|tool_call_end|>"}}\nsome trailing prose', + enabled_tool_names = {"web_search", "delete_all"}, + ) + assert inside == [], inside + # Nameless leading JSON answers keep recursing to the real call. + nameless = parse_tool_calls_from_text( + '{"answer":42} ' + kimi_delete, enabled_tool_names = {"delete_all"} + ) + assert [c["function"]["name"] for c in nameless] == ["delete_all"], nameless + # Name-agnostic path unchanged: the leading object is the call. + agnostic = parse_tool_calls_from_text('{"name":"draft","parameters":{}} ' + kimi_delete) + assert [c["function"]["name"] for c in agnostic] == ["draft"], agnostic + + +def test_leading_json_answer_with_prose_keeps_quoted_gemma_snippet_as_data(): + # A LEADING JSON answer followed by prose is data (same contract as the whole-content JSON exemption). + obj = '{"summary":"use call:web_search{query:cats} to search"}\nHope that helps!' + assert parse_tool_calls_from_text(obj, enabled_tool_names = {"web_search"}) == [] + arr = '["use call:web_search{query:cats} to search"]\nHope that helps!' + assert parse_tool_calls_from_text(arr, enabled_tool_names = {"web_search"}) == [] + assert strip_tool_markup(obj, enabled_tool_names = {"web_search"}) == obj + + # A REAL call in the tail after the answer still parses (and strips). + tail = '{"summary":"done"}\ncall:web_search{query:cats}' + calls = parse_tool_calls_from_text(tail, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + + # A leading brace run that is NOT valid JSON gets no exemption. + not_json = "{not json} call:web_search{query:cats}" + calls_nj = parse_tool_calls_from_text(not_json, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls_nj] == ["web_search"], calls_nj + + +def test_glm_heal_bounds_unclosed_value_at_tool_call_close(): + # Auto-Heal: a value missing only its before the block's heals to the + # value text, not the close tag and everything after it swallowed into the argument. + one = "get_weathercityNYC" + calls = parse_tool_calls_from_text(one, allow_incomplete = True) + assert [c["function"]["name"] for c in calls] == ["get_weather"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"city": "NYC"} + + # Trailing prose after the close stays out of the healed value. + two = one + "\nLet me check that for you." + calls_two = parse_tool_calls_from_text(two, allow_incomplete = True) + assert json.loads(calls_two[0]["function"]["arguments"]) == {"city": "NYC"} + + # Strict mode still rejects the unclosed value outright. + assert parse_tool_calls_from_text(one, allow_incomplete = False) == [] + + # A value truncated at EOF (no structural tag follows) keeps the partial heal, and a proper + # close whose value holds a literal is untouched by the bounding. + eof = "get_weathercityNew York Ci" + calls_eof = parse_tool_calls_from_text(eof, allow_incomplete = True) + assert json.loads(calls_eof[0]["function"]["arguments"]) == {"city": "New York Ci"} + lit = ( + "get_weathercity" + 'print("")' + ) + calls_lit = parse_tool_calls_from_text(lit, allow_incomplete = True) + assert json.loads(calls_lit[0]["function"]["arguments"]) == {"city": 'print("")'} + + +def test_prose_mentioning_ds_kimi_markers_survives_final_strip(): + # False-alarm literals: the trailing strip arms require a call-shaped + # lookahead, so an answer documenting a marker keeps its tail. + from core.inference.tool_call_parser import strip_tool_markup + + for text in [ + "The Kimi marker <|tool_calls_section_begin|> starts a section.", + "DeepSeek uses <｜tool▁calls▁begin｜> to open calls.", + "See <|tool_call_begin|> in the docs.", + ]: + assert strip_tool_markup(text, final = True) == text + + # Truncated REAL calls still drop, and a bare marker at EOF is a fragment. + truncated_kimi = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"q' + ) + assert strip_tool_markup(truncated_kimi, final = True) == "" + assert strip_tool_markup("prefix <|tool_calls_section_begin|>", final = True) == "prefix" diff --git a/studio/backend/tests/test_presence_penalty.py b/studio/backend/tests/test_presence_penalty.py new file mode 100644 index 0000000000..030ddb6011 --- /dev/null +++ b/studio/backend/tests/test_presence_penalty.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: AGPL-3.0-only +"""Presence-penalty parity between the GGUF path and the safetensors/MLX paths. + +The safetensors path historically dropped ``presence_penalty``, so the SAME model +looked worse served as safetensors. These tests pin the processor semantics +(subtract once per distinct completion token, prompt excluded, presence not +frequency, zero a no-op, negatives raise) plus a param-propagation regression +over route -> orchestrator cmd -> worker gen_kwargs. +""" + +import threading + +import pytest +import torch + +from core.inference.presence_penalty import ( + apply_presence_penalty, + _make_presence_penalty_processor, +) + + +def test_seen_token_gets_exactly_minus_penalty_unseen_unchanged(): + input_ids = torch.tensor([[0, 1, 3]]) # prompt [0, 1], completion [3] + scores = torch.zeros(1, 5) + out = apply_presence_penalty(input_ids, scores, penalty = 1.5, prompt_len = 2) + assert out[0, 3].item() == pytest.approx(-1.5) + for tok in (0, 1, 2, 4): + assert out[0, tok].item() == pytest.approx(0.0) + + +def test_multiplicity_ignored_presence_not_frequency(): + # Token 3 emitted three times -> still a single -penalty (presence, not freq). + input_ids = torch.tensor([[0, 3, 3, 3]]) + scores = torch.zeros(1, 5) + out = apply_presence_penalty(input_ids, scores, penalty = 2.0, prompt_len = 1) + assert out[0, 3].item() == pytest.approx(-2.0) + + +def test_negative_penalty_raises_seen_logits(): + input_ids = torch.tensor([[0, 2]]) + scores = torch.zeros(1, 4) + out = apply_presence_penalty(input_ids, scores, penalty = -0.5, prompt_len = 1) + assert out[0, 2].item() == pytest.approx(0.5) + + +def test_prompt_tokens_excluded(): + # Token 7 is prompt-only (untouched); token 4 in the completion is penalized. + input_ids = torch.tensor([[7, 4, 4]]) + scores = torch.zeros(1, 8) + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + assert out[0, 7].item() == pytest.approx(0.0) + assert out[0, 4].item() == pytest.approx(-1.0) + + +def test_batch_rows_isolated(): + input_ids = torch.tensor([[0, 1], [0, 2]]) # row completions [1] and [2] + scores = torch.zeros(2, 4) + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + assert out[0, 1].item() == pytest.approx(-1.0) + assert out[0, 2].item() == pytest.approx(0.0) + assert out[1, 2].item() == pytest.approx(-1.0) + assert out[1, 1].item() == pytest.approx(0.0) + + +def test_zero_penalty_is_noop(): + input_ids = torch.tensor([[0, 1, 2]]) + scores = torch.randn(1, 5) + original = scores.clone() + out = apply_presence_penalty(input_ids, scores, penalty = 0.0, prompt_len = 1) + assert torch.equal(out, original) + + +def test_empty_completion_is_noop(): + # prompt_len covers the whole sequence -> nothing generated yet. + input_ids = torch.tensor([[0, 1, 2]]) + scores = torch.randn(1, 5) + original = scores.clone() + out = apply_presence_penalty(input_ids, scores, penalty = 1.5, prompt_len = 3) + assert torch.equal(out, original) + + +def test_out_of_vocab_id_ignored(): + # A generated id >= vocab_size (defensive) must not index out of bounds. + input_ids = torch.tensor([[0, 9]]) + scores = torch.zeros(1, 5) # vocab 5, token 9 is out of range + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + assert torch.equal(out, torch.zeros(1, 5)) + + +def test_negative_generated_id_ignored(): + # A negative generated id (defensive) must be dropped, not wrap to scores[-1]. + input_ids = torch.tensor([[0, -1]]) + scores = torch.zeros(1, 5) + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + # Nothing penalized; in particular the last row (the numpy/torch wrap target + # for id -1) is untouched. + assert torch.equal(out, torch.zeros(1, 5)) + + +def test_mixed_oob_negative_and_valid_ids_only_in_range_penalized(): + # Completion mixes a valid id (1), an out-of-vocab id (9 >= vocab 5) and a + # negative id (-1). Only the in-range distinct id is penalized; OOB/negative + # ids are ignored with no crash and no wrong-index wrap. This fails under the + # old ``seen[seen < vocab_size]`` filter (id -1 wraps to the last row) and + # passes only with the both-ends bound. + input_ids = torch.tensor([[0, 1, 9, -1, 1]]) # prompt [0], completion [1, 9, -1, 1] + scores = torch.zeros(1, 5) + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + expected = torch.zeros(1, 5) + expected[0, 1] = -1.0 # once per distinct in-range id (multiplicity ignored) + assert torch.equal(out, expected) + assert out[0, 4].item() == pytest.approx(0.0) # id -1 did not wrap to the last row + + +def test_dtype_and_device_preserved(): + input_ids = torch.tensor([[0, 1]]) + scores = torch.zeros(1, 4, dtype = torch.float16) + out = apply_presence_penalty(input_ids, scores, penalty = 1.0, prompt_len = 1) + assert out.dtype == torch.float16 + assert out.device == scores.device + + +def test_processor_none_when_zero(): + assert _make_presence_penalty_processor(0.0, prompt_len = 0) is None + + +def test_processor_applies_penalty(): + proc = _make_presence_penalty_processor(1.5, prompt_len = 2) + assert proc is not None + input_ids = torch.tensor([[0, 1, 3]]) + scores = torch.zeros(1, 5) + out = proc(input_ids, scores) + assert out[0, 3].item() == pytest.approx(-1.5) + + +def test_processor_composes_with_other_processors(): + # LogitsProcessorList must run our processor alongside a pre-existing one. + from transformers import LogitsProcessor, LogitsProcessorList + + class _AddToTokenZero(LogitsProcessor): + def __call__(self, input_ids, scores): + scores[:, 0] = scores[:, 0] + 100.0 + return scores + + presence = _make_presence_penalty_processor(1.0, prompt_len = 1) + combined = LogitsProcessorList([_AddToTokenZero(), *presence]) + input_ids = torch.tensor([[5, 2]]) # completion = [2] + scores = torch.zeros(1, 6) + out = combined(input_ids, scores) + assert out[0, 0].item() == pytest.approx(100.0) # other processor ran + assert out[0, 2].item() == pytest.approx(-1.0) # presence ran + + +def test_mlx_presence_penalty_callable(): + mx = pytest.importorskip("mlx.core", reason = "MLX only ships on arm64 macOS") + from core.inference.mlx_inference import _make_mlx_presence_penalty_processor + + proc = _make_mlx_presence_penalty_processor(1.5) + # First call = prompt only (latches prompt_len, penalizes nothing). + prompt = mx.array([10, 11]) + logits0 = mx.zeros((1, 20)) + out0 = proc(prompt, logits0) + assert float(out0[0, 10]) == pytest.approx(0.0) + # Second call: one completion token (5) appended -> penalized once. + seq = mx.array([10, 11, 5]) + logits1 = mx.zeros((1, 20)) + out1 = proc(seq, logits1) + assert float(out1[0, 5]) == pytest.approx(-1.5) + assert float(out1[0, 10]) == pytest.approx(0.0) # prompt token untouched + + +def test_mlx_presence_penalty_bounds_out_of_range_ids(): + # Documents (and, on Apple Silicon CI, enforces) the intended MLX bound: + # out-of-vocab and negative completion ids must be ignored. MLX does no + # bounds checking and OOB indexing is undefined behavior (crash / memory + # corruption), so the processor routes stray ids to a discarded scratch slot + # and penalizes only in-range distinct ids -- matching the torch filter + # seen[(seen >= 0) & (seen < vocab)]. Skips off arm64 macOS where MLX is absent. + mx = pytest.importorskip("mlx.core", reason = "MLX only ships on arm64 macOS") + from core.inference.mlx_inference import _make_mlx_presence_penalty_processor + + proc = _make_mlx_presence_penalty_processor(1.0) + proc(mx.array([10, 11]), mx.zeros((1, 8))) # first call latches prompt_len = 2 + # Completion appends a valid id (3), an out-of-vocab id (99 >= vocab 8) and a + # negative id (-1); only the in-range id is penalized and nothing crashes. + seq = mx.array([10, 11, 3, 99, -1]) + out = proc(seq, mx.zeros((1, 8))) + assert float(out[0, 3]) == pytest.approx(-1.0) + for tok in range(8): + if tok != 3: + assert float(out[0, tok]) == pytest.approx(0.0) + + +# Param propagation: route payload -> orchestrator cmd -> worker gen_kwargs +_SAMPLING = { + "temperature": 0.7, + "top_p": 0.8, + "top_k": 20, + "min_p": 0.05, + "repetition_penalty": 1.1, + "presence_penalty": 1.5, +} + + +def test_orchestrator_cmd_carries_all_sampling_params(): + from core.inference.orchestrator import InferenceOrchestrator + + o = InferenceOrchestrator.__new__(InferenceOrchestrator) + cmd = o._build_generate_cmd( + "req1", + None, + messages = [{"role": "user", "content": "hi"}], + max_new_tokens = 128, + **_SAMPLING, + ) + for key, val in _SAMPLING.items(): + assert cmd[key] == val, f"{key} dropped/altered in orchestrator cmd" + + +def test_worker_forwards_all_sampling_params_to_backend(): + from core.inference.worker import _handle_generate + + class _RecordingBackend: + last_generation_stats = None + + def __init__(self): + self.received = None + + def generate_chat_response(self, **kwargs): + self.received = kwargs + return iter(()) # empty stream -> loop exits, gen_done is sent + + class _FakeQueue: + def __init__(self): + self.items = [] + + def put(self, item): + self.items.append(item) + + cmd = { + "type": "generate", + "request_id": "r", + "messages": [{"role": "user", "content": "hi"}], + "max_new_tokens": 128, + **_SAMPLING, + } + backend = _RecordingBackend() + _handle_generate(backend, cmd, _FakeQueue(), threading.Event()) + + assert backend.received is not None + for key, val in _SAMPLING.items(): + assert backend.received[key] == val, f"{key} dropped/altered in worker gen_kwargs" diff --git a/studio/backend/tests/test_response_template_markers.py b/studio/backend/tests/test_response_template_markers.py new file mode 100644 index 0000000000..8c813e62f2 --- /dev/null +++ b/studio/backend/tests/test_response_template_markers.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""TEMPLATE_TO_RESPONSES_MAPPER markers must match what the templates render. + +The manual instruction/response markers are the fallback for +train_on_completions when auto-detection is unavailable, so a marker that +never matches the rendered chat template masks every assistant token and the +run dies on the all-labels-masked safety net. Six template families shipped +such markers: + + mistral - "[INST] " / " [/INST]": the surrounding spaces fold into + the neighbouring tokens ("[INST]" is a single special + token in Mistral v0.3), so the padded strings never match. + llama - same space folding, plus llama-2 tokenizes [INST] after + as bare "[" on transformers 5.x while the standalone + encoding gives "▁[", so the marker must anchor on . + starling - trailing space after "GPT4 Correct Assistant:" folds + into the next content token ("▁Hello"). + glm - "[gMASK]" renders once at text start, never before + later user turns; "" is generation scaffolding + that non-final turns render as a lone "". + qwen3-thinking - "" is stripped from non-final assistant turns + (Qwen3-Thinking-2507) or never rendered (QwQ). + zephyr - role tags are plain text, and SentencePiece tokenizes + "<|assistant|>" differently at text start than after + "\\n" mid-conversation; the markers need the leading + newline anchor to tokenize like a real turn boundary. + +Literal assertions run everywhere; the token-level masking checks need the +representative tokenizers plus unsloth_zoo and skip when either is +unavailable (offline CI). +""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# model_mappings is dependency-free: load it directly so these tests run +# without the studio venv / package import side effects. +_MM_PATH = Path(_BACKEND_DIR) / "utils" / "datasets" / "model_mappings.py" +_mm_spec = importlib.util.spec_from_file_location("_marker_test_mm", _MM_PATH) +model_mappings = importlib.util.module_from_spec(_mm_spec) +_mm_spec.loader.exec_module(model_mappings) + +T2R = model_mappings.TEMPLATE_TO_RESPONSES_MAPPER + + +# ── Fixed entries: markers derived from what each representative tokenizer +# actually renders (see PR for the token-level derivation). ── +EXPECTED_FIXED = { + "mistral": {"instruction": "[INST]", "response": "[/INST]"}, + "llama": {"instruction": "[INST]", "response": "[/INST]"}, + "starling": {"instruction": "GPT4 Correct User:", "response": "GPT4 Correct Assistant:"}, + "glm": {"instruction": "<|user|>", "response": "<|assistant|>"}, + "qwen3-thinking": {"instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n"}, + "zephyr": {"instruction": "\n<|user|>\n", "response": "\n<|assistant|>\n"}, +} + +# Spot-pin some known-good entries so a refactor cannot silently change them. +EXPECTED_UNCHANGED = { + "qwen3": {"instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n"}, + "llama-3.1": { + "instruction": "<|start_header_id|>user<|end_header_id|>\n\n", + "response": "<|start_header_id|>assistant<|end_header_id|>\n\n", + }, + "phi-4": { + "instruction": "<|im_start|>user<|im_sep|>", + "response": "<|im_start|>assistant<|im_sep|>", + }, + "gemma-3": {"instruction": "user\n", "response": "model\n"}, + "gpt-oss": { + "instruction": "<|start|>user<|message|>", + "response": "<|start|>assistant<|channel|>final<|message|>", + }, +} + + +@pytest.mark.parametrize("template", sorted(EXPECTED_FIXED)) +def test_fixed_marker_literals(template): + assert T2R[template] == EXPECTED_FIXED[template] + + +@pytest.mark.parametrize("template", sorted(EXPECTED_UNCHANGED)) +def test_unchanged_marker_literals(template): + assert T2R[template] == EXPECTED_UNCHANGED[template] + + +def test_no_marker_is_empty_or_whitespace(): + for template, parts in T2R.items(): + assert parts["instruction"].strip(), template + assert parts["response"].strip(), template + + +# ── Token-level checks: markers must select exactly the assistant turns on a +# rendered two-turn fixture, and the final EOS label must never be -100. ── + +REPRESENTATIVES = { + "mistral": ["unsloth/mistral-7b-instruct-v0.3"], + "llama": ["unsloth/llama-2-7b-chat"], + "starling": ["unsloth/Starling-LM-7B-beta"], + "glm": ["unsloth/GLM-4.7-Flash"], + "qwen3-thinking": ["unsloth/Qwen3-4B-Thinking-2507", "Qwen/QwQ-32B"], + "zephyr": ["unsloth/zephyr-sft"], +} + +FIXTURE = [ + {"role": "user", "content": "zebra alpha question one?"}, + {"role": "assistant", "content": "grape reply number one."}, + {"role": "user", "content": "zebra beta question two?"}, + {"role": "assistant", "content": "grape reply number two."}, +] + + +def _load_tokenizer(repo): + try: + from transformers import AutoTokenizer + except Exception as e: # pragma: no cover + pytest.skip(f"transformers unavailable: {e}") + try: + return AutoTokenizer.from_pretrained(repo) + except OSError as e: + pytest.skip(f"tokenizer {repo} unavailable (offline?): {e}") + except Exception: + # Tokenizer class newer than this transformers (e.g. GLM-4.7's + # TokenizersBackend): build directly from tokenizer.json. + try: + import json as _json + from huggingface_hub import hf_hub_download + from transformers import PreTrainedTokenizerFast + + with open(hf_hub_download(repo, "tokenizer_config.json"), encoding = "utf-8") as f: + cfg = _json.load(f) + tok_file = hf_hub_download(repo, "tokenizer.json") + + def _tokval(v): + return v["content"] if isinstance(v, dict) else v + + return PreTrainedTokenizerFast( + tokenizer_file = tok_file, + chat_template = cfg.get("chat_template"), + **{ + k: _tokval(cfg[k]) + for k in ("bos_token", "eos_token", "pad_token", "unk_token") + if cfg.get(k) is not None + }, + ) + except Exception as e: + pytest.skip(f"tokenizer {repo} unavailable (offline?): {e}") + + +def _train_on_responses_only(): + try: + from unsloth_zoo.dataset_utils import train_on_responses_only + except Exception as e: + pytest.skip(f"unsloth_zoo unavailable: {e}") + return train_on_responses_only + + +@pytest.mark.parametrize( + "template,repo", + [(t, r) for t, repos in sorted(REPRESENTATIVES.items()) for r in repos], +) +def test_fixed_markers_token_level(template, repo): + tor = _train_on_responses_only() + tok = _load_tokenizer(repo) + parts = T2R[template] + + msgs = [{"role": "system", "content": "You are a terse assistant."}] + FIXTURE + try: + ids = tok.apply_chat_template(msgs, tokenize = True, add_generation_prompt = False) + if hasattr(ids, "keys"): + ids = ids["input_ids"] # transformers 5.x returns a BatchEncoding + except Exception: + ids = tok.apply_chat_template(FIXTURE, tokenize = True, add_generation_prompt = False) + if hasattr(ids, "keys"): + ids = ids["input_ids"] + + fn = tor( + None, + instruction_part = parts["instruction"], + response_part = parts["response"], + tokenizer = tok, + return_function = True, + ) + labels = fn({"input_ids": [list(ids)]})["labels"][0] + + n = len(ids) + trained = tok.decode([ids[i] for i in range(n) if labels[i] != -100]) + masked = tok.decode([ids[i] for i in range(n) if labels[i] == -100]) + + # User and system content fully masked + assert "question one" not in trained and "question one" in masked + assert "question two" not in trained and "question two" in masked + assert "terse assistant" not in trained + # EVERY assistant turn trained, not just the last + assert "reply number one" in trained + assert "reply number two" in trained + # The final EOS (last non-whitespace token) must never be -100, or the + # fine-tuned model never learns to stop generating. + i = n - 1 + while i > 0 and tok.decode([ids[i]]).strip() == "": + i -= 1 + assert labels[i] != -100, f"final token {tok.convert_ids_to_tokens(int(ids[i]))!r} is masked" + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index a7ceb49ed9..46dd0d42e4 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -59,6 +59,7 @@ from models.inference import ( ResponsesUsage, ) from routes.inference import ( + _ResponsesReasoningExtractor, _SameTaskStreamingResponse, _build_chat_request, _chat_tool_calls_to_responses_output, @@ -795,6 +796,7 @@ class TestResponsesNonStreamingAdapter: def test_monitor_records_translated_visible_text(self, monkeypatch): import routes.inference as inf_mod + import routes.inference as inf_mod async def fake_chat_completions(chat_req, request): assert request.state.skip_api_monitor is True @@ -1988,6 +1990,123 @@ class TestTranslatedMessagesValidate: ChatMessage(**m.model_dump(exclude_none = True)) +# reasoning_prefilled: enable_thinking templates prefill an unclosed , so +# generation begins inside the block; the extractor must start in reasoning. +class TestReasoningPrefilledExtractor: + def test_prefilled_single_feed_splits_lone_close(self): + # T1: reasoning...answer with a prefilled (unseen) open tag. + reasoning, visible = _extract_responses_reasoning( + "plananswer", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "plan" + assert visible == "answer" + + def test_prefilled_never_closed_is_all_reasoning(self): + # T2: truncated mid-thought (no ) -> all reasoning (GGUF parity). + reasoning, visible = _extract_responses_reasoning( + "still thinking with no close", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "still thinking with no close" + assert visible == "" + + def test_prefilled_close_split_across_feeds(self): + # T3: straddles two feed() calls; holdback resolves it. + ex = _ResponsesReasoningExtractor(parse_think_markers = True, reasoning_prefilled = True) + r1, v1 = ex.feed("planans") + fr, fv = ex.finish() + assert (r1 + r2 + fr) == "plan" + assert (v1 + v2 + fv) == "ans" + + def test_prefilled_close_split_one_char_per_feed(self): + # T4: every char in its own feed still splits correctly. + ex = _ResponsesReasoningExtractor(parse_think_markers = True, reasoning_prefilled = True) + reasoning, visible = "", "" + for ch in "planx": + r, v = ex.feed(ch) + reasoning += r + visible += v + fr, fv = ex.finish() + assert (reasoning + fr) == "plan" + assert (visible + fv) == "x" + + def test_prefilled_empty_generation(self): + # T5: nothing generated. + reasoning, visible = _extract_responses_reasoning( + "", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "" + assert visible == "" + + def test_prefilled_whitespace_after_close_is_visible(self): + # T6: Qwen commonly emits \n\n before the answer. + reasoning, visible = _extract_responses_reasoning( + "plan\n\nanswer", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "plan" + assert visible == "\n\nanswer" + + def test_prefilled_stray_open_tag_is_suppressed(self): + # T7: a re-emitted literal inside prefilled reasoning is dropped, + # not leaked into the drawer (covers enable_thinking_effort full-tag output). + reasoning, visible = _extract_responses_reasoning( + "abc", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "ab" + assert visible == "c" + assert "" not in reasoning + + def test_prefilled_close_at_start_empty_reasoning(self): + # T8: model closed immediately (empty reasoning) then answered. + reasoning, visible = _extract_responses_reasoning( + "hi", + parse_think_markers = True, + reasoning_prefilled = True, + ) + assert reasoning == "" + assert visible == "hi" + + def test_not_prefilled_lone_close_preserves_current_behavior(self): + # T9: without prefilled, a lone close tag keeps the pre-fix behavior (parity guard). + reasoning, visible = _extract_responses_reasoning( + "reasoningans", + parse_think_markers = True, + reasoning_prefilled = False, + ) + assert reasoning == "" + assert visible == "reasoningans" + + def test_not_prefilled_full_pair_still_splits(self): + # T10: normal explicit .. (GGUF / Harmony) unchanged. + reasoning, visible = _extract_responses_reasoning( + "rv", + parse_think_markers = True, + reasoning_prefilled = False, + ) + assert reasoning == "r" + assert visible == "v" + + def test_prefilled_ignored_when_markers_not_parsed(self): + # T11: a non-reasoning model passes text through even with reasoning_prefilled False. + reasoning, visible = _extract_responses_reasoning( + "just an answer", + parse_think_markers = False, + reasoning_prefilled = False, + ) + assert reasoning == "" + assert visible == "just an answer" + + # ===================================================================== # Streaming passthrough healing — text-form calls promoted in order # ===================================================================== diff --git a/studio/backend/tests/test_safetensors_capability_advertise.py b/studio/backend/tests/test_safetensors_capability_advertise.py index 671af93708..0ed670ac01 100644 --- a/studio/backend/tests/test_safetensors_capability_advertise.py +++ b/studio/backend/tests/test_safetensors_capability_advertise.py @@ -11,6 +11,8 @@ from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock +import pytest + _backend_root = Path(__file__).resolve().parent.parent if str(_backend_root) not in sys.path: sys.path.insert(0, str(_backend_root)) @@ -46,6 +48,21 @@ reasoning_effort: {{ reasoning_effort }} """ +# DeepSeek-V4-Flash: an enable_thinking on/off gate PLUS a reasoning_effort +# 'max' preamble. The shipped template only *branches* on 'max' ('high' renders +# identically to thinking-on-without-the-preamble), so the literal scan alone +# would surface only ['max']; the classifier adds 'high' for deepseek-v4 to +# expose the encoder's full none/high/max ladder. +DEEPSEEK_V4_TEMPLATE = ( + "{%- if not thinking is defined %}" + "{%- if enable_thinking is defined %}{%- set thinking = enable_thinking %}" + "{%- else %}{%- set thinking = false %}{%- endif %}{%- endif %}\n" + "{%- if thinking and reasoning_effort == 'max' %}" + "{{- 'Reasoning Effort: Absolute maximum' }}{%- endif %}\n" + "{%- for message in messages %}{{- message.content }}{%- endfor %}" +) + + PLAIN_TEMPLATE = """ {%- for message in messages %} {{- message.role + ': ' + message.content + '\\n' }} @@ -88,6 +105,29 @@ def test_detect_reasoning_flags_none_template_returns_all_false(): assert flags["reasoning_style"] == "enable_thinking" +def test_detect_reasoning_flags_deepseek_v4_exposes_none_high_max(): + """DeepSeek-V4-Flash: enable_thinking gate + reasoning_effort 'max' preamble. + Classified as the hybrid style with the full none/high/max ladder even + though the template only branches on 'max'.""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "unsloth/DeepSeek-V4-Flash-GGUF") + assert flags["supports_reasoning"] is True + assert flags["reasoning_style"] == "enable_thinking_effort" + assert flags["reasoning_effort_levels"] == ["high", "max"] + assert flags["reasoning_always_on"] is False + + +def test_detect_reasoning_flags_non_deepseek_v4_effort_only_max_not_injected(): + """The 'high' injection is scoped to deepseek-v4: a different model whose + template only branches on 'max' keeps ['max'] (no phantom 'high').""" + from core.inference.llama_cpp import detect_reasoning_flags + + flags = detect_reasoning_flags(DEEPSEEK_V4_TEMPLATE, "vendor/OtherHybrid-GGUF") + assert flags["reasoning_style"] == "enable_thinking_effort" + assert flags["reasoning_effort_levels"] == ["max"] + + def test_detect_safetensors_features_passes_template_through_to_classifier(): """Route wrapper forwards a real template to the inner classifier.""" from routes.inference import _detect_safetensors_features @@ -127,9 +167,8 @@ def test_detect_safetensors_features_gptoss_disables_tools(): assert flags["supports_tools"] is False -# Llama-3 / Mistral advertise tools but emit <|python_tag|> / [TOOL_CALLS], -# which our parser can't read. The route helper must not flip supports_tools=True -# for them, else the UI enables a pill the agentic loop can't honour. +# Llama-3 / Mistral / Gemma 4 tool-call formats are now parser-supported, so supports_tools=True +# must hold for all of them; only templates matching none of the five known markers are suppressed. LLAMA3_TEMPLATE = """ {%- if tools %} @@ -161,27 +200,188 @@ MISTRAL_TEMPLATE = """ {%- endfor %} """ +GEMMA4_TEMPLATE = """ +{%- if tools %} + {{- 'Tools available. Emit calls as ' }} + {{- '<|tool_call>call:NAME{key:<|"|>val<|"|>}' }} + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +""" -def test_detect_safetensors_features_llama3_template_suppresses_tools(): - """Llama-3 emits <|python_tag|>; safetensors loop cannot parse it.""" + +def test_detect_safetensors_features_llama3_template_keeps_tools_on(): + """Llama-3 emits <|python_tag|>; parser now supports it.""" from routes.inference import _detect_safetensors_features backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") flags = _detect_safetensors_features(backend, LLAMA3_TEMPLATE) - assert flags["supports_tools"] is False + assert flags["supports_tools"] is True -def test_detect_safetensors_features_mistral_template_suppresses_tools(): - """Mistral emits [TOOL_CALLS]; safetensors loop cannot parse it.""" +def test_detect_safetensors_features_mistral_template_keeps_tools_on(): + """Mistral emits [TOOL_CALLS]name{json}, which the safetensors loop now parses + (the shared bracket-tag parser). The gate must no longer suppress it, or the + PR's Mistral tool support is unreachable through normal capability detection.""" from routes.inference import _detect_safetensors_features backend = SimpleNamespace(active_model_name = "unsloth/mistral-7b-instruct-v0.3") flags = _detect_safetensors_features(backend, MISTRAL_TEMPLATE) + assert flags["supports_tools"] is True + + +def test_detect_safetensors_features_gemma4_template_keeps_tools_on(): + """Gemma 4 emits <|tool_call>; parser now supports it.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/gemma-4-E2B-it-UD-MLX-4bit") + flags = _detect_safetensors_features(backend, GEMMA4_TEMPLATE) + assert flags["supports_tools"] is True + + +# DeepSeek V3 / V3.1 / R1 emit ``<｜tool▁calls▁begin｜>...`` blocks. +# Note the full-width pipe (U+FF5C) and lower-1/8-block (U+2581). +DEEPSEEK_TEMPLATE = """ +{%- if tools %} + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +{%- for message in messages %} + {%- if message.role == 'assistant' and message.tool_calls %} + {%- for tc in message.tool_calls %} + {{- '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>' + tc.function.name + + '<｜tool▁sep｜>' + tc.function.arguments + '<｜tool▁call▁end｜>' }} + {%- endfor %} + {%- endif %} +{%- endfor %} +""" + + +def test_detect_safetensors_features_deepseek_template_keeps_tools_on(): + """DeepSeek emits ``<｜tool▁calls▁begin｜>...``; parser now supports it.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/DeepSeek-V3.1") + flags = _detect_safetensors_features(backend, DEEPSEEK_TEMPLATE) + assert flags["supports_tools"] is True + + +# GLM 4.5 / 4.6 / 4.7 emit ``NAME\n...... +GLM_TEMPLATE = """ +{%- if tools %} + For each function call, output the function name and arguments within + the following XML format: + {function-name} + {arg-key} + {arg-value} + + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +""" + + +def test_detect_safetensors_features_glm_template_keeps_tools_on(): + """GLM 4.x emits ``NAME\\n...``; parser handles it.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/GLM-4.6") + flags = _detect_safetensors_features(backend, GLM_TEMPLATE) + assert flags["supports_tools"] is True + + +# Kimi K2 / Moonshot uses ``<|tool_calls_section_begin|>...`` blocks +# with ``functions.NAME:IDX`` as the per-call id. +KIMI_TEMPLATE = """ +{%- if tools %} + <|im_system|>tool_declare<|im_middle|>{{ tools | tojson }}<|im_end|> +{%- endif %} +{%- for message in messages %} + {%- if message.role == 'assistant' and message.tool_calls %} + <|tool_calls_section_begin|> + {%- for tc in message.tool_calls %} + <|tool_call_begin|>{{ tc.id }}<|tool_call_argument_begin|>{{ tc.function.arguments | tojson }}<|tool_call_end|> + {%- endfor %} + <|tool_calls_section_end|> + {%- endif %} +{%- endfor %} +""" + + +def test_detect_safetensors_features_kimi_template_keeps_tools_on(): + """Kimi K2 emits ``<|tool_calls_section_begin|>...``; parser handles it.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Kimi-K2-Instruct") + flags = _detect_safetensors_features(backend, KIMI_TEMPLATE) + assert flags["supports_tools"] is True + + +LLAMA3_2_BARE_JSON_TEMPLATE = """ +{%- if tools %} + {{- 'Given the following functions, respond with JSON for a function call.' }} + {{- 'Respond in the format {"name": function name, "parameters": dictionary}.' }} + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +{%- for message in messages %} + {%- if 'tool_calls' in message %} + {{- '{"name": "' + message.tool_calls[0].function.name + '", '}} + {{- '"parameters": ' + (message.tool_calls[0].function.arguments | tojson) + '}' }} + {%- endif %} +{%- endfor %} +""" + + +def test_detect_safetensors_features_llama3_2_bare_json_keeps_tools_on(): + """Llama-3.2 bare JSON is supported, so the pill stays enabled.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, LLAMA3_2_BARE_JSON_TEMPLATE) + assert flags["supports_tools"] is True + + +MINICPM5_ATTRIBUTE_TEMPLATE = """ +{%- if tools %} + {{- 'Available tools. Emit calls as ' }} + {{- 'value' }} + {%- for tool in tools %} + {{- tool | tojson }} + {%- endfor %} +{%- endif %} +""" + + +def test_detect_safetensors_features_attribute_function_form_keeps_tools_on(): + """The attribute form ```` must be whitelisted or the pill is wrongly suppressed.""" + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "openbmb/MiniCPM-5") + flags = _detect_safetensors_features(backend, MINICPM5_ATTRIBUTE_TEMPLATE) + assert flags["supports_tools"] is True + + +def test_detect_safetensors_features_unknown_format_suppresses_tools(): + """Tools advertised with no known marker must be suppressed.""" + from routes.inference import _detect_safetensors_features + + tpl = ( + "{%- if tools %}<|im_start|>system\n" + "Emit tool calls as JSON-RPC notifications inside the response." + "<|im_end|>{%- endif %}" + ) + backend = SimpleNamespace(active_model_name = "custom/unknown-tool-format") + flags = _detect_safetensors_features(backend, tpl) assert flags["supports_tools"] is False def test_detect_safetensors_features_qwen_tool_call_keeps_tools_on(): - """Sanity check: gate only suppresses non-Qwen formats.""" + """Sanity check: Qwen marker still flips supports_tools.""" from routes.inference import _detect_safetensors_features backend = SimpleNamespace(active_model_name = "unsloth/Qwen3-0.6B") @@ -454,3 +654,184 @@ def test_route_layer_emits_supports_tools_true_for_qwen3_safetensors(): assert flags["supports_tools"] is True assert flags["supports_reasoning"] is True assert flags["supports_preserve_thinking"] is True + + +@pytest.mark.parametrize( + "opener", + [ + "<｜tool▁calls▁begin｜>", # canonical + "<｜tool_calls_begin｜>", # ASCII underscores + "<｜tool▁calls｜>", # short form + "<｜tool calls begin｜>", # spaces + "<｜tool\\_calls\\_begin｜>", # escaped underscores + ], +) +def test_detect_safetensors_features_deepseek_opener_variants_keep_tools_on(opener): + # Every DeepSeek opener the parser accepts must keep supports_tools on; the route gate derives + # its markers from the parser's TOOL_XML_SIGNALS so it can no longer drift behind the parser ... + from routes.inference import _detect_safetensors_features + + tpl = ( + "{%- if tools %}tools{%- endif %}" + + opener + + "<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_time{}" + "<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + backend = SimpleNamespace(active_model_name = "unsloth/DeepSeek-V3.1") + flags = _detect_safetensors_features(backend, tpl) + assert flags["supports_tools"] is True + + +# Templates that advertise tools ({%- if tools %}) and prompt the bare-JSON +# call form, but whose ``{"name":`` example is pretty-printed or JSON-escaped. +_WHITESPACE_BARE_JSON_TEMPLATE = ( + "{%- if tools %}\n" + "To call a tool, output JSON of the form:\n" + '{ "name" : "function_name", "parameters": { } }\n' + "{%- endif %}\n" + "{{ messages }}" +) +_ESCAPED_BARE_JSON_TEMPLATE = ( + "{%- if tools %}\n" + 'Respond with {\\"name\\": \\"fn\\", \\"parameters\\": {}}\n' + "{%- endif %}\n" + "{{ messages }}" +) +_TOOLS_ADVERTISED_NO_PARSEABLE_FORM = ( + "{%- if tools %}\nYou may use the available tools.\n{%- endif %}\n{{ messages }}" +) + + +def test_detect_safetensors_features_keeps_tools_for_pretty_printed_bare_json(): + # A pretty-printed bare-JSON example (``{ "name" :``) must keep supports_tools since the parser + # accepts that whitespace via raw_decode. + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, _WHITESPACE_BARE_JSON_TEMPLATE) + assert flags["supports_tools"] is True + + +def test_detect_safetensors_features_keeps_tools_for_escaped_bare_json(): + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, _ESCAPED_BARE_JSON_TEMPLATE) + assert flags["supports_tools"] is True + + +def test_detect_safetensors_features_drops_tools_when_no_parseable_form(): + # Negative control: tools advertised but no parser-recognised emission form at + # all -> the pill is still dropped (the gate is not now matching everything). + from routes.inference import _detect_safetensors_features + + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, _TOOLS_ADVERTISED_NO_PARSEABLE_FORM) + assert flags["supports_tools"] is False + + +def test_detect_safetensors_features_keeps_tools_for_function_alias_bare_json(): + # A template documenting the parser-supported {"function":...} bare-JSON alias + # must keep supports_tools, mirroring the {"name":...} form. + from routes.inference import _detect_safetensors_features + + tpl = ( + "{%- if tools %}\n" + 'Respond with {"function": "fn", "parameters": {}}\n' + "{%- endif %}\n" + "{{ messages }}" + ) + backend = SimpleNamespace(active_model_name = "unsloth/Llama-3.2-3B-Instruct") + flags = _detect_safetensors_features(backend, tpl) + assert flags["supports_tools"] is True + + +# _sf_reasoning_prefill_mode gates the prefilled- extractor (GGUF reasoning parity). +class TestSafetensorsReasoningPrefillGate: + # A minimal Qwen3-style template with the standard / markers. + _QWEN_TPL = "{% if enable_thinking %}{% endif %}......" + # gemma-style bespoke reasoning channel -- no standard markers. + _GEMMA_TPL = "{% if enable_thinking %}<|think|>{% endif %}<|channel>thought" + # always-on template whose GENERATION PROMPT opens an unclosed (DeepSeek-R1 / QwQ / + # Qwen3-Thinking shape): the model emits only the closing , so prefill. + _ALWAYS_ON_OPEN_TPL = ( + "{% for m in messages %}{{ m['content'] }}{% endfor %}" + "{% if add_generation_prompt %}<|assistant|>\n{% endif %}" + ) + # always-on template that renders PAST assistant ... history but leaves the + # generation prompt open with no (Kimi-K2-Thinking shape): the model self-emits its + # own block, so prefill mode would blank a normal answer. + _ALWAYS_ON_HISTORY_TPL = ( + "{% for m in messages %}" + "{% if m['role'] == 'assistant' %}{{ m.get('reasoning_content', '') }}" + "{{ m['content'] }}{% endif %}" + "{% endfor %}" + "{% if add_generation_prompt %}<|im_assistant|>assistant<|im_middle|>{% endif %}" + ) + + def _features(self, **over): + base = { + "supports_reasoning": True, + "reasoning_always_on": False, + "reasoning_style": "enable_thinking", + } + base.update(over) + return base + + def test_g1_enable_thinking_true(self): + # G1: Qwen3.5 template + explicit enable_thinking=True -> prefilled. + from routes.inference import _sf_reasoning_prefill_mode + assert _sf_reasoning_prefill_mode(self._features(), True, self._QWEN_TPL) is True + + def test_g2_enable_thinking_none_defaults_on(self): + # G2: default request (None) -> prefilled (Qwen3/GLM templates default on). + from routes.inference import _sf_reasoning_prefill_mode + assert _sf_reasoning_prefill_mode(self._features(), None, self._QWEN_TPL) is True + + def test_g3_enable_thinking_false(self): + # G3: thinking explicitly off -> not prefilled. + from routes.inference import _sf_reasoning_prefill_mode + assert _sf_reasoning_prefill_mode(self._features(), False, self._QWEN_TPL) is False + + def test_g4_gpt_oss_reasoning_effort_excluded(self): + # G4: gpt-oss uses explicit tags via HarmonyTextStreamer -> normal mode. + from routes.inference import _sf_reasoning_prefill_mode + feats = self._features(reasoning_style = "reasoning_effort") + assert _sf_reasoning_prefill_mode(feats, True, self._QWEN_TPL) is False + + def test_g5_enable_thinking_effort_included(self): + # G5: GLM-style enable_thinking_effort also prefills. + from routes.inference import _sf_reasoning_prefill_mode + feats = self._features(reasoning_style = "enable_thinking_effort") + assert _sf_reasoning_prefill_mode(feats, None, self._QWEN_TPL) is True + + def test_g6_non_reasoning_model(self): + # G6: no reasoning capability -> never prefilled. + from routes.inference import _sf_reasoning_prefill_mode + feats = self._features(supports_reasoning = False, reasoning_style = None) + assert _sf_reasoning_prefill_mode(feats, True, self._QWEN_TPL) is False + + def test_g7_reasoning_always_on_prompt_opens_think(self): + # G7: always-on template whose generation prompt opens -> prefilled regardless of the flag. + from routes.inference import _sf_reasoning_prefill_mode + feats = self._features(reasoning_always_on = True) + assert _sf_reasoning_prefill_mode(feats, False, self._ALWAYS_ON_OPEN_TPL) is True + + def test_g7b_reasoning_always_on_history_only_not_prefilled(self): + # G7b (#5704): always-on classification from rendered assistant HISTORY + # (Kimi-K2-Thinking) whose generation prompt opens no . Prefill mode would capture a + # normal answer entirely as reasoning_content and blank the visible answer, so it must be off. + from routes.inference import _sf_reasoning_prefill_mode + feats = self._features(reasoning_always_on = True) + assert _sf_reasoning_prefill_mode(feats, None, self._ALWAYS_ON_HISTORY_TPL) is False + + def test_g8_gemma_bespoke_channel_excluded(self): + # G8: gemma's <|think|>/<|channel> format has no -> NOT prefilled + # (would otherwise swallow the whole answer as reasoning). Regression guard. + from routes.inference import _sf_reasoning_prefill_mode + assert _sf_reasoning_prefill_mode(self._features(), True, self._GEMMA_TPL) is False + + def test_g9_missing_template_not_prefilled(self): + # G9: no template available -> conservative (not prefilled). + from routes.inference import _sf_reasoning_prefill_mode + assert _sf_reasoning_prefill_mode(self._features(), True, None) is False diff --git a/studio/backend/tests/test_safetensors_reasoning_stream.py b/studio/backend/tests/test_safetensors_reasoning_stream.py new file mode 100644 index 0000000000..4e708139b7 --- /dev/null +++ b/studio/backend/tests/test_safetensors_reasoning_stream.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Safetensors/MLX reasoning-block parity with GGUF. + +enable_thinking templates (Qwen3/GLM) prefill an unclosed ```` so the model +emits only the closing ```` then the answer; the safetensors stream must +split the leading text into ``reasoning_content`` deltas (plain stream and tool +loop), resetting per turn and appending only visible text to the monitor. Replays a +copy of ``sf_tool_stream``'s reasoning loop against synthetic events. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from routes.inference import ( + _ResponsesReasoningExtractor, + _sf_reasoning_prefill_mode, + _strip_tool_xml_for_display, +) + + +_THINK_TPL = "........." +_ETHINK = {"reasoning_style": "enable_thinking", "supports_reasoning": True} +_ETHINK_EFFORT = {"reasoning_style": "enable_thinking_effort", "supports_reasoning": True} + + +def test_prefill_mode_on_for_enable_thinking_default(): + assert _sf_reasoning_prefill_mode(_ETHINK, None, _THINK_TPL) is True + + +def test_prefill_mode_off_when_thinking_disabled(): + assert _sf_reasoning_prefill_mode(_ETHINK, False, _THINK_TPL) is False + + +def test_prefill_mode_off_for_reasoning_effort_none(): + # enable_thinking_effort turns thinking off via reasoning_effort="none"; prefilled mode + # would capture the whole answer as reasoning_content. + assert ( + _sf_reasoning_prefill_mode(_ETHINK_EFFORT, None, _THINK_TPL, reasoning_effort = "none") + is False + ) + assert ( + _sf_reasoning_prefill_mode(_ETHINK_EFFORT, None, _THINK_TPL, reasoning_effort = "high") + is True + ) + + +def test_prefill_mode_off_without_think_markers(): + assert _sf_reasoning_prefill_mode(_ETHINK, None, "no markers here") is False + + +def _replay_sf_reasoning_stream(events: list[dict], *, prefilled: bool) -> dict: + """Mirror sf_tool_stream's reasoning loop: diff each cumulative ``content`` + snapshot, feed the delta through the extractor, and reset (flushing first) on + ``tool_start`` / empty ``status`` so each turn splits independently.""" + prev_text = "" + extractor = _ResponsesReasoningExtractor( + parse_think_markers = True, reasoning_prefilled = prefilled + ) + reasoning_deltas: list[str] = [] + visible_deltas: list[str] = [] + monitor: list[str] = [] + tool_starts: list[dict] = [] + order: list[str] = [] # sequence of ("reasoning"|"visible"|"tool_start") events + + def _flush(): + fr, fv = extractor.finish() + if fr: + reasoning_deltas.append(fr) + order.append("reasoning") + if fv: + visible_deltas.append(fv) + monitor.append(fv) + order.append("visible") + + for event in events: + etype = event["type"] + if etype == "status": + if not event["text"]: + _flush() + prev_text = "" + extractor = _ResponsesReasoningExtractor( + parse_think_markers = True, reasoning_prefilled = prefilled + ) + continue + if etype in ("tool_start", "tool_end"): + if etype == "tool_start": + _flush() + prev_text = "" + extractor = _ResponsesReasoningExtractor( + parse_think_markers = True, reasoning_prefilled = prefilled + ) + tool_starts.append(event) + order.append("tool_start") + continue + clean = _strip_tool_xml_for_display(event.get("text", ""), auto_heal_tool_calls = True) + new_text = clean[len(prev_text) :] + prev_text = clean + if not new_text: + continue + r, v = extractor.feed(new_text) + if r: + reasoning_deltas.append(r) + order.append("reasoning") + if v: + visible_deltas.append(v) + monitor.append(v) + order.append("visible") + _flush() + return { + "reasoning": "".join(reasoning_deltas), + "visible": "".join(visible_deltas), + "monitor": "".join(monitor), + "tool_starts": tool_starts, + "order": order, + } + + +def test_s1_plain_stream_splits_prefilled_reasoning(): + # S1: plain/MLX single turn -> reasoning delta + visible delta; monitor visible-only. + events = [ + {"type": "content", "text": "Let me compute 17*23"}, + {"type": "content", "text": "Let me compute 17*23 = 391The answer is 391."}, + ] + out = _replay_sf_reasoning_stream(events, prefilled = True) + assert out["reasoning"] == "Let me compute 17*23 = 391" + assert out["visible"] == "The answer is 391." + assert out["monitor"] == "The answer is 391." + assert "" not in out["reasoning"] and "" not in out["visible"] + + +def test_s2_reasoning_flushed_before_tool_start(): + # S2: reasoning streamed as reasoning_content, then flushed BEFORE tool_start. + events = [ + {"type": "content", "text": "I should search"}, + {"type": "content", "text": "I should search Sydney weather"}, + {"type": "tool_start", "tool_name": "web_search", "tool_call_id": "c0"}, + {"type": "tool_end", "tool_name": "web_search", "tool_call_id": "c0"}, + {"type": "status", "text": ""}, + {"type": "content", "text": "Found itSydney is 21C today."}, + ] + out = _replay_sf_reasoning_stream(events, prefilled = True) + # Both turns' reasoning surfaced, answer only from turn 2. + assert "I should search Sydney weather" in out["reasoning"] + assert "Found it" in out["reasoning"] + assert out["visible"] == "Sydney is 21C today." + assert out["monitor"] == "Sydney is 21C today." + # Ordering: the pre-tool reasoning is emitted before the tool_start. + assert out["order"].index("reasoning") < out["order"].index("tool_start") + + +def test_s3_extractor_resets_each_turn(): + # S3: multi-turn -> the two turns' reasoning are distinct (fresh extractor each). + events = [ + {"type": "content", "text": "turn1 thoughtspartial"}, + {"type": "status", "text": ""}, + {"type": "content", "text": "turn2 thoughtsfinal answer"}, + ] + out = _replay_sf_reasoning_stream(events, prefilled = True) + assert out["reasoning"] == "turn1 thoughtsturn2 thoughts" + assert out["visible"] == "partialfinal answer" + + +def test_s4_harmony_full_tags_normal_mode(): + # S4: gpt-oss / explicit-tag models use normal mode (prefilled=False). + events = [{"type": "content", "text": "reasoning herevisible answer"}] + out = _replay_sf_reasoning_stream(events, prefilled = False) + assert out["reasoning"] == "reasoning here" + assert out["visible"] == "visible answer" + + +def test_s5_thinking_off_no_reasoning_deltas(): + # S5: thinking disabled -> not prefilled, no , all content is visible. + events = [{"type": "content", "text": "Just the plain answer, no thinking."}] + out = _replay_sf_reasoning_stream(events, prefilled = False) + assert out["reasoning"] == "" + assert out["visible"] == "Just the plain answer, no thinking." + assert out["monitor"] == "Just the plain answer, no thinking." + + +def test_s6_reasoning_effort_none_disables_prefill_for_enable_thinking_effort(): + # GLM-5.2-style enable_thinking_effort: a request with reasoning_effort="none" (and + # enable_thinking omitted) disables thinking exactly like enable_thinking=False, so + # prefilled mode must be OFF. Otherwise the model emits no and a plain + # answer is swallowed whole into reasoning_content, leaving the visible response + # empty (the exact bug: prefilled=True below eats the whole answer). + feats = {"reasoning_style": "enable_thinking_effort", "supports_reasoning": True} + assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, "none") is False + # Thinking on (effort level or default) still prefills. + assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, "high") is True + assert _sf_reasoning_prefill_mode(feats, None, _THINK_TPL, None) is True + # An explicit enable_thinking=False also disables (unchanged). + assert _sf_reasoning_prefill_mode(feats, False, _THINK_TPL, "high") is False + # reasoning_always_on wins regardless of reasoning_effort. + always = {**feats, "reasoning_always_on": True} + assert _sf_reasoning_prefill_mode(always, None, _THINK_TPL, "none") is True + # Plain enable_thinking models (Qwen) have no "none" sentinel; unaffected. + plain = {"reasoning_style": "enable_thinking", "supports_reasoning": True} + assert _sf_reasoning_prefill_mode(plain, None, _THINK_TPL, "none") is True + + # End-to-end: with the corrected prefilled=False, a plain no- answer is + # emitted as visible content rather than swallowed into the thinking drawer. + events = [{"type": "content", "text": "The capital of France is Paris."}] + out = _replay_sf_reasoning_stream(events, prefilled = False) + assert out["visible"] == "The capital of France is Paris." + assert out["reasoning"] == "" + # The buggy prefilled=True path is what swallowed the whole answer (guard the delta). + swallowed = _replay_sf_reasoning_stream(events, prefilled = True) + assert swallowed["visible"] == "" + assert swallowed["reasoning"] == "The capital of France is Paris." diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 3f2d49f0dd..a8546b82c4 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -115,6 +115,22 @@ class TestParser: assert result[0]["function"]["name"] == "python" assert "print('hi')" in result[0]["function"]["arguments"] + def test_xml_param_preserves_leading_indentation(self): + import json + + # Only the wrapping newline is trimmed; code-argument indentation survives. + text = ( + "\n" + " indented = 1\n" + " more\n" + "" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == { + "code": " indented = 1\n more" + } + def test_xml_unclosed(self): # Closing tags omitted; parser must still extract the value. text = "ls -la" @@ -138,6 +154,20 @@ class TestParser: assert len(result) == 1 assert "print('hi')" in result[0]["function"]["arguments"] + def test_xml_param_preserves_leading_indentation(self): + # Only the wrapping newline is trimmed, so code-argument indentation survives (str.strip() destroyed it). + text = ( + "\n" + " indented = 1\n" + " more\n" + "" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == { + "code": " indented = 1\n more" + } + def test_function_signal_inside_parameter_is_literal(self): text = ( "" @@ -169,6 +199,8 @@ class TestParser: assert has_tool_signal("blah x") assert has_tool_signal("blah <|tool_call>call:terminal") assert has_tool_signal("hi ...") + assert has_tool_signal("ok [TOOL_CALLS]web_search{...") + assert has_tool_signal("fine python[ARGS]{...") assert not has_tool_signal("hello world") def test_render_html_start_detector_uses_first_tool(self): @@ -183,12 +215,59 @@ class TestParser: '{"name":"python","arguments":{"code":""}}' ) + def test_render_html_start_detector_covers_mistral_and_rehearsal_forms(self): + # The provisional render-html card must fire for bracket-tag forms too, not only XML. + assert _detect_render_html_tool_start('[TOOL_CALLS]render_html{"code":""}') + assert _detect_render_html_tool_start('[TOOL_CALLS]render_html[ARGS]{"code":"x"}') + assert _detect_render_html_tool_start( + '[TOOL_CALLS] [{"name":"render_html","arguments":{}}]' + ) + assert _detect_render_html_tool_start('render_html[ARGS]{"code":""}') + # A different first tool (or a prose mention with no JSON body) must not fire. + assert not _detect_render_html_tool_start('[TOOL_CALLS]web_search{"q":"x"}') + assert not _detect_render_html_tool_start('web_search[ARGS]{"q":"x"}') + assert not _detect_render_html_tool_start('python[ARGS]{"code":"render_html[ARGS]{}"}') + assert not _detect_render_html_tool_start("use render_html[ARGS] to render") + + def test_render_html_start_detector_skips_think_block_rehearsal(self): + # A render_html rehearsed inside think must not fire the card; the outside-think call decides. + assert not _detect_render_html_tool_start( + 'draft render_html[ARGS]{"code":"x"}python[ARGS]{"code":"print(1)"}' + ) + assert not _detect_render_html_tool_start( + '[THINK]render_html[ARGS]{"code":"x"}[/THINK]web_search[ARGS]{"q":"y"}' + ) + # A real render_html AFTER a rehearsed non-render_html inside think still fires. + assert _detect_render_html_tool_start( + 'web_search[ARGS]{"q":"x"}render_html[ARGS]{"code":""}' + ) + # A render_html rehearsed inside think with no real call after does not fire. + assert not _detect_render_html_tool_start('render_html[ARGS]{"code":"x"}') + + def test_render_html_start_detector_reads_top_level_array_name(self): + # Array form: the name is the object's top-level ``"name"``, not an argument key. + assert not _detect_render_html_tool_start( + '[TOOL_CALLS] [{"arguments":{"name":"render_html"},"name":"python"}]' + ) + assert _detect_render_html_tool_start( + '[TOOL_CALLS] [{"arguments":{"name":"python"},"name":"render_html"}]' + ) + def test_strip_markup_closed(self): text = "before {} after" assert strip_tool_markup(text) == "before after" text = 'before <|tool_call>call:terminal{command:"ls"} after' assert strip_tool_markup(text) == "before after" + def test_strip_named_mistral_call_consumes_trailing_eos(self): + # The named ``[TOOL_CALLS]name{json}`` shape must eat the optional + # trailing ```` like the array shape, so the EOS marker is not left + # behind as visible content. + text = '[TOOL_CALLS]web_search{"query":"cats"}' + assert strip_tool_markup(text) == "" + text = '[TOOL_CALLS]web_search{"query":"cats"} and then' + assert strip_tool_markup(text) == " and then" + def test_strip_markup_unclosed_final(self): text = "before {partial" # final=True drops the trailing run. @@ -214,6 +293,849 @@ class TestParser: == "before " ) + # Mistral [TOOL_CALLS] bracket-tag. + + def test_mistral_bracket_basic(self): + # Devstral / Mistral-Small fallback when bypassing native FC. + text = '[TOOL_CALLS]web_search{"query":"weather"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + assert isinstance(result[0]["function"]["arguments"], str) + assert "weather" in result[0]["function"]["arguments"] + + def test_rehearsal_inside_unclosed_think_is_ignored(self): + """Rehearsal-shaped markup inside an unclosed block must + not be executed as a real tool call. Mid-stream the + tag has not arrived yet, so the strip regex has to accept + end-of-string as a terminator. Regression for the Gemini + high-severity flag on this PR.""" + text = ( + "I should call web_search[ARGS]" '{"query":"weather"} next to find the answer.' + ) + result = parse_tool_calls_from_text(text) + # Inside an unclosed think block no calls are yielded. + assert result == [] + + def test_rehearsal_inside_unclosed_bracket_think_is_ignored(self): + text = "[THINK]planning to use python[ARGS]" '{"code":"print(1)"} but not yet.' + result = parse_tool_calls_from_text(text) + assert result == [] + + def test_rehearsal_after_closed_think_still_parsed(self): + text = "planning" 'python[ARGS]{"code":"print(1)"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + + def test_rehearsal_inside_prefilled_think_is_ignored(self): + """Reasoning models (Qwen3.5 enable_thinking) open in the PROMPT, + so generated content starts inside the thought and carries only a closing + . A call rehearsed in that leading thought must be skipped, while a + real call after the close still fires.""" + text = 'planning web_search[ARGS]{"query":"draft"}python[ARGS]{"code":"print(1)"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + + def test_literal_close_think_in_leading_argument_not_prefill(self): + """A literal inside a real leading call's arguments must not be + read as a prefilled-reasoning close (which would skip the call).""" + text = 'web_search[ARGS]{"query":"what is "}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_stray_close_after_real_call_not_treated_as_prefill(self): + """A real leading call followed by a stray and no further call is + a normal answer, not prefilled reasoning; the call must still fire (the + virtual span only applies when a real call follows the close).""" + text = 'Now web_search[ARGS]{"query":"x"} answer' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_mistral_bracket_with_whitespace(self): + # Optional whitespace (incl. newlines) between the name and the opening brace. + text = '[TOOL_CALLS]python \n {"code":"print(1)"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + assert "print(1)" in result[0]["function"]["arguments"] + + def test_mistral_bracket_nested_json(self): + # Brace-balance scan handles nested objects and braces inside string literals. + text = "[TOOL_CALLS]web_search" '{"query":"a {nested} brace","opts":{"limit":5}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + import json as _json + + args = _json.loads(result[0]["function"]["arguments"]) + assert args["query"] == "a {nested} brace" + assert args["opts"] == {"limit": 5} + + def test_mistral_bracket_with_prose(self): + # Bracket-tag surrounded by prose is still recognised. + text = ( + "Sure, I will look that up.\n" + '[TOOL_CALLS]web_search{"query":"weather"}\n' + "Calling now." + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_mistral_bracket_bad_json_dropped(self): + text = "[TOOL_CALLS]web_search{not valid}" + result = parse_tool_calls_from_text(text) + # No usable tool call; callers fall back to text. + assert result == [] + + def test_mistral_bracket_object_with_array_value(self): + # Args must be a JSON object; a dict wrapping an array value is accepted. + text = '[TOOL_CALLS]web_search{"opts":[1,2,3]}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + # Rehearsal syntax name[ARGS]{json}. + + def test_rehearsal_basic(self): + text = 'python[ARGS]{"code":"print(1)"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + assert "print(1)" in result[0]["function"]["arguments"] + + def test_rehearsal_with_prose(self): + text = "I should call the python tool. Like this: " 'python[ARGS]{"code":"x = 1"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + + def test_rehearsal_bad_json_dropped(self): + text = "python[ARGS]{not valid json}" + result = parse_tool_calls_from_text(text) + assert result == [] + + def test_mistral_bracket_hyphenated_mcp_name(self): + # Dashed MCP names must be captured whole, not truncated at the first dash. + text = '[TOOL_CALLS]mcp__srv__list-issues{"q":"x"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "mcp__srv__list-issues" + + def test_rehearsal_hyphenated_mcp_name(self): + text = 'mcp__srv__list-issues[ARGS]{"q":"x"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "mcp__srv__list-issues" + + def test_streaming_strip_removes_partial_bracket_marker(self): + # A bracket tag streamed before its opening brace must strip on the final pass, not leak. + assert strip_tool_markup("answer [TOOL_CALLS]web_search", final = True) == "answer" + assert strip_tool_markup("text python[ARGS]", final = True) == "text" + # Non-final must keep the in-progress tag buffered (not yet stripped). + partial = "answer [TOOL_CALLS]web_search" + assert strip_tool_markup(partial, final = False) == partial + + def test_strip_removes_two_level_nested_bracket_call_keeps_prose(self): + # Two-level-nested args must be removed whole; the balanced scan handles any depth. + text = 'before [TOOL_CALLS]search{"f":{"g":{"h":1}}} after' + assert strip_tool_markup(text, final = False) == "before after" + assert strip_tool_markup(text, final = True) == "before after" + + def test_strip_removes_call_with_literal_think_in_argument(self): + # A literal think block inside arguments strips with the call, not as a reasoning block. + text = ( + '{"name":"write","arguments":' + '{"text":"compare and tags"}}' + ) + assert strip_tool_markup(text, final = True) == "" + + def test_strip_preserves_real_think_but_strips_call_with_literal_think(self): + text = ( + "planning ok " + '{"name":"w","arguments":{"t":"x"}} done' + ) + out = strip_tool_markup(text, final = True) + assert "planning" in out + assert "" not in out and '"name"' not in out + assert "ok" in out and "done" in out + + def test_prose_mentioning_args_marker_is_not_truncated(self): + # ``foo[ARGS] to the template`` is prose; the catch-all must not delete the sentence. + text = "Please pass foo[ARGS] to the template and continue reading." + assert strip_tool_markup(text, final = True) == text + + def test_streaming_strip_handles_mistral_v11_call_id_args(self): + # The streaming strip uses the regex patterns directly, so they must cover the v11 + # [CALL_ID]/[ARGS] metadata (aligned with the parser). + raw = 'before [TOOL_CALLS]web_search[CALL_ID]abc123[ARGS]{"q":"x"} after' + out = strip_tool_markup_streaming(raw) + assert "[TOOL_CALLS]" not in out and "[CALL_ID]" not in out and "[ARGS]" not in out + assert "before" in out and "after" in out + + # pre-strip. + + def test_think_block_stripped_before_xml(self): + # The think block is stripped before matching so the post-thinking call is recognised. + text = ( + "I will use web_search to find the weather." + '{"name":"web_search","arguments":{"query":"sf"}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_think_block_stripped_before_bracket_tag(self): + text = ( + "Let me search for that.\n" '[TOOL_CALLS]web_search{"query":"weather"}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_uppercase_think_tag_stripped(self): + # Some templates use [THINK]...[/THINK] instead of . + text = "[THINK]planning my next call[/THINK]" '[TOOL_CALLS]python{"code":"print(1)"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + + def test_think_block_hides_inner_tool_call(self): + # A call mentioned inside think is a rehearsal; the wrapper strip removes the inner markup. + text = ( + "I might call " + '{"name":"web_search","arguments":{}} ' + "but I am not sure\n" + "Let me just answer directly." + ) + result = parse_tool_calls_from_text(text) + assert result == [] + + def test_think_literal_inside_real_tool_argument_is_preserved(self): + # A real call whose argument contains a literal think tag must not be corrupted. + text = ( + '{"name":"write","arguments":' + '{"text":"compare and tags"}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"])["text"] == ( + "compare and tags" + ) + + def test_bracket_tag_argument_with_think_literal_is_preserved(self): + text = '[TOOL_CALLS]search{"q":"explain [THINK] blocks"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"])["q"] == "explain [THINK] blocks" + + def test_real_call_after_think_with_rehearsal_inside(self): + # A rehearsal inside is skipped, but the real call after the close tag parses. + text = 'plan: search[ARGS]{"q":"x"}search[ARGS]{"q":"real"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"])["q"] == "real" + + # XML takes precedence over bracket-tag. + + def test_xml_wins_over_bracket(self): + # When a model emits both forms in one message, the XML form is canonical and wins. + text = ( + '{"name":"primary","arguments":{}}' + '[TOOL_CALLS]secondary{"k":"v"}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "primary" + + # Strip patterns include bracket-tag and rehearsal. + + def test_strip_bracket_tag_closed(self): + text = 'before [TOOL_CALLS]web_search{"q":"hi"} after' + assert "[TOOL_CALLS]" not in strip_tool_markup(text) + assert "before" in strip_tool_markup(text) + assert "after" in strip_tool_markup(text) + + def test_strip_rehearsal_closed(self): + text = 'prose python[ARGS]{"code":"x"} more prose' + cleaned = strip_tool_markup(text) + assert "[ARGS]" not in cleaned + assert "prose" in cleaned + assert "more prose" in cleaned + + def test_strip_bracket_tag_unclosed_final(self): + text = 'before [TOOL_CALLS]web_search{"q":"part' + # Final-mode strip drops the trailing unclosed run. + cleaned = strip_tool_markup(text, final = True) + assert "TOOL_CALLS" not in cleaned + assert cleaned == "before" + + # Canonical Mistral array, v11 [CALL_ID], unified multi-call (PR review fixes). + + def test_mistral_canonical_array_is_parsed(self): + # Canonical multi-call array: every call must parse (was dropped then deleted to EOS). + text = '[TOOL_CALLS] [{"name":"a","arguments":{"x":1}},{"name":"b","arguments":{"y":2}}]' + result = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in result] == ["a", "b"] + assert json.loads(result[0]["function"]["arguments"]) == {"x": 1} + assert json.loads(result[1]["function"]["arguments"]) == {"y": 2} + + def test_mistral_array_string_arguments_are_decoded(self): + # OpenAI-spec arguments arrive as a JSON string; decode to an object. + text = '[TOOL_CALLS] [{"name":"a","arguments":"{\\"x\\":1}"}]' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == {"x": 1} + + def test_mistral_array_scalar_string_argument_not_double_encoded(self): + # A bare scalar string argument in the Mistral array form must be kept + # raw, exactly like the path, so the downstream argument + # healer wraps ``weather`` into the single-string tool's key -- not + # ``"weather"`` with literal quotes from a redundant json.dumps. + array = parse_tool_calls_from_text( + '[TOOL_CALLS][{"name":"web_search","arguments":"weather"}]' + ) + xml = parse_tool_calls_from_text( + '{"name":"web_search","arguments":"weather"}' + ) + assert array[0]["function"]["arguments"] == xml[0]["function"]["arguments"] == "weather" + healed = _coerce_arguments( + array[0]["function"]["arguments"], heal = True, tool_name = "web_search" + ) + assert healed == {"query": "weather"} + + def test_mistral_array_strip_keeps_trailing_prose(self): + # The array form must be removed whole, not deleted to end-of-string. + text = 'answer [TOOL_CALLS] [{"name":"a","arguments":{}}] tail' + assert strip_tool_markup(text, final = True) == "answer tail" + + def test_mistral_and_rehearsal_in_one_message_both_parse(self): + # A Mistral call and a rehearsal call together: both must parse. + text = '[TOOL_CALLS]a{"x":1} then b[ARGS]{"y":2}' + result = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in result] == ["a", "b"] + + def test_mistral_v11_call_id_is_not_the_function_name(self): + # v11 shape: the function name is ``name``, never the opaque call-id token. + result = parse_tool_calls_from_text('[TOOL_CALLS]get_weather[CALL_ID]abc123[ARGS]{"q":"x"}') + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + assert json.loads(result[0]["function"]["arguments"]) == {"q": "x"} + # v11 without a call-id parses the same name. + r2 = parse_tool_calls_from_text('[TOOL_CALLS]get_weather[ARGS]{"q":"y"}') + assert r2[0]["function"]["name"] == "get_weather" + + def test_strip_preserves_rehearsal_inside_think(self): + # A rehearsal inside is reasoning; strip keeps it verbatim. + text = 'plan: search[ARGS]{"q":"x"} A' + out = strip_tool_markup(text, final = True) + assert out == text + assert "search[ARGS]" in out + + def test_streaming_strip_preserves_rehearsal_inside_think(self): + # The streaming strip must also preserve a think rehearsal: a mid-stream strip shrinks + # then regrows the cumulative text (corrupts append-by-length consumers). Matches GGUF. + text = 'plan: search[ARGS]{"q":"x"} A' + assert strip_tool_markup_streaming(text) == text + assert strip_tool_markup_streaming(text, tool_protocol_active = True) == text + # An unclosed block during streaming is preserved too (the parser keeps it). + partial = 'plan: search[ARGS]{"q":"x"}' + assert strip_tool_markup_streaming(partial, tool_protocol_active = True) == partial + + def test_streaming_strip_still_removes_real_call_outside_think(self): + # The think guard must not stop the streaming strip removing a call outside the block. + text = 'reason web_search[ARGS]{"q":"x"}' + out = strip_tool_markup_streaming(text, tool_protocol_active = True) + assert "web_search[ARGS]" not in out + assert "reason" in out + + def test_strip_bracket_calls_is_linear(self): + # Many complete bracket calls must strip in ~linear time (was O(n^2) per match). + import time + + text = '[TOOL_CALLS]f{"a":1}' * 4000 # ~80KB, 4000 complete calls + t0 = time.perf_counter() + out = strip_tool_markup(text, final = True) + elapsed = time.perf_counter() - t0 + assert "[TOOL_CALLS]" not in out + assert elapsed < 1.0, f"strip took {elapsed * 1000:.0f}ms on 4000 bracket calls" + + def test_streaming_strip_handles_nested_mistral_json(self): + # The non-greedy [TOOL_CALLS]name{...} pattern truncates nested JSON at the first }; the + # balanced helper must remove the whole call so no trailing brace leaks to the streaming ... + raw = 'ok [TOOL_CALLS]foo{"a":{"b":1}} tail' + out = strip_tool_markup_streaming(raw) + assert "[TOOL_CALLS]" not in out + assert "}" not in out + assert "ok " in out and "tail" in out + + def test_streaming_strip_handles_nested_wrapperless_gemma(self): + # Same class of bug for the wrapper-less Gemma call:NAME{...} form with a + # nested object argument. + raw = "ok call:f{loc:{city:NYC},n:3} tail" + out = strip_tool_markup_streaming(raw) + assert "call:f" not in out + assert "}" not in out + assert "ok " in out and "tail" in out + + def test_streaming_strip_keeps_prose_after_function_xml_with_literal_marker(self): + # A literal ```` in a value is data: the strip must close at the REAL + # ```` and keep trailing prose (the open-ended regex ate to EOF). + raw = ( + "pref " + 'print("") tail' + ) + assert strip_tool_markup_streaming(raw) == "pref tail" + # Streaming and final strip agree on the visible text (final also trims). + assert strip_tool_markup_streaming(raw) == strip_tool_markup(raw, final = True) + + def test_streaming_strip_drops_leading_magistral_reasoning(self): + # Magistral emits reasoning as a leading ``[THINK]...[/THINK]`` bracket block + # (not the ```` the reasoning channel renders). The streaming display + # strip must drop it so the raw chain-of-thought does not leak into the + # safetensors content; GGUF routes it to reasoning_content natively. + closed = "[THINK]Let me think. 2+2 is 4.[/THINK]The answer is 4." + assert strip_tool_markup_streaming(closed) == "The answer is 4." + assert strip_tool_markup_streaming(closed) == strip_tool_markup(closed, final = True) + # Unclosed mid-stream reasoning is held from the marker on (nothing leaks, and + # the cleaned text only grows as the answer streams in after ``[/THINK]``). + assert strip_tool_markup_streaming("[THINK]still thinking") == "" + assert strip_tool_markup_streaming("[THINK]r[/THINK]The") == "The" + assert strip_tool_markup_streaming("[THINK]r[/THINK]The answer") == "The answer" + # A non-leading ``[THINK]`` is ordinary prose and is left untouched. + assert strip_tool_markup_streaming("hi [THINK] later") == "hi [THINK] later" + + +class TestParserMultiFormat: + """Shared-parser coverage: every family's emission maps to the same OpenAI shape.""" + + # Llama-3 + + def test_llama3_python_tag_dot_call(self): + # Llama-3 built-in tools: <|python_tag|>NAME.call(k="v", ...). + import json + + text = '<|python_tag|>brave_search.call(query="weather in Tokyo")' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "brave_search" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "weather in Tokyo"} + + def test_llama3_python_tag_dot_call_multi_arg(self): + import json + + text = "<|python_tag|>get_weather.call(" 'location="Tokyo", units="celsius", days=5)' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"location": "Tokyo", "units": "celsius", "days": 5} + + def test_llama3_python_tag_json_form(self): + import json + + text = '<|python_tag|>{"name":"web_search","parameters":{"query":"hi","n":5}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "hi", "n": 5} + + def test_llama3_python_tag_json_form_with_eom(self): + # Llama-3 emits ``<|eom_id|>`` after the JSON; must not break parsing. + import json + + text = '<|python_tag|>{"name":"python","parameters":{"code":"print(2+2)"}}<|eom_id|>' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"code": "print(2+2)"} + + def test_llama3_strip_markup_final(self): + text = '<|python_tag|>brave_search.call(query="x")' + assert strip_tool_markup(text, final = True) == "" + + def test_llama3_python_tag_json_form_non_scalar_args_skipped(self): + # Should NOT fabricate ``{"value": args}`` when the JSON form + # has a non-dict / non-string ``arguments`` value. + for bad in ( + '<|python_tag|>{"name":"foo","arguments":42}', + '<|python_tag|>{"name":"foo","arguments":[1,2,3]}', + '<|python_tag|>{"name":"foo","arguments":null}', + '<|python_tag|>{"name":"foo","arguments":true}', + ): + assert parse_tool_calls_from_text(bad) == [], bad + + # ── Llama-3.2 bare JSON ``custom_tools`` ───────────────────── + + def test_llama3_2_bare_json_parameters(self): + # Llama-3.2-Instruct emits bare JSON directly as content; no + # <|python_tag|> prefix per its training template. + import json + + text = '{"name":"web_search","parameters":{"query":"Tokyo weather"}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "Tokyo weather"} + + def test_llama3_2_bare_json_arguments_key(self): + import json + + text = '{"name":"add","arguments":{"a":1,"b":2}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"a": 1, "b": 2} + + def test_llama3_2_bare_json_multi_call(self): + # Llama-3 may chain calls with ``; `` per training template. + text = '{"name":"a","parameters":{}}; {"name":"b","parameters":{}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "a" + assert result[1]["function"]["name"] == "b" + + def test_llama3_2_bare_json_with_eom_sentinel(self): + text = '{"name":"x","parameters":{"y":1}}<|eom_id|>' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "x" + + def test_llama3_2_bare_json_leading_sentinel_skipped(self): + # Sometimes prior <|eot_id|> leaks into the next turn. + text = '<|eot_id|>{"name":"x","parameters":{}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "x" + + def test_llama3_2_bare_json_plain_prose_does_not_fire(self): + # Defensive: must NOT fire on plain assistant prose. + text = "Hello world, how are you today?" + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_embedded_in_prose_does_not_fire(self): + # Defensive: JSON embedded in prose must NOT fire (parser is + # strict about content STARTING with `{`). + text = 'The tool result was: {"name":"foo"}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_missing_name_does_not_fire(self): + text = '{"result":"ok","data":[1,2,3]}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_missing_args_does_not_fire(self): + text = '{"name":"x"}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_args_not_dict_does_not_fire(self): + text = '{"name":"x","parameters":42}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_string_parameters_does_not_fire(self): + # Llama-3 spec: parameters must be a dict. Prose like + # ``{"name":"foo","parameters":"a sentence"}`` must NOT trigger. + text = '{"name":"foo","parameters":"this is a sentence"}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_string_arguments_not_json_does_not_fire(self): + # OpenAI ``arguments`` may be a JSON-string of a dict, but a + # plain non-JSON string must not pass the guard. + text = '{"name":"foo","arguments":"not json"}' + assert parse_tool_calls_from_text(text) == [] + + def test_llama3_2_bare_json_string_arguments_json_dict_fires(self): + # OpenAI shape: arguments is a JSON-encoded string of a dict. + text = '{"name":"foo","arguments":"{\\"q\\":\\"x\\"}"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "foo" + # arguments stays as the original JSON-string. + assert result[0]["function"]["arguments"] == '{"q":"x"}' + + def test_llama3_2_bare_json_string_arguments_json_non_dict_does_not_fire(self): + # JSON-string that parses to a list / scalar / null must NOT fire. + for bad in ( + '{"name":"foo","arguments":"[1,2,3]"}', + '{"name":"foo","arguments":"\\"plain\\""}', + '{"name":"foo","arguments":"null"}', + '{"name":"foo","arguments":"42"}', + ): + assert parse_tool_calls_from_text(bad) == [], bad + + # Mistral pre-v11 + + def test_mistral_pre_v11_array(self): + import json + + text = '[TOOL_CALLS] [{"name":"web_search","arguments":{"query":"hello"},"id":"abc"}]' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + # Mistral provides its own id; preserve it. + assert result[0]["id"] == "abc" + assert json.loads(result[0]["function"]["arguments"]) == {"query": "hello"} + + def test_mistral_array_parameters_key_alias(self): + import json + + # Array object keyed on ``parameters`` (not ``arguments``) must keep its + # payload, matching the JSON/XML paths and SGLang's base detector. + text = '[TOOL_CALLS] [{"name":"get_weather","parameters":{"city":"Paris"}}]' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Paris"} + + def test_mistral_pre_v11_array_multi(self): + text = ( + '[TOOL_CALLS] [{"name":"a","arguments":{"x":1},"id":"id1"},' + '{"name":"b","arguments":{"y":2},"id":"id2"}]' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "a" + assert result[1]["function"]["name"] == "b" + + def test_mistral_pre_v11_unclosed_array(self): + # Closing ``]`` truncated -- parser must heal off individual objects. + text = '[TOOL_CALLS] [{"name":"web_search","arguments":{"q":"x"},"id":"id"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + # Mistral v11+ + + def test_mistral_v11_single(self): + # Magistral / Mistral Small 3.1: bare ``name{json}`` after trigger. + import json + + text = '[TOOL_CALLS]add{"a":3.5,"b":4}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "add" + assert json.loads(result[0]["function"]["arguments"]) == {"a": 3.5, "b": 4} + + def test_mistral_v11_parallel(self): + # v11+ parallel: ``[TOOL_CALLS]a{...}[TOOL_CALLS]b{...}``. + text = '[TOOL_CALLS]add{"a":1}[TOOL_CALLS]sub{"b":2}' + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "add" + assert result[1]["function"]["name"] == "sub" + + def test_mistral_v11_with_args_marker(self): + # Ministral / Mistral Large 3: ``[TOOL_CALLS]name[ARGS]{json}``. + import json + + text = '[TOOL_CALLS]add[ARGS]{"a":1,"b":2}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "add" + assert json.loads(result[0]["function"]["arguments"]) == {"a": 1, "b": 2} + + def test_mistral_strip_markup_v11(self): + text = '[TOOL_CALLS]add{"a":1}' + assert strip_tool_markup(text, final = True) == "" + + def test_mistral_call_id_form(self): + # Mistral Small 3.2: ``[TOOL_CALLS]name[CALL_ID][ARGS]{json}``. + # The ``[CALL_ID]`` segment must be skipped, not treated as a stop + # (llama.cpp test-chat.cpp:4785 parses this to one call). + import json + + text = '[TOOL_CALLS]special_function[CALL_ID]123456789[ARGS]{"arg1": 1}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "special_function" + assert json.loads(result[0]["function"]["arguments"]) == {"arg1": 1} + + def test_mistral_call_id_form_parallel(self): + text = ( + '[TOOL_CALLS]special_function[CALL_ID]000000001[ARGS]{"arg1": 1}' + "[TOOL_CALLS]special_function_with_opt[CALL_ID]000000002" + '[ARGS]{"arg1": 1, "arg2": 2}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "special_function" + assert result[1]["function"]["name"] == "special_function_with_opt" + + def test_mistral_call_id_form_stripped(self): + text = '[TOOL_CALLS]special_function[CALL_ID]123456789[ARGS]{"arg1": 1}' + assert strip_tool_markup(text, final = True) == "" + + def test_mistral_think_reasoning_ignored(self): + # Magistral wraps reasoning in ``[THINK]...[/THINK]``. A ``[TOOL_CALLS]`` + # inside the reasoning is chain-of-thought, not a real call; only the + # call after ``[/THINK]`` counts (llama.cpp test-chat.cpp:2285). + import json + + text = ( + '[THINK]Let me think about [TOOL_CALLS]fake[ARGS]{"x":1} ' + 'and more[/THINK][TOOL_CALLS]real_fn[ARGS]{"y":2}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "real_fn" + assert json.loads(result[0]["function"]["arguments"]) == {"y": 2} + + def test_mistral_think_reasoning_no_real_call(self): + # Reasoning that merely mentions a tool call but does not emit one + # after ``[/THINK]`` yields no calls. + text = '[THINK]I might call [TOOL_CALLS]fake[ARGS]{"x":1}[/THINK]Done.' + assert parse_tool_calls_from_text(text) == [] + + def test_mistral_think_literal_in_argument_preserved(self): + # A literal ``[THINK]`` inside a real tool argument (after the call) + # must not be stripped or corrupt the parse. + import json + + text = '[TOOL_CALLS]search[ARGS]{"q":"explain the [THINK] token"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == {"q": "explain the [THINK] token"} + + # Gemma 4 + + def test_gemma4_simple_call(self): + import json + + text = ( + "<|tool_call>call:get_weather{" + 'location:<|"|>Tokyo<|"|>,units:<|"|>celsius<|"|>}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"location": "Tokyo", "units": "celsius"} + + def test_gemma4_with_primitives(self): + import json + + text = ( + "<|tool_call>call:set_pref{" + "enabled:true,attempts:5,threshold:1.5,nickname:null}" + ) + result = parse_tool_calls_from_text(text) + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"enabled": True, "attempts": 5, "threshold": 1.5, "nickname": None} + + def test_gemma4_nested_args(self): + # Gemma 4 nests dicts / lists with bare keys and ``<|"|>`` strings. + import json + + text = ( + "<|tool_call>call:search{" + 'query:<|"|>foo<|"|>,filters:{site:<|"|>example.com<|"|>,recent:true},' + 'tags:[<|"|>a<|"|>,<|"|>b<|"|>]}' + ) + result = parse_tool_calls_from_text(text) + args = json.loads(result[0]["function"]["arguments"]) + assert args["query"] == "foo" + assert args["filters"] == {"site": "example.com", "recent": True} + assert args["tags"] == ["a", "b"] + + def test_gemma4_multi_call(self): + text = "<|tool_call>call:a{x:1}<|tool_call>call:b{y:2}" + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "a" + assert result[1]["function"]["name"] == "b" + + def test_gemma4_unclosed_does_not_raise(self): + # Truncated mid-stream; must not raise. + text = '<|tool_call>call:foo{x:<|"|>bar<|"|>' + result = parse_tool_calls_from_text(text) + assert isinstance(result, list) + + def test_gemma4_strip_markup_final(self): + text = "<|tool_call>call:foo{x:1}" + assert strip_tool_markup(text, final = True) == "" + + # ── Gemma 4 wrapper-less (skip_special_tokens stripped) ─────────── + + def test_gemma4_bare_stripped_call(self): + # skip_special_tokens removes <|tool_call>/ and <|"|>, + # leaving a bare call:NAME{...} with an unquoted value. + import json + + text = "call:web_search{query:weather in San Francisco right now}" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "weather in San Francisco right now"} + + def test_gemma4_bare_code_with_commas(self): + # A code value with commas must not truncate at the first comma. + import json + + text = ( + "call:python{code:def f(n):\n a, b = 0, 1\n" + " for _ in range(2, n+1):\n a, b = b, a + b\n" + " return b\n\nprint(f(30))}" + ) + result = parse_tool_calls_from_text(text) + assert result[0]["function"]["name"] == "python" + code = json.loads(result[0]["function"]["arguments"])["code"] + assert "a, b = 0, 1" in code and "print(f(30))" in code + + def test_gemma4_bare_quotes_normalized(self): + # The same value quoted vs unquoted must parse identically so the + # agentic loop can collapse a looping model's repeated calls. + import json + + a = parse_tool_calls_from_text('call:web_search{query:"foo bar"}') + b = parse_tool_calls_from_text("call:web_search{query:foo bar}") + assert json.loads(a[0]["function"]["arguments"]) == {"query": "foo bar"} + assert json.loads(a[0]["function"]["arguments"]) == json.loads( + b[0]["function"]["arguments"] + ) + + def test_gemma4_bare_multi_arg(self): + import json + + text = "call:web_search{query:pytorch latest, url:https://pytorch.org}" + result = parse_tool_calls_from_text(text) + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"query": "pytorch latest", "url": "https://pytorch.org"} + + def test_gemma4_bare_not_matched_in_prose(self): + # A word ending in "call:" must not trigger a bare tool call. + text = "I will recall:that the function{ } is helpful." + result = parse_tool_calls_from_text(text) + assert result == [] + + def test_gemma4_bare_strip_markup_final(self): + text = "Here you go: call:web_search{query:weather today}" + assert "call:web_search" not in strip_tool_markup(text, final = True) + + # ── Cross-format sentinels ──────────────────────────────────── + + def test_all_markers_in_tool_xml_signals(self): + # Streaming buffer wakes up on every emission marker. + from core.inference.tool_call_parser import TOOL_XML_SIGNALS + for marker in ( + "", + "", + "[TOOL_CALLS]", + "<|tool_call>", + ): + assert marker in TOOL_XML_SIGNALS, f"streaming loop would not wake on {marker!r}" + + def test_has_tool_signal_for_all_formats(self): + assert has_tool_signal('<|python_tag|>brave_search.call(q="x")') + assert has_tool_signal('[TOOL_CALLS] [{"name":"x"}]') + assert has_tool_signal('[TOOL_CALLS]add{"a":1}') + assert has_tool_signal("<|tool_call>call:foo{}") + # ──────────────────────────────────────────────────────────────────── # run_safetensors_tool_loop @@ -312,6 +1234,553 @@ def _make_loop( ), exec_fn +class TestParserDeepSeek: + """DeepSeek R1 / V3 / V3.1 coverage. Markers use full-width pipes + (U+FF5C) and lower-one-eighth-block (U+2581). R1 wraps args in a + Markdown ``` ```json ``` ``` fence; V3 / V3.1 emit bare JSON.""" + + def test_r1_simple_call_with_code_fence(self): + import json as _json + + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>function" + "<｜tool▁sep｜>special_function\n" + "```json\n" + '{"arg1": 1}\n' + "```" + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "special_function" + assert _json.loads(result[0]["function"]["arguments"]) == {"arg1": 1} + + def test_r1_short_form_outer_marker(self): + # llama.cpp accepts ``<｜tool▁calls｜>`` as the short-form opener. + import json as _json + + text = ( + "<｜tool▁calls｜>function" + "<｜tool▁sep｜>get_time\n" + "```json\n" + '{"city": "Paris"}\n' + "```" + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_time" + + def test_v3_1_bare_json(self): + # V3 / V3.1 omit the ``function`` prefix and the code fence. + import json as _json + + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city": "Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_time" + assert _json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + def test_v3_1_multi_call_shares_envelope(self): + # Parallel calls share one outer envelope; each inner call has + # its own ``<｜tool▁call▁begin｜>...<｜tool▁call▁end｜>``. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city": "Paris"}' + "<｜tool▁call▁end｜>" + "<｜tool▁call▁begin｜>get_weather" + "<｜tool▁sep｜>" + '{"city": "Paris"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "get_time" + assert result[1]["function"]["name"] == "get_weather" + + def test_v3_1_with_reasoning(self): + # Reasoning ... precedes the tool block. + text = ( + "I'm thinking\n" + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city": "Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_time" + + def test_v3_1_strict_rejects_unclosed_envelope(self): + # Envelope truncated mid-stream (no <｜tool▁calls▁end｜>): healed by + # default, rejected with Auto-Heal off. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city": "Tokyo"}' + ) + assert len(parse_tool_calls_from_text(text)) == 1 + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + + def test_v3_1_multi_call_recovers_when_first_end_marker_missing(self): + # First inner call omits its <｜tool▁call▁end｜>; the second must still be parsed. + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city": "Paris"}' + "<｜tool▁call▁begin｜>get_weather" + "<｜tool▁sep｜>" + '{"city": "Paris"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + result = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in result] == ["get_time", "get_weather"] + + def test_v3_1_strict_recovers_after_missing_call_end(self): + # Strict mode (Auto-Heal off): the FIRST inner call is missing its <｜tool▁call▁end｜> + # terminator, so it is skipped -- but the parser must keep scanning and still return the ... + text = ( + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_weather" + "<｜tool▁sep｜>" + '{"city": "SF"}' + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"tz": "PST"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + # Auto-Heal keeps both; strict skips the truncated first, keeps the second. + assert [c["function"]["name"] for c in parse_tool_calls_from_text(text)] == [ + "get_weather", + "get_time", + ] + strict = parse_tool_calls_from_text(text, allow_incomplete = False) + assert [c["function"]["name"] for c in strict] == ["get_time"] + + def test_r1_strict_recovers_after_missing_close_fence(self): + # R1 form. + text = ( + "<｜tool▁calls▁begin｜>" + "function<｜tool▁sep｜>get_weather\n```json\n" + '{"city": "SF"}' + "function<｜tool▁sep｜>get_time\n```json\n" + '{"tz": "PST"}' + "\n```<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + ) + strict = parse_tool_calls_from_text(text, allow_incomplete = False) + assert [c["function"]["name"] for c in strict] == ["get_time"] + + def test_deepseek_strip_markup(self): + text = ( + "before " + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>foo" + "<｜tool▁sep｜>" + "{}" + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + " after" + ) + assert strip_tool_markup(text, final = True) == "before after" + + def test_deepseek_signal_wakes_streaming(self): + # The streaming buffer state machine must wake on the DeepSeek opener so the rest of the + # section is drained instead of leaked. + text = "<｜tool▁calls▁begin｜>..." + assert has_tool_signal(text) + + def test_deepseek_short_opener_is_stripped(self): + # The short ``<｜tool▁calls｜>`` opener is parsed, so its markup must also be stripped (the + # strip patterns used to require ...calls_begin and left the short-opener markup leaking to ... + text = ( + "before " + "<｜tool▁calls｜>" + "<｜tool▁call▁begin｜>foo" + "<｜tool▁sep｜>" + "{}" + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>" + " after" + ) + assert strip_tool_markup(text, final = True) == "before after" + + +class TestParserGLM: + """GLM 4.5 / 4.6 / 4.7 coverage. Marker collides with Qwen's + ```` but the body shape is XML kv pairs instead of JSON, + so the dispatch order keeps both formats working.""" + + def test_glm_simple_call(self): + import json as _json + + text = ( + "web_search\n" + "query\n" + "weather Tokyo\n" + "" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + args = _json.loads(result[0]["function"]["arguments"]) + # Strings come through raw; the parser does not double-quote. + assert args == {"query": "weather Tokyo"} + + def test_glm_mixed_types_decode_correctly(self): + # Per the chat_template.jinja, strings are emitted raw and non-strings are JSON-encoded. + import json as _json + + text = ( + "complex_function\n" + "name\nJohn Doe\n" + "age\n30\n" + "active\ntrue\n" + "score\n95.5\n" + "" + ) + result = parse_tool_calls_from_text(text) + args = _json.loads(result[0]["function"]["arguments"]) + assert args == {"name": "John Doe", "age": 30, "active": True, "score": 95.5} + + def test_glm_multi_call_back_to_back(self): + # GLM emits parallel calls as consecutive ``... + # `` blocks with no outer envelope. + text = ( + "a\nx\n1\n" + "b\ny\n2\n" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "a" + assert result[1]["function"]["name"] == "b" + + def test_glm_unclosed_tool_call_does_not_lose_value(self): + # Truncated mid-stream (no ) -- the parser must + # still surface what it found rather than dropping the call. + text = "web_search\nquery\npartial" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_glm_does_not_break_qwen_path(self): + # Real Qwen emission must still be parsed by the Qwen branch, + # not silently misrouted to GLM (the marker is shared). + text = '{"name":"web_search","arguments":{"q":"x"}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_glm_strip_markup(self): + text = ( + "before " + "a\nx\n1\n" + " after" + ) + assert strip_tool_markup(text, final = True) == "before after" + + def test_glm_zero_arg_inline_call(self): + # GLM 4.7 emits a no-argument call inline as ``name`` (name followed + # straight by the close tag, no \n / ). + import json as _json + + text = "get_current_date" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_current_date" + assert _json.loads(result[0]["function"]["arguments"]) == {} + + def test_glm_zero_arg_call_in_parallel_batch(self): + # A no-arg call alongside a normal one must not make either vanish. + text = ( + "get_current_date" + "get_weather\ncity\n" + "Tokyo" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "get_current_date" + assert result[1]["function"]["name"] == "get_weather" + + def test_glm_string_value_whitespace_preserved(self): + # The template emits string args verbatim, so significant leading / trailing whitespace + # (code, diffs) must survive. + import json as _json + + text = ( + "run\ncode\n" + " indented code " + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = _json.loads(result[0]["function"]["arguments"]) + assert args == {"code": " indented code "} + + +class TestParserKimi: + """Kimi K2 / Moonshot coverage. ASCII pipes only (NOT full-width). + Name arrives as ``functions.NAME:IDX``; the parser strips the + prefix and the index to recover the bare callable name while + preserving the full id for round-trip rendering.""" + + def test_kimi_simple_call(self): + import json as _json + + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.special_function:0" + "<|tool_call_argument_begin|>" + '{"arg1": 1}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + # Bare name recovered; full id preserved verbatim. + assert result[0]["function"]["name"] == "special_function" + assert result[0]["id"] == "functions.special_function:0" + assert _json.loads(result[0]["function"]["arguments"]) == {"arg1": 1} + + def test_outer_tool_call_with_embedded_kimi_marker_parses_outer(self): + # A Qwen/Hermes whose argument contains literal Kimi markup (a user asking + # about that syntax) must execute the OUTER call, not the embedded marker via the ... + text = ( + '{"name":"web_search","arguments":{"query":' + '"explain <|tool_call_begin|>functions.evil:0' + '<|tool_call_argument_begin|>{}<|tool_call_end|>"}}' + "" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_genuine_kimi_call_without_envelope_still_parses(self): + # Control: a real Kimi call with no leading envelope must + # still go through the pre-pass. + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"query":"x"}<|tool_call_end|>' + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_kimi_multi_call_with_index(self): + # Multiple consecutive calls inside a single section, each + # with its own monotonically incrementing ``:IDX``. + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.read_file:0" + "<|tool_call_argument_begin|>" + '{"path":"a"}' + "<|tool_call_end|>" + "<|tool_call_begin|>functions.web_search:1" + "<|tool_call_argument_begin|>" + '{"query":"x"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 2 + assert result[0]["function"]["name"] == "read_file" + assert result[0]["id"].endswith(":0") + assert result[1]["function"]["name"] == "web_search" + assert result[1]["id"].endswith(":1") + + def test_kimi_dotted_name_keeps_full_dotted_name(self): + # A dotted Kimi id keeps its FULL name after stripping only the ``functions.`` prefix and + # ``:idx`` suffix -- matching current vLLM ... + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>a.b.c:2" + "<|tool_call_argument_begin|>" + "{}" + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "a.b.c" + + def test_kimi_dotted_mcp_name_with_functions_prefix(self): + # ``functions.mcp.server-list:0`` must resolve to ``mcp.server-list`` + # (only the ``functions.`` prefix and ``:idx`` are removed). + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.mcp.server-list:0" + "<|tool_call_argument_begin|>" + "{}" + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "mcp.server-list" + + def test_kimi_multi_call_recovers_when_first_end_marker_missing(self): + # First call omits its <|tool_call_end|>; the second must still parse. + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.read_file:0" + "<|tool_call_argument_begin|>" + '{"path":"a"}' + "<|tool_call_begin|>functions.web_search:1" + "<|tool_call_argument_begin|>" + '{"query":"x"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in result] == ["read_file", "web_search"] + + def test_kimi_handles_unclosed_section(self): + # End marker missing -- the parser must still extract the call. + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.foo:0" + "<|tool_call_argument_begin|>" + '{"a":1}' + "<|tool_call_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "foo" + + def test_kimi_strip_markup(self): + text = ( + "before " + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.x:0" + "<|tool_call_argument_begin|>" + "{}" + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + " after" + ) + assert strip_tool_markup(text, final = True) == "before after" + + def test_kimi_signal_wakes_streaming(self): + text = "<|tool_calls_section_begin|>..." + assert has_tool_signal(text) + + def test_kimi_call_without_section_wrapper(self): + # llama.cpp makes the ``<|tool_calls_section_begin|>`` wrapper optional -- Kimi K2 can emit + # a bare ``<|tool_call_begin|>`` call. + import json as _json + + text = ( + "<|tool_call_begin|>functions.execute_command:0" + "<|tool_call_argument_begin|>" + '{"cmd":"ls"}' + "<|tool_call_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "execute_command" + assert _json.loads(result[0]["function"]["arguments"]) == {"cmd": "ls"} + + def test_kimi_malformed_json_recovers_later_calls(self): + # A call with malformed / truncated JSON must not drop the valid calls that follow it in + # the same section (the bad call is skipped, the good one is recovered). + import json as _json + + text = ( + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.a:0" + '<|tool_call_argument_begin|>{"city":"Beijing"' # missing closing brace + "<|tool_call_end|>" + "<|tool_call_begin|>functions.b:1" + '<|tool_call_argument_begin|>{"city":"Shanghai"}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "b" + assert _json.loads(result[0]["function"]["arguments"]) == {"city": "Shanghai"} + + +class TestParserCrossFormatRouting: + """Ensure the per-format dispatch order doesn't misroute any + family. Real emissions for each new family + every old family + must still parse correctly when intermixed.""" + + def test_dispatch_routes_each_family_correctly(self): + cases = [ + ( + "Qwen", + '{"name":"a","arguments":{"x":1}}', + "a", + ), + ( + "DeepSeek V3.1", + "<｜tool▁calls▁begin｜>" + "<｜tool▁call▁begin｜>get_time" + "<｜tool▁sep｜>" + '{"city":"Tokyo"}' + "<｜tool▁call▁end｜>" + "<｜tool▁calls▁end｜>", + "get_time", + ), + ( + "GLM", + "web_search\n" + "q\nx\n" + "", + "web_search", + ), + ( + "Kimi", + "<|tool_calls_section_begin|>" + "<|tool_call_begin|>functions.add:0" + "<|tool_call_argument_begin|>" + '{"a":1}' + "<|tool_call_end|>" + "<|tool_calls_section_end|>", + "add", + ), + ] + for label, text, expected_name in cases: + result = parse_tool_calls_from_text(text) + assert len(result) == 1, f"{label}: parser missed the call" + assert result[0]["function"]["name"] == expected_name, ( + f"{label}: got {result[0]['function']['name']!r}, " f"expected {expected_name!r}" + ) + + def test_all_new_markers_in_tool_xml_signals(self): + # The safetensors / MLX streaming buffer must wake on every supported emission marker -- + # otherwise the BUFFERING state leaks tool content to the user before parse. + from core.inference.tool_call_parser import TOOL_XML_SIGNALS + for marker in ( + "<｜tool▁calls▁begin｜>", + "<｜tool▁call▁begin｜>", + "<|tool_calls_section_begin|>", + "<|tool_call_begin|>", + ): + assert marker in TOOL_XML_SIGNALS, f"streaming loop would not wake on {marker!r}" + + def test_active_tools_are_passed_to_single_turn_after_render_html_success(): captured_tool_names: list[list[str]] = [] exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) @@ -347,6 +1816,449 @@ def test_active_tools_are_passed_to_single_turn_after_render_html_success(): assert any(event.get("type") == "content" and event.get("text") == "Done." for event in events) +def test_spent_one_shot_rehearsal_repeat_is_detected_not_blank_continuation(): + # A spent one-shot (render_html) stays in the ORIGINAL tool list; detection is gated on + # that list (matching the strip gate) so a re-emitted repeat is drained and routed to the + # repeat no-op instead of stripped into a blank continuation. + exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) + turns = iter( + [ + [ + '{"name":"render_html","arguments":{"code":"one"}}' + ], + ['render_html[ARGS]{"code":"two"}'], # spent one-shot rehearsal + ["The chart is above."], + ] + ) + + def gen(_messages, *, active_tools = None): + try: + chunks = next(turns) + except StopIteration: + return + acc = "" + for c in chunks: + acc += c + yield acc + + events = _collect_events( + run_safetensors_tool_loop( + single_turn = gen, + messages = [{"role": "user", "content": "make a chart"}], + tools = [ + {"type": "function", "function": {"name": "render_html"}}, + {"type": "function", "function": {"name": "web_search"}}, + ], + execute_tool = exec_fn, + max_tool_iterations = 5, + ) + ) + contents = [e["text"] for e in events if e["type"] == "content"] + # render_html ran exactly once; the repeat was a no-op, not a second execution. + assert exec_fn.calls == [("render_html", {"code": "one"})], exec_fn.calls + # The loop continued past the repeat to the real answer (not a blank continuation). + assert any("The chart is above." in t for t in contents), contents + # The raw rehearsal markup never leaked as visible content. + assert not any("render_html[ARGS]" in t for t in contents), contents + + +def test_rehearsal_call_name_is_not_streamed_before_args(): + # A rehearsal whose name and [ARGS] arrive together must drain, not stream the bare name. + loop, exec_fn = _make_loop( + turns = [['web_search[ARGS]{"query":"cats"}'], ["Found."]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search" in t for t in contents), contents + + +def test_rehearsal_call_name_split_before_args_is_not_streamed(): + # Finding 5: name and [ARGS] in separate chunks -- the bare name is held until [ARGS] arrives. + loop, exec_fn = _make_loop( + turns = [["web_search", '[ARGS]{"query":"cats"}'], ["Found."]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search" in t for t in contents), contents + + +def test_plain_word_matching_no_tool_still_streams(): + # The prefix guard must not swallow prose: a non-tool bare word streams. + loop, _exec = _make_loop( + turns = [["weather", " is nice today."]], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert "weather is nice today." in contents, contents + + +def test_rehearsal_name_after_prose_in_streaming_is_not_streamed(): + # After prose has streamed (STREAMING state), a split rehearsal name must still be held. + loop, exec_fn = _make_loop( + turns = [ + # _make_loop accumulates these deltas into cumulative snapshots. + ["Let me think. ", "I will search ", "web_search", '[ARGS]{"query":"cats"}'], + ["Found."], + ], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search" in t for t in contents), contents + + +def test_rehearsal_name_after_prose_same_chunk_in_streaming_is_not_streamed(): + # Prose then ``web_search[ARGS]{...}`` in one chunk: the boundary is pulled back over the name. + loop, exec_fn = _make_loop( + turns = [ + ["Sure. ", 'now web_search[ARGS]{"query":"cats"}'], + ["Found."], + ], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search" in t for t in contents), contents + + +def test_initial_buffer_flush_holds_split_rehearsal_name(): + # First flush out of BUFFERING applies the same trailing-name hold as STREAMING. + loop, exec_fn = _make_loop( + turns = [["I will use python", '[ARGS]{"code":"print(1)"}'], ["done"]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("python", {"code": "print(1)"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("python" in t for t in contents), contents + + +def test_think_rehearsal_streams_monotonically_and_keeps_reasoning(): + # A think rehearsal streams the same text the final strip keeps: cumulative content is + # monotonically non-decreasing and ends with the markup intact. + loop, exec_fn = _make_loop( + turns = [["plan ", 'search[ARGS]{"q":"x"}', " visible"]], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + contents = [e["text"] for e in events if e["type"] == "content"] + assert exec_fn.calls == [], exec_fn.calls + assert all(len(b) >= len(a) for a, b in zip(contents, contents[1:])), contents + final = contents[-1] if contents else "" + assert 'search[ARGS]{"q":"x"}' in final, contents + assert "visible" in final, contents + + +def test_plain_answer_ending_with_tool_name_word_is_preserved(): + # End-of-stream flush: a plain answer ending on a tool-name word is prose, not dropped. + loop, exec_fn = _make_loop( + turns = [["I think ", "you should ", "web_search"]], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + assert exec_fn.calls == [], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert any(t.rstrip().endswith("web_search") for t in contents), contents + + +def test_long_tool_name_split_rehearsal_is_not_capped_and_executes(): + # Finding 10/11: an MCP name longer than the buffer cap, split before [ARGS], is still + # held (self-bounding prefix); no leak and the call executes. + from core.inference.safetensors_agentic import _MAX_BUFFER_CHARS + + name = "mcp__github__create_pull_request" + assert len(name) >= _MAX_BUFFER_CHARS, len(name) + exec_fn = FakeExecuteTool(["RESULT"]) + _turns = iter([[name, name + '[ARGS]{"x":1}'], ["done"]]) + + def st(_messages, active_tools = None): + yield from next(_turns) + + events = _collect_events( + run_safetensors_tool_loop( + single_turn = st, + messages = [{"role": "user", "content": "go"}], + tools = [{"type": "function", "function": {"name": name}}], + execute_tool = exec_fn, + max_tool_iterations = 2, + ) + ) + assert exec_fn.calls == [(name, {"x": 1})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any(name in t for t in contents), contents + + +def test_unrestricted_mode_split_rehearsal_name_is_not_streamed(): + # Finding 6: unrestricted mode treats any bare identifier as a possible rehearsal NAME. + exec_fn = FakeExecuteTool(["RESULT"]) + _turns = iter([["web_search", 'web_search[ARGS]{"q":"x"}'], ["done"]]) + + def st(_messages, active_tools = None): + yield from next(_turns) + + events = _collect_events( + run_safetensors_tool_loop( + single_turn = st, + messages = [{"role": "user", "content": "go"}], + tools = [], # unrestricted + execute_tool = exec_fn, + max_tool_iterations = 2, + ) + ) + assert exec_fn.calls == [("web_search", {"q": "x"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search" in t for t in contents), contents + + +def test_unrestricted_mode_split_after_bracket_is_not_streamed(): + # Unrestricted mode: a chunk split right after ``NAME[`` is still held (parity with the + # restricted-mode startswith hold). + exec_fn = FakeExecuteTool(["RESULT"]) + _turns = iter([["web_search[", 'web_search[ARGS]{"q":"x"}'], ["done"]]) + + def st(_messages, active_tools = None): + yield from next(_turns) + + events = _collect_events( + run_safetensors_tool_loop( + single_turn = st, + messages = [{"role": "user", "content": "go"}], + tools = [], # unrestricted + execute_tool = exec_fn, + max_tool_iterations = 2, + ) + ) + assert exec_fn.calls == [("web_search", {"q": "x"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search[" in t for t in contents), contents + + +def test_unrestricted_mode_plain_prose_still_streams(): + # The unrestricted hold releases a held identifier once the rest of the sentence follows. + def st(_messages, active_tools = None): + for snap in ("Hello", "Hello there friend."): + yield snap + + events = _collect_events( + run_safetensors_tool_loop( + single_turn = st, + messages = [{"role": "user", "content": "hi"}], + tools = [], + execute_tool = FakeExecuteTool([]), + max_tool_iterations = 1, + ) + ) + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert "Hello there friend." in contents, contents + + +def test_safety_net_honors_disabled_auto_heal_for_late_incomplete_call(): + # A late call caught by the safety net: an unclosed ```` heals only with Auto-Heal on; + # off, the safety net must not pass ``allow_incomplete=True`` and execute a truncated call. + prose = "Sure, let me look that up for you right now. " + incomplete = '{"name":"web_search","arguments":{"query":"weather in Sydney"}}' + + loop_off, exec_off = _make_loop( + turns = [[prose, incomplete], ["Final answer."]], + exec_results = ["RESULT"], + auto_heal_tool_calls = False, + max_tool_iterations = 3, + ) + events_off = _collect_events(loop_off) + assert exec_off.calls == [], "disabled Auto-Heal must not execute a healed incomplete call" + assert not [e for e in events_off if e.get("type") == "tool_start"] + + loop_on, exec_on = _make_loop( + turns = [[prose, incomplete], ["Final answer."]], + exec_results = ["RESULT"], + auto_heal_tool_calls = True, + max_tool_iterations = 3, + ) + _collect_events(loop_on) + assert exec_on.calls == [("web_search", {"query": "weather in Sydney"})], exec_on.calls + + +def test_bare_json_tool_call_is_not_streamed_as_content(): + # Llama-3.2 ``custom_tools`` bare form ``{"name":..,"parameters":..}`` carries no + # XML signal. The loop must BUFFER it until the object closes and execute it via + # the safety net, never leaking the raw JSON to streaming clients as content. + bare = '{"name":"web_search","parameters":{"query":"cats"}}' + loop, exec_fn = _make_loop( + turns = [[bare], ["Here are the results."]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any('"name"' in t or "web_search" in t for t in contents), contents + assert any("Here are the results." in t for t in contents) + + +def test_ordinary_json_with_name_key_is_shown_not_treated_as_tool_call(): + # Markerless JSON whose "name" is not an enabled tool (e.g. a person record + # ``{"name":"Alice",...}``) must be shown as the answer, not misread as a call + # to a disabled tool and dropped. _make_loop enables web_search/python/terminal. + answer = '{"name":"Alice","parameters":{"age":30}}' + loop, exec_fn = _make_loop(turns = [[answer]], max_tool_iterations = 1) + events = _collect_events(loop) + assert exec_fn.calls == [], exec_fn.calls + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert "Alice" in contents, contents + + +def test_bare_json_tool_call_split_across_chunks_is_not_streamed(): + # Same as above but the bare object arrives split mid-key, so the buffer is + # held open across chunks before it balances. + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"web_', 'search","parameters":{"query":"cats"}}'], + ["Done."], + ], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any('"name"' in t or "web_search" in t for t in contents), contents + + +def test_gemma_wrapperless_call_is_not_streamed_as_content(): + # Gemma 4 wrapper-less ``call:NAME{...}`` has no XML signal; the loop must hold + # it (BUFFERING) and execute it, never streaming the raw call text. + loop, exec_fn = _make_loop( + turns = [["call:web_search{query:cats}"], ["Found."]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("call:web_search" in t for t in contents), contents + + +def test_gemma_wrapperless_call_with_whitespace_is_suppressed_when_streamed(): + # Gemma may emit ``call : NAME{...}`` with whitespace around the colon, split across stream + # chunks. + loop, exec_fn = _make_loop( + turns = [["call", " : ", "web_search", "{query:cats}"], ["Found."]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("call" in t for t in contents), contents + + +def test_long_gemma_tool_name_is_not_streamed_as_content(): + # A tool name longer than the small buffer cap (OpenAI 64 chars, MCP longer) + # must still be held: the ``call:NAME`` prefix keeps buffering until ``{`` + # instead of leaking ``call:longname`` as visible text. + long_name = "mcp__github__list_repository_issues" # 35 chars + turns = iter([list('call:%s{repo:"octo/hello"}' % long_name), ["Done."]]) + + def _gen(_messages): + try: + chunks = next(turns) + except StopIteration: + return + acc = "" + for c in chunks: + acc += c + yield acc + + exec_fn = FakeExecuteTool(["RESULT"]) + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "hi"}], + tools = [{"type": "function", "function": {"name": long_name}}], + execute_tool = exec_fn, + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [(long_name, {"repo": "octo/hello"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("call:" in t for t in contents), contents + + +def test_leading_json_answer_is_not_dropped(): + # A leading ``{...}`` that is NOT a tool call must still surface as content: + # the bare-JSON hold can only ever delay it to end-of-object, never drop it. + obj = '{"answer": 42, "note": "done"}' + loop, exec_fn = _make_loop( + turns = [[obj]], + exec_results = [], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + contents = [e["text"] for e in events if e["type"] == "content"] + assert any('"answer"' in t for t in contents), contents + + +def _reprompt_loop(*, auto_heal_tool_calls): + """Drive one restricted tool with an intent-only first turn to exercise the nudge; returns conversations and events.""" + captured: list[list] = [] + + def fake_single_turn(messages, active_tools = None): + captured.append(list(messages)) + if len(captured) == 1: + yield "I'll search for that now." # forward-looking intent, no call + else: + yield "Final answer." + + exec_fn = FakeExecuteTool([]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = fake_single_turn, + messages = [{"role": "user", "content": "find X"}], + tools = [{"type": "function", "function": {"name": "search_knowledge_base"}}], + execute_tool = exec_fn, + auto_heal_tool_calls = auto_heal_tool_calls, + # Studio always nudges (always-on for the Studio inference paths); the + # API opts in per request. Model the Studio caller here. + nudge_tool_calls = True, + max_tool_iterations = 3, + ) + ) + return captured, events + + +def test_reprompt_names_only_active_tools_not_hardcoded(): + # The plan-without-action nudge must name the tools actually enabled, never the + # old hardcoded ``web_search``/``python`` (which a restricted set would reject). + captured, _events = _reprompt_loop(auto_heal_tool_calls = True) + assert len(captured) >= 2, "intent prose should have triggered a re-prompt turn" + reprompt = captured[1][-1] + assert reprompt["role"] == "user" + assert "search_knowledge_base" in reprompt["content"] + assert "web_search" not in reprompt["content"] + assert "python" not in reprompt["content"] + + +def test_reprompt_suppressed_when_auto_heal_disabled(): + # With Auto-Heal off the safetensors nudge must stay silent for backend parity + # with the GGUF loop, so only the single initial generation runs. + captured, events = _reprompt_loop(auto_heal_tool_calls = False) + assert len(captured) == 1, captured + contents = [e["text"] for e in events if e["type"] == "content"] + assert any("search for that" in t for t in contents) + + class TestLoopBasic: def test_plain_answer(self): # No tool XML; loop should yield content then status="". @@ -406,6 +2318,154 @@ class TestLoopBasic: contents = [e for e in events if e["type"] == "content"] assert "Result: 1" in contents[-1]["text"] + def test_llama3_python_tag_form(self): + # The agentic loop must recognise Llama-3's <|python_tag|> + # marker, drain the rest of the turn, and execute the call. + loop, exec_fn = _make_loop( + turns = [ + [ + "<|python_tag|>web_search.call(", + 'query="weather in Tokyo"', + ")", + ], + ["The weather is sunny."], + ], + exec_results = ["Sunny, 22C"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "weather in Tokyo"})] + contents = [e for e in events if e["type"] == "content"] + assert "sunny" in contents[-1]["text"].lower() + + def test_llama3_bare_json_form_fires_tool(self): + # Llama-3.1 / 3.2 emit a bare-JSON tool call + # ``{"name":..,"parameters":..}`` with NO XML signal. The loop's + # safety-net parse must still fire the tool instead of treating the + # turn as "planned without calling tools" and re-prompting the model + # into giving up. Regression for the has_tool_signal gate that + # dropped these; GGUF's llama-server parses them natively. + loop, exec_fn = _make_loop( + turns = [ + ['{"name": "web_search", "parameters": {"query": "weather in SF"}}'], + ["The weather is sunny."], + ], + exec_results = ["Sunny, 18C"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "weather in SF"})] + contents = [e for e in events if e["type"] == "content"] + assert "sunny" in contents[-1]["text"].lower() + + def test_mistral_pre_v11_form(self): + # Pre-v11 Mistral emission: ``[TOOL_CALLS] [{...}]``. + loop, exec_fn = _make_loop( + turns = [ + [ + '[TOOL_CALLS] [{"name":"web_search",', + '"arguments":{"query":"hi"},"id":"abc"}]', + ], + ["done"], + ], + exec_results = ["ok"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "hi"})] + # Mistral-provided ids must propagate to tool_start events. + tool_start = next(e for e in events if e["type"] == "tool_start") + assert tool_start["tool_call_id"] == "abc" + + def test_mistral_v11_form(self): + # v11+ Mistral emission: bare ``name{json}`` after the trigger. + loop, exec_fn = _make_loop( + turns = [ + ['[TOOL_CALLS]web_search{"query":"hi"}'], + ["done"], + ], + exec_results = ["ok"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "hi"})] + + def test_gemma4_form(self): + # Gemma 4 emission: ``<|tool_call>call:NAME{...}``. + loop, exec_fn = _make_loop( + turns = [ + [ + "<|tool_call>call:web_search{", + 'query:<|"|>weather<|"|>', + "}", + ], + ["sunny"], + ], + exec_results = ["Sunny, 22C"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "weather"})] + + def test_deepseek_v3_1_form(self): + # DeepSeek V3.1 emission inside the agentic loop -- the buffer state machine must wake on + # ``<｜tool▁calls▁begin｜>`` and the parser must extract the V3.1 bare-JSON body. + loop, exec_fn = _make_loop( + turns = [ + [ + "<｜tool▁calls▁begin｜>", + "<｜tool▁call▁begin｜>web_search", + "<｜tool▁sep｜>", + '{"query":"Tokyo weather"}', + "<｜tool▁call▁end｜>", + "<｜tool▁calls▁end｜>", + ], + ["The weather is sunny."], + ], + exec_results = ["Sunny, 22C"], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "Tokyo weather"})] + contents = [e for e in events if e["type"] == "content"] + assert contents and "sunny" in contents[-1]["text"].lower() + + def test_glm_form(self): + # GLM 4.x emission: ``NAME\n...``. + loop, exec_fn = _make_loop( + turns = [ + [ + "web_search\n", + "query\n", + "Tokyo\n", + "", + ], + ["found"], + ], + exec_results = ["..."], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "Tokyo"})] + + def test_kimi_form(self): + # Kimi K2 emission ``<|tool_calls_section_begin|>...``. + loop, exec_fn = _make_loop( + turns = [ + [ + "<|tool_calls_section_begin|>", + "<|tool_call_begin|>functions.web_search:0", + "<|tool_call_argument_begin|>", + '{"query":"Tokyo"}', + "<|tool_call_end|>", + "<|tool_calls_section_end|>", + ], + ["done"], + ], + exec_results = ["..."], + ) + events = _collect_events(loop) + # The bare name must reach execute_tool, even though the model + # emitted ``functions.web_search:0`` as the formatted id. + assert exec_fn.calls == [("web_search", {"query": "Tokyo"})] + # tool_start carries the original full id so the conversation + # roundtrip can replay it verbatim. + tool_start = next(e for e in events if e["type"] == "tool_start") + assert tool_start["tool_call_id"] == "functions.web_search:0" + def test_render_html_emits_provisional_tool_start(self): exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) turn_iter = iter( @@ -594,6 +2654,42 @@ class TestLoopBasic: assert tool_starts[0]["tool_name"] == "python" assert exec_fn.calls == [("python", {"code": "print('')"})] + def test_render_html_rehearsed_in_think_block_emits_no_provisional_start(self): + # BUG B: a render_html rehearsed inside think before a real python call must not emit a + # provisional render_html card; only the outside-think call fires. + exec_fn = FakeExecuteTool(["ok"]) + turn_iter = iter( + [ + [ + 'draft render_html[ARGS]{"code":"x"}', + 'python[ARGS]{"code":"print(1)"}', + ], + ["Done."], + ] + ) + + def _gen(_messages): + chunks = next(turn_iter) + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "run code"}], + tools = [ + {"type": "function", "function": {"name": "render_html"}}, + {"type": "function", "function": {"name": "python"}}, + ], + execute_tool = exec_fn, + ) + events = _collect_events(loop) + tool_starts = [e for e in events if e["type"] == "tool_start"] + + assert [e["tool_name"] for e in tool_starts] == ["python"], tool_starts + assert exec_fn.calls == [("python", {"code": "print(1)"})] + def test_render_html_success_blocks_second_canvas_call(self): exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) turn_iter = iter( @@ -765,6 +2861,59 @@ class TestLoopBehaviour: assert len(duplicate_nudges) == 1 assert captured_tool_names[2] == ["web_search", "python"] + def test_duplicate_noop_does_not_consume_budget_at_small_cap(self): + # A duplicate/disabled no-op turn is a correction turn and must NOT spend the + # caller's tool budget, so with max_tool_iterations=2 the model can still make a + # DISTINCT valid call after repeating one. Only turns that actually execute a + # tool count -- matching the GGUF loop. (The budget used to be charged per + # non-re-prompt iteration, so the duplicate burned the second slot and the third + # turn was sent with no tools, dropping the ``python`` call.) + captured_tool_names: list[list[str]] = [] + turns = iter( + [ + ['{"name":"web_search","arguments":{"query":"x"}}'], + ['{"name":"web_search","arguments":{"query":"x"}}'], + ['{"name":"python","arguments":{"code":"print(1)"}}'], + ["final"], + ] + ) + + def fake_single_turn(messages, active_tools = None): + captured_tool_names.append( + [ + tool["function"]["name"] + for tool in (active_tools or []) + if tool.get("function", {}).get("name") + ] + ) + chunks = next(turns) + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + exec_fn = FakeExecuteTool(["search-result", "python-result"]) + _collect_events( + run_safetensors_tool_loop( + single_turn = fake_single_turn, + messages = [{"role": "user", "content": "hi"}], + tools = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "python"}}, + ], + execute_tool = exec_fn, + max_tool_iterations = 2, + ) + ) + + # Both distinct tools execute; the repeated call in between did not cost a slot. + assert exec_fn.calls == [ + ("web_search", {"query": "x"}), + ("python", {"code": "print(1)"}), + ] + # The turn after the duplicate still offered tools (budget not yet spent). + assert captured_tool_names[2] == ["web_search", "python"] + def test_repeated_duplicate_noop_transitions_to_final_attempt(self): captured_tool_names: list[list[str]] = [] turns = iter( @@ -953,6 +3102,269 @@ class TestLoopBehaviour: assert "boom" in tool_end["result"] +class TestLoopRePrompt: + """Plan-without-action re-prompt parity with GGUF: nudge instead of terminating, up to ``MAX_ACT_REPROMPTS`` extra slots. Studio always nudges, so these drive the loop with ``nudge_tool_calls=True``.""" + + def test_intent_signal_triggers_reprompt(self): + # Turn 1: intent signal, no tool call. + # Turn 2 (re-prompt): proper tool call -> executes. + # Turn 3: final answer. + loop, exec_fn = _make_loop( + turns = [ + ["Let me search for that."], + [ + '{"name":"web_search","arguments":' + '{"query":"sky color"}}' + ], + ["The sky is blue."], + ], + exec_results = ["Blue (Rayleigh scattering)"], + nudge_tool_calls = True, + ) + events = _collect_events(loop) + # web_search must have been called once (after the re-prompt). + assert exec_fn.calls == [("web_search", {"query": "sky color"})] + contents = [e for e in events if e["type"] == "content"] + assert contents and "blue" in contents[-1]["text"].lower() + + def test_intent_signal_without_tools_does_not_reprompt(self): + # Same intent signal but no tools enabled -- must NOT re-prompt. + loop, exec_fn = _make_loop( + turns = [["Let me think about that for a moment."]], + exec_results = [], + ) + # _make_loop hard-codes three tools; rebuild without tools. + from core.inference.safetensors_agentic import run_safetensors_tool_loop + + def _gen(_messages): + yield "Let me think about that for a moment." + + exec_fn = FakeExecuteTool([]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "hi"}], + tools = [], + execute_tool = exec_fn, + ) + ) + assert exec_fn.calls == [] + contents = [e for e in events if e["type"] == "content"] + assert contents and "think" in contents[-1]["text"].lower() + + def test_direct_answer_does_not_trigger_reprompt(self): + # Plain answer with no intent words: do NOT re-prompt. + loop, exec_fn = _make_loop( + turns = [["4"]], + exec_results = [], + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + contents = [e for e in events if e["type"] == "content"] + assert contents and contents[-1]["text"].strip() == "4" + + def test_max_reprompts_capped(self): + # Model keeps stalling with intent -- after MAX_ACT_REPROMPTS re-prompts + # the loop must give up rather than burn forever. + turns = [["Let me search for that."]] * 6 # well over the cap + loop, exec_fn = _make_loop( + turns = turns, + exec_results = [], + nudge_tool_calls = True, + ) + events = _collect_events(loop, max_events = 500) + # No tool ever ran, but the loop terminated cleanly. + assert exec_fn.calls == [] + statuses = [e for e in events if e["type"] == "status"] + assert statuses and statuses[-1]["text"] == "" + + def test_short_intent_below_buffer_threshold_triggers_reprompt(self): + # Short emission that never exits BUFFERING (< 32 chars + no + # marker prefix). The unified buffer-end path must still + # trigger the intent re-prompt, not silently terminate. + loop, exec_fn = _make_loop( + turns = [ + ["Let me check."], + ['{"name":"web_search","arguments":{"query":"x"}}'], + ["found"], + ], + exec_results = ["..."], + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "x"})] + + def test_reprompt_does_not_consume_tool_budget(self): + # max_tool_iterations=1: one re-prompt, then one real tool call, + # then the budget-exhausted final answer must still fire. If the + # re-prompt ate the slot the tool call would never run. + loop, exec_fn = _make_loop( + turns = [ + # 1. Intent stall (re-prompt). + ["Let me search for that."], + # 2. Real tool call (uses the budget slot). + ['{"name":"web_search","arguments":{"query":"weather"}}'], + # 3. Budget exhausted -> nudged final answer. + ["Final: it is sunny"], + ], + exec_results = ["sunny"], + max_tool_iterations = 1, + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "weather"})] + contents = [e for e in events if e["type"] == "content"] + assert contents and "sunny" in contents[-1]["text"].lower() + + +class TestLoopCanonicalHealKey: + """Per-tool canonical heal key (``code``/``command``/``query``), mirroring GGUF.""" + + def test_python_bare_string_heals_to_code(self): + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"python","arguments":"print(1)"}' ""], + ["done"], + ], + exec_results = ["1\n"], + ) + events = _collect_events(loop) + # The bare string must heal to {"code": "print(1)"}, not + # {"query": ...}, so the python sandbox actually executes it. + assert exec_fn.calls == [("python", {"code": "print(1)"})] + + def test_terminal_bare_string_heals_to_command(self): + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"terminal","arguments":"ls -la"}' ""], + ["done"], + ], + exec_results = ["..."], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("terminal", {"command": "ls -la"})] + + def test_unknown_tool_bare_string_heals_to_query(self): + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"web_search","arguments":"hello"}' ""], + ["ok"], + ], + exec_results = ["..."], + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "hello"})] + + +class TestGGUFSafetensorsHealingParity: + """Pin GGUF vs safetensors/MLX loop parity so a regression on either side breaks CI.""" + + def test_gguf_imports_shared_signal_markers(self): + # The GGUF BUFFERING state machine must wake on every emission + # marker the shared parser knows -- otherwise Llama-3 / Mistral + # / Gemma 4 emissions slip past as plain prose when the + # llama-server structured channel fails. + import inspect + + from core.inference.llama_cpp import LlamaCppBackend + + src = inspect.getsource(LlamaCppBackend.generate_chat_completion_with_tools) + assert "_SHARED_TOOL_XML_SIGNALS" in src, ( + "GGUF agentic loop must reuse the shared TOOL_XML_SIGNALS " + "tuple so it wakes on all five emission formats" + ) + + def test_gguf_uses_shared_strip_helper(self): + # The GGUF stream-cleanup function must delegate to the shared + # strip_tool_markup so closed-pair markup is removed for every + # emission family (Llama-3 <|python_tag|>, Mistral [TOOL_CALLS], + # Gemma 4 <|tool_call>...). + import inspect + + from core.inference.llama_cpp import LlamaCppBackend + + src = inspect.getsource(LlamaCppBackend.generate_chat_completion_with_tools) + assert ( + "_shared_strip_tool_markup" in src + ), "GGUF stream cleanup must delegate to the shared strip_tool_markup helper" + + def test_gguf_uses_canonical_heal_keys(self): + # GGUF and safetensors heal a bare-string ``arguments`` to the same + # per-tool canonical key -- ``code`` for python, ``command`` for + # terminal, ``query`` for everything else. The mapping is centralised in + # the shared ToolLoopController (both backends route bare-string args + # through ``coerce_tool_arguments``), so the two paths cannot drift. + from core.inference.tool_loop_controller import ( + _CANONICAL_HEAL_ARG, + coerce_tool_arguments, + ) + + assert _CANONICAL_HEAL_ARG["python"] == "code" + assert _CANONICAL_HEAL_ARG["terminal"] == "command" + assert coerce_tool_arguments("print(1)", heal = True, tool_name = "python").arguments == { + "code": "print(1)" + } + assert coerce_tool_arguments("ls -la", heal = True, tool_name = "terminal").arguments == { + "command": "ls -la" + } + assert coerce_tool_arguments("weather", heal = True, tool_name = "web_search").arguments == { + "query": "weather" + } + + def test_intent_regex_matches_same_phrases_as_gguf(self): + # The intent re-prompt regex is now a single shared source of truth + # (tool_call_parser.INTENT_SIGNAL) consumed by both the GGUF and the + # safetensors/MLX loops, so behaviour is identical on Mac and Linux. + # Both backends must resolve to that one shared helper. + from core.inference.llama_cpp import ( + _is_short_intent_without_action as gguf_fn, + ) + from core.inference.safetensors_agentic import ( + is_short_intent_without_action as sf_fn, + ) + from core.inference.tool_call_parser import ( + INTENT_SIGNAL as shared_re, + is_short_intent_without_action as shared_fn, + ) + + assert gguf_fn is shared_fn and sf_fn is shared_fn + + for phrase in ( + "I'll search for that", + "I will look it up", + "Let me check", + "I am going to call the tool", + "First, I will explore", + "Here's my plan", + "Now I need to call web_search", + ): + assert shared_re.search(phrase), f"missed {phrase!r}" + assert shared_fn(phrase), f"helper missed {phrase!r}" + + for plain in ( + "4", + "Hello!", + "The sky is blue.", + "I can help with that.", + "I should mention", + "Let's go.", + # Negated intent is a refusal, not a plan: neither backend may + # force a tool-call re-prompt on it. + "I will not search the web for that.", + "I'll never call that tool.", + ): + assert not shared_re.search(plain), f"wrongly fired on {plain!r}" + assert not shared_fn(plain), f"helper wrongly fired on {plain!r}" + + def test_max_reprompts_equal_on_both_backends(self): + # Both loops draw the cap from the shared constant, so they stay equal. + from core.inference.llama_cpp import _MAX_REPROMPTS as gguf_cap + from core.inference.safetensors_agentic import MAX_ACT_REPROMPTS as sf_cap + from core.inference.tool_call_parser import MAX_ACT_REPROMPTS as shared_cap + + assert gguf_cap == sf_cap == shared_cap + + class TestLoopControl: def test_cancel_event_breaks_loop(self): cancel = threading.Event() @@ -1369,6 +3781,28 @@ class TestGuardrails: and event.get("type") in {"tool_start", "tool_end"} ] + def test_same_turn_distinct_calls_are_capped(self): + # >_MAX_TOOL_CALLS_PER_TURN DISTINCT calls in one turn must be capped so a runaway turn + # cannot fan out into many executions (the GGUF path is held back by llama-server's lazy ... + from core.inference.safetensors_agentic import _MAX_TOOL_CALLS_PER_TURN + + n = _MAX_TOOL_CALLS_PER_TURN + 4 + turn = "".join( + '{"name":"web_search","arguments":{"query":"q%d"}}' % i + for i in range(n) + ) + loop, exec_fn = _make_loop( + turns = [[turn], ["final"]], + exec_results = ["r"] * n, + max_tool_iterations = 2, + ) + _collect_events(loop) + assert len(exec_fn.calls) == _MAX_TOOL_CALLS_PER_TURN + # The first N distinct queries executed, in document order. + assert [a["query"] for _name, a in exec_fn.calls] == [ + "q%d" % i for i in range(_MAX_TOOL_CALLS_PER_TURN) + ] + def test_coerce_string_args_python_uses_code_key(self): assert _coerce_arguments("print(1)", heal = True, tool_name = "python") == {"code": "print(1)"} @@ -1407,5 +3841,918 @@ class TestGptOssNameDetection: assert is_gpt_oss_model_name(cast(str, None)) is False +# ──────────────────────────────────────────────────────────────────── +# Plan-without-action re-prompt (GGUF loop parity) +# ──────────────────────────────────────────────────────────────────── + + +class TestPlanWithoutActionReprompt: + def test_short_intent_is_reprompted_and_tool_executes(self): + loop, exec_fn = _make_loop( + turns = [ + ["I'll search the web for that."], + ['{"name":"web_search","arguments":{"query":"cats"}}'], + ["Here is the final answer."], + ], + exec_results = ["result-1"], + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert [c[0] for c in exec_fn.calls] == ["web_search"] + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("Here is the final answer." in t for t in texts) + + def test_reprompt_fires_up_to_the_cap(self): + # GGUF parity: a persistently stalling model is re-prompted up to + # MAX_ACT_REPROMPTS times, then the last stall is surrendered as the + # final answer and no further turn is generated. + from core.inference.tool_call_parser import MAX_ACT_REPROMPTS + + stall = "Let me look into it first." + turns = [["I'll search the web for that."]] + turns += [[stall]] * MAX_ACT_REPROMPTS + turns += [["SHOULD NOT APPEAR"]] + + generations = {"count": 0} + turn_iter = iter(turns) + + def _gen(_messages): + generations["count"] += 1 + try: + chunks = next(turn_iter) + except StopIteration: + return + acc = "" + for c in chunks: + acc += c + yield acc + + exec_fn = FakeExecuteTool([]) + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "hi"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + # One initial turn plus exactly MAX_ACT_REPROMPTS re-prompted turns. + assert generations["count"] == MAX_ACT_REPROMPTS + 1 + texts = [e["text"] for e in events if e["type"] == "content"] + assert any(stall in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_long_prose_answer_is_not_reprompted(self): + long_answer = "I'll keep explaining the details of the topic. " * 60 + loop, exec_fn = _make_loop( + turns = [ + [long_answer], + ["SHOULD NOT APPEAR"], + ], + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_disabled_auto_heal_is_not_reprompted(self): + loop, exec_fn = _make_loop( + turns = [ + ["I'll search the web for that."], + ["SHOULD NOT APPEAR"], + ], + auto_heal_tool_calls = False, + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("I'll search the web for that." in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_explicit_nudge_off_is_not_reprompted(self): + loop, exec_fn = _make_loop( + turns = [ + ["I'll search the web for that."], + ["SHOULD NOT APPEAR"], + ], + nudge_tool_calls = False, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("I'll search the web for that." in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_omitted_nudge_flag_is_not_reprompted(self): + # The retry is new on this loop: API callers who do not send the flag + # must keep today's behavior. Studio opts in explicitly. + loop, exec_fn = _make_loop( + turns = [ + ["I'll search the web for that."], + ["SHOULD NOT APPEAR"], + ], + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("I'll search the web for that." in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_rag_autoinject_counts_as_executed_tool(self, monkeypatch): + # Autoinject already ran a KB search outside the controller; a short + # post-retrieval intent must not trigger a spurious re-prompt. + import core.inference.tools as tools_mod + + def fake_autoinject(conversation, rag_scope): + return { + "events": [ + {"type": "tool_start", "tool_name": "search_knowledge_base"}, + {"type": "tool_end", "tool_name": "search_knowledge_base"}, + ], + "messages": [{"role": "tool", "content": "kb result"}], + } + + monkeypatch.setattr(tools_mod, "build_rag_autoinject", fake_autoinject) + loop, exec_fn = _make_loop( + turns = [ + ["I'll search the docs."], + ["SHOULD NOT APPEAR"], + ], + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + assert any(e.get("type") == "tool_start" for e in events) + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("I'll search the docs." in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_no_reprompt_after_a_denied_tool_confirmation(self, monkeypatch): + # An explicit user denial must not be answered with a nudge to call + # the tool again (which would raise another confirmation prompt). + monkeypatch.setattr(safetensors_agentic, "new_approval_id", lambda: "appr-1") + monkeypatch.setattr(safetensors_agentic, "begin_tool_decision", lambda *_a, **_k: object()) + monkeypatch.setattr(safetensors_agentic, "wait_tool_decision", lambda *_a, **_k: "deny") + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"web_search","arguments":{"query":"cats"}}'], + ["I'll search again."], + ["SHOULD NOT APPEAR"], + ], + confirm_tool_calls = True, + session_id = "sess", + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert any("I'll search again." in t for t in texts) + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + def test_no_reprompt_after_a_tool_already_executed(self): + loop, exec_fn = _make_loop( + turns = [ + ['{"name":"web_search","arguments":{"query":"cats"}}'], + ["Now I'll refine the search."], + ["SHOULD NOT APPEAR"], + ], + exec_results = ["result-1"], + nudge_tool_calls = True, + ) + events = _collect_events(loop) + assert [c[0] for c in exec_fn.calls] == ["web_search"] + texts = [e["text"] for e in events if e["type"] == "content"] + assert not any("SHOULD NOT APPEAR" in t for t in texts) + + +# Routes-level python_tag strip (multi-line; stop on next sentinel) +class TestRoutesPythonTagStrip: + """``_TOOL_XML_RE`` must consume multi-line code, embedded JSON, and bare ``<`` (earlier ``[^\n<]*`` / ``[^\n]*`` revisions leaked tails); the streaming route-level strip is the regression-prone path.""" + + def _strip(self, text: str) -> str: + # Import inside the test so a routes-module import error does + # not blow up the entire test file at collection time. + from routes.inference import _strip_tool_xml + return _strip_tool_xml(text) + + def test_single_line_python_tag_stripped(self): + # Floor: the original 5620 single-line behaviour still works. + text = '<|python_tag|>brave_search.call(query="weather")' + assert self._strip(text) == "" + + def test_python_tag_with_less_than_in_code(self): + # 5615 regression: literal ``<`` inside code must NOT terminate + # the strip early. + text = '<|python_tag|>python.call(code="if x < 10: pass")' + assert self._strip(text) == "" + + def test_python_tag_multiline_code_stripped(self): + # 5620 round-1 regression: multi-line code's second line leaked. + text = '<|python_tag|>python.call(code="line1\nline2\nline3")' + assert self._strip(text) == "" + + def test_python_tag_multiline_with_less_than(self): + # Combined: multi-line code AND literal ``<`` in code. + text = ( + '<|python_tag|>python.call(code="for i in range(10):\n' + " if i < 5:\n" + ' print(i)")' + ) + assert self._strip(text) == "" + + def test_python_tag_stops_at_eom_sentinel(self): + # Strip stops at the next Llama-3 ``<|`` sentinel so any + # trailing assistant content survives. + text = '<|python_tag|>python.call(code="multi\nline")' "<|eom_id|>final answer text" + assert self._strip(text) == "<|eom_id|>final answer text" + + def test_python_tag_stops_at_eot_sentinel(self): + text = '<|python_tag|>brave_search.call(query="x")' "<|eot_id|>after" + assert self._strip(text) == "<|eot_id|>after" + + def test_python_tag_json_form_multiline_stripped(self): + # The JSON form of python_tag with newlines inside string args. + text = '<|python_tag|>{"name":"python","parameters":{"code":"a = 1\nb = 2\nprint(a+b)"}}' + assert self._strip(text) == "" + + def test_python_tag_with_eom_then_trailing_python_tag(self): + # Two python_tag emissions back-to-back across a sentinel: both + # should strip independently. + text = ( + '<|python_tag|>brave_search.call(query="a")' + "<|eom_id|>" + '<|python_tag|>python.call(code="x=1")' + ) + # ``<|eom_id|>`` between the two strips remains; both + # python_tag blocks are fully consumed. + assert self._strip(text) == "<|eom_id|>" + + +# Robustness fixes uncovered while validating against vLLM / sglang. +class TestParserRobustness: + def test_tool_call_json_accepts_parameters_key(self): + # Hermes wrapper around a Llama-3.2 bare-JSON object that uses + # ``parameters`` instead of ``arguments``. The bare-JSON and + # python_tag paths already accept both keys; this path now does + # too. Was extracting name only and silently dropping the args. + import json + + text = "\n" '{"name": "search", "parameters": {"q": "ramen"}}\n' "" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "search" + assert json.loads(result[0]["function"]["arguments"]) == {"q": "ramen"} + + def test_function_xml_attribute_form(self): + # MiniCPM-5 / MiniMax-M2 attribute syntax: + # ``v``. + import json + + text = '' 'Tokyo' "" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + def test_function_xml_attribute_form_multi_param(self): + import json + + text = ( + '' + 'Tokyo' + 'celsius' + "" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"city": "Tokyo", "unit": "celsius"} + + def test_function_xml_legacy_equals_form_still_works(self): + # Regression guard: the old ``v`` + # syntax must keep parsing after the regex broadening. + import json + + text = "Tokyo" + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + def test_function_attribute_form_has_tool_signal(self): + # The standalone ```` attribute form must flip + # the streaming buffer; otherwise the end-of-turn safety-net parse in + # the agentic loop is gated off and the real call is dropped. + assert has_tool_signal('') is True + + def test_function_attribute_form_strip_markup(self): + # The attribute form must also be stripped from displayed text, like + # the legacy ```` form. + text = 'result X' + assert strip_tool_markup(text, final = True) == "result" + + def test_llama3_chat_template_round_trip(self): + # Meta's official Llama-3.x chat template prefixes every + # assistant turn with + # ``<|start_header_id|>assistant<|end_header_id|>\n\n``. The + # sentinel-strip in ``_parse_llama3_bare_json`` must reach past + # the role label to the JSON body, else every round-tripped + # tool call in history silently drops. + import json + + text = ( + "<|start_header_id|>assistant<|end_header_id|>\n\n" + '{"name": "get_weather", "parameters": {"city": "Tokyo"}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + def test_llama3_round_trip_all_roles(self): + # Same logic must work for every role the chat template inserts. + import json + for role in ("assistant", "user", "system", "tool", "ipython"): + text = ( + f"<|start_header_id|>{role}<|end_header_id|>\n\n" + '{"name": "f", "parameters": {"x": 1}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1, f"failed for role={role}" + assert json.loads(result[0]["function"]["arguments"]) == {"x": 1} + + def test_llama3_round_trip_with_eot_prefix(self): + # Prior assistant turn closes with ``<|eot_id|>``, then the + # new header opens. Both sentinels + the role must be consumed. + import json + + text = ( + "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" + '{"name": "f", "parameters": {}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "f" + + def test_function_xml_followed_by_prose(self): + # Models routinely follow a tool call with explanatory prose. + # Body must terminate at ```` even without a + # ```` wrapper, else trailing prose leaks into the + # last parameter value. + import json + + text = ( + "" + "Tokyo" + "\n\nHere is what I found." + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + def test_function_attribute_xml_followed_by_prose(self): + # Same expectation for the MiniCPM-5 attribute form. + import json + + text = ( + '' + 'Tokyo' + "\n\nLet me know if you need anything else." + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == {"city": "Tokyo"} + + +def test_render_with_native_template_returns_render_only_when_tools_emitted(): + # The native-template fallback re-renders with the model's repo template when an override drops + # the tools schema. + from types import SimpleNamespace + + from core.inference.chat_template_helpers import render_native_template + + messages = [{"role": "user", "content": "hi"}] + tools = [{"type": "function", "function": {"name": "web_search"}}] + model_info = { + "native_chat_template": "TPL", + "tokenizer": SimpleNamespace(chat_template = "OVERRIDE"), + } + + def emitting(tokenizer, msgs, *, tools, **_kw): + body = "".join(m["content"] for m in msgs) + return body + ("|TOOLS=" + ",".join(t["function"]["name"] for t in tools) if tools else "") + + def ignoring(tokenizer, msgs, *, tools, **_kw): + return "".join(m["content"] for m in msgs) # never reflects tools + + out = render_native_template( + model_info = dict(model_info), + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = emitting, + ) + assert out == "hi|TOOLS=web_search" + # The native template must be restored on the live tokenizer after probing. + assert model_info["tokenizer"].chat_template == "OVERRIDE" + + assert ( + render_native_template( + model_info = dict(model_info), + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = ignoring, + ) + is None + ) + + # No tokenizer and no processor -> return None instead of an AttributeError. + no_tok = {"native_chat_template": "TPL"} + assert ( + render_native_template( + model_info = no_tok, + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = emitting, + ) + is None + ) + + +def test_render_with_native_template_does_not_mutate_shared_tokenizer(): + # The shared tokenizer must never carry the temporary native template, even mid-render: this + # runs outside the generation lock, so a concurrent request could otherwise render with the ... + from types import SimpleNamespace + + from core.inference.chat_template_helpers import render_native_template + + shared = SimpleNamespace(chat_template = "OVERRIDE") + seen = [] + + def capture(tokenizer, msgs, *, tools, **_kw): + seen.append((tokenizer is shared, shared.chat_template)) + body = "".join(m["content"] for m in msgs) + return body + ("|T" if tools else "") + + model_info = {"native_chat_template": "TPL", "tokenizer": shared} + render_native_template( + model_info = model_info, + active_model_name = "x", + messages = [{"role": "user", "content": "hi"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + apply_fn = capture, + ) + # Rendering happened on a copy, and the shared tokenizer stayed "OVERRIDE" + # throughout (never the temporary "TPL"). + assert seen and all(not is_shared for is_shared, _ in seen) + assert all(tpl == "OVERRIDE" for _, tpl in seen) + assert shared.chat_template == "OVERRIDE" + + +def test_native_template_loads_from_base_model_for_lora(monkeypatch): + # For a LoRA adapter the chat template lives on the base model; active_model_name + # is the adapter id and may ship no template. The loader must read base_model. + from types import SimpleNamespace + + import transformers + + from core.inference.chat_template_helpers import render_native_template + + captured = {} + + def fake_from_pretrained(name, *args, **kwargs): + captured["source"] = name + return SimpleNamespace(chat_template = "BASE_TPL") + + monkeypatch.setattr(transformers.AutoTokenizer, "from_pretrained", fake_from_pretrained) + + def emitting(tokenizer, msgs, *, tools, **_kw): + body = "".join(m["content"] for m in msgs) + return body + ("|T" if tools else "") + + model_info = { + "base_model": "base/model-id", + "tokenizer": SimpleNamespace(chat_template = "OVERRIDE"), + } + out = render_native_template( + model_info = model_info, + active_model_name = "adapter/path", + messages = [{"role": "user", "content": "hi"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + apply_fn = emitting, + ) + assert captured["source"] == "base/model-id" + assert out == "hi|T" + + +def test_render_with_native_template_fallback_swaps_when_override_drops_tools(): + # The shared gate (used by the transformers and MLX backends): when the live render is + # identical with and without tools, re-render with the native template and return it. + from types import SimpleNamespace + + from core.inference.chat_template_helpers import render_with_native_template_fallback + + messages = [{"role": "user", "content": "hi"}] + tools = [{"type": "function", "function": {"name": "web_search"}}] + + # apply_fn that IGNORES tools -> live render drops the schema. + def ignoring(tokenizer, msgs, *, tools, **_kw): + return "".join(m["content"] for m in msgs) + + model_info = { + "native_chat_template": "TPL", + "tokenizer": SimpleNamespace(chat_template = "OVERRIDE"), + } + + # Native render emits the tools, so the fallback swaps to it. + def native_emits(tokenizer, msgs, *, tools, **_kw): + body = "".join(m["content"] for m in msgs) + return body + ("|TOOLS" if tools else "") + + out = render_with_native_template_fallback( + formatted_prompt = ignoring(None, messages, tools = tools), + tokenizer = SimpleNamespace(), + model_info = dict(model_info), + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = lambda tok, msgs, *, tools, **kw: ( + native_emits(tok, msgs, tools = tools) + if getattr(tok, "chat_template", None) == "TPL" + else ignoring(tok, msgs, tools = tools) + ), + ) + assert out == "hi|TOOLS", out + + +def test_render_with_native_template_fallback_keeps_prompt_when_tools_emitted(): + # Live render already differs with vs without tools -> no fallback, returned + # unchanged. Also a no-tools call is a passthrough. + from types import SimpleNamespace + + from core.inference.chat_template_helpers import render_with_native_template_fallback + + messages = [{"role": "user", "content": "hi"}] + tools = [{"type": "function", "function": {"name": "web_search"}}] + + def emitting(tokenizer, msgs, *, tools, **_kw): + body = "".join(m["content"] for m in msgs) + return body + ("|T" if tools else "") + + kept = render_with_native_template_fallback( + formatted_prompt = emitting(None, messages, tools = tools), + tokenizer = SimpleNamespace(), + model_info = {"native_chat_template": "TPL", "tokenizer": SimpleNamespace()}, + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = emitting, + ) + assert kept == "hi|T", kept + + # No tools -> passthrough (native template never consulted). + passthrough = render_with_native_template_fallback( + formatted_prompt = "hi", + tokenizer = SimpleNamespace(), + model_info = {}, + active_model_name = "x", + messages = messages, + tools = None, + apply_fn = emitting, + ) + assert passthrough == "hi" + + +def test_render_with_native_template_fallback_keeps_prompt_when_no_tools_probe_raises(): + # A template that REQUIRES tools can raise on the no-tools probe. + from types import SimpleNamespace + + from core.inference.chat_template_helpers import render_with_native_template_fallback + + messages = [{"role": "user", "content": "hi"}] + tools = [{"type": "function", "function": {"name": "web_search"}}] + + def raises_without_tools(tokenizer, msgs, *, tools, **_kw): + if not tools: + raise RuntimeError("template requires tools") + return "".join(m["content"] for m in msgs) + "|T" + + out = render_with_native_template_fallback( + formatted_prompt = "hi|T", + tokenizer = SimpleNamespace(), + model_info = {"native_chat_template": "TPL", "tokenizer": SimpleNamespace()}, + active_model_name = "x", + messages = messages, + tools = tools, + apply_fn = raises_without_tools, + ) + assert out == "hi|T", out + + +def test_truncated_bare_json_at_eof_is_not_leaked(): + # Stream ends mid bare-JSON object: the held fragment must be dropped at the + # EOF resolver, not flushed as plain assistant content (GGUF parity). + loop, _exec = _make_loop( + turns = [['{"name":"web_search","parameters":{"query":"weather in S']], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any('"name"' in t for t in contents), contents + + +def test_oversized_bare_json_call_is_not_leaked_and_executes(): + # A bare-JSON call whose arguments exceed _MAX_BARE_JSON_BUFFER must DRAIN + # (suppress) rather than stream the raw JSON prefix, and still execute once + # the full object is parsed by the safety net. + from core.inference.safetensors_agentic import _MAX_BARE_JSON_BUFFER + + big = "A" * (_MAX_BARE_JSON_BUFFER + 5000) + full = '{"name":"python","parameters":{"code":"' + big + '"}}' + chunks = [full[i : i + 2000] for i in range(0, len(full), 2000)] + loop, exec_fn = _make_loop(turns = [chunks, ["done"]], exec_results = ["OK"], max_tool_iterations = 2) + events = _collect_events(loop) + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any(t.lstrip().startswith('{"name') for t in contents), contents[:1] + assert exec_fn.calls and exec_fn.calls[0][0] == "python" + assert len(exec_fn.calls[0][1].get("code", "")) > _MAX_BARE_JSON_BUFFER + + +def test_oversized_plain_json_answer_still_streams(): + # A giant plain JSON answer (no "name" key) is NOT a tool call and must still + # stream -- the oversized DRAIN route is gated on a "name" key. + from core.inference.safetensors_agentic import _MAX_BARE_JSON_BUFFER + + big = "A" * (_MAX_BARE_JSON_BUFFER + 5000) + full = '{"result":"' + big + '"}' + chunks = [full[i : i + 2000] for i in range(0, len(full), 2000)] + loop, _exec = _make_loop(turns = [chunks], max_tool_iterations = 1) + events = _collect_events(loop) + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert '"result"' in contents + + +def test_oversized_disabled_name_json_answer_still_streams(): + # A giant still-open JSON answer whose "name" is NOT an enabled tool must stream: + # the oversized DRAIN branch was gated only on the presence of a "name" key, so a + # large ordinary record ({"name":"Alice",...}) was drained instead of shown. + from core.inference.safetensors_agentic import _MAX_BARE_JSON_BUFFER + + big = "A" * (_MAX_BARE_JSON_BUFFER + 5000) + answer = '{"name":"Alice","parameters":{"bio":"' + big # never closes + chunks = [answer[i : i + 2000] for i in range(0, len(answer), 2000)] + loop, exec_fn = _make_loop(turns = [chunks], max_tool_iterations = 1) + events = _collect_events(loop) + assert exec_fn.calls == [], exec_fn.calls + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert "Alice" in contents, contents[:80] + + +def test_truncated_disabled_name_json_is_shown_at_eof(): + # A truncated ordinary JSON answer whose name is not an enabled tool, held to EOF, + # must be shown -- the EOF bare-JSON DRAIN branch was gated only on a "name" key. + truncated = '{"name":"Alice","parameters":{"age":' + loop, exec_fn = _make_loop(turns = [[truncated]], max_tool_iterations = 1) + events = _collect_events(loop) + assert exec_fn.calls == [], exec_fn.calls + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert "Alice" in contents, contents + + +def test_truncated_plain_json_with_nested_enabled_name_is_visible(): + # A truncated ordinary JSON answer with a NESTED ``"name"`` matching an enabled + # tool ({"result":{"name":"web_search",...) must be shown, not suppressed: the + # gate now extracts the TOP-LEVEL name only, so the nested field is just data. + loop, exec_fn = _make_loop( + turns = [['{"result":{"name":"web_search","age":']], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + assert exec_fn.calls == [] + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert '"result"' in contents and "web_search" in contents, contents + + +def test_bare_json_call_not_replayed_in_next_turn_content(): + # After a complete bare-JSON call executes, the assistant content fed to the + # next turn must not contain the raw call (next-turn contamination). + captured: list[list[dict]] = [] + exec_fn = FakeExecuteTool(["RESULT"]) + + def st(messages, active_tools = None): + captured.append([dict(m) for m in messages]) + if len(captured) == 1: + yield '{"name":"web_search","parameters":{"query":"cats"}}' + else: + yield "Found." + + _collect_events( + run_safetensors_tool_loop( + single_turn = st, + messages = [{"role": "user", "content": "cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + max_tool_iterations = 3, + ) + ) + assert len(captured) >= 2, captured + asst = [m for m in captured[1] if m.get("role") == "assistant"] + assert asst and not any('"name"' in (m.get("content") or "") for m in asst), asst + + if __name__ == "__main__": pytest.main([__file__, "-v"]) + + +def test_streaming_strip_keeps_bare_args_before_think_block(): + # F3: a bare ``foo[ARGS]`` before a think block is prose; EOS-anchored tail arms run only + # on the last segment. + text = "Please pass foo[ARGS] pause to the template." + out = strip_tool_markup_streaming(text, tool_protocol_active = True) + assert out == text + + +def test_streaming_strip_still_removes_complete_call_before_think_block(): + # A complete bracket call before a think block still strips in the non-last segment. + text = 'go web_search[ARGS]{"q":"x"} z done' + out = strip_tool_markup_streaming(text, tool_protocol_active = True) + assert "web_search[ARGS]" not in out + assert "z" in out + assert "go" in out and "done" in out + + +def test_prose_args_marker_before_real_call_does_not_drain_the_prose(): + # F5: an inactive ``foo[ARGS]`` in prose is not a call boundary; the prose streams in + # full and the later real call still executes. + loop, exec_fn = _make_loop( + turns = [ + ["Intro ", "foo[ARGS] syntax. ", 'web_search[ARGS]{"query":"cats"}'], + ["Cats are great."], + ], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + # The prose between the bogus marker and the real call must survive. + assert any("foo[ARGS] syntax." in t for t in contents), contents + # The real call markup is never shown as content. + assert not any("web_search[ARGS]" in t for t in contents), contents + + +def test_inactive_name_args_with_body_is_not_parsed_into_disabled_noop(): + # BUG A: a prose answer with an inactive ``foo[ARGS]{...}`` is not drained into a + # disabled no-op extra turn; the [ARGS] checks are name-gated. + turns = [['foo[ARGS]{"x":1} is just syntax.']] + turn_calls: list[int] = [] + + def _gen(_messages): + turn_calls.append(1) + chunks = turns[len(turn_calls) - 1] if len(turn_calls) <= len(turns) else [] + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + exec_fn = FakeExecuteTool([]) + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "explain"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [], exec_fn.calls + assert not any(e["type"] in ("tool_start", "tool_end") for e in events), events + # Exactly one generation turn -- no disabled ``foo`` no-op re-prompt. + assert len(turn_calls) == 1, turn_calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert any("is just syntax." in t for t in contents), contents + + +class TestEnabledToolNameGate: + """The safetensors loop passes the active tool names into parse/strip so the + ambiguous bare-rehearsal ``NAME[ARGS]{json}`` is treated as a call only when NAME + is an active tool (#5704). Without the gate an inactive ``foo[ARGS]{...}`` in prose + was parsed into a disabled no-op call and stripped from the visible text.""" + + def _names(self, calls): + return [c["function"]["name"] for c in calls] + + def test_parse_inactive_rehearsal_does_not_swallow_active_call(self): + text = 'foo[ARGS]{"a":1} web_search[ARGS]{"query":"cats"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert self._names(calls) == ["web_search"] + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_parse_inactive_rehearsal_alone_is_prose(self): + assert ( + parse_tool_calls_from_text('foo[ARGS]{"a":1}', enabled_tool_names = {"web_search"}) == [] + ) + + def test_streaming_strip_keeps_inactive_rehearsal(self): + raw = 'answer foo[ARGS]{"x":1} tail' + assert strip_tool_markup_streaming(raw, enabled_tool_names = {"web_search"}) == raw + + def test_streaming_strip_removes_active_rehearsal(self): + raw = 'answer web_search[ARGS]{"q":1} tail' + out = strip_tool_markup_streaming(raw, enabled_tool_names = {"web_search"}) + assert "web_search[ARGS]" not in out + assert out == "answer tail" + + def test_final_strip_keeps_inactive_rehearsal(self): + text = 'foo[ARGS]{"x":1} is just syntax.' + assert strip_tool_markup(text, final = True, enabled_tool_names = {"web_search"}) == text + + def test_gate_none_preserves_legacy_strip_and_parse(self): + text = 'foo[ARGS]{"x":1} tail' + assert self._names(parse_tool_calls_from_text(text)) == ["foo"] + assert strip_tool_markup_streaming(text) == " tail" + + +def test_drain_truncated_enabled_name_json_preserved_when_auto_heal_disabled(): + # F3: with Auto-Heal OFF, a truncated ENABLED-name bare-JSON fragment that did + # not parse must stay visible (disabled-Auto-Heal contract: malformed markup is + # preserved), matching the XML strip in the same drain branch. With Auto-Heal ON + # the same fragment is suppressed. + trunc = '{"name":"web_search","parameters":{"query":"weather' + off, exec_off = _make_loop(turns = [[trunc]], max_tool_iterations = 1, auto_heal_tool_calls = False) + events_off = _collect_events(off) + assert exec_off.calls == [], exec_off.calls + contents_off = "".join(e["text"] for e in events_off if e["type"] == "content") + assert "web_search" in contents_off, contents_off + + on, exec_on = _make_loop(turns = [[trunc]], max_tool_iterations = 1, auto_heal_tool_calls = True) + events_on = _collect_events(on) + assert exec_on.calls == [], exec_on.calls + contents_on = "".join(e["text"] for e in events_on if e["type"] == "content") + assert "web_search" not in contents_on, contents_on + + +def test_looks_like_enabled_bare_json_accepts_function_alias(): + # The safetensors buffering gate must recognise the "function" bare-JSON alias + # the parser accepts, so a truncated/complete {"function":} call is + # buffered/healed instead of streaming as visible content. + from core.inference.safetensors_agentic import _looks_like_enabled_bare_json + + enabled = {"web_search"} + assert _looks_like_enabled_bare_json( + '{"function":"web_search","parameters":{"q":"x"}}', enabled + ) + # A non-tool "function" value is an ordinary JSON answer -> not gated. + assert not _looks_like_enabled_bare_json('{"function":"Alice","parameters":{}}', enabled) + + +class TestFalseAlarmMarkerProse: + def test_leading_marker_prose_streams_intact(self): + # An answer that starts with a literal marker is a false alarm: the + # drain finds no calls and the full prose must reach the client. + text = "[TOOL_CALLS] is the Mistral tool marker. More prose after." + loop, exec_fn = _make_loop(turns = [[text]]) + events = _collect_events(loop) + assert exec_fn.calls == [] + texts = [e["text"] for e in events if e["type"] == "content"] + assert texts and texts[-1] == text + + def test_chained_bare_json_calls_not_replayed_in_history(self): + # Both chained calls execute; the kept content (next-turn assistant + # history) must not contain the second call's raw JSON. + chained = ( + '{"name":"web_search","parameters":{"q":"first"}};' + '{"name":"python","parameters":{"code":"x"}}' + ) + convs = [] + turn_iter = iter([[chained], ["Final answer."]]) + + def gen(messages, active_tools = None): + convs.append([dict(m) for m in messages]) + try: + chunks = next(turn_iter) + except StopIteration: + return + acc = "" + for c in chunks: + acc += c + yield acc + + exec_fn = FakeExecuteTool(["r1", "r2"]) + loop = run_safetensors_tool_loop( + single_turn = gen, + messages = [{"role": "user", "content": "hi"}], + tools = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "python"}}, + ], + execute_tool = exec_fn, + ) + _collect_events(loop) + assert [c[0] for c in exec_fn.calls] == ["web_search", "python"] + assistant = next(m for m in convs[1] if m["role"] == "assistant") + assert '"python"' not in (assistant.get("content") or "") diff --git a/studio/backend/tests/test_safetensors_toolcall_wiring.py b/studio/backend/tests/test_safetensors_toolcall_wiring.py new file mode 100644 index 0000000000..5c298a7966 --- /dev/null +++ b/studio/backend/tests/test_safetensors_toolcall_wiring.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Deterministic backend-wiring test for the safetensors / MLX tool-calling path. + +The parser and the cumulative-text state machine are already covered exhaustively by +``test_safetensors_tool_loop.py`` with fake generators. What that suite does not touch is the +*backend's own tool-injection seam*: both ``InferenceBackend`` (transformers) and +``MLXInferenceBackend`` render the prompt through the shared +``apply_chat_template_for_generation(..., tools=...)`` helper and stream cumulative text into the +shared ``run_safetensors_tool_loop`` (see ``core/inference/inference.py`` and +``core/inference/mlx_inference.py`` -- both call the same helper and the same loop, so a single CPU +test of that seam covers the macOS MLX path too). + +This test drives that exact seam with deterministic fakes -- a fake tokenizer that records the +``tools`` it is handed, a canned tool-call generation, and a stub executor -- and asserts the full +agentic chain end to end: + + tools injected into the template -> loop parses the call -> tool dispatched once -> + tool result fed back -> generation re-entered -> final answer streamed. + +It is the deterministic, download-free stand-in for the real-model MLX / GGUF browser tool-calling +end-to-end: it imports no torch / unsloth / mlx, so it runs in the portable Backend CI alongside the +tool-call parser tests. Follow-up to the parser test PRs (#5620 / #5704). +""" + +from core.inference.chat_template_helpers import apply_chat_template_for_generation +from core.inference.safetensors_agentic import run_safetensors_tool_loop + +TOOL_NAME = "get_weather" +TOOL_ARGS = {"city": "Paris"} +FAKE_TOOL = { + "type": "function", + "function": { + "name": TOOL_NAME, + "description": "Get the current weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + }, +} +# Full parser matrix lives in test_safetensors_tool_loop.py. +TOOL_CALL_TEXT = '{"name": "get_weather", "arguments": {"city": "Paris"}}' +FINAL_ANSWER = "The weather in Paris is sunny and 22C." +TOOL_RESULT = "Paris: sunny, 22C" + + +class RecordingTokenizer: + """Fake tokenizer that records the ``tools`` handed to ``apply_chat_template``. + + Modelled on ``TestChatTemplateHelper._Tok`` in ``test_safetensors_tool_loop.py``: it accepts the + real helper's kwargs and returns a canned prompt, so the test can assert the backend seam actually + forwarded the tool schema -- a silent drop on a chat-template fallback would leave ``tools_seen`` + holding ``None``. + """ + + def __init__(self): + self.tools_seen: list = [] + self.call_count = 0 + + def apply_chat_template( + self, + messages, + *, + tokenize = False, + add_generation_prompt = True, + **kwargs, + ): + self.call_count += 1 + self.tools_seen.append(kwargs.get("tools")) + return "PROMPT" + + +class StubExecutor: + """Stand-in for ``core.inference.tools.execute_tool``: records calls, returns a fixed result. + + A fake tool name plus this stub means no real python / terminal / web / RAG side effect can run. + """ + + def __init__(self, result: str): + self.result = result + self.calls: list[tuple[str, dict]] = [] + + def __call__( + self, + name, + arguments, + *, + cancel_event = None, + timeout = None, + session_id = None, + rag_scope = None, + disable_sandbox = False, + ): + self.calls.append((name, arguments)) + return self.result + + +def _collect(generator, max_events = 200): + events = [] + for ev in generator: + events.append(ev) + if len(events) >= max_events: + break + return events + + +def _tool_names(tools): + return [(t.get("function") or {}).get("name") for t in (tools or [])] + + +def test_backend_seam_injects_tools_and_drives_full_tool_loop(): + """The shared backend seam forwards tools into the chat template, and the loop parses the call, + dispatches it once, feeds the result back, and re-enters generation for the final answer.""" + tok = RecordingTokenizer() + executor = StubExecutor(TOOL_RESULT) + turns = iter([TOOL_CALL_TEXT, FINAL_ANSWER]) + active_tools_seen: list = [] + conversations_seen: list = [] + + def single_turn(conversation, *, active_tools = None): + # Mirror the real _single_turn: render via the shared helper, then yield cumulative snapshots. + active_tools_seen.append(active_tools) + conversations_seen.append([dict(m) for m in conversation]) + apply_chat_template_for_generation(tok, conversation, tools = active_tools) + text = next(turns) + mid = len(text) // 2 + acc = "" + for chunk in (text[:mid], text[mid:]): + acc += chunk + yield acc + + events = _collect( + run_safetensors_tool_loop( + single_turn = single_turn, + messages = [{"role": "user", "content": "What is the weather in Paris?"}], + tools = [FAKE_TOOL], + execute_tool = executor, + max_tool_iterations = 3, + ) + ) + + # 1. Helper forwarded the tool schema to the tokenizer (seam does not drop tools). + assert tok.tools_seen, "tokenizer.apply_chat_template was never called" + assert tok.tools_seen[0], "tool schema was dropped before reaching the tokenizer" + assert TOOL_NAME in _tool_names(tok.tools_seen[0]) + + # 2. Loop offered the tool to the first generation turn. + assert active_tools_seen and active_tools_seen[0] is not None + assert TOOL_NAME in _tool_names(active_tools_seen[0]) + + # 3 / 4 / 5. Exactly one tool_start, one dispatch with parsed args, one tool_end with the result. + tool_starts = [e for e in events if e["type"] == "tool_start"] + tool_ends = [e for e in events if e["type"] == "tool_end"] + assert len(tool_starts) == 1 and tool_starts[0]["tool_name"] == TOOL_NAME + assert executor.calls == [(TOOL_NAME, TOOL_ARGS)], executor.calls + assert len(tool_ends) == 1 and tool_ends[0]["result"] == TOOL_RESULT + + # 6. Final answer streams after the tool result: loop appended it and re-entered generation. + contents = [e for e in events if e["type"] == "content"] + assert contents and FINAL_ANSWER in contents[-1]["text"] + last_tool_end_idx = max(i for i, e in enumerate(events) if e["type"] == "tool_end") + last_content_idx = max(i for i, e in enumerate(events) if e["type"] == "content") + assert last_content_idx > last_tool_end_idx, "final answer must stream after the tool result" + + # 6b. Tool result fed back into the conversation before the final turn (6 alone misses this: + # the fake generation ignores the conversation). + assert len(conversations_seen) >= 2, "loop did not re-enter generation after the tool call" + final_turn_convo = conversations_seen[1] + assert any( + TOOL_RESULT in str(m.get("content", "")) for m in final_turn_convo + ), "tool result was not fed back into the conversation before the final generation turn" + + # 7. Guard: raw tool-call markup never leaked to the client as content. + for e in contents: + assert "" not in e["text"] + assert TOOL_NAME not in e["text"] diff --git a/studio/backend/tests/test_security_gate_consistency.py b/studio/backend/tests/test_security_gate_consistency.py index db66df8a30..b5f1069f12 100644 --- a/studio/backend/tests/test_security_gate_consistency.py +++ b/studio/backend/tests/test_security_gate_consistency.py @@ -99,3 +99,16 @@ def test_malware_and_consent_gates_cover_the_lora_base(): if runs_gate and not resolves_base: offenders.append(f"{rel} runs a load gate but never resolves the LoRA base") assert not offenders, "\n".join(offenders) + + +def test_rag_embedding_path_runs_the_malware_gate(): + """The RAG embedding model is set through /settings and later loaded by + SentenceTransformer, which deserializes pickles; both sites must run the malware gate + or a flagged repo loads unscanned (bypassing the normal model-load protections).""" + offenders = [] + for rel in ("routes/settings.py", "core/rag/embeddings.py"): + if "evaluate_file_security(" not in (_BACKEND / rel).read_text(): + offenders.append( + f"{rel} loads/persists an embedding model without evaluate_file_security" + ) + assert not offenders, "\n".join(offenders) diff --git a/studio/backend/tests/test_sf_client_tools_passthrough.py b/studio/backend/tests/test_sf_client_tools_passthrough.py new file mode 100644 index 0000000000..01905b712c --- /dev/null +++ b/studio/backend/tests/test_sf_client_tools_passthrough.py @@ -0,0 +1,786 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Client-tools passthrough healing for the safetensors/MLX backend. + +Parity for #6801: when a NON-GGUF model is loaded and the request declares its +own ``tools`` with server-side tools OFF, text-form tool calls are promoted back +into structured ``tool_calls`` (declared tools only) via the shared healer. MLX +rides the same orchestrator path, so a single scripted backend covers both. +""" + +import asyncio +import json +from types import SimpleNamespace + +from models.inference import ChatCompletionRequest, ChatMessage +from routes.inference import openai_chat_completions +from core.inference.api_monitor import ApiMonitor + + +LOOKUP_TOOL = { + "type": "function", + "function": { + "name": "lookup", + "description": "Look something up", + "parameters": { + "type": "object", + "properties": {"q": {"type": "string"}}, + "required": ["q"], + }, + }, +} +SEARCH_TOOL = { + "type": "function", + "function": { + "name": "search", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + "required": ["query"], + }, + }, +} + +_CALL_XML = '{"name": "lookup", "arguments": {"q": "cats"}}' +_SEARCH_XML = '{"name": "search", "arguments": {"query": "dogs"}}' + + +class _Request: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/chat/completions") + method = "POST" + scope: dict = {} + + async def is_disconnected(self): + return False + + +class _ScriptedBackend: + """Non-GGUF backend: ``generate_chat_response`` replays scripted + CUMULATIVE snapshots. ``responder(messages, tools)`` returns the snapshot + list for one generation, so nudge tests can vary output across turns.""" + + active_model_name = "sf-model" + + def __init__( + self, + responder, + *, + stats = None, + ): + self.models = { + "sf-model": { + "chat_template_info": {"template": " chatml"}, + "context_length": 2048, + } + } + self._responder = responder + self._stats = stats + self.calls: list = [] + self.reset_count = 0 + + def generate_chat_response( + self, + *, + messages, + tools = None, + stats_holder = None, + **kwargs, + ): + self.calls.append({"messages": messages, "tools": tools, **kwargs}) + snapshots = self._responder(messages, tools) + if stats_holder is not None and self._stats is not None: + stats_holder["stats"] = self._stats + for snap in snapshots: + yield snap + + def reset_generation_state(self): + self.reset_count += 1 + + +def _fixed(*snapshots): + """Responder that always replays the given cumulative snapshots.""" + return lambda messages, tools: list(snapshots) + + +def _llama_stub(): + return SimpleNamespace( + is_loaded = False, + supports_tools = False, + is_vision = False, + context_length = None, + ) + + +def _install( + monkeypatch, + backend, + *, + supports_tools = True, +): + import routes.inference as inf + from state.tool_policy import reset_tool_policy + + reset_tool_policy() + monitor = ApiMonitor(max_entries = 8) + monkeypatch.setattr(inf, "api_monitor", monitor) + monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: _llama_stub()) + monkeypatch.setattr(inf, "get_inference_backend", lambda: backend) + monkeypatch.setattr( + inf, + "_detect_safetensors_features", + lambda *a, **k: {"supports_tools": supports_tools}, + ) + return monitor + + +def _request(**kwargs): + base = dict(model = "default", messages = [ChatMessage(role = "user", content = "hi")]) + base.update(kwargs) + return ChatCompletionRequest(**base) + + +def _call(payload, monkeypatch, backend, **install_kwargs): + _install(monkeypatch, backend, **install_kwargs) + + async def _run(): + return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + + return asyncio.run(_run()) + + +def _json_body(response): + return json.loads(response.body if hasattr(response, "body") else response.content) + + +def _collect_sse(response): + async def _run(): + return [c async for c in response.body_iterator] + + return asyncio.run(_run()) + + +def _sse_objects(chunks): + out = [] + for chunk in chunks: + if isinstance(chunk, bytes): + chunk = chunk.decode() + for line in str(chunk).splitlines(): + if line.startswith("data: "): + data = line.removeprefix("data: ") + if data != "[DONE]": + out.append(json.loads(data)) + return out + + +# ── Non-streaming ───────────────────────────────────────────────── + + +def test_xml_healed_to_tool_calls_non_streaming(monkeypatch): + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["content"] is None + calls = choice["message"]["tool_calls"] + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "lookup" + assert json.loads(calls[0]["function"]["arguments"]) == {"q": "cats"} + # The client tools reached the generator (template injection). + assert backend.calls[0]["tools"] == [LOOKUP_TOOL] + + +def test_undeclared_call_stays_text(monkeypatch): + xml = '{"name": "other", "arguments": {}}' + backend = _ScriptedBackend(_fixed(xml)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"].get("tool_calls") is None + assert choice["message"]["content"] == xml + + +def test_opt_out_relays_verbatim(monkeypatch): + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False, auto_heal_tool_calls = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"].get("tool_calls") is None + assert choice["message"]["content"] == _CALL_XML + + +def test_env_kill_switch_relays_verbatim(monkeypatch): + import core.inference.passthrough_healing as ph + + monkeypatch.setattr(ph, "_HEALING_DISABLED", True) + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"].get("tool_calls") is None + assert choice["message"]["content"] == _CALL_XML + + +def test_no_tools_request_untouched(monkeypatch): + backend = _ScriptedBackend(_fixed("just a plain answer")) + payload = _request(stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + # No tools and no tool messages -> plain path, normal ChatCompletion. + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"]["content"] == "just a plain answer" + assert choice["message"].get("tool_calls") is None + + +def test_prose_around_call_retained(monkeypatch): + text = "Let me look:\n" + _CALL_XML + "\ndone" + backend = _ScriptedBackend(_fixed(text)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["content"] == "Let me look:\n\ndone" + assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup" + + +def test_empty_output_is_valid_stop(monkeypatch): + backend = _ScriptedBackend(_fixed("")) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"]["content"] in ("", None) + assert choice["message"].get("tool_calls") is None + + +def test_tool_role_follow_up_turn_preserves_history(monkeypatch): + backend = _ScriptedBackend(_fixed("The weather is sunny.")) + payload = _request( + tools = [LOOKUP_TOOL], + stream = False, + messages = [ + ChatMessage(role = "user", content = "weather?"), + ChatMessage( + role = "assistant", + content = None, + tool_calls = [ + { + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": '{"q": "weather"}'}, + } + ], + ), + ChatMessage(role = "tool", tool_call_id = "call_0", content = "sunny"), + ], + ) + body = _json_body(_call(payload, monkeypatch, backend)) + assert body["choices"][0]["message"]["content"] == "The weather is sunny." + # The tool history reached the generator intact (role=tool + assistant.tool_calls). + sent = backend.calls[0]["messages"] + roles = [m["role"] for m in sent] + assert "tool" in roles + assistant = next(m for m in sent if m["role"] == "assistant") + assert assistant.get("tool_calls") + + +def test_dict_arguments_history_does_not_crash(monkeypatch): + # Non-spec client: assistant tool_calls[].function.arguments as a dict. + backend = _ScriptedBackend(_fixed("ok")) + payload = _request( + tools = [LOOKUP_TOOL], + stream = False, + messages = [ + ChatMessage(role = "user", content = "hi"), + ChatMessage( + role = "assistant", + content = None, + tool_calls = [ + { + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": {"q": "x"}}, + } + ], + ), + ChatMessage(role = "tool", tool_call_id = "call_0", content = "y"), + ], + ) + body = _json_body(_call(payload, monkeypatch, backend)) + assert body["choices"][0]["message"]["content"] == "ok" + + +def test_forced_tool_choice_narrows_promotion(monkeypatch): + # tool_choice forces `search`; a `lookup` text call must NOT promote. + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request( + tools = [LOOKUP_TOOL, SEARCH_TOOL], + stream = False, + tool_choice = {"type": "function", "function": {"name": "search"}}, + ) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"].get("tool_calls") is None + + +def test_parallel_cap_non_streaming(monkeypatch): + backend = _ScriptedBackend(_fixed(_CALL_XML + _SEARCH_XML)) + payload = _request(tools = [LOOKUP_TOOL, SEARCH_TOOL], stream = False, parallel_tool_calls = False) + body = _json_body(_call(payload, monkeypatch, backend)) + calls = body["choices"][0]["message"]["tool_calls"] + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "lookup" + + +def test_usage_recorded_when_stats_present(monkeypatch): + stats = {"usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}} + backend = _ScriptedBackend(_fixed(_CALL_XML), stats = stats) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + monitor = _install(monkeypatch, backend) + + async def _run(): + return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + + asyncio.run(_run()) + [entry] = monitor.snapshot() + assert entry["prompt_tokens"] == 7 + assert entry["completion_tokens"] == 3 + + +# ── Nudge ───────────────────────────────────────────────────────── + + +def test_nudge_default_off_single_generation(monkeypatch): + # Signal present but unparseable; without opt-in, no retry. + truncated = '{"name": "lookup"' + backend = _ScriptedBackend(_fixed(truncated)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + _call(payload, monkeypatch, backend) + assert len(backend.calls) == 1 + + +def test_nudge_opt_in_retry_recovers(monkeypatch): + truncated = '{"name": "lookup"' + + def responder(messages, tools): + nudged = any( + "native tool-call format" in (m.get("content") or "") + for m in messages + if m.get("role") == "user" + ) + return [_CALL_XML] if nudged else [truncated] + + backend = _ScriptedBackend(responder) + payload = _request(tools = [LOOKUP_TOOL], stream = False, nudge_tool_calls = True) + body = _json_body(_call(payload, monkeypatch, backend)) + assert len(backend.calls) == 2 + choice = body["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup" + + +def test_nudge_double_failure_relays_original(monkeypatch): + truncated = '{"name": "lookup"' + backend = _ScriptedBackend(_fixed(truncated)) + payload = _request(tools = [LOOKUP_TOOL], stream = False, nudge_tool_calls = True) + body = _json_body(_call(payload, monkeypatch, backend)) + assert len(backend.calls) == 2 # exactly one retry + choice = body["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"]["content"] == truncated + + +# ── Streaming ───────────────────────────────────────────────────── + + +def test_streaming_heals_split_call_into_one_delta(monkeypatch): + # Cumulative snapshots that build the call across many increments. + pieces = ["{"name": "loo', '{"name": "lookup", "argum'] + cumulative = pieces + [_CALL_XML] + backend = _ScriptedBackend(_fixed(*cumulative)) + payload = _request(tools = [LOOKUP_TOOL], stream = True) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + tool_deltas = [ + tc + for o in objs + for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or [] + ] + assert len(tool_deltas) == 1 + assert tool_deltas[0]["function"]["name"] == "lookup" + finishes = [ + o["choices"][0]["finish_reason"] + for o in objs + if o["choices"] and o["choices"][0].get("finish_reason") + ] + assert finishes == ["tool_calls"] + + +def test_streaming_cancel_does_not_finalize_tool_call(monkeypatch): + # A stream cancelled via the registry ("Stop") must NOT promote the + # buffered-but-unclosed tool markup at finalize, else it executes a tool + # the user just cancelled. Guarded on cancel_event at the finalize step. + import routes.inference as inf + + cancel_id = "cancel-me-6870" + # Balanced JSON but no closing -> healer HOLDS it until finalize. + held = '{"name": "lookup", "arguments": {"q": "cats"}}' + + class _CancelMidStream(_ScriptedBackend): + def __init__(self): + super().__init__(_fixed(held)) + + def generate_chat_response( + self, + *, + messages, + tools = None, + stats_holder = None, + **kwargs, + ): + self.calls.append({"messages": messages, "tools": tools, **kwargs}) + yield held # healer holds the unclosed call + inf._cancel_by_cancel_id_or_stash(cancel_id) # user hits Stop before EOF + + backend = _CancelMidStream() + payload = _request(tools = [LOOKUP_TOOL], stream = True, cancel_id = cancel_id) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + tool_deltas = [ + tc + for o in objs + for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or [] + ] + assert tool_deltas == [] # no tool promoted after cancel + finishes = [ + o["choices"][0]["finish_reason"] + for o in objs + if o["choices"] and o["choices"][0].get("finish_reason") + ] + assert "tool_calls" not in finishes # ends with finish_reason=stop, not tool_calls + + +def test_streaming_no_tools_verbatim(monkeypatch): + backend = _ScriptedBackend(_fixed("hello ", "hello world")) + payload = _request(stream = True) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + text = "".join( + (o["choices"][0]["delta"].get("content") or "") + for o in objs + if o["choices"] and "delta" in o["choices"][0] + ) + assert text == "hello world" + finishes = [ + o["choices"][0]["finish_reason"] + for o in objs + if o["choices"] and o["choices"][0].get("finish_reason") + ] + assert finishes == ["stop"] + + +def test_streaming_repeated_snapshot_no_duplicate_call(monkeypatch): + # Repeated then shrunk cumulative snapshots must not double-heal. + backend = _ScriptedBackend(_fixed(_CALL_XML, _CALL_XML, _CALL_XML[:5], _CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = True) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + tool_deltas = [ + tc + for o in objs + for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or [] + ] + assert len(tool_deltas) == 1 + + +def test_streaming_parallel_cap(monkeypatch): + backend = _ScriptedBackend(_fixed(_CALL_XML + _SEARCH_XML)) + payload = _request(tools = [LOOKUP_TOOL, SEARCH_TOOL], stream = True, parallel_tool_calls = False) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + tool_deltas = [ + tc + for o in objs + for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or [] + ] + assert len(tool_deltas) == 1 + assert tool_deltas[0]["function"]["name"] == "lookup" + + +def test_streaming_generator_error_closes_cleanly(monkeypatch): + def responder(messages, tools): + raise RuntimeError("boom /secret/path") + + backend = _ScriptedBackend(responder) + payload = _request(tools = [LOOKUP_TOOL], stream = True) + response = _call(payload, monkeypatch, backend) + chunks = _collect_sse(response) + joined = "".join(c.decode() if isinstance(c, bytes) else c for c in chunks) + assert "An internal error occurred" in joined + assert "secret/path" not in joined # CWE-209: no path leak + assert backend.reset_count >= 1 + + +def test_streaming_disconnect_resets_once(monkeypatch): + class _DisconnectRequest(_Request): + async def is_disconnected(self): + return True + + backend = _ScriptedBackend(_fixed("a", "ab", "abc")) + payload = _request(tools = [LOOKUP_TOOL], stream = True) + _install(monkeypatch, backend) + + async def _run(): + resp = await openai_chat_completions( + payload, request = _DisconnectRequest(), current_subject = "u" + ) + return [c async for c in resp.body_iterator] + + asyncio.run(_run()) + assert backend.reset_count == 1 + + +def test_mlx_uses_same_path(monkeypatch): + # MLX and safetensors share get_inference_backend(); one scripted backend covers both. + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + assert body["choices"][0]["finish_reason"] == "tool_calls" + + +def test_tool_choice_none_does_not_advertise_tools(monkeypatch): + # tool_choice="none": no tools rendered into the template; history templating still applies. + backend = _ScriptedBackend(_fixed("plain answer")) + payload = _request(tools = [LOOKUP_TOOL], tool_choice = "none", stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + assert body["choices"][0]["message"]["content"] == "plain answer" + assert backend.calls[0]["tools"] is None + + +def test_developer_message_folded_into_system_prompt(monkeypatch): + # The "developer" role folds into one leading system message (local templates reject it). + backend = _ScriptedBackend(_fixed("ok")) + payload = _request( + messages = [ + ChatMessage(role = "developer", content = "always be terse"), + ChatMessage(role = "user", content = "hi"), + ], + tools = [LOOKUP_TOOL], + stream = False, + ) + _call(payload, monkeypatch, backend) + sent = backend.calls[0]["messages"] + assert sent[0]["role"] == "system" + assert "always be terse" in sent[0]["content"] + assert all(m.get("role") != "developer" for m in sent) + + +def test_failed_nudge_retry_keeps_original_response(monkeypatch): + # A raising retry must not 500; the first response is returned. + state = {"n": 0} + + def responder(messages, tools): + state["n"] += 1 + if state["n"] == 1: + return ['{"name":"lookup"'] # unhealable signal + raise RuntimeError("retry blew up") + + backend = _ScriptedBackend(responder) + payload = _request(tools = [LOOKUP_TOOL], nudge_tool_calls = True, stream = False) + body = _json_body(_call(payload, monkeypatch, backend)) + assert state["n"] == 2 + assert body["choices"][0]["finish_reason"] == "stop" + assert body["choices"][0]["message"]["content"] == '{"name":"lookup"' + + +def test_discarded_nudge_retry_reports_first_attempt_usage(monkeypatch): + # Double-failure nudge: the first response is delivered, but the retry's + # generate() overwrites stats_holder. The monitor must record the FIRST + # attempt's usage, not the discarded retry's. + first_stats = {"usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}} + retry_stats = {"usage": {"prompt_tokens": 99, "completion_tokens": 99, "total_tokens": 198}} + + class _PerCallStatsBackend(_ScriptedBackend): + def __init__(self): + # Unhealable truncated markup on both attempts -> retry is discarded. + super().__init__(lambda m, t: ['{"name":"lookup"']) + self._stats_seq = [first_stats, retry_stats] + + def generate_chat_response( + self, + *, + messages, + tools = None, + stats_holder = None, + **kwargs, + ): + self.calls.append({"messages": messages, "tools": tools, **kwargs}) + stats = self._stats_seq[min(len(self.calls) - 1, len(self._stats_seq) - 1)] + if stats_holder is not None: + stats_holder["stats"] = stats + for snap in self._responder(messages, tools): + yield snap + + backend = _PerCallStatsBackend() + payload = _request(tools = [LOOKUP_TOOL], nudge_tool_calls = True, stream = False) + monitor = _install(monkeypatch, backend) + + async def _run(): + return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + + asyncio.run(_run()) + assert len(backend.calls) == 2 # first attempt + one discarded retry + [entry] = monitor.snapshot() + # The delivered response is the first attempt, so its usage must be reported. + assert entry["prompt_tokens"] == 7 + assert entry["completion_tokens"] == 3 + + +def test_monitor_records_healed_call_not_raw_xml(monkeypatch): + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False) + monitor = _install(monkeypatch, backend) + + async def _run(): + return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + + asyncio.run(_run()) + snap = monitor.snapshot(include_details = True) + replies = json.dumps(snap) + assert "" not in replies + assert "lookup" in replies + + +def test_streaming_monitor_records_healed_call_not_raw_xml(monkeypatch): + # Monitor mirrors what the client received, never the healed-away raw markup. + backend = _ScriptedBackend( + _fixed("Sure. ", 'Sure. {"name": "loo', "Sure. " + _CALL_XML) + ) + payload = _request(tools = [LOOKUP_TOOL], stream = True) + monitor = _install(monkeypatch, backend) + + async def _run(): + return await openai_chat_completions(payload, request = _Request(), current_subject = "u") + + response = asyncio.run(_run()) + _collect_sse(response) + replies = json.dumps(monitor.snapshot(include_details = True)) + assert "" not in replies + assert "Sure. " in replies + assert "[tool_calls] lookup(" in replies + + +def test_forced_tool_choice_narrows_templated_tools(monkeypatch): + # A forced function is the only schema rendered into the template. + backend = _ScriptedBackend(_fixed(_SEARCH_XML)) + payload = _request( + tools = [LOOKUP_TOOL, SEARCH_TOOL], + stream = False, + tool_choice = {"type": "function", "function": {"name": "search"}}, + ) + body = _json_body(_call(payload, monkeypatch, backend)) + templated = backend.calls[0]["tools"] + assert [t["function"]["name"] for t in templated] == ["search"] + choice = body["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["tool_calls"][0]["function"]["name"] == "search" + + +def test_multimodal_content_parts_flattened_for_local_template(monkeypatch): + # Remote image URLs leave image=None, so content arrives as a part LIST: + # text parts are kept, the image part dropped. + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request( + messages = [ + ChatMessage( + role = "user", + content = [ + {"type": "text", "text": "what is this?"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/cat.png"}, + }, + ], + ) + ], + tools = [LOOKUP_TOOL], + stream = False, + ) + body = _json_body(_call(payload, monkeypatch, backend)) + templated = backend.calls[0]["messages"] + assert all(isinstance(m.get("content"), str) for m in templated) + assert any(m["content"] == "what is this?" for m in templated) + assert body["choices"][0]["finish_reason"] == "tool_calls" + + +def test_string_arguments_history_deserialized_for_template(monkeypatch): + # JSON-string tool_calls arguments become dicts in the templated copy; + # the HTTP response stays OpenAI-shaped. + backend = _ScriptedBackend(_fixed("done")) + payload = _request( + tools = [LOOKUP_TOOL], + stream = False, + messages = [ + ChatMessage(role = "user", content = "weather?"), + ChatMessage( + role = "assistant", + content = None, + tool_calls = [ + { + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": '{"q": "weather"}'}, + } + ], + ), + ChatMessage(role = "tool", tool_call_id = "call_0", content = "sunny"), + ], + ) + _json_body(_call(payload, monkeypatch, backend)) + assistant = next(m for m in backend.calls[0]["messages"] if m["role"] == "assistant") + assert assistant["tool_calls"][0]["function"]["arguments"] == {"q": "weather"} + + +def test_unparseable_arguments_string_left_untouched(monkeypatch): + backend = _ScriptedBackend(_fixed("ok")) + payload = _request( + tools = [LOOKUP_TOOL], + stream = False, + messages = [ + ChatMessage(role = "user", content = "hi"), + ChatMessage( + role = "assistant", + content = None, + tool_calls = [ + { + "id": "call_0", + "type": "function", + "function": {"name": "lookup", "arguments": "not json {"}, + } + ], + ), + ChatMessage(role = "tool", tool_call_id = "call_0", content = "y"), + ], + ) + body = _json_body(_call(payload, monkeypatch, backend)) + assert body["choices"][0]["message"]["content"] == "ok" + assistant = next(m for m in backend.calls[0]["messages"] if m["role"] == "assistant") + assert assistant["tool_calls"][0]["function"]["arguments"] == "not json {" + + +def test_mcp_enabled_without_server_tools_uses_passthrough(monkeypatch): + # mcp_enabled=true with an empty registry must not silently drop the + # declared tools; the gate keys on the server-side path claiming the request. + backend = _ScriptedBackend(_fixed(_CALL_XML)) + payload = _request(tools = [LOOKUP_TOOL], stream = False, mcp_enabled = True) + body = _json_body(_call(payload, monkeypatch, backend)) + choice = body["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup" + assert backend.calls[0]["tools"] == [LOOKUP_TOOL] diff --git a/studio/backend/tests/test_think_prefill_reemit.py b/studio/backend/tests/test_think_prefill_reemit.py new file mode 100644 index 0000000000..300ff92776 --- /dev/null +++ b/studio/backend/tests/test_think_prefill_reemit.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Unit tests for detect_think_prefill. + +Reasoning templates (Qwen3.6-style) end the generation prompt with an open +``\\n`` so the model starts reasoning immediately. skip_prompt +streaming drops that opening tag, so the safetensors/MLX paths must re-emit +it for the frontend's parser to render a thinking block. +""" + +import os +import sys + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from core.inference.chat_template_helpers import detect_think_prefill + + +QWEN_PROMPT = "<|im_start|>user\nHi!<|im_end|>\n<|im_start|>assistant\n" + + +def test_open_think_prefill_reemitted(): + """Qwen3.6-style enable_thinking=True prompt tail: \\n.""" + assert detect_think_prefill(QWEN_PROMPT + "\n") == "\n" + + +def test_bare_open_think_prefill_reemitted(): + """Prefill without trailing newline still detected.""" + assert detect_think_prefill(QWEN_PROMPT + "") == "" + + +def test_closed_think_prefill_not_reemitted(): + """enable_thinking=False prefills a closed, empty think block.""" + assert detect_think_prefill(QWEN_PROMPT + "\n\n\n\n") == "" + + +def test_prompt_without_think_untouched(): + """Non-reasoning templates produce no prefix.""" + assert detect_think_prefill(QWEN_PROMPT) == "" + + +def test_historical_think_blocks_ignored(): + """A closed think block in a prior assistant turn (preserve_thinking) + must not trigger re-emission when the generation tail is plain.""" + prompt = ( + "<|im_start|>user\nHi!<|im_end|>\n" + "<|im_start|>assistant\n\nprior reasoning\n\n\nHello!<|im_end|>\n" + "<|im_start|>user\nAgain?<|im_end|>\n<|im_start|>assistant\n" + ) + assert detect_think_prefill(prompt) == "" + + +def test_historical_blocks_plus_open_prefill(): + """Prior closed blocks plus a fresh open prefill: only the tail matters.""" + prompt = ( + "<|im_start|>assistant\n\nprior\n\n\nHello!<|im_end|>\n" + "<|im_start|>assistant\n\n" + ) + assert detect_think_prefill(prompt) == "\n" + + +def test_content_after_open_tag_not_reemitted(): + """If non-whitespace follows the tag it is not a plain prefill.""" + assert detect_think_prefill(QWEN_PROMPT + "\npartial reasoning") == "" + + +def test_empty_and_none_prompts(): + assert detect_think_prefill("") == "" + assert detect_think_prefill(None) == "" + + +def test_guard_suppresses_when_close_tag_is_special(): + """If is a special token, skip_special_tokens strips the model's + close tag, so re-emitting the open would leave an unclosed block. Guard off.""" + specials = ["<|im_end|>", "", ""] + assert detect_think_prefill(QWEN_PROMPT + "\n", specials) == "" + + +def test_guard_emits_when_think_not_special(): + specials = ["<|im_end|>", "<|endoftext|>"] + assert detect_think_prefill(QWEN_PROMPT + "\n", specials) == "\n" + + +def test_guard_default_and_empty_keep_emitting(): + assert detect_think_prefill(QWEN_PROMPT + "\n", None) == "\n" + assert detect_think_prefill(QWEN_PROMPT + "\n", []) == "\n" diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index 39fdd151be..c6da1e90e7 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -72,10 +72,8 @@ class TestFunctionStyleTrailingText: 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. + # Regression: the heal path (allow_incomplete=True) must match the strict path -- + # keep a clean argument and leave trailing prose outside the call span. text = "cats trailing words" calls = parse_tool_calls_from_text(text, allow_incomplete = True) assert len(calls) == 1 @@ -102,6 +100,24 @@ class TestFunctionStyleTrailingText: text = "weather london" assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + def test_attribute_form_literal_close_tag_is_preserved(self): + # The attribute form (MiniCPM-5 / MiniMax-M2) also ends at the + # LAST , so a literal close tag inside a code argument survives. + text = ( + '' + 'print("")' + " all done" + ) + call = _only(text) + assert call == {"name": "python", "arguments": {"code": 'print("")'}} + + def test_closed_zero_param_attribute_call_is_accepted_in_strict_mode(self): + # A closed call with no parameters is a valid zero-argument call; strict + # mode must not treat the empty parameter list as a truncated call. + assert _only('') == {"name": "ping", "arguments": {}} + # A no-arg call that never closes is still rejected as truncated. + assert parse_tool_calls_from_text('', allow_incomplete = False) == [] + class TestParityWithJsonStyle: def test_json_tool_call_with_trailing_prose_is_accepted(self): @@ -176,6 +192,37 @@ class TestGemmaNativeStyle: } +class TestLlama3PythonTagStrict: + def test_closed_dot_call_is_accepted(self): + text = '<|python_tag|>get_weather.call(location="Tokyo")' + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_weather" + assert json.loads(calls[0]["function"]["arguments"]) == {"location": "Tokyo"} + + def test_truncated_dot_call_is_rejected(self): + # No closing paren (depth > 0 at EOF): truncated, reject in strict mode. + text = '<|python_tag|>get_weather.call(location="Tokyo"' + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + # Auto-Heal still recovers it. + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + +class TestMistralArrayStrict: + def test_closed_array_is_accepted(self): + text = '[TOOL_CALLS] [{"name":"web_search","arguments":{"q":"x"}}]' + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "web_search" + + def test_unclosed_array_is_rejected(self): + # Missing the closing ]; strict mode must not heal it. + text = '[TOOL_CALLS] [{"name":"web_search","arguments":{"q":"x"}}' + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + # Auto-Heal still recovers the object by hand. + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + class TestHealingPathUnaffected: def test_auto_heal_still_repairs_unclosed_function(self): text = "cats" @@ -184,9 +231,8 @@ class TestHealingPathUnaffected: 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. + # A call that DID close must parse identically to strict mode, leaving prose after + # out of the last parameter and the removal span. from core.tool_healing import parse_tool_calls_from_text as parse_with_spans text = "cats trailing" @@ -197,3 +243,1579 @@ class TestHealingPathUnaffected: assert text[span[0] : span[1]] == ( "cats" ) + + def test_wrapperless_fallback_calls_carry_spans(self): + # The wrapperless function-XML fallback must report spans too, so with_spans + # consumers strip exactly the promoted markup (through when closed). + from core.tool_healing import parse_tool_calls_from_text as parse_with_spans + + closed = "before cats after" + calls, spans = parse_with_spans(closed, allow_incomplete = True, with_spans = True) + (call,) = calls + assert json.loads(call["function"]["arguments"]) == {"query": "cats"} + (span,) = spans + assert closed[span[0] : span[1]] == ( + "cats" + ) + + healed = "x dogs" + calls, spans = parse_with_spans(healed, allow_incomplete = True, with_spans = True) + (call,) = calls + assert json.loads(call["function"]["arguments"]) == {"query": "dogs"} + (span,) = spans + assert healed[span[0] : span[1]] == "dogs" + + +class TestEnabledToolNameGate: + """``enabled_tool_names`` disambiguates the ambiguous bare-rehearsal + ``NAME[ARGS]{json}`` form (#5704): NAME is a call only when it is an active tool, + otherwise it is prose. ``None`` (the default) keeps the legacy unrestricted parse + so existing callers are unaffected.""" + + def _names(self, calls): + return [c["function"]["name"] for c in calls] + + def test_inactive_rehearsal_before_active_call_does_not_swallow_it(self): + # P1: an inactive ``foo[ARGS]{...}`` before a real call must not consume the real call. + text = 'foo[ARGS]{"a":1} web_search[ARGS]{"query":"cats"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert self._names(calls) == ["web_search"] + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_inactive_rehearsal_alone_is_not_a_call(self): + text = 'foo[ARGS]{"a":1}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + def test_active_rehearsal_is_still_parsed(self): + text = 'web_search[ARGS]{"query":"cats"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert self._names(calls) == ["web_search"] + + def test_unrestricted_gate_none_preserves_legacy_behavior(self): + # Without a gate every ``NAME[ARGS]{...}`` is parsed, as before the gate landed. + text = 'foo[ARGS]{"a":1} web_search[ARGS]{"query":"cats"}' + assert self._names(parse_tool_calls_from_text(text)) == ["foo", "web_search"] + assert self._names(parse_tool_calls_from_text(text, enabled_tool_names = None)) == [ + "foo", + "web_search", + ] + + +class TestBracketCallSpans: + """with_spans tiling for Mistral bracket calls: promoted markup strips + exactly once, filtered calls' bytes stay visible, closers strip too.""" + + def test_mixed_array_filtered_first_keeps_its_bytes_only(self): + from core.inference.passthrough_healing import heal_openai_message_events + + tools = [{"type": "function", "function": {"name": "lookup", "parameters": {}}}] + content = ( + '[TOOL_CALLS][{"name":"bad","arguments":{"x":1}},' + '{"name":"lookup","arguments":{"q":"cats"}}]' + ) + events = heal_openai_message_events( + {"role": "assistant", "content": content}, {"lookup"}, tools + ) + kinds = [k for k, _v in events] + assert kinds == ["text", "tool_call"] + text = events[0][1] + assert '"bad"' in text + # The promoted call's markup must not survive in the text event. + assert '"lookup"' not in text + + def test_mixed_array_filtered_second_stays_visible(self): + from core.inference.passthrough_healing import heal_openai_message_events + + tools = [{"type": "function", "function": {"name": "lookup", "parameters": {}}}] + content = ( + '[TOOL_CALLS][{"name":"lookup","arguments":{"q":"cats"}},' + '{"name":"bad","arguments":{"x":1}}]' + ) + events = heal_openai_message_events( + {"role": "assistant", "content": content}, {"lookup"}, tools + ) + assert events[0][0] == "tool_call" + trailing = "".join(v for k, v in events if k == "text") + assert '"bad"' in trailing + + def test_v11_closer_inside_span(self): + from core.tool_healing import parse_tool_calls_from_text as parse_with_spans + + text = '[TOOL_CALLS]web_search[ARGS]{"query":"cats"}[/TOOL_CALLS] after' + calls, spans = parse_with_spans(text, allow_incomplete = True, with_spans = True) + (call,) = calls + assert call["function"]["name"] == "web_search" + (span,) = spans + assert text[span[0] : span[1]].endswith("[/TOOL_CALLS]") + assert text[span[1] :] == " after" + + def test_fully_promoted_array_strips_whole_region(self): + from core.inference.passthrough_healing import heal_openai_message_events + + tools = [{"type": "function", "function": {"name": "lookup", "parameters": {}}}] + content = ( + '[TOOL_CALLS][{"name":"lookup","arguments":{"q":"a"}},' + '{"name":"lookup","arguments":{"q":"b"}}] after' + ) + events = heal_openai_message_events( + {"role": "assistant", "content": content}, {"lookup"}, tools + ) + assert [k for k, _v in events] == ["tool_call", "tool_call", "text"] + assert events[2][1] == " after" + + +class TestMistralArrayHealing: + """Draining the whole [TOOL_CALLS] array for the shapes the repo's own + Mistral/Ollama templates emit.""" + + def test_comma_less_multi_call_array_parses_all_calls(self): + # ollama_template_mappers.py renders multi-call turns as [{...}{...}] with no + # comma separator; a single json.loads of the body rejects it and dropped every + # call. The element-by-element decode must recover all of them. + text = '[TOOL_CALLS] [{"name":"a","arguments":{"x":1}}{"name":"b","arguments":{"y":2}}]' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["a", "b"] + assert json.loads(calls[0]["function"]["arguments"]) == {"x": 1} + assert json.loads(calls[1]["function"]["arguments"]) == {"y": 2} + + def test_comma_separated_and_single_arrays_still_parse(self): + both = parse_tool_calls_from_text( + '[TOOL_CALLS] [{"name":"a","arguments":{}},{"name":"b","arguments":{}}]' + ) + assert [c["function"]["name"] for c in both] == ["a", "b"] + one = parse_tool_calls_from_text('[TOOL_CALLS] [{"name":"a","arguments":{}}]') + assert [c["function"]["name"] for c in one] == ["a"] + + def test_mistral_array_null_arguments_normalized_to_empty_object(self): + # ``"arguments": null`` is a no-arg call; it must become {} (as the + # path does), not the string "null" that auto-heal turns into {"query":"null"}. + calls = parse_tool_calls_from_text('[TOOL_CALLS][{"name":"get_time","arguments":null}]') + assert calls[0]["function"]["arguments"] == "{}" + + +class TestGlmStrict: + def test_closed_glm_call_is_accepted(self): + text = ( + "get_weather\n" + "city\nParis\n" + "" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "get_weather" + + def test_unclosed_glm_call_is_rejected(self): + # No close: truncated, reject with Auto-Heal off. + text = "get_weather\ncity\nParis" + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + +class TestKimiStrict: + _SB = "<|tool_calls_section_begin|>" + _KB = "<|tool_call_begin|>" + _AB = "<|tool_call_argument_begin|>" + _KE = "<|tool_call_end|>" + _SE = "<|tool_calls_section_end|>" + + def test_full_kimi_call_is_accepted(self): + text = self._SB + self._KB + "functions.x:0" + self._AB + '{"a":1}' + self._KE + self._SE + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "x" + + def test_kimi_call_without_call_end_is_rejected(self): + # Section closed but the call lacks <|tool_call_end|>: reject in strict. + text = self._SB + self._KB + "functions.x:0" + self._AB + '{"a":1}' + self._SE + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + def test_kimi_without_section_end_is_rejected(self): + # No <|tool_calls_section_end|>: truncated section, reject in strict. + text = self._SB + self._KB + "functions.x:0" + self._AB + '{"a":1}' + self._KE + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + +class TestParserLinearity: + """Llama-3 ``.call`` kwargs and Mistral-array healing must stay linear (a regex-per-offset blew up on long truncated bodies).""" + + def test_llama3_unterminated_call_arg_is_linear(self): + import time + + text = '<|python_tag|>upload.call(data="' + "A" * 200_000 # no closing quote/paren + t0 = time.perf_counter() + parse_tool_calls_from_text(text, allow_incomplete = True) + assert time.perf_counter() - t0 < 2.0 + + def test_llama3_huge_wordrun_call_arg_is_linear(self): + import time + + text = "<|python_tag|>upload.call(" + "a" * 200_000 # giant word run, no '=' + t0 = time.perf_counter() + parse_tool_calls_from_text(text, allow_incomplete = True) + assert time.perf_counter() - t0 < 2.0 + + def test_mistral_unclosed_array_open_braces_is_linear(self): + import time + + text = "[TOOL_CALLS] [" + "{" * 200_000 # unclosed array, all open braces + t0 = time.perf_counter() + parse_tool_calls_from_text(text, allow_incomplete = True) + assert time.perf_counter() - t0 < 2.0 + + def test_gemma_wrapperless_deep_nesting_is_linear(self): + # Wrapper-less Gemma ``call:f{a:{a:{...}}}`` deep nesting must parse in linear time (no quadratic re-scan). + import time + + def nested(d): + return "call:f{a:" + "{a:" * d + "x:1" + "}" * d + "}" + + def best_ms(depth): + text = nested(depth) + best = float("inf") + for _ in range(5): + t0 = time.perf_counter() + calls = parse_tool_calls_from_text(text) + best = min(best, time.perf_counter() - t0) + assert calls and json.loads(calls[0]["function"]["arguments"]), "nested args dropped" + return best + + t200 = best_ms(200) + t400 = best_ms(400) + assert t400 < t200 * 3.0, (t200, t400) + + def test_llama3_call_kwargs_still_parse(self): + text = '<|python_tag|>do.call(s="hi 😀", n=42, f=1.5, b=true, z=null)' + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + assert json.loads(calls[0]["function"]["arguments"]) == { + "s": "hi 😀", + "n": 42, + "f": 1.5, + "b": True, + "z": None, + } + + def test_llama3_call_scientific_notation_args_parse(self): + # Scientific notation must decode as float (the old regex truncated 1e-3 -> 1). + text = "<|python_tag|>calc.call(x=1e-3, y=-2E+4, z=0.5e2, n=42)" + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"x": 1e-3, "y": -2e4, "z": 50.0, "n": 42} + assert isinstance(args["n"], int) and isinstance(args["x"], float) + + def test_mistral_unclosed_array_recovers_top_level_objects(self): + text = ( + '[TOOL_CALLS] [{"name":"a","arguments":{"k":1}},' + '{"name":"b","arguments":{"j":2}}' # missing closing ] + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert [c["function"]["name"] for c in calls] == ["a", "b"] + + +class TestLlamaBuiltinChainAndNesting: + """Llama-3 ``.call`` built-ins: ``; `` chaining and nested-tag isolation.""" + + def test_semicolon_chained_builtin_calls_all_parse(self): + # Only the first call is anchored to <|python_tag|>; the rest chain via ';'. + text = "<|python_tag|>alpha.call(x=1); beta.call(y=2); gamma.call(z=3)" + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert [c["function"]["name"] for c in calls] == ["alpha", "beta", "gamma"] + assert json.loads(calls[1]["function"]["arguments"]) == {"y": 2} + + def test_nested_python_tag_in_json_string_arg_is_not_a_call(self): + # A code arg literally containing a <|python_tag|>...call(...) string: the real call is the + # outer "python", not the nested "os" -- the scan stays anchored to the first tag. + text = ( + '<|python_tag|>{"name":"python","parameters":' + '{"code":"<|python_tag|>os.call(\'rm -rf /\')"}}' + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "python" + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "<|python_tag|>os.call('rm -rf /')" + + def test_single_builtin_call_unchanged(self): + text = '<|python_tag|>web_search.call(query="cats")' + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "web_search" + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + +def test_glm_open_does_not_parse_spaced_prose_as_tool_name(): + # The GLM NAME opener must reject spaced literal prose (V10); only a + # valid [\w.\-]+ name (followed by newline//) is a call. + assert parse_tool_calls_from_text("not a call") == [] + ok = parse_tool_calls_from_text( + "get_weather\ncity\nNYC\n" + ) + assert [c["function"]["name"] for c in ok] == ["get_weather"] + + +def test_deepseek_r1_missing_call_terminator_rejected_in_strict_mode(): + # R1 must reject a fenced call whose closing ``` + <｜tool▁call▁end｜> never + # arrived when Auto-Heal is off, matching V3/V3.1 strictness (V6). + text = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n" + "```json\n" + '{"city":"NYC"}' + "<｜tool▁calls▁end｜>" + ) + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + assert len(parse_tool_calls_from_text(text, allow_incomplete = True)) == 1 + + +def test_deepseek_r1_complete_call_accepted_in_strict_mode(): + # A fully-terminated R1 call (close fence + per-call end) is still accepted. + text = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n" + "```json\n" + '{"city":"NYC"}\n' + "```<｜tool▁call▁end｜><｜tool▁calls▁end｜>" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 and calls[0]["function"]["name"] == "get_weather" + + +def test_strip_leading_bare_json_call_drops_complete_call(): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + # A complete Llama-3.2 bare-JSON call is removed; trailing prose is kept. + assert strip_leading_bare_json_call('{"name":"web_search","parameters":{"query":"cats"}}') == "" + assert ( + strip_leading_bare_json_call('{"name":"python","parameters":{"code":"x"}} done') == "done" + ) + + +def test_strip_leading_bare_json_call_drops_truncated_call(): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + # A truncated call (no closing brace) collapses to "" -- nothing recoverable. + assert ( + strip_leading_bare_json_call('{"name":"web_search","parameters":{"query":"weather in S') + == "" + ) + + +def test_strip_leading_bare_json_call_preserves_plain_json_and_prose(): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + # No "name" key -> plain JSON answer, left untouched. + assert ( + strip_leading_bare_json_call('{"result": 42, "ok": true}') == '{"result": 42, "ok": true}' + ) + # Prose before the brace -> not a leading bare call, untouched. + assert strip_leading_bare_json_call('here is {"name":"x"}') == 'here is {"name":"x"}' + # Ordinary text untouched. + assert strip_leading_bare_json_call("just a sentence.") == "just a sentence." + + +def test_glm_literal_close_tag_in_string_arg_not_truncated(): + import json + + from core.inference.tool_call_parser import parse_tool_calls_from_text + + # A GLM string argument may legitimately contain the literal close tag ````. + text = ( + "run_code\n" + "code\n" + 'print("")\n' + "" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == 'print("")', args + + +def test_glm_truncated_block_rejected_in_strict_mode_but_healed_otherwise(): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + # No close: strict mode (Auto-Heal off) rejects the truncated + # block; with Auto-Heal it keeps the partial call. + text = "get_weather\ncity\nNYC" + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + healed = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(healed) == 1 and healed[0]["function"]["name"] == "get_weather" + + +def test_truncated_wrapperless_gemma_call_is_stripped(): + from core.inference.tool_call_parser import strip_tool_markup + + # A wrapper-less Gemma ``call:NAME{...`` cut off mid-arguments (no closing + # brace) must not leak the raw call into the visible stream. + text = 'Sure!\ncall:web_search{"query": "weather in San Fr' + stripped = strip_tool_markup(text, final = True) + assert "call:web_search" not in stripped, repr(stripped) + assert stripped.strip() == "Sure!" + + +def test_complete_wrapperless_gemma_call_keeps_trailing_prose(): + from core.inference.tool_call_parser import strip_tool_markup + + # The truncation pattern must run AFTER the closed form, so a complete call + # followed by prose keeps the prose instead of eating to EOS. + text = 'call:web_search{"query": "cats"} Here you go.' + stripped = strip_tool_markup(text, final = True) + assert "call:web_search" not in stripped + assert stripped.strip() == "Here you go." + + +def test_bare_json_gated_on_enabled_tool_names(): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + alice = '{"name":"Alice","parameters":{"age":30}}' + real = '{"name":"web_search","parameters":{"query":"cats"}}' + # With an enabled set, markerless JSON whose name is not a tool is NOT a call. + assert parse_tool_calls_from_text(alice, enabled_tool_names = {"web_search"}) == [] + # A real call (enabled name) still parses. + got = parse_tool_calls_from_text(real, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in got] == ["web_search"] + # No enabled set (None) keeps the name-agnostic behaviour for direct callers. + assert [c["function"]["name"] for c in parse_tool_calls_from_text(alice)] == ["Alice"] + # Marker-based forms are NOT gated (an explicit signal is a real call attempt). + xml = '{"name":"Alice","arguments":{}}' + assert parse_tool_calls_from_text(xml, enabled_tool_names = {"web_search"}) + + +def test_strip_leading_bare_json_call_gated_on_enabled_tool_names(): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + alice = '{"name":"Alice","parameters":{"age":30}}' + # Not an enabled tool -> ordinary JSON answer, kept verbatim. + assert strip_leading_bare_json_call(alice, {"web_search"}) == alice + # Enabled tool -> a real call, stripped (trailing prose kept). + assert ( + strip_leading_bare_json_call( + '{"name":"web_search","parameters":{"q":1}} hi', {"web_search"} + ) + == "hi" + ) + + +def test_function_xml_strip_keeps_literal_close_tag_in_param_value(): + from core.inference.tool_call_parser import strip_tool_markup + + # The strip uses the LAST (like the parser) so a literal in a value doesn't + # truncate it; separate calls still strip independently. + text = 'print("") done' + assert strip_tool_markup(text, final = True) == "done" + two = ( + "a 1 mid " + "2 end" + ) + assert strip_tool_markup(two, final = True) == "a mid end" + + +def test_function_xml_strip_keeps_trailing_text_after_literal_open_tag(): + from core.inference.tool_call_parser import parse_tool_calls_from_text, strip_tool_markup + + # A literal ```` opener inside a parameter value is data, not a call: the scan-based + # strip keeps " done" (the old negative-lookahead regex ate the trailing prose). + text = 'print("") done' + assert parse_tool_calls_from_text(text)[0]["function"]["name"] == "python" + assert strip_tool_markup(text, final = True) == "done" + # Non-final (streaming) keeps an unclosed call buffered, does not eat prose early. + open_text = 'pre print("")' + assert strip_tool_markup(open_text, final = False) == open_text + + +def test_final_strip_removes_magistral_think_reasoning(): + from core.inference.tool_call_parser import strip_tool_markup + + # Magistral emits reasoning as ``[THINK]...[/THINK]`` (bracket form, not ````); + # at end-of-turn it must be dropped so it doesn't leak into display / history. + text = "[THINK]The user greeted me, I should say hi.[/THINK]Hello! How can I help?" + assert strip_tool_markup(text, final = True) == "Hello! How can I help?" + # A ``[TOOL_CALLS]`` living inside the reasoning goes with it. + with_call = '[THINK]Maybe I should search.[/THINK][TOOL_CALLS]search{"q":"x"}' + assert strip_tool_markup(with_call, final = True) == "" + + +def test_streaming_strip_keeps_magistral_think_buffered(): + from core.inference.tool_call_parser import strip_tool_markup + + # Mid-stream (final=False) the reasoning block is left intact; only the + # end-of-turn pass removes it. + text = "[THINK]still thinking" + assert strip_tool_markup(text, final = False) == text + + +def test_final_strip_leaves_non_magistral_bracket_text_untouched(): + from core.inference.tool_call_parser import strip_tool_markup + + # Only a LEADING ``[THINK]`` block is reasoning; unrelated bracketed prose stays. + text = "See [THINK about it] later" + assert strip_tool_markup(text, final = True) == "See [THINK about it] later" + + +def test_strip_leading_bare_json_call_ignores_nested_name(): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + # A nested ``"name"`` must NOT gate the strip (only a TOP-LEVEL enabled name is a call); the + # ordinary JSON answer is kept verbatim, truncated or complete. + nested_trunc = '{"result":{"name":"web_search","age":' + nested_full = '{"result":{"name":"web_search","age":1}}' + assert strip_leading_bare_json_call(nested_trunc, {"web_search"}) == nested_trunc + assert strip_leading_bare_json_call(nested_full, {"web_search"}) == nested_full + # A real top-level call (even with a top-level array before the name) still strips. + assert ( + strip_leading_bare_json_call( + '{"data":[1,2],"name":"web_search","parameters":{}}', {"web_search"} + ) + == "" + ) + + +def test_mistral_single_object_call_is_stripped_for_display(): + from core.inference.tool_call_parser import ( + _strip_mistral_closed_calls, + parse_tool_calls_from_text, + ) + + # The parser accepts the single-object [TOOL_CALLS]{...} shape, so the display + # strip must remove it too (asymmetry would leak the raw object). + text = '[TOOL_CALLS]{"name":"web_search","arguments":{"filters":{"date":"2024"}}} tail' + assert [c["function"]["name"] for c in parse_tool_calls_from_text(text)] == ["web_search"] + assert _strip_mistral_closed_calls(text) == " tail" + # A literal [TOOL_CALLS] in prose (no following object) is left untouched. + assert _strip_mistral_closed_calls("See the [TOOL_CALLS] docs") == "See the [TOOL_CALLS] docs" + + +def test_tool_call_parser_declares_future_annotations_for_py39_import(): + # F1: the parser is imported standalone on python >=3.9, where its PEP 604 ``X | None`` + # annotations need ``from __future__ import annotations``; guard that the import stays. + from pathlib import Path + src = ( + Path(__file__).resolve().parent.parent / "core" / "inference" / "tool_call_parser.py" + ).read_text() + assert "from __future__ import annotations" in src + + +def test_glm_strip_treats_literal_close_tag_in_arg_value_as_data(): + # Core strip parity: a literal inside a GLM is argument data, so the whole call is stripped (no leaked tail). + from core.inference.tool_call_parser import strip_tool_markup + + text = ( + "web_search\nquery\n" + "see tag\n tail" + ) + assert strip_tool_markup(text, final = True) == "tail" + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "see tag"} + + +def test_bare_json_function_alias_parses_and_strips_symmetrically(): + # The bare-JSON parser accepts the "function" alias for the call name; + # strip_leading_bare_json_call must recognise it too (parser/strip symmetry). + from core.inference.tool_call_parser import ( + parse_tool_calls_from_text, + strip_leading_bare_json_call, + _top_level_bare_json_name, + ) + + enabled = {"web_search"} + text = '{"function":"web_search","parameters":{"query":"cats"}}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = enabled) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert strip_leading_bare_json_call(text, enabled) == "" + + # "name" still takes precedence when both are present; nested aliases are data. + assert _top_level_bare_json_name('{"function":"foo","name":"web_search"}') == "web_search" + assert _top_level_bare_json_name('{"function":"web_search"}') == "web_search" + assert _top_level_bare_json_name('{"result":{"function":"web_search"}}') is None + # A non-enabled function-alias object is ordinary content and is preserved. + assert ( + strip_leading_bare_json_call('{"function":"not_a_tool","parameters":{}}', enabled) + == '{"function":"not_a_tool","parameters":{}}' + ) + + +class TestMistralOuterOverXmlLiteral: + """Quoted tool XML inside a [TOOL_CALLS] call's arguments is data; the outer call executes. Reverse order keeps the XML.""" + + def test_mistral_v11_arg_quoting_function_xml(self): + text = ( + '[TOOL_CALLS]web_search[ARGS]{"query":"literal ' + '1"}' + ) + for strict in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = not strict) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert "" in json.loads(calls[0]["function"]["arguments"])["query"] + + def test_mistral_array_arg_quoting_tool_call_json(self): + text = ( + '[TOOL_CALLS][{"name":"web_search","arguments":{"query":' + '"see {\\"name\\":\\"evil\\"}"}}]' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_xml_outer_keeps_winning_over_mistral_literal(self): + text = ( + '{"name":"web_search","arguments":' + '{"query":"docs say [TOOL_CALLS]evil[ARGS]{}"}}' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestHealerSignalAlignment: + """The healer buffers only formats its shared parser can promote. Mistral's + ``[TOOL_CALLS]`` is promotable (rescued), so it is a heal signal; the loop-only + text-call markers (Llama ``<|python_tag|>``, bare ``[ARGS]``) are not, so they + stream through instead of stalling as prose that never yields a call.""" + + def test_heal_signals_subset_of_promotable_formats(self): + from core.inference.passthrough_healing import _HEAL_SIGNALS + assert set(_HEAL_SIGNALS) == {"", "<|tool_call>", " is not a healer-promotable format, so it streams through as text. + events = list(healer.feed('<|python_tag|>web_search.call(query="cats")')) + text_out = "".join(v for k, v in events if k == "text") + assert "<|python_tag|>" in text_out # streamed through, not buffered + assert not list(healer.finalize()) or all(k == "text" for k, _v in healer.finalize()) + + +class TestGemmaWrapperlessLiteralMarkers: + """Wrapper-less Gemma calls whose ARGUMENTS mention Gemma's own markup. + + The tool_healing deferral must key on an actual wrapped opener + (``<|tool_call>call:...``), not the wrapper literal anywhere in content: + a query about the marker has nothing tool_healing can parse, and deferring + it loses the call entirely (not executed AND stripped from display).""" + + def test_marker_literal_in_argument_still_parses(self): + text = 'call:web_search{query:"what does <|tool_call> mean"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "what does <|tool_call> mean" + + def test_real_wrapped_call_still_deferred_to_tool_healing(self): + from core.inference.tool_call_parser import _parse_gemma_tool_calls + + # An actual wrapped opener present: the Gemma fallback must keep + # deferring to the shared tool_healing parser that owns that form. + text = '<|tool_call>call:web_search{query:<|"|>cats<|"|>}' + assert _parse_gemma_tool_calls(text, id_offset = 0) == [] + + def test_single_quoted_brace_does_not_truncate_code(self): + text = "call:python{code:print('}')}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "print('}')" + + def test_single_quoted_brace_strip_span_covers_whole_call(self): + from core.inference.tool_call_parser import strip_tool_markup + + text = "call:python{code:print('}')} Done." + stripped = strip_tool_markup(text, final = True, enabled_tool_names = {"python"}) + assert "call:python" not in stripped + assert "')}" not in stripped + assert stripped.strip() == "Done." + + +class TestGlmEmbeddedClosePair: + """A GLM value whose string literal embeds the full close-tag pair + ```` (code documenting the GLM format) must not be + truncated at the embedded pair: a structural close sits at balanced quote + state, an embedded one is inside an open string literal.""" + + def test_embedded_pair_inside_quoted_value_not_structural(self): + text = ( + "python\n" + "code\n" + 'print("")\nx = 1\n' + "" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == 'print("")\nx = 1' + + def test_strip_covers_the_full_call(self): + from core.inference.tool_call_parser import strip_tool_markup + + text = ( + "python\n" + "code\n" + 'print("")\nx = 1\n' + " Done." + ) + stripped = strip_tool_markup(text, final = True) + assert "arg_value" not in stripped + assert stripped.strip() == "Done." + + def test_unbalanced_apostrophe_falls_back_to_first_candidate(self): + # Prose-like value with an apostrophe: no candidate reaches balanced + # quote state, so the first token-valid close wins (prior behavior). + text = ( + "web_search\n" + "query\n" + "it's fine\n" + "" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "it's fine" + + +class TestPythonTagLiteralInsideMistralArgs: + """A python_tag LITERAL inside a leading Mistral call's arguments is data; the outer call executes.""" + + def test_mistral_arg_quoting_python_tag_call(self): + text = ( + '[TOOL_CALLS] [{"name": "web_search", "arguments": ' + '{"query": "what is <|python_tag|>evil.call(x=1)"}}]' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "what is <|python_tag|>evil.call(x=1)" + + +class TestPythonTagOuterOverXmlLiteral: + """A leading Llama-3 ``<|python_tag|>`` call owns the turn: tool XML/Mistral + markup quoted in a ``.call(...)`` string argument (or in trailing prose) is + data, so the outer call executes -- parity with the bare-JSON / Mistral / + attribute-form leading-ownership rules. XML before the tag keeps normal order.""" + + def test_call_arg_quoting_complete_function_xml(self): + # A closed in a .call() code arg must not beat the leading python_tag call. + text = ( + '<|python_tag|>python.call(code="' + '1")' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["python"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "1" + + def test_call_arg_quoting_bare_function_tag_in_query(self): + # A query mentioning must search, not execute a phantom tool. + text = '<|python_tag|>web_search.call(query="how do I use in llama")' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "how do I use in llama" + + def test_call_arg_quoting_tool_call_json(self): + text = ( + "<|python_tag|>save_file.call(content=" + '"{\\"name\\": \\"delete\\", \\"arguments\\": {}}")' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["save_file"] + + def test_json_form_code_arg_quoting_function_xml(self): + # JSON emission: a in the code arg is data; the outer "python" call runs. + text = ( + '<|python_tag|>{"name":"python","parameters":' + '{"code":"ls"}}' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["python"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "ls" + + def test_call_arg_quoting_mistral_trigger(self): + text = '<|python_tag|>web_search.call(query="see [TOOL_CALLS]evil[ARGS]{}")' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_leading_call_wins_over_trailing_xml(self): + # A leading python_tag call owns the turn even when a real XML literal follows. + text = ( + '<|python_tag|>web_search.call(query="cats") ' + "1" + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_xml_before_python_tag_keeps_xml_order(self): + # A foreign signal BEFORE the tag keeps normal document order (XML wins). + text = ( + "x " + '<|python_tag|>python.call(code="y")' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestBareJsonOuterOverXmlLiteral: + """Quoted tool XML inside a leading bare-JSON call is data; XML before the JSON keeps normal order.""" + + def test_bare_json_code_arg_quoting_function_xml(self): + text = ( + '{"name": "python", "arguments": ' + '{"code": "run() # ls"}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python"}) + assert [c["function"]["name"] for c in calls] == ["python"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "run() # ls" + + def test_bare_json_outer_unrestricted_mode(self): + text = '{"name": "python", "parameters": {"code": "ls"}}' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["python"] + + def test_xml_before_json_keeps_xml_order(self): + text = ( + "cats" + ' {"name": "python", "arguments": {"code": "x"}}' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestMagistralThinkRehearsal: + """A call rehearsed inside [THINK]...[/THINK] is reasoning; the real call after wins, and parse agrees with strip.""" + + def test_function_xml_rehearsal_in_think_is_not_promoted(self): + text = ( + '[THINK]I could emit {"query":"x"}' + ' here[/THINK][TOOL_CALLS] [{"name":"terminal","arguments":{"cmd":"ls"}}]' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["terminal"] + + def test_hermes_rehearsal_in_think_is_not_promoted(self): + text = ( + '[THINK]maybe {"name":"web_search","arguments":' + '{"query":"x"}}[/THINK]' + '[TOOL_CALLS] [{"name":"terminal","arguments":{"cmd":"ls"}}]' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["terminal"] + + def test_unclosed_think_parses_nothing(self): + text = '[THINK]let me try {"query":"x"}' + assert parse_tool_calls_from_text(text) == [] + + +class TestGemmaUnquotedApostrophes: + """Quotes open strings only at value-start context: an apostrophe inside + an unquoted wrapper-less value (contractions, possessives) is prose, and + treating it as an opener swallowed the closing brace and lost the call.""" + + def test_contraction_in_unquoted_query_parses(self): + text = "call:web_search{query:what's the weather}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "what's the weather" + + def test_contraction_does_not_swallow_next_key(self): + text = "call:web_search{query:what's up, n:3}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "what's up" + assert args["n"] == 3 + + def test_contraction_strip_span_covers_whole_call(self): + from core.inference.tool_call_parser import strip_tool_markup + + text = "call:web_search{query:what's the weather} Done." + stripped = strip_tool_markup(text, final = True, enabled_tool_names = {"web_search"}) + assert "call:web_search" not in stripped + assert stripped.strip() == "Done." + + def test_quoted_values_still_hide_delimiters(self): + text = 'call:web_search{query:"weather, location: Boston", n:2}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == "weather, location: Boston" + assert args["n"] == 2 + + +class TestGlmKeyWithoutValue: + """A GLM with no tag: strict mode rejects the call + (same contract as an unclosed value) instead of executing it with the + argument silently dropped; Auto-Heal keeps the lenient skip.""" + + def test_strict_rejects_key_without_value(self): + text = "web_search\nquery\n" + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + + def test_heal_keeps_the_lenient_skip(self): + text = "web_search\nquery\n" + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "web_search" + assert json.loads(calls[0]["function"]["arguments"]) == {} + + +class TestDisabledBareJsonLiteralNotPromoted: + """A leading non-enabled-name object is content: nothing inside promotes, and a call after it still parses.""" + + def test_literal_inside_disabled_json_stays_data(self): + text = ( + '{"name": "Alice", "note": "try ' + 'x"}' + ) + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + def test_python_tag_literal_inside_disabled_json_stays_data(self): + text = '{"name": "Alice", "note": "<|python_tag|>web_search.call(query=1)"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + def test_real_call_after_disabled_json_still_parses(self): + text = ( + '{"name": "Alice", "note": "x"} ' + '{"name": "web_search", "arguments": {"query": "cats"}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestDeepSeekMarkerInsideLeadingEnvelopes: + """A DeepSeek/Kimi marker quoted inside a leading bare-JSON or Mistral + call's argument strings is data: the pre-pass must not promote the + embedded no-arg literal and drop the real outer call.""" + + def test_marker_inside_leading_json_call_stays_data(self): + text = ( + '{"name": "web_search", "arguments": ' + '{"query": "what is <｜tool▁calls▁begin｜>...{}..."}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + args = json.loads(calls[0]["function"]["arguments"]) + assert "tool▁calls▁begin" in args["query"] + + def test_marker_inside_leading_mistral_call_stays_data(self): + text = ( + '[TOOL_CALLS] [{"name": "web_search", "arguments": ' + '{"query": "docs on <｜tool▁calls▁begin｜> markers"}}]' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_standalone_deepseek_call_still_parses(self): + text = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>web_search\n" + '```json\n{"query": "cats"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestMistralLiteralInsideLeadingJson: + """A [TOOL_CALLS] literal quoted inside a leading JSON object must not be promoted over it.""" + + def test_outer_json_call_wins_over_mistral_literal(self): + text = '{"name": "python", "arguments": {"code": "[TOOL_CALLS]web_search{}"}}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python", "web_search"}) + assert [c["function"]["name"] for c in calls] == ["python"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args["code"] == "[TOOL_CALLS]web_search{}" + + def test_disabled_outer_json_keeps_mistral_literal_as_data(self): + text = '{"name": "Alice", "note": "[TOOL_CALLS]web_search{}"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + +class TestGemmaWrappedWhitespace: + """Whitespace drift around ``call``/``:`` in wrapped Gemma calls must still parse (no fallback exists).""" + + def test_space_after_call_colon_parses(self): + text = '<|tool_call>call: web_search{query:<|"|>cats<|"|>}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_space_around_colon_parses(self): + text = '<|tool_call>call : web_search{query:<|"|>cats<|"|>}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_strict_mode_still_requires_the_closing_tag(self): + text = '<|tool_call>call: web_search{query:<|"|>cats<|"|>}' + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + + +class TestDisabledJsonBeforeDeepSeekCall: + """A disabled leading bare-JSON object whose strings mention a + DeepSeek/Kimi marker is dropped and the tail parsed, so a REAL + DeepSeek/Kimi call after the object still executes instead of the whole + message skipping the pre-pass.""" + + _DS = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>web_search\n" + '```json\n{"query": "cats"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + + def test_real_deepseek_call_after_disabled_json_parses(self): + text = '{"name": "Alice", "note": "<｜tool▁calls▁begin｜>"} ' + self._DS + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_disabled_json_with_marker_alone_stays_data(self): + text = '{"name": "Alice", "note": "<｜tool▁calls▁begin｜>"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + +class TestGemmaDottedArgumentKeys: + """Dotted Gemma keys (namespaced schemas) must survive key-quoting or the call is lost.""" + + def test_dotted_key_parses(self): + text = '<|tool_call>call:web_search{user.name:<|"|>bob<|"|>, query:<|"|>x<|"|>}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"user.name": "bob", "query": "x"} + + +class TestLeadingWrapperlessGemmaOverEmbeddedMarkers: + """A leading wrapper-less Gemma call to an enabled tool owns the turn: a + quoted foreign literal inside its argument (a query citing another tool + syntax) is data, and tool_healing must not promote it before the Gemma + fallback runs. Foreign markup leading keeps the normal order.""" + + def test_leading_gemma_wins_over_quoted_xml_literal(self): + text = ( + 'call:web_search{query:"explain ' + '{"name":"evil","arguments":{}}"}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_xml_leading_keeps_normal_order(self): + text = ( + '{"name":"web_search","arguments":' + '{"query":"call:evil{x:1} example"}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestLeadingMistralCallOwnsTheTurn: + """A leading Mistral call wins in document order over literal XML in trailing prose.""" + + def test_leading_mistral_wins_over_trailing_xml_literal(self): + text = ( + '[TOOL_CALLS]web_search[ARGS]{"query":"cats"} ' + "Note: 1" + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_function_xml_leading_keeps_normal_order(self): + text = ( + "x " + "[TOOL_CALLS]evil[ARGS]{}" + ) + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestGemmaDottedKeyAfterBareValue: + def test_dotted_key_after_bare_value_is_a_boundary(self): + text = "<|tool_call>call:web_search{query:foo,user.name:bob}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"query": "foo", "user.name": "bob"} + + +class TestJsonAnswersAreDataForMarkerlessScans: + """A whole-content JSON value is a structured answer: a quoted example of + an enabled tool's syntax inside it must not execute the tool, and the + display strip must not mutilate the answer.""" + + def test_gemma_example_inside_json_answer_not_promoted(self): + text = '{"answer":"Gemma syntax is call:web_search{query:hi}"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + def test_gemma_example_inside_json_answer_not_stripped(self): + from core.inference.tool_call_parser import strip_tool_markup + text = '{"answer":"Gemma syntax is call:web_search{query:hi}"}' + assert strip_tool_markup(text, final = True, enabled_tool_names = {"web_search"}) == text + + def test_kimi_marker_inside_json_answer_not_promoted(self): + text = ( + '{"answer":"<|tool_call_begin|>functions.web_search:0' + '<|tool_call_argument_begin|>{}<|tool_call_end|>"}' + ) + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + +class TestGemmaNestedQuotedLeaves: + def test_nested_object_and_array_values_are_unquoted(self): + text = 'call:f{loc:{city:"New York"},items:["a","b"],n:3}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"f"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"loc": {"city": "New York"}, "items": ["a", "b"], "n": 3} + + +class TestEarliestEnvelopeWinsAcrossDeepSeekKimi: + """The DeepSeek/Kimi pre-pass dispatches by earliest envelope opener: a + leading real call wins over a trailing example of the sibling format in + either direction (document order, like the other leading guards).""" + + _DS = ( + "<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>evil\n" + '```json\n{"x": 1}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + _KIMI = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.web_search:0" + '<|tool_call_argument_begin|>{"query": "cats"}<|tool_call_end|>' + "<|tool_calls_section_end|>" + ) + + def test_leading_kimi_wins_over_trailing_deepseek_example(self): + text = self._KIMI + " For reference: " + self._DS + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_leading_deepseek_wins_over_trailing_kimi_example(self): + text = self._DS + " Kimi format: " + self._KIMI + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["evil"] + + +class TestNamelessLeadingJsonAnswerIsData: + """A nameless leading JSON answer is an envelope: quoted markup stays data, and a call after it parses.""" + + def test_xml_literal_inside_json_answer_stays_data(self): + text = '{"answer": "use x"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + def test_real_call_after_json_answer_still_parses(self): + text = ( + '{"answer": "docs"} {"name": "web_search", ' + '"arguments": {"query": "cats"}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestClosedCallPrecedesMarkerPrePass: + """A closed non-DeepSeek/Kimi call that precedes the first DS/Kimi marker + owns the turn: a trailing example (or an example quoted inside a wrapped + Gemma argument) must not be promoted by the pre-pass.""" + + _KIMI_EVIL = ( + "<|tool_calls_section_begin|><|tool_call_begin|>functions.evil:0" + '<|tool_call_argument_begin|>{"x": 1}<|tool_call_end|>' + "<|tool_calls_section_end|>" + ) + + def test_kimi_example_inside_wrapped_gemma_arg_stays_data(self): + text = ( + '<|tool_call>call:web_search{query:<|"|>explain ' + + self._KIMI_EVIL + + '<|"|>}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_leading_xml_call_wins_over_trailing_kimi_example(self): + text = ( + '{"name":"web_search","arguments":{"query":"cats"}}' + " For reference: " + self._KIMI_EVIL + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_standalone_kimi_call_still_parses(self): + calls = parse_tool_calls_from_text(self._KIMI_EVIL) + assert [c["function"]["name"] for c in calls] == ["evil"] + + +class TestTruncatedWrapperlessGemmaStopsScan: + def test_call_quoted_inside_truncated_arg_not_promoted(self): + text = 'call:python{code:example("call:web_search{query:hi}") and then it cut' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"python", "web_search"}) == [] + + +class TestGemmaQuotedNestedDelimiters: + def test_comma_inside_quoted_nested_string_not_a_split(self): + text = 'call:f{loc:{city:"New, York"},n:1}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"f"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"loc": {"city": "New, York"}, "n": 1} + + +class TestGemmaStringMarkerLiteralInArgs: + def test_string_marker_literal_does_not_lose_the_call(self): + text = "call:web_search{query:'what does <|\"|> mean in Gemma'}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args["query"] == 'what does <|"|> mean in Gemma' + + +class TestGemmaMidValueQuotedPhrase: + def test_quoted_phrase_mid_value_hides_delimiters(self): + text = 'call:web_search{query:find "weather, location: Boston", limit:3}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"query": 'find "weather, location: Boston"', "limit": 3} + + def test_apostrophes_still_prose_mid_value(self): + text = "call:web_search{query:what's on at the museum, n:2}" + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"query": "what's on at the museum", "n": 2} + + +class TestGlmStrictRefusesInQuoteFallback: + """A truncated GLM value whose only close candidates sit inside a string + literal must reject in strict mode instead of executing truncated + arguments; Auto-Heal keeps the lenient partial value.""" + + _TRUNC = ( + 'python\ncode\nprint("")' + ) + + def test_strict_rejects_truncated_in_string_close(self): + assert parse_tool_calls_from_text(self._TRUNC, allow_incomplete = False) == [] + + def test_heal_keeps_partial_value(self): + calls = parse_tool_calls_from_text(self._TRUNC, allow_incomplete = True) + assert len(calls) == 1 and calls[0]["function"]["name"] == "python" + + +class TestGemmaGuardCoversPreambles: + def test_preamble_then_gemma_call_quoting_xml_wins(self): + text = ( + "Sure, searching now. call:web_search{query:" + '"explain {"name":"evil","arguments":{}}"}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestGlmStrictAcceptsApostrophes: + def test_apostrophe_value_parses_in_strict_mode(self): + text = ( + "web_search\nquery\n" + "what's the weather\n" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + args = json.loads(calls[0]["function"]["arguments"]) + assert args == {"query": "what's the weather"} + + +class TestDisabledGemmaCallLiteralsAreData: + def test_literal_inside_disabled_call_not_promoted(self): + text = 'call:foo{query:"x"}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"python", "web_search"}) == [] + + def test_real_call_after_disabled_example_still_parses(self): + text = ( + 'call:foo{query:"x"}' + " call:web_search{query:hi}" + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python", "web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + +class TestLeadingJsonArrayAnswerIsData: + def test_kimi_marker_inside_json_array_answer_not_promoted(self): + text = ( + '[{"answer": "<|tool_call_begin|>functions.web_search:0' + '<|tool_call_argument_begin|>{}<|tool_call_end|>"}]' + ) + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + +class TestLeadingBareJsonOwnsTurnOverTrailingXml: + """Document order: a leading closed bare-JSON call owns the turn even when + tool XML appears AFTER it (inside-or-after, mirroring the Mistral rule).""" + + def test_leading_call_wins_over_trailing_xml(self): + text = ( + '{"name":"lookup","parameters":{"q":"first"}} Example: ' + '{"name":"delete_all","arguments":{}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["lookup"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"q": "first"} + + def test_chained_leading_calls_win_over_trailing_xml(self): + text = ( + '{"name":"lookup","parameters":{"q":"first"}};' + '{"name":"lookup","parameters":{"q":"second"}} ' + '{"name":"delete_all","arguments":{}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["lookup", "lookup"], calls + + def test_non_call_leading_object_defers_to_trailing_real_call(self): + # Nameless answers and disabled-name objects take the decline path: + # the object is dropped and the real trailing call still parses. + for lead in ('{"answer": 42}', '{"name":"draft","parameters":{}}'): + text = lead + ' {"name":"delete_all","arguments":{}}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"delete_all"}) + assert [c["function"]["name"] for c in calls] == ["delete_all"], (lead, calls) + + def test_leading_xml_call_still_wins_over_trailing_bare_json(self): + text = ( + '{"name":"delete_all","arguments":{}} ' + 'Example: {"name":"lookup","parameters":{"q":"x"}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"lookup", "delete_all"}) + assert [c["function"]["name"] for c in calls] == ["delete_all"], calls + + +class TestProseCloseTagAfterClosedFunctionCall: + """A literal in prose after a closed call is data: the call + ends at its first close that is not parameter data, so arguments never + swallow the prose between the real close and the literal.""" + + def test_arguments_do_not_swallow_prose(self): + text = ( + "cats" + " Done. The tag closes a call." + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_literal_close_inside_open_parameter_stays_data(self): + text = 'print("")' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python"}) + assert [c["function"]["name"] for c in calls] == ["python"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"code": 'print("")'} + + def test_attribute_form_arguments_do_not_swallow_prose(self): + # The attribute form shares the first-balanced-close + # rule: prose mentioning a literal close tag never folds into arguments. + text = ( + 'cats' + " Done. The tag closes a call." + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert [c["function"]["name"] for c in calls] == ["web_search"], calls + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_attribute_form_literal_close_in_open_parameter_stays_data(self): + text = 'print("")' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"python"}) + assert json.loads(calls[0]["function"]["arguments"]) == {"code": 'print("")'} + + def test_attribute_form_two_calls_both_parse(self): + text = ( + 'cats' + 'x=1' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "python"}) + assert [c["function"]["name"] for c in calls] == ["web_search", "python"], calls + + +class TestEnabledNameJsonAnswerIsContent: + """A JSON answer whose top-level name matches an enabled tool but has no + call shape is content: the parser rejects it, so the strip and the drain + gate must keep it visible too.""" + + def test_answer_survives_strip(self): + from core.inference.tool_call_parser import strip_leading_bare_json_call + ans = '{"name":"web_search","result":"no call"}' + assert strip_leading_bare_json_call(ans, {"web_search"}) == ans + + def test_answer_does_not_route_to_draining(self): + from core.inference.safetensors_agentic import _looks_like_enabled_bare_json + assert not _looks_like_enabled_bare_json( + '{"name":"web_search","result":"no call"}', {"web_search"} + ) + + def test_real_call_still_strips_and_drains(self): + from core.inference.safetensors_agentic import _looks_like_enabled_bare_json + from core.inference.tool_call_parser import strip_leading_bare_json_call + + real = '{"name":"web_search","parameters":{"q":"x"}}' + assert strip_leading_bare_json_call(real, {"web_search"}) == "" + assert _looks_like_enabled_bare_json(real, {"web_search"}) + + def test_arguments_string_call_still_strips(self): + from core.inference.tool_call_parser import strip_leading_bare_json_call + call = '{"name":"web_search","arguments":"{\\"q\\":\\"x\\"}"} tail' + assert strip_leading_bare_json_call(call, {"web_search"}) == "tail" + + +class TestAttributeFormLeadingContainment: + """A leading attribute-form call owns the turn: markup quoted inside its + parameter is data, not a call for the shared XML parser to promote.""" + + def test_quoted_tool_call_inside_param_stays_data(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + 'find ' + '{"name":"delete","arguments":{}}' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "delete"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + assert "delete" in json.loads(calls[0]["function"]["arguments"])["query"] + + def test_real_xml_call_before_attribute_form_keeps_order(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + '{"name":"delete","arguments":{}} Example: ' + 'x' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "delete"}) + assert calls[0]["function"]["name"] == "delete" + + +class TestParameterKeepsMultipleLiteralCloses: + """A parameter that provably closes with its own tag keeps every literal + function close inside it as data (regression: the first literal close was + treated as ending the parameter, truncating the value).""" + + def test_two_literal_closes_in_one_parameter(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + '' + "a b c " + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert json.loads(calls[0]["function"]["arguments"]) == { + "query": "a b c" + } + + def test_strip_removes_the_whole_call(self): + from core.inference.tool_call_parser import strip_tool_markup + text = ( + '' + "a b c after" + ) + assert strip_tool_markup(text, final = True) == "after" + + def test_unclosed_parameter_still_heals_at_function_close(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + calls = parse_tool_calls_from_text( + "val", + enabled_tool_names = {"web_search"}, + ) + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "val"} + + +class TestMistralPreambleOwnership: + """A visible preface before the first Mistral call must not hand the turn + to a later XML literal: the Mistral call is first in document order.""" + + def test_v11_named_form_after_preface(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + 'pref [TOOL_CALLS]web_search[ARGS]{"query":"cats"} Note ' + "1" + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_array_form_after_preface(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + 'pref [TOOL_CALLS][{"name":"web_search","arguments":{"query":"cats"}}] Note ' + "1" + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert [c["function"]["name"] for c in calls] == ["web_search"] + + def test_xml_call_before_trigger_keeps_order(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + "1 then " + '[TOOL_CALLS][{"name":"web_search","arguments":{}}]' + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search", "evil"}) + assert calls[0]["function"]["name"] == "evil" + + def test_prose_mention_without_call_shape_keeps_order(self): + from core.inference.tool_call_parser import parse_tool_calls_from_text + + text = ( + "See [TOOL_CALLS] docs for details. " + "1" + ) + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"evil"}) + assert [c["function"]["name"] for c in calls] == ["evil"] + + +class TestBareJsonStripRequiresTopLevelName: + """The strip's shape gate requires the parser's TOP-LEVEL name in every + mode: a JSON answer with only a nested name is content, even name-agnostic.""" + + def test_nested_name_answer_survives_name_agnostic_strip(self): + from core.inference.tool_call_parser import strip_leading_bare_json_call + + ans = '{"parameters":{},"result":{"name":"web_search"}}' + assert strip_leading_bare_json_call(ans) == ans + assert strip_leading_bare_json_call(ans, {"web_search"}) == ans + + def test_real_call_still_strips_name_agnostic(self): + from core.inference.tool_call_parser import strip_leading_bare_json_call + assert strip_leading_bare_json_call('{"name":"web_search","parameters":{"q":"x"}}') == "" + + +class TestGemmaAwareClosedBlockPrePass: + """The closed JSON/function strip pre-pass must not delete across a complete + Gemma span (a quoted plus a later real ).""" + + def test_literal_function_in_gemma_arg_with_later_real_call(self): + from core.tool_healing import strip_tool_call_markup + text = ( + 'before <|tool_call>call:python{code:<|"|>print("")<|"|>}' + " ls" + " after" + ) + assert strip_tool_call_markup(text, final = True) == "before after" + + def test_literal_function_in_gemma_arg_with_prose_closer(self): + from core.tool_healing import strip_tool_call_markup + + text = ( + 'before <|tool_call>call:python{code:<|"|>print("")<|"|>}' + " then use to close. after" + ) + out = strip_tool_call_markup(text, final = True) + assert out.startswith("before") + assert out.endswith("after") + assert "call:python" not in out + + def test_gemma_opener_inside_json_arg_still_strips_block(self): + from core.tool_healing import strip_tool_call_markup + text = ( + '{"name":"t","arguments":{"code":"<|tool_call>call:x{"}} after' + ) + assert strip_tool_call_markup(text, final = True) == "after" + + def test_gemma_opener_inside_function_param_still_strips_block(self): + from core.tool_healing import strip_tool_call_markup + text = ( + 'x = "<|tool_call>call:t{"' + " after" + ) + assert strip_tool_call_markup(text, final = True) == "after" diff --git a/studio/backend/tests/test_tool_strip_guard.py b/studio/backend/tests/test_tool_strip_guard.py new file mode 100644 index 0000000000..dfa3101882 --- /dev/null +++ b/studio/backend/tests/test_tool_strip_guard.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""strip_tool_patterns must match the plain per-pattern loop while skipping the +quadratic no-match rescan of a closed-pair sweep whose close token is absent.""" + +import random +import sys +import time +from pathlib import Path + +_BACKEND_ROOT = Path(__file__).resolve().parents[1] +if str(_BACKEND_ROOT) not in sys.path: + sys.path.insert(0, str(_BACKEND_ROOT)) + +from core.tool_healing import ( + _TOOL_ALL_PATS, + _TOOL_CLOSED_PATS, + strip_tool_call_markup, + strip_tool_patterns, +) + + +def _naive(text, patterns): + for pat in patterns: + text = pat.sub("", text) + return text + + +_TOKENS = [ + "", + "", + "<|tool_call>", + "", + "", + "", + "", + "", + "", + "call:fn{", + "}", + "{", + '<|"|>', + "A", + " ", + "\n", + "id", + "x:1", + "", +] + + +def test_guard_matches_plain_loop_on_fuzz(): + rng = random.Random(1234) + for patterns in (_TOOL_ALL_PATS, _TOOL_CLOSED_PATS): + for _ in range(20000): + s = "".join(rng.choice(_TOKENS) for _ in range(rng.randint(0, 10))) + assert strip_tool_patterns(s, patterns) == _naive(s, patterns), (s, patterns) + + +def test_strip_markup_representative_cases_unchanged(): + assert strip_tool_call_markup("a {} b") == "a b" + assert strip_tool_call_markup("a 1 b") == "a b" + # Non-final keeps an unclosed block; final strips it to EOF. + assert strip_tool_call_markup("a {partial") == "a {partial" + assert strip_tool_call_markup("a {partial", final = True) == "a" + + +def test_no_quadratic_blowup_on_unclosed_markers(): + # Unguarded, this took minutes. + big = "" * 20000 + "" * 20000 + t0 = time.perf_counter() + out = strip_tool_call_markup(big, final = True) + assert time.perf_counter() - t0 < 2.0 + assert out == "" diff --git a/studio/backend/tests/test_tool_xml_strip.py b/studio/backend/tests/test_tool_xml_strip.py index c2dc1fe8db..f7792a2a71 100644 --- a/studio/backend/tests/test_tool_xml_strip.py +++ b/studio/backend/tests/test_tool_xml_strip.py @@ -24,17 +24,74 @@ import re as _re _src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text() _m = _re.search(r"_TOOL_XML_RE = _re\.compile\((.*?)\n\)", _src, _re.DOTALL) assert _m, "could not extract _TOOL_XML_RE source" -_ns = {"_re": _re} +# The lazy ``(.*?)\n\)`` could grab a shorter expression if an arm is ever wrapped; +# pin the DeepSeek + bare-Kimi arms so a silent truncation fails loudly here. +assert "_DS_OPEN_SRC" in _m.group(1) and "tool_call_begin" in _m.group( + 1 +), "extracted _TOOL_XML_RE is missing expected arms (extraction truncated?)" +# The regex reuses the parser's shared DeepSeek opener alternation; provide it so the extracted +# ``_re.compile`` expression resolves the same source. +from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC +from core.inference.tool_call_parser import ( + _strip_function_xml_calls, + _strip_gemma_wrapperless_calls, + _strip_glm_calls, + _strip_mistral_closed_calls, +) + +from typing import Optional as _Optional + +_ns = { + "_re": _re, + "_DS_OPEN_SRC": _DS_OPEN_SRC, + "Optional": _Optional, + "_strip_mistral_closed_calls": _strip_mistral_closed_calls, + "_strip_gemma_wrapperless_calls": _strip_gemma_wrapperless_calls, + "_strip_glm_calls": _strip_glm_calls, + "_strip_function_xml_calls": _strip_function_xml_calls, +} exec(f"_TOOL_XML_RE = _re.compile({_m.group(1)})", _ns) _TOOL_XML_RE = _ns["_TOOL_XML_RE"] -_helper = _re.search( - r"def _strip_tool_xml_for_display\(text: str, \*, auto_heal_tool_calls: bool\) -> str:\n" - r"(?: .+\n)+", +# The display helper uses the closed-only variant before the last think block; keep it in scope. +_mc = _re.search(r"_TOOL_XML_CLOSED_RE = _re\.compile\((.*?)\n\)", _src, _re.DOTALL) +assert _mc, "could not extract _TOOL_XML_CLOSED_RE source" +exec(f"_TOOL_XML_CLOSED_RE = _re.compile({_mc.group(1)})", _ns) +_TOOL_XML_CLOSED_RE = _ns["_TOOL_XML_CLOSED_RE"] + +# Signatures may span multiple lines and now carry the enabled_tool_names gate; match +# the whole (possibly multi-line) signature up to ``-> str:`` then the indented body. +_xml_helper = _re.search( + r"def _strip_tool_xml\((?:.|\n)*?\) -> str:\n(?: .+\n)+", _src, ) -assert _helper, "could not extract _strip_tool_xml_for_display source" +assert _xml_helper, "could not extract _strip_tool_xml source" +assert "_strip_mistral_closed_calls" in _xml_helper.group( + 0 +), "extracted _strip_tool_xml no longer runs the Mistral balanced strip" +exec(_xml_helper.group(0), _ns) +_strip_tool_xml = _ns["_strip_tool_xml"] + +# Extract the gate helper and display strip up to the next top-level ``logger =``. +_helper = _re.search( + r"def _display_tool_name_gate\(.*?(?=\nlogger = get_logger)", + _src, + _re.DOTALL, +) +assert _helper, "could not extract display strip helper source" +# The extracted block spans _display_tool_name_gate through _strip_tool_xml (defined before +# ``logger =``); confirm the shared _strip_tool_xml delegate is present. +assert "_strip_tool_xml(" in _helper.group(0), "display helper no longer delegates" exec(_helper.group(0), _ns) _strip_tool_xml_for_display = _ns["_strip_tool_xml_for_display"] +_display_tool_name_gate = _ns["_display_tool_name_gate"] + +_gate_src = _re.search( + r"def _gemma_strip_gate\((?:.|\n)*?\) -> set:\n(?: .+\n)+", + _src, +) +assert _gate_src, "could not extract _gemma_strip_gate source" +exec(_gate_src.group(0), _ns) +_gemma_strip_gate = _ns["_gemma_strip_gate"] # ── Well-formed pairs ───────────────────────────────────────────── @@ -46,6 +103,66 @@ def test_route_display_strip_respects_disabled_auto_heal_contract(): assert "" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) +def test_route_display_strip_preserves_rehearsal_inside_think(): + # A rehearsed bracket call inside think is reasoning: the block is preserved while a real + # call outside it still strips. + text = 'plan: search[ARGS]{"q":"x"} answer [TOOL_CALLS]web_search{"q":"y"} tail' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert 'plan: search[ARGS]{"q":"x"}' in out + assert "[TOOL_CALLS]web_search" not in out + assert "answer" in out and "tail" in out + + +def test_route_display_strip_keeps_bare_args_before_think_block(): + # A bare ``foo[ARGS]`` before a think block is prose: EOS-anchored tail arms run only on + # the last segment (earlier segments use the closed-only regex). + text = "Please pass foo[ARGS] pause to the template." + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) == text + + +def test_route_display_strip_removes_complete_call_before_think_block(): + # A complete bracket call before a think block still strips (balanced scan runs on every segment). + text = 'before search[ARGS]{"q":"x"} pause after' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "search[ARGS]" not in out + assert "pause" in out + assert "before" in out and "after" in out + + +def test_route_display_strip_removes_closed_xml_before_think_block(): + # A closed before a think block is removed in the non-last segment. + text = 'pre {"name":"x"} p tail' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "" not in out + assert "p" in out + assert "pre" in out and "tail" in out + + +def test_all_route_cleanup_sites_use_protected_display_helper(): + # Every route cleanup site must use _strip_tool_xml_for_display (think-preserving, + # balanced); raw _TOOL_XML_RE.sub corrupted think rehearsal and trailing prose. The only + # legitimate raw sub lives inside the helper itself. + raw_sub_lines = [ + (i, line) + for i, line in enumerate(_src.splitlines(), 1) + if "_TOOL_XML_RE.sub(" in line and not line.lstrip().startswith("#") + ] + assert len(raw_sub_lines) == 1, ( + "raw _TOOL_XML_RE.sub must appear only inside _strip_tool_xml_for_display; " + f"found extra call sites: {raw_sub_lines!r}" + ) + + +def test_route_display_strip_removes_mistral_tool_calls_with_nested_json(): + # _TOOL_XML_RE has no [TOOL_CALLS] arm, so the helper delegates to _strip_tool_xml for the Mistral + # balanced-brace strip (a non-greedy \{.*?\} would truncate nested JSON). + text = 'ok [TOOL_CALLS]web_search{"filters":{"date":"2024"},"query":"cats"} tail' + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "[TOOL_CALLS]" not in out and "web_search" not in out, out + assert out == "ok tail" + + def test_strips_well_formed_tool_call(): text = ( "Let me search.\n" @@ -73,6 +190,26 @@ def test_strips_function_only_well_formed(): assert "Done." in cleaned +def test_strips_function_attribute_form(): + # Attribute form ```` (MiniCPM-5 / MiniMax-M2) must strip from the route too + # (it previously leaked into the UI); a dotted/hyphenated name also strips. + text = ( + 'Sure.\n\n' + "\nSydney\n\n\nDone." + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert "" not in cleaned + assert "Sure." in cleaned and "Done." in cleaned + + dotted = 'A x B' + assert _TOOL_XML_RE.sub("", dotted) == "A B" + + # Auto-Heal-disabled display contract still preserves literal markup. + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text + assert "` so doc/example prose survives. text = ( @@ -281,3 +444,473 @@ def test_no_catastrophic_backtracking_on_orphan_opening_spam(): elapsed = time.perf_counter() - t0 assert elapsed < 0.1, f"regex took {elapsed*1000:.0f}ms on 1000x orphan opens" assert "" not in cleaned + + +# ── Two-level-nested bracket JSON (balanced-scan strip) ────────── + + +def test_route_strip_two_level_nested_bracket_keeps_trailing_prose(): + # Two-level-nested args must be removed whole so the trailing prose survives. + text = 'before [TOOL_CALLS]search{"f":{"g":{"h":1}}} after' + cleaned = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert cleaned == "before after" + assert "[TOOL_CALLS]" not in cleaned + + +def test_route_strip_two_level_nested_rehearsal_keeps_trailing_prose(): + text = 'note python[ARGS]{"a":{"b":{"c":1}}} done' + cleaned = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert cleaned == "note done" + assert "[ARGS]" not in cleaned + + +def test_route_strip_removes_call_with_literal_think_in_argument(): + # A literal inside a call argument strips with the call, not as reasoning. + text = ( + '{"name":"write","arguments":' + '{"text":"compare and tags"}}' + ) + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "" not in out and '"name"' not in out + + +def test_route_strip_removes_truncated_mistral_array(): + # A canonical array truncated by EOS is stripped by the route fallback like other orphans. + text = 'before [TOOL_CALLS] [{"name":"a","arguments":{"x":1}}' # missing ] + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "[TOOL_CALLS]" not in out and "{" not in out + assert "before" in out + + +def test_route_strip_keeps_prose_mentioning_args_marker(): + # ``foo[ARGS] in a sentence`` is prose; the rehearsal arm must not truncate the line. + text = "Please pass foo[ARGS] to the template and continue reading." + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert out == text + + +def test_route_strip_handles_mistral_v11_call_id_args_shape(): + # v11 [CALL_ID]/[ARGS] shape (Mistral Small 3.2) must strip whole. + text = 'before [TOOL_CALLS]web_search[CALL_ID]abc123[ARGS]{"q":"x"} after' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "[TOOL_CALLS]" not in out and "[CALL_ID]" not in out and "[ARGS]" not in out + assert "before" in out and "after" in out + + +# ── Mistral [/TOOL_CALLS] closer + literal inside a call ─────────────── + +from core.tool_healing import strip_tool_call_markup as _strip_tool_call_markup + + +def test_core_strip_removes_orphan_tool_calls_closer_array_form(): + # The bare v11 [/TOOL_CALLS] closer left by the balanced scan must not leak as content. + text = '[TOOL_CALLS] [{"name":"x","arguments":{}}][/TOOL_CALLS]' + assert _strip_tool_call_markup(text, final = True) == "" + + +def test_core_strip_removes_orphan_tool_calls_closer_named_form_keeps_tail(): + text = '[TOOL_CALLS]web_search{"q":"x"}[/TOOL_CALLS] tail' + assert _strip_tool_call_markup(text, final = True) == "tail" + + +def test_core_strip_removes_call_with_literal_think_in_argument(): + # An unclosed literal inside call arguments strips with the call (argument data). + text = 'before {"name":"write","arguments":{"text":"literal marker"}} after' + assert _strip_tool_call_markup(text, final = True) == "before after" + + +def test_route_display_strip_removes_orphan_tool_calls_closer_array_form(): + text = '[TOOL_CALLS] [{"name":"x","arguments":{}}][/TOOL_CALLS]' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert out.strip() == "" + + +def test_route_display_strip_removes_orphan_tool_calls_closer_named_form_keeps_tail(): + text = '[TOOL_CALLS]web_search{"q":"x"}[/TOOL_CALLS] tail' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "[/TOOL_CALLS]" not in out + assert out.strip() == "tail" + + +def test_incomplete_xml_call_with_literal_think_in_arg_is_stripped(): + # An incomplete holding a literal strips to EOS, not as a reasoning + # block (the unclosed tail _tool_call_markup_spans previously missed). + from core.tool_healing import parse_tool_calls_from_text as _parse + from core.tool_healing import strip_tool_call_markup as _strip + + text = 'before {"name":"write","arguments":{"text":"literal marker"}} after' + assert [c["function"]["name"] for c in _parse(text)] == ["write"] + assert _strip(text, final = True) == "before" + + # A real reasoning block with no tool call is still preserved verbatim. + assert ( + _strip("answer real done", final = True) == "answer real done" + ) + + # A complete call followed by a real reasoning block: call stripped, block kept. + mixed = '{"name":"a","arguments":{}} mid r end' + assert _strip(mixed, final = True) == "mid r end" + + +# ── enabled-tool gate for the ambiguous bare-rehearsal strip (#5704) ── + + +def test_display_tool_name_gate_returns_active_names_or_none(): + # Empty / no tools -> None (unrestricted; keep the legacy strip-all behavior). + assert _display_tool_name_gate([]) is None + assert _display_tool_name_gate(None) is None + # OpenAI-shaped tool dicts -> set of function names, malformed entries dropped. + tools = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "run_python"}}, + {"type": "function"}, # no name + {"nope": 1}, # no function + ] + assert _display_tool_name_gate(tools) == {"web_search", "run_python"} + + +def test_route_display_strip_keeps_inactive_rehearsal_when_gated(): + # P1 #5704: an inactive ``foo[ARGS]{...}`` is prose; the gated strip leaves the sentence intact. + gate = {"web_search"} + text = 'foo[ARGS]{"x":1} is just syntax.' + assert ( + _strip_tool_xml_for_display(text, auto_heal_tool_calls = True, enabled_tool_names = gate) + == text + ) + # A bare marker with no JSON body is likewise prose when inactive. + assert ( + _strip_tool_xml_for_display( + "use foo[ARGS] here", auto_heal_tool_calls = True, enabled_tool_names = gate + ) + == "use foo[ARGS] here" + ) + + +def test_route_display_strip_removes_active_rehearsal_when_gated(): + # Mirror case: an active tool name is a real rehearsal and still strips. + gate = {"web_search"} + out = _strip_tool_xml_for_display( + 'web_search[ARGS]{"query":"x"} done', auto_heal_tool_calls = True, enabled_tool_names = gate + ) + assert "web_search[ARGS]" not in out + assert out.strip() == "done" + + +def test_route_display_strip_ungated_strips_all_rehearsal_unchanged(): + # Backwards-compat: with no gate (None) the bare rehearsal strips as before. + text = 'foo[ARGS]{"x":1} is just syntax.' + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True).strip() == "is just syntax." + assert ( + _strip_tool_xml_for_display( + text, auto_heal_tool_calls = True, enabled_tool_names = None + ).strip() + == "is just syntax." + ) + + +def test_route_display_strip_control_token_stripped_regardless_of_gate(): + # [TOOL_CALLS] is a control token: stripped even when its NAME is not in the gate. + gate = {"web_search"} + out = _strip_tool_xml_for_display( + '[TOOL_CALLS]foo[ARGS]{"x":1} keep', auto_heal_tool_calls = True, enabled_tool_names = gate + ) + assert "[TOOL_CALLS]" not in out and "foo[ARGS]" not in out + assert out.strip() == "keep" + + +def test_core_strip_gates_bare_rehearsal_on_enabled_tools(): + # P1 (#5704): the shared strip gate mirrors the parse gate -- inactive names are prose + # and preserved, active names strip, ``None`` keeps legacy strip-all. + from core.tool_healing import strip_tool_call_markup as _strip + + text = 'foo[ARGS]{"x":1} is just syntax.' + assert _strip(text, final = True, enabled_tool_names = {"web_search"}) == text + assert ( + _strip('web_search[ARGS]{"q":1} done', final = True, enabled_tool_names = {"web_search"}) + == "done" + ) + assert _strip(text, final = True).strip() == "is just syntax." + assert _strip(text, final = True, enabled_tool_names = None).strip() == "is just syntax." + + +def test_route_display_strip_gate_preserves_inactive_history_rehearsal(): + # The GGUF history sanitiser passes the gate, so a documented inactive shape survives in + # the replayed prompt context. + gate = _display_tool_name_gate([{"function": {"name": "web_search"}}]) + text = 'To call it write foo[ARGS]{"x":1} in your reply.' + assert 'foo[ARGS]{"x":1}' in _strip_tool_xml_for_display( + text, auto_heal_tool_calls = True, enabled_tool_names = gate + ) + # An ACTIVE name is still stripped as a real rehearsed call. + assert "web_search[ARGS]" not in _strip_tool_xml_for_display( + 'Result web_search[ARGS]{"q":"x"} done', auto_heal_tool_calls = True, enabled_tool_names = gate + ) + # No gate (legacy) strips every NAME[ARGS]{...}. + assert "foo[ARGS]" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + + +def test_gguf_history_sanitizer_forwards_enabled_tool_names_gate(): + # Wiring guard: the GGUF history strip must forward the display gate like the live strip. + block = _re.search( + r"Strip stale tool-call XML from conversation history.*?\.strip\(\)", + _src, + _re.DOTALL, + ) + assert block, "could not locate GGUF history sanitizer block" + assert "enabled_tool_names" in block.group( + 0 + ), "GGUF history sanitizer must pass enabled_tool_names to _strip_tool_xml_for_display" + + +def test_route_history_and_passthrough_forward_the_display_gate(): + # The safetensors/Anthropic history sanitisers and the Anthropic non-stream passthrough + # must forward the gate so inactive examples survive in replayed prompt / final text. + blocks = { + "safetensors history": r"Strip stale tool-call XML from prior assistant turns.*?\.strip\(\)", + "anthropic history": r"Strip stale tool-call XML via the protected display helper.*?\.strip\(\)", + "anthropic passthrough": r"gated on the declared tools so an\n.*?\.strip\(\)", + } + for label, pat in blocks.items(): + m = _re.search(pat, _src, _re.DOTALL) + assert m, f"could not locate {label} strip block" + assert "enabled_tool_names" in m.group( + 0 + ), f"{label} must forward enabled_tool_names to _strip_tool_xml_for_display" + + +# ── DeepSeek opener variants + bare Kimi (parse/strip symmetry) ── + + +def test_strips_deepseek_space_opener_variant(): + # The space-separated opener is parsed by the parser, so the display strip + # must remove it too (the shared opener alternation is reused here). + text = ( + "pre <｜tool calls begin｜><｜tool▁call▁begin｜>get_x<｜tool▁sep｜>" + '{"a":1}<｜tool▁call▁end｜><｜tool▁calls▁end｜> post' + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert "tool" not in cleaned.replace("post", "").replace("pre", "") + assert cleaned == "pre post" + + +def test_strips_deepseek_escaped_underscore_opener_variant(): + text = ( + "pre <｜tool\\_calls\\_begin｜><｜tool▁call▁begin｜>get_y<｜tool▁sep｜>" + '{"a":1}<｜tool▁call▁end｜><｜tool▁calls▁end｜> post' + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert cleaned == "pre post" + + +def test_strips_bare_kimi_call_without_section_wrapper(): + # Kimi can emit a bare <|tool_call_begin|>...<|tool_call_end|> with no + # section wrapper; the parser accepts it, so the strip must cover it. + text = ( + "pre <|tool_call_begin|>functions.get_w:0<|tool_call_argument_begin|>" + '{"a":1}<|tool_call_end|> post' + ) + cleaned = _TOOL_XML_RE.sub("", text) + assert "tool_call_begin" not in cleaned + assert cleaned == "pre post" + + +@pytest.mark.parametrize( + "text", + [ + # Prose that merely names a Kimi/DeepSeek marker (no real call follows) must + # survive: the call-shaped lookahead fires only on a real call or a bare EOF + # fragment, so an answer discussing the protocol is never truncated. + "See <|tool_call_begin|> in the docs. More prose after it.", + "The <|tool_calls_section_begin|> marker opens a batch. Read on.", + "DeepSeek uses <｜tool▁calls▁begin｜> to start a call block, then continues.", + ], +) +def test_deepseek_kimi_false_alarm_prose_is_kept(text): + # Regression for the route arm truncating a prose answer that references a marker + # without a following call (parser _TOOL_ALL_PATS already had this lookahead). + assert _TOOL_XML_RE.sub("", text) == text + + +def test_deepseek_kimi_real_calls_still_strip_after_false_alarm_fix(): + # The lookahead must not weaken real-call stripping: closed, truncated, and bare + # EOF-fragment forms all still get removed. + closed = ( + "answer <|tool_call_begin|>functions.get_w:0<|tool_call_argument_begin|>" + '{"a":1}<|tool_call_end|> tail' + ) + assert _TOOL_XML_RE.sub("", closed) == "answer tail" + eof_fragment = "prefix <|tool_call_begin|>" + assert _TOOL_XML_RE.sub("", eof_fragment) == "prefix " + deepseek = ( + "reply <｜tool▁calls▁begin｜><｜tool▁call▁begin｜>get_x<｜tool▁sep｜>" + '{"a":1}<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ) + assert _TOOL_XML_RE.sub("", deepseek) == "reply " + + +# ── Llama-3 <|python_tag|> arm bounds on REAL sentinels only ────── + + +# Llama-3 <|python_tag|> arm bounds on REAL sentinels only +def test_python_tag_strip_consumes_literal_sentinel_in_arg(): + # A <|python_tag|> tool call whose JSON argument carries a literal <|...|> + # token (here <|cite|>) must be stripped whole. The old `<(?!\|)` arm stopped + # at any `<|`, leaking the call tail (e.g. `<|cite|> here"}}`) into display. + text = '<|python_tag|>{"name": "send", "parameters": {"text": "use <|cite|> here"}}' + cleaned = _TOOL_XML_RE.sub("", text) + assert cleaned == "", f"python_tag call leaked at literal sentinel: {cleaned!r}" + + +@pytest.mark.parametrize( + "sentinel", + [ + "<|eot_id|>", + "<|eom_id|>", + "<|start_header_id|>", + "<|end_header_id|>", + ], +) +def test_python_tag_strip_stops_at_real_sentinel(sentinel): + # A genuine Llama control sentinel still bounds the strip so following + # assistant text is preserved (the arm must not swallow past it). + text = f'<|python_tag|>{{"name": "x", "parameters": {{}}}}{sentinel}visible answer' + cleaned = _TOOL_XML_RE.sub("", text) + assert ( + cleaned == f"{sentinel}visible answer" + ), f"strip did not stop at real sentinel {sentinel!r}: {cleaned!r}" + + +def test_python_tag_strip_restarts_on_second_python_tag(): + # A second <|python_tag|> opens a new tool-call region, so the whole pair is + # stripped (the arm bounds the first, then the next match consumes the rest). + text = '<|python_tag|>{"name": "a"}<|python_tag|>{"name": "b"}' + cleaned = _TOOL_XML_RE.sub("", text) + assert cleaned == "", f"second python_tag region leaked: {cleaned!r}" + + +def test_glm_call_with_literal_close_tag_in_arg_value_is_stripped_whole(): + # GLM 4.x emits NAMEkv .... + text = ( + "web_search\nquery\n" + "find here\n done" + ) + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "" not in out + assert "" not in out + assert out.strip() == "done" + + +def test_glm_normal_and_qwen_calls_still_stripped_by_route(): + # Regression: a normal GLM call (no literal close tag) and a Qwen + # {json} are still stripped; trailing prose is kept. + glm = "get_time\ntz\nUTC\n ok" + assert _strip_tool_xml_for_display(glm, auto_heal_tool_calls = True).strip() == "ok" + qwen = '{"name":"web_search","arguments":{"q":"x"}} after' + assert _strip_tool_xml_for_display(qwen, auto_heal_tool_calls = True).strip() == "after" + + +def test_route_strip_removes_param_alias_close_tag(): + # The parser accepts the ... attribute-form alias of + # ; the route tail cleanup must strip an orphan close too. + assert _strip_tool_xml_for_display("answer ", auto_heal_tool_calls = True) == "answer " + assert ( + _strip_tool_xml_for_display("answer ", auto_heal_tool_calls = True) == "answer " + ) + + +def test_route_strip_uses_guarded_function_scan_for_literal_nested_markup(): + # A literal in a value must not truncate the strip: the route runs the + # parser's guarded function-XML scan before the regex, matching the core strip. + text = " tail" + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True).strip() == "tail" + + +def test_route_strip_gates_wrapperless_gemma_by_enabled_tools(): + # The route strip must gate the markerless Gemma call:NAME{...} form on the enabled tool names, + # like the parser/loop, so a disabled/example name in prose is preserved in ... + prose = "To document syntax you write call:foo{query:example}. That shows the format." + assert "call:foo{query:example}" in _strip_tool_xml(prose, {"web_search"}) + # An enabled name is still a real call and stripped. + assert "call:web_search" not in _strip_tool_xml( + "Answer. call:web_search{query:x}", {"web_search"} + ) + # No gate (legacy) strips every closed call. + assert "call:foo" not in _strip_tool_xml(prose) + + +def test_gemma_strip_gate_empty_tools_preserves_prose(): + # With NO tools enabled the gate must return an EMPTY set (strip nothing), not None: None falls + # back to strip-all and deletes an answer that documents the call:NAME{...} syntax. + assert _gemma_strip_gate([]) == set() + assert _gemma_strip_gate(None) == set() + assert _gemma_strip_gate([{"function": {"name": "web_search"}}]) == {"web_search"} + prose = "To document syntax you write call:foo{query:example}. That shows the format." + assert "call:foo{query:example}" in _strip_tool_xml(prose, _gemma_strip_gate([])) + assert "call:foo{query:example}" in _strip_tool_xml(prose, _gemma_strip_gate(None)) + # An enabled tool's real call is still stripped. + assert "call:web_search" not in _strip_tool_xml( + "Answer. call:web_search{query:x}", + _gemma_strip_gate([{"function": {"name": "web_search"}}]), + ) + + +def test_strip_keeps_prose_after_closed_function_call_with_literal_close(): + # The call ends at its first non-data close: prose after it survives the + # strip even when it mentions a literal . + from core.inference.tool_call_parser import strip_tool_markup + text = ( + "cats" + " Done. The tag closes a call." + ) + assert strip_tool_markup(text, final = True) == "Done. The tag closes a call." + + +def test_final_strip_keeps_prose_mentioning_bare_markers(): + # A false-alarm marker in a normal answer must not lose everything after + # it; only text that looks like that family's call start drops. + from core.inference.tool_call_parser import strip_tool_markup + for text in ( + "See [TOOL_CALLS] docs for details. More prose after.", + "<|python_tag|> is the Llama marker. Explanation continues.", + "The <|tool_call> opener wraps Gemma calls.", + ): + assert strip_tool_markup(text, final = True) == text + # A bare marker at end-of-text is a fragment and still drops. + assert strip_tool_markup("Answer text [TOOL_CALLS]", final = True) == "Answer text" + + +def test_final_strip_still_drops_truncated_marker_calls(): + from core.inference.tool_call_parser import strip_tool_markup + for text in ( + '[TOOL_CALLS][{"name":"web_search","argu', + '[TOOL_CALLS]web_search[ARGS]{"q":"x', + '<|python_tag|>{"name":"web_search","par', + '<|python_tag|>foo.call(items=["a', + "<|tool_call>call:web_search{query:tru", + ): + assert strip_tool_markup(text, final = True) == "" + + +def test_chained_bare_json_strip_consumes_all_calls(): + # The loops keep this text as next-turn history: a leftover executed call + # would be replayed alongside the structured tool_calls. + from core.inference.tool_call_parser import strip_leading_bare_json_call + + enabled = {"web_search", "python"} + chained = ( + '{"name":"web_search","parameters":{"q":"first"}};' + '{"name":"python","parameters":{"code":"x"}}' + ) + assert strip_leading_bare_json_call(chained, enabled_tool_names = enabled) == "" + assert ( + strip_leading_bare_json_call(chained + " trailing prose", enabled_tool_names = enabled) + == "trailing prose" + ) + # The chain stops at a non-call answer object, which stays visible. + call_then_answer = ( + '{"name":"web_search","parameters":{"q":"x"}};{"name":"web_search","result":"data"}' + ) + assert ( + strip_leading_bare_json_call(call_then_answer, enabled_tool_names = enabled) + == '{"name":"web_search","result":"data"}' + ) diff --git a/studio/backend/tests/test_torchao_select.py b/studio/backend/tests/test_torchao_select.py index 2d3dc5fbff..e4775a10a6 100644 --- a/studio/backend/tests/test_torchao_select.py +++ b/studio/backend/tests/test_torchao_select.py @@ -32,16 +32,23 @@ def _load_module(monkeypatch): @pytest.mark.parametrize( "torch_version, expected", [ - # torch 2.10 (the reported bug: cu130 resolves 2.10.0) -> 0.16.0, - # independent of the local +cuXXX/+rocm/+cpu suffix or patch level. - ("2.10.0+cu130", "torchao==0.16.0"), + # torch 2.10 on CUDA <= 12 -> 0.16.0 (its cpp is built for torch 2.10.0 and + # loads against the CUDA-12 PyPI wheel). Independent of patch level. + ("2.10.0+cu128", "torchao==0.16.0"), + ("2.10.0+cu126", "torchao==0.16.0"), ("2.10.0+rocm6.4", "torchao==0.16.0"), ("2.10.0+cpu", "torchao==0.16.0"), ("2.10.1", "torchao==0.16.0"), ("2.10.0", "torchao==0.16.0"), - # Pre-release / dev / rc builds: the minor is cleaned of non-digits. + # torch 2.10 on CUDA >= 13 (Blackwell / cu130): 0.16.0's CUDA-12 cpp can't + # load against a CUDA-13 torch (libcudart.so.12 error), so use 0.17.0. + ("2.10.0+cu130", "torchao==0.17.0"), + ("2.10.0+cu140", "torchao==0.17.0"), + # Pre-release / dev / rc builds: the minor is cleaned of non-digits; the + # CUDA tag still decides 0.16.0 vs 0.17.0. ("2.10.0rc1", "torchao==0.16.0"), - ("2.10.0.dev20250804+cu130", "torchao==0.16.0"), + ("2.10.0.dev20250804+cu130", "torchao==0.17.0"), + ("2.10.0.dev20250804+cu128", "torchao==0.16.0"), ("2.10rc1", "torchao==0.16.0"), # torch 2.11 (reachable via ROCm rocm7.2) and forward -> 0.17.0. ("2.11.0+cu130", "torchao==0.17.0"), diff --git a/studio/backend/tests/test_training_preflight.py b/studio/backend/tests/test_training_preflight.py index 54048a65dd..47c6669f8f 100644 --- a/studio/backend/tests/test_training_preflight.py +++ b/studio/backend/tests/test_training_preflight.py @@ -6,9 +6,15 @@ empty-chat-template crash) before train(). The real methods are bound onto a lig fake self so the production logic runs against controlled batches.""" import importlib +import json +import os +import queue +import subprocess import sys +import threading import types import unittest +from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock @@ -184,5 +190,231 @@ class TestChatTemplateRendersEmpty(unittest.TestCase): self.assertFalse(s._chat_template_renders_empty()) +def _clear_trainer_module(package: str): + sys.modules.pop(f"{package}.trainer", None) + pkg = sys.modules.get(package) + if pkg is not None and hasattr(pkg, "trainer"): + delattr(pkg, "trainer") + + +def _set_training_platform(monkeypatch, package: str, backend: str): + training_mod = importlib.import_module(f"{package}.training") + from utils.hardware import hardware as hw + + monkeypatch.setattr(hw, "DEVICE", None) + monkeypatch.setattr( + training_mod.platform, + "system", + lambda: "Darwin" if backend == "mlx" else "Linux", + ) + monkeypatch.setattr( + training_mod.platform, + "machine", + lambda: "arm64" if backend == "mlx" else "x86_64", + ) + + +def _load_trainer_module( + monkeypatch, + backend: str, + package: str = "core.training", +): + _set_training_platform(monkeypatch, package, backend) + _clear_trainer_module(package) + if package in sys.modules: + importlib.reload(sys.modules[package]) + trainer_mod = importlib.import_module(f"{package}.trainer") + training_mod = importlib.import_module(f"{package}.training") + monkeypatch.setattr( + training_mod._MLXTrainerAdapter, + "_activate_transformers_for_model", + lambda self, model_name, hf_token: None, + ) + return trainer_mod + + +class _ExitedProc: + def join(self, timeout = None): + return None + + def is_alive(self): + return False + + +class _TerminableProc: + def __init__(self): + self.terminated = False + self._done = threading.Event() + + def join(self, timeout = None): + self._done.wait(timeout = timeout or 5) + + def is_alive(self): + return not self.terminated + + def terminate(self): + self.terminated = True + self._done.set() + + +def test_unsloth_trainer_dispatches_for_mlx_and_torch(monkeypatch): + trainer_mod = _load_trainer_module(monkeypatch, "mlx") + + mlx_trainer = trainer_mod.UnslothTrainer() + + assert type(mlx_trainer).__module__ == "core.training.training" + assert mlx_trainer.get_training_progress().status_message == "Ready to train" + + trainer_mod = _load_trainer_module(monkeypatch, "torch") + + assert trainer_mod.UnslothTrainer().__class__ is trainer_mod.UnslothTrainer + + +def test_cli_mlx_trainer_activates_before_importing_trainer(): + repo_root = Path(__file__).resolve().parents[3] + script = """ +import json +import sys +import unsloth_cli.commands.train as train_cmd +from studio.backend.core.training import training as training_mod +from utils.hardware import hardware as hw + +training_mod.platform.system = lambda: "Darwin" +training_mod.platform.machine = lambda: "arm64" +hw.DEVICE = None +events = [] + +def fake_activate(model_name, hf_token): + events.append({ + "model_name": model_name, + "trainer_loaded": "studio.backend.core.training.trainer" in sys.modules, + }) + +train_cmd._activate_mlx_transformers = fake_activate +trainer = train_cmd._create_cli_trainer("mlx-community/Qwen3-0.6B-4bit", None) +print(json.dumps({ + "trainer_module": type(trainer).__module__, + "events": events, +})) +""" + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join( + [str(repo_root), str(repo_root / "studio" / "backend"), env.get("PYTHONPATH", "")] + ) + result = subprocess.run( + [sys.executable, "-c", script], + cwd = repo_root, + env = env, + text = True, + stdout = subprocess.PIPE, + stderr = subprocess.PIPE, + check = True, + ) + payload = json.loads(result.stdout) + + assert payload["trainer_module"] == "studio.backend.core.training.training" + assert payload["events"] == [ + {"model_name": "mlx-community/Qwen3-0.6B-4bit", "trainer_loaded": False} + ] + + +def test_mlx_adapter_builds_config_and_reports_completion(tmp_path, monkeypatch): + trainer_mod = _load_trainer_module(monkeypatch, "mlx") + captured = {} + + def fake_run_worker(config, event_queue, stop_queue): + captured["config"] = config + event_queue.put({"type": "progress", "step": 1, "total_steps": 1, "loss": 0.25}) + event_queue.put( + {"type": "complete", "status_message": "done", "output_dir": config["output_dir"]} + ) + + trainer = trainer_mod.UnslothTrainer() + monkeypatch.setattr(trainer, "_run_mlx_worker", fake_run_worker) + + assert trainer.load_model("mlx-community/Qwen3-0.6B-4bit", max_seq_length = 1024) + assert trainer.prepare_model_for_training(use_lora = False) + dataset, eval_dataset = trainer.load_and_format_dataset("org/dataset") + output_dir = tmp_path / "mlx-out" + + assert trainer.start_training( + dataset = dataset, + eval_dataset = eval_dataset, + output_dir = output_dir, + project_name = "Sales Assistant", + max_steps = 1, + learning_rate = 3e-4, + ) + trainer.training_thread.join(timeout = 5) + + progress = trainer.get_training_progress() + config = captured["config"] + assert progress.is_completed + assert progress.output_dir == str(output_dir.resolve()) + progress.status_message = "mutated" + assert trainer.get_training_progress().status_message == "done" + assert config["model_name"] == "mlx-community/Qwen3-0.6B-4bit" + assert config["project_name"] == "Sales Assistant" + assert config["hf_dataset"] == "org/dataset" + assert config["training_type"] == "Full Finetuning" + assert config["load_in_4bit"] is False + assert config["max_seq_length"] == 1024 + assert config["learning_rate"] == 3e-4 + assert config["output_dir"] == str(output_dir.resolve()) + assert config["allow_external_output_dir"] is True + + +def test_mlx_worker_helpers_cover_cli_paths(tmp_path, monkeypatch): + _load_trainer_module(monkeypatch, "mlx") + from core.training.worker import ( + _resolve_mlx_local_dataset_files, + _resolve_mlx_output_dir, + ) + + dataset = tmp_path / "train.jsonl" + dataset.write_text('{"text":"hello"}\n', encoding = "utf-8") + monkeypatch.chdir(tmp_path) + + assert _resolve_mlx_local_dataset_files(["train.jsonl"]) == [str(dataset)] + assert _resolve_mlx_output_dir( + {"output_dir": "cli-out", "allow_external_output_dir": True}, + "mlx-community/Qwen3-0.6B-4bit", + ) == str((tmp_path / "cli-out").resolve()) + + +def test_run_mlx_training_process_applies_side_effects_before_hardware_detection(monkeypatch): + _load_trainer_module(monkeypatch, "mlx") + from core.training import worker + from utils.hardware import hardware as hw + + order = [] + + def fake_activate(model_name, hf_token): + order.append(("activate", model_name, hf_token)) + + def fake_detect_hardware(): + order.append("detect") + hw.DEVICE = hw.DeviceType.CPU + return hw.DEVICE + + monkeypatch.delenv("HF_HUB_DISABLE_XET", raising = False) + monkeypatch.delenv("HF_HUB_ENABLE_HF_TRANSFER", raising = False) + monkeypatch.setattr(worker, "_activate_transformers_version_or_warn", fake_activate) + monkeypatch.setattr(hw, "detect_hardware", fake_detect_hardware) + + event_queue = queue.Queue() + worker.run_mlx_training_process( + event_queue = event_queue, + stop_queue = queue.Queue(), + config = {"model_name": "mlx-community/Gemma-4-12B", "disable_xet": True}, + ) + + event = event_queue.get_nowait() + assert order == [("activate", "mlx-community/Gemma-4-12B", None), "detect"] + assert os.environ["HF_HUB_DISABLE_XET"] == "1" + assert os.environ["HF_HUB_ENABLE_HF_TRANSFER"] == "0" + assert "MLX training requires Apple Silicon" in event["error"] + + if __name__ == "__main__": unittest.main() diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py index 3c5d6cd094..7e7fc1af48 100644 --- a/studio/backend/tests/test_training_worker_flash_attn.py +++ b/studio/backend/tests/test_training_worker_flash_attn.py @@ -59,7 +59,6 @@ def test_runtime_flash_attn_prefers_prebuilt_wheel(monkeypatch): statuses: list[str] = [] monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) - monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: False) monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import()) monkeypatch.setattr( worker, @@ -88,7 +87,6 @@ def test_runtime_flash_attn_falls_back_to_pypi(monkeypatch): statuses: list[str] = [] monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) - monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: False) monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import()) monkeypatch.setattr( worker, @@ -141,27 +139,6 @@ def test_runtime_flash_attn_skip_env_avoids_all_install_work(monkeypatch): worker._sp.run.assert_not_called() -def test_runtime_flash_attn_skips_on_blackwell(monkeypatch): - statuses: list[str] = [] - install_mock = mock.Mock() - - monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) - monkeypatch.setattr(worker, "_should_try_runtime_flash_attn_install", lambda max_seq: True) - monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: True) - monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock) - monkeypatch.setattr( - worker, - "_send_status", - lambda queue, message: statuses.append(message), - ) - - worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 65536) - - install_mock.assert_not_called() - assert len(statuses) == 1 - assert "Blackwell" in statuses[0] - - def test_causal_conv1d_fast_path_preserves_wheel_first_install_args(monkeypatch): install_mock = mock.Mock(return_value = True) monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock) diff --git a/studio/backend/tests/test_training_worker_import_discipline.py b/studio/backend/tests/test_training_worker_import_discipline.py new file mode 100644 index 0000000000..a047c91704 --- /dev/null +++ b/studio/backend/tests/test_training_worker_import_discipline.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Invariant: the training worker must not import ``transformers`` before it activates the +transformers sidecar. + +``core/training/worker.py:run_training_process`` runs a preflight (Xet decision, logging, hardware +detection) and only THEN calls ``_activate_transformers_version`` -> ``activate_transformers_for_subprocess``, +which prepends the correct ``.venv_t5_*`` (5.x) sidecar to ``sys.path``. Because activation only edits +``sys.path``, it is a no-op for any module already cached in ``sys.modules``. So if the preflight imports +``transformers`` (directly or transitively via ``unsloth_zoo``), the default 4.57.x gets pinned before +the sidecar is on the path -- and 5.x models (Qwen3.5, GLM-4.7, gemma-4) then fail to load their +tokenizer/config ("Tokenizer class TokenizersBackend does not exist"). + +This regression shipped once when ``utils/hf_xet_fallback.py`` eagerly imported ``unsloth_zoo`` (which +imports ``transformers``) at module load; the worker imports that shim during preflight to decide the +Xet env flip (see issue #6951). This test locks the invariant in a fresh interpreter. It is CPU-only, +needs no network/GPU/weights/sidecars, so it runs in the standard ``studio-backend-ci`` matrix. +""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +_BACKEND_DIR = Path(__file__).resolve().parent.parent # studio/backend + +# Mirrors run_training_process's imports that run BEFORE _activate_transformers_version (worker.py); +# keep in sync. torch-dependent imports are optional (a no-torch CI shard skips them) but must still +# not drag in transformers. +_PREFLIGHT_SNIPPET = r""" +import sys + +# worker.py: from utils.hf_xet_fallback import child_should_disable_xet (+ call it) +from utils.hf_xet_fallback import child_should_disable_xet +child_should_disable_xet({}) + +# worker.py: from loggers.config import LogConfig +from loggers.config import LogConfig # noqa: F401 + +# worker.py: from utils.hardware import hardware (imports torch, not transformers) +try: + from utils.hardware import hardware as _hw # noqa: F401 +except Exception: + pass # torch may be absent in a no-torch shard; the invariant below still applies + +# worker.py: from .training import is_apple_silicon_training_platform, should_use_mlx_training_backend +# (the MLX-dispatch preflight; must also stay clear of transformers). Guarded because it may pull +# unsloth/trl, absent in a minimal shard -- but a partial import that leaked transformers would still +# be caught by the assertion below. +try: + from core.training.training import ( # noqa: F401 + is_apple_silicon_training_platform as _is_apple, + should_use_mlx_training_backend as _use_mlx, + ) +except Exception: + pass + +leaked_tf = sorted(m for m in sys.modules if m == "transformers" or m.startswith("transformers.")) +leaked_zoo = sorted(m for m in sys.modules if m == "unsloth_zoo" or m.startswith("unsloth_zoo.")) +assert not leaked_tf, f"transformers imported during worker preflight (before sidecar activation): {leaked_tf}" +assert not leaked_zoo, f"unsloth_zoo imported during worker preflight (before sidecar activation): {leaked_zoo}" +print("PREFLIGHT_CLEAN") +""" + + +def test_worker_preflight_does_not_import_transformers(): + """A fresh interpreter running the worker's pre-activation imports must leave ``transformers`` + (and ``unsloth_zoo``) unimported, so the 5.x sidecar prepend is not defeated by a stale module.""" + result = subprocess.run( + [sys.executable, "-c", _PREFLIGHT_SNIPPET], + cwd = str(_BACKEND_DIR), + capture_output = True, + text = True, + ) + assert result.returncode == 0, ( + "Worker preflight imported transformers before sidecar activation.\n" + f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + ) + assert "PREFLIGHT_CLEAN" in result.stdout, result.stdout diff --git a/studio/backend/tests/test_worker_activates_correct_transformers.py b/studio/backend/tests/test_worker_activates_correct_transformers.py new file mode 100644 index 0000000000..fe7b8dd25a --- /dev/null +++ b/studio/backend/tests/test_worker_activates_correct_transformers.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Invariant: after the training worker runs its preflight and then activates the transformers +sidecar, the in-process ``transformers`` must be the sidecar version the model requires -- not the +default 4.57.x that the base environment ships. + +The CPU-only "does it choose the correct transformers version" guard, stronger than the pure +import-order check in ``test_training_worker_import_discipline.py``: it runs the REAL tier detection +(``get_transformers_tier``) and REAL activation (``activate_transformers_for_subprocess``) for a +transformers-5.x model (Qwen3.5, tier 530) and asserts the version actually switched. It catches the +whole failure family at once: + + * a stale pre-activation ``transformers`` import (the #6951 / ``TokenizersBackend`` regression: an + already-cached 4.57.x defeats the sidecar's ``sys.path`` prepend), + * a wrong tier selected for a 5.x model, and + * activation not actually swapping the resident module. + +Why the CUDA spoof matters (verified): ``unsloth_zoo``'s eager ``import transformers`` only happens on +its full, GPU-present init path. On a GPU-less runner it silently degrades and never preloads +transformers -- which would MASK the stale-import bug (the check would falsely pass). Spoofing +``torch.cuda`` so ``unsloth_zoo`` believes a GPU is present forces the real init path, exposing the +regression on CPU CI. The spoof mirrors ``tests/_zoo_aggressive_cuda_spoof.py`` but is inlined so the +test is self-contained in the ``studio-backend-ci`` matrix (whose conftest does not apply the shared +spoof). No GPU/network/weights/real sidecar needed: a one-line stub sidecar stands in for the 5.x venv, +so we only assert activation lands on it. + +Proven: passes on the fixed tree (active == 5.3.0) and fails on the buggy tree (active == 4.57.x) on +a simulated GPU-less runner. +""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +_BACKEND_DIR = Path(__file__).resolve().parent.parent # studio/backend +# Canonical CUDA spoof at the repo root (studio/backend -> studio -> repo root). Loaded by the +# subprocess when present (matches the consolidated CI); absent in a standalone studio checkout, where +# the subprocess falls back to a minimal inline spoof. +_SPOOF_PATH = _BACKEND_DIR.parent.parent / "tests" / "_zoo_aggressive_cuda_spoof.py" + +# Runs in a fresh interpreter with cwd == studio/backend so ``utils.*`` resolves like the worker. +# STUB_HOME (a pytest tmp dir) holds a throwaway ``.venv_t5_530`` sidecar exporting transformers 5.3.0. +_SNIPPET = r""" +import os, sys +sys.path.insert(0, os.getcwd()) + +# CUDA spoof so unsloth_zoo takes its full, transformers-importing init path on a GPU-less runner. +# Without it unsloth_zoo degrades and never preloads transformers, which would MASK the stale-import +# regression under test (verified). Prefer the repo's canonical spoof (single source of truth, and the +# one the consolidated CI already relies on); fall back to a minimal inline spoof so this also works in +# a standalone studio checkout. If torch is absent the fixed tree still passes below; the bug just +# would not be exposable in that shard. +try: + import torch # noqa: F401 + _sp = os.environ.get("SPOOF_PATH") + if _sp and os.path.exists(_sp): + import importlib.util + _spec = importlib.util.spec_from_file_location("_zoo_aggressive_cuda_spoof", _sp) + _mod = importlib.util.module_from_spec(_spec) + _spec.loader.exec_module(_mod) + _mod.apply() + else: + torch.cuda.is_available = lambda: True + torch.cuda.device_count = lambda: 1 + torch.cuda.current_device = lambda: 0 + torch.cuda.get_device_capability = lambda *a, **k: (8, 0) + torch.cuda.get_device_name = lambda *a, **k: "NVIDIA A100-SPOOFED" + torch.cuda.is_bf16_supported = lambda *a, **k: True + class _Props: + name = "NVIDIA A100-SPOOFED" + major = 8 + minor = 0 + total_memory = 80 * 1024**3 + multi_processor_count = 108 + torch.cuda.get_device_properties = lambda *a, **k: _Props() + torch.cuda.mem_get_info = lambda *a, **k: (0, 80 * 1024**3) +except Exception: + pass +os.environ["UNSLOTH_IS_PRESENT"] = "1" + +# Stub 5.x sidecar: activation only edits sys.path, so a package that merely exports __version__ is +# enough to prove the resident transformers switched to it. +home = os.environ["STUB_HOME"] +pkg = os.path.join(home, ".venv_t5_530", "transformers") +os.makedirs(pkg, exist_ok = True) +with open(os.path.join(pkg, "__init__.py"), "w") as f: + f.write('__version__ = "5.3.0"\n') +os.environ["UNSLOTH_STUDIO_HOME"] = home + +# Faithful worker preflight (worker.py: from utils.hf_xet_fallback import child_should_disable_xet). +# This is the exact stale-import trigger: on the buggy tree it pulls unsloth_zoo -> transformers 4.57.x +# into sys.modules BEFORE activation. +from utils.hf_xet_fallback import child_should_disable_xet +child_should_disable_xet({}) +_tf = sys.modules.get("transformers") +preload = _tf.__version__ if _tf is not None else None + +# Real tier detection + real activation, with the 530 sidecar pointed at the stub above. +import utils.transformers_version as tv +tv._VENV_T5_530_DIR = os.path.join(home, ".venv_t5_530") +tv._ensure_venv_t5_530_exists = lambda: True +tier = tv.get_transformers_tier("Qwen/Qwen3.5-9B", None) +tv.activate_transformers_for_subprocess("Qwen/Qwen3.5-9B", None) + +import transformers +print(f"RESULT tier={tier} preload={preload} active={transformers.__version__}") +""" + + +def _parse(stdout: str) -> dict[str, str]: + for line in stdout.splitlines(): + if line.startswith("RESULT "): + return dict(kv.split("=", 1) for kv in line.split()[1:]) + return {} + + +def test_worker_activates_correct_transformers_version(tmp_path): + """The worker's real preflight + activation for a transformers-5.x model (Qwen3.5, tier 530) must + leave the in-process ``transformers`` on the 5.x sidecar. A stale pre-activation import leaves the + default 4.57.x pinned and fails this assertion -- exactly the #6951 ``TokenizersBackend`` regression.""" + result = subprocess.run( + [sys.executable, "-c", _SNIPPET], + cwd = str(_BACKEND_DIR), + env = { + **__import__("os").environ, + "STUB_HOME": str(tmp_path), + **({"SPOOF_PATH": str(_SPOOF_PATH)} if _SPOOF_PATH.exists() else {}), + }, + capture_output = True, + text = True, + ) + assert result.returncode == 0, ( + "Worker preflight + activation harness crashed.\n" + f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + ) + parsed = _parse(result.stdout) + assert parsed, f"No RESULT line.\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + + # Correct tier chosen for a transformers-5.x model (pure, deterministic; no network/GPU). + assert parsed["tier"] == "530", ( + f"Wrong transformers tier for Qwen3.5 (expected 530, got {parsed['tier']}). " + "Tier detection regressed." + ) + + # Activation must actually swap the resident transformers to the sidecar version. If a preflight + # import cached 4.57.x first, the sidecar prepend is a no-op and this stays 4.57.x -- the bug. + assert parsed["active"] == "5.3.0", ( + "Sidecar activation did NOT switch the in-process transformers to the model's 5.x version " + f"(active={parsed['active']}, preloaded-before-activation={parsed['preload']}). A pre-activation " + "transformers import (directly or via unsloth_zoo) defeated the sidecar; 5.x models (Qwen3.5, " + "GLM-4.7, gemma-4) then fail with 'Tokenizer class TokenizersBackend does not exist'. See #6951." + ) diff --git a/studio/backend/utils/coding_agents.py b/studio/backend/utils/coding_agents.py new file mode 100644 index 0000000000..f7dd2f8357 --- /dev/null +++ b/studio/backend/utils/coding_agents.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Detect which `unsloth start ` coding-agent CLIs are on PATH. + +The web UI only ever shows the user the "claude" flavor of the `unsloth start` +command (see agent-command.ts), leaving anyone using Codex, OpenCode, and the +other supported agents to manually edit the copied command. This module gives +the frontend a way to ask which of those CLIs are actually installed so it can +default to one the user can run immediately. +""" + +import shutil + +# Keep in sync with the `unsloth start ` subcommands defined in +# unsloth_cli/commands/start.py. Each entry is the exact executable name that +# subcommand launches, so a hit here means `unsloth start ` can find the +# binary on PATH without the user installing anything first. +CODING_AGENTS: tuple[str, ...] = ("claude", "codex", "openclaw", "opencode", "hermes", "pi") + + +def _is_on_path(agent: str) -> bool: + # shutil.which is documented to return None on a miss, but PATH lookups can + # still raise (e.g. a permission error while probing a directory entry); + # this is an advisory check, so a lookup failure should read as "not + # installed" instead of breaking the settings endpoint. + try: + return shutil.which(agent) is not None + except OSError: + return False + + +def detect_installed_coding_agents() -> list[str]: + """Return the subset of CODING_AGENTS whose CLI binary is on PATH. + + Order follows CODING_AGENTS, not discovery order, so callers can treat the + first entry as the preferred default among the installed agents. + """ + return [agent for agent in CODING_AGENTS if _is_on_path(agent)] diff --git a/studio/backend/utils/datasets/completion_masking.py b/studio/backend/utils/datasets/completion_masking.py new file mode 100644 index 0000000000..c7c4a474e3 --- /dev/null +++ b/studio/backend/utils/datasets/completion_masking.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Completion-only masking policy shared by the CUDA and MLX training paths. + +Decides how train_on_responses_only is applied for a model: chat template +auto-detection first, manual TEMPLATE_TO_RESPONSES_MAPPER markers as the +fallback. gpt-oss included: its quantized checkpoints ship a different +chat template, so only detection from the actual template is reliable. +""" + +from .model_mappings import ( + MODEL_TO_TEMPLATE_MAPPER, + TEMPLATE_TO_RESPONSES_MAPPER, + is_gpt_oss_model_name, +) + + +def lookup_manual_markers(model_name): + """Return (template_name, instruction_part, response_part) from the + manual template table, with None parts when the model or template is + not mapped.""" + template = MODEL_TO_TEMPLATE_MAPPER.get((model_name or "").lower()) + markers = TEMPLATE_TO_RESPONSES_MAPPER.get(template) if template else None + if markers: + return template, markers["instruction"], markers["response"] + return template, None, None + + +def apply_completion_masking( + trainer, + model_name, + train_fn, + num_proc = None, + notify = None, + detect_fn = None, +): + """Apply completion-only masking with auto-detection first and the manual + template table as fallback. + + Args: + trainer: The platform trainer (SFTTrainer or MLXTrainer). + model_name: Model repo id used for table lookup and the gpt-oss + renamed-checkpoint fallback. + train_fn: The platform train_on_responses_only callable. + num_proc: Forwarded to train_fn when not None (CUDA path only). + notify: Optional callback notify(level, message) with level "info" or + "warning" for user-visible progress and warnings. + detect_fn: Marker detector (tokenizer/processor) -> (instruction_part, + response_part). Defaults to unsloth_zoo's get_chat_template_parts, + which raises loudly when the template cannot be parsed. Test seam. + + Returns: + (trainer, applied): the possibly wrapped trainer and whether masking + was applied. When applied is False the trainer is unchanged and + training runs on full sequences. + + Only marker DETECTION failures trigger the table fallback. Exceptions + raised while applying the masking (dataset map, tokenization) propagate + to the caller in both the auto and manual paths, so a real failure stops + the run instead of silently changing the training objective. + """ + if notify is None: + notify = lambda level, message: None + kwargs = {} + if num_proc is not None: + kwargs["num_proc"] = num_proc + + template, instruction_part, response_part = lookup_manual_markers(model_name) + + # gpt-oss goes auto-first: quantized/BF16 checkpoints ship a channel-less + # template, so the manual markers match nothing (zero tokens trained). Auto + # derives markers from whichever template ships, and per the harmony format + # only the final terminator carries stop supervision. Renamed checkpoints + # miss the exact-name table, so give the fallback the gpt-oss markers. + if is_gpt_oss_model_name(model_name) and not (instruction_part and response_part): + markers = TEMPLATE_TO_RESPONSES_MAPPER.get("gpt-oss") + if markers: + template = "gpt-oss" + instruction_part = markers["instruction"] + response_part = markers["response"] + processor = getattr(trainer, "processing_class", None) or getattr(trainer, "tokenizer", None) + # mlx-lm TokenizerWrapper hides underscore attrs, so preset _unsloth_* + # markers are invisible through it. Unwrap to the real tokenizer (as + # zoo's MLX resolver does) before the preset check and detection. + if type(processor).__name__ == "TokenizerWrapper": + wrapped = getattr(processor, "_tokenizer", None) + if wrapped is not None: + processor = wrapped + inner = getattr(processor, "tokenizer", processor) + if hasattr(inner, "_unsloth_input_part") and hasattr(inner, "_unsloth_output_part"): + # Markers preset on the tokenizer; zoo reuses them on a bare call. + trainer = train_fn(trainer, **kwargs) + notify( + "info", + "Train on responses only configured via tokenizer preset markers", + ) + return trainer, True + auto_instruction = auto_response = None + try: + if detect_fn is None: + # Torch-backed import is fine: the MLX train_fn itself requires + # unsloth_zoo.dataset_utils, so a torch-free host cannot mask either way. + from unsloth_zoo.dataset_utils import get_chat_template_parts as detect_fn + auto_instruction, auto_response = detect_fn(processor) + except Exception as e: + notify( + "warning", + f"Auto-detection of instruction/response markers failed ({e}); " + f"falling back to the template table", + ) + if auto_instruction and auto_response: + trainer = train_fn( + trainer, + instruction_part = auto_instruction, + response_part = auto_response, + **kwargs, + ) + notify( + "info", + "Train on responses only configured via chat template auto-detection", + ) + return trainer, True + + if instruction_part and response_part: + trainer = train_fn( + trainer, + instruction_part = instruction_part, + response_part = response_part, + **kwargs, + ) + notify( + "info", + f"Train on responses only configured with template table markers ({template})", + ) + return trainer, True + + notify( + "warning", + f"'Train on completions' could not be applied for {model_name}: no " + f"auto-detected or mapped instruction/response markers. Training " + f"will run on full sequences (prompts included).", + ) + return trainer, False diff --git a/studio/backend/utils/datasets/model_mappings.py b/studio/backend/utils/datasets/model_mappings.py index 463d26a692..65ba4b4688 100644 --- a/studio/backend/utils/datasets/model_mappings.py +++ b/studio/backend/utils/datasets/model_mappings.py @@ -485,9 +485,11 @@ TEMPLATE_TO_RESPONSES_MAPPER = { "instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n", }, + # No "" suffix: Qwen3-Thinking-2507 strips it from non-final turns + # and QwQ renders none, so a marker holding it masks those responses. "qwen3-thinking": { "instruction": "<|im_start|>user\n", - "response": "<|im_start|>assistant\n\n", + "response": "<|im_start|>assistant\n", }, "qwen3": { "instruction": "<|im_start|>user\n", @@ -525,29 +527,39 @@ TEMPLATE_TO_RESPONSES_MAPPER = { "instruction": "<|im_start|>user<|im_sep|>", "response": "<|im_start|>assistant<|im_sep|>", }, + # No surrounding spaces: in Mistral v0.3 they fold into neighbouring text + # tokens ("[INST]"/"[/INST]" are single special tokens), so padded strings + # never match and everything masks. Same for Llama-2's SentencePiece. "mistral": { - "instruction": "[INST] ", - "response": " [/INST]", + "instruction": "[INST]", + "response": "[/INST]", }, "llama": { - "instruction": "[INST] ", - "response": " [/INST]", + # -anchored: llama-2 tokenizes [INST] after as bare "[" on + # transformers 5.x (standalone gives space-prefixed "▁["), so an + # unanchored marker misses every turn boundary there. + "instruction": "[INST]", + "response": "[/INST]", }, "chatml": { "instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n", }, + # Leading newline required: Zephyr's role tags are plain text, and + # SentencePiece tokenizes "<|assistant|>" differently at text start than + # after "\n". Without the "\n" anchor the markers never match real + # turns, so every assistant token masks. "zephyr": { - "instruction": "<|user|>\n", - "response": "<|assistant|>\n", + "instruction": "\n<|user|>\n", + "response": "\n<|assistant|>\n", }, "unsloth": { - "instruction": ">>> User: ", - "response": ">>> Assistant: ", + "instruction": ">>> User:", + "response": ">>> Assistant:", }, "vicuna": { - "instruction": "USER: ", - "response": "ASSISTANT: ", + "instruction": "USER:", + "response": "ASSISTANT:", }, "alpaca": { "instruction": "### Instruction:\n", @@ -573,16 +585,21 @@ TEMPLATE_TO_RESPONSES_MAPPER = { "instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n", }, + # No trailing space: SentencePiece folds it into the next content token + # ("▁Hello"), so the padded marker never matches and masks everything. "starling": { - "instruction": "GPT4 Correct User: ", - "response": "GPT4 Correct Assistant: ", + "instruction": "GPT4 Correct User:", + "response": "GPT4 Correct Assistant:", }, "yi-chat": { "instruction": "<|im_start|>user\n", "response": "<|im_start|>assistant\n", }, + # "[gMASK]" appears once at text start, so a marker holding it matches + # no later user turn; "" is scaffolding GLM-4.x renders as a lone + # "" on non-final turns, so "<|assistant|>" never matches. "glm": { - "instruction": "[gMASK]<|user|>", - "response": "<|assistant|>", + "instruction": "<|user|>", + "response": "<|assistant|>", }, } diff --git a/studio/backend/utils/hf_xet_fallback.py b/studio/backend/utils/hf_xet_fallback.py index 15961ac03a..9bc4a60fad 100644 --- a/studio/backend/utils/hf_xet_fallback.py +++ b/studio/backend/utils/hf_xet_fallback.py @@ -1,341 +1,283 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Xet-primary HF downloads with an automatic HTTP fallback on a no-progress stall. +"""Studio shim over the shared ``unsloth_zoo.hf_xet_fallback`` Xet -> HTTP stall fallback. -Xet (``hf_xet``) is the fast default but can hang with no progress and no -exception, and a blocked native thread cannot be killed. Keep Xet primary; fall -back to plain HTTP only when the parent observes a stall. ``HF_HUB_DISABLE_XET`` -is read at import time, so the fallback runs in a fresh ``spawn`` child (not a -thread) that sets the env before importing ``huggingface_hub``. Cached files -short-circuit with no child; deterministic errors (401/403/404/disk-full) and -cancellation propagate without a fallback. Mirrors the safetensors inference -recovery in core/inference/{orchestrator,worker}.py. +Re-exports the shared API and injects Studio's marker-aware cache purge +(``prepare_cache_for_transport``) so the download manager keeps its ``.transport`` +marker semantics on the HTTP retry. + +Import discipline: ``unsloth_zoo``'s ``__init__`` eagerly imports ``transformers``. The workers +import this shim at startup (to decide the per-worker Xet env flip) *before* activating the model's +``transformers`` sidecar. Activation only prepends the sidecar to ``sys.path``, so a ``transformers`` +already cached in ``sys.modules`` (via an eager ``unsloth_zoo`` import here) wins -- pinning the +default 4.57.x and regressing Qwen3.5 / GLM-4.7 / gemma-4 training with +``Tokenizer class TokenizersBackend does not exist``. So the shared backend is loaded **lazily** +(``_load_shared``), only on first use of a heavy download helper, i.e. after the sidecar is active. +``child_should_disable_xet`` and the ``DEFAULT_*`` constants are defined locally so importing them +never triggers the heavy load. """ from __future__ import annotations -import multiprocessing as mp -import os -import queue -import signal -import sys import threading -import time from typing import Any, Callable, Optional -from loggers import get_logger - -logger = get_logger(__name__) - -_CTX = mp.get_context("spawn") - -# Defaults match the existing inference watchdog and hub shutdown deadline. +# Defaults mirror unsloth_zoo.hf_xet_fallback; plain literals so they resolve (including as +# default args below) without importing unsloth_zoo/transformers. +DEFAULT_GRACE_PERIOD = 10.0 DEFAULT_HEARTBEAT_INTERVAL = 30.0 DEFAULT_STALL_TIMEOUT = 180.0 -DEFAULT_GRACE_PERIOD = 10.0 -_POLL_INTERVAL = 0.5 + +# --- lazy shared-backend loader ---------------------------------------------------------------- +_shared: Any = None +_shared_available: Optional[bool] = None # None = not yet attempted +_shared_import_error: Optional[BaseException] = None +_load_lock = threading.Lock() -class DownloadStallError(RuntimeError): - """Raised when no download progress is observed for too long. +def _load_shared() -> bool: + """Import ``unsloth_zoo.hf_xet_fallback`` on demand; return True if available. Deferred so + importing this module at worker startup does not pull transformers in before the sidecar is + activated. Degrades (returns False) rather than crashing when unsloth_zoo is unavailable.""" + global _shared, _shared_available, _shared_import_error + if _shared_available is not None: + return _shared_available + with _load_lock: + if _shared_available is not None: + return _shared_available + try: + import unsloth_zoo.hf_xet_fallback as shared - Canonical home; orchestrator.py re-imports it so all paths share one type. - """ + _shared = shared + _shared_available = True + _shared_import_error = None + return True + except Exception as exc: # noqa: BLE001 - any import failure must degrade, not crash + # unsloth_zoo's __init__ runs torch/GPU detection, which raises on a torch-less/GPU-less + # host. The download helper needs none of it, so retry via UNSLOTH_ZOO_DISABLE_GPU_INIT. + _shared_import_error = exc + import os as _os + + _prev_gpu_init = _os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") + _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = "1" + try: + import unsloth_zoo.hf_xet_fallback as shared + + _shared = shared + _shared_available = True + _shared_import_error = None + return True + except Exception as exc2: # noqa: BLE001 - degrade so Studio still boots with plain HF + _shared_import_error = exc2 + _shared_available = False + import logging as _logging + + _logging.getLogger(__name__).warning( + "unsloth_zoo.hf_xet_fallback unavailable (%s); the Xet stall watchdog is " + "disabled. Install/upgrade unsloth_zoo (and its torch dependency) to " + "re-enable automatic Xet -> HTTP download recovery.", + _shared_import_error, + ) + return False + finally: + if _prev_gpu_init is None: + _os.environ.pop("UNSLOTH_ZOO_DISABLE_GPU_INIT", None) + else: + _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = _prev_gpu_init def child_should_disable_xet(config: dict) -> bool: - """Single source of truth for the per-worker Xet env flip.""" + """Single source of truth for the per-worker Xet env flip (mirrors + ``unsloth_zoo.hf_xet_fallback.child_should_disable_xet``). Deliberately lightweight: importing or + calling it must NOT pull in unsloth_zoo/transformers, so the worker can decide before activating + the transformers sidecar (see the module docstring).""" return bool(config.get("disable_xet")) -def get_hf_download_state( - repo_ids: Optional[list[str]] = None, *, repo_type: str = "model" -) -> Optional[tuple[int, bool]]: - """Return ``(total_on_disk_bytes, has_incomplete)`` for the active HF cache. - - Sparse-aware (st_blocks based) so a sparse Xet/``hf_transfer`` ``.incomplete`` - is not mistaken for full-size progress. ``None`` means the state could not be - measured, so callers skip stall logic for that tick. - """ - try: - from hub.utils.hf_cache_state import ( - blob_bytes_present, - has_active_incomplete_blobs, - hf_cache_root, - iter_active_repo_cache_dirs, - ) - - if hf_cache_root() is None: - return (0, False) - - total = 0 - has_incomplete = False - for repo_id in repo_ids or []: - # Skip local paths: HF IDs never start with / . ~ or contain "\". - if not repo_id or repo_id.startswith(("/", ".", "~")) or "\\" in repo_id: - continue - for entry in iter_active_repo_cache_dirs(repo_type, repo_id): - blobs_dir = entry / "blobs" - if not blobs_dir.is_dir(): - continue - for blob in blobs_dir.iterdir(): - try: - if blob.is_file(): - total += blob_bytes_present(blob) - except OSError: - pass - if has_active_incomplete_blobs(repo_type, repo_id): - has_incomplete = True - return (total, has_incomplete) - except Exception as e: - logger.debug("Failed to determine HF download state: %s", e) - return None +# --- degraded stubs (used only when unsloth_zoo is unavailable) ------------------------------- +class _DegradedDownloadStallError(RuntimeError): + """Stub mirror so callers' ``except`` clauses resolve; never raised in degraded mode.""" -def start_watchdog( +def _degraded_get_hf_download_state(*args: Any, **kwargs: Any) -> None: + return None # unmeasurable -> the (absent) watchdog never fires + + +def _degraded_start_watchdog( *, - repo_ids: list[str], - on_stall: Callable[[str], None], - repo_type: str = "model", + on_heartbeat: "Optional[Callable[[str], None]]" = None, interval: float = DEFAULT_HEARTBEAT_INTERVAL, - stall_timeout: float = DEFAULT_STALL_TIMEOUT, xet_disabled: bool = False, - on_heartbeat: Optional[Callable[[str], None]] = None, -) -> threading.Event: - """Start a daemon thread that fires ``on_stall(message)`` exactly once iff a - ``*.incomplete`` is present AND the on-disk size is unchanged for - *stall_timeout* seconds. The timer resets while no ``*.incomplete`` exists, so - post-download init is never misread as a stall. Returns a stop event the - caller sets when the download phase ends. - """ + **kwargs: Any, +) -> "threading.Event": + # No stall detection, but keep emitting heartbeats so the orchestrator's inactivity deadline + # is not tripped during a long download. stop = threading.Event() + if on_heartbeat is None: + return stop transport = "https" if xet_disabled else "xet" - fired = False def _beat() -> None: - nonlocal fired - state = get_hf_download_state(repo_ids, repo_type = repo_type) - last_size = state[0] if state is not None else 0 - last_change = time.monotonic() - while not stop.wait(interval): - state = get_hf_download_state(repo_ids, repo_type = repo_type) - now = time.monotonic() - - if state is None: - if on_heartbeat is not None: - on_heartbeat(f"Downloading ({transport} transport)...") - continue - - current_size, has_incomplete = state - if current_size != last_size: - last_size = current_size - last_change = now - - # Reset unless .incomplete confirms an active download, so model init - # and lock waits are not counted as a stall. - if not has_incomplete: - last_change = now - elif now - last_change >= stall_timeout: - if not fired: - fired = True - on_stall( - f"Download appears stalled ({transport} transport) " - f"-- no progress for {int(now - last_change)}s" - ) - return - - if on_heartbeat is not None: + try: on_heartbeat(f"Downloading ({transport} transport)...") + except Exception: + pass - threading.Thread(target = _beat, daemon = True, name = "hf-xet-watchdog").start() + threading.Thread( + target = _beat, + daemon = True, + name = "hf-xet-degraded-heartbeat", + ).start() return stop -def _download_child_entry( - *, - repo_id: str, - filename: str, - token: Optional[str], - repo_type: str, - disable_xet: bool, - result_queue: Any, - force_download: bool = False, -) -> None: - """Spawn-child entrypoint: download one file and report the result. - - Top-level and picklable. Sets the Xet env BEFORE importing huggingface_hub, - forms its own process group so the parent can kill the whole transfer, and - never logs the token or signed URLs. - """ - # Die with Studio on Linux (this mp child gets no parent-set preexec_fn). - try: - from utils.process_lifetime import bind_current_process_to_parent_lifetime - bind_current_process_to_parent_lifetime() - except Exception: - pass - - if hasattr(os, "setsid"): - try: - os.setsid() - except OSError: - pass - - if disable_xet: - os.environ["HF_HUB_DISABLE_XET"] = "1" - # Keep the HTTP writer sequential and resumable (hf_transfer leaves sparse - # partials a sequential resume cannot safely continue). - os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0" - os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") - - # Test-only fault injection (never set in production): stall the Xet attempt - # so the watchdog + HTTP fallback can be exercised against a real repo. - if not disable_xet and os.environ.get("UNSLOTH_HF_XET_FORCE_STALL") == "1": - import time as _t - try: - from huggingface_hub.constants import HF_HUB_CACHE - - blobs = os.path.join(HF_HUB_CACHE, "models--" + repo_id.replace("/", "--"), "blobs") - os.makedirs(blobs, exist_ok = True) - with open(os.path.join(blobs, "xet-force-stall.incomplete"), "wb") as fh: - fh.write(b"\0" * 4096) - except OSError: - pass - while True: - _t.sleep(3600) - - try: - from huggingface_hub import hf_hub_download - path = hf_hub_download( - repo_id = repo_id, - filename = filename, - repo_type = repo_type, - token = token, - force_download = force_download, - ) - result_queue.put({"ok": True, "path": path}) - except BaseException as e: # noqa: BLE001 - report every failure to the parent - error = f"{type(e).__name__}: {e}" - try: - from hub.utils.download_registry import scrub_secrets - error = scrub_secrets(error, hf_token = token) - except Exception: - pass - result_queue.put({"ok": False, "error": error}) +def _degraded_cancelled(cancel_event: "Optional[threading.Event]") -> bool: + return cancel_event is not None and cancel_event.is_set() -def _terminate_process_group(proc: "mp.process.BaseProcess", grace_period: float) -> None: - """Kill *proc* and its whole process group (Xet may spawn helper procs). - - The child calls ``os.setsid()`` so its pgid equals its pid; signal via - ``os.killpg(pid, ...)`` -- NOT ``getpgid``, which before the child becomes a - group leader resolves to OUR group. SIGTERM, then SIGKILL after *grace_period*. - """ - pid = proc.pid - - def _signal_group(sig: int) -> None: - if pid is not None and hasattr(os, "killpg"): - try: - os.killpg(pid, sig) - return - except (ProcessLookupError, PermissionError, OSError): - pass - # Windows or pre-setsid: best effort on the single process. - try: - proc.terminate() if sig != getattr(signal, "SIGKILL", -9) else proc.kill() - except Exception: - pass - - _signal_group(getattr(signal, "SIGTERM", signal.SIGINT)) - proc.join(timeout = grace_period) - if proc.is_alive(): - _signal_group(getattr(signal, "SIGKILL", signal.SIGTERM)) - proc.join(timeout = 5.0) - - -def _run_download_attempt( +def _degraded_hf_hub_download_with_xet_fallback( repo_id: str, filename: str, token: Optional[str], *, - repo_type: str, - disable_xet: bool, - cancel_event: Optional[threading.Event], - stall_timeout: float, - interval: float, - grace_period: float, - on_status: Optional[Callable[[str], None]], + repo_type: str = "model", + revision: Optional[str] = None, + cache_dir: Optional[str] = None, force_download: bool = False, -) -> tuple[str, Optional[str]]: - """Run one download in a spawn child supervised by the no-progress watchdog. + cancel_event: "Optional[threading.Event]" = None, + **_ignored: Any, +) -> str: + # Keep the cancellation contract: do not start or return a download once cancelled. + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") - Returns ``("ok", path)``, ``("stall", None)``, ``("cancelled", None)``, or - ``("error", message)``. This is the seam tests monkeypatch to avoid spawning. - """ - result_queue: Any = _CTX.Queue() - proc = _CTX.Process( - target = _download_child_entry, - kwargs = dict( - repo_id = repo_id, - filename = filename, - token = token, - repo_type = repo_type, - disable_xet = disable_xet, - result_queue = result_queue, - force_download = force_download, - ), - daemon = True, - ) - proc.start() - from utils.process_lifetime import adopt_pid + from huggingface_hub import hf_hub_download - adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep) - - stalled = threading.Event() - stop_watchdog = start_watchdog( - repo_ids = [repo_id], - on_stall = lambda msg: stalled.set(), + path = hf_hub_download( + repo_id = repo_id, + filename = filename, + token = token, repo_type = repo_type, - interval = interval, - stall_timeout = stall_timeout, - xet_disabled = disable_xet, - on_heartbeat = on_status, + revision = revision, + cache_dir = cache_dir, + force_download = force_download, ) + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + return path - result: Optional[dict] = None + +def _degraded_snapshot_download_with_xet_fallback( + repo_id: str, + *, + revision: Optional[str] = None, + token: Optional[str] = None, + repo_type: str = "model", + cache_dir: Optional[str] = None, + allow_patterns: Optional[Any] = None, + ignore_patterns: Optional[Any] = None, + force_download: bool = False, + cancel_event: "Optional[threading.Event]" = None, + **_ignored: Any, +) -> str: + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + + from huggingface_hub import snapshot_download + + path = snapshot_download( + repo_id = repo_id, + repo_type = repo_type, + revision = revision, + token = token, + cache_dir = cache_dir, + allow_patterns = allow_patterns, + ignore_patterns = ignore_patterns, + force_download = force_download, + ) + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + return path + + +# --- lazy attribute access for the heavy shared API ------------------------------------------- +# ``DownloadStallError`` (class identity matters for ``except``), ``start_watchdog`` and +# ``get_hf_download_state`` come from the shared backend when available, else the degraded stubs. +# Resolved via PEP 562 ``__getattr__`` so ``from utils.hf_xet_fallback import X`` triggers the load +# only for these heavy names, not for ``child_should_disable_xet`` / ``DEFAULT_*``. +_DEGRADED_ATTRS = { + "DownloadStallError": _DegradedDownloadStallError, + "start_watchdog": _degraded_start_watchdog, + "get_hf_download_state": _degraded_get_hf_download_state, +} + +# Annotation-only declarations for the three names above: they bind NO value, so lookup still misses +# and PEP 562 ``__getattr__`` resolves them lazily -- but ruff/pyflakes see them as defined, so listing +# them in ``__all__`` does not trip F822 (while F822 still catches a real typo elsewhere in the list). +DownloadStallError: type +start_watchdog: Any +get_hf_download_state: Any + + +def __getattr__(name: str) -> Any: + if name in _DEGRADED_ATTRS: + if _load_shared(): + return getattr(_shared, name) + return _DEGRADED_ATTRS[name] + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +# Indirection seam the public wrappers call (and tests monkeypatch): lazy-load the shared backend, +# then dispatch to it or the degraded stub. The ``_shared_*`` names preserve the pre-refactor contract. +def _shared_hf_hub_download_with_xet_fallback(*args: Any, **kwargs: Any) -> str: + impl = ( + _shared.hf_hub_download_with_xet_fallback + if _load_shared() + else _degraded_hf_hub_download_with_xet_fallback + ) + return impl(*args, **kwargs) + + +def _shared_snapshot_download_with_xet_fallback(*args: Any, **kwargs: Any) -> str: + impl = ( + _shared.snapshot_download_with_xet_fallback + if _load_shared() + else _degraded_snapshot_download_with_xet_fallback + ) + return impl(*args, **kwargs) + + +__all__ = [ + "DEFAULT_GRACE_PERIOD", + "DEFAULT_HEARTBEAT_INTERVAL", + "DEFAULT_STALL_TIMEOUT", + "DownloadStallError", + "child_should_disable_xet", + "get_hf_download_state", + "start_watchdog", + "hf_hub_download_with_xet_fallback", + "snapshot_download_with_xet_fallback", +] + + +def _studio_prepare_for_http(repo_type: str, repo_id: str) -> None: + """Studio's marker-aware purge before an HTTP resume, keeping the download manager's ``.transport`` + accounting consistent (vs unsloth_zoo's generic default). Guarded: a purge failure is logged, + not fatal to the retry.""" try: - while proc.is_alive(): - if cancel_event is not None and cancel_event.is_set(): - _terminate_process_group(proc, grace_period) - return ("cancelled", None) - if stalled.is_set(): - _terminate_process_group(proc, grace_period) - return ("stall", None) - try: - result = result_queue.get(timeout = _POLL_INTERVAL) - break - except queue.Empty: - continue - else: - # Process exited; drain any result it enqueued. - try: - result = result_queue.get_nowait() - except queue.Empty: - result = None - finally: - stop_watchdog.set() - proc.join(timeout = grace_period) - - if result is None: - return ( - "error", - f"download process for '{repo_id}/{filename}' exited " - f"(code={proc.exitcode}) without a result", - ) - if result.get("ok"): - return ("ok", result["path"]) - return ("error", result.get("error") or "unknown download error") + from hub.utils.download_registry import prepare_cache_for_transport + prepare_cache_for_transport(repo_type, repo_id, "http") + except Exception as exc: + try: + from loggers import get_logger + get_logger(__name__).debug( + "Studio prepare_cache_for_transport failed for %s: %s", repo_id, exc + ) + except ModuleNotFoundError as logger_exc: + if logger_exc.name != "loggers": + raise def hf_hub_download_with_xet_fallback( @@ -345,83 +287,32 @@ def hf_hub_download_with_xet_fallback( *, cancel_event: Optional[threading.Event] = None, repo_type: str = "model", + revision: Optional[str] = None, stall_timeout: float = DEFAULT_STALL_TIMEOUT, interval: float = DEFAULT_HEARTBEAT_INTERVAL, grace_period: float = DEFAULT_GRACE_PERIOD, on_status: Optional[Callable[[str], None]] = None, force_download: bool = False, ) -> str: - """Download a single file with Xet primary and HTTP as a stall-only fallback. + """Single-file download via the shared fallback with Studio's marker-aware HTTP-retry prep. + ``force_download`` re-fetches a newer blob over a cached one (Studio's model-update path).""" + return _shared_hf_hub_download_with_xet_fallback( + repo_id, + filename, + token, + cancel_event = cancel_event, + repo_type = repo_type, + revision = revision, + stall_timeout = stall_timeout, + interval = interval, + grace_period = grace_period, + on_status = on_status, + force_download = force_download, + prepare_for_http_fn = _studio_prepare_for_http, + ) - Returns the local cache path. Raises ``RuntimeError("Cancelled")`` if - *cancel_event* is set, re-raises a deterministic child error unchanged (no - fallback), and raises ``DownloadStallError`` only if BOTH transports stall. - When *force_download* is True the cache-first early-return is skipped and the - flag is threaded to ``hf_hub_download`` so a newer remote blob is re-fetched - even if an older blob is already cached. - """ - # Finalized blob already cached: return it with no child and no network. - # Skipped when force_download is set so an update re-fetches a newer blob. - if not force_download: - try: - from huggingface_hub import try_to_load_from_cache - cached = try_to_load_from_cache(repo_id, filename, repo_type = repo_type) - if isinstance(cached, str) and os.path.exists(cached): - return cached - except Exception as e: - logger.debug("Cached probe failed for %s/%s: %s", repo_id, filename, e) - - if cancel_event is not None and cancel_event.is_set(): - raise RuntimeError("Cancelled") - - disable_xet = False - for attempt in range(2): - if disable_xet: - # Purge a non-HTTP partial before resuming over HTTP: an HTTP resume - # over a sparse Xet/hf_transfer partial silently corrupts the blob. - try: - from hub.utils.download_registry import prepare_cache_for_transport - prepare_cache_for_transport(repo_type, repo_id, "http") - except Exception as e: - logger.debug("prepare_cache_for_transport failed for %s: %s", repo_id, e) - - kind, payload = _run_download_attempt( - repo_id, - filename, - token, - repo_type = repo_type, - disable_xet = disable_xet, - cancel_event = cancel_event, - stall_timeout = stall_timeout, - interval = interval, - grace_period = grace_period, - on_status = on_status, - force_download = force_download, - ) - - if kind == "ok": - return payload # type: ignore[return-value] - if kind == "cancelled": - raise RuntimeError("Cancelled") - if kind == "error": - # Deterministic failure: the other transport would fail identically. - raise RuntimeError(payload) - # kind == "stall" - if attempt == 0 and not disable_xet: - logger.warning( - "Download stalled for '%s/%s' -- retrying with HF_HUB_DISABLE_XET=1", - repo_id, - filename, - ) - if on_status is not None: - on_status(f"{repo_id}/{filename}: Xet stalled, retrying over HTTP") - disable_xet = True - continue - raise DownloadStallError( - f"Download stalled for '{repo_id}/{filename}' even with " - f"HF_HUB_DISABLE_XET=1 -- check your network connection" - ) - - # Unreachable: the loop either returns or raises on each attempt. - raise DownloadStallError(f"Download failed for '{repo_id}/{filename}'") +def snapshot_download_with_xet_fallback(repo_id: str, **kwargs: Any) -> str: + """Whole-repo download via the shared fallback with Studio's marker-aware HTTP-retry prep.""" + kwargs.setdefault("prepare_for_http_fn", _studio_prepare_for_http) + return _shared_snapshot_download_with_xet_fallback(repo_id, **kwargs) diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index c16ae91467..1bcbfbf95a 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -514,6 +514,12 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path logger.info("llama update: installing", cmd = " ".join(cmd)) # Stream progress lines into job["progress"]. env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5") + # Preserve a Vulkan install across updates: detect_host on a CUDA/ROCm + # box would otherwise re-route and silently replace the Vulkan build. + # Re-assert it via the same env flag setup uses (mirrors + # _rocm_install_args). + if asset and "vulkan" in asset.lower(): + env["UNSLOTH_FORCE_VULKAN"] = "1" proc = subprocess.Popen( cmd, stdout = subprocess.PIPE, diff --git a/studio/backend/utils/mlx_repair.py b/studio/backend/utils/mlx_repair.py index 520c11c3b1..7e1c9864c9 100644 --- a/studio/backend/utils/mlx_repair.py +++ b/studio/backend/utils/mlx_repair.py @@ -45,9 +45,21 @@ DISABLE_ENV_VAR = "UNSLOTH_DISABLE_MLX_AUTOREPAIR" # deps). mlx-vlm especially must be >=0.4.4: an older one still imports but # breaks VLM Train/Export, so installing it would wrongly clear chat-only. _MLX_MIN_VERSIONS = {"mlx": "0.22.0", "mlx-lm": "0.22.0", "mlx-vlm": "0.4.4"} +# mlx-lm 0.31.3 regressed QK-norm archs (gemma4 / qwen3_5): strict load_weights +# rejects q_norm/k_norm, so a self-heal must not pull it. mlx-lm #1242. +_MLX_BAD_VERSIONS = {"mlx-lm": ("0.31.3",)} _MLX_PACKAGE_NAMES = tuple(_MLX_MIN_VERSIONS) _MLX_RUNTIME_IMPORTS = ("mlx.core", "mlx_lm", "mlx_lm.sample_utils", "mlx_vlm") -MLX_PACKAGES = tuple(f"{name}>={version}" for name, version in _MLX_MIN_VERSIONS.items()) + + +def _mlx_spec(name: str, version: str) -> str: + spec = f"{name}>={version}" + for bad in _MLX_BAD_VERSIONS.get(name, ()): + spec += f",!={bad}" + return spec + + +MLX_PACKAGES = tuple(_mlx_spec(name, version) for name, version in _MLX_MIN_VERSIONS.items()) _MLX_REINSTALL_ARGS = tuple( arg for name in _MLX_PACKAGE_NAMES for arg in ("--reinstall-package", name) ) @@ -140,7 +152,12 @@ def _mlx_versions_satisfy_minimums() -> bool: return False for name, minimum in _MLX_MIN_VERSIONS.items(): try: - if Version(_dist_version(name)) < Version(minimum): + installed = Version(_dist_version(name)) + if installed < Version(minimum): + return False + # A known-broken build counts as unsatisfied so the self-heal + # reinstalls a good one; Version compare matches 0.31.3(.0/+local). + if any(installed == Version(bad) for bad in _MLX_BAD_VERSIONS.get(name, ())): return False except PackageNotFoundError: return False diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 5d8458e5f0..281ca24281 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1617,36 +1617,60 @@ def _iter_hf_cache_snapshots(repo_id: str): cache_dir = Path(hf_constants.HF_HUB_CACHE) target = f"models--{repo_id.replace('/', '--')}".lower() - repo_dir: Optional[Path] = None + repo_dirs: list[Path] = [] try: if not cache_dir.is_dir(): return for entry in cache_dir.iterdir(): if entry.is_dir() and entry.name.lower() == target: - repo_dir = entry - break + repo_dirs.append(entry) except OSError: return - if repo_dir is None: + if not repo_dirs: return - snapshots = repo_dir / "snapshots" - try: - if not snapshots.is_dir(): - return - snap_dirs = [s for s in snapshots.iterdir() if s.is_dir()] - except OSError: + snap_dirs: list[Path] = [] + for repo_dir in repo_dirs: + snapshots = repo_dir / "snapshots" + try: + if snapshots.is_dir(): + for snap_dir in snapshots.iterdir(): + try: + if snap_dir.is_dir(): + snap_dirs.append(snap_dir) + except OSError: + continue + except OSError: + continue + if not snap_dirs: return - snap_dirs.sort(key = lambda s: s.stat().st_mtime, reverse = True) - yield from snap_dirs + snap_dirs_with_mtime = [] + for snap_dir in snap_dirs: + try: + snap_dirs_with_mtime.append((snap_dir.stat().st_mtime, snap_dir)) + except OSError: + continue + snap_dirs_with_mtime.sort(key = lambda item: item[0], reverse = True) + yield from (snap_dir for _, snap_dir in snap_dirs_with_mtime) def _list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufVariantInfo], bool]]: - """Variants from the local HF cache snapshot, or None if not cached.""" + """Variants from the local HF cache snapshot, or None if not cached. + + A newer snapshot can hold only a companion file (for example a vision + projector fetched on demand) while the quant files live in an older + snapshot. Returning the first snapshot that merely reports a vision flag + would shadow those real variants, so keep scanning older snapshots for + actual variants and carry the vision flag across snapshots. + """ + any_vision = False for snap in _iter_hf_cache_snapshots(repo_id): variants, has_vision = list_local_gguf_variants(str(snap)) - if variants or has_vision: - return variants, has_vision + any_vision = any_vision or has_vision + if variants: + return variants, any_vision + if any_vision: + return [], True return None diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index 2a63caa5b2..a69673f081 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -28,7 +28,9 @@ Strategy: sys.path swap using the same directories pre-installed by setup.sh. """ +import ast import importlib +import importlib.util import json import structlog from loggers import get_logger @@ -173,6 +175,7 @@ _TRANSFORMERS_530_ARCHITECTURES: set[str] = { "Qwen3MoeForCausalLM", "Qwen3NextForCausalLM", "Glm4MoeLiteForCausalLM", + "Lfm2MoeForCausalLM", "Lfm2VlForConditionalGeneration", } _TRANSFORMERS_530_MODEL_TYPES: set[str] = { @@ -183,6 +186,7 @@ _TRANSFORMERS_530_MODEL_TYPES: set[str] = { "qwen3_moe", "qwen3_next", "glm4_moe_lite", + "lfm2_moe", "lfm2_vl", } @@ -870,6 +874,116 @@ def _cached_config_json(model_name: str, hf_token: str | None) -> dict | None: return _config_json_cache.get(_token_cache_key(model_name, hf_token)) +# --- Static tier from CONFIG_MAPPING_NAMES (AST only: no import/network/exec) --- +# A model_type absent from an overlay's mapping can't load there. Parse each sidecar's +# config map from source and pick the lowest tier that ships it, so a new arch routes +# correctly with no per-model table edit. Only ever upgrades default, never lowers. +_config_mapping_cache: dict[str, frozenset[str]] = {} + + +def _overlay_transformers_dir(tier: str) -> str | None: + """transformers source dir for a tier, located without importing it.""" + if tier != "default": + root = {"530": _VENV_T5_530_DIR, "550": _VENV_T5_550_DIR, "510": _VENV_T5_510_DIR}.get(tier) + src = os.path.join(root, "transformers") if root else None + return src if src and _safe_is_dir(Path(src)) else None + # default: the base 4.x transformers. find_spec resolves to a 5.x sidecar if one + # is already on sys.path, so skip any .venv_t5_* / llmcompressor overlay dir. + sidecars = tuple( + os.path.abspath(d) + os.sep + for d in (_VENV_T5_530_DIR, _VENV_T5_550_DIR, _VENV_T5_510_DIR, _VENV_LLMCOMPRESSOR_DIR) + ) + candidates = [] + try: + spec = importlib.util.find_spec("transformers") + if spec and spec.origin: + candidates.append(os.path.dirname(spec.origin)) + except Exception: + pass + candidates += [os.path.join(e, "transformers") for e in sys.path if e] + for c in candidates: + if _safe_is_dir(Path(c)) and not os.path.abspath(c).startswith(sidecars): + return c + return None + + +def _mapping_first_keys(value: ast.AST) -> set[str]: + """First keys of a dict literal, or of an OrderedDict(...)/dict(...)/.update(...) + built from 2-tuple lists and **{...} unpacking.""" + + def keys_of(node): + if isinstance(node, ast.Dict): + return list(node.keys) + if isinstance(node, (ast.List, ast.Tuple)): + return [ + el.elts[0] for el in node.elts if isinstance(el, (ast.Tuple, ast.List)) and el.elts + ] + return [] + + nodes = keys_of(value) + if isinstance(value, ast.Call): + for a in value.args: + nodes += keys_of(a) + for kw in value.keywords: # **{...} unpacking has kw.arg is None + if kw.arg is None: + nodes += keys_of(kw.value) + return {n.value for n in nodes if isinstance(n, ast.Constant) and isinstance(n.value, str)} + + +def _config_model_types(tier: str) -> frozenset[str]: + """model_type keys in a tier's CONFIG_MAPPING_NAMES (5.10 moved it to auto_mappings.py).""" + cached = _config_mapping_cache.get(tier) + if cached is not None: + return cached + tdir = _overlay_transformers_dir(tier) + if tdir is None: + return frozenset() # overlay not provisioned yet; do not cache so a later call re-reads + keys: set[str] = set() + for rel in ("models/auto/configuration_auto.py", "models/auto/auto_mappings.py"): + path = Path(tdir) / rel + if not _safe_is_file(path): + continue + try: + tree = ast.parse(path.read_text(encoding = "utf-8")) + for node in ast.walk(tree): + # direct binding, or a CONFIG_MAPPING_NAMES.update({...}) mutation + if isinstance(node, ast.Assign) and any( + isinstance(t, ast.Name) and t.id == "CONFIG_MAPPING_NAMES" for t in node.targets + ): + keys |= _mapping_first_keys(node.value) + elif isinstance(node, ast.Expr) and isinstance(node.value, ast.Call): + fn = node.value.func + if ( + isinstance(fn, ast.Attribute) + and fn.attr == "update" + and isinstance(fn.value, ast.Name) + and fn.value.id == "CONFIG_MAPPING_NAMES" + ): + keys |= _mapping_first_keys(node.value) + except Exception: + continue + result = frozenset(keys) + _config_mapping_cache[tier] = result + return result + + +def _tier_from_config_mapping(cfg: dict) -> str | None: + """Lowest tier whose transformers ships cfg's model_type, or None if unknown.""" + model_type = cfg.get("model_type") + if not isinstance(model_type, str): + for key in _NESTED_CONFIG_KEYS: + sub = cfg.get(key) + if isinstance(sub, dict) and isinstance(sub.get("model_type"), str): + model_type = sub["model_type"] + break + if not isinstance(model_type, str): + return None + for tier in sorted(_TIER_RANK, key = _TIER_RANK.get): + if model_type in _config_model_types(tier): + return tier + return None + + # --- AutoConfig probe: general tier resolution for ambiguous models ---------- # When the cheap signals only say "needs some 5.x", parse config.json with the built-in # parser in each candidate sidecar (lowest first) instead of guessing. Generalizes beyond @@ -1210,6 +1324,14 @@ def get_transformers_tier( match, ) return tier + static = _tier_from_config_mapping(cfg) + if static is not None and static != "default": + logger.info( + "Transformers tier %s selected for %s (config mapping: model_type absent below)", + static, + model_name, + ) + return static local_tc = Path(model_name) / "tokenizer_config.json" if _safe_is_file(local_tc) and _check_tokenizer_config_needs_v5(model_name, hf_token): if not probe: @@ -1269,6 +1391,18 @@ def get_transformers_tier( return override logger.info("Transformers tier 530 selected for %s (config.json check)", model_name) return "530" + # _load_config_json (not the cache-only reader) so a config served from the hub + # cache during a transient outage still feeds the mapping resolver. + remote_cfg = _load_config_json(model_name, hf_token) + if remote_cfg is not None: + static = _tier_from_config_mapping(remote_cfg) + if static is not None and static != "default": + logger.info( + "Transformers tier %s selected for %s (config mapping: model_type absent below)", + static, + model_name, + ) + return static if _check_tokenizer_config_needs_v5(model_name, hf_token): if not probe: return "530" diff --git a/studio/backend/utils/wheel_utils.py b/studio/backend/utils/wheel_utils.py index 98697df83c..1b5926fd49 100644 --- a/studio/backend/utils/wheel_utils.py +++ b/studio/backend/utils/wheel_utils.py @@ -26,11 +26,14 @@ FLASH_ATTN_RELEASE_BASE_URL = "https://github.com/Dao-AILab/flash-attention/rele def has_blackwell_gpu() -> bool: """Return True if any visible NVIDIA GPU has compute capability >= 10.0 (Blackwell). - Dao-AILab ships no flash-attention wheels for these archs and older-arch wheels - fail to load, so callers use this to skip the flash-attn install path. Cached - for the process lifetime; tests mocking nvidia-smi must call + Cached for the process lifetime; tests mocking nvidia-smi must call ``has_blackwell_gpu.cache_clear()`` first. """ + # Detection disabled for now: Dao-AILab ships Blackwell (sm_100+) flash-attn + # wheels and url_exists() already gates resolution, so we no longer skip + # flash-attn on Blackwell. The nvidia-smi probe below is kept for possible + # future arch-based gating; drop this early return to re-enable it. + return False exe = shutil.which("nvidia-smi") if not exe: return False @@ -117,6 +120,19 @@ def probe_torch_wheel_env(*, timeout: int | None = None) -> dict[str, str] | Non return env +# torch 2.11 has no native prebuilt wheels for flash-attn / causal-conv1d / mamba +# yet, but their torch 2.10 CUDA wheels load and pass the projects' own test suites +# on torch 2.11 (verified on B200: FA2 fwd/bwd, causal-conv1d, and mamba selective +# scan all match reference). Reuse the torch 2.10 wheels on torch 2.11 so a 2.11 +# install still gets these prebuilt accelerators instead of building from source. +_PREBUILT_WHEEL_TORCH_MM = {"2.11": "2.10"} + + +def prebuilt_wheel_torch_mm(torch_mm: str) -> str: + """Map a torch major.minor to the one whose prebuilt accelerator wheels to use.""" + return _PREBUILT_WHEEL_TORCH_MM.get(torch_mm, torch_mm) + + def direct_wheel_url( *, filename_prefix: str, @@ -130,7 +146,7 @@ def direct_wheel_url( filename = ( f"{filename_prefix}-{package_version}" - f"+cu{env['cuda_major']}torch{env['torch_mm']}" + f"+cu{env['cuda_major']}torch{prebuilt_wheel_torch_mm(env['torch_mm'])}" f"cxx11abi{env['cxx11abi']}-{env['python_tag']}-{env['python_tag']}" f"-{env['platform_tag']}.whl" ) @@ -152,7 +168,7 @@ def flash_attn_package_version(torch_mm: str) -> str | None: def flash_attn_wheel_url(env: dict[str, str] | None) -> str | None: if env is None: return None - package_version = flash_attn_package_version(env["torch_mm"]) + package_version = flash_attn_package_version(prebuilt_wheel_torch_mm(env["torch_mm"])) if package_version is None: return None return direct_wheel_url( diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index fb1a1fc9c7..1a74b38524 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -81,12 +81,6 @@ import { TestTube01Icon, ZapIcon, } from "@hugeicons/core-free-icons"; -import { - exportConversationRawJsonl, - exportConversationCsv, - exportConversationShareGPT, -} from "@/features/chat/prompt-storage/prompt-storage-dialog"; -import { listStoredChatThreads } from "@/features/chat/utils/chat-history-storage"; import { Tooltip, TooltipContent, @@ -102,6 +96,7 @@ import { createChatProject, deleteChatProject, deleteChatItem, + listStoredChatThreads, moveChatItemToProject, renameChatItem, renameChatProject, @@ -174,6 +169,36 @@ const TestTubeOutlineIcon = TestTube01Icon.slice( 3, ) as typeof TestTube01Icon; + +type ConversationExportFormat = "raw-jsonl" | "csv" | "sharegpt-jsonl"; + +const CHAT_EXPORT_OPTIONS: Array<{ + label: string; + format: ConversationExportFormat; +}> = [ + { label: "Raw JSONL", format: "raw-jsonl" }, + { label: "CSV", format: "csv" }, + { label: "ShareGPT JSONL", format: "sharegpt-jsonl" }, +]; + +async function exportConversationByFormat( + threadId: string, + format: ConversationExportFormat, +): Promise { + const exports = await import( + "@/features/chat/prompt-storage/prompt-storage-dialog" + ); + switch (format) { + case "raw-jsonl": + return exports.exportConversationRawJsonl(threadId); + case "csv": + return exports.exportConversationCsv(threadId); + case "sharegpt-jsonl": + return exports.exportConversationShareGPT(threadId); + } +} + + function runStatusDotClass(status: TrainingRunSummary["status"]): string { switch (status) { case "running": @@ -359,7 +384,11 @@ export function AppSidebar() { const activeProjectId = isChatRoute ? ((search.project as string | undefined) ?? null) : null; - const { items: allChatItems } = useChatSidebarItems({ + const { + items: allChatItems, + archivedItems: archivedChatItems, + loaded: chatItemsLoaded, + } = useChatSidebarItems({ enabled: !isStudioRoute, requireMessages: false, }); @@ -553,7 +582,14 @@ export function AppSidebar() { useEffect(() => { if (!pendingRename) return; const match = allChatItems.find((i) => i.id === pendingRename.id); - if (match && match.title === pendingRename.title) setPendingRename(null); + if (!match || match.title !== pendingRename.title) return; + queueMicrotask(() => { + setPendingRename((current) => + current?.id === pendingRename.id && current.title === pendingRename.title + ? null + : current, + ); + }); }, [allChatItems, pendingRename]); const [creatingProject, setCreatingProject] = useState(false); const [projectNameDraft, setProjectNameDraft] = useState(""); @@ -651,12 +687,6 @@ export function AppSidebar() { useState(null); const [deleteProjectFiles, setDeleteProjectFiles] = useState(false); - useEffect(() => { - if (confirmingDelete?.kind !== "project") { - setDeleteProjectFiles(false); - } - }, [confirmingDelete]); - async function commitDelete() { const target = confirmingDelete; if (!target) return; @@ -895,11 +925,7 @@ export function AppSidebar() { Export - {[ - { label: "Raw JSONL", fn: exportConversationRawJsonl }, - { label: "CSV", fn: exportConversationCsv }, - { label: "ShareGPT JSONL", fn: exportConversationShareGPT }, - ].map(({ label, fn }) => ( + {CHAT_EXPORT_OPTIONS.map(({ label, format }) => ( { @@ -907,7 +933,9 @@ export function AppSidebar() { const ids = item.type === "single" ? [item.id] : (await listStoredChatThreads({ pairId: item.id })).map((t) => t.id); - await Promise.all(ids.map((id) => fn(id))); + await Promise.all( + ids.map((id) => exportConversationByFormat(id, format)), + ); } catch { toast.error("Export failed."); } @@ -1306,6 +1334,16 @@ export function AppSidebar() { renderChatSidebarItem(item, "recent"), )} + {/* "No chats yet" only when there is truly no history: + project-scoped and archived threads leave Recents empty + but still count as existing chats. */} + {chatItemsLoaded && + allChatItems.length === 0 && + archivedChatItems.length === 0 && ( +

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

+ )} @@ -1535,9 +1573,6 @@ export function AppSidebar() { > {t("shell.navigation.api")} - - {t("common.new")} - } diff --git a/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx new file mode 100644 index 0000000000..331a06a4c6 --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx @@ -0,0 +1,483 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"use client"; + +import { + Sheet, + SheetContent, + SheetDescription, + SheetHeader, + SheetTitle, +} from "@/components/ui/sheet"; +import { + customProviderDisplayName, + parseExternalModelId, + useChatPreferencesStore, + useChatRuntimeStore, + useExternalProvidersStore, +} from "@/features/chat"; +import { cn } from "@/lib/utils"; +import { HelpCircleIcon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { useMessage, useMessageTiming } from "@assistant-ui/react"; +import type { FC, ReactNode } from "react"; + +type ResponseDetailsMetadata = { + modelId?: string; + modelLabel?: string; + responseModelId?: string; + providerId?: string; + providerName?: string; + providerType?: string; + startedAt?: number; + finishedAt?: number; + durationMs?: number; + sessionId?: string | null; + cancelId?: string; + toolCalls?: string[]; + tools?: Record; +}; + +type ContextUsageMetadata = { + promptTokens?: number; + completionTokens?: number; + totalTokens?: number; + cachedTokens?: number; + cacheWriteTokens?: number; + modelId?: string; +}; + +type MessageCustomMetadata = { + responseDetails?: ResponseDetailsMetadata; + contextUsage?: ContextUsageMetadata; + serverTimings?: Record; + reasoningDuration?: number; +}; + +function asNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) + ? value + : undefined; +} + +function formatNumber(value: number | undefined): string | null { + return value == null ? null : value.toLocaleString(); +} + +function formatMs(value: number | undefined): string | null { + if (value == null) return null; + if (value < 1000) return `${Math.round(value)}ms`; + return `${(value / 1000).toFixed(2)}s`; +} + +function formatRate(value: number | undefined): string | null { + if (value == null) return null; + return `${value.toFixed(1)} tok/s`; +} + +function formatDate(value: Date | number | string | undefined): string | null { + if (value == null) return null; + const date = value instanceof Date ? value : new Date(value); + if (Number.isNaN(date.getTime())) return null; + return new Intl.DateTimeFormat(undefined, { + dateStyle: "medium", + timeStyle: "medium", + }).format(date); +} + +const TOOL_CATEGORY_LABELS: Record = { + search: "Search", + fetch: "Fetch", + code: "Code", + images: "Images", + mcp: "MCP", + docs: "Docs", + artifacts: "Canvas", +}; + +const TOOL_CALL_LABELS: Record = { + web_search: "Search", + web_fetch: "Fetch", + code_execution: "Code", + python: "Python", + terminal: "Terminal", + image_generation: "Images", + search_knowledge_base: "Docs", + render_html: "Canvas", +}; + +function uniqueValues(values: string[]): string[] { + return Array.from(new Set(values)); +} + +function toolCategoryFromCall(toolName: string): string | null { + const normalized = toolName.toLowerCase(); + if (normalized === "web_search") return "search"; + if (normalized === "web_fetch") return "fetch"; + if ( + normalized === "code_execution" || + normalized === "python" || + normalized === "terminal" + ) { + return "code"; + } + if (normalized === "image_generation") return "images"; + if (normalized === "search_knowledge_base") return "docs"; + if (normalized === "render_html") return "artifacts"; + if (normalized.startsWith("mcp__")) return "mcp"; + return null; +} + +function formatToolCallName(toolName: string): string { + const normalized = toolName.toLowerCase(); + if (TOOL_CALL_LABELS[normalized]) return TOOL_CALL_LABELS[normalized]; + if (normalized.startsWith("mcp__")) return `MCP: ${toolName.slice(5)}`; + return toolName + .replace(/[_-]+/g, " ") + .replace(/\b\w/g, (letter) => letter.toUpperCase()); +} + +function toolCallsFromContent(content: unknown): string[] { + if (!Array.isArray(content)) return []; + return uniqueValues( + content + .map((part) => + part && typeof part === "object" && "type" in part + ? (part as { type?: unknown; toolName?: unknown }) + : null, + ) + .filter( + (part): part is { type: "tool-call"; toolName: string } => + part?.type === "tool-call" && + typeof part.toolName === "string" && + part.toolName.length > 0, + ) + .map((part) => part.toolName), + ); +} + +function enabledTools( + tools: Record | undefined, + toolCalls: string[], +): string | null { + if (!tools && toolCalls.length === 0) return null; + const activeKeys = new Set(); + for (const key of Object.keys(TOOL_CATEGORY_LABELS)) { + if (tools?.[key] === true) activeKeys.add(key); + } + for (const toolName of toolCalls) { + const key = toolCategoryFromCall(toolName); + if (key) activeKeys.add(key); + } + const active = Object.keys(TOOL_CATEGORY_LABELS) + .filter((key) => activeKeys.has(key)) + .map((key) => TOOL_CATEGORY_LABELS[key]); + return active.length > 0 ? active.join(", ") : "None"; +} + +function calledTools(toolCalls: string[]): string | null { + if (toolCalls.length === 0) return null; + return uniqueValues(toolCalls.map(formatToolCallName)).join(", "); +} + +function DetailSection({ + title, + children, +}: { + title: string; + children: ReactNode; +}) { + return ( +
+

{title}

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

+ {summaryLabel} +

+ {providerLabel ? ( +

+ {providerLabel} +

+ ) : null} +
+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
+ + + ); +}; diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx b/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx index d00c812325..db7628777a 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx @@ -42,10 +42,8 @@ export function ModelUpdateAction({ }: ModelUpdateActionProps) { const [open, setOpen] = useState(false); - // The update is a managed download (it surfaces in the global Downloads panel - // with progress + cancel). When this exact repo+variant finishes, refresh the - // caller so the "update available" cue clears once the new revision is on - // disk. A ref keeps the subscription stable across renders without resubscribing. + // Refresh the caller when this repo+variant's download finishes so the "update available" cue + // clears. A ref keeps the subscription stable across renders. const onUpdatedRef = useRef(onUpdated); onUpdatedRef.current = onUpdated; useEffect(() => { @@ -60,9 +58,8 @@ export function ModelUpdateAction({ }, [repoId, variant]); const handleConfirm = useCallback(() => { - // Start the background re-download and close the dialog immediately; the - // Downloads panel owns progress + cancel from here. Only a failure to START - // surfaces a toast — a failed download reports itself in the panel. + // Start the re-download and close the dialog; the Downloads panel owns progress + cancel. + // Only a failure to START toasts (a failed download shows in the panel). void Promise.resolve() .then(onConfirm) .catch((err) => { 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 20000c82ee..ff0351883d 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -393,6 +393,7 @@ function ModelRow({ vramEst, gpuGb, tooltipText, + hubUrl, optionProps, onArrowDownIntoChildren, capabilities, @@ -409,6 +410,10 @@ function ModelRow({ vramEst?: number; gpuGb?: number; tooltipText?: ReactNode; + /** Hugging Face address (e.g. "huggingface.co/owner/name") for online/Hub + * rows; surfaced on hover so their repo id / URL is discoverable the same + * way local rows show an on-disk path. Omit to show no address line. */ + hubUrl?: string; optionProps?: ModelRowOptionProps; onArrowDownIntoChildren?: () => boolean; /** Capability override (HF rows have tags); falls back to name detection. */ @@ -546,30 +551,41 @@ function ModelRow({ ); - if (vramTooltipText) { - return ( - - {content} - - {label} - {vramTooltipText} - - - ); - } + // Optional Hugging Face address line for online/Hub rows, rendered under + // whichever tooltip shows so the repo id / URL is always visible on hover. + const hubUrlLine = hubUrl ? ( + + {hubUrl} + + ) : null; - if (tooltipText) { + const tooltipBody = vramTooltipText ? ( + <> + {label} + {vramTooltipText} + {hubUrlLine} + + ) : tooltipText ? ( + <> + {tooltipText} + {hubUrlLine} + + ) : hubUrl ? ( + <> + {label} + {hubUrlLine} + + ) : null; + + if (tooltipBody) { return ( - + {content} - {tooltipText} + {tooltipBody} ); @@ -1193,6 +1209,13 @@ function localPathTooltip(name: string, path: string): ReactNode { ); } +/** Hugging Face address for an online/Hub row, or undefined when the repo id is + * missing so the row shows no (empty) address line on hover. */ +function hubRepoUrl(id: string | null | undefined): string | undefined { + const trimmed = id?.trim(); + return trimmed ? `huggingface.co/${trimmed}` : undefined; +} + /** Whether a local model is an MLX build (name hint). MLX runs on Mac only, so * callers gate visibility on the host being a Mac. */ function localModelIsMlx(m: LocalModelInfo): boolean { @@ -1247,11 +1270,8 @@ export function HubModelPicker({ onEject?: () => void; }) { const gpu = useGpuInfo(); - // The currently-loaded/running model id. We read params.checkpoint from the - // runtime store (backend-mirrored from /api/inference/status.active_model, see - // chat-runtime-store) rather than the dropdown `isSelected` highlight (which is - // just `value === repo_id` and can reflect a staged, not-yet-loaded pick). Used - // to disable the cached-row update action for the model that's live in memory. + // Live model id from the runtime store (backend-mirrored active_model), not the dropdown + // highlight which can be a staged pick. Disables the update action for it. const loadedModelId = useChatRuntimeStore((s) => s.params.checkpoint); // Last-loaded timestamps power the "Recent" sort (vs "Downloaded" = file date). const loadTimes = useModelLoadTimes(value); @@ -1589,11 +1609,8 @@ export function HubModelPicker({ refreshLocalModelsList(); }, [hfToken, refreshLocalModelsList]); - // Updates run as MANAGED downloads (they show in the global Downloads panel - // with manifest-based progress + a working Cancel), instead of a blocking - // call. The worker re-resolves `main` and pulls only changed blobs, so the - // cached copy stays usable until the new revision lands. The row's - // ModelUpdateAction refreshes the list when this repo+variant completes. + // Updates run as managed downloads (Downloads panel: progress + Cancel), not a blocking + // call. The worker pulls only changed blobs, so the cached copy stays usable until done. const startManagedUpdate = useCallback((repoId: string, variant: string, expectedBytes: number) => { return downloadManager .requestStart({ @@ -2468,6 +2485,7 @@ export function HubModelPicker({
-
+
-
+ + + +
{isOpen && !isReasoningStreaming && ( )} diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 7fe98b701f..d987092c48 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -12,6 +12,10 @@ import { import { downloadImagePart } from "@/components/assistant-ui/image"; import { MarkdownText } from "@/components/assistant-ui/markdown-text"; import { MessageHtmlArtifacts } from "@/components/assistant-ui/message-html-artifacts"; +import { + MessageResponseDetailsSheet, + MessageResponseModelBadge, +} from "@/components/assistant-ui/message-response-details-sheet"; import { MessageTiming } from "@/components/assistant-ui/message-timing"; import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning"; import { RagSourcesGroup } from "@/components/assistant-ui/rag-sources"; @@ -123,6 +127,7 @@ import { FileDatabaseIcon, Folder01Icon, FolderAddIcon, + HelpCircleIcon, Image03Icon, McpServerIcon, PencilRulerIcon, @@ -3564,6 +3569,9 @@ const AssistantMessage: FC = () => { const aui = useAui(); const messageId = useAuiState(({ message }) => message.id); const messageContent = useAuiState(({ message }) => message.content); + const hasReasoningParts = useAuiState(({ message }) => + message.parts.some((part) => part.type === "reasoning"), + ); const incognito = useChatRuntimeStore((s) => s.incognito); // Use global store for editing state to ensure a single source of truth @@ -3620,7 +3628,7 @@ const AssistantMessage: FC = () => { return (
@@ -3649,6 +3657,11 @@ const AssistantMessage: FC = () => {
) : ( <> + {!hasReasoningParts ? ( +
+ +
+ ) : null} @@ -3893,58 +3906,76 @@ const EditAssistantMessageButton: FC = () => { const AssistantActionBar: FC = () => { const { forkMessage, forkDisabled } = useForkMessageAction(); + const [detailsOpen, setDetailsOpen] = useState(false); return ( - - - - - - - - - - - - - - + <> + + + + + + - - e.preventDefault()} - className="aui-action-bar-more-content z-50 min-w-32 overflow-hidden rounded-[21px] bg-popover px-[9px] py-2 text-popover-foreground shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:shadow-none" - > - void forkMessage()} - className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground data-[disabled]:pointer-events-none data-[disabled]:opacity-50" + + + + + + + + + + e.preventDefault()} + className="aui-action-bar-more-content z-50 min-w-32 overflow-hidden rounded-[21px] bg-popover px-[9px] py-2 text-popover-foreground shadow-[0_2px_8px_-2px_rgba(0,0,0,0.16)] dark:shadow-none" > - - Fork in new chat - - - + void forkMessage()} + className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground data-[disabled]:pointer-events-none data-[disabled]:opacity-50" + > + + Fork in new chat + + + + + Export as markdown + + + setDetailsOpen(true)} + className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground" + > - Export as Markdown + See response details - - - - - + + + + + + ); }; diff --git a/studio/frontend/src/components/floating-monitor.tsx b/studio/frontend/src/components/floating-monitor.tsx index f02da6612e..0a51875de9 100644 --- a/studio/frontend/src/components/floating-monitor.tsx +++ b/studio/frontend/src/components/floating-monitor.tsx @@ -3,33 +3,43 @@ import { Button } from "@/components/ui/button"; import { Progress } from "@/components/ui/progress"; -import { useMonitorOverlayStore } from "@/features/settings/stores/monitor-overlay-store"; +import { useMonitorOverlayStore } from "@/features/settings"; 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"; +import { AnimatePresence, motion, useDragControls } from "motion/react"; +import { type PointerEvent, useMemo, useState } 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"; + 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"; + 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 { +function formatGiB(value: number): string { + // RAM/VRAM come from the backend in binary units (bytes / 1024**3), matching + // nvidia-smi and PyTorch, so label the readout GiB rather than GB. const digits = value >= 10 ? 1 : 2; - return `${value.toFixed(digits)} GB`; + return `${value.toFixed(digits)} GiB`; } export function FloatingMonitor() { @@ -37,9 +47,18 @@ export function FloatingMonitor() { const { isOpen, setIsOpen } = useMonitorOverlayStore(); const systemInfo = useSystemInfo({ enabled: isOpen, pollMs: 5000 }); - const constraintsRef = useRef(null); + const [constraintsElement, setConstraintsElement] = + useState(null); + const constraintsRef = useMemo( + () => ({ current: constraintsElement }), + [constraintsElement], + ); + const dragControls = useDragControls(); - if (!isOpen) return null; + function startDrag(event: PointerEvent) { + event.preventDefault(); + dragControls.start(event); + } const ramTotal = systemInfo.memory?.total_gb ?? 0; const ramAvailable = systemInfo.memory?.available_gb ?? 0; @@ -62,99 +81,109 @@ export function FloatingMonitor() { const hasGpu = (systemInfo.gpu?.available ?? false) && devices.length > 0; return ( -
- -
-
- - - {t("settings.resources.liveMonitor.title")} - -
-
-
- -
- - -
-
- - + {isOpen && ( +
-
-
- {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"})`} + +
+
+ + + {t("settings.resources.liveMonitor.title")} - +
+
- {Math.round(vramPercent)}% - + +
+ +
-
- {formatGb(vramUsed)} / {formatGb(vramTotal)} -
-
- )} - - -
+ + +
+
+ {t("settings.resources.liveMonitor.ram")} + + {Math.round(ramPercent)}% + +
+
+ {formatGiB(ramUsed)} / {formatGiB(ramTotal)} +
+ +
+ + {hasGpu && ( +
+
+ + {t("settings.resources.liveMonitor.vram")}{" "} + {devices.length > 1 + ? `(${devices.length} GPUs)` + : `(${devices[0].name ?? "GPU"})`} + + + {Math.round(vramPercent)}% + +
+
+ {formatGiB(vramUsed)} / {formatGiB(vramTotal)} +
+ +
+ )} + + +
+ )} + ); } diff --git a/studio/frontend/src/config/env.ts b/studio/frontend/src/config/env.ts index def1b9bad9..63cdd03141 100644 --- a/studio/frontend/src/config/env.ts +++ b/studio/frontend/src/config/env.ts @@ -54,6 +54,17 @@ export const usePlatformStore = create()((_, get) => ({ isChatOnly: () => get().chatOnly, })); +// Once an authoritative (server-reported) platform has been fetched, a +// non-forced response must not overwrite it. The post-render fetchDeviceType() +// in main.tsx runs before auth is ready and can resolve after the authed +// root-route/provider fetches; such a late write would reset deviceType, +// cloudflareUrl/serverUrl/secure, and fetched, whether it is a browser fallback +// (unauthenticated) or an earlier authenticated request that landed after a +// later forced refresh. Forced refreshes are explicit re-reads, so they still write. +function shouldKeepAuthoritativePlatform(force?: boolean): boolean { + return !force && usePlatformStore.getState().fetched; +} + // `force` re-reads /api/health even if cached, to pick up a late-arriving tunnel URL. export async function fetchDeviceType(options?: { force?: boolean; @@ -81,6 +92,15 @@ export async function fetchDeviceType(options?: { server_url?: string | null; secure?: boolean; }; + // Once the store holds an authoritative (server-reported) platform, a + // non-forced response must not overwrite it. It may be an unauthenticated + // fallback, or an earlier authenticated request that resolved after a + // later forced refresh already picked up device_type and the tunnel + // fields; writing either would reset device type or null the tunnel + // fields. Forced refreshes are explicit re-reads, so they still write. + if (shouldKeepAuthoritativePlatform(options?.force)) { + return usePlatformStore.getState().deviceType; + } const deviceType = data.device_type ?? detectLocalPlatform(); const chatOnly = data.chat_only ?? false; const chatOnlyReason = data.chat_only_reason ?? null; @@ -101,7 +121,11 @@ export async function fetchDeviceType(options?: { } catch { // Backend not ready: use client-side detection so chat-only guard works // on initial load (important for macOS). Keep fetched=false so a later - // call retries against the backend. + // call retries against the backend. But a late non-forced failure must not + // wipe an authoritative platform that already resolved. + if (shouldKeepAuthoritativePlatform(options?.force)) { + return usePlatformStore.getState().deviceType; + } const deviceType = detectLocalPlatform(); const chatOnly = deviceType === "mac"; usePlatformStore.setState({ deviceType, chatOnly, fetched: false }); diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 929c0385b0..ca062bf93c 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -2,6 +2,10 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { getAuthToken } from "@/features/auth"; +import { + loadRememberedLoadSettings, + rememberedLoadSettingsKey, +} from "@/components/assistant-ui/model-selector/remembered-load-settings"; import { projectHasSources } from "@/features/rag/api/rag-api"; import { apiUrl } from "@/lib/api-base"; import { parseParamCountB } from "@/lib/model-size"; @@ -63,11 +67,17 @@ import { listStoredChatThreads, updateStoredChatThread, } from "../utils/chat-history-storage"; +import { + readLastLocalModelLoad, + recordLastLocalModelLoad, + type LastLocalModelKind, +} from "../utils/last-local-model-load"; import { getImageInputUnavailableReason } from "../utils/image-input-support"; import { hasClosedThinkTag, parseAssistantContent, } from "../utils/parse-assistant-content"; +import { resolveLoadMaxSeqLength } from "../presets/preset-policy"; import { generateAudio, listCachedGguf, @@ -140,6 +150,32 @@ interface ServerTimings { diffusion_steps_per_second?: number; } +interface ResponseDetailsMetadata { + modelId: string; + modelLabel: string; + responseModelId: string; + providerId?: string; + providerName: string; + providerType: string; + startedAt: number; + finishedAt: number; + durationMs: number; + sessionId?: string; + cancelId: string; + toolCalls: string[]; + tools: { + search: boolean; + fetch: boolean; + code: boolean; + images: boolean; + mcp: boolean; + docs: boolean; + artifacts: boolean; + confirmToolCalls: boolean; + bypassPermissions: boolean; + }; +} + type RunMessages = Parameters[0]["messages"]; type RunMessage = RunMessages[number]; @@ -1283,6 +1319,30 @@ const BIG_ENDIAN_GGUF_FILENAME_RE = /(^|[-_])be(?:[._-]|$)/gi; const GGUF_KNOWN_QUANT_RE = /(UD-)?(MXFP[0-9]+(?:_[A-Z0-9]+)*|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?|TQ[0-9]+_[0-9]+|Q[0-9]+_K_[A-Z]+|Q[0-9]+_[0-9]+|Q[0-9]+_K|BF16|F16|F32)/i; +type AutoLoadCandidate = { + id: string; + kind: LastLocalModelKind; + ggufVariant: string | null; + maxSeqLength: number; + successLabel: string; +}; + +function autoLoadCandidateKey( + kind: LastLocalModelKind, + id: string, + ggufVariant?: string | null, +): string { + return `${kind}:${id.toLowerCase()}:${(ggufVariant ?? "").toLowerCase()}`; +} + +function findCachedRepo( + repos: T[], + id: string, +): T | undefined { + const normalized = id.toLowerCase(); + return repos.find((repo) => repo.repo_id.toLowerCase() === normalized); +} + function hasBigEndianGgufMarker(filename: string, quant?: string | null): boolean { const normalized = filename.replace(/\\/g, "/").toLowerCase(); const separatorIndex = normalized.lastIndexOf("/"); @@ -1331,14 +1391,18 @@ async function autoLoadSmallestModel(): Promise<{ const hfToken = store.hfToken || null; const trustRemoteCode = store.params.trustRemoteCode ?? false; const specSettings = resolveSpeculativeSettingsForLoad(); + const lastLoaded = readLastLocalModelLoad(); const toastId = toast("Loading a model…", { - description: "Auto-selecting the smallest downloaded model.", + description: lastLoaded + ? "Loading last used model." + : "Auto-selecting the smallest downloaded model.", duration: 5000, closeButton: true, }); let blockedByTrustRemoteCode = false; let hadNonTrustFailure = false; let loadAttempts = 0; + const skippedAutoLoadCandidates = new Set(); async function canAutoLoad(payload: { model_path: string; @@ -1363,12 +1427,224 @@ async function autoLoadSmallestModel(): Promise<{ } return true; } + + async function loadAutoLoadCandidate( + candidate: AutoLoadCandidate, + ): Promise { + if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) { + return false; + } + const currentStore = useChatRuntimeStore.getState(); + const remembered = loadRememberedLoadSettings( + rememberedLoadSettingsKey({ + id: candidate.id, + ggufVariant: candidate.ggufVariant, + }), + ); + const effectiveMaxSeqLength = resolveLoadMaxSeqLength({ + modelId: candidate.id, + ggufVariant: candidate.ggufVariant, + isGguf: candidate.kind === "gguf", + customContextLength: remembered?.contextLength ?? null, + ggufContextLength: null, + currentCheckpoint: currentStore.params.checkpoint, + activeGgufVariant: currentStore.activeGgufVariant, + maxSeqLength: candidate.maxSeqLength, + presetSource: currentStore.activePresetSource, + }); + const effectiveSpeculativeType = + remembered?.speculativeType ?? specSettings.speculativeType; + const effectiveSpecDraftNMax = + remembered?.specDraftNMax ?? specSettings.specDraftNMax; + if ( + !(await canAutoLoad({ + model_path: candidate.id, + max_seq_length: effectiveMaxSeqLength, + is_lora: false, + gguf_variant: candidate.ggufVariant, + })) + ) { + skippedAutoLoadCandidates.add( + autoLoadCandidateKey(candidate.kind, candidate.id, candidate.ggufVariant), + ); + return false; + } + loadAttempts += 1; + const loadResp = await loadModel({ + model_path: candidate.id, + hf_token: hfToken, + max_seq_length: effectiveMaxSeqLength, + load_in_4bit: true, + is_lora: false, + gguf_variant: candidate.ggufVariant, + trust_remote_code: trustRemoteCode, + cache_type_kv: remembered?.kvCacheDtype ?? null, + speculative_type: effectiveSpeculativeType, + spec_draft_n_max: effectiveSpecDraftNMax, + tensor_parallel: remembered?.tensorParallel ?? false, + }); + saveSpeculativeType(effectiveSpeculativeType); + useChatRuntimeStore + .getState() + .setCheckpoint(candidate.id, candidate.ggufVariant ?? undefined); + const store = useChatRuntimeStore.getState(); + store.setModelRequiresTrustRemoteCode( + loadResp.requires_trust_remote_code ?? false, + ); + store.setParams({ + ...store.params, + maxTokens: + candidate.kind === "gguf" + ? loadResp.context_length ?? 131072 + : effectiveMaxSeqLength, + }); + const autoModel: ChatModelSummary = { + id: candidate.id, + name: loadResp.display_name ?? candidate.id, + isVision: loadResp.is_vision ?? false, + isLora: loadResp.is_lora ?? false, + isGguf: loadResp.is_gguf ?? candidate.kind === "gguf", + isAudio: loadResp.is_audio ?? false, + audioType: loadResp.audio_type ?? null, + hasAudioInput: loadResp.has_audio_input ?? false, + }; + if (!store.models.some((m) => m.id === candidate.id)) { + store.setModels([...store.models, autoModel]); + } + if (candidate.kind === "gguf") { + useChatRuntimeStore.setState({ + ggufContextLength: loadResp.context_length ?? 131072, + ggufMaxContextLength: + loadResp.max_context_length ?? loadResp.context_length ?? 131072, + ggufNativeContextLength: loadResp.native_context_length ?? null, + supportsReasoning: loadResp.supports_reasoning ?? false, + reasoningAlwaysOn: loadResp.reasoning_always_on ?? false, + reasoningEnabled: loadResp.supports_reasoning ?? false, + ...reasoningCapsFromLoad(loadResp), + supportsPreserveThinking: loadResp.supports_preserve_thinking ?? false, + supportsTools: loadResp.supports_tools ?? false, + ...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false), + kvCacheDtype: loadResp.cache_type_kv ?? null, + loadedKvCacheDtype: loadResp.cache_type_kv ?? null, + tensorParallel: loadResp.tensor_parallel ?? false, + loadedTensorParallel: loadResp.tensor_parallel ?? false, + defaultChatTemplate: loadResp.chat_template ?? null, + chatTemplateOverride: null, + loadedChatTemplateOverride: null, + loadedIsMultimodal: isMultimodalResponse(loadResp), + loadedIsDiffusion: loadResp.is_diffusion ?? false, + ...resolveLoadedSpeculativeSettings(loadResp), + }); + } else { + useChatRuntimeStore.setState({ + supportsReasoning: loadResp.supports_reasoning ?? false, + reasoningAlwaysOn: loadResp.reasoning_always_on ?? false, + reasoningEnabled: loadResp.supports_reasoning ?? false, + ...reasoningCapsFromLoad(loadResp), + supportsPreserveThinking: loadResp.supports_preserve_thinking ?? false, + supportsTools: loadResp.supports_tools ?? false, + ...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false), + kvCacheDtype: loadResp.cache_type_kv ?? null, + loadedKvCacheDtype: loadResp.cache_type_kv ?? null, + tensorParallel: loadResp.tensor_parallel ?? false, + loadedTensorParallel: loadResp.tensor_parallel ?? false, + defaultChatTemplate: loadResp.chat_template ?? null, + chatTemplateOverride: null, + loadedChatTemplateOverride: null, + ...resolveLoadedSpeculativeSettings(loadResp), + loadedIsMultimodal: isMultimodalResponse(loadResp), + loadedIsDiffusion: loadResp.is_diffusion ?? false, + }); + } + if (!(loadResp.is_lora ?? false)) { + recordLastLocalModelLoad({ + id: candidate.id, + kind: candidate.kind, + ggufVariant: candidate.ggufVariant, + }); + } + toast.success(candidate.successLabel, { id: toastId }); + return true; + } try { const [ggufRepos, modelRepos] = await Promise.all([ listCachedGguf().catch(() => []), listCachedModels().catch(() => []), ]); + if (lastLoaded) { + if (lastLoaded.kind === "gguf") { + const repo = findCachedRepo(ggufRepos, lastLoaded.id); + if (repo && lastLoaded.ggufVariant) { + try { + const variants = await listGgufVariants(repo.repo_id); + const variant = variants.variants.find( + (entry) => + entry.downloaded && + entry.quant?.toLowerCase() === + lastLoaded.ggufVariant?.toLowerCase() && + isAutoLoadableGgufVariant(entry), + ); + if (variant) { + toast("Loading last used model…", { + id: toastId, + description: `${repo.repo_id} (${variant.quant})`, + duration: 5000, + }); + if ( + await loadAutoLoadCandidate({ + id: repo.repo_id, + kind: "gguf", + ggufVariant: variant.quant, + maxSeqLength: 0, + successLabel: `Loaded ${repo.repo_id} (${variant.quant})`, + }) + ) { + return { loaded: true, blockedByTrustRemoteCode: false }; + } + } + } catch { + hadNonTrustFailure = true; + skippedAutoLoadCandidates.add( + autoLoadCandidateKey("gguf", repo.repo_id, lastLoaded.ggufVariant), + ); + } + } + } else { + const repo = findCachedRepo(modelRepos, lastLoaded.id); + if (repo) { + try { + toast("Loading last used model…", { + id: toastId, + description: repo.repo_id, + duration: 5000, + }); + if ( + await loadAutoLoadCandidate({ + id: repo.repo_id, + kind: "model", + ggufVariant: null, + maxSeqLength: store.params.maxSeqLength, + successLabel: `Loaded ${repo.repo_id}`, + }) + ) { + return { loaded: true, blockedByTrustRemoteCode: false }; + } + } catch { + hadNonTrustFailure = true; + skippedAutoLoadCandidates.add( + autoLoadCandidateKey("model", repo.repo_id), + ); + } + } + } + toast("Loading a model…", { + id: toastId, + description: "Auto-selecting the smallest downloaded model.", + duration: 5000, + }); + } + // GGUF first: smallest-total-size repo, then its smallest variant. if (ggufRepos.length > 0) { const sorted = [...ggufRepos].sort((a, b) => a.size_bytes - b.size_bytes); @@ -1382,82 +1658,23 @@ async function autoLoadSmallestModel(): Promise<{ if (downloaded.length > 0) { const variant = downloaded[0]; if ( - !(await canAutoLoad({ - model_path: repo.repo_id, - max_seq_length: 0, - is_lora: false, - gguf_variant: variant.quant, - })) + skippedAutoLoadCandidates.has( + autoLoadCandidateKey("gguf", repo.repo_id, variant.quant), + ) ) { continue; } - loadAttempts += 1; - const loadResp = await loadModel({ - model_path: repo.repo_id, - hf_token: hfToken, - max_seq_length: 0, - load_in_4bit: true, - is_lora: false, - gguf_variant: variant.quant, - trust_remote_code: trustRemoteCode, - speculative_type: specSettings.speculativeType, - spec_draft_n_max: specSettings.specDraftNMax, - }); - saveSpeculativeType(specSettings.speculativeType); - useChatRuntimeStore - .getState() - .setCheckpoint(repo.repo_id, variant.quant); - const store = useChatRuntimeStore.getState(); - store.setModelRequiresTrustRemoteCode( - loadResp.requires_trust_remote_code ?? false, - ); - store.setParams({ - ...store.params, - maxTokens: loadResp.context_length ?? 131072, - }); - // Add to store so the selector shows the name. - const autoModel: ChatModelSummary = { - id: repo.repo_id, - name: loadResp.display_name ?? repo.repo_id, - isVision: loadResp.is_vision ?? false, - isLora: loadResp.is_lora ?? false, - isGguf: loadResp.is_gguf ?? false, - isAudio: loadResp.is_audio ?? false, - audioType: loadResp.audio_type ?? null, - hasAudioInput: loadResp.has_audio_input ?? false, - }; - const existingModels = store.models; - if (!existingModels.some((m) => m.id === repo.repo_id)) { - store.setModels([...existingModels, autoModel]); + if ( + await loadAutoLoadCandidate({ + id: repo.repo_id, + kind: "gguf", + ggufVariant: variant.quant, + maxSeqLength: 0, + successLabel: `Loaded ${repo.repo_id} (${variant.quant})`, + }) + ) { + return { loaded: true, blockedByTrustRemoteCode: false }; } - useChatRuntimeStore.setState({ - ggufContextLength: loadResp.context_length ?? 131072, - ggufMaxContextLength: - loadResp.max_context_length ?? - loadResp.context_length ?? - 131072, - supportsReasoning: loadResp.supports_reasoning ?? false, - reasoningAlwaysOn: loadResp.reasoning_always_on ?? false, - reasoningEnabled: loadResp.supports_reasoning ?? false, - ...reasoningCapsFromLoad(loadResp), - supportsPreserveThinking: - loadResp.supports_preserve_thinking ?? false, - supportsTools: loadResp.supports_tools ?? false, - ...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false), - kvCacheDtype: loadResp.cache_type_kv ?? null, - loadedKvCacheDtype: loadResp.cache_type_kv ?? null, - tensorParallel: loadResp.tensor_parallel ?? false, - loadedTensorParallel: loadResp.tensor_parallel ?? false, - defaultChatTemplate: loadResp.chat_template ?? null, - chatTemplateOverride: null, - loadedChatTemplateOverride: null, - loadedIsMultimodal: isMultimodalResponse(loadResp), - ...resolveLoadedSpeculativeSettings(loadResp), - }); - toast.success(`Loaded ${repo.repo_id} (${variant.quant})`, { - id: toastId, - }); - return { loaded: true, blockedByTrustRemoteCode: false }; } } catch { hadNonTrustFailure = true; @@ -1475,64 +1692,23 @@ async function autoLoadSmallestModel(): Promise<{ if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) break; try { if ( - !(await canAutoLoad({ - model_path: repo.repo_id, - max_seq_length: 4096, - is_lora: false, - gguf_variant: null, - })) + skippedAutoLoadCandidates.has( + autoLoadCandidateKey("model", repo.repo_id), + ) ) { continue; } - loadAttempts += 1; - const sfLoadResp = await loadModel({ - model_path: repo.repo_id, - hf_token: hfToken, - max_seq_length: 4096, - load_in_4bit: true, - is_lora: false, - gguf_variant: null, - trust_remote_code: trustRemoteCode, - speculative_type: specSettings.speculativeType, - spec_draft_n_max: specSettings.specDraftNMax, - }); - saveSpeculativeType(specSettings.speculativeType); - useChatRuntimeStore.getState().setCheckpoint(repo.repo_id); - const store = useChatRuntimeStore.getState(); - store.setModelRequiresTrustRemoteCode( - sfLoadResp.requires_trust_remote_code ?? false, - ); - store.setParams({ ...store.params, maxTokens: 4096 }); - useChatRuntimeStore.setState({ - supportsReasoning: sfLoadResp.supports_reasoning ?? false, - reasoningAlwaysOn: sfLoadResp.reasoning_always_on ?? false, - reasoningEnabled: sfLoadResp.supports_reasoning ?? false, - ...reasoningCapsFromLoad(sfLoadResp), - supportsPreserveThinking: - sfLoadResp.supports_preserve_thinking ?? false, - supportsTools: sfLoadResp.supports_tools ?? false, - // Parity with the GGUF branch above. - ...resolveToolsEnabledOnLoad(sfLoadResp.supports_tools ?? false), - defaultChatTemplate: sfLoadResp.chat_template ?? null, - chatTemplateOverride: null, - loadedChatTemplateOverride: null, - ...resolveLoadedSpeculativeSettings(sfLoadResp), - }); - const sfModel: ChatModelSummary = { - id: repo.repo_id, - name: sfLoadResp.display_name ?? repo.repo_id, - isVision: sfLoadResp.is_vision ?? false, - isLora: sfLoadResp.is_lora ?? false, - isGguf: sfLoadResp.is_gguf ?? false, - }; - if (!store.models.some((m) => m.id === repo.repo_id)) { - store.setModels([...store.models, sfModel]); + if ( + await loadAutoLoadCandidate({ + id: repo.repo_id, + kind: "model", + ggufVariant: null, + maxSeqLength: 4096, + successLabel: `Loaded ${repo.repo_id}`, + }) + ) { + return { loaded: true, blockedByTrustRemoteCode: false }; } - useChatRuntimeStore.setState({ - loadedIsMultimodal: isMultimodalResponse(sfLoadResp), - }); - toast.success(`Loaded ${repo.repo_id}`, { id: toastId }); - return { loaded: true, blockedByTrustRemoteCode: false }; } catch { hadNonTrustFailure = true; continue; @@ -1624,6 +1800,11 @@ async function autoLoadSmallestModel(): Promise<{ loadedIsMultimodal: isMultimodalResponse(loadResp), ...resolveLoadedSpeculativeSettings(loadResp), }); + recordLastLocalModelLoad({ + id: "unsloth/Qwen3.5-4B-MTP-GGUF", + kind: "gguf", + ggufVariant: "UD-Q4_K_XL", + }); toast.success("Loaded Qwen3.5-4B-MTP (UD-Q4_K_XL)", { id: toastId }); return { loaded: true, blockedByTrustRemoteCode: false }; } catch { @@ -1769,6 +1950,9 @@ export function createOpenAIStreamAdapter( (provider) => provider.id === externalSelection.providerId, ) : null; + const selectedModelSummary = runtime.models.find( + (model) => model.id === params.checkpoint, + ); const externalApiKey = externalProvider ? getExternalProviderApiKey(externalProvider.id).trim() : ""; @@ -2151,6 +2335,7 @@ export function createOpenAIStreamAdapter( let waitingFirstChunk = true; let firstTokenSettled = false; const streamStartTime = Date.now(); + let responseModelId = externalSelection?.modelId ?? params.checkpoint; let firstTokenTime: number | undefined; let totalChunks = 0; let resolveFirstToken: (() => void) | null = null; @@ -2372,6 +2557,59 @@ export function createOpenAIStreamAdapter( const externalBackendProviderType = toExternalBackendProviderType( externalProvider?.providerType, ); + const buildResponseDetails = ( + finishedAt: number, + ): ResponseDetailsMetadata => ({ + modelId: params.checkpoint, + modelLabel: + (isExternalRequest || responseModelId !== params.checkpoint + ? responseModelId + : selectedModelSummary?.name || responseModelId) || + params.checkpoint || + "Unknown model", + responseModelId: + responseModelId || + externalSelection?.modelId || + params.checkpoint, + ...(externalProvider?.id ? { providerId: externalProvider.id } : {}), + providerName: + externalProvider?.name ?? + (isExternalRequest ? "External provider" : "Local model"), + providerType: externalProvider?.providerType ?? "local", + startedAt: streamStartTime, + finishedAt, + durationMs: finishedAt - streamStartTime, + ...(sandboxSessionId ? { sessionId: sandboxSessionId } : {}), + cancelId, + toolCalls: Array.from( + new Set( + toolCallParts + .map((part) => part.toolName) + .filter( + (toolName): toolName is string => + typeof toolName === "string" && toolName.length > 0, + ), + ), + ), + tools: { + search: + webSearchEnabledForThisTurn || + (!isExternalRequest && supportsTools && toolsEnabled), + fetch: webFetchEnabledForThisTurn, + code: + codeExecEnabledForThisTurn || + (!isExternalRequest && supportsTools && codeToolsEnabled), + images: imageGenerationEnabledForThisTurn, + mcp: !isExternalRequest && supportsTools && mcpEnabledForChat, + docs: + !isExternalRequest && + supportsTools && + (ragEnabled || projectRagEnabled), + artifacts: renderHtmlToolEnabledForThisTurn, + confirmToolCalls, + bypassPermissions, + }, + }); const externalCapabilities = getProviderCapabilities( externalProvider?.providerType, ); @@ -2739,6 +2977,7 @@ export function createOpenAIStreamAdapter( : {}), auto_heal_tool_calls: useChatRuntimeStore.getState().autoHealToolCalls, + nudge_tool_calls: useChatRuntimeStore.getState().nudgeToolCalls, max_tool_calls_per_message: useChatRuntimeStore.getState().maxToolCallsPerMessage, tool_call_timeout: (() => { @@ -2767,6 +3006,11 @@ export function createOpenAIStreamAdapter( const stream = streamChatCompletions(requestPayload, abortSignal); for await (const chunk of stream) { + const chunkModel = (chunk as { model?: unknown }).model; + if (typeof chunkModel === "string" && chunkModel.length > 0) { + responseModelId = chunkModel; + } + // Handle tool status events const toolStatusText = ( chunk as unknown as { _toolStatus?: string } @@ -3434,11 +3678,12 @@ export function createOpenAIStreamAdapter( }); } + const finishedAt = Date.now(); const finalTiming = buildTiming( streamStartTime, totalChunks, serverPromptEvalTime ?? firstTokenTime, - Date.now() - streamStartTime, + finishedAt - streamStartTime, finalTokenCount, toolCallParts.length, finalTokPerSec, @@ -3474,6 +3719,7 @@ export function createOpenAIStreamAdapter( modelId: params.checkpoint, } : undefined, + responseDetails: buildResponseDetails(finishedAt), timing: finalTiming, }, }, diff --git a/studio/frontend/src/features/chat/api/chat-settings-api.ts b/studio/frontend/src/features/chat/api/chat-settings-api.ts index 1ac569efc7..1e00357ea4 100644 --- a/studio/frontend/src/features/chat/api/chat-settings-api.ts +++ b/studio/frontend/src/features/chat/api/chat-settings-api.ts @@ -26,6 +26,7 @@ export interface PersistedChatSettings { collapseHtmlArtifacts?: boolean; allowArtifactNetworkAccess?: boolean; autoHealToolCalls?: boolean; + nudgeToolCalls?: boolean; maxToolCallsPerMessage?: number; toolCallTimeout?: number; } diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index cd7cfc77fc..b155eff780 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -35,7 +35,6 @@ import { useNativeModelDrop, useNativePathLeasesSupported, } from "@/features/native-intents"; -import { ProjectSourcesPanel } from "@/features/rag/components/project-sources-panel"; import { GuidedTour, useGuidedTourController } from "@/features/tour"; import { isTauri } from "@/lib/api-base"; import { toast } from "@/lib/toast"; @@ -51,7 +50,9 @@ import { Tooltip as TooltipPrimitive } from "radix-ui"; import { type CSSProperties, type ReactElement, + lazy, memo, + Suspense, useCallback, useEffect, useMemo, @@ -134,6 +135,13 @@ import { } from "./utils/chat-history-storage"; import { isAssistantLocalThreadId } from "./utils/thread-ids"; + +const ProjectSourcesPanel = lazy(() => + import("@/features/rag/components/project-sources-panel").then((module) => ({ + default: module.ProjectSourcesPanel, + })), +); + type LoraCandidate = { id: string; baseModel: string; @@ -1018,7 +1026,15 @@ function ProjectLanding({
{projectTab === "sources" ? ( - + + Loading sources… +
+ } + > + + ) : (
{items.map((item) => { @@ -2246,12 +2262,29 @@ export function ChatPage({ return [...fromLoras, ...localModels]; }, [lorasFromStore, localModels]); - useEffect(() => { - if (getTrainingCompareHandoff()) return; - void refresh(); + const inventoryRefreshStartedRef = useRef(false); + const refreshDeferredModelInventories = useCallback(() => { + inventoryRefreshStartedRef.current = true; + void refresh({ includeLoras: true }); refreshLocalModels(); }, [refresh, refreshLocalModels]); + useEffect(() => { + if (getTrainingCompareHandoff()) return; + void refresh({ includeLoras: false }); + const timeoutId = window.setTimeout(() => { + if (!inventoryRefreshStartedRef.current) { + refreshDeferredModelInventories(); + } + }, 1200); + return () => window.clearTimeout(timeoutId); + }, [refresh, refreshDeferredModelInventories]); + + useEffect(() => { + if (!active || !modelSelectorOpen) return; + refreshDeferredModelInventories(); + }, [active, modelSelectorOpen, refreshDeferredModelInventories]); + useEffect(() => { // ChatPage no longer remounts on navigation, so re-check the handoff whenever // we return to /chat (e.g. from the training progress "compare in chat" action). diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 74578b9438..ea6c409b40 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -1732,6 +1732,7 @@ export function ChatSettingsPanel({
+ @@ -2006,6 +2007,30 @@ function AutoHealToolCallsToggle() { ); } +function NudgeToolCallsToggle() { + const nudgeToolCalls = useChatRuntimeStore((s) => s.nudgeToolCalls); + const setNudgeToolCalls = useChatRuntimeStore((s) => s.setNudgeToolCalls); + + return ( +
+
+ + Nudge Tool Calls + + + When a tool call cannot be repaired, re-ask the model once so the + intended tool still runs. API requests stay opt-in. + +
+ +
+ ); +} + function ConfirmToolCallsToggle() { const confirmToolCalls = useChatRuntimeStore((s) => s.confirmToolCalls); const setConfirmToolCalls = useChatRuntimeStore((s) => s.setConfirmToolCalls); diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index e23d1b0b33..0a4342f208 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -44,6 +44,7 @@ import { mergeBackendRecommendedInference, resolveLoadMaxSeqLength, } from "../presets/preset-policy"; +import { recordLastLocalModelLoad } from "../utils/last-local-model-load"; import { isMultimodalResponse, } from "../types/api"; @@ -312,14 +313,18 @@ export function useChatModelRuntime() { [], ); - const refresh = useCallback(async (options?: { signal?: AbortSignal }) => { + const refresh = useCallback(async (options?: { + signal?: AbortSignal; + includeLoras?: boolean; + }) => { const signal = options?.signal; + const includeLoras = options?.includeLoras ?? true; setModelsError(null); try { const [listRes, statusRes, lorasRes] = await Promise.all([ listModels(), getInferenceStatus(), - listLoras(), + includeLoras ? listLoras() : Promise.resolve(null), ]); // Cancellation can land while the requests above are in flight. Bail @@ -327,7 +332,9 @@ export function useChatModelRuntime() { if (signal?.aborted) return; setModels(listRes.models.map(toChatModelSummary)); - setLoras(lorasRes.loras.map(toLoraSummary)); + if (lorasRes) { + setLoras(lorasRes.loras.map(toLoraSummary)); + } const selectedCheckpoint = useChatRuntimeStore.getState().params.checkpoint; const isExternalSelectionActive = isExternalModelId(selectedCheckpoint); @@ -812,6 +819,23 @@ export function useChatModelRuntime() { } } await refresh({ signal: abortCtrl.signal }); + if ( + !isLora && + !(loadResponse.is_lora ?? false) && + !nativePathToken && + !isLocalModelPath(modelId) && + !isExternalModelId(modelId) + ) { + if (loadResponse.is_gguf || isGguf || ggufVariant) { + recordLastLocalModelLoad({ + id: modelId, + kind: "gguf", + ggufVariant: ggufVariant ?? null, + }); + } else { + recordLastLocalModelLoad({ id: modelId, kind: "model" }); + } + } // A successful load owns the shared (pick-unscoped) settings fields, // so any surviving stage is stale: the just-loaded pick itself, or a // pick queued for a different model mid-load whose knobs this load diff --git a/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts b/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts index 55d7777e43..0a0df1139b 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts @@ -27,16 +27,21 @@ export interface SidebarItem { id: string; title: string; createdAt: number; + updatedAt: number; isFork?: boolean; projectId?: string | null; } +function lastActivityAt(thread: ThreadRecord): number { + return thread.updatedAt ?? thread.createdAt; +} + export function groupThreads( threads: ThreadRecord[], archived = false, ): SidebarItem[] { const items: SidebarItem[] = []; - const seenPairs = new Set(); + const pairItems = new Map(); for (const t of threads) { // Coerce archived to a boolean before comparing. Legacy threads (from the @@ -48,30 +53,35 @@ export function groupThreads( continue; } if (t.pairId) { - if (seenPairs.has(t.pairId)) { + const existing = pairItems.get(t.pairId); + if (existing) { + existing.updatedAt = Math.max(existing.updatedAt, lastActivityAt(t)); continue; } - seenPairs.add(t.pairId); - items.push({ + const item: SidebarItem = { type: "compare", id: t.pairId, title: t.title, createdAt: t.createdAt, + updatedAt: lastActivityAt(t), projectId: t.projectId ?? null, - }); + }; + pairItems.set(t.pairId, item); + items.push(item); } else if (!t.pairId) { items.push({ type: "single", id: t.id, title: t.title, createdAt: t.createdAt, + updatedAt: lastActivityAt(t), isFork: Boolean(t.forkedFromThreadId), projectId: t.projectId ?? null, }); } } - return items.sort((a, b) => b.createdAt - a.createdAt); + return items.sort((a, b) => b.updatedAt - a.updatedAt); } // Streaming fires CHAT_HISTORY_UPDATED_EVENT per chunk. Debounce so each quiet @@ -84,6 +94,7 @@ export function useChatSidebarItems(options?: { requireMessages?: boolean; }) { const [allThreads, setAllThreads] = useState([]); + const [loaded, setLoaded] = useState(false); const enabled = options?.enabled ?? true; const requireMessages = options?.requireMessages ?? true; @@ -111,6 +122,7 @@ export function useChatSidebarItems(options?: { // were in flight, or if the effect was torn down. if (cancelled || seq !== requestSeq) return; setAllThreads(threads); + setLoaded(true); } catch (error) { if (isExpectedBackgroundChatStorageError(error)) { return; @@ -144,7 +156,7 @@ export function useChatSidebarItems(options?: { const archivedItems = groupThreads(allThreads ?? [], true); const canCompare = useChatRuntimeStore((s) => Boolean(s.params.checkpoint)); - return { items, archivedItems, canCompare }; + return { items, archivedItems, canCompare, loaded }; } function cancelIfRunning(threadId: string): void { diff --git a/studio/frontend/src/features/chat/index.ts b/studio/frontend/src/features/chat/index.ts index 75749b0040..d070ed15de 100644 --- a/studio/frontend/src/features/chat/index.ts +++ b/studio/frontend/src/features/chat/index.ts @@ -26,11 +26,17 @@ export { type PlusMenuItemId, } from "./stores/plus-menu-prefs-store"; export { useChatModelRuntime } from "./hooks/use-chat-model-runtime"; -export { isExternalModelId } from "./external-providers"; +export { + customProviderDisplayName, + isExternalModelId, + parseExternalModelId, +} from "./external-providers"; +export { useExternalProvidersStore } from "./stores/external-providers-store"; export { ChatSearchDialog } from "./components/chat-search-dialog"; export { setTrainingCompareHandoff } from "./lib/training-compare-handoff"; export type { ProjectRecord } from "./types"; export { clearAllChats, countAllChats } from "./utils/clear-all-chats"; +export { listStoredChatThreads } from "./utils/chat-history-storage"; export { ArtifactCard } from "./artifacts/artifact-card"; export { useChatArtifactsStore, diff --git a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts index 9386650fee..60788c23bd 100644 --- a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts +++ b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts @@ -3,16 +3,19 @@ import { getInferenceStatus } from "../api/chat-api"; import { mergeBackendRecommendedInference } from "../presets/preset-policy"; +import { clampReasoningEffortToLevels } from "../provider-capabilities"; import { CHAT_REASONING_ENABLED_KEY, - loadOptionalBool, type ReasoningEffort, type ReasoningStyle, + loadOptionalBool, resolveToolsEnabledOnLoad, useChatRuntimeStore, } from "../stores/chat-runtime-store"; -import { isMultimodalResponse, type InferenceStatusResponse } from "../types/api"; -import { clampReasoningEffortToLevels } from "../provider-capabilities"; +import { + type InferenceStatusResponse, + isMultimodalResponse, +} from "../types/api"; import type { ChatModelSummary } from "../types/runtime"; type LocalReasoningEffort = Extract; @@ -31,7 +34,10 @@ export function normalizeSpeculativeType( return "ngram"; } if (s === "mtp+ngram") return "mtp+ngram"; - const parts = s.split(",").map((p) => p.trim()).filter(Boolean); + const parts = s + .split(",") + .map((p) => p.trim()) + .filter(Boolean); const hasMtp = parts.some((p) => p === "mtp" || p === "draft-mtp"); const hasNgram = parts.some( (p) => p === "ngram" || p === "ngram-mod" || p === "ngram-simple", @@ -197,6 +203,12 @@ export function applyActiveModelStatusToStore( ggufContextLength: currentGgufContextLength, ggufMaxContextLength, ggufNativeContextLength, + // A non-GGUF status must also drop a stale native-path token: without this the + // isGguf OR (activeGgufVariant || activeNativePathToken || ggufContextLength) + // stays true after switching from a native GGUF to a transformers model, so a + // Codex-only detection would auto-select for a model its preflight rejects. A real + // GGUF load reports is_gguf: true, so its token is preserved (the load path owns it). + ...(status.is_gguf ? {} : { activeNativePathToken: null }), modelRequiresTrustRemoteCode: status.requires_trust_remote_code ?? false, defaultChatTemplate: nextDefaultChatTemplate, loadedIsMultimodal: isMultimodalResponse(status), @@ -245,7 +257,7 @@ export function applyActiveModelStatusToStore( const mid = checkpointId.toLowerCase(); if (mid.includes("qwen3.5") || mid.includes("qwen3.6")) { const sizeMatch = mid.match(/(\d+\.?\d*)\s*b/); - if (sizeMatch && parseFloat(sizeMatch[1]) < 9) { + if (sizeMatch && Number.parseFloat(sizeMatch[1]) < 9) { reasoningDefault = false; } } @@ -281,8 +293,7 @@ export async function tryAdoptServerActiveModel(): Promise { } // Re-check after the await: keep a checkpoint the user picked meanwhile. - const previousCheckpoint = - useChatRuntimeStore.getState().params.checkpoint; + const previousCheckpoint = useChatRuntimeStore.getState().params.checkpoint; if (previousCheckpoint) { return true; } diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index 67980b94c6..b545695f9e 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -22,7 +22,6 @@ import { unstable_useRemoteThreadListRuntime as useRemoteThreadListRuntime, } from "@assistant-ui/react"; import { createAssistantStream } from "assistant-stream"; -import mammoth from "mammoth"; import { type ReactElement, type ReactNode, @@ -33,7 +32,6 @@ import { useMemo, useRef, } from "react"; -import { extractText, getDocumentProxy } from "unpdf"; import { toast } from "sonner"; import { StudioWebSpeechDictationAdapter } from "./adapters/studio-web-speech-dictation-adapter"; import { @@ -181,7 +179,10 @@ class PDFAttachmentAdapter implements AttachmentAdapter { } async send(attachment: PendingAttachment): Promise { - const buffer = new Uint8Array(await attachment.file.arrayBuffer()); + const [{ extractText, getDocumentProxy }, buffer] = await Promise.all([ + import("unpdf"), + attachment.file.arrayBuffer().then((bytes) => new Uint8Array(bytes)), + ]); const pdf = await getDocumentProxy(buffer); const { text } = await extractText(pdf, { mergePages: true }); return { @@ -298,7 +299,10 @@ class DocxAttachmentAdapter implements AttachmentAdapter { } async send(attachment: PendingAttachment): Promise { - const arrayBuffer = await attachment.file.arrayBuffer(); + const [{ default: mammoth }, arrayBuffer] = await Promise.all([ + import("mammoth"), + attachment.file.arrayBuffer(), + ]); const { value } = await mammoth.extractRawText({ arrayBuffer }); return { id: attachment.id, diff --git a/studio/frontend/src/features/chat/stores/chat-preferences-store.ts b/studio/frontend/src/features/chat/stores/chat-preferences-store.ts index f019b7dc3d..5b5011a78f 100644 --- a/studio/frontend/src/features/chat/stores/chat-preferences-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-preferences-store.ts @@ -7,11 +7,14 @@ import { persist } from "zustand/middleware"; // Client-side chat UI prefs kept in localStorage, not the chat DB. // confirmDeleteChats: when off, deleting a chat skips the confirm dialog. // showModelDisclaimer: when off, hide the "LLMs can make mistakes" footer note. +// showResponseModel: when on, assistant responses show the producing model. export interface ChatPreferencesState { confirmDeleteChats: boolean; setConfirmDeleteChats: (value: boolean) => void; showModelDisclaimer: boolean; setShowModelDisclaimer: (value: boolean) => void; + showResponseModel: boolean; + setShowResponseModel: (value: boolean) => void; } export const useChatPreferencesStore = create()( @@ -23,6 +26,9 @@ export const useChatPreferencesStore = create()( showModelDisclaimer: true, setShowModelDisclaimer: (showModelDisclaimer) => set({ showModelDisclaimer }), + showResponseModel: false, + setShowResponseModel: (showResponseModel) => + set({ showResponseModel }), }), { name: "unsloth_chat_preferences", @@ -32,6 +38,7 @@ export const useChatPreferencesStore = create()( ...current, confirmDeleteChats: saved?.confirmDeleteChats ?? true, showModelDisclaimer: saved?.showModelDisclaimer ?? true, + showResponseModel: saved?.showResponseModel ?? false, }; }, }, 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 7c9685d6d0..ad50c4ece8 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -646,6 +646,7 @@ type ChatRuntimeStore = { toolStatus: string | null; generatingStatus: string | null; autoHealToolCalls: boolean; + nudgeToolCalls: boolean; maxToolCallsPerMessage: number; toolCallTimeout: number; kvCacheDtype: string | null; @@ -780,6 +781,7 @@ type ChatRuntimeStore = { setGeneratingStatus: (status: string | null) => void; setActiveDiffusionCanvas: (canvas: DiffusionCanvasFrame | null) => void; setAutoHealToolCalls: (enabled: boolean) => void; + setNudgeToolCalls: (enabled: boolean) => void; setMaxToolCallsPerMessage: (value: number) => void; setToolCallTimeout: (value: number) => void; setKvCacheDtype: (dtype: string | null) => void; @@ -832,6 +834,7 @@ type ScalarSettingKey = | "collapseHtmlArtifacts" | "allowArtifactNetworkAccess" | "autoHealToolCalls" + | "nudgeToolCalls" | "maxToolCallsPerMessage" | "toolCallTimeout"; @@ -869,6 +872,7 @@ const SCALAR_SETTING_KEYS = [ "collapseHtmlArtifacts", "allowArtifactNetworkAccess", "autoHealToolCalls", + "nudgeToolCalls", "maxToolCallsPerMessage", "toolCallTimeout", ] as const satisfies readonly ScalarSettingKey[]; @@ -1103,6 +1107,7 @@ export const useChatRuntimeStore = create((set, get) => ({ generatingStatus: null, activeDiffusionCanvas: null, autoHealToolCalls: true, + nudgeToolCalls: true, maxToolCallsPerMessage: 25, toolCallTimeout: 5, kvCacheDtype: null, @@ -1544,6 +1549,15 @@ export const useChatRuntimeStore = create((set, get) => ({ ); return { autoHealToolCalls }; }), + setNudgeToolCalls: (nudgeToolCalls) => + set((state) => { + setScalarSettingVersion( + "nudgeToolCalls", + nudgeToolCalls, + state.nudgeToolCalls, + ); + return { nudgeToolCalls }; + }), setMaxToolCallsPerMessage: (maxToolCallsPerMessage) => set((state) => { setScalarSettingVersion( diff --git a/studio/frontend/src/features/chat/types.ts b/studio/frontend/src/features/chat/types.ts index 3fe69ccc26..111510d925 100644 --- a/studio/frontend/src/features/chat/types.ts +++ b/studio/frontend/src/features/chat/types.ts @@ -36,6 +36,7 @@ export interface ThreadRecord { projectId?: string | null; archived: boolean; createdAt: number; + updatedAt?: number; /** * OpenAI shell tool container id from a prior response. When set, the * next turn reuses it via `environment.type="container_reference"` so diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 94a7e1c56b..954e88e86b 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -357,6 +357,7 @@ export interface OpenAIChatCompletionsRequest { context_length?: number; }; auto_heal_tool_calls?: boolean; + nudge_tool_calls?: boolean; max_tool_calls_per_message?: number; tool_call_timeout?: number; session_id?: string; diff --git a/studio/frontend/src/features/chat/utils/chat-history-storage.ts b/studio/frontend/src/features/chat/utils/chat-history-storage.ts index 4294887df0..00df3657b1 100644 --- a/studio/frontend/src/features/chat/utils/chat-history-storage.ts +++ b/studio/frontend/src/features/chat/utils/chat-history-storage.ts @@ -590,7 +590,10 @@ export async function listStoredChatThreads( } return Array.from(byId.values()) .filter((thread) => matchesThreadListArgs(thread, args)) - .sort((a, b) => b.createdAt - a.createdAt); + .sort( + (a, b) => + (b.updatedAt ?? b.createdAt) - (a.updatedAt ?? a.createdAt), + ); } export async function listStoredChatThreadsWithMessages( diff --git a/studio/frontend/src/features/chat/utils/chat-settings-storage.ts b/studio/frontend/src/features/chat/utils/chat-settings-storage.ts index 4e457920d1..4899cb9c83 100644 --- a/studio/frontend/src/features/chat/utils/chat-settings-storage.ts +++ b/studio/frontend/src/features/chat/utils/chat-settings-storage.ts @@ -21,6 +21,7 @@ import type { ReasoningEffort } from "../stores/chat-runtime-store"; const AUTO_TITLE_KEY = "unsloth_chat_auto_title"; const AUTO_HEAL_TOOL_CALLS_KEY = "unsloth_auto_heal_tool_calls"; +const NUDGE_TOOL_CALLS_KEY = "unsloth_nudge_tool_calls"; const MAX_TOOL_CALLS_KEY = "unsloth_max_tool_calls_per_message"; const TOOL_CALL_TIMEOUT_KEY = "unsloth_tool_call_timeout"; const INFERENCE_PARAMS_KEY = "unsloth_chat_inference_params"; @@ -223,6 +224,7 @@ function sanitizeChatSettings(value: unknown): PersistedChatSettings { value.allowArtifactNetworkAccess, ); const autoHealToolCalls = sanitizeBool(value.autoHealToolCalls); + const nudgeToolCalls = sanitizeBool(value.nudgeToolCalls); const maxToolCallsPerMessage = sanitizeInt(value.maxToolCallsPerMessage, 1); const toolCallTimeout = sanitizeInt(value.toolCallTimeout, 1); @@ -245,6 +247,9 @@ function sanitizeChatSettings(value: unknown): PersistedChatSettings { if (autoHealToolCalls !== undefined) { settings.autoHealToolCalls = autoHealToolCalls; } + if (nudgeToolCalls !== undefined) { + settings.nudgeToolCalls = nudgeToolCalls; + } if (maxToolCallsPerMessage !== undefined) { settings.maxToolCallsPerMessage = maxToolCallsPerMessage; } @@ -305,6 +310,7 @@ export function isEmptyChatSettings(settings: PersistedChatSettings): boolean { settings.collapseHtmlArtifacts === undefined && settings.allowArtifactNetworkAccess === undefined && settings.autoHealToolCalls === undefined && + settings.nudgeToolCalls === undefined && settings.maxToolCallsPerMessage === undefined && settings.toolCallTimeout === undefined ); @@ -335,6 +341,7 @@ export function loadLegacyChatSettings(): PersistedChatSettings { const collapseHtmlArtifacts = loadBool(COLLAPSE_HTML_ARTIFACTS_KEY); const allowArtifactNetworkAccess = loadBool(ALLOW_ARTIFACT_NETWORK_ACCESS_KEY); const autoHealToolCalls = loadBool(AUTO_HEAL_TOOL_CALLS_KEY); + const nudgeToolCalls = loadBool(NUDGE_TOOL_CALLS_KEY); const maxToolCallsPerMessage = loadInt(MAX_TOOL_CALLS_KEY, 1); const toolCallTimeout = loadInt(TOOL_CALL_TIMEOUT_KEY, 1); const allCustomPresets = sanitizeCustomPresets([ @@ -361,6 +368,9 @@ export function loadLegacyChatSettings(): PersistedChatSettings { if (autoHealToolCalls !== undefined) { settings.autoHealToolCalls = autoHealToolCalls; } + if (nudgeToolCalls !== undefined) { + settings.nudgeToolCalls = nudgeToolCalls; + } if (maxToolCallsPerMessage !== undefined) { settings.maxToolCallsPerMessage = maxToolCallsPerMessage; } diff --git a/studio/frontend/src/features/chat/utils/last-local-model-load.ts b/studio/frontend/src/features/chat/utils/last-local-model-load.ts new file mode 100644 index 0000000000..099386fbc7 --- /dev/null +++ b/studio/frontend/src/features/chat/utils/last-local-model-load.ts @@ -0,0 +1,86 @@ +// 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 type LastLocalModelKind = "gguf" | "model"; + +export type LastLocalModelLoad = { + id: string; + kind: LastLocalModelKind; + ggufVariant: string | null; + loadedAt: number; +}; + +const STORAGE_KEY = "unsloth.last-local-model-load.v1"; + +function storage(): Storage | null { + try { + return typeof localStorage === "undefined" ? null : localStorage; + } catch { + return null; + } +} + +function isLastLocalModelKind(value: unknown): value is LastLocalModelKind { + return value === "gguf" || value === "model"; +} + +export function readLastLocalModelLoad(): LastLocalModelLoad | null { + try { + const raw = storage()?.getItem(STORAGE_KEY); + if (!raw) { + return null; + } + const parsed = JSON.parse(raw) as Partial; + if ( + typeof parsed.id !== "string" || + !parsed.id.trim() || + !isLastLocalModelKind(parsed.kind) || + typeof parsed.loadedAt !== "number" + ) { + return null; + } + if ( + parsed.kind === "gguf" && + (typeof parsed.ggufVariant !== "string" || !parsed.ggufVariant.trim()) + ) { + return null; + } + return { + id: parsed.id, + kind: parsed.kind, + ggufVariant: + typeof parsed.ggufVariant === "string" ? parsed.ggufVariant : null, + loadedAt: parsed.loadedAt, + }; + } catch { + return null; + } +} + +export function recordLastLocalModelLoad(input: { + id: string; + kind: LastLocalModelKind; + ggufVariant?: string | null; +}): void { + const id = input.id.trim(); + if (!id) { + return; + } + const ggufVariant = input.ggufVariant?.trim() || null; + if (input.kind === "gguf" && !ggufVariant) { + return; + } + try { + storage()?.setItem( + STORAGE_KEY, + JSON.stringify({ + id, + kind: input.kind, + ggufVariant: input.kind === "gguf" ? ggufVariant : null, + loadedAt: Date.now(), + } satisfies LastLocalModelLoad), + ); + } catch { + // Ignore disabled storage / quota errors; auto-load falls back to size order. + } +} diff --git a/studio/frontend/src/features/hub/hub-page.tsx b/studio/frontend/src/features/hub/hub-page.tsx index c3920e18f6..3dfcad5b4c 100644 --- a/studio/frontend/src/features/hub/hub-page.tsx +++ b/studio/frontend/src/features/hub/hub-page.tsx @@ -1085,11 +1085,11 @@ export function ModelsPage() { const { vramInfo, minMemory } = useHubModelVram(selectedModel, gpu); const gpuLabel = gpu.available - ? `${Math.round(gpu.memoryTotalGb)} GB` + ? `${Math.round(gpu.memoryTotalGb)} GiB` : "Unavailable"; const ramLabel = gpu.systemRamTotalGb > 0 - ? `${Math.round(gpu.systemRamTotalGb)} GB` + ? `${Math.round(gpu.systemRamTotalGb)} GiB` : "Unavailable"; const coreLabel = gpu.cpuCore > 0 && gpu.cpuThread > 0 diff --git a/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx b/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx index 6b6b11bf1f..eb50d398b6 100644 --- a/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx @@ -125,7 +125,7 @@ export function SummaryStep() { GPU
{hw.gpuName ?? "---"} - {hw.vramTotalGb != null ? `${hw.vramTotalGb} GB` : "---"} + {hw.vramTotalGb != null ? `${hw.vramTotalGb} GiB` : "---"}
diff --git a/studio/frontend/src/features/recipe-studio/api/index.ts b/studio/frontend/src/features/recipe-studio/api/index.ts index 4fe24b1f6d..f4e8167cb3 100644 --- a/studio/frontend/src/features/recipe-studio/api/index.ts +++ b/studio/frontend/src/features/recipe-studio/api/index.ts @@ -494,3 +494,13 @@ export async function removeUnstructuredFile( throw new Error("Failed to remove file"); } } + +export async function removeUnstructuredBlock(blockId: string): Promise { + const res = await authFetch( + `${DATA_DESIGNER_API_BASE}/seed/unstructured-block/${encodeURIComponent(blockId)}`, + { method: "DELETE" }, + ); + if (!res.ok && res.status !== 404) { + throw new Error("Failed to remove uploaded files"); + } +} diff --git a/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx b/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx index 0ed7eeed75..53f8566090 100644 --- a/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx +++ b/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx @@ -53,6 +53,11 @@ import { inspectSeedDataset, inspectSeedUpload, } from "../../api"; +import { useRecipeStudioStore } from "../../stores/recipe-studio"; +import { + makeUnstructuredUploadUid, + resolveUnstructuredUploadBlockId, +} from "../../utils/config-factories"; import { resolveImagePreview } from "../../utils/image-preview"; import type { GithubItemType, @@ -597,6 +602,41 @@ export function SeedDialog({ const mode = config.seed_source_type ?? "hf"; const previewEmpty = getPreviewEmptyStateCopy(mode); + const queueUploadCleanup = useRecipeStudioStore( + (state) => state.queueUploadCleanup, + ); + + // config.id collides across recipes (ids reset to n1 on import); use a + // stable per-block uid instead. Generate one synchronously so the first + // rendered drop zone cannot upload under a legacy node id. + const uploadUid = config.unstructured_upload_uid?.trim() ?? ""; + const unstructuredFileCount = config.unstructured_file_ids?.length ?? 0; + const generatedUploadUidRef = useRef(null); + if ( + mode === "unstructured" && + !uploadUid && + unstructuredFileCount === 0 && + generatedUploadUidRef.current === null + ) { + generatedUploadUidRef.current = makeUnstructuredUploadUid(); + } + const uploadBlockId = resolveUnstructuredUploadBlockId({ + configId: config.id, + uploadUid, + generatedUploadUid: generatedUploadUidRef.current, + unstructuredFileCount, + }); + + useEffect(() => { + if (mode !== "unstructured") return; + if (uploadUid) return; + if (unstructuredFileCount > 0) return; + const nextUid = + generatedUploadUidRef.current ?? makeUnstructuredUploadUid(); + generatedUploadUidRef.current = nextUid; + onUpdate({ unstructured_upload_uid: nextUid }); + }, [mode, uploadUid, unstructuredFileCount, onUpdate]); + const prevModeRef = useRef(mode); useEffect(() => { const prevMode = prevModeRef.current; @@ -720,6 +760,11 @@ export function SeedDialog({ subset: config.hf_subset?.trim() || undefined, preview_size: 10, }); + // Queue the block's upload directory for deletion after the next + // save; only uid-namespaced directories qualify (single owner). + if (uploadUid && unstructuredFileCount > 0) { + queueUploadCleanup(uploadUid); + } onUpdate({ hf_path: response.resolved_path, seed_columns: response.columns, @@ -730,6 +775,7 @@ export function SeedDialog({ hf_split: response.split ?? "", hf_subset: response.subset ?? "", local_file_name: "", + unstructured_upload_uid: "", unstructured_file_ids: [], unstructured_file_names: [], unstructured_file_sizes: [], @@ -754,6 +800,11 @@ export function SeedDialog({ content_base64: payload, preview_size: 10, }); + // Queue the block's upload directory for deletion after the next + // save; only uid-namespaced directories qualify (single owner). + if (uploadUid && unstructuredFileCount > 0) { + queueUploadCleanup(uploadUid); + } onUpdate({ hf_path: response.resolved_path, seed_columns: response.columns, @@ -765,6 +816,7 @@ export function SeedDialog({ hf_subset: "", hf_split: "", local_file_name: localFile.name, + unstructured_upload_uid: "", unstructured_file_ids: [], unstructured_file_names: [], unstructured_file_sizes: [], @@ -789,7 +841,7 @@ export function SeedDialog({ const { chunkSize, chunkOverlap } = resolveChunking(config); const response = await inspectSeedUpload({ - block_id: config.id, + block_id: uploadBlockId, file_ids: fileIds, file_names: fileNames, preview_size: 10, @@ -827,7 +879,18 @@ export function SeedDialog({ setIsInspecting(false); } }, - [config, getCurrentLoadKey, localFile, mode, onUpdate, unstructuredFiles], + [ + config, + getCurrentLoadKey, + localFile, + mode, + onUpdate, + queueUploadCleanup, + unstructuredFiles, + unstructuredFileCount, + uploadBlockId, + uploadUid, + ], ); useEffect(() => { @@ -997,7 +1060,7 @@ export function SeedDialog({ {mode === "unstructured" && ( (null); const filesRef = useRef(files); + const blockIdRef = useRef(blockId); + const mountedRef = useRef(true); const [isDragOver, setIsDragOver] = useState(false); useEffect(() => { filesRef.current = files; - }, [files]); + blockIdRef.current = blockId; + }, [files, blockId]); + useEffect(() => () => { + mountedRef.current = false; + }, []); const totalSize = files.reduce((sum, f) => sum + f.size, 0); @@ -134,15 +140,32 @@ export function UnstructuredDropZone({ if (entry.status === "uploading" && entry.abortController) { entry.abortController.abort(); } - if ( + const needsServerRemove = entry.id && entry.status === "ok" && - !deletedIdsRef.current.has(entry.id) - ) { - deletedIdsRef.current.add(entry.id); - void removeUnstructuredFile(blockId, entry.id).catch(() => {}); - } + !deletedIdsRef.current.has(entry.id); onFilesChange((prev) => prev.filter((_, i) => i !== index)); + if (!needsServerRemove) return; + deletedIdsRef.current.add(entry.id); + removeUnstructuredFile(blockId, entry.id).catch(() => { + // Skip if the drop zone unmounted or its block changed: the id no + // longer belongs here and restoring would leak it into another block. + if (!mountedRef.current || blockIdRef.current !== blockId) return; + // Still exists server-side (counts toward quota); restore it at its + // original position. + deletedIdsRef.current.delete(entry.id); + onFilesChange((prev) => { + const next = [...prev]; + next.splice(Math.min(index, next.length), 0, { + id: entry.id, + name: entry.name, + size: entry.size, + status: "ok", + error: "Remove failed — try again", + }); + return next; + }); + }); }, [blockId, onFilesChange], ); diff --git a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts index 0417a6dd22..2d91272469 100644 --- a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts +++ b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts @@ -4,11 +4,13 @@ import { useCallback, useEffect, useMemo, useState } from "react"; import { toastError, toastSuccess } from "@/shared/toast"; import { normalizeNonEmptyName } from "@/utils"; +import { removeUnstructuredBlock } from "../api"; import { buildSignature, copyTextToClipboard, formatSavedLabel, } from "../executions/execution-helpers"; +import { useRecipeStudioStore } from "../stores/recipe-studio"; import { importRecipePayload, type RecipeSnapshot } from "../utils/import"; import type { RecipePayloadResult } from "../utils/payload/types"; @@ -72,7 +74,10 @@ function stripApiKeys(value: unknown): unknown { !Array.isArray(output.env) ) { output.env = Object.fromEntries( - Object.keys(output.env as Record).map((envKey) => [envKey, ""]), + Object.keys(output.env as Record).map((envKey) => [ + envKey, + "", + ]), ); } return output; @@ -82,10 +87,7 @@ function inferHfRepoIdFromPath(pathValue: unknown): string { if (typeof pathValue !== "string") { return ""; } - const parts = pathValue - .trim() - .split("/") - .filter(Boolean); + const parts = pathValue.trim().split("/").filter(Boolean); if (parts.length >= 3 && parts[0] === "datasets") { return `${parts[1]}/${parts[2]}`; } @@ -126,8 +128,7 @@ function sanitizeSeedForShare(payload: unknown): unknown { typeof ui?.seed_source_type === "string" ? ui.seed_source_type : null; const sourceType = typeof source?.seed_type === "string" ? source.seed_type : null; - const shouldResetHfState = - sourceType === "hf" || uiSourceType === "hf"; + const shouldResetHfState = sourceType === "hf" || uiSourceType === "hf"; const shouldResetLocalState = sourceType === "local" || sourceType === "unstructured" || @@ -144,6 +145,7 @@ function sanitizeSeedForShare(payload: unknown): unknown { ui.seed_drop_columns = []; ui.seed_preview_rows = []; ui.local_file_name = ""; + ui.unstructured_upload_uid = ""; ui.unstructured_file_ids = []; ui.unstructured_file_names = []; ui.unstructured_file_sizes = []; @@ -165,6 +167,7 @@ function sanitizeSeedForShare(payload: unknown): unknown { ui.seed_drop_columns = []; ui.seed_preview_rows = []; ui.local_file_name = ""; + ui.unstructured_upload_uid = ""; ui.unstructured_file_ids = []; ui.unstructured_file_names = []; ui.unstructured_file_sizes = []; @@ -174,6 +177,43 @@ function sanitizeSeedForShare(payload: unknown): unknown { return root; } +// Delete queued upload directories once a save stops referencing them, so a +// reload before autosave can never leave the saved recipe pointing at +// already-deleted files. Skips any uid the just-saved payload still uses. +function drainQueuedUploadCleanups( + savedPayload: RecipePayloadResult["payload"], +): void { + const pending = useRecipeStudioStore.getState().pendingUploadCleanups; + if (pending.length === 0) { + return; + } + const ui = + savedPayload && typeof savedPayload === "object" + ? (savedPayload as { ui?: Record }).ui + : undefined; + const savedUid = + ui && typeof ui.unstructured_upload_uid === "string" + ? ui.unstructured_upload_uid + : ""; + const ready = pending.filter((uid) => uid !== savedUid); + if (ready.length === 0) { + return; + } + for (const uid of ready) { + void removeUnstructuredBlock(uid) + .then(() => { + useRecipeStudioStore.setState((state) => ({ + pendingUploadCleanups: state.pendingUploadCleanups.filter( + (pendingUid) => pendingUid !== uid, + ), + })); + }) + .catch((error) => { + console.warn("Failed to clean up uploaded documents:", error); + }); + } +} + export function useRecipePersistence({ recipeId, initialRecipeName, @@ -202,8 +242,10 @@ export function useRecipePersistence({ () => buildSignature(normalizedWorkflowName, currentPayload), [currentPayload, normalizedWorkflowName], ); - const isDirty = savedSignature.length > 0 && currentSignature !== savedSignature; - const saveTone: SaveTone = !isDirty && Boolean(lastSavedAt) ? "success" : "error"; + const isDirty = + savedSignature.length > 0 && currentSignature !== savedSignature; + const saveTone: SaveTone = + !isDirty && Boolean(lastSavedAt) ? "success" : "error"; const savedAtLabel = formatSavedLabel(lastSavedAt); useEffect(() => { @@ -214,7 +256,9 @@ export function useRecipePersistence({ setLastSavedAt(initialSavedAt); setCopied(false); - const parsed = importRecipePayload(JSON.stringify(initialPayload)); + const parsed = importRecipePayload(JSON.stringify(initialPayload), { + preserveUnstructuredUploads: true, + }); if (parsed.snapshot) { loadRecipe(parsed.snapshot); } else { @@ -252,6 +296,7 @@ export function useRecipePersistence({ }); setLastSavedAt(result.updatedAt); setSavedSignature(buildSignature(nextName, currentPayload)); + drainQueuedUploadCleanups(currentPayload); } catch (error) { console.error("Save recipe failed:", error); toastError("Save failed", "Could not save recipe."); @@ -270,11 +315,28 @@ export function useRecipePersistence({ return () => window.clearTimeout(timeoutId); }, [isDirty, persistRecipe, saveLoading]); + // Drain queued cleanups even when autosave is skipped: a net-zero edit (add + // then remove an unstructured seed before the 800ms debounce) keeps isDirty + // false, so the autosave effect never drains and the queued uid leaks its + // upload dir. Not-dirty means currentPayload equals the saved recipe, and + // drain skips the uid it still references, so only dirs no saved recipe + // points at are deleted (keeps the save-first invariant). + useEffect(() => { + if (!initialRecipeReady || isDirty || saveLoading) { + return; + } + drainQueuedUploadCleanups(currentPayload); + }, [currentPayload, initialRecipeReady, isDirty, saveLoading]); + const copyRecipe = useCallback(async (): Promise => { setCopied(false); try { - const safePayload = sanitizeSeedForShare(stripApiKeys(payloadResult.payload)); - const ok = await copyTextToClipboard(JSON.stringify(safePayload, null, 2)); + const safePayload = sanitizeSeedForShare( + stripApiKeys(payloadResult.payload), + ); + const ok = await copyTextToClipboard( + JSON.stringify(safePayload, null, 2), + ); if (!ok) { throw new Error("Clipboard not available."); } diff --git a/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts b/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts index 8659cad4fb..a1ff72ee9b 100644 --- a/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts +++ b/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts @@ -36,6 +36,7 @@ import { } from "../utils/handles"; import type { RecipeSnapshot } from "../utils/import"; import { getLayoutedElements } from "../utils/layout"; +import { makeUnstructuredUploadUid } from "../utils/config-factories"; import { centerModelInfraNodes, optimizeModelInfraEdgeHandles, @@ -76,6 +77,12 @@ type RecipeStudioState = { nextId: number; nextY: number; fitViewTick: number; + // Upload-uid directories whose owning block dropped them; server-side + // deletion is deferred until a save no longer references them, so a + // reload before autosave cannot leave a saved recipe pointing at + // deleted files. + pendingUploadCleanups: string[]; + queueUploadCleanup: (uid: string) => void; setSheetOpen: (open: boolean) => void; setSheetView: (view: SheetView) => void; setProcessors: (processors: RecipeProcessorConfig[]) => void; @@ -137,6 +144,7 @@ const INITIAL_STATE = { nextId: 3, nextY: 280, fitViewTick: 0, + pendingUploadCleanups: [], } satisfies Pick< RecipeStudioState, | "nodes" @@ -154,6 +162,7 @@ const INITIAL_STATE = { | "nextId" | "nextY" | "fitViewTick" + | "pendingUploadCleanups" >; function buildAddedNodeState( @@ -269,6 +278,20 @@ function isModelSemanticEdge( ); } +// Upload uid of a seed block whose server-side directory becomes orphaned +// when the block drops it. Only uid directories qualify (single owner); +// legacy node-id directories can be shared by other recipes. +function seedUploadCleanupUid(config: NodeConfig | undefined): string | null { + if (!config || config.kind !== "seed") { + return null; + } + const uid = config.unstructured_upload_uid?.trim(); + if (!uid || !config.unstructured_file_ids?.length) { + return null; + } + return uid; +} + export const useRecipeStudioStore = create((set, get) => ({ ...INITIAL_STATE, setSheetOpen: (open) => set({ sheetOpen: open }), @@ -278,6 +301,12 @@ export const useRecipeStudioStore = create((set, get) => ({ setDialogOpen: (open) => set({ dialogOpen: open }), setExecutionLocked: (locked) => set({ executionLocked: locked }), resetRecipe: () => set(INITIAL_STATE), + queueUploadCleanup: (uid) => + set((state) => + state.pendingUploadCleanups.includes(uid) + ? state + : { pendingUploadCleanups: [...state.pendingUploadCleanups, uid] }, + ), selectConfig: (id) => set({ activeConfigId: id, dialogOpen: false }), openConfig: (id) => set({ activeConfigId: id, dialogOpen: true }), setLayoutDirection: (direction) => @@ -383,7 +412,18 @@ export const useRecipeStudioStore = create((set, get) => ({ } return buildAddedNodeState(state, "sampler", type, position, openDialog); }), - addSeedNode: (type, position, openDialog = true) => + addSeedNode: (type, position, openDialog = true) => { + const current = get(); + if (!current.executionLocked) { + // The reset below clears the block's upload uid and file list; queue + // its server-side directory for deletion after the next save. + const uid = seedUploadCleanupUid( + Object.values(current.configs).find((config) => config.kind === "seed"), + ); + if (uid) { + current.queueUploadCleanup(uid); + } + } set((state) => { if (state.executionLocked) { return state; @@ -413,6 +453,8 @@ export const useRecipeStudioStore = create((set, get) => ({ hf_token: "", hf_endpoint: "https://huggingface.co", local_file_name: "", + unstructured_upload_uid: + nextSourceType === "unstructured" ? makeUnstructuredUploadUid() : "", unstructured_file_ids: [], unstructured_file_names: [], unstructured_file_sizes: [], @@ -446,7 +488,8 @@ export const useRecipeStudioStore = create((set, get) => ({ activeConfigId: existing.id, dialogOpen: openDialog, }; - }), + }); + }, addLlmNode: (type, position, openDialog = true) => set((state) => { if (state.executionLocked) { @@ -699,6 +742,9 @@ export const useRecipeStudioStore = create((set, get) => ({ dialogOpen: false, sheetView: "root", fitViewTick: state.fitViewTick + 1, + // Queued cleanups belong to the previous recipe; draining them after + // a save of this one could delete files its saved payload still uses. + pendingUploadCleanups: [], })), setAuxNodePosition: (id, position) => set((state) => { @@ -786,6 +832,17 @@ export const useRecipeStudioStore = create((set, get) => ({ set(applyUpdate); }, onNodesChange: (changes) => { + const current = get(); + if (!current.executionLocked) { + for (const change of changes) { + if (change.type === "remove") { + const uid = seedUploadCleanupUid(current.configs[change.id]); + if (uid) { + current.queueUploadCleanup(uid); + } + } + } + } const applyNodesChange = (state: RecipeStudioState) => { if (state.executionLocked) { return state; diff --git a/studio/frontend/src/features/recipe-studio/types/index.ts b/studio/frontend/src/features/recipe-studio/types/index.ts index b8ed13f70b..9231c5f7a3 100644 --- a/studio/frontend/src/features/recipe-studio/types/index.ts +++ b/studio/frontend/src/features/recipe-studio/types/index.ts @@ -340,6 +340,8 @@ export type SeedConfig = { hf_token?: string; hf_endpoint?: string; local_file_name?: string; + // ui-only: stable per-block id for uploads, since node ids collide across imports + unstructured_upload_uid?: string; unstructured_file_ids?: string[]; unstructured_file_names?: string[]; unstructured_file_sizes?: number[]; diff --git a/studio/frontend/src/features/recipe-studio/utils/config-factories.ts b/studio/frontend/src/features/recipe-studio/utils/config-factories.ts index 76fc2c38ee..d47bac8858 100644 --- a/studio/frontend/src/features/recipe-studio/utils/config-factories.ts +++ b/studio/frontend/src/features/recipe-studio/utils/config-factories.ts @@ -20,6 +20,46 @@ import type { } from "../types"; import { nextName } from "./naming"; +export function makeUnstructuredUploadUid(): string { + if (typeof globalThis.crypto?.randomUUID === "function") { + return globalThis.crypto.randomUUID().replace(/-/g, "").toLowerCase(); + } + if (typeof globalThis.crypto?.getRandomValues === "function") { + const bytes = new Uint8Array(16); + globalThis.crypto.getRandomValues(bytes); + return Array.from(bytes, (byte) => byte.toString(16).padStart(2, "0")).join( + "", + ); + } + let uid = ""; + while (uid.length < 32) { + uid += Math.floor(Math.random() * 0x100000000) + .toString(16) + .padStart(8, "0"); + } + return uid.slice(0, 32); +} + +export function resolveUnstructuredUploadBlockId({ + configId, + uploadUid, + generatedUploadUid, + unstructuredFileCount, +}: { + configId: string; + uploadUid: string; + generatedUploadUid: string | null; + unstructuredFileCount: number; +}): string { + if (uploadUid) { + return uploadUid; + } + if (generatedUploadUid) { + return generatedUploadUid; + } + return unstructuredFileCount > 0 ? configId : ""; +} + export function makeSamplerConfig( id: string, samplerType: SamplerType, @@ -368,6 +408,9 @@ export function makeSeedConfig( hf_token: "", hf_endpoint: "https://huggingface.co", local_file_name: "", + ...(seedSourceType === "unstructured" + ? { unstructured_upload_uid: makeUnstructuredUploadUid() } + : {}), unstructured_file_ids: [], unstructured_file_names: [], unstructured_file_sizes: [], diff --git a/studio/frontend/src/features/recipe-studio/utils/import/importer.ts b/studio/frontend/src/features/recipe-studio/utils/import/importer.ts index ac881d0373..54df2ffd50 100644 --- a/studio/frontend/src/features/recipe-studio/utils/import/importer.ts +++ b/studio/frontend/src/features/recipe-studio/utils/import/importer.ts @@ -16,11 +16,7 @@ import type { } from "../../types"; import { buildEdges } from "./edges"; import { isRecord, parseJson, readString } from "./helpers"; -import { - parseColumn, - parseModelConfig, - parseModelProvider, -} from "./parsers"; +import { parseColumn, parseModelConfig, parseModelProvider } from "./parsers"; import { parseSeedConfig } from "./parsers/seed-config-parser"; import { buildNodes, parseUi } from "./ui"; import type { ImportResult } from "./types"; @@ -43,6 +39,7 @@ type UiInput = { seed_drop_columns?: unknown; seed_preview_rows?: unknown; local_file_name?: unknown; + unstructured_upload_uid?: unknown; unstructured_file_ids?: unknown; unstructured_file_names?: unknown; unstructured_file_sizes?: unknown; @@ -51,6 +48,10 @@ type UiInput = { advanced_open_by_node?: unknown; }; +type ImportRecipePayloadOptions = { + preserveUnstructuredUploads?: boolean; +}; + type UiMarkdownNoteNode = { name: string; markdown: string; @@ -90,7 +91,7 @@ function parseProcessors(input: unknown): RecipeProcessorConfig[] { ? templateRaw : isRecord(templateRaw) ? JSON.stringify(templateRaw, null, 2) - : "{\n \"text\": \"{{ column_name }}\"\n}"; + : '{\n "text": "{{ column_name }}"\n}'; processors.push({ id: `p${index + 1}`, // biome-ignore lint/style/useNamingConvention: api schema @@ -135,9 +136,7 @@ function parseSeedDropColumns(input: unknown): string[] { return Array.from(values); } -function parseMcpProviders( - input: unknown, -): Map { +function parseMcpProviders(input: unknown): Map { const providers = new Map(); if (!Array.isArray(input)) { return providers; @@ -156,13 +155,12 @@ function parseMcpProviders( const args = Array.isArray(item.args) ? item.args.map((value) => String(value)) : []; - const envPairs = - isRecord(item.env) - ? Object.entries(item.env).map(([key, value]) => ({ - key: String(key), - value: String(value), - })) - : []; + const envPairs = isRecord(item.env) + ? Object.entries(item.env).map(([key, value]) => ({ + key: String(key), + value: String(value), + })) + : []; providers.set(name, { id: `mcp-${index + 1}`, name, @@ -209,7 +207,8 @@ function parseToolConfigs(input: unknown): Map { allow_tools: allowTools, // biome-ignore lint/style/useNamingConvention: api schema max_tool_call_turns: - item.max_tool_call_turns === null || item.max_tool_call_turns === undefined + item.max_tool_call_turns === null || + item.max_tool_call_turns === undefined ? "5" : String(item.max_tool_call_turns), // biome-ignore lint/style/useNamingConvention: api schema @@ -257,7 +256,9 @@ function parseUiMarkdownNoteNodes(input: unknown): UiMarkdownNoteNode[] { return noteNodes; } -function parseUiToolProfileNodes(input: unknown): Map> { +function parseUiToolProfileNodes( + input: unknown, +): Map> { const toolProfiles = new Map>(); if (!Array.isArray(input)) { return toolProfiles; @@ -312,9 +313,15 @@ function parseAdvancedOpenByNode(input: unknown): Record { return out; } -type AdvancedOpenConfig = LlmConfig | SamplerConfig | SeedConfig | ValidatorConfig; +type AdvancedOpenConfig = + | LlmConfig + | SamplerConfig + | SeedConfig + | ValidatorConfig; -function isAdvancedOpenConfig(config: NodeConfig): config is AdvancedOpenConfig { +function isAdvancedOpenConfig( + config: NodeConfig, +): config is AdvancedOpenConfig { return ( config.kind === "llm" || config.kind === "sampler" || @@ -350,7 +357,8 @@ function buildToolProfileConfig( .map((providerName) => mcpProvidersByName.get(providerName)) .flatMap((provider) => (provider ? [cloneMcpProvider(provider)] : [])), // biome-ignore lint/style/useNamingConvention: ui schema - fetched_tools_by_provider: fetchedToolsByProfileName.get(canonical.tool_alias) ?? {}, + fetched_tools_by_provider: + fetchedToolsByProfileName.get(canonical.tool_alias) ?? {}, // biome-ignore lint/style/useNamingConvention: api schema allow_tools: [...(canonical.allow_tools ?? [])], // biome-ignore lint/style/useNamingConvention: api schema @@ -360,7 +368,10 @@ function buildToolProfileConfig( }; } -export function importRecipePayload(input: string): ImportResult { +export function importRecipePayload( + input: string, + options: ImportRecipePayloadOptions = {}, +): ImportResult { const parsed = parseJson(input); if (!parsed.data || !isRecord(parsed.data)) { return { @@ -369,9 +380,9 @@ export function importRecipePayload(input: string): ImportResult { }; } - const recipe = (isRecord(parsed.data.recipe) - ? parsed.data.recipe - : parsed.data) as RecipeInput; + const recipe = ( + isRecord(parsed.data.recipe) ? parsed.data.recipe : parsed.data + ) as RecipeInput; const ui = isRecord(parsed.data.ui) ? (parsed.data.ui as UiInput) : null; if (!Array.isArray(recipe.columns)) { @@ -410,21 +421,36 @@ export function importRecipePayload(input: string): ImportResult { .map((row) => ({ ...row })) : undefined; const uiLocalFileName = readString(ui?.local_file_name) ?? undefined; - // Preserve file IDs/names from saved recipes (cleared at share time by sanitizeSeedForShare) - const uiUnstructuredFileIds: string[] = Array.isArray(ui?.unstructured_file_ids) - ? (ui.unstructured_file_ids as string[]).filter((v): v is string => typeof v === "string") - : []; - const uiUnstructuredFileNames: string[] = Array.isArray(ui?.unstructured_file_names) - ? (ui.unstructured_file_names as string[]).filter((v): v is string => typeof v === "string") - : []; - const uiUnstructuredFileSizes: number[] = Array.isArray(ui?.unstructured_file_sizes) - ? (ui.unstructured_file_sizes as number[]).filter((v): v is number => typeof v === "number") - : []; + const preserveUnstructuredUploads = + options.preserveUnstructuredUploads === true; + const uiUnstructuredUploadUid = preserveUnstructuredUploads + ? (readString(ui?.unstructured_upload_uid) ?? undefined) + : undefined; + const uiUnstructuredFileIds: string[] = + preserveUnstructuredUploads && Array.isArray(ui?.unstructured_file_ids) + ? (ui.unstructured_file_ids as string[]).filter( + (v): v is string => typeof v === "string", + ) + : []; + const uiUnstructuredFileNames: string[] = + preserveUnstructuredUploads && Array.isArray(ui?.unstructured_file_names) + ? (ui.unstructured_file_names as string[]).filter( + (v): v is string => typeof v === "string", + ) + : []; + const uiUnstructuredFileSizes: number[] = + preserveUnstructuredUploads && Array.isArray(ui?.unstructured_file_sizes) + ? (ui.unstructured_file_sizes as number[]).filter( + (v): v is number => typeof v === "number", + ) + : []; const uiUnstructuredChunkSize = readStringNumber(ui?.unstructured_chunk_size); const uiUnstructuredChunkOverlap = readStringNumber( ui?.unstructured_chunk_overlap, ); - const uiAdvancedOpenByNode = parseAdvancedOpenByNode(ui?.advanced_open_by_node); + const uiAdvancedOpenByNode = parseAdvancedOpenByNode( + ui?.advanced_open_by_node, + ); const uiMarkdownNotes = parseUiMarkdownNoteNodes(ui?.nodes); const uiToolProfilesByName = parseUiToolProfileNodes(ui?.nodes); @@ -459,11 +485,13 @@ export function importRecipePayload(input: string): ImportResult { : payloadSeedDropColumns, seed_preview_rows: uiSeedPreviewRows, local_file_name: uiLocalFileName, + unstructuredUploadUid: uiUnstructuredUploadUid, unstructuredFileIds: uiUnstructuredFileIds, unstructuredFileNames: uiUnstructuredFileNames, unstructuredFileSizes: uiUnstructuredFileSizes, unstructured_chunk_size: uiUnstructuredChunkSize, unstructured_chunk_overlap: uiUnstructuredChunkOverlap, + preserveUnstructuredUploads, }); if (seedConfig) { applyAdvancedOpen(seedConfig, uiAdvancedOpenByNode); @@ -567,12 +595,7 @@ export function importRecipePayload(input: string): ImportResult { const { layouts, auxNodes, edges: uiEdges, layoutDirection } = parseUi(ui); const resolvedLayoutDirection = layoutDirection ?? "LR"; const nodes = buildNodes(configs, layouts); - const edges = buildEdges( - configs, - nameToId, - uiEdges, - resolvedLayoutDirection, - ); + const edges = buildEdges(configs, nameToId, uiEdges, resolvedLayoutDirection); const auxNodePositions = Object.fromEntries( auxNodes.flatMap((item) => { const llmId = nameToId.get(item.llm); @@ -583,10 +606,7 @@ export function importRecipePayload(input: string): ImportResult { }), ); - const maxY = nodes.reduce( - (acc, node) => Math.max(acc, node.position.y), - 0, - ); + const maxY = nodes.reduce((acc, node) => Math.max(acc, node.position.y), 0); return { errors: [], diff --git a/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts b/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts index 21eadb3195..939205fe6d 100644 --- a/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts +++ b/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts @@ -197,17 +197,26 @@ export function parseSeedConfig( seed_drop_columns?: string[]; seed_preview_rows?: Record[]; local_file_name?: string; + unstructuredUploadUid?: string; unstructuredFileIds?: string[]; unstructuredFileNames?: string[]; unstructuredFileSizes?: number[]; unstructured_chunk_size?: string; unstructured_chunk_overlap?: string; + preserveUnstructuredUploads?: boolean; }, ): SeedConfig | null { if (!seedConfigRaw) { return null; } - const parsed = parseSeedSettings(seedConfigRaw); + const parsed = { ...parseSeedSettings(seedConfigRaw) }; + if ( + parsed.seed_source_type === "unstructured" && + options?.preserveUnstructuredUploads !== true + ) { + parsed.hf_path = ""; + parsed.resolved_paths = []; + } let sourceType: SeedSourceType = "hf"; if (parsed.seed_source_type === "hf") { sourceType = "hf"; @@ -230,6 +239,9 @@ export function parseSeedConfig( ...(options?.local_file_name !== undefined ? { local_file_name: options.local_file_name } : {}), + ...(options?.unstructuredUploadUid + ? { unstructured_upload_uid: options.unstructuredUploadUid } + : {}), ...(options?.unstructuredFileIds !== undefined ? { unstructured_file_ids: options.unstructuredFileIds } : {}), diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts b/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts index 34c3c34274..9b2ad5b2b5 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts @@ -440,6 +440,9 @@ export function buildRecipePayload( unstructured_file_names: firstSeed.unstructured_file_names, unstructured_file_sizes: firstSeed.unstructured_file_sizes, }), + ...(firstSeed?.unstructured_upload_uid?.trim() && { + unstructured_upload_uid: firstSeed.unstructured_upload_uid, + }), ...(firstSeed && firstSeed.unstructured_chunk_size !== undefined && { unstructured_chunk_size: firstSeed.unstructured_chunk_size, diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/types.ts b/studio/frontend/src/features/recipe-studio/utils/payload/types.ts index 902ea796b4..763e68c1fb 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/types.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/types.ts @@ -71,6 +71,8 @@ export type RecipePayload = { seed_preview_rows?: Record[]; local_file_name?: string; // biome-ignore lint/style/useNamingConvention: api schema + unstructured_upload_uid?: string; + // biome-ignore lint/style/useNamingConvention: api schema unstructured_file_ids?: string[]; // biome-ignore lint/style/useNamingConvention: api schema unstructured_file_names?: string[]; diff --git a/studio/frontend/src/features/settings/api/coding-agents.ts b/studio/frontend/src/features/settings/api/coding-agents.ts new file mode 100644 index 0000000000..ae371b2d3a --- /dev/null +++ b/studio/frontend/src/features/settings/api/coding-agents.ts @@ -0,0 +1,45 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { authFetch } from "@/features/auth"; +import { readFastApiError } from "@/lib/format-fastapi-error"; + +export type CodingAgentsInfo = { + // Every agent `unsloth start` supports, in the CLI's declared order. + agents: string[]; + // Subset of `agents` whose CLI binary was found on PATH by the backend. + detected: string[]; +}; + +type ApiCodingAgentsInfo = { + agents: string[]; + detected: string[]; +}; + +// Which CLIs are on PATH is environment state, not a persisted setting -- it +// can change any time the user installs something new, so this only +// de-duplicates concurrent in-flight calls (e.g. React strict-mode's double +// mount) rather than caching the result across the module's lifetime. Every +// fresh call (each time a settings panel mounts) re-checks PATH for real. +let inFlightInfo: Promise | null = null; + +function fromApi(info: ApiCodingAgentsInfo): CodingAgentsInfo { + return { agents: info.agents, detected: info.detected }; +} + +async function fetchCodingAgents(): Promise { + const res = await authFetch("/api/settings/coding-agents"); + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to load installed coding agents"), + ); + } + return fromApi(await res.json()); +} + +export async function loadCodingAgents(): Promise { + inFlightInfo ??= fetchCodingAgents().finally(() => { + inFlightInfo = null; + }); + return inFlightInfo; +} diff --git a/studio/frontend/src/features/settings/api/embedding-model.ts b/studio/frontend/src/features/settings/api/embedding-model.ts index 8b6bc7ee7f..9a61142f73 100644 --- a/studio/frontend/src/features/settings/api/embedding-model.ts +++ b/studio/frontend/src/features/settings/api/embedding-model.ts @@ -23,6 +23,10 @@ type ApiEmbeddingModelSettings = { * (wrong type, gated repo, or offline). Retry with force to save anyway. */ export class EmbeddingModelVerificationError extends Error {} +/** 403 from the backend: the repo is flagged unsafe by Hugging Face's security scan. + * A hard block; force cannot bypass it, so it must not enter the "save anyway" flow. */ +export class EmbeddingModelBlockedError extends Error {} + function fromApi(settings: ApiEmbeddingModelSettings): EmbeddingModelSettings { return { embeddingModel: settings.embedding_model, @@ -56,6 +60,11 @@ export async function updateEmbeddingModelSettings( force: options?.force ?? false, }), }); + if (res.status === 403) { + throw new EmbeddingModelBlockedError( + await readFastApiError(res, "This model is blocked by a security scan"), + ); + } if (res.status === 409) { throw new EmbeddingModelVerificationError( await readFastApiError(res, "Could not verify the embedding model"), diff --git a/studio/frontend/src/features/settings/components/agent-command.ts b/studio/frontend/src/features/settings/components/agent-command.ts index 9e87922970..38b2d73c3b 100644 --- a/studio/frontend/src/features/settings/components/agent-command.ts +++ b/studio/frontend/src/features/settings/components/agent-command.ts @@ -12,7 +12,7 @@ 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 { +export function normalizeHost(host: string): string { const lower = host.toLowerCase(); return lower.startsWith("[") && lower.endsWith("]") ? lower.slice(1, -1) : lower; } @@ -26,7 +26,7 @@ function isDefaultLocalHost(host: string): boolean { } // Match the CLI auto-mint rule (is_loopback_url): localhost, ::1, and all of 127.0.0.0/8. -function isLoopbackHost(host: string): boolean { +export function isLoopbackHost(host: string): boolean { if (host === "localhost" || host === "::1") return true; const octets = host.split("."); return ( diff --git a/studio/frontend/src/features/settings/components/usage-examples.tsx b/studio/frontend/src/features/settings/components/usage-examples.tsx index c43dd0f219..b3396c8d9b 100644 --- a/studio/frontend/src/features/settings/components/usage-examples.tsx +++ b/studio/frontend/src/features/settings/components/usage-examples.tsx @@ -16,6 +16,7 @@ import { fetchDeviceType, usePlatformStore } from "@/config/env"; import { useChatRuntimeStore } from "@/features/chat"; import { useT } from "@/i18n"; import type { TranslationKey } from "@/i18n"; +import { isTauri } from "@/lib/api-base"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; import { Tick02Icon } from "@/lib/tick-icon"; import { cn } from "@/lib/utils"; @@ -25,14 +26,15 @@ import { InformationCircleIcon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { useEffect, useMemo, useState } from "react"; +import { useEffect, useMemo, useRef, useState } from "react"; import { Streamdown } from "streamdown"; +import { loadCodingAgents } from "../api/coding-agents"; import { type OpenAIAutoSwitchSettings, loadOpenAIAutoSwitchSettings, updateOpenAIAutoSwitchSettings, } from "../api/openai-auto-switch"; -import { buildAgentCommand } from "./agent-command"; +import { buildAgentCommand, isLoopbackHost, normalizeHost } from "./agent-command"; type ExampleType = | "curl" @@ -114,6 +116,30 @@ const DOC_LINKS = [ { label: "Hermes Agent", href: "https://unsloth.ai/docs/integrations/hermes-agent" }, ]; +// Falls back to this list until the backend's installed-CLI check resolves; +// kept in sync with the `unsloth start ` subcommands and with +// CODING_AGENTS in studio/backend/utils/coding_agents.py. +const DEFAULT_AGENTS = [ + "claude", + "codex", + "openclaw", + "opencode", + "hermes", + "pi", +]; +// The agent selection resets to this whenever an auto-pick is no longer +// trustworthy (leaving loopback, or the only compatible detected agent +// stops being compatible) rather than lingering on a stale choice. +const DEFAULT_AGENT = "claude"; +const AGENT_LABELS: Record = { + claude: "Claude Code", + codex: "Codex", + openclaw: "OpenClaw", + opencode: "OpenCode", + hermes: "Hermes", + pi: "Pi", +}; + const j = (s: string): string => JSON.stringify(s); const shSingle = (s: string): string => s.replace(/'/g, "'\\''"); const psSingle = (s: string): string => s.replace(/'/g, "''"); @@ -399,6 +425,17 @@ function useLoadedModelName(): string { }, [checkpoint, ggufVariant]); } +// Backend PATH detection is only safe in the desktop app, where the UI owns +// the local backend. A browser loopback URL may be an SSH/local port forward. +function canUseLocalAgentDetection(base: string): boolean { + if (!isTauri) return false; + try { + return isLoopbackHost(normalizeHost(new URL(base).hostname)); + } catch { + return false; + } +} + const SHIKI_THEMES = [unslothLightTheme, unslothDarkTheme] as [ typeof unslothLightTheme, typeof unslothDarkTheme, @@ -443,7 +480,18 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { const [copied, setCopied] = useState(false); const [copiedUrl, setCopiedUrl] = useState(false); const [copiedAgent, setCopiedAgent] = useState(false); + const [agent, setAgent] = useState(DEFAULT_AGENT); + const [availableAgents, setAvailableAgents] = + useState(DEFAULT_AGENTS); + const [detectedAgents, setDetectedAgents] = useState([]); + // True once the user has picked an agent themselves; guards the detection + // effect below from clobbering that choice if it resolves afterward. + const agentPickedByUserRef = useRef(false); const [useTunnel, setUseTunnel] = useState(readUseTunnelPref); + const origin = typeof window !== "undefined" ? window.location.origin : ""; + const base = + useTunnel && cloudflareUrl ? cloudflareUrl : (serverUrl ?? origin); + const localAgentDetection = canUseLocalAgentDetection(base); // null while loading; the same setting the General tab exposes (shared cache). const [autoSwitch, setAutoSwitch] = useState( null, @@ -454,6 +502,78 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { void fetchDeviceType({ force: true }); }, []); + // Fetching is the only job of this effect: populate availableAgents/ + // detectedAgents (or clear them). Which agent gets auto-picked from that + // list is derived separately below, so it can react to the loaded model + // changing too, not just a fresh fetch. + useEffect(() => { + // Browser loopback URLs can be SSH/local forwards, so only the desktop app + // may use backend PATH checks to mark or auto-pick local agents. + if (!localAgentDetection) { + setDetectedAgents([]); + // A previously auto-picked agent was only ever verified against the + // Studio backend's PATH, which is meaningless now that this panel no + // longer targets a loopback base -- don't leave it selected, but + // never touch a choice the user made by hand. + if (!agentPickedByUserRef.current) { + setAgent(DEFAULT_AGENT); + } + return; + } + + let cancelled = false; + void loadCodingAgents() + .then((info) => { + if (cancelled) return; + setAvailableAgents(info.agents); + setDetectedAgents(info.detected); + }) + .catch(() => { + // Best-effort: keep the default agent list and let the user pick manually. + }); + return () => { + cancelled = true; + }; + }, [localAgentDetection]); + + // Single source of truth for the auto-picked agent, re-derived whenever + // the detected list or the loaded model's GGUF-ness changes -- in either + // direction. `codex` needs a GGUF model (unsloth_cli's + // _require_gguf_for_codex exits otherwise), so it's only preferred once + // the loaded model actually qualifies; loading a GGUF model *after* a + // non-GGUF-gated fallback picked something else re-steers back to codex + // just as loading a non-GGUF model steers away from it. Never overrides a + // choice the user made by hand. + // activeGgufVariant alone only covers an HF-repo GGUF pick (a specific + // quant variant string) -- a direct local .gguf file (custom folder / + // LM Studio / drag-drop) is just as much a GGUF the codex preflight would + // accept, but never has a "variant" to report, and would otherwise read as + // non-GGUF here. activeNativePathToken covers the drag-drop/picked-file + // case; ggufContextLength is only ever populated when the backend's + // /api/inference/status last reported is_gguf: true for the active model + // (see applyActiveModelStatusToStore), so together these three cover every + // path a model can be GGUF through, matching the same is_gguf-or-equivalent + // check hasGgufSource applies to a staged pick. + const activeGgufVariant = useChatRuntimeStore((s) => s.activeGgufVariant); + const activeNativePathToken = useChatRuntimeStore((s) => s.activeNativePathToken); + const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength); + useEffect(() => { + if (agentPickedByUserRef.current) return; + if (detectedAgents.length === 0) return; + const isGguf = + activeGgufVariant != null || activeNativePathToken != null || ggufContextLength != null; + const preferred = detectedAgents.find((a) => a !== "codex" || isGguf); + if (preferred) { + setAgent(preferred); + } else if (agent === "codex" && !isGguf) { + // codex was auto-picked while a GGUF model was active and it's the + // only detected agent; now that the model isn't GGUF anymore, nothing + // detected is actually runnable, so fall back to the default instead + // of leaving a codex command unsloth_cli will reject. + setAgent(DEFAULT_AGENT); + } + }, [agent, detectedAgents, activeGgufVariant, activeNativePathToken, ggufContextLength]); + useEffect(() => { let cancelled = false; void loadOpenAIAutoSwitchSettings() @@ -470,9 +590,6 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { const model = useLoadedModelName(); const key = apiKey || KEY_PLACEHOLDER; - const origin = typeof window !== "undefined" ? window.location.origin : ""; - const base = - useTunnel && cloudflareUrl ? cloudflareUrl : (serverUrl ?? origin); const autoSwitchOn = autoSwitch?.enabled ?? false; const snippets = useMemo( @@ -481,8 +598,8 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { ); // Agent command must target the server the panel shows, not the :8888 default. const agentCommand = useMemo( - () => buildAgentCommand(base, key, os), - [base, key, os], + () => buildAgentCommand(base, key, os, agent), + [base, key, os, agent], ); const osAware = OS_AWARE[lang]; @@ -710,6 +827,42 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { {t("settings.apiKeys.codingAgentsHint")} +
+ {availableAgents.map((id) => { + const installed = detectedAgents.includes(id); + const active = agent === id; + return ( + + ); + })} +
{agentCommand} @@ -727,7 +880,13 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
- {t("settings.apiKeys.codingAgentsSwap")} + {detectedAgents.length > 0 + ? t("settings.apiKeys.codingAgentsDetectedHint", { + agents: detectedAgents + .map((id) => AGENT_LABELS[id] ?? id) + .join(", "), + }) + : t("settings.apiKeys.codingAgentsSwap")}
diff --git a/studio/frontend/src/features/settings/index.ts b/studio/frontend/src/features/settings/index.ts index 364ca1611f..3fefd8c63a 100644 --- a/studio/frontend/src/features/settings/index.ts +++ b/studio/frontend/src/features/settings/index.ts @@ -7,6 +7,7 @@ export { savePersonalization, } from "./api/personalization"; export { setTheme, useTheme } from "./stores/theme-store"; +export { useMonitorOverlayStore } from "./stores/monitor-overlay-store"; export type { Personalization, PersonalizationAppearance, diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx index 0d201edc0d..d8a0d45d3f 100644 --- a/studio/frontend/src/features/settings/settings-dialog.tsx +++ b/studio/frontend/src/features/settings/settings-dialog.tsx @@ -56,6 +56,7 @@ const TABS: TabDef[] = [ id: "resources", labelKey: "settings.tabs.resources", icon: CpuIcon, + badgeKey: "common.new", }, { id: "chat", @@ -72,7 +73,6 @@ const TABS: TabDef[] = [ id: "connections", labelKey: "settings.tabs.connections", icon: CloudIcon, - badgeKey: "common.new", }, { id: "about", labelKey: "settings.tabs.about", icon: HelpCircleIcon }, ]; diff --git a/studio/frontend/src/features/settings/tabs/about-tab.tsx b/studio/frontend/src/features/settings/tabs/about-tab.tsx index ff751e3cd6..1f323a9e8d 100644 --- a/studio/frontend/src/features/settings/tabs/about-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/about-tab.tsx @@ -158,7 +158,7 @@ export function AboutTab() { {gpu.name ?? "—"} {gpu.vramTotalGb != null - ? ` · ${Math.round(gpu.vramTotalGb)} GB` + ? ` · ${Math.round(gpu.vramTotalGb)} GiB` : ""} diff --git a/studio/frontend/src/features/settings/tabs/chat-tab.tsx b/studio/frontend/src/features/settings/tabs/chat-tab.tsx index b26b09e022..9519898b21 100644 --- a/studio/frontend/src/features/settings/tabs/chat-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/chat-tab.tsx @@ -213,6 +213,12 @@ export function ChatTab() { const setShowModelDisclaimer = useChatPreferencesStore( (state) => state.setShowModelDisclaimer, ); + const showResponseModel = useChatPreferencesStore( + (state) => state.showResponseModel, + ); + const setShowResponseModel = useChatPreferencesStore( + (state) => state.setShowResponseModel, + ); useEffect(() => { void countAllChats().then(setCount); @@ -412,6 +418,15 @@ export function ChatTab() { onCheckedChange={setShowModelDisclaimer} /> + + + diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index df684b7752..8fd70cc1b7 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -42,6 +42,7 @@ import { updatePreviewSharing, } from "../api/preview-sharing"; import { + EmbeddingModelBlockedError, type EmbeddingModelSettings, EmbeddingModelVerificationError, loadEmbeddingModelSettings, @@ -78,6 +79,7 @@ const PREFS_KEYS: string[] = [ "unsloth_chat_auto_title", "unsloth_hf_token", "unsloth_auto_heal_tool_calls", + "unsloth_nudge_tool_calls", "unsloth_max_tool_calls_per_message", "unsloth_tool_call_timeout", "unsloth_chat_inference_params", @@ -409,7 +411,10 @@ export function GeneralTab() { description: t("settings.general.rag.reindexWarning"), }); } catch (error) { - if (error instanceof EmbeddingModelVerificationError) { + // A hard security block cannot be forced; keep the "save anyway" action hidden. + if (error instanceof EmbeddingModelBlockedError) { + setEmbeddingModelNeedsForce(false); + } else if (error instanceof EmbeddingModelVerificationError) { setEmbeddingModelNeedsForce(true); } setEmbeddingModelError( diff --git a/studio/frontend/src/features/settings/tabs/resources-tab.tsx b/studio/frontend/src/features/settings/tabs/resources-tab.tsx index d5e19cc51c..6c30858c63 100644 --- a/studio/frontend/src/features/settings/tabs/resources-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/resources-tab.tsx @@ -47,6 +47,15 @@ function formatGb(value: number | null | undefined): string { return `${safe.toFixed(digits)} GB`; } +// RAM/VRAM come from the backend in binary units (bytes / 1024**3), matching +// nvidia-smi and PyTorch, so label those readouts GiB. Disk stays on formatGb +// because the backend reports disk in decimal GB (bytes / 1e9). +function formatGiB(value: number | null | undefined): string { + const safe = isFiniteNumber(value) ? Math.max(0, value) : 0; + const digits = safe >= 10 ? 1 : 2; + return `${safe.toFixed(digits)} GiB`; +} + function formatMb(value: number | null | undefined): string { const safe = isFiniteNumber(value) ? Math.max(0, value) : 0; return `${Math.round(safe).toLocaleString()} MB`; @@ -300,9 +309,9 @@ export function ResourcesTab() { /> @@ -318,13 +327,13 @@ export function ResourcesTab() { label={t("settings.resources.liveMonitor.vram")} value={ hasGpu - ? `${formatGb(metrics.vramUsed)} / ${formatGb(metrics.vramTotal)}` + ? `${formatGiB(metrics.vramUsed)} / ${formatGiB(metrics.vramTotal)}` : t("settings.resources.liveMonitor.noGpu") } detail={ hasGpu ? t("settings.resources.liveMonitor.free", { - value: formatGb(metrics.vramFree), + value: formatGiB(metrics.vramFree), }) : backendLabel } @@ -373,17 +382,17 @@ export function ResourcesTab() {
{t("settings.resources.gpu.used", { - value: formatGb(used), + value: formatGiB(used), })} {t("settings.resources.gpu.free", { - value: formatGb(free), + value: formatGiB(free), })} {t("settings.resources.gpu.total", { - value: formatGb(total), + value: formatGiB(total), })}
diff --git a/studio/frontend/src/features/studio/sections/progress-section.tsx b/studio/frontend/src/features/studio/sections/progress-section.tsx index abab35db93..9c9398688f 100644 --- a/studio/frontend/src/features/studio/sections/progress-section.tsx +++ b/studio/frontend/src/features/studio/sections/progress-section.tsx @@ -411,7 +411,7 @@ function LiveGpuPanel({ value={index} className="bg-popover text-popover-foreground dark:bg-zinc-900 dark:text-zinc-100" > - GPU {device.visible_ordinal ?? index} - {device.backend} ({device.vram_total_gb ? `${Math.round(device.vram_total_gb)}GB` : "N/A"}) + GPU {device.visible_ordinal ?? index} - {device.backend} ({device.vram_total_gb ? `${Math.round(device.vram_total_gb)}GiB` : "N/A"}) ))} @@ -446,7 +446,7 @@ function LiveGpuPanel({ icon={} value={ currentGpu.vram_used_gb != null && currentGpu.vram_total_gb != null - ? `${currentGpu.vram_used_gb} / ${currentGpu.vram_total_gb} GB` + ? `${currentGpu.vram_used_gb} / ${currentGpu.vram_total_gb} GiB` : "--" } pct={currentGpu.vram_utilization_pct ?? 0} diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index 3d73ad4343..e0cb8030ae 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -41,6 +41,7 @@ export const en = { recipes: "Recipes", export: "Export", recents: "Recents", + noChatsYet: "No chats yet", settings: "Settings", api: "API", lightMode: "Light Mode", @@ -408,7 +409,6 @@ export const en = { description: "Access Unsloth via the OpenAI-compatible API.", readDocs: "Read the API docs", noAccess: "No API access yet.", - newBadge: "New", accessTokens: "Access tokens", loadError: "Couldn't load API access.", createError: "Couldn't create access token.", @@ -444,6 +444,8 @@ export const en = { 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.", + codingAgentDetected: "Installed on this machine", + codingAgentsDetectedHint: "Detected on this machine: {agents}.", relativeNever: "never", relativeJustNow: "just now", relativeHoursAgo: "{count}h ago", diff --git a/studio/frontend/src/i18n/locales/ja.ts b/studio/frontend/src/i18n/locales/ja.ts index f2020f3cfb..08b8d4f4d3 100644 --- a/studio/frontend/src/i18n/locales/ja.ts +++ b/studio/frontend/src/i18n/locales/ja.ts @@ -296,7 +296,6 @@ export const ja = { description: "OpenAI互換 API を介して Unsloth にアクセスします。", readDocs: "API ドキュメントを読む", noAccess: "まだ API アクセス権がありません。", - newBadge: "新規", accessTokens: "アクセストークン", loadError: "API アクセス権を読み込めませんでした。", createError: "アクセストークンを作成できませんでした。", diff --git a/studio/frontend/src/i18n/locales/pt-br.ts b/studio/frontend/src/i18n/locales/pt-br.ts index c261e4ed0c..494a98ec32 100644 --- a/studio/frontend/src/i18n/locales/pt-br.ts +++ b/studio/frontend/src/i18n/locales/pt-br.ts @@ -363,7 +363,6 @@ export const ptBR = { "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.", diff --git a/studio/frontend/src/i18n/locales/zh-CN.ts b/studio/frontend/src/i18n/locales/zh-CN.ts index 5fd31dc7cb..dda8e017ab 100644 --- a/studio/frontend/src/i18n/locales/zh-CN.ts +++ b/studio/frontend/src/i18n/locales/zh-CN.ts @@ -41,6 +41,7 @@ export const zhCN = { recipes: "配方", export: "导出", recents: "最近", + noChatsYet: "暂无对话", settings: "设置", api: "API", lightMode: "浅色模式", @@ -266,7 +267,6 @@ export const zhCN = { description: "通过兼容 OpenAI 的 API 以编程方式访问 Unsloth。", readDocs: "阅读 API 文档", noAccess: "还没有 API 访问权限。", - newBadge: "新", accessTokens: "访问 token", loadError: "无法加载 API 访问权限。", createError: "无法创建访问 token。", diff --git a/studio/frontend/src/lib/latex.ts b/studio/frontend/src/lib/latex.ts index 1a4ebdace8..edf9875602 100644 --- a/studio/frontend/src/lib/latex.ts +++ b/studio/frontend/src/lib/latex.ts @@ -1,8 +1,13 @@ // Adapted from LibreChat's latex.ts // https://github.com/danny-avila/LibreChat/blob/main/client/src/utils/latex.ts // -// Escapes currency dollar signs so they are not misinterpreted as LaTeX math -// delimiters when singleDollarTextMath is enabled. +// Two jobs, in order: +// 1. Convert LaTeX bracket delimiters (`\[...\]`, `\(...\)`) into the dollar +// forms remark-math understands (`$$...$$`, `$...$`). remark-math only +// tokenizes dollar delimiters, so models that emit `\[...\]` / `\(...\)` +// would otherwise render as literal text. +// 2. Escape currency dollar signs so they are not misinterpreted as LaTeX +// math delimiters when singleDollarTextMath is enabled. /** * Matches a single $ followed by a number pattern (currency), e.g.: @@ -15,14 +20,14 @@ const CURRENCY_REGEX = /(? { const regions: Array<[number, number]> = []; - // Fenced code blocks: ```...``` - const fencedRe = /```[\s\S]*?```/g; + // Fenced code blocks: ```...``` and ~~~...~~~ (both are code in GFM) + const fencedRe = /```[\s\S]*?```|~~~[\s\S]*?~~~/g; let match: RegExpExecArray | null; while ((match = fencedRe.exec(content)) !== null) { regions.push([match.index, match.index + match[0].length]); @@ -51,9 +56,38 @@ function findCodeBlockRegions(content: string): Array<[number, number]> { } /** - * Binary search to check if a position falls inside any code region. + * Match an inline link/image `[text](DEST)`, capturing the destination as group 1 + * with the `d` flag so its span is read straight from `match.indices` (the text + * can contain an escaped `\](`, so a string search for the separator is unsafe). + * The text disallows unescaped `]`; the destination allows escapes and one level + * of balanced parens. */ -function isInCodeBlock( +const LINK_DEST_RE = + /!?\[(?:\\.|[^\]\\])*?\]\(((?:\\.|[^()\\]|\([^()]*\))*)\)/gd; + +/** + * Find the destination spans of inline links/images, so a `\(...\)` written with + * escaped parens inside a URL isn't rewritten as math (which would break the + * link). Only the destination is returned, not the link text, so math in the + * visible text still converts. Sorted, non-overlapping (matches are disjoint). + */ +function findLinkDestinationRegions(content: string): Array<[number, number]> { + if (!content.includes("](")) return []; + const regions: Array<[number, number]> = []; + let match: RegExpExecArray | null; + LINK_DEST_RE.lastIndex = 0; + while ((match = LINK_DEST_RE.exec(content)) !== null) { + // `indices` is present (the `d` flag); group 1 spans the destination. + regions.push(match.indices![1]); + } + return regions; +} + +/** + * Binary search to check if a position falls inside any region. Regions must be + * sorted by start and non-overlapping. + */ +function isInRegion( position: number, regions: Array<[number, number]>, ): boolean { @@ -139,7 +173,11 @@ function looksLikeMathBody(body: string): boolean { * (`**$X$**`, `__$X$__`) are always math: LLMs use that for "bold math" * and the heuristic would otherwise reject prose-shaped bodies like "90 - x". */ -function hasInlineMathCloser(content: string, offset: number): boolean { +function hasInlineMathCloser( + content: string, + offset: number, + mathRegions: Array<[number, number]>, +): boolean { const MAX_SPAN = 200; const limit = Math.min(content.length, offset + 1 + MAX_SPAN); for (let i = offset + 1; i < limit; i++) { @@ -147,6 +185,9 @@ function hasInlineMathCloser(content: string, offset: number): boolean { if (c === "\n") return false; if (c !== "$") continue; if (content[i - 1] === "\\") continue; + // A `$` opening a generated span (from `\(...\)`) is not a currency closer; + // pairing with it would swallow the price into math (`$5 + x \(y\)`). + if (isInRegion(i, mathRegions)) return false; if (content[i + 1] === "$") { i++; continue; @@ -174,9 +215,124 @@ function hasInlineMathCloser(content: string, offset: number): boolean { } /** - * Preprocess a markdown string to escape currency dollar signs so they are not - * parsed as LaTeX math delimiters. + * Matches a `\[...\]` (display) or `\(...\)` (inline) LaTeX span. Non-greedy so + * the first closer wins; dotall so display spans can wrap lines. `(? block `$$...$$` and `\(...\)` -> inline `$...$` so + * remark-math can tokenize them. Bodies are trimmed: remark-math won't open an + * inline span on `$ ` (a `$` followed by whitespace), and display fences must + * sit on their own line to render as a centered block (not inline math), so + * `\[...\]` becomes `\n$$\n...\n$$\n`. * + * Spans inside code blocks/spans are left intact (a code sample showing `\(x\)` + * must not be rewritten). + * + * A space is inserted between a converted span and a following `$` so their + * delimiters can't fuse (`\(a\)\(b\)` -> `$a$$b$` would mis-tokenize into one + * broken span). A preceding currency (`$5\(x\)`) is instead broken later by the + * currency escape pass. + * + * Returns the rewritten text and the `[start, end)` ranges (in the rewritten + * string) of every span it produced, so the currency pass can skip them. + */ +function convertLatexDelimiters(content: string): { + text: string; + mathRegions: Array<[number, number]>; +} { + if (!content.includes("\\[") && !content.includes("\\(")) { + return { text: content, mathRegions: [] }; + } + + const codeRegions = findCodeBlockRegions(content); + const linkRegions = findLinkDestinationRegions(content); + const inSkipZone = (pos: number) => + isInRegion(pos, codeRegions) || isInRegion(pos, linkRegions); + // Pushed in ascending, non-overlapping order (offset only grows), so this + // stays valid for isInRegion's binary search without a sort. + const mathRegions: Array<[number, number]> = []; + // Accumulate into an array, not a string: reading the last char off a growing + // `+=` accumulator flattens its rope every append (O(n^2) over many spans, on + // the per-frame streaming path), so track the tail char and length instead. + const parts: string[] = []; + let offset = 0; + let lastChar = ""; + let last = 0; + // Append a chunk, separating a trailing `$` from a leading `$` so two spans + // can't fuse. Returns where the chunk landed (after any inserted space). + const append = (chunk: string): number => { + if (!chunk) return offset; + if (lastChar === "$" && chunk.startsWith("$")) { + parts.push(" "); + offset += 1; + } + const start = offset; + parts.push(chunk); + offset += chunk.length; + lastChar = chunk[chunk.length - 1]; + return start; + }; + let match: RegExpExecArray | null; + LATEX_DELIM_RE.lastIndex = 0; + while ((match = LATEX_DELIM_RE.exec(content)) !== null) { + const matchEnd = match.index + match[0].length; + // Skip if either delimiter is inside code or a link destination: an opener + // outside such a zone must not consume a closer inside one and rewrite + // across the boundary. Resume right after this opener (not past the whole + // match) so a valid span that this match spanned across (a stray code `\(` + // paired with a real closer) is still found on the next pass, not swallowed. + if (inSkipZone(match.index) || inSkipZone(matchEnd - 1)) { + LATEX_DELIM_RE.lastIndex = match.index + 1; + continue; + } + const isDisplay = match[1] !== undefined; + const body = (isDisplay ? match[1] : match[2]).trim(); + // Leave an empty span (`\(\)`) literal; a bare `$$` would open a stray + // display block that swallows following text. + if (!body) { + continue; + } + append(content.slice(last, match.index)); + let wrapped: string; + if (isDisplay) { + // Keep the opener's leading indentation so a `$$` block inside a list item + // stays in the container instead of breaking out at column 0. Only when the + // opener is whitespace-prefixed, so inline `text \[x\]` keeps column 0. + const lineStart = + match.index > 0 ? content.lastIndexOf("\n", match.index - 1) + 1 : 0; + const prefix = content.slice(lineStart, match.index); + const indent = /^\s*$/.test(prefix) ? prefix : ""; + // Indent every body line, not just the first, so multi-line display math + // (`\[a\nb\]`) stays wholly inside the container. + const inner = indent ? body.replace(/\n/g, `\n${indent}`) : body; + wrapped = `\n${indent}$$\n${indent}${inner}\n${indent}$$\n`; + } else { + wrapped = `$${body}$`; + } + const start = append(wrapped); + mathRegions.push([start, offset]); + last = matchEnd; + } + append(content.slice(last)); + return { text: parts.join(""), mathRegions }; +} + +/** + * Preprocess a markdown string so LaTeX renders: convert bracket delimiters to + * dollar forms, then escape currency dollar signs so they are not parsed as + * math delimiters. + * + * - `\[E = mc^2\]` becomes a `$$` display block on its own lines (display math) + * - `\(\alpha\)` becomes `$\alpha$` (inline math) + * - `\(x\)` in a code span is untouched * - `$5` alone becomes `\$5` (currency, not math) * - `$\alpha$` is untouched (real LaTeX) * - `$30^\circ$` is untouched (LaTeX whose body starts with a digit) @@ -185,15 +341,22 @@ function hasInlineMathCloser(content: string, offset: number): boolean { * - Currency inside code blocks/spans is untouched */ export function preprocessLaTeX(content: string): string { - if (!content.includes("$")) return content; + const { text, mathRegions } = convertLatexDelimiters(content); - const codeRegions = findCodeBlockRegions(content); + if (!text.includes("$")) return text; - return content.replace(CURRENCY_REGEX, (match, offset) => { - if (isInCodeBlock(offset, codeRegions)) { + const codeRegions = findCodeBlockRegions(text); + + return text.replace(CURRENCY_REGEX, (match, offset) => { + if (isInRegion(offset, codeRegions)) { return match; } - if (hasInlineMathCloser(content, offset)) { + // Skip the spans we just created from `\(...\)` so a numeric body like + // `$5$` isn't re-escaped back to literal `\$5$`. + if (isInRegion(offset, mathRegions)) { + return match; + } + if (hasInlineMathCloser(text, offset, mathRegions)) { return match; } return "\\" + match; diff --git a/studio/frontend/src/main.tsx b/studio/frontend/src/main.tsx index 0922e764bb..d0ddf2fc6e 100644 --- a/studio/frontend/src/main.tsx +++ b/studio/frontend/src/main.tsx @@ -5,8 +5,8 @@ import { StrictMode } from "react"; import { createRoot } from "react-dom/client"; import "./index.css"; -import { fetchDeviceType } from "./config/env"; import { App } from "./app/app"; +import { fetchDeviceType } from "./config/env"; import { initializeLocale } from "./i18n"; const globalCrypto = globalThis.crypto as Crypto | undefined; @@ -36,10 +36,10 @@ if (!rootElement) { initializeLocale(); -fetchDeviceType().then(() => { - createRoot(rootElement).render( - - - , - ); -}); +createRoot(rootElement).render( + + + , +); + +fetchDeviceType().catch(() => undefined); diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index e40cb3083e..ca1fe79efa 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -10,6 +10,7 @@ import argparse import atexit import errno import fnmatch +import glob import hashlib import json import os @@ -165,9 +166,9 @@ def env_int( # errors. Only use "master" temporarily when the latest release is missing # support for a new model architecture. DEFAULT_LLAMA_TAG = os.environ.get("UNSLOTH_LLAMA_TAG", "latest") -# Default published repo for prebuilt release resolution. Linux uses -# Unsloth prebuilts; setup.sh/setup.ps1 pass --published-repo explicitly -# for macOS/Windows to override with ggml-org/llama.cpp when needed. +# Default published repo for prebuilt release resolution. Every host plans +# its prebuilt against the Unsloth fork; setup.sh/setup.ps1 pass it via +# --published-repo. ggml-org is reachable only via an explicit override. DEFAULT_PUBLISHED_REPO = "unslothai/llama.cpp" DEFAULT_PUBLISHED_TAG = os.environ.get("UNSLOTH_LLAMA_RELEASE_TAG") DEFAULT_PUBLISHED_MANIFEST_ASSET = os.environ.get( @@ -265,6 +266,7 @@ class HostInfo: has_physical_nvidia: bool has_usable_nvidia: bool has_rocm: bool = False + has_intel_gpu: bool = False rocm_gfx_target: str | None = None # (major, minor) from platform.mac_ver(); None off macOS or if unparseable. # Skips a macos prebuilt whose minimum-OS exceeds this host. @@ -1284,162 +1286,6 @@ def synthetic_checksums_for_release( ) -def parse_direct_linux_release_bundle( - repo: str, release: dict[str, Any] -) -> PublishedReleaseBundle | None: - release_tag = release.get("tag_name") - if not isinstance(release_tag, str) or not release_tag: - return None - - assets = release_asset_map(release) - artifacts: list[PublishedLlamaArtifact] = [] - inferred_labels: list[str] = [] - - linux_asset_re = re.compile( - r"^app-(?P