diff --git a/.github/workflows/release-desktop.yml b/.github/workflows/release-desktop.yml index 188e078f90..4daafae35d 100644 --- a/.github/workflows/release-desktop.yml +++ b/.github/workflows/release-desktop.yml @@ -19,6 +19,19 @@ on: permissions: contents: read +env: + DESKTOP_RELEASE_NOTES: | + Desktop app for Unsloth Studio. + + **macOS**: Download the Apple Silicon `.dmg`. + **Windows**: Download the `-setup.exe` installer. + **Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental. + + > Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package. + > Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available. + > Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64` + > First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually. + concurrency: group: release-desktop-${{ github.repository }} cancel-in-progress: false @@ -295,14 +308,6 @@ jobs: PY build: - # TODO: split into a "build (no secrets)" + "publish (secrets)" job pair - # with actions/upload-artifact handoff so the matrix build cannot - # publish a Release on its own. The current matrix runs across - # Linux/macOS/Windows in a single job, so the split needs artefact - # collection across the OS matrix and is out of scope for this - # hardening pass. - permissions: - contents: write # tauri-apps/tauri-action creates / uploads a GitHub Release strategy: fail-fast: false max-parallel: 1 @@ -311,15 +316,21 @@ jobs: - platform: macos-latest args: '--target aarch64-apple-darwin' label: macOS (Apple Silicon) + artifact: macos-aarch64 + release_arch: aarch64 # - platform: macos-latest # args: '--target x86_64-apple-darwin' # label: macOS (Intel) - platform: ubuntu-22.04 args: '' label: Linux (x64) + artifact: linux-x64 + release_arch: x64 - platform: windows-latest args: '' label: Windows (x64) + artifact: windows-x64 + release_arch: x64 name: Build ${{ matrix.label }} needs: prepare-version @@ -465,41 +476,18 @@ jobs: if (chmodIdx !== -1 && sha256Idx > chmodIdx) { throw new Error('Desktop Linux release must verify the linuxdeploy digest before chmod +x'); } - const releaseBodies = []; - for (let i = 0; i < lines.length; i += 1) { - const match = lines[i].match(/^(\s*)releaseBody:\s*\|\s*$/); - if (!match) continue; - const baseIndent = match[1].length; - const bodyLines = []; - i += 1; - for (; i < lines.length; i += 1) { - const line = lines[i]; - if (line.trim() === '') { - bodyLines.push(''); - continue; - } - const indent = line.match(/^\s*/)[0].length; - if (indent <= baseIndent) { - i -= 1; - break; - } - bodyLines.push(line.slice(baseIndent + 2)); - } - releaseBodies.push(bodyLines.join('\n')); + const releaseBody = process.env.DESKTOP_RELEASE_NOTES; + if (!releaseBody) { + throw new Error('DESKTOP_RELEASE_NOTES must not be empty'); } - if (releaseBodies.length === 0) { - throw new Error('Expected at least one desktop release body'); + if (/\brpm\b|\.rpm/i.test(releaseBody)) { + throw new Error('Desktop release body must not advertise RPM packages'); } - for (const body of releaseBodies) { - if (/\brpm\b|\.rpm/i.test(body)) { - throw new Error('Desktop release body must not advertise RPM packages'); - } - if (/AppImage.*universal|universal.*AppImage/i.test(body)) { - throw new Error('Desktop release body must not advertise AppImage as universal'); - } - if (!/AppImage.*experimental/i.test(body)) { - throw new Error('Desktop release body must mark AppImage as experimental'); - } + if (/AppImage.*universal|universal.*AppImage/i.test(releaseBody)) { + throw new Error('Desktop release body must not advertise AppImage as universal'); + } + if (!/AppImage.*experimental/i.test(releaseBody)) { + throw new Error('Desktop release body must mark AppImage as experimental'); } JS @@ -644,48 +632,33 @@ jobs: dest="$tools_dir/linuxdeploy-x86_64.AppImage" curl -fsSL "$LINUXDEPLOY_URL" -o "$dest" # Verify the digest BEFORE the binary is ever marked executable. The - # next step builds the AppImage with the Tauri signing key and a - # contents:write GITHUB_TOKEN in scope, so a substituted linuxdeploy - # that ran here could exfiltrate signing material or tamper with - # published release artifacts. Fail closed on any mismatch. + # next step builds the AppImage with the Tauri signing key, so a + # substituted linuxdeploy that ran here could exfiltrate signing + # material or tamper with release artifacts. Fail closed on any + # mismatch. echo "${LINUXDEPLOY_SHA256} ${dest}" | sha256sum -c - chmod +x "$dest" - # ── Linux: build + sign + upload ── + # ── Linux: build + sign ── - name: Build Linux app + id: build_linux if: matrix.platform == 'ubuntu-22.04' uses: tauri-apps/tauri-action@84b9d35b5fc46c1e45415bdb6144030364f7ebc5 env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} TAURI_SIGNING_PRIVATE_KEY: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY }} TAURI_SIGNING_PRIVATE_KEY_PASSWORD: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY_PASSWORD }} XDG_CACHE_HOME: ${{ runner.temp }}/tauri-tools-cache with: projectPath: studio tauriScript: npx --prefix . tauri - tagName: ${{ needs.prepare-version.outputs.desktop_release_tag }} - releaseName: 'Unsloth Studio (Desktop) ${{ needs.prepare-version.outputs.studio_version }}' - releaseBody: | - Desktop app for Unsloth Studio. - - **macOS**: Download the Apple Silicon `.dmg`. - **Windows**: Download the `-setup.exe` installer. - **Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental. - - > Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package. - > Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available. - > Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64` - > First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually. - releaseDraft: ${{ inputs.draft }} - prerelease: ${{ needs.prepare-version.outputs.prerelease }} args: -v ${{ matrix.args }} - # ── macOS: build + sign + notarize + upload ── + # ── macOS: build + sign + notarize ── - name: Build macOS app + id: build_macos if: matrix.platform == 'macos-latest' uses: tauri-apps/tauri-action@84b9d35b5fc46c1e45415bdb6144030364f7ebc5 env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} TAURI_SIGNING_PRIVATE_KEY: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY }} TAURI_SIGNING_PRIVATE_KEY_PASSWORD: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY_PASSWORD }} APPLE_SIGNING_IDENTITY: ${{ secrets.APPLE_SIGNING_IDENTITY }} @@ -695,29 +668,14 @@ jobs: with: projectPath: studio tauriScript: npx --prefix . tauri - tagName: ${{ needs.prepare-version.outputs.desktop_release_tag }} - releaseName: 'Unsloth Studio (Desktop) ${{ needs.prepare-version.outputs.studio_version }}' - releaseBody: | - Desktop app for Unsloth Studio. - - **macOS**: Download the Apple Silicon `.dmg`. - **Windows**: Download the `-setup.exe` installer. - **Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental. - - > Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package. - > Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available. - > Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64` - > First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually. - releaseDraft: ${{ inputs.draft }} - prerelease: ${{ needs.prepare-version.outputs.prerelease }} args: -v ${{ matrix.args }} - # ── Windows: build + sign + upload ── + # ── Windows: build + sign ── - name: Build Windows app + id: build_windows if: matrix.platform == 'windows-latest' uses: tauri-apps/tauri-action@84b9d35b5fc46c1e45415bdb6144030364f7ebc5 env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} TAURI_SIGNING_PRIVATE_KEY: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY }} TAURI_SIGNING_PRIVATE_KEY_PASSWORD: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY_PASSWORD }} AZURE_CLIENT_ID: ${{ secrets.AZURE_CLIENT_ID }} @@ -728,35 +686,83 @@ jobs: with: projectPath: studio tauriScript: npx --prefix . tauri - tagName: ${{ needs.prepare-version.outputs.desktop_release_tag }} - releaseName: 'Unsloth Studio (Desktop) ${{ needs.prepare-version.outputs.studio_version }}' - releaseBody: | - Desktop app for Unsloth Studio. - - **macOS**: Download the Apple Silicon `.dmg`. - **Windows**: Download the `-setup.exe` installer. - **Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental. - - > Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package. - > Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available. - > Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64` - > First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually. - releaseDraft: ${{ inputs.draft }} - prerelease: ${{ needs.prepare-version.outputs.prerelease }} args: -v ${{ matrix.args }} - # Release process note: only non-draft workflow runs advance the public - # desktop-latest updater channel. Draft builds are for private review; if a - # draft is manually published later, this channel intentionally remains - # unchanged until a narrow manual channel-publish flow is added or a public - # desktop release is created by running this workflow with draft=false. - publish-updater-channel: - name: Publish desktop updater channel + - name: Stage release assets + shell: bash + env: + ARTIFACT_PATHS: ${{ steps.build_linux.outputs.artifactPaths || steps.build_macos.outputs.artifactPaths || steps.build_windows.outputs.artifactPaths }} + RELEASE_ARCH: ${{ matrix.release_arch }} + run: | + set -euo pipefail + if command -v python3 >/dev/null 2>&1; then + PYTHON=python3 + else + PYTHON=python + fi + "$PYTHON" <<'PY' + import json + import os + import pathlib + import re + import shutil + import sys + import unicodedata + + raw_paths = os.environ.get('ARTIFACT_PATHS', '') + try: + artifact_paths = json.loads(raw_paths) + except json.JSONDecodeError as error: + sys.exit(f'Invalid tauri-action artifactPaths output: {error}') + if not isinstance(artifact_paths, list) or not artifact_paths: + sys.exit('tauri-action did not return any release artifacts') + + destination = pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-release-assets') + destination.mkdir(parents=True, exist_ok=True) + staged = [] + for raw_path in artifact_paths: + source = pathlib.Path(raw_path) + if not source.is_file(): + continue + name = source.name + for extension in ('.app.tar.gz.sig', '.app.tar.gz'): + if name.endswith(extension): + name = f'{name[:-len(extension)]}_{os.environ["RELEASE_ARCH"]}{extension}' + break + name = unicodedata.normalize('NFD', name) + name = ''.join(character for character in name if not unicodedata.combining(character)) + name = re.sub(r'[ ()\[\]{}]', '.', name) + while '..' in name: + name = name.replace('..', '.') + target = destination / name + if target.exists(): + sys.exit(f'Duplicate staged release asset name: {name}') + shutil.copy2(source, target) + staged.append(name) + + if not staged: + sys.exit('No release files were staged') + print('Staged release assets:') + print('\n'.join(sorted(staged))) + PY + + - name: Upload signed release assets + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: desktop-release-${{ matrix.artifact }} + path: ${{ runner.temp }}/desktop-release-assets/* + if-no-files-found: error + compression-level: 0 + retention-days: 1 + + # Only this job gets write access; builds hand off signed files via artifacts. + # Draft runs do not advance the public desktop-latest channel. + publish-release: + name: Publish desktop release needs: [prepare-version, build] - if: ${{ !inputs.draft }} runs-on: ubuntu-latest permissions: - contents: write + contents: write # create the versioned Release and replace updater-channel metadata env: GH_REPO: ${{ github.repository }} APP_VERSION: ${{ needs.prepare-version.outputs.app_version }} @@ -765,7 +771,164 @@ jobs: DESKTOP_PRERELEASE: ${{ needs.prepare-version.outputs.prerelease }} steps: + - name: Harden runner (audit) + uses: step-security/harden-runner@a5ad31d6a139d249332a2605b85202e8c0b78450 # v2.19.1 + with: + egress-policy: audit + + - name: Download signed release assets + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + pattern: desktop-release-* + path: ${{ runner.temp }}/desktop-release-assets + merge-multiple: true + + - name: Validate release asset set + shell: bash + run: | + set -euo pipefail + python3 <<'PY' + import pathlib + import os + import sys + + asset_dir = pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-release-assets') + files = [path for path in asset_dir.iterdir() if path.is_file()] + required_suffixes = ( + '.dmg', + '.app.tar.gz', + '.app.tar.gz.sig', + '.deb', + '.AppImage', + '.AppImage.sig', + '-setup.exe', + '-setup.exe.sig', + ) + for suffix in required_suffixes: + matches = [path for path in files if path.name.endswith(suffix)] + if len(matches) != 1: + sys.exit(f'Expected exactly one {suffix} release asset, found {len(matches)}') + if any(path.name == 'latest.json' for path in files): + sys.exit('Build artifacts must not supply latest.json') + print('\n'.join(sorted(path.name for path in files))) + PY + + - name: Create or validate versioned release + shell: bash + env: + GH_TOKEN: ${{ github.token }} + RELEASE_DRAFT: ${{ inputs.draft }} + run: | + set -euo pipefail + notes_file="$RUNNER_TEMP/desktop-release-notes.md" + printf '%s\n' "$DESKTOP_RELEASE_NOTES" > "$notes_file" + + release_json="$RUNNER_TEMP/versioned-release.json" + # REST tag lookup omits drafts; `gh release view` also checks pending tags. + if gh release view "$DESKTOP_RELEASE_TAG" \ + --json tagName,isDraft,isPrerelease > "$release_json" 2>/dev/null; then + python3 <<'PY' + import json + import os + import pathlib + import sys + + release = json.loads(pathlib.Path(os.environ['RUNNER_TEMP'], 'versioned-release.json').read_text()) + expected_draft = os.environ['RELEASE_DRAFT'].lower() == 'true' + expected_prerelease = os.environ['DESKTOP_PRERELEASE'].lower() == 'true' + if release.get('tagName') != os.environ['DESKTOP_RELEASE_TAG']: + sys.exit('Existing desktop release tag does not match the requested tag') + if bool(release.get('isDraft')) != expected_draft: + sys.exit('Existing desktop release draft state does not match the workflow input') + if bool(release.get('isPrerelease')) != expected_prerelease: + sys.exit('Existing desktop release prerelease state does not match the requested version') + PY + else + release_flags=( + --title "Unsloth Studio (Desktop) ${STUDIO_VERSION}" + --notes-file "$notes_file" + --target "$GITHUB_SHA" + ) + if [ "$RELEASE_DRAFT" = "true" ]; then + release_flags+=(--draft) + fi + if [ "$DESKTOP_PRERELEASE" = "true" ]; then + release_flags+=(--prerelease) + fi + gh release create "$DESKTOP_RELEASE_TAG" "${release_flags[@]}" + fi + + - name: Publish versioned release assets + shell: bash + env: + GH_TOKEN: ${{ github.token }} + run: | + set -euo pipefail + gh release upload "$DESKTOP_RELEASE_TAG" "$RUNNER_TEMP/desktop-release-assets"/* --clobber + + - name: Generate and publish versioned updater metadata + shell: bash + env: + GH_TOKEN: ${{ github.token }} + run: | + set -euo pipefail + python3 <<'PY' + import datetime + import json + import os + import pathlib + import sys + import urllib.parse + + asset_dir = pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-release-assets') + files = [path for path in asset_dir.iterdir() if path.is_file()] + + def exactly_one(suffix: str) -> pathlib.Path: + matches = [path for path in files if path.name.endswith(suffix)] + if len(matches) != 1: + sys.exit(f'Expected exactly one {suffix} updater asset, found {len(matches)}') + return matches[0] + + def entry(signature_suffix: str) -> dict[str, str]: + signature_path = exactly_one(signature_suffix) + bundle_name = signature_path.name.removesuffix('.sig') + bundle_path = asset_dir / bundle_name + if not bundle_path.is_file(): + sys.exit(f'Missing updater bundle for {signature_path.name}: {bundle_name}') + encoded_tag = urllib.parse.quote(os.environ['DESKTOP_RELEASE_TAG'], safe='') + encoded_name = urllib.parse.quote(bundle_name, safe='') + return { + 'signature': signature_path.read_text(), + 'url': ( + f'https://github.com/{os.environ["GITHUB_REPOSITORY"]}/releases/download/' + f'{encoded_tag}/{encoded_name}' + ), + } + + darwin = entry('.app.tar.gz.sig') + linux = entry('.AppImage.sig') + windows = entry('.exe.sig') + notes = pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-release-notes.md').read_text() + metadata = { + 'version': os.environ['APP_VERSION'], + 'notes': notes, + 'pub_date': datetime.datetime.now(datetime.timezone.utc).isoformat(timespec='milliseconds').replace('+00:00', 'Z'), + 'platforms': { + 'darwin-aarch64': darwin, + 'darwin-aarch64-app': darwin, + 'linux-x86_64': linux, + 'linux-x86_64-appimage': linux, + 'windows-x86_64': windows, + 'windows-x86_64-nsis': windows, + }, + } + output = pathlib.Path(os.environ['RUNNER_TEMP'], 'latest.json') + output.write_text(json.dumps(metadata, indent=2) + '\n') + PY + gh release upload "$DESKTOP_RELEASE_TAG" "$RUNNER_TEMP/latest.json" --clobber + - name: Download versioned updater metadata + if: ${{ !inputs.draft }} shell: bash env: GH_TOKEN: ${{ github.token }} @@ -790,6 +953,7 @@ jobs: test -s "$RUNNER_TEMP/desktop-updater/latest.json" - name: Validate versioned updater metadata + if: ${{ !inputs.draft }} shell: bash run: | python3 <<'PY' @@ -849,6 +1013,7 @@ jobs: PY - name: Ensure desktop updater channel release + if: ${{ !inputs.draft }} shell: bash env: GH_TOKEN: ${{ github.token }} @@ -881,6 +1046,7 @@ jobs: PY - name: Prevent updater channel downgrade + if: ${{ !inputs.draft }} shell: bash env: GH_TOKEN: ${{ github.token }} @@ -971,6 +1137,7 @@ jobs: PY - name: Publish desktop updater channel metadata + if: ${{ !inputs.draft }} shell: bash env: GH_TOKEN: ${{ github.token }} diff --git a/.github/workflows/studio-inference-smoke.yml b/.github/workflows/studio-inference-smoke.yml index cf8c021e38..58ef2558f3 100644 --- a/.github/workflows/studio-inference-smoke.yml +++ b/.github/workflows/studio-inference-smoke.yml @@ -729,6 +729,7 @@ jobs: content, events = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": prompt}], "enable_tools": True, + "permission_mode": "full", "enabled_tools": enabled, "session_id": f"{session}-att{attempt_i}", "temperature": TOOL_PROBE_TEMP, @@ -818,6 +819,7 @@ jobs: content, events = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": "Search the web for 'unsloth ai github' and summarise."}], "enable_tools": True, + "permission_mode": "full", "enabled_tools": ["web_search"], "session_id": "ci-tool-calling-web", "temperature": 0.0, diff --git a/.github/workflows/studio-mac-inference-smoke.yml b/.github/workflows/studio-mac-inference-smoke.yml index d3d765aa84..946681706a 100644 --- a/.github/workflows/studio-mac-inference-smoke.yml +++ b/.github/workflows/studio-mac-inference-smoke.yml @@ -612,6 +612,7 @@ jobs: content = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": "What is 123 * 456? Use the python tool to compute it and tell me the number."}], "enable_tools": True, + "permission_mode": "full", "enabled_tools": ["python"], "session_id": "ci-tool-calling-py", "temperature": TEMP, @@ -647,6 +648,7 @@ jobs: content = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": "Search the web for 'unsloth ai github' and summarise."}], "enable_tools": True, + "permission_mode": "full", "enabled_tools": ["web_search"], "session_id": "ci-tool-calling-web", "temperature": TEMP, diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index 233292f7a3..63a7e9dc8f 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -791,6 +791,7 @@ jobs: content = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": "What is 123 * 456? Use the python tool to compute it and tell me the number."}], "enable_tools": True, + "permission_mode": "full", "enabled_tools": ["python"], "session_id": "ci-tool-calling-py", "temperature": TEMP, @@ -816,6 +817,7 @@ jobs: content = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": "Use the terminal tool to run `echo hello-bash-tool` and tell me the exact output."}], "enable_tools": True, + "permission_mode": "full", "enabled_tools": ["terminal"], "session_id": "ci-tool-calling-bash", "temperature": TEMP, @@ -840,6 +842,7 @@ jobs: content = post_sse("/v1/chat/completions", { "messages": [{"role": "user", "content": "Search the web for 'unsloth ai github' and summarise."}], "enable_tools": True, + "permission_mode": "full", "enabled_tools": ["web_search"], "session_id": "ci-tool-calling-web", "temperature": TEMP, diff --git a/studio/MCP.md b/studio/MCP.md new file mode 100644 index 0000000000..91b39fcc69 --- /dev/null +++ b/studio/MCP.md @@ -0,0 +1,34 @@ +# Unsloth Studio MCP server + +Studio can expose a local MCP server so an MCP client can inspect models and +GPU state, validate recipes, start or stop training, inspect recipe output, and +export a loaded model. + +The server is disabled by default. Enable it for a local Studio process with: + +```bash +UNSLOTH_STUDIO_ENABLE_MCP=1 \ +UNSLOTH_STUDIO_MCP_TOKEN='use-a-local-secret' \ +unsloth studio +``` + +The endpoint is `http://127.0.0.1:8888/mcp/` when Studio uses its default port +(a request to `/mcp` redirects to the canonical `/mcp/`). Use the actual Studio +port when it is configured differently. + +The high-impact tools are: + +- `studio_status` and `list_local_models` for discovery +- `get_training_status`, `start_training`, `stop_training`, and `list_training_runs` +- `validate_recipe`, `get_recipe_job_status`, and `get_recipe_job_dataset` +- `load_checkpoint` and `export_gguf` + +`start_training` accepts the same fields as the Studio `TrainingStartRequest`. +The request is validated by the existing Pydantic model before a subprocess is +started. Export paths use the existing Studio validation as well. + +The endpoint always requires `UNSLOTH_STUDIO_MCP_TOKEN` and checks an exact +Bearer token for both HTTP and WebSocket connections. Keep it on localhost +unless the deployment has an authenticated reverse proxy. The MCP endpoint is +intentionally opt-in because tools can consume GPU memory, write model +artifacts, and stop active work. \ No newline at end of file diff --git a/studio/backend/core/inference/chat_template_helpers.py b/studio/backend/core/inference/chat_template_helpers.py index 897db8262d..5113eebb36 100644 --- a/studio/backend/core/inference/chat_template_helpers.py +++ b/studio/backend/core/inference/chat_template_helpers.py @@ -10,10 +10,242 @@ native-chat-template fallback used by the transformers and MLX backends. import copy import json import logging +from dataclasses import dataclass from typing import Optional _THINK_OPEN = "" _THINK_CLOSE = "" +_GEMMA_CHANNEL_START = "<|channel>" +_GEMMA_THOUGHT_OPEN = "<|channel>thought" +_GEMMA_THOUGHT_CLOSE = "" +_GEMMA_TEMPLATE_OPENERS = ( + _GEMMA_THOUGHT_OPEN + "\n", + _GEMMA_THOUGHT_OPEN + "\\n", + _GEMMA_THOUGHT_OPEN + _GEMMA_THOUGHT_CLOSE, +) + + +def _tokenizer_objects(tokenizer) -> tuple: + """Return a processor/tokenizer and its distinct nested tokenizer.""" + if tokenizer is None: + return () + nested = getattr(tokenizer, "tokenizer", None) + return (tokenizer,) if nested is None or nested is tokenizer else (tokenizer, nested) + + +def _selected_template_strings_from_value( + template, + tools = None, + *, + prefer_tool_use: bool = True, +) -> tuple[str, ...]: + """Return the named chat template matching HF's default selection rules.""" + tools = tools or None + if isinstance(template, str): + return (template,) + if not isinstance(template, dict): + return () + if prefer_tool_use and tools and isinstance(template.get("tool_use"), str): + return (template["tool_use"],) + if isinstance(template.get("default"), str): + return (template["default"],) + values = tuple(value for value in template.values() if isinstance(value, str)) + return values if len(values) == 1 else () + + +def _selected_chat_template_strings(tokenizer, tools = None) -> tuple[str, ...]: + """Return the active chat template selected for this request.""" + tools = tools or None + getter = getattr(tokenizer, "get_chat_template", None) + if callable(getter): + for kwargs in ({"chat_template": None, "tools": tools}, {"tools": tools}, {}): + try: + selected = getter(**kwargs) + except Exception: + continue + if isinstance(selected, str): + return (selected,) + # ProcessorMixin.apply_chat_template does not switch to "tool_use" implicitly; + # it uses "default" unless chat_template= names another template. + is_processor = getattr(tokenizer, "tokenizer", None) is not None and callable( + getattr(tokenizer, "apply_chat_template", None) + ) + return _selected_template_strings_from_value( + getattr(tokenizer, "chat_template", None), + tools, + prefer_tool_use = not is_processor, + ) + + +def _detect_reasoning_channel_markers_from_templates( + templates: tuple[str, ...], +) -> Optional[tuple[str, str]]: + """Return Gemma native reasoning markers only when a template emits them.""" + if any(opener in template for template in templates for opener in _GEMMA_TEMPLATE_OPENERS): + return _GEMMA_THOUGHT_OPEN, _GEMMA_THOUGHT_CLOSE + return None + + +def detect_reasoning_channel_markers(tokenizer, tools = None) -> Optional[tuple[str, str]]: + """Return native Gemma thought-channel markers supported by a tokenizer. + + Detection uses the active chat template rather than model names or vocabulary + membership. Some models expose Gemma control tokens without using the native + thought-channel response protocol, and those must keep normal + ``skip_special_tokens`` streaming. + """ + for obj in _tokenizer_objects(tokenizer): + templates = _selected_chat_template_strings(obj, tools) + if templates: + return _detect_reasoning_channel_markers_from_templates(templates) + return None + + +def detect_reasoning_channel_markers_from_template( + template, tools = None +) -> Optional[tuple[str, str]]: + """Return native Gemma thought-channel markers from a raw template value.""" + return _detect_reasoning_channel_markers_from_templates( + _selected_template_strings_from_value(template, tools) + ) + + +def detect_reasoning_channel_markers_from_model_info( + tokenizer, + model_info: Optional[dict] = None, + tools = None, +) -> Optional[tuple[str, str]]: + """Return reasoning markers from the active or cached native template.""" + markers = detect_reasoning_channel_markers(tokenizer, tools = tools) + if markers is not None or not isinstance(model_info, dict): + return markers + + native_templates = ( + model_info.get("native_chat_template"), + (model_info.get("chat_template_info") or {}).get("template"), + ) + for template in native_templates: + markers = detect_reasoning_channel_markers_from_template(template, tools) + if markers is not None: + return markers + return None + + +@dataclass(frozen = True) +class ChatTemplateRenderResult: + """Prompt plus response-protocol metadata selected by the renderer.""" + + prompt: str + reasoning_channel_markers: Optional[tuple[str, str]] = None + + +def _split_partial_marker(text: str, marker: str) -> tuple[str, str]: + """Hold the longest suffix that may become ``marker`` in the next chunk.""" + for length in range(min(len(text), len(marker) - 1), 0, -1): + if text.endswith(marker[:length]): + return text[:-length], text[-length:] + return text, "" + + +class ReasoningChannelNormalizer: + """Incrementally convert one native reasoning channel to ````. + + The parser follows mlx-vlm's streaming boundary behavior but emits Studio's + established canonical text contract. Only the configured opening and + closing markers are consumed; tool-call and other control markers remain + available to downstream parsers. + """ + + def __init__(self, opening_marker: str, closing_marker: str): + self._opening_marker = opening_marker + self._closing_marker = closing_marker + self._buffer = "" + self._in_reasoning = False + self._reasoning_done = False + self._skip_opening_newline = False + + def feed(self, text: str) -> str: + """Consume a raw text delta and return the stable canonical delta.""" + self._buffer += text or "" + output: list[str] = [] + while self._buffer: + if self._reasoning_done: + output.append(self._buffer) + self._buffer = "" + break + + if self._in_reasoning and self._skip_opening_newline: + if self._buffer.startswith("\n"): + self._buffer = self._buffer[1:] + self._skip_opening_newline = False + if not self._buffer: + break + + marker = self._closing_marker if self._in_reasoning else self._opening_marker + index = self._buffer.find(marker) + if index < 0: + stable, self._buffer = _split_partial_marker(self._buffer, marker) + output.append(stable) + break + + output.append(self._buffer[:index]) + self._buffer = self._buffer[index + len(marker) :] + if self._in_reasoning: + output.append(_THINK_CLOSE) + self._in_reasoning = False + self._reasoning_done = True + else: + output.append(_THINK_OPEN) + self._in_reasoning = True + self._skip_opening_newline = True + return "".join(output) + + def finish(self) -> str: + """Flush a naturally completed stream and close an open think block.""" + output = self.drain() + if self._in_reasoning: + output += _THINK_CLOSE + self._in_reasoning = False + self._reasoning_done = True + return output + + def drain(self) -> str: + """Flush buffered literal text without synthesizing a closing tag.""" + output = self._buffer + self._buffer = "" + return output + + +def normalize_reasoning_snapshots( + stream, + tokenizer = None, + cancel_event = None, + markers: Optional[tuple[str, str]] = None, + tools = None, +): + """Normalize a prefix-monotonic cumulative text stream when supported.""" + markers = markers or detect_reasoning_channel_markers(tokenizer, tools = tools) + if markers is None: + yield from stream + return + + normalizer = ReasoningChannelNormalizer(*markers) + raw_output = "" + normalized_output = "" + for snapshot in stream: + if not snapshot.startswith(raw_output): + raise RuntimeError("Reasoning normalization requires cumulative text snapshots") + delta = normalizer.feed(snapshot[len(raw_output) :]) + raw_output = snapshot + if delta: + normalized_output += delta + yield normalized_output + + cancelled = cancel_event is not None and cancel_event.is_set() + tail = normalizer.drain() if cancelled else normalizer.finish() + if tail: + normalized_output += tail + yield normalized_output def detect_think_prefill(prompt: Optional[str], special_tokens = None) -> str: @@ -166,7 +398,8 @@ def render_native_template( preserve_thinking: Optional[bool] = None, apply_fn = None, hf_token: Optional[str] = None, -) -> Optional[str]: + return_metadata: bool = False, +): """Render ``messages`` + ``tools`` with the model's NATIVE chat template. Some Unsloth override templates (e.g. ``mistral``, ``gemma-4``) do not emit @@ -175,7 +408,9 @@ def render_native_template( 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``. + differs with vs without tools); otherwise ``None``. With ``return_metadata``, + returns ``ChatTemplateRenderResult`` so callers can stream with the response + protocol selected by this request's template. ``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 @@ -261,7 +496,16 @@ def render_native_template( exc, ) return None - return with_tools if with_tools != no_tools else None + if with_tools == no_tools: + return None + if return_metadata: + return ChatTemplateRenderResult( + with_tools, + _detect_reasoning_channel_markers_from_templates( + _selected_template_strings_from_value(native_tpl, tools) + ), + ) + return with_tools def render_with_native_template_fallback( @@ -277,7 +521,8 @@ def render_with_native_template_fallback( preserve_thinking: Optional[bool] = None, apply_fn = None, hf_token: Optional[str] = None, -) -> str: + return_metadata: bool = False, +): """Return ``formatted_prompt``, swapping in a native-template render when an override template dropped the ``tools`` schema. @@ -285,9 +530,27 @@ def render_with_native_template_fallback( 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.""" + gated/private model's native template can still be fetched. With + ``return_metadata``, returns the selected prompt plus reasoning-channel markers + for the exact template used by this request.""" + live_markers = detect_reasoning_channel_markers(tokenizer, tools = tools) + + def _result(prompt: str, markers = live_markers): + if return_metadata: + return ChatTemplateRenderResult(prompt, markers) + return prompt + if not tools: - return formatted_prompt + # Gemma 4 can emit its native reasoning protocol even when a generation-time + # Unsloth override rendered a marker-free prompt. Preserve the live-verified + # no-tools thinking behavior without letting cached native metadata describe + # unrelated tool prompts that kept the active override. + markers = live_markers + if markers is None: + markers = detect_reasoning_channel_markers_from_model_info( + tokenizer, model_info, tools = None + ) + return _result(formatted_prompt, markers) if apply_fn is None: apply_fn = apply_chat_template_for_generation # Probe whether the live template dropped the schema. A tools-requiring template @@ -307,9 +570,9 @@ def render_with_native_template_fallback( active_model_name, exc, ) - return formatted_prompt + return _result(formatted_prompt) if formatted_prompt != probe_no_tools: - return formatted_prompt # template already emits the tools schema + return _result(formatted_prompt) # template already emits the tools schema native_prompt = render_native_template( model_info = model_info, active_model_name = active_model_name, @@ -320,6 +583,7 @@ def render_with_native_template_fallback( preserve_thinking = preserve_thinking, apply_fn = apply_fn, hf_token = hf_token, + return_metadata = return_metadata, ) if native_prompt: logger.info( @@ -328,4 +592,4 @@ def render_with_native_template_fallback( active_model_name, ) return native_prompt - return formatted_prompt + return _result(formatted_prompt) diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 172e9e5546..8d262bbb0f 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -5,7 +5,7 @@ from unsloth import FastLanguageModel, FastVisionModel from unsloth.chat_templates import get_chat_template -from transformers import TextStreamer +from transformers import TextIteratorStreamer, TextStreamer from peft import PeftModel, PeftModelForCausalLM import json @@ -32,6 +32,11 @@ from core.inference.chat_eos import ( chat_eos_repair, resolve_chat_turn_end_eos_ids_using, ) +from core.inference.chat_template_helpers import ( + ReasoningChannelNormalizer, + detect_reasoning_channel_markers, + detect_think_prefill, +) from core.inference.presence_penalty import _make_presence_penalty_processor from io import StringIO import structlog @@ -187,6 +192,53 @@ class HarmonyTextStreamer: self._queue.put(new_content) +class ReasoningTextIteratorStreamer(TextIteratorStreamer): + """TextIteratorStreamer that preserves native channel tokens until parsed.""" + + def __init__( + self, + tokenizer, + *, + markers: tuple[str, str], + skip_prompt: bool = True, + timeout: float = 0.2, + cancel_event = None, + **decode_kwargs, + ): + decode_kwargs["skip_special_tokens"] = False + super().__init__(tokenizer, skip_prompt = skip_prompt, timeout = timeout, **decode_kwargs) + self._normalizer = ReasoningChannelNormalizer(*markers) + self._cancel_event = cancel_event + self._aborted = False + + def abort(self): + """Mark generation as failed so ``end`` drains without closing.""" + self._aborted = True + + def on_finalized_text( + self, + text: str, + stream_end: bool = False, + ): + """Queue canonical deltas, closing only on natural stream completion.""" + delta = self._normalizer.feed(text) + if delta: + self.text_queue.put(delta, timeout = self.timeout) + + if stream_end: + cancelled = self._aborted or ( + self._cancel_event is not None and self._cancel_event.is_set() + ) + tail = self._normalizer.drain() if cancelled else self._normalizer.finish() + if tail: + self.text_queue.put(tail, timeout = self.timeout) + self.text_queue.put(self.stop_signal, timeout = self.timeout) + + +class _GenerationThreadError(RuntimeError): + """Generation worker failures that should propagate through stream routes.""" + + class InferenceBackend: """Unified inference backend supporting text, vision, and LoRA models""" @@ -836,6 +888,7 @@ class InferenceBackend: thread_id: Optional[str] = None, rag_scope: Optional[dict] = None, presence_penalty: float = 0.0, + reasoning_prefilled: bool = False, ): """Run an agentic tool loop on top of ``generate_chat_response``. @@ -889,6 +942,7 @@ class InferenceBackend: session_id = session_id, thread_id = thread_id, rag_scope = rag_scope, + reasoning_prefilled = reasoning_prefilled, ) def generate_chat_response( @@ -960,8 +1014,7 @@ class InferenceBackend: thread can toggle adapters under the generation lock. """ if not self.active_model_name: - yield "Error: No active model" - return + raise RuntimeError("No active model") model_info = self.models[self.active_model_name] is_vision = model_info.get("is_vision", False) @@ -1049,6 +1102,7 @@ class InferenceBackend: template_messages = [{"role": "system", "content": system_prompt}] + messages else: template_messages = messages + reasoning_channel_markers_resolved = False try: if not (hasattr(tokenizer, "chat_template") and tokenizer.chat_template): raise ValueError( @@ -1058,6 +1112,7 @@ class InferenceBackend: f"Please use a model that includes a chat template, or manually set " f"one via tokenizer.chat_template before inference." ) + reasoning_channel_markers = None formatted_prompt = self._apply_chat_template_for_generation( tokenizer, template_messages, @@ -1073,7 +1128,7 @@ class InferenceBackend: render_with_native_template_fallback, ) - formatted_prompt = render_with_native_template_fallback( + render_result = render_with_native_template_fallback( formatted_prompt = formatted_prompt, tokenizer = tokenizer, model_info = model_info, @@ -1085,13 +1140,19 @@ class InferenceBackend: preserve_thinking = preserve_thinking, apply_fn = self._apply_chat_template_for_generation, hf_token = model_info.get("hf_token"), + return_metadata = True, ) + formatted_prompt = render_result.prompt + reasoning_channel_markers = render_result.reasoning_channel_markers + reasoning_channel_markers_resolved = True logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...") except Exception as e: logger.error(f"Error applying chat template: {e}") # Fall back to manual formatting formatted_prompt = self.format_chat_prompt(messages, system_prompt) + reasoning_channel_markers = None + reasoning_channel_markers_resolved = True # Step 3: generate yield from self.generate_stream( @@ -1105,6 +1166,8 @@ class InferenceBackend: cancel_event = cancel_event, _adapter_state = _adapter_state, presence_penalty = presence_penalty, + reasoning_channel_markers = reasoning_channel_markers, + reasoning_channel_markers_resolved = reasoning_channel_markers_resolved, ) def _generate_vision_response( @@ -1190,21 +1253,27 @@ class InferenceBackend: # 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 - streamer = TextIteratorStreamer( + streamer = self._make_text_streamer( raw_tokenizer, + protocol_source = processor, + # The text-only VLM fallback above did not render with the + # processor template, so its native markers do not describe + # this request's response protocol. + reasoning_channel_markers = detect_reasoning_channel_markers(processor) + if image + else None, + reasoning_channel_markers_resolved = True, skip_prompt = True, - skip_special_tokens = True, timeout = 0.2, + cancel_event = cancel_event, + use_harmony = self._is_gpt_oss_model(), ) generation_kwargs = dict( @@ -1226,6 +1295,10 @@ class InferenceBackend: ) if _pp is not None: generation_kwargs["logits_processor"] = _pp + stopping_criteria = self._cancel_stopping_criteria(cancel_event) + if stopping_criteria is not None: + generation_kwargs["stopping_criteria"] = stopping_criteria + active_stop_token_ids = self._generation_stop_token_ids(model, generation_kwargs) err: dict[str, str] = {} @@ -1235,6 +1308,8 @@ class InferenceBackend: model.generate(**generation_kwargs) except Exception as e: err["msg"] = str(e) + if hasattr(streamer, "abort"): + streamer.abort() logger.error(f"Vision generation error in thread: {e}") finally: try: @@ -1251,12 +1326,17 @@ class InferenceBackend: if think_prefix: yield think_prefix from queue import Empty + import time generation_complete = False + cancel_deadline = None try: while True: if cancel_event is not None and cancel_event.is_set(): - break + if cancel_deadline is None: + cancel_deadline = time.monotonic() + 10 + elif time.monotonic() >= cancel_deadline: + break try: new_token = next(streamer) except StopIteration: @@ -1265,27 +1345,48 @@ class InferenceBackend: except Empty: if not thread.is_alive(): generation_complete = True + output = yield from self._drain_streamer_tail( + streamer, output, active_stop_token_ids + ) + break + if cancel_deadline is not None: + remaining = cancel_deadline - time.monotonic() + if remaining <= 0: + break + thread.join(timeout = remaining) + if thread.is_alive(): + break + generation_complete = True + output = yield from self._drain_streamer_tail( + streamer, output, active_stop_token_ids + ) break continue if new_token: - output += new_token - cleaned = self._clean_generated_text(output) + output, cleaned = self._append_stream_delta( + output, new_token, active_stop_token_ids + ) yield cleaned finally: if cancel_event is not None and not generation_complete: cancel_event.set() - thread.join(timeout = 10) + join_timeout = 10 + if cancel_deadline is not None: + join_timeout = max(0, cancel_deadline - time.monotonic()) + thread.join(timeout = join_timeout) if thread.is_alive(): logger.warning( "Vision generation thread did not exit after cancel/join timeout" ) if err.get("msg"): - yield f"Error: {err['msg']}" + raise _GenerationThreadError(err["msg"]) + except _GenerationThreadError: + raise except Exception as e: logger.error(f"Vision generation error: {e}") - yield f"Error: {str(e)}" + raise def generate_audio_input_response( self, @@ -1410,11 +1511,13 @@ class InferenceBackend: ) if err.get("msg"): - yield f"Error: {err['msg']}" + raise _GenerationThreadError(err["msg"]) + except _GenerationThreadError: + raise except Exception as e: logger.error(f"Audio input generation error: {e}") - yield f"Error: {str(e)}" + raise def generate_whisper_response( self, @@ -1447,6 +1550,86 @@ class InferenceBackend: from utils.datasets import is_gpt_oss_model_name return is_gpt_oss_model_name(model_name or self.active_model_name or "") + def _make_text_streamer( + self, + tokenizer, + *, + protocol_source = None, + reasoning_channel_markers = None, + reasoning_channel_markers_resolved: bool = False, + skip_prompt: bool = True, + timeout: float = 0.2, + cancel_event = None, + use_harmony: bool = False, + ): + """Create the streamer matching this model's native response protocol.""" + if use_harmony: + try: + return HarmonyTextStreamer( + tokenizer, + skip_prompt = skip_prompt, + timeout = timeout, + ) + except Exception as e: + logger.warning(f"HarmonyTextStreamer init failed, falling back: {e}") + return TextIteratorStreamer( + tokenizer, + skip_prompt = skip_prompt, + skip_special_tokens = True, + timeout = timeout, + ) + + markers = ( + reasoning_channel_markers + if reasoning_channel_markers_resolved + else reasoning_channel_markers + or detect_reasoning_channel_markers(protocol_source or tokenizer) + ) + if markers is not None: + return ReasoningTextIteratorStreamer( + tokenizer, + markers = markers, + skip_prompt = skip_prompt, + timeout = timeout, + cancel_event = cancel_event, + ) + return TextIteratorStreamer( + tokenizer, + skip_prompt = skip_prompt, + skip_special_tokens = True, + timeout = timeout, + ) + + def _append_stream_delta( + self, + output: str, + new_token: str, + stop_token_ids = None, + ): + """Append a streamer delta and apply response-boundary cleanup.""" + output += new_token + return output, self._clean_generated_text(output, stop_token_ids = stop_token_ids) + + def _drain_streamer_tail( + self, + streamer, + output: str, + stop_token_ids = None, + ): + """Drain queued streamer text after the producer exits.""" + while True: + try: + new_token = next(streamer) + except StopIteration: + return output + except Exception: + return output + if new_token: + output, cleaned = self._append_stream_delta( + output, new_token, stop_token_ids = stop_token_ids + ) + yield cleaned + def generate_stream( self, prompt: str, @@ -1459,6 +1642,8 @@ class InferenceBackend: cancel_event = None, _adapter_state = None, presence_penalty: float = 0.0, + reasoning_channel_markers = None, + reasoning_channel_markers_resolved: bool = False, ) -> Generator[str, None, None]: """Generate a streaming text response (text models only). @@ -1467,8 +1652,7 @@ class InferenceBackend: ``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" - return + raise RuntimeError("No active model") model_info = self.models[self.active_model_name] model = model_info["model"] @@ -1481,9 +1665,7 @@ class InferenceBackend: try: inputs = tokenizer(prompt, return_tensors = "pt").to(model.device) - 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. @@ -1494,30 +1676,16 @@ class InferenceBackend: else detect_think_prefill(prompt, getattr(tokenizer, "all_special_tokens", None)) ) - # gpt-oss models: HarmonyTextStreamer parses the multi-channel - # harmony protocol into tags - if self._is_gpt_oss_model(): - try: - streamer = HarmonyTextStreamer( - tokenizer, - skip_prompt = True, - timeout = 0.2, - ) - except Exception as e: - logger.warning(f"HarmonyTextStreamer init failed, falling back: {e}") - streamer = TextIteratorStreamer( - tokenizer, - skip_prompt = True, - skip_special_tokens = True, - timeout = 0.2, - ) - else: - streamer = TextIteratorStreamer( - tokenizer, - skip_prompt = True, - skip_special_tokens = True, - timeout = 0.2, - ) + streamer = self._make_text_streamer( + tokenizer, + protocol_source = model_info.get("tokenizer"), + reasoning_channel_markers = reasoning_channel_markers, + reasoning_channel_markers_resolved = reasoning_channel_markers_resolved, + skip_prompt = True, + timeout = 0.2, + cancel_event = cancel_event, + use_harmony = self._is_gpt_oss_model(), + ) generation_kwargs = dict( **inputs, @@ -1535,27 +1703,16 @@ class InferenceBackend: if tokenizer.pad_token_id is None else tokenizer.pad_token_id, ) + active_stop_token_ids = self._generation_stop_token_ids(model, generation_kwargs) # 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, - StoppingCriteriaList, - ) - class _CancelCriteria(StoppingCriteria): - def __init__(self, ev): - self.ev = ev - - def __call__(self, input_ids, scores, **kwargs): - return self.ev.is_set() - - generation_kwargs["stopping_criteria"] = StoppingCriteriaList( - [_CancelCriteria(cancel_event)] - ) + stopping_criteria = self._cancel_stopping_criteria(cancel_event) + if stopping_criteria is not None: + generation_kwargs["stopping_criteria"] = stopping_criteria def generate_fn(): with self._generation_lock: @@ -1565,6 +1722,8 @@ class InferenceBackend: model.generate(**generation_kwargs) except Exception as e: err["msg"] = str(e) + if hasattr(streamer, "abort"): + streamer.abort() logger.error(f"Generation error: {e}") finally: try: @@ -1582,12 +1741,17 @@ class InferenceBackend: if think_prefix: yield think_prefix from queue import Empty + import time generation_complete = False + cancel_deadline = None try: while True: if cancel_event is not None and cancel_event.is_set(): - break + if cancel_deadline is None: + cancel_deadline = time.monotonic() + 10 + elif time.monotonic() >= cancel_deadline: + break try: new_token = next(streamer) except StopIteration: @@ -1596,11 +1760,27 @@ class InferenceBackend: except Empty: if not thread.is_alive(): generation_complete = True + output = yield from self._drain_streamer_tail( + streamer, output, active_stop_token_ids + ) + break + if cancel_deadline is not None: + remaining = cancel_deadline - time.monotonic() + if remaining <= 0: + break + thread.join(timeout = remaining) + if thread.is_alive(): + break + generation_complete = True + output = yield from self._drain_streamer_tail( + streamer, output, active_stop_token_ids + ) break continue if new_token: - output += new_token - cleaned = self._clean_generated_text(output) + output, cleaned = self._append_stream_delta( + output, new_token, active_stop_token_ids + ) yield cleaned finally: # Set cancel_event only on early exit (user cancel), NOT on @@ -1609,16 +1789,21 @@ class InferenceBackend: # disrupt the next serialized request (e.g. compare mode). if cancel_event is not None and not generation_complete: cancel_event.set() - thread.join(timeout = 10) + join_timeout = 10 + if cancel_deadline is not None: + join_timeout = max(0, cancel_deadline - time.monotonic()) + thread.join(timeout = join_timeout) if thread.is_alive(): logger.warning("Generation thread did not exit after cancel/join timeout") if err.get("msg"): - yield f"Error: {err['msg']}" + raise _GenerationThreadError(err["msg"]) + except _GenerationThreadError: + raise except Exception as e: logger.error(f"Error during generation: {e}") - yield f"Error: {str(e)}" + raise # ── Audio (TTS) Generation ──────────────────────────────────── @@ -2107,8 +2292,42 @@ class InferenceBackend: return img.resize(new_size, Image.Resampling.LANCZOS) return img - def _clean_generated_text(self, text: str) -> str: - """Strip leaked special tokens using the tokenizer's own token list.""" + def _generation_stop_token_ids(self, model, generation_kwargs: dict): + """Return the stop-token ids active for a ``generate`` call.""" + if "eos_token_id" in generation_kwargs: + return generation_kwargs.get("eos_token_id") + generation_config = getattr(model, "generation_config", None) + eos_token_id = getattr(generation_config, "eos_token_id", None) + if eos_token_id is not None: + return eos_token_id + config = getattr(model, "config", None) + return getattr(config, "eos_token_id", None) + + def _cancel_stopping_criteria(self, cancel_event): + """Build a Transformers stopping criteria list for user cancellation.""" + if cancel_event is None: + return None + from transformers.generation.stopping_criteria import ( + StoppingCriteria, + StoppingCriteriaList, + ) + + class _CancelCriteria(StoppingCriteria): + def __init__(self, ev): + self.ev = ev + + def __call__(self, input_ids, scores, **kwargs): + return self.ev.is_set() + + return StoppingCriteriaList([_CancelCriteria(cancel_event)]) + + def _clean_generated_text( + self, + text: str, + *, + stop_token_ids = None, + ) -> str: + """Strip leaked response-boundary tokens after streaming.""" if self._is_gpt_oss_model(): # HarmonyTextStreamer emits clean .... Strip any # harmony protocol tokens and other gpt-oss tokens (e.g. @@ -2118,10 +2337,28 @@ class InferenceBackend: return text.strip() tokenizer = self.models.get(self.active_model_name, {}).get("tokenizer") + tokenizer = getattr(tokenizer, "tokenizer", tokenizer) if tokenizer: - for token in getattr(tokenizer, "all_special_tokens", []): - if token in text: - text = text.replace(token, "") + if stop_token_ids is None: + stop_token_ids = self.models.get(self.active_model_name, {}).get( + "chat_turn_end_eos_ids" + ) + if isinstance(stop_token_ids, int): + stop_token_ids = (stop_token_ids,) + for token_id in stop_token_ids or (): + try: + token = tokenizer.convert_ids_to_tokens(int(token_id)) + except Exception: + token = None + if isinstance(token, str) and token and text.endswith(token): + text = text[: -len(token)] + elif ( + isinstance(token, str) + and token + and text.endswith("") + and text[: -len("")].endswith(token) + ): + text = text[: -len("") - len(token)] + "" return text.strip() def _load_chat_template_info(self, model_name: str): diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index e797cfe22a..fec5214cf9 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -9,6 +9,7 @@ OpenAI-compatible /v1/chat/completions endpoint. import atexit import contextlib +import functools import json import os import re @@ -82,6 +83,7 @@ from core.inference.tool_call_parser import ( ) from core.inference.tool_loop_controller import ( ToolLoopController, + append_deferred_nudges, tool_event_provenance, ) from state.tool_approvals import ( @@ -998,6 +1000,283 @@ def _cached_colocated_split_main( return None +def _cached_variant_resolution(repo_id: str, hf_variant: str) -> tuple[Optional[str], list[str]]: + """Find a cached main GGUF and its shards for a variant.""" + candidate = next(_cached_variant_candidates(repo_id, hf_variant), None) + if candidate is None: + return None, [] + _, main, shards, _ = candidate + return main, shards + + +def _cached_variant_candidates( + repo_id: str, + hf_variant: str, + *, + require_mmproj: bool = False, +) -> Generator[tuple[str, str, list[str], Path], None, None]: + """Yield complete cached variant copies in snapshot preference order.""" + try: + from utils.models.model_config import _iter_hf_cache_snapshots + for snap in _iter_hf_cache_snapshots(repo_id): + cached_files = _gguf_snapshot_files(snap) + matches = _gguf_files_for_variant(cached_files, hf_variant) + if not matches: + continue + main = matches[0] + shards = _gguf_extra_shards(matches, main) + split = _SHARD_FULL_RE.match(main) + if split: + numbers = { + int(match.group(2)) + for path in [main, *shards] + if (match := _SHARD_FULL_RE.match(path)) + } + if numbers != set(range(1, int(split.group(3)) + 1)): + continue + main_path = snap.joinpath(*main.replace("\\", "/").split("/")) + if not main_path.is_file() or not _snapshot_has_all_shards( + str(main_path), main, shards, {} + ): + continue + if require_mmproj and not _pick_mmproj(cached_files): + continue + yield str(main_path), main, shards, snap + except Exception as e: + logger.debug(f"Cache lookup for variant failed: {e}") + + +def _cached_candidate_matches_revision_size( + repo_id: str, candidate: tuple[str, str, list[str], Path], hf_token: Optional[str] +) -> bool: + """Check cached byte sizes against the snapshot's own Hub revision. + + A snapshot pointer is normally published only after its blob is complete. + When the old revision is still queryable, also compare every weight file's + size so a manually truncated cache entry is not treated as reusable. If + metadata cannot be reached, retain the cache's normal offline semantics. + """ + main_path, main, shards, snap = candidate + paths = [main, *shards] + try: + from huggingface_hub import get_paths_info + infos = list( + get_paths_info( + repo_id, + paths, + revision = snap.name, + token = hf_token, + ) + ) + except Exception as e: + logger.debug( + "Could not size-check cached GGUF %s at revision %s: %s", + repo_id, + snap.name, + e, + ) + return True + + if not infos: + # The Hub answers an unknown (e.g. force-pushed away) revision with an + # empty result, not an error; treat it like unreachable metadata. + return True + expected_sizes = {info.path: info.size for info in infos if info.size is not None} + if any(path not in expected_sizes for path in paths): + return False + try: + if os.path.getsize(main_path) < expected_sizes[main]: + return False + except OSError: + return False + return _snapshot_has_all_shards(main_path, main, shards, expected_sizes) + + +def _cached_complete_candidate( + repo_id: str, gguf_filename: Optional[str], shards: list[str] +) -> Optional[tuple[str, str, list[str], Path]]: + """Return one complete exact-filename cache candidate with snapshot context.""" + if not gguf_filename: + return None + if shards: + main_path = _cached_colocated_split_main(repo_id, gguf_filename, shards, {}) + else: + m = _SHARD_FULL_RE.match(gguf_filename) + if m and int(m.group(3)) > 1: + return None + main_path = _cached_hf_snapshot_file(repo_id, gguf_filename) + if main_path is None: + return None + snap = _snapshot_dir_of(main_path) + if snap is None: + return None + return main_path, gguf_filename, shards, snap + + +def cached_gguf_for_load( + hf_repo: str, + hf_variant: Optional[str], + *, + require_mmproj: bool = False, + verify_sizes: bool = False, + hf_token: Optional[str] = None, +) -> Optional[str]: + """Return a cached GGUF that can be loaded without downloading.""" + if not hf_variant: + return None + hf_repo = _resolve_repo_id_casing(hf_repo) + for candidate in _cached_variant_candidates( + hf_repo, + hf_variant, + require_mmproj = require_mmproj, + ): + if verify_sizes and not _cached_candidate_matches_revision_size( + hf_repo, candidate, hf_token + ): + continue + return candidate[0] + return None + + +def _snapshot_dir_of(path: str) -> Optional[Path]: + """Return the HF cache snapshot containing path, if any.""" + try: + p = Path(os.path.abspath(path)) + except OSError: + return None + for ancestor in p.parents: + if ancestor.parent.name == "snapshots": + return ancestor + return None + + +def _companion_snapshot_sibling( + near_path: str, pick: Callable[[list[str]], Optional[str]] +) -> Optional[str]: + """Find a companion in the same snapshot as near_path.""" + snap = _snapshot_dir_of(near_path) + if snap is None: + return None + try: + sibling = pick(_gguf_snapshot_files(snap)) + except Exception: + return None + if not sibling: + return None + candidate = snap / sibling + return str(candidate) if candidate.is_file() else None + + +def _pick_mmproj(candidates: list[str]) -> Optional[str]: + mmproj_files = sorted( + f for f in candidates if f.lower().endswith(".gguf") and "mmproj" in Path(f).name.lower() + ) + if not mmproj_files: + return None + return next((f for f in mmproj_files if f.lower().endswith("-f16.gguf")), mmproj_files[0]) + + +def _hub_download_in_flight(hf_repo: str) -> bool: + try: + from hub.utils.download_registry import get_models_registry + return bool(get_models_registry().active_job_refs(hf_repo)) + except Exception: + return False + + +def _hub_download_blocks_gguf_load( + hf_repo: str, + hf_variant: Optional[str], + *, + require_mmproj: bool = False, + hf_token: Optional[str] = None, +) -> bool: + """Whether an active Hub job makes this GGUF load unsafe. + + Same-variant jobs can reclaim the stale snapshot a load would reuse, so + they always block. Other jobs block only when this load lacks a complete + cached copy and would write to the shared cache itself. + """ + try: + from hub.utils.download_registry import get_models_registry + + registry = get_models_registry() + if not registry.active_job_refs(hf_repo): + return False + if registry.has_active_variant(hf_repo, hf_variant): + return True + except Exception: + return False + return ( + cached_gguf_for_load( + hf_repo, + hf_variant, + require_mmproj = require_mmproj, + verify_sizes = True, + hf_token = hf_token, + ) + is None + ) + + +# Active GGUF loads by normalized repo ID. +_LOADS_IN_FLIGHT: dict[str, int] = {} +_LOADS_IN_FLIGHT_LOCK = threading.Lock() + + +@contextlib.contextmanager +def gguf_load_in_flight(hf_repo: Optional[str]): + """Track an HF GGUF load until the context exits.""" + key = (hf_repo or "").strip().lower() + if not key: + yield + return + with _LOADS_IN_FLIGHT_LOCK: + _LOADS_IN_FLIGHT[key] = _LOADS_IN_FLIGHT.get(key, 0) + 1 + try: + yield + finally: + with _LOADS_IN_FLIGHT_LOCK: + remaining = _LOADS_IN_FLIGHT.get(key, 1) - 1 + if remaining <= 0: + _LOADS_IN_FLIGHT.pop(key, None) + else: + _LOADS_IN_FLIGHT[key] = remaining + + +def hf_gguf_load_in_flight(hf_repo: str) -> bool: + """Return whether a GGUF load is active for hf_repo.""" + key = (hf_repo or "").strip().lower() + if not key: + return False + with _LOADS_IN_FLIGHT_LOCK: + return _LOADS_IN_FLIGHT.get(key, 0) > 0 + + +def _with_gguf_load_marker(load: Callable): + """Keep an HF repo marked for the full synchronous load call.""" + + @functools.wraps(load) + def wrapped(self, *args, **kwargs): + hf_repo = kwargs.get("hf_repo") + with gguf_load_in_flight(hf_repo): + if hf_repo and _hub_download_blocks_gguf_load( + hf_repo, + kwargs.get("hf_variant"), + require_mmproj = bool( + kwargs.get("is_vision") + and not extra_args_disable_mmproj(kwargs.get("extra_args")) + ), + hf_token = kwargs.get("hf_token"), + ): + raise RuntimeError( + f"'{hf_repo}' is currently being downloaded by the download manager" + ) + return load(self, *args, **kwargs) + + return wrapped + + def _gguf_extra_shards(files: Iterable[str], first_shard: str) -> list[str]: m = _SHARD_FULL_RE.match(first_shard) if not m: @@ -1541,6 +1820,31 @@ def _is_external_link(path: Path) -> bool: return False +# Inkling's template takes a numeric thinking-effort dial (0..0.99) and its +# float() coercion turns unrecognized named levels into 0, i.e. no thinking. +# Map OpenAI-style names to the values the model was trained on. Module-level +# so duck-typed engine stand-ins in tests do not need the attribute. +_INKLING_REASONING_EFFORT = { + "none": 0.0, + "minimal": 0.1, + "low": 0.2, + "medium": 0.7, + "high": 0.9, + "xhigh": 0.99, + "max": 0.99, +} + + +def _coerce_reasoning_effort(architecture, kwargs: dict) -> dict: + if architecture == "inkling": + effort = kwargs.get("reasoning_effort") + if isinstance(effort, str): + mapped = _INKLING_REASONING_EFFORT.get(effort.strip().lower()) + if mapped is not None: + kwargs["reasoning_effort"] = mapped + return kwargs + + class LlamaCppBackend: """Manages a llama-server subprocess for GGUF model inference. @@ -1941,36 +2245,15 @@ class LlamaCppBackend: def reasoning_default(self) -> bool: return self._reasoning_default - # Inkling's template takes a numeric thinking-effort dial (0..0.99) and its - # float() coercion turns unrecognized named levels into 0, i.e. no thinking. - # Map OpenAI-style names to the values the model was trained on. - _INKLING_REASONING_EFFORT = { - "none": 0.0, - "minimal": 0.2, - "low": 0.2, - "medium": 0.7, - "high": 0.9, - "xhigh": 0.99, - "max": 0.99, - } - - def _coerce_reasoning_effort(self, kwargs: dict) -> dict: - if getattr(self, "_architecture", None) == "inkling": - effort = kwargs.get("reasoning_effort") - if isinstance(effort, str): - mapped = self._INKLING_REASONING_EFFORT.get(effort.strip().lower()) - if mapped is not None: - kwargs["reasoning_effort"] = mapped - return kwargs - def _reasoning_kwargs(self, enable_thinking: bool) -> dict: if self._reasoning_style == "enable_thinking_effort": # GLM-5.2-style: enable_thinking is the on/off gate; when on, leave # the template's default effort (max) in place. return {"enable_thinking": enable_thinking} if self._reasoning_style == "reasoning_effort": - return self._coerce_reasoning_effort( - {"reasoning_effort": "high" if enable_thinking else "low"} + return _coerce_reasoning_effort( + getattr(self, "_architecture", None), + {"reasoning_effort": "high" if enable_thinking else "low"}, ) return {"enable_thinking": enable_thinking} @@ -2019,7 +2302,7 @@ class LlamaCppBackend: kwargs["enable_thinking"] = enable_thinking if self._supports_preserve_thinking and preserve_thinking is not None: kwargs["preserve_thinking"] = preserve_thinking - self._coerce_reasoning_effort(kwargs) + _coerce_reasoning_effort(getattr(self, "_architecture", None), kwargs) return kwargs or None @property @@ -3801,13 +4084,13 @@ class LlamaCppBackend: hf_repo: str, free_bytes: int, hf_token: Optional[str] = None, - ) -> Optional[tuple[str, int]]: + ) -> Optional[tuple[str, int, list[str]]]: """Find the smallest GGUF variant (including all shards) that fits. Groups split shards by variant prefix and sums their sizes (e.g. UD-Q4_K_XL with 9 shards of 50 GB each = 450 GB total). - Returns (first_shard_filename, total_size_bytes) or None. + Returns (first_shard_filename, total_size_bytes, extra_shards) or None. """ try: from huggingface_hub import get_paths_info, list_repo_files @@ -3843,9 +4126,13 @@ class LlamaCppBackend: # Smallest that fits variant_sizes.sort(key = lambda x: x[1]) - for first_file, total_size, _ in variant_sizes: + for first_file, total_size, shard_files in variant_sizes: if total_size > 0 and total_size <= free_bytes: - return first_file, total_size + return ( + first_file, + total_size, + [path for path in sorted(shard_files) if path != first_file], + ) return None except Exception: @@ -4441,42 +4728,52 @@ class LlamaCppBackend: except Exception as e: logger.warning(f"Could not list repo files: {e}") - # Offline: resolve variant -> filename from the local HF cache. - # The heuristic below assumes filenames echo the repo name, which - # breaks for e.g. Qwen3.6-27B-MTP-GGUF (no "MTP" in file). Match - # against the rel path (not just basename) so subdir layouts like - # ``BF16/foo.gguf`` are findable. + # Fall back to the local cache when the repo listing is unavailable. if not gguf_filename: - try: - from utils.models.model_config import _iter_hf_cache_snapshots - for snap in _iter_hf_cache_snapshots(hf_repo): - cached_files = _gguf_snapshot_files(snap) - matches = _gguf_files_for_variant(cached_files, hf_variant) - if not matches: - continue - gguf_filename = matches[0] - gguf_extra_shards = _gguf_extra_shards(matches, gguf_filename) - logger.info( - "Resolved variant %s -> %s from local HF cache", - hf_variant, - gguf_filename, - ) - break - except Exception as e: - logger.debug(f"Offline cache lookup for variant failed: {e}") + cached_name, cached_shards = _cached_variant_resolution(hf_repo, hf_variant) + if cached_name: + gguf_filename = cached_name + gguf_extra_shards = cached_shards + logger.info( + "Resolved variant %s -> %s from local HF cache", + hf_variant, + gguf_filename, + ) if not gguf_filename: repo_name = hf_repo.split("/")[-1].replace("-GGUF", "") gguf_filename = f"{repo_name}-{hf_variant}.gguf" + # Prefer the existing model. Updates use force=True to fetch a new revision. + if not force: + if hf_variant: + # Resolve by variant so a newer revision's filename does not hide + # the complete older copy. Size-check against that older snapshot's + # own revision when its metadata remains available. + cached_main = cached_gguf_for_load( + hf_repo, + hf_variant, + verify_sizes = True, + hf_token = hf_token, + ) + else: + candidate = _cached_complete_candidate(hf_repo, gguf_filename, gguf_extra_shards) + cached_main = ( + candidate[0] + if candidate is not None + and _cached_candidate_matches_revision_size(hf_repo, candidate, hf_token) + else None + ) + if cached_main is not None: + logger.info(f"Reusing cached GGUF: {cached_main}") + return cached_main + # 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 @@ -4485,25 +4782,10 @@ class LlamaCppBackend: # cold whenever free disk is below the full weight footprint, # even though nothing needs downloading. already_cached_bytes = 0 - # 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. + # Count only files that can resume this download. 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 + # Offline split sets are reusable only when every shard shares a snapshot. + split_needs_refetch = bool(offline and not force and gguf_extra_shards) if not force and not split_needs_refetch: for p in path_infos: if not p.size: @@ -4564,32 +4846,26 @@ class LlamaCppBackend: hf_token, ) if smaller: - fallback_file, fallback_size = smaller + fallback_file, fallback_size, fallback_shards = smaller logger.info( f"Selected variant too large ({total_gb:.1f} GB), " f"falling back to {fallback_file} ({fallback_size / (1024**3):.1f} GB)" ) gguf_filename = fallback_file - _m = _SHARD_RE.match(gguf_filename) - _prefix = _m.group(1) if _m else None - if _prefix: - prefix_lower = _prefix.lower() - gguf_extra_shards = sorted( - f - for f in all_gguf_files - if f.lower().startswith(prefix_lower) - and f != gguf_filename - and not _is_companion_gguf_path(f) + gguf_extra_shards = fallback_shards + + # The selected fallback is a new load target. Apply the + # same any-revision reuse policy before starting a fetch. + fallback_candidate = _cached_complete_candidate( + hf_repo, gguf_filename, gguf_extra_shards + ) + if fallback_candidate is not None and ( + _cached_candidate_matches_revision_size( + hf_repo, fallback_candidate, hf_token ) - 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 + ): + logger.info(f"Reusing cached fallback GGUF: {fallback_candidate[0]}") + return fallback_candidate[0] else: raise RuntimeError( f"Not enough disk space to download any variant. " @@ -4609,45 +4885,25 @@ class LlamaCppBackend: raise RuntimeError("Cancelled") dl_start = time.monotonic() # Xet primary, HTTP fallback on stall; per-file so finished shards stay cached. - 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( + 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( hf_repo, - gguf_filename, + shard, 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 @@ -4670,10 +4926,12 @@ class LlamaCppBackend: pick: Callable[[list[str]], Optional[str]], label: str, cancel_event: Optional[threading.Event] = None, + near_path: Optional[str] = None, ) -> Optional[str]: """Resolve and fetch a companion GGUF (mmproj / MTP drafter) by name. - Tries the live repo file list, then the local HF cache snapshots + Prefers a companion co-located with ``near_path``'s cache snapshot, + then tries the live repo file list, then the local HF cache snapshots (offline, same fallback as _download_gguf), then hf_hub_download. Runs WITHOUT self._lock (like _download_gguf); honors _cancel_event so an /unload between the main download and here skips the fetch. @@ -4683,6 +4941,17 @@ class LlamaCppBackend: if cancel_event.is_set(): return None + # Keep companion files in the main GGUF's snapshot. + if near_path: + cached = _companion_snapshot_sibling(near_path, pick) + if cached: + logger.info("Reusing cached %s: %s", label, cached) + return cached + + if _hub_download_in_flight(hf_repo): + logger.info("Skipping %s download while a hub download is active", label) + return None + target: Optional[str] = None from huggingface_hub import list_repo_files @@ -4755,33 +5024,23 @@ class LlamaCppBackend: hf_repo: str, hf_token: Optional[str] = None, cancel_event: Optional[threading.Event] = None, + near_path: Optional[str] = None, ) -> Optional[str]: """Download the mmproj (vision projection) file from a GGUF repo. Prefers mmproj-F16.gguf, else any mmproj*.gguf. Returns the local path, or None if none exists. ``cancel_event`` overrides - ``self._cancel_event`` (defaults to it). + ``self._cancel_event`` (defaults to it). ``near_path`` prefers a + copy co-located with the main GGUF's cache snapshot. """ - def _pick_mmproj(candidates: list[str]) -> Optional[str]: - mmproj_files = sorted( - f - for f in candidates - if f.lower().endswith(".gguf") and "mmproj" in Path(f).name.lower() - ) - if not mmproj_files: - return None - for f in mmproj_files: - if f.lower().endswith("-f16.gguf"): - return f - return mmproj_files[0] - return self._download_companion_gguf( hf_repo = hf_repo, hf_token = hf_token, pick = _pick_mmproj, label = "mmproj", cancel_event = cancel_event, + near_path = near_path, ) def _cached_repo_mtp_drafter(self, hf_repo: str) -> Optional[str]: @@ -4812,6 +5071,7 @@ class LlamaCppBackend: *, hf_repo: str, hf_token: Optional[str] = None, + near_path: Optional[str] = None, ) -> Optional[str]: """Download the separate MTP drafter (speculative head) from a GGUF repo. @@ -4823,16 +5083,6 @@ 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). @@ -4845,11 +5095,28 @@ class LlamaCppBackend: ) return mtp_files[0] if mtp_files else None + if near_path: + cached = _companion_snapshot_sibling(near_path, _pick_mtp) + if cached: + logger.info("Reusing cached MTP drafter: %s", cached) + return cached + + # 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 + return self._download_companion_gguf( hf_repo = hf_repo, hf_token = hf_token, pick = _pick_mtp, label = "MTP drafter", + near_path = near_path, ) def _resolve_launch_mmproj_path( @@ -5431,6 +5698,7 @@ class LlamaCppBackend: ) self._stdout_thread.start() + @_with_gguf_load_marker def load_model( self, *, @@ -5576,6 +5844,7 @@ class LlamaCppBackend: mmproj_path = self._download_mmproj( hf_repo = hf_repo, hf_token = hf_token, + near_path = model_path, ) # Auto-download the separate MTP drafter (e.g. Gemma) when # the requested spec mode can use it. Repos with the head @@ -5593,6 +5862,7 @@ class LlamaCppBackend: mtp_draft_path = self._download_mtp( hf_repo = hf_repo, hf_token = hf_token, + near_path = model_path, ) elif gguf_path: if not Path(gguf_path).is_file(): @@ -8227,6 +8497,11 @@ class LlamaCppBackend: if not is_ours: continue + # A live parent means a running Studio (or the user's + # shell) still owns it -- not an orphan. + if LlamaCppBackend._pid_parent_is_alive(proc.info["pid"]): + continue + proc.kill() killed += 1 logger.info( @@ -8279,6 +8554,9 @@ class LlamaCppBackend: if not owned: continue + if LlamaCppBackend._pid_parent_is_alive(pid): + continue + try: os.kill(pid, signal.SIGKILL) killed += 1 @@ -10107,6 +10385,9 @@ class LlamaCppBackend: assistant_msg: dict = {"role": "assistant", "content": content_text} assistant_appended = False + # Collect no-op nudges and flush them after the batch, so a no-op + # doesn't abort it and drop the parallel calls that follow. + deferred_noop_msgs: list = [] # The text-path provisional card uses the parser's default id ("call_0"); # a Mistral-style call carries its own id and would open a duplicate. Reuse @@ -10149,14 +10430,14 @@ class LlamaCppBackend: "provenance": decision.provenance, } completion = tool_controller.record_noop(decision) - conversation.append(completion.model_message()) + deferred_noop_msgs.append(completion.model_message()) if _forced_tool_call_pending: _forced_tool_call_pending = False logger.info( "Suppressed local GGUF tool call as internal no-op: " f"action={decision.action} tool={decision.tool_name}" ) - break + continue if not assistant_appended: assistant_msg["tool_calls"] = [decision.as_assistant_tool_call()] @@ -10275,6 +10556,8 @@ class LlamaCppBackend: if _forced_tool_call_pending: _forced_tool_call_pending = False + append_deferred_nudges(conversation, deferred_noop_msgs) + # Close provisional cards not resolved by execution/no-op handling. for _pid, _pname in provisional_started_tool_calls.items(): if _pid not in resolved_provisional_tool_call_ids: diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 163ade10c4..e7a90b4307 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -11,6 +11,10 @@ import threading from typing import Optional, Generator from core.inference.message_content import content_to_text from core.inference.runtime_context import runtime_context_length +from core.inference.chat_template_helpers import ( + ReasoningChannelNormalizer, + normalize_reasoning_snapshots, +) from loggers import get_logger logger = get_logger(__name__) @@ -533,7 +537,7 @@ class MLXInferenceBackend: break if self._is_vlm: - yield from self._generate_vlm( + stream = self._generate_vlm( full_messages, image, temperature, @@ -550,7 +554,7 @@ class MLXInferenceBackend: presence_penalty = presence_penalty, ) else: - yield from self._generate_text( + stream = self._generate_text( full_messages, temperature, top_p, @@ -565,6 +569,7 @@ class MLXInferenceBackend: preserve_thinking = preserve_thinking, presence_penalty = presence_penalty, ) + yield from stream def _generate_text( self, @@ -609,7 +614,7 @@ class MLXInferenceBackend: # 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( + render_result = render_with_native_template_fallback( formatted_prompt = prompt, tokenizer = self._tokenizer, model_info = model_info, @@ -620,7 +625,10 @@ class MLXInferenceBackend: reasoning_effort = reasoning_effort, preserve_thinking = preserve_thinking, hf_token = model_info.get("hf_token"), + return_metadata = True, ) + prompt = render_result.prompt + reasoning_channel_markers = render_result.reasoning_channel_markers # An open prefilled by the template lives in the prompt, not # the generated tokens; re-emit it so the frontend renders the block. @@ -654,7 +662,17 @@ class MLXInferenceBackend: if not logits_processors: logits_processors = None + preserve_native_channels = reasoning_channel_markers is not None token_ids = [] + normalizer = ( + ReasoningChannelNormalizer(*reasoning_channel_markers) + if reasoning_channel_markers is not None + else None + ) + # MLX consumers diff cumulative snapshots. Keep a prompt-prefilled + # prefix on every native-protocol snapshot just as the normal + # decoding path does below. + normalized_output = think_prefix logger.info( "Generating: prompt_len=%d, max_tokens=%d, model=%s, tokenizer=%s", len(prompt), @@ -678,12 +696,19 @@ class MLXInferenceBackend: **gen_kwargs, ): final_response = response - token_ids.append(response.token) - cumulative = self._tokenizer.decode( - token_ids, - skip_special_tokens = True, - ) - yield think_prefix + cumulative + if preserve_native_channels: + piece = getattr(response, "text", None) or "" + delta = normalizer.feed(piece) + if delta: + normalized_output += delta + yield normalized_output + else: + token_ids.append(response.token) + cumulative = self._tokenizer.decode( + token_ids, + skip_special_tokens = True, + ) + yield think_prefix + cumulative if cancel_event and cancel_event.is_set(): break @@ -700,6 +725,12 @@ class MLXInferenceBackend: getattr(final_response, "generation_tokens", 0), getattr(final_response, "generation_tps", 0.0), ) + if normalizer is not None: + cancelled = cancel_event is not None and cancel_event.is_set() + tail = normalizer.drain() if cancelled else normalizer.finish() + if tail: + normalized_output += tail + yield normalized_output def _generate_vlm( self, @@ -858,31 +889,37 @@ class MLXInferenceBackend: elif _rep_active: vlm_kwargs["repetition_penalty"] = float(repetition_penalty) - with self._generation_lock: - final_response = None - try: - for response in vlm_stream( - self._model, - self._processor, - prompt, - images, - **vlm_kwargs, - ): - final_response = response - token_text = response.text if hasattr(response, "text") else str(response) - cumulative += token_text - yield cumulative - if cancel_event and cancel_event.is_set(): - break - finally: - # mlx_vlm exposes the same stats fields as mlx_lm. - if final_response is not None: - self.last_generation_stats = _build_generation_stats( - getattr(final_response, "prompt_tokens", 0), - getattr(final_response, "prompt_tps", 0.0), - getattr(final_response, "generation_tokens", 0), - getattr(final_response, "generation_tps", 0.0), - ) + def _stream_vlm_snapshots(): + nonlocal cumulative + with self._generation_lock: + final_response = None + try: + for response in vlm_stream( + self._model, + self._processor, + prompt, + images, + **vlm_kwargs, + ): + final_response = response + token_text = response.text if hasattr(response, "text") else str(response) + cumulative += token_text + yield cumulative + if cancel_event and cancel_event.is_set(): + break + finally: + # mlx_vlm exposes the same stats fields as mlx_lm. + if final_response is not None: + self.last_generation_stats = _build_generation_stats( + getattr(final_response, "prompt_tokens", 0), + getattr(final_response, "prompt_tps", 0.0), + getattr(final_response, "generation_tokens", 0), + getattr(final_response, "generation_tps", 0.0), + ) + + yield from normalize_reasoning_snapshots( + _stream_vlm_snapshots(), chat_target, cancel_event, tools = tools + ) def generate_with_adapter_control( self, diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index c2082bc198..3afda74411 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -59,7 +59,32 @@ class GenStreamError(str): "Error:" by checking isinstance(chunk, GenStreamError). """ - __slots__ = () + __slots__ = ("public",) + + def __new__( + cls, + value, + *, + public: bool = False, + ): + obj = str.__new__(cls, value) + obj.public = bool(public) + return obj + + +class GenStreamErrorRaised(RuntimeError): + """Internal exception form of ``GenStreamError`` for generator boundaries.""" + + __slots__ = ("public",) + + def __init__( + self, + value, + *, + public: bool = False, + ): + super().__init__(value) + self.public = bool(public) class InferenceOrchestrator: @@ -531,13 +556,19 @@ class InferenceOrchestrator: 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)}") + yield GenStreamError( + f"Error: {self._subprocess_crash_message(crash_context)}", + public = True, + ) return resp = read_one(read_timeout) if resp is None: # Check subprocess health if not self._ensure_subprocess_alive(): - yield GenStreamError(f"Error: {self._subprocess_crash_message(crash_context)}") + yield GenStreamError( + f"Error: {self._subprocess_crash_message(crash_context)}", + public = True, + ) return continue @@ -689,11 +720,11 @@ class InferenceOrchestrator: GPU work stays serialized; this only avoids orchestrator lock contention. """ if not self._ensure_subprocess_alive(): - yield GenStreamError("Error: Inference subprocess is not running") + yield GenStreamError("Error: Inference subprocess is not running", public = True) return if not self.active_model_name: - yield GenStreamError("Error: No active model") + yield GenStreamError("Error: No active model", public = True) 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 @@ -704,7 +735,7 @@ class InferenceOrchestrator: # 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") + yield GenStreamError("Error: model is being unloaded", public = True) return # Ensure the dispatcher runs. _start_dispatcher serializes concurrent starters under @@ -776,7 +807,7 @@ class InferenceOrchestrator: # _stop_dispatcher joins the dispatcher, which itself takes that lock. if orphaned_dispatcher: self._stop_dispatcher() - yield GenStreamError("Error: model is being unloaded") + yield GenStreamError("Error: model is being unloaded", public = True) return try: @@ -1376,6 +1407,7 @@ class InferenceOrchestrator: use_adapter: Optional[Union[bool, str]] = None, stats_holder: Optional[dict] = None, presence_penalty: float = 0.0, + reasoning_prefilled: bool = False, **_unused, ): """Run the safetensors agentic tool loop in the parent process, @@ -1414,12 +1446,27 @@ class InferenceOrchestrator: presence_penalty = presence_penalty, ) if use_adapter is not None: - yield from self.generate_with_adapter_control( + stream = self.generate_with_adapter_control( use_adapter = use_adapter, **common_kwargs, ) else: - yield from self.generate_chat_response(**common_kwargs) + stream = self.generate_chat_response(**common_kwargs) + close_stream = False + try: + for chunk in stream: + if isinstance(chunk, GenStreamError): + close_stream = True + raise GenStreamErrorRaised(str(chunk), public = chunk.public) + yield chunk + finally: + if close_stream: + close = getattr(stream, "close", None) + if callable(close): + try: + close() + except Exception: + logger.debug("failed to close errored generation stream", exc_info = True) initial = list(messages) if system_prompt: @@ -1441,6 +1488,7 @@ class InferenceOrchestrator: confirm_tool_calls = confirm_tool_calls, bypass_permissions = bypass_permissions, permission_mode = permission_mode, + reasoning_prefilled = reasoning_prefilled, ) def generate_with_adapter_control( @@ -1489,11 +1537,11 @@ class InferenceOrchestrator: readers don't consume each other's tokens off the shared resp_queue. """ if not self._ensure_subprocess_alive(): - yield GenStreamError("Error: Inference subprocess is not running") + yield GenStreamError("Error: Inference subprocess is not running", public = True) return if not self.active_model_name: - yield GenStreamError("Error: No active model") + yield GenStreamError("Error: No active model", public = True) return expected_model = self.active_model_name @@ -1510,7 +1558,7 @@ class InferenceOrchestrator: # 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") + yield GenStreamError("Error: model is being unloaded", public = True) return request_id = str(uuid.uuid4()) image_b64 = self._pil_to_base64(image) if image is not None else None @@ -1695,10 +1743,10 @@ class InferenceOrchestrator: ) -> Generator[str, None, None]: """Shared inner logic for audio input generation (Whisper + ASR).""" if not self._ensure_subprocess_alive(): - yield GenStreamError("Error: Inference subprocess is not running") + yield GenStreamError("Error: Inference subprocess is not running", public = True) return if not self.active_model_name: - yield GenStreamError("Error: No active model") + yield GenStreamError("Error: No active model", public = True) return expected_model = self.active_model_name @@ -1707,7 +1755,7 @@ class InferenceOrchestrator: # 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") + yield GenStreamError("Error: model is being unloaded", public = True) return request_id = str(uuid.uuid4()) diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index f4c243d1bf..9110315815 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -50,6 +50,7 @@ from core.inference.tool_call_parser import ( # pattern lists, so the safetensors streaming strip stays aligned with the parser. from core.tool_healing import ( _REHEARSAL_TAIL_STRIP_RE, + _THINK_CLOSE_RE, _strip_bracket_tag_calls, _think_spans_outside_tool_markup, apply_tool_strip_patterns, @@ -57,6 +58,7 @@ from core.tool_healing import ( ) from core.inference.tool_loop_controller import ( ToolLoopController, + append_deferred_nudges, coerce_tool_arguments, status_for_tool, tool_event_provenance, @@ -303,6 +305,45 @@ def _status_for_tool(tool_name: str, arguments: dict) -> str: return status_for_tool(tool_name, arguments) +def _reprompt_intent_text(text: str, *, reasoning_prefilled: bool = False) -> str: + """Return visible answer text for the plan-without-action classifier. + + Safetensors reasoning shares the cumulative text channel with the answer. + Forward-looking phrases inside ```` / ``[THINK]`` are private + planning, not a user-visible promise to call a tool. Match GGUF's behavior: + classify visible content when present and fall back to reasoning only for a + reasoning-only stall. + """ + prefilled_reasoning = "" + if reasoning_prefilled: + close = _THINK_CLOSE_RE.search(text) + if close is None: + return text.strip() + prefilled_reasoning = text[: close.end()].strip() + text = text[close.end() :].strip() + if not text: + return prefilled_reasoning + + spans = _think_spans_outside_tool_markup(text) + if not spans: + return text.strip() + + visible: list[str] = [] + reasoning: list[str] = [] + cursor = 0 + for start, end in spans: + visible.append(text[cursor:start]) + reasoning.append(text[start:end]) + cursor = end + visible.append(text[cursor:]) + + visible_text = "".join(visible).strip() + reasoning_text = "".join(reasoning).strip() + if visible_text: + return visible_text + return "\n".join(part for part in (prefilled_reasoning, reasoning_text) if part).strip() + + 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()) @@ -447,6 +488,7 @@ def run_safetensors_tool_loop( confirm_tool_calls: bool = False, bypass_permissions: bool = False, permission_mode: Optional[str] = None, + reasoning_prefilled: bool = False, ) -> Generator[dict, None, None]: """Drive an agentic tool loop on top of a cumulative-text generator. @@ -955,7 +997,10 @@ def run_safetensors_tool_loop( # (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() + intent_text = _reprompt_intent_text( + content_accum, + reasoning_prefilled = reasoning_prefilled, + ) if ( auto_heal_tool_calls and nudge_tool_calls @@ -964,7 +1009,7 @@ def run_safetensors_tool_loop( 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) + and is_short_intent_without_action(intent_text) ): reprompt_count += 1 logger.info( @@ -972,9 +1017,9 @@ def run_safetensors_tool_loop( "calling tools (%d chars)", reprompt_count, MAX_ACT_REPROMPTS, - len(stripped_answer), + len(intent_text), ) - conversation.append({"role": "assistant", "content": stripped_answer}) + conversation.append({"role": "assistant", "content": intent_text}) tool_hint = " or ".join(_active_tool_names(active_tools)) or "an available tool" conversation.append( { @@ -1099,6 +1144,9 @@ def run_safetensors_tool_loop( assistant_msg: dict = {"role": "assistant", "content": content_text} assistant_appended = False + # Collect no-op nudges and flush them after the batch, so a no-op doesn't + # abort it and drop the parallel calls that follow. + deferred_noop_msgs: list = [] for tc in tool_calls or []: func = tc.get("function", {}) or {} @@ -1127,12 +1175,12 @@ def run_safetensors_tool_loop( "provenance": decision.provenance, } completion = tool_controller.record_noop(decision) - conversation.append(completion.model_message()) + deferred_noop_msgs.append(completion.model_message()) logger.info( "Suppressed local safetensors tool call as internal no-op: " f"action={decision.action} tool={decision.tool_name}" ) - break + continue if not assistant_appended: assistant_msg["tool_calls"] = [decision.as_assistant_tool_call()] @@ -1243,6 +1291,8 @@ def run_safetensors_tool_loop( yield completion.tool_end_event() conversation.append(completion.tool_message()) + append_deferred_nudges(conversation, deferred_noop_msgs) + # Clear the status badge before the next turn. yield {"type": "status", "text": ""} diff --git a/studio/backend/core/inference/tool_loop_controller.py b/studio/backend/core/inference/tool_loop_controller.py index f595531b90..f7ed450d11 100644 --- a/studio/backend/core/inference/tool_loop_controller.py +++ b/studio/backend/core/inference/tool_loop_controller.py @@ -266,6 +266,17 @@ def strip_result_for_model(result: str) -> str: return result +def append_deferred_nudges(conversation: list, msgs: Sequence[dict]) -> None: + """Append a batch's no-op nudges as one deduped ``role=user`` message. + + Deferred to after the batch's tool results so a no-op never splits an + assistant's ``tool_calls`` from their ``role=tool`` results. + """ + contents = list(dict.fromkeys(msg["content"] for msg in msgs)) + if contents: + conversation.append({"role": "user", "content": "\n\n".join(contents)}) + + def _tool_name_from_schema(tool: Mapping[str, Any]) -> str: function = tool.get("function") if not isinstance(function, Mapping): @@ -277,8 +288,9 @@ def _tool_name_from_schema(tool: Mapping[str, Any]) -> str: def _noop_result(reason: NoopReason, tool_name: str) -> str: if reason == "duplicate": return ( - "The previous tool request was not executed because this exact " - "tool call already completed successfully. Do not repeat the same " + f"One earlier request to call tool '{tool_name}' in this batch was " + "not executed because an identical call had already completed " + "successfully. Do not repeat the same " "tool call. Continue with a different enabled tool if that would " "materially help, or provide the final answer if you have enough " "information." @@ -291,8 +303,8 @@ def _noop_result(reason: NoopReason, tool_name: str) -> str: "the requested final note or answer." ) return ( - f"The previous tool request was not executed because tool " - f"'{tool_name}' is not enabled for this request. Provide the " + f"One earlier request to call tool '{tool_name}' in this batch was " + "not executed because that tool is not enabled for this request. Provide the " "final answer now without calling more tools." ) diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index dd268a6bb7..5fd57e1b2c 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -3600,14 +3600,18 @@ _MAX_PAGE_CHARS = 16000 # cap fetched page text (after HTML-to-MD conversion) # Raw download cap > _MAX_PAGE_CHARS since SSR pages embed large sections # stripped during conversion; 512 KB still reaches article content. _MAX_FETCH_BYTES = 512 * 1024 +# PDF cross-reference data lives at EOF, so extraction needs the whole body. +_MAX_PDF_FETCH_BYTES = 10 * 1024 * 1024 +_MAX_WEB_PDF_PAGES = 50 # Control/undecodable chars, excluding text whitespace and ESC (for ANSI logs). # Binary when they exceed 12.5%, after allowing 16 minor encoding glitches. _BINARY_CHAR_RE = re.compile("[\\x00-\\x08\\x0b\\x0c\\x0e-\\x1a\\x1c-\\x1f\\x7f-\\x9f\\ufffd]") _MIN_BINARY_CHARS = 16 _BINARY_CHAR_DIVISOR = 8 # Common binary signatures that can otherwise look text-heavy when mislabeled. +_PDF_MAGIC = b"%PDF-" _BINARY_MAGIC = ( - b"%PDF-", # PDF + _PDF_MAGIC, b"PK\x03\x04", # zip / docx / xlsx / pptx / epub / jar b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1", # OLE / legacy Office b"\x89PNG\r\n\x1a\n", # PNG @@ -3641,14 +3645,22 @@ def _looks_binary(text: str) -> bool: ) -def _has_binary_magic(data: bytes) -> bool: - """Whether a common binary signature follows optional BOM or whitespace.""" +def _magic_head(data: bytes) -> bytes: head = data[:1024].lstrip() for bom, _codec in _UNICODE_BOM_CODECS: if head.startswith(bom): head = head.removeprefix(bom).lstrip() break - return head.startswith(_BINARY_MAGIC) + return head + + +def _has_pdf_magic(data: bytes) -> bool: + return _magic_head(data).startswith(_PDF_MAGIC) + + +def _has_binary_magic(data: bytes) -> bool: + """Whether a common binary signature follows optional BOM or whitespace.""" + return _magic_head(data).startswith(_BINARY_MAGIC) def _has_single_byte_text_evidence(data: bytes) -> bool: @@ -3659,6 +3671,45 @@ def _has_single_byte_text_evidence(data: bytes) -> bool: return ascii_text_bytes / len(data) >= _MIN_SINGLE_BYTE_ASCII_RATIO +def _extract_pdf_text(data: bytes) -> str: + """Extract page-delimited text with the same parser used by RAG ingestion.""" + from ..rag.parsers import parse_pdf_bytes + + pages, total_pages = parse_pdf_bytes(data, max_pages = _MAX_WEB_PDF_PAGES) + page_limit_reached = total_pages > _MAX_WEB_PDF_PAGES + parts: list[str] = [] + length = 0 + text_limited = False + for page in pages: + page_text = page.text.strip() + if not page_text: + continue + section = f"## Page {page.page_number}\n\n{page_text}" + piece = ("\n\n" if parts else "") + section + remaining = _MAX_PAGE_CHARS - length + if len(piece) > remaining: + parts.append(piece[:remaining]) + text_limited = True + break + parts.append(piece) + length += len(piece) + + text = "".join(parts).rstrip() + if not text: + if page_limit_reached: + return f"(PDF contains no extractable text in the first {_MAX_WEB_PDF_PAGES} pages)" + return "" + limits = [] + if text_limited: + limits.append(f"text limited to {_MAX_PAGE_CHARS:,} characters") + if page_limit_reached: + limits.append(f"page processing capped at {_MAX_WEB_PDF_PAGES} pages") + if limits: + marker = f"\n\n... (PDF extraction {'; '.join(limits)})" + text = text[: _MAX_PAGE_CHARS - len(marker)].rstrip() + marker + return text + + _USER_AGENTS = ( "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36", @@ -4054,29 +4105,71 @@ def _fetch_url_raw( return reason2, "", "" current_host = rp.hostname continue + + # get_content_type() defaults to "text/plain" when the header is + # absent (RFC 2045); report "" instead so callers can tell a missing + # header apart from a server that really declared text/plain. + if resp.headers.get("Content-Type") is None: + content_type = "" + else: + content_type = (resp.headers.get_content_type() or "").lower() + # Success: read the capped body enforcing the budget between chunks # (see _read_capped_body), so a slow-drip server can't stretch a # single resp.read past the deadline. + declared_pdf = content_type == "application/pdf" + read_limit = _MAX_PDF_FETCH_BYTES + 1 if declared_pdf else max_bytes body_error, raw_bytes = _read_capped_body( resp, - max_bytes, + read_limit, timeout, deadline, cancel_event, ) if body_error is not None: return body_error, "", "" + + # A missing or wrong PDF MIME type is common: once the initial text-sized + # read identifies PDF magic, finish the bounded download to reach the EOF xref. + if not declared_pdf and len(raw_bytes) == max_bytes and _has_pdf_magic(raw_bytes): + tail_error, tail = _read_capped_body( + resp, + _MAX_PDF_FETCH_BYTES - max_bytes + 1, + timeout, + deadline, + cancel_event, + ) + if tail_error is not None: + return tail_error, "", "" + raw_bytes += tail break else: return "Failed to fetch URL: too many redirects.", "", "" - # get_content_type() defaults to "text/plain" when the header is - # absent (RFC 2045); report "" instead so callers can tell a missing - # header apart from a server that really declared text/plain. - if resp.headers.get("Content-Type") is None: - content_type = "" - else: - content_type = (resp.headers.get_content_type() or "").lower() + is_pdf = declared_pdf or _has_pdf_magic(raw_bytes) + if is_pdf: + if len(raw_bytes) > _MAX_PDF_FETCH_BYTES: + return ( + "(PDF content exceeds the download limit; not readable as text)", + "", + content_type, + ) + budget_error = _fetch_budget_exceeded(deadline, cancel_event) + if budget_error is not None: + return budget_error, "", content_type + try: + pdf_text = _extract_pdf_text(raw_bytes) + except Exception as exc: + logger.debug("web PDF text extraction failed (%s)", type(exc).__name__) + return "(PDF content could not be read as text)", "", content_type + budget_error = _fetch_budget_exceeded(deadline, cancel_event) + if budget_error is not None: + return budget_error, "", content_type + if not pdf_text: + pdf_text = "(PDF contains no extractable text)" + # Report the true type even for a mislabeled body so the caller's "html" + # check routes the extracted text to the plain-text path, not html_to_markdown. + return None, pdf_text, "application/pdf" # Reject known-binary MIME types before decoding. Binary is returned as the # error string so the caller surfaces the placeholder, not replacement chars. diff --git a/studio/backend/core/rag/parsers.py b/studio/backend/core/rag/parsers.py index 9afddf1d9e..0b42906b85 100644 --- a/studio/backend/core/rag/parsers.py +++ b/studio/backend/core/rag/parsers.py @@ -103,7 +103,7 @@ def _markdown_incomplete(markdown: str, plain: str) -> bool: return markdown_letters < _PDF_INCOMPLETE_RATIO * plain_letters -def _pdf_markdown(doc) -> list[str] | None: +def _pdf_markdown(doc, pages: range | None = None) -> list[str] | None: """Per-page layout-aware Markdown (tables, headings, lists) via pymupdf4llm; index i maps to page i+1. Returns None when the lib is missing, extraction fails, or the page count does not line up, so the caller falls back to plain PyMuPDF text.""" @@ -112,28 +112,44 @@ def _pdf_markdown(doc) -> list[str] | None: except Exception: return None try: - chunks = pymupdf4llm.to_markdown( - doc, - page_chunks = True, - show_progress = False, - ) + kwargs = {"page_chunks": True, "show_progress": False} + if pages is not None: + kwargs["pages"] = list(pages) + chunks = pymupdf4llm.to_markdown(doc, **kwargs) except Exception: # noqa: BLE001 - never let Markdown extraction break ingestion logger.warning("pymupdf4llm extraction failed; using plain text", exc_info = True) return None - if not isinstance(chunks, list) or len(chunks) != doc.page_count: + expected_pages = doc.page_count if pages is None else len(pages) + if not isinstance(chunks, list) or len(chunks) != expected_pages: return None return [str(c.get("text") or "") for c in chunks] -def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]: +def _pdf( + source: str | bytes, + want_images: bool, + max_pages: int | None = None, +) -> tuple[list[Page], list[ParsedImage], int]: import fitz # PyMuPDF pages: list[Page] = [] images: list[ParsedImage] = [] - doc = fitz.open(path) + doc = ( + fitz.open(stream = source, filetype = "pdf") if isinstance(source, bytes) else fitz.open(source) + ) try: - md = _pdf_markdown(doc) if config.PDF_MARKDOWN else None - for i, page in enumerate(doc): + if doc.needs_pass: + raise ValueError("encrypted PDF requires a password") + total_pages = doc.page_count + page_numbers = range(total_pages if max_pages is None else min(total_pages, max_pages)) + if not config.PDF_MARKDOWN: + md = None + elif max_pages is None: + md = _pdf_markdown(doc) + else: + md = _pdf_markdown(doc, page_numbers) + for i, page_number in enumerate(page_numbers): + page = doc[page_number] plain = page.get_text("text") or "" candidate = md[i] if md else "" # Prefer layout-aware Markdown (keeps tables/headings legible for retrieval), @@ -147,7 +163,7 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]: text = candidate else: text = plain - pages.append(_page(text, i + 1)) + pages.append(_page(text, page_number + 1)) if want_images: for img in page.get_images(full = True): xref = img[0] @@ -161,13 +177,22 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]: images.append( ParsedImage( image_bytes = image_bytes, - page_number = i + 1, + page_number = page_number + 1, xref = xref, ) ) finally: doc.close() - return pages, images + return pages, images, total_pages + + +def parse_pdf_bytes(data: bytes, *, max_pages: int | None = None) -> tuple[list[Page], int]: + """Extract PDF pages from an in-memory download using the ingestion parser. + + Returns the (capped) pages plus the document's full page count, so a caller + that set ``max_pages`` can tell a fully-read short PDF from a truncated one.""" + pages, _images, total_pages = _pdf(data, want_images = False, max_pages = max_pages) + return pages, total_pages def _merge_rects(boxes: list) -> list: @@ -416,7 +441,7 @@ def parse(path: str, *, want_images: bool = False): ext = os.path.splitext(path)[1].lower() if ext == ".pdf": - pages, images = _pdf(path, want_images) + pages, images, _total = _pdf(path, want_images) return (pages, images) if want_images else pages if ext == ".docx": diff --git a/studio/backend/hub/services/download_lifecycle.py b/studio/backend/hub/services/download_lifecycle.py index 50e20fc13f..e5d48872c1 100644 --- a/studio/backend/hub/services/download_lifecycle.py +++ b/studio/backend/hub/services/download_lifecycle.py @@ -8,6 +8,7 @@ import os import signal import subprocess import sys +import time import threading from pathlib import Path from typing import Callable, Optional @@ -210,7 +211,9 @@ def finalize_worker_exit( repo_type: Optional[RepoType] = None, repo_id: Optional[str] = None, transport: Optional[str] = None, -) -> None: + cancel_marker_transport: Optional[str] = None, + defer_error: bool = False, +) -> str: """Block until *proc* exits, then record the job's terminal state in *registry*. Drains and scrubs stderr first, then classifies the exit code. A no-op when the process was already dropped (e.g. superseded). @@ -222,7 +225,7 @@ def finalize_worker_exit( rc = proc.wait() cancel_requested = registry.cancel_requested(key) if not registry.drop_process(key, proc): - return + return "idle" stderr_text = download_registry.scrub_secrets( (stderr_data or b"").decode("utf-8", "replace").strip(), hf_token = hf_token, @@ -230,6 +233,8 @@ def finalize_worker_exit( state = classify_exit(rc, cancel_requested = cancel_requested) if state == "complete": registry.set_job(key, "complete") + if transport == download_registry.TRANSPORT_HTTP: + registry.update_job_transport(key, download_registry.TRANSPORT_HTTP) if stderr_text: if download_manifest.MANIFEST_DEGRADED_MARKER in stderr_text: logger.warning( @@ -262,18 +267,226 @@ def finalize_worker_exit( metadata.variant if metadata is not None and metadata.variant else download_registry.variant_from_key(key), - transport, + cancel_marker_transport or transport, logger = logger, ) else: - registry.set_job( - key, - "error", - stderr_text or f"worker exited with code {rc}", - ) + if not defer_error: + registry.set_job( + key, + "error", + stderr_text or f"worker exited with code {rc}", + ) logger.error( f"{log_prefix} failed for {label} (rc={rc}): {stderr_text}", ) + return state + + +def _set_retry_failure_state( + registry: download_registry.DownloadRegistry, + key: str, + error: str, + *, + repo_type: RepoType, + repo_id: str, + fallback_variant: Optional[str], + fallback_transport: Optional[str], + logger, +) -> str: + state, metadata = registry.set_error_unless_cancelled(key, error) + if state == "cancelled": + download_registry.persist_cancel_marker( + repo_type, + repo_id, + metadata.variant if metadata is not None and metadata.variant else fallback_variant, + metadata.transport + if metadata is not None and metadata.transport + else fallback_transport, + logger = logger, + ) + return state + + +def _try_http_retry( + registry: download_registry.DownloadRegistry, + key: str, + *, + hf_token: Optional[str], + label: str, + log_prefix: str, + logger, + repo_type: RepoType, + repo_id: str, + watch_name: str, +) -> bool: + """Reclaim *key* with HTTP transport and spawn a recovery worker. + + Returns ``True`` when the HTTP worker was successfully registered. + Caller is responsible for ensuring this is only called when: the job is + in ``"error"`` state, the original transport was XET, and HTTP is available. + + Derives variant and blob-hash metadata from the registry entry written by + the original XET claim so callers do not re-construct worker arguments. + Re-queries peer protection hashes at spawn time to reflect any concurrent + sibling changes between the XET failure and this call. + """ + original_metadata = registry.get_job_metadata(key) + if original_metadata is None: + logger.debug("%s XET retry skipped for %s; metadata unavailable", log_prefix, label) + _set_retry_failure_state( + registry, + key, + "XET retry skipped: metadata unavailable", + repo_type = repo_type, + repo_id = repo_id, + fallback_variant = download_registry.variant_from_key(key), + fallback_transport = download_registry.TRANSPORT_XET, + logger = logger, + ) + return False + if original_metadata.transport != download_registry.TRANSPORT_XET: + logger.debug( + "%s XET retry skipped for %s; original transport was %s", + log_prefix, + label, + original_metadata.transport, + ) + _set_retry_failure_state( + registry, + key, + f"XET retry skipped: original transport was {original_metadata.transport}", + repo_type = repo_type, + repo_id = repo_id, + fallback_variant = original_metadata.variant, + fallback_transport = original_metadata.transport, + logger = logger, + ) + return False + variant = original_metadata.variant + blob_hashes = original_metadata.blob_hashes + progress_blob_hashes = original_metadata.progress_blob_hashes + completed_baseline_bytes = ( + download_registry.completed_blob_bytes( + repo_type, + repo_id, + progress_blob_hashes, + ) + if progress_blob_hashes + else 0 + ) + generation = registry.current_generation(key) + registry.release_active_slot(key) + while True: + if registry.cancel_requested(key): + _set_retry_failure_state( + registry, + key, + "HTTP retry cancelled before reclaiming the download slot", + repo_type = repo_type, + repo_id = repo_id, + fallback_variant = variant, + fallback_transport = original_metadata.transport, + logger = logger, + ) + return False + + claimed, conflict_state = registry.claim( + key, + download_registry.TRANSPORT_HTTP, + repo_type = repo_type, + repo_id = repo_id, + variant = variant, + blob_hashes = blob_hashes, + progress_blob_hashes = progress_blob_hashes, + completed_baseline_bytes = completed_baseline_bytes, + generation = generation, + replace_active = True, + cancel_marker_transport = original_metadata.transport, + ) + if claimed: + break + if conflict_state == "deleting": + logger.debug( + "%s XET retry claim rejected for %s; repo is being deleted", + log_prefix, + label, + ) + _set_retry_failure_state( + registry, + key, + "HTTP retry could not reclaim the download slot", + repo_type = repo_type, + repo_id = repo_id, + fallback_variant = variant, + fallback_transport = original_metadata.transport, + logger = logger, + ) + return False + logger.debug( + "%s XET retry claim blocked for %s by active sibling state %s; waiting", + log_prefix, + label, + conflict_state, + ) + time.sleep(0.05) + + args: list[str] = ["--repo-id", repo_id] + if repo_type == "dataset": + args.append("--dataset") + elif variant: + args.extend(["--variant", variant]) + + # Re-query at spawn time: sibling state may have changed since XET failed. + peer_hashes = registry.peer_blob_hashes(key) if variant else frozenset() + + logger.warning( + "%s XET worker failed for %s; retrying over HTTP", + log_prefix, + label, + ) + try: + proc = spawn_worker( + args, + hf_token, + use_xet = False, + protected_blob_hashes = peer_hashes or None, + ) + except Exception as exc: + scrubbed = download_registry.scrub_secrets(str(exc), hf_token = hf_token) + logger.error( + "%s HTTP retry spawn failed for %s: %s", + log_prefix, + label, + scrubbed, + ) + registry.update_job_transport(key, original_metadata.transport) + _set_retry_failure_state( + registry, + key, + scrubbed, + repo_type = repo_type, + repo_id = repo_id, + fallback_variant = variant, + fallback_transport = original_metadata.transport, + logger = logger, + ) + return False + + return register_worker( + registry, + key, + proc, + hf_token = hf_token, + label = label, + log_prefix = log_prefix, + logger = logger, + repo_type = repo_type, + repo_id = repo_id, + transport = download_registry.TRANSPORT_HTTP, + cancel_marker_transport = original_metadata.transport, + watch_name = watch_name, + ) def kill_and_reap_process( @@ -309,6 +522,7 @@ def register_worker( repo_type: RepoType, repo_id: str, transport: str, + cancel_marker_transport: Optional[str] = None, watch_name: str, ) -> bool: if not registry.register_process(key, proc): @@ -319,7 +533,14 @@ def register_worker( def _watch() -> None: try: - finalize_worker_exit( + can_retry_http = ( + transport == download_registry.TRANSPORT_XET + and download_registry.download_transport_unavailable_reason( + download_registry.TRANSPORT_HTTP + ) + is None + ) + state = finalize_worker_exit( registry, key, proc, @@ -330,7 +551,25 @@ def register_worker( repo_type = repo_type, repo_id = repo_id, transport = transport, + cancel_marker_transport = cancel_marker_transport, + defer_error = can_retry_http, ) + # XET-to-HTTP recovery: when a non-cancelled XET worker fails and + # HTTP is available, attempt one automatic retry over HTTP. The + # transport check is the recursion guard: an HTTP worker that errors + # never satisfies `transport == TRANSPORT_XET`, so it stays terminal. + if can_retry_http and state == "error": + _try_http_retry( + registry, + key, + hf_token = worker_token, + label = label, + log_prefix = log_prefix, + logger = logger, + repo_type = repo_type, + repo_id = repo_id, + watch_name = watch_name, + ) except Exception: # finalize_worker_exit is the only thing that clears running/cancelling; # if it raises, force a terminal state so claim() isn't blocked until restart. @@ -426,8 +665,19 @@ def cancel_worker( return "cancelling" return registry.get_job(key).state # Worker already exited; let its watcher classify the real return code. - # Arming a pending cancel here could mislabel a genuine failure as a cancel. if proc.poll() is not None: + get_metadata = getattr(registry, "get_job_metadata", None) + metadata = get_metadata(key) if get_metadata is not None else None + can_retry_http = ( + metadata is not None + and metadata.transport == download_registry.TRANSPORT_XET + and download_registry.download_transport_unavailable_reason( + download_registry.TRANSPORT_HTTP + ) + is None + ) + if can_retry_http and registry.mark_pending_cancel(key, generation): + return "cancelling" return registry.get_job(key).state if not registry.request_cancel(key, proc, generation): diff --git a/studio/backend/hub/services/models/downloads.py b/studio/backend/hub/services/models/downloads.py index c2ffe7bffc..862af0141a 100644 --- a/studio/backend/hub/services/models/downloads.py +++ b/studio/backend/hub/services/models/downloads.py @@ -60,6 +60,30 @@ def _job_status( return DownloadJobStatus(state = state, error = error, generation = generation) +def _load_in_flight(repo_id: str) -> bool: + try: + from core.inference.llama_cpp import hf_gguf_load_in_flight + return hf_gguf_load_in_flight(repo_id) + except Exception: + return False + + +def _load_in_flight_error(repo_id: str) -> HTTPException: + return HTTPException( + status_code = 409, + detail = ( + f"A model load for '{repo_id}' is in progress and may be " + "downloading it. Wait for the load to finish (or cancel it), " + "then start the download." + ), + ) + + +def _reject_if_load_in_flight(repo_id: str) -> None: + if _load_in_flight(repo_id): + raise _load_in_flight_error(repo_id) + + def _spawn_download_worker( repo_id: str, variant: Optional[str], @@ -89,6 +113,9 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional # Canonicalize so two different-cased paste-ins share one job + cache dir. repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model") + # Avoid concurrent writers to the same HF cache files. + _reject_if_load_in_flight(repo_id) + variant = (body.gguf_variant or "").strip() or None if variant is not None and not _is_valid_gguf_variant(variant): raise HTTPException( @@ -147,9 +174,12 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional blob_hashes = variant_blob_hashes, progress_blob_hashes = variant_progress_blob_hashes, completed_baseline_bytes = completed_baseline_bytes, + admission_check = lambda: not _load_in_flight(repo_id), ) generation = _registry.current_generation(key) if not claimed: + if claim_state == "admission_blocked": + raise _load_in_flight_error(repo_id) # claim_state is the blocking job's state. The client can attach only # when the blocker is this key's own in-flight job (adoptable); a # cross-variant conflict or in-progress delete is not accepted. diff --git a/studio/backend/hub/tests/test_download_lifecycle.py b/studio/backend/hub/tests/test_download_lifecycle.py index a4baafa317..87346573b0 100644 --- a/studio/backend/hub/tests/test_download_lifecycle.py +++ b/studio/backend/hub/tests/test_download_lifecycle.py @@ -1,27 +1,147 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +import io +import logging + from hub.services import download_lifecycle +from hub.utils import download_registry, state_dir -def _set_xet_reason(monkeypatch, reason): +class _Proc: + pid = 4242 + + def __init__( + self, + rc, + stderr = b"", + ): + self.rc = rc + self.stderr = io.BytesIO(stderr) + self.waited = False + + def poll(self): + return self.rc if self.waited else None + + def wait(self, timeout = None): + self.waited = True + return self.rc + + def kill(self): + pass + + +class _ImmediateThread: + def __init__(self, *, target, **_kwargs): + self.target = target + + def start(self): + self.target() + + +def test_resolve_effective_use_xet(monkeypatch): + for requested, unavailable_reason, expected in ( + (False, "unused", False), + (True, None, True), + (True, "hf_xet is not installed", False), + ): + monkeypatch.setattr( + download_lifecycle.download_registry, + "download_transport_unavailable_reason", + lambda _transport, reason = unavailable_reason: reason, + ) + assert download_lifecycle.resolve_effective_use_xet(requested) is expected + + +def test_xet_failure_retries_over_http_for_model_and_dataset(monkeypatch, tmp_path): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + monkeypatch.setattr(download_lifecycle.threading, "Thread", _ImmediateThread) + register_worker = download_lifecycle.register_worker + + for repo_type, repo_id, variant, expected_args in ( + ("model", "Org/Model", "Q4_K_M", ["--repo-id", "Org/Model", "--variant", "Q4_K_M"]), + ("dataset", "Org/Data", None, ["--repo-id", "Org/Data", "--dataset"]), + ): + registry = download_registry.DownloadRegistry() + key = download_registry.normalize_job_key(f"{repo_id}::{variant}" if variant else repo_id) + assert registry.claim( + key, + download_registry.TRANSPORT_XET, + repo_type = repo_type, + repo_id = repo_id, + variant = variant, + blob_hashes = frozenset({"blob"}), + )[0] + generation = registry.current_generation(key) + spawned = [] + + def fake_spawn( + args, + _token, + *, + use_xet, + protected_blob_hashes = None, + ): + spawned.append((args, use_xet, protected_blob_hashes)) + return _Proc(0) + + def fake_retry_register(*_args, **kwargs): + assert kwargs["transport"] == download_registry.TRANSPORT_HTTP + return True + + monkeypatch.setattr(download_lifecycle, "spawn_worker", fake_spawn) + monkeypatch.setattr(download_lifecycle, "register_worker", fake_retry_register) + assert register_worker( + registry, + key, + _Proc(1, b"xet failed"), + hf_token = None, + label = repo_id, + log_prefix = "Download", + logger = logging.getLogger("test"), + repo_type = repo_type, + repo_id = repo_id, + transport = download_registry.TRANSPORT_XET, + watch_name = f"{repo_type}-watch", + ) + + metadata = registry.get_job_metadata(key) + assert spawned == [(expected_args, False, None)] + assert metadata.transport == download_registry.TRANSPORT_HTTP + assert metadata.blob_hashes == frozenset({"blob"}) + assert registry.current_generation(key) == generation + + +def test_http_failure_remains_terminal(monkeypatch, tmp_path): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state") + monkeypatch.setattr(download_lifecycle.threading, "Thread", _ImmediateThread) + register_worker = download_lifecycle.register_worker + registry = download_registry.DownloadRegistry() + key = download_registry.normalize_repo_key("Org/Data") + assert registry.claim( + key, + download_registry.TRANSPORT_HTTP, + repo_type = "dataset", + repo_id = "Org/Data", + )[0] monkeypatch.setattr( - download_lifecycle.download_registry, - "download_transport_unavailable_reason", - lambda _transport: reason, + download_lifecycle, + "register_worker", + lambda *_args, **_kwargs: (_ for _ in ()).throw( + AssertionError("HTTP failures must not retry") + ), ) - - -def test_resolve_effective_use_xet_keeps_http_when_not_requested(monkeypatch): - _set_xet_reason(monkeypatch, "should not be consulted") - assert download_lifecycle.resolve_effective_use_xet(False) is False - - -def test_resolve_effective_use_xet_keeps_xet_when_available(monkeypatch): - _set_xet_reason(monkeypatch, None) - assert download_lifecycle.resolve_effective_use_xet(True) is True - - -def test_resolve_effective_use_xet_downgrades_when_xet_unavailable(monkeypatch): - _set_xet_reason(monkeypatch, "Xet transport is unavailable because hf_xet is not installed.") - assert download_lifecycle.resolve_effective_use_xet(True) is False + assert register_worker( + registry, + key, + _Proc(1, b"http failed"), + hf_token = None, + label = "Org/Data", + log_prefix = "Download", + logger = logging.getLogger("test"), + repo_type = "dataset", + repo_id = "Org/Data", + transport = download_registry.TRANSPORT_HTTP, + watch_name = "dataset-watch", + ) + assert registry.get_job(key).state == "error" diff --git a/studio/backend/hub/utils/download_registry.py b/studio/backend/hub/utils/download_registry.py index 777d63e1b5..b6bdee3bce 100644 --- a/studio/backend/hub/utils/download_registry.py +++ b/studio/backend/hub/utils/download_registry.py @@ -45,9 +45,9 @@ import sys import threading import time import weakref -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from pathlib import Path -from typing import Iterator, Literal, Optional +from typing import Callable, Iterator, Literal, Optional from loggers import get_logger @@ -126,6 +126,9 @@ def write_worker_breadcrumb(key: str, pid: int, metadata: Optional["DownloadMeta "repo_id": metadata.repo_id if metadata is not None else None, "variant": metadata.variant if metadata is not None else None, "transport": metadata.transport if metadata is not None else None, + "cancel_marker_transport": metadata.cancel_marker_transport + if metadata is not None + else None, } tmp = path.with_name(f".{path.name}.tmp-{pid}") try: @@ -305,7 +308,7 @@ def reap_orphan_workers() -> None: data.get("repo_type"), repo_id, data.get("variant"), - data.get("transport"), + data.get("cancel_marker_transport") or data.get("transport"), ) except Exception as exc: logger.debug("Reaper failed for breadcrumb %s: %s", entry, exc) @@ -699,6 +702,7 @@ class DownloadMetadata: repo_id: str variant: Optional[str] transport: Optional[str] + cancel_marker_transport: Optional[str] = None # GGUF variant main/writable hashes, identifying the variant-specific shards # for concurrency decisions. blob_hashes: frozenset[str] = field(default_factory = frozenset) @@ -801,6 +805,7 @@ class DownloadRegistry: self._processes: dict[str, subprocess.Popen] = {} self._repo_active: dict[str, set[str]] = {} self._metadata: dict[str, DownloadMetadata] = {} + self._cancel_marker_transports: dict[str, str] = {} self._pending_cancel: dict[str, Optional[int]] = {} self._generations: dict[str, int] = {} # Monotonic across keys so an evicted then re-claimed key never reuses a @@ -839,6 +844,7 @@ class DownloadRegistry: if state in TERMINAL_STATES: self._put_terminal_job_locked(key, state, error) self._pending_cancel.pop(key, None) + self._cancel_marker_transports.pop(key, None) repo = _repo_of_key(key) active = self._repo_active.get(repo) if active is not None: @@ -848,6 +854,57 @@ class DownloadRegistry: else: self._jobs[key] = DownloadState(state, error) + def set_error_unless_cancelled( + self, key: str, error: str + ) -> tuple[JobState, Optional[DownloadMetadata]]: + key = normalize_job_key(key) + with self._lock: + current = self._jobs.get(key, DownloadState("idle")).state + has_pending_cancel = key in self._pending_cancel + pending_generation = self._pending_cancel.get(key) + metadata = self._metadata.get(key) + should_cancel = current == "cancelling" or ( + has_pending_cancel and self._generation_matches_locked(key, pending_generation) + ) + terminal_state: JobState = "cancelled" if should_cancel else "error" + marker_transport = self._cancel_marker_transports.pop(key, None) + if marker_transport is None and metadata is not None: + marker_transport = metadata.cancel_marker_transport + self._put_terminal_job_locked( + key, + terminal_state, + None if should_cancel else error, + ) + self._pending_cancel.pop(key, None) + repo = _repo_of_key(key) + active = self._repo_active.get(repo) + if active is not None: + active.discard(key) + if not active: + self._repo_active.pop(repo, None) + if should_cancel and metadata is not None and marker_transport is not None: + metadata = replace(metadata, transport = marker_transport) + return terminal_state, metadata + + def update_job_transport(self, key: str, transport: str) -> None: + key = normalize_job_key(key) + with self._lock: + metadata = self._metadata.get(key) + if metadata is None or metadata.transport == transport: + return + self._metadata[key] = replace(metadata, transport = transport) + + def release_active_slot(self, key: str) -> None: + key = normalize_job_key(key) + repo = _repo_of_key(key) + with self._lock: + active = self._repo_active.get(repo) + if active is None: + return + active.discard(key) + if not active: + self._repo_active.pop(repo, None) + def get_job(self, key: str) -> DownloadState: key = normalize_job_key(key) with self._lock: @@ -884,6 +941,14 @@ class DownloadRegistry: ): self._put_terminal_job_locked(key, "cancelled") metadata_to_persist = self._metadata.pop(key, None) + marker_transport = self._cancel_marker_transports.pop(key, None) + if marker_transport is None and metadata_to_persist is not None: + marker_transport = metadata_to_persist.cancel_marker_transport + if metadata_to_persist is not None and marker_transport is not None: + metadata_to_persist = replace( + metadata_to_persist, + transport = marker_transport, + ) repo = _repo_of_key(key) active = self._repo_active.get(repo) if active is not None: @@ -963,12 +1028,24 @@ class DownloadRegistry: blob_hashes: Optional[frozenset[str]] = None, progress_blob_hashes: Optional[frozenset[str]] = None, completed_baseline_bytes: int = 0, + admission_check: Optional[Callable[[], bool]] = None, + generation: Optional[int] = None, + replace_active: bool = False, + metadata_transport: Optional[str] = None, + cancel_marker_transport: Optional[str] = None, ) -> tuple[bool, str]: key = normalize_job_key(key) repo = _repo_of_key(key) requested_hashes = blob_hashes or frozenset() requested_progress_hashes = progress_blob_hashes or frozenset() with self._lock: + # Run the final external admission check while the registry lock is + # held, immediately before inspecting and publishing active state. + # The GGUF load path establishes its marker before calling + # its active-job probe, so either this claim observes that marker + # or the load's later probe observes this claim. + if admission_check is not None and not admission_check(): + return False, "admission_blocked" deleting_scopes = self._deleting.get(repo) if deleting_scopes is not None and ( None in deleting_scopes or variant_from_key(key) in deleting_scopes @@ -1007,10 +1084,13 @@ class DownloadRegistry: if conflict_state is not None: return False, conflict_state current = self._jobs.get(key, DownloadState("idle")).state - if current in _ACTIVE_STATES: + if current in _ACTIVE_STATES and not replace_active: return False, current - self._generation_seq += 1 - self._generations[key] = self._generation_seq + if generation is None: + self._generation_seq += 1 + self._generations[key] = self._generation_seq + else: + self._generations[key] = generation self._jobs[key] = DownloadState("running") self._repo_active.setdefault(repo, active).add(key) if repo_type and repo_id: @@ -1018,7 +1098,8 @@ class DownloadRegistry: repo_type = repo_type, repo_id = repo_id, variant = variant, - transport = transport, + transport = metadata_transport if metadata_transport is not None else transport, + cancel_marker_transport = cancel_marker_transport, blob_hashes = requested_hashes, progress_blob_hashes = requested_progress_hashes, completed_baseline_bytes = max( @@ -1026,8 +1107,13 @@ class DownloadRegistry: int(completed_baseline_bytes or 0), ), ) + if cancel_marker_transport is not None: + self._cancel_marker_transports[key] = cancel_marker_transport + else: + self._cancel_marker_transports.pop(key, None) else: self._metadata.pop(key, None) + self._cancel_marker_transports.pop(key, None) return True, "running" def adoptable(self, key: str) -> bool: @@ -1053,7 +1139,8 @@ class DownloadRegistry: download. A variant delete conflicts only with that same variant or a whole-repo download writing the shared snapshot; other quantizations download concurrently and never block it.""" - for key in self._repo_active.get(repo_id, set()): + active_keys = self._repo_active.get(repo_id, set()) + for key in active_keys: job = self._jobs.get(key) if job is None or job.state not in _ACTIVE_STATES: continue @@ -1062,6 +1149,16 @@ class DownloadRegistry: other_variant = self._active_job_variant_locked(key) if other_variant is None or other_variant == variant: return True + for key, job in self._jobs.items(): + if key in active_keys or _repo_of_key(key) != repo_id: + continue + if job.state not in _ACTIVE_STATES: + continue + if variant is None: + return True + other_variant = self._active_job_variant_locked(key) + if other_variant is None or other_variant == variant: + return True return False def peer_blob_hashes(self, key: str) -> frozenset[str]: @@ -1108,6 +1205,16 @@ class DownloadRegistry: candidate_keys = list(self._repo_active.get(repo_key, set())) else: candidate_keys = [key for active in self._repo_active.values() for key in active] + # An XET->HTTP retry handoff briefly drops its key from _repo_active + # while its job stays active; include those released-but-active jobs + # so the waiting retry still lists and can be adopted or cancelled. + seen = set(candidate_keys) + for key, job in self._jobs.items(): + if key in seen or job.state not in _ACTIVE_STATES: + continue + if repo_key is not None and _repo_of_key(key) != repo_key: + continue + candidate_keys.append(key) refs: list[ActiveDownloadRef] = [] for key in candidate_keys: job = self._jobs.get(key) @@ -1123,6 +1230,23 @@ class DownloadRegistry: ) return refs + def has_active_variant(self, repo_id: str, variant: Optional[str]) -> bool: + """Whether an active model job targets this exact GGUF variant. + + Scans the job table rather than only ``_repo_active`` so an XET-to-HTTP + retry handoff remains visible while it has temporarily released its + active slot. + """ + repo_key = normalize_repo_key(repo_id) + target = (variant or "").strip().lower() or None + with self._lock: + for key, job in self._jobs.items(): + if _repo_of_key(key) != repo_key or job.state not in _ACTIVE_STATES: + continue + if self._active_job_variant_locked(key) == target: + return True + return False + def begin_delete( self, repo_id: str, @@ -1169,12 +1293,25 @@ class DownloadRegistry: repo_id = normalize_repo_key(repo_id) target = (variant or "").strip().lower() or None with self._lock: - for key in self._repo_active.get(repo_id, set()): + active_keys = self._repo_active.get(repo_id, set()) + for key in active_keys: job = self._jobs.get(key) if job is None or job.state not in _ACTIVE_STATES: continue if self._active_job_variant_locked(key) != target: return True + # An XET->HTTP retry peer between release_active_slot() and its reclaim + # is briefly absent from _repo_active while its job stays active and + # still owns the shared companion; mirror the released-but-active scan + # used by _delete_blocked_by_active_locked so it still blocks companion + # deletion of a different variant. + for key, job in self._jobs.items(): + if key in active_keys or _repo_of_key(key) != repo_id: + continue + if job.state not in _ACTIVE_STATES: + continue + if self._active_job_variant_locked(key) != target: + return True return False def request_cancel( @@ -1198,17 +1335,58 @@ class DownloadRegistry: return True def terminate_all(self, kind: str = "download") -> None: + settled_no_proc: list[Optional[DownloadMetadata]] = [] with self._lock: live = [ (key, proc, self._metadata.get(key)) for key, proc in self._processes.items() if proc.poll() is None ] + live_keys = {key for key, _proc, _metadata in live} # Flag as an intentional stop so the watcher's exit classification # reports them cancelled rather than an OOM/crash once SIGKILL lands. for key, _proc, _metadata in live: if self._jobs.get(key, DownloadState("idle")).state == "running": self._jobs[key] = DownloadState("cancelling") + # Settle active jobs without a live worker too. Two cases: an + # XET->HTTP retry parked in the reclaim wait loop has dropped its + # worker and slot guard, so it is absent from `live`; and a + # registered worker that already exited with an error but whose + # watcher has not yet run would otherwise stay `running` and spawn an + # HTTP retry after this shutdown snapshot. Skip a registered worker + # that exited cleanly (rc == 0): it completed and the watcher will + # mark it done, so marking it cancelling would strand a stale marker. + for key, job in list(self._jobs.items()): + if job.state not in _ACTIVE_STATES or key in live_keys: + continue + proc = self._processes.get(key) + if proc is not None: + if proc.poll() == 0: + continue + # A registered worker that exited nonzero on its own over HTTP + # is a genuine terminal download failure, not a shutdown cancel + # and not retry-capable: leave its error status intact rather + # than persisting a cancel marker that would read as + # cancelled/resumable after restart. Only an exited XET worker + # could still spawn a post-shutdown HTTP retry, so only that + # needs settling here. + metadata = self._metadata.get(key) + if metadata is not None and metadata.transport == TRANSPORT_HTTP: + continue + self._pending_cancel[key] = self._generations.get(key) + self._jobs[key] = DownloadState("cancelling") + settled_no_proc.append(self._metadata.get(key)) + # Persist a cancel marker for each settled no-live-worker job outside the + # lock (mirroring the reaped path) so shutdown records resumable/cancelled + # state even if it returns before the daemon watcher wakes to do so. + for metadata in settled_no_proc: + if metadata is not None: + persist_cancel_marker( + metadata.repo_type, + metadata.repo_id, + metadata.variant, + metadata.cancel_marker_transport or metadata.transport, + ) reaped: list[tuple[str, subprocess.Popen, Optional[DownloadMetadata]]] = [] for key, proc, metadata in live: try: @@ -1222,7 +1400,7 @@ class DownloadRegistry: metadata.repo_type, metadata.repo_id, metadata.variant, - metadata.transport, + metadata.cancel_marker_transport or metadata.transport, ) continue reaped.append((key, proc, metadata)) @@ -1242,7 +1420,7 @@ class DownloadRegistry: metadata.repo_type, metadata.repo_id, metadata.variant, - metadata.transport, + metadata.cancel_marker_transport or metadata.transport, ) diff --git a/studio/backend/main.py b/studio/backend/main.py index e64048dc00..6e16dc00ca 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -612,6 +612,22 @@ app = FastAPI( lifespan = lifespan, ) +# The MCP surface is opt-in because it can start GPU jobs and write model +# artifacts. Mount it only when explicitly enabled by the Studio process. +if os.environ.get("UNSLOTH_STUDIO_ENABLE_MCP") == "1": + from fastmcp.utilities.lifespan import combine_lifespans + + from mcp_server import BearerTokenMiddleware, create_studio_mcp + + _studio_mcp_app = create_studio_mcp().http_app(path = "/") + _studio_mcp_lifespan = _studio_mcp_app.lifespan + _mcp_token = os.environ.get("UNSLOTH_STUDIO_MCP_TOKEN") + if not _mcp_token: + raise RuntimeError("UNSLOTH_STUDIO_MCP_TOKEN is required when MCP is enabled") + _studio_mcp_app = BearerTokenMiddleware(_studio_mcp_app, _mcp_token) + app.router.lifespan_context = combine_lifespans(lifespan, _studio_mcp_lifespan) + app.mount("/mcp", _studio_mcp_app) + from loggers.config import LogConfig from loggers.handlers import LoggingMiddleware @@ -752,6 +768,7 @@ _BODY_PROTECTED_PREFIXES = ( "/api/settings", "/api/train", "/api/export", + "/mcp", ) _DATASET_UPLOAD_PASSTHROUGH_PREFIX = "/api/datasets/upload" _DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX = ( diff --git a/studio/backend/mcp_server.py b/studio/backend/mcp_server.py new file mode 100644 index 0000000000..f837f46425 --- /dev/null +++ b/studio/backend/mcp_server.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Curated MCP tools for driving an Unsloth Studio instance. + +The MCP surface deliberately wraps the existing Studio services instead of +duplicating training or export logic. It is opt-in because several tools can +start GPU work or write model artifacts. +""" + +from __future__ import annotations + +import hmac +from typing import Any + +from fastmcp import FastMCP + + +class BearerTokenMiddleware: + """Require an exact bearer token when Studio MCP is exposed remotely.""" + + def __init__(self, app: Any, token: str) -> None: + if not token or not token.strip(): + raise ValueError("Studio MCP bearer token must be a non-empty value") + if not token.isascii(): + # A non-ASCII token cannot be sent in an HTTP header; reject it here. + raise ValueError("Studio MCP bearer token must contain ASCII characters only") + self.app = app + # Compare on raw header bytes: str hmac.compare_digest raises on non-ASCII + # input, which would surface as a 500 instead of a clean 401. + self.expected = token.encode("utf-8") + + async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None: + scope_type = scope.get("type") + if scope_type not in ("http", "websocket"): + await self.app(scope, receive, send) + return + + headers = dict(scope.get("headers", [])) + raw_auth = headers.get(b"authorization", b"") + scheme, _, supplied = raw_auth.partition(b" ") + if scheme.lower() != b"bearer" or not hmac.compare_digest(supplied, self.expected): + await _send_unauthorized(send, scope_type) + return + + await self.app(scope, receive, send) + + +async def _send_unauthorized(send: Any, scope_type: str) -> None: + if scope_type == "websocket": + await send({"type": "websocket.close", "code": 4401}) + return + + await send( + { + "type": "http.response.start", + "status": 401, + "headers": [(b"content-type", b"application/json"), (b"www-authenticate", b"Bearer")], + } + ) + await send( + { + "type": "http.response.body", + "body": b'{"detail":"MCP bearer token required"}', + } + ) + + +def _dump(value: Any) -> Any: + """Convert Pydantic responses to plain JSON values for MCP clients.""" + if hasattr(value, "model_dump"): + return value.model_dump(mode = "json") + return value + + +def _clamp(value: int, low: int, high: int) -> int: + """Clamp an MCP-supplied integer into an inclusive range. + + MCP tools call the Studio route functions directly, which skips FastAPI's + Query(ge=, le=) validation, so we re-apply the same bounds here. + """ + return max(low, min(value, high)) + + +def create_studio_mcp() -> FastMCP: + """Create the Studio MCP server and register the high-value tools.""" + mcp = FastMCP( + "Unsloth Studio", + instructions = ( + "Use read tools to inspect the local Studio state before starting GPU work. " + "Training and export tools can consume substantial VRAM and write files. " + "Never expose tokens or local paths from tool results unless the user asks." + ), + ) + + @mcp.tool + async def studio_status() -> dict[str, Any]: + """Return the current training, export, inference, and GPU state.""" + from routes.export import get_export_status + from routes.inference import get_status as get_inference_status + from routes.training import get_training_status + + from utils.hardware import get_gpu_utilization + + training, export, inference = await _gather_status( + get_training_status(current_subject = "mcp"), + get_export_status(current_subject = "mcp"), + get_inference_status(current_subject = "mcp"), + ) + return { + "training": _dump(training), + "export": _dump(export), + "inference": _dump(inference), + "hardware": get_gpu_utilization(), + } + + @mcp.tool + async def list_local_models(models_dir: str = "./models") -> dict[str, Any]: + """List local and cached models available to Studio.""" + from routes.models import list_local_models as list_models + return _dump(await list_models(models_dir = models_dir, current_subject = "mcp")) + + @mcp.tool + async def get_training_status() -> dict[str, Any]: + """Read the active training job, phase, progress, and recent metrics.""" + from routes.training import get_training_status as get_status + return _dump(await get_status(current_subject = "mcp")) + + @mcp.tool + async def start_training(config: dict[str, Any]) -> dict[str, Any]: + """Start a validated Studio training job from a TrainingStartRequest-shaped object. + + The config is validated by the same Pydantic model used by the Studio UI. + Call get_training_status first and do not start work while another job runs. + """ + from models import TrainingStartRequest + from routes.training import start_training as start + + request = TrainingStartRequest.model_validate(config) + # Pass via_api_key explicitly (a direct call leaves it a Depends object). + # MCP drives Studio like the UI session, so it coexists and frees VRAM. + return _dump(await start(request, current_subject = "mcp", via_api_key = False)) + + @mcp.tool + async def stop_training(save: bool = True) -> dict[str, Any]: + """Ask the active training process to stop at its next safe checkpoint.""" + from routes.training import TrainingStopRequest, stop_training as stop + return _dump(await stop(TrainingStopRequest(save = save), current_subject = "mcp")) + + @mcp.tool + async def list_training_runs(limit: int = 50, offset: int = 0) -> dict[str, Any]: + """List completed and stopped training runs, newest first.""" + from routes.training_history import list_training_runs as list_runs + + # Clamp here (direct call skips Query bounds); a negative LIMIT = no limit. + limit = _clamp(limit, 1, 200) + offset = max(0, offset) + return _dump(await list_runs(limit = limit, offset = offset, current_subject = "mcp")) + + @mcp.tool + def validate_recipe(recipe: dict[str, Any]) -> dict[str, Any]: + """Validate a Data Recipe with the same validator used by Studio.""" + from models.data_recipe import RecipePayload + from routes.data_recipe.validate import validate + + return _dump(validate(RecipePayload(recipe = recipe))) + + @mcp.tool + def get_recipe_job_status(job_id: str) -> dict[str, Any]: + """Read the status of a Data Recipe job.""" + from routes.data_recipe.jobs import job_status + return _dump(job_status(job_id)) + + @mcp.tool + def get_recipe_job_dataset( + job_id: str, + limit: int = 20, + offset: int = 0, + ) -> dict[str, Any]: + """Read a bounded page of generated Data Recipe rows.""" + from routes.data_recipe.jobs import job_dataset + + # Clamp here (direct call skips FastAPI's Query bounds). + limit = _clamp(limit, 1, 500) + offset = max(0, offset) + return _dump(job_dataset(job_id, limit = limit, offset = offset)) + + @mcp.tool + async def load_checkpoint( + checkpoint_path: str, + max_seq_length: int = 2048, + load_in_4bit: bool = True, + trust_remote_code: bool = False, + approved_remote_code_fingerprint: str | None = None, + hf_token: str | None = None, + ) -> dict[str, Any]: + """Load a checkpoint into the export backend. + + Export runs in its own subprocess and coexists with training and + inference; it does not unload them, so a load can fail with a clear + out-of-memory error if the GPU is already full. Pass hf_token to load a + gated checkpoint, and approved_remote_code_fingerprint to retry a + trust_remote_code load that was blocked pending review. + """ + from models import LoadCheckpointRequest + from routes.export import load_checkpoint as load + + request = LoadCheckpointRequest( + checkpoint_path = checkpoint_path, + max_seq_length = max_seq_length, + load_in_4bit = load_in_4bit, + trust_remote_code = trust_remote_code, + approved_remote_code_fingerprint = approved_remote_code_fingerprint, + hf_token = hf_token, + ) + return _dump(await load(request, current_subject = "mcp")) + + @mcp.tool + async def export_gguf( + save_directory: str, + quantization_method: str | list[str] = "Q4_K_M", + push_to_hub: bool = False, + repo_id: str | None = None, + hf_token: str | None = None, + imatrix: bool = False, + imatrix_path: str | None = None, + ) -> dict[str, Any]: + """Export the loaded model to GGUF using Studio's existing path validation. + + quantization_method may be a single method or a list to produce several + GGUFs from one load. Pass hf_token when push_to_hub is set (the backend + rejects a Hub upload without it). Set imatrix (or imatrix_path) for the + IQ low-bit quants that require an importance matrix. + """ + from models import ExportGGUFRequest + from routes.export import export_gguf as export + + request = ExportGGUFRequest( + save_directory = save_directory, + quantization_method = quantization_method, + push_to_hub = push_to_hub, + repo_id = repo_id, + hf_token = hf_token, + imatrix = imatrix, + imatrix_path = imatrix_path, + ) + return _dump(await export(request, current_subject = "mcp")) + + return mcp + + +async def _gather_status(*coroutines: Any) -> tuple[Any, ...]: + """Gather independent status calls without letting one optional backend fail all state.""" + import asyncio + + results = await asyncio.gather(*coroutines, return_exceptions = True) + return tuple( + {"error": str(result)} if isinstance(result, Exception) else result for result in results + ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 0541a81ff2..52ea1f86a3 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -20,6 +20,7 @@ from loggers import get_logger import asyncio import threading import weakref +from contextlib import ExitStack import re as _re @@ -28,6 +29,7 @@ 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.orchestrator import GenStreamError, GenStreamErrorRaised from core.inference.llama_admission import ( LlamaAdmissionCancelled, LlamaAdmissionConfig, @@ -212,6 +214,14 @@ def _friendly_error(exc: Exception) -> str: return "An internal error occurred" +def _friendly_gen_stream_error(value) -> str: + """Return a client-safe message for typed local generation errors.""" + text = str(value) + if getattr(value, "public", False): + return text + return safe_error_detail(RuntimeError(text), fallback = "An internal error occurred.") + + def _friendly_upstream_error(text: str) -> str: """Rewrite a raw llama-server error body into an actionable message where we can. @@ -998,6 +1008,7 @@ try: from core.inference.llama_server_args import ( _effective_tensor_parallel, _tensor_parallel_matches_loaded, + extra_args_disable_mmproj, parse_split_mode_override, resolve_tensor_parallel, strip_shadowing_flags, @@ -1035,6 +1046,7 @@ except ImportError: from core.inference.llama_server_args import ( _effective_tensor_parallel, _tensor_parallel_matches_loaded, + extra_args_disable_mmproj, parse_split_mode_override, resolve_tensor_parallel, strip_shadowing_flags, @@ -1915,16 +1927,57 @@ async def artifact_preview_frame(allow_network: bool = False): _BARE_JSON_NAME_MARKER_RE = _re.compile(r'\{\s*\\?"(?:name|function)\\?"\s*:') -def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: +def _detect_safetensors_features( + backend, + chat_template: Optional[str], + tools = None, +) -> dict: """Classify reasoning/tool capabilities via the GGUF classifier so flags match across backends. gpt-oss is overridden: Harmony routes reasoning and tools through tokenizer channels, not template markup.""" model_id = getattr(backend, "active_model_name", None) + feature_template = chat_template + try: + from core.inference.chat_template_helpers import _selected_template_strings_from_value + selected_templates = _selected_template_strings_from_value(chat_template, tools) + if selected_templates: + feature_template = selected_templates[0] + except Exception: + logger.debug("safetensors_named_template_selection_failed", exc_info = True) flags = detect_reasoning_flags( - chat_template, + feature_template, model_identifier = model_id, log_source = "safetensors", ) + if not flags.get("supports_reasoning"): + try: + from core.inference.chat_template_helpers import ( + detect_reasoning_channel_markers_from_template, + ) + + templates = [chat_template] + models = getattr(backend, "models", None) + model_info = ( + models.get(model_id, {}) + if isinstance(models, dict) and model_id is not None + else {} + ) + if isinstance(model_info, dict): + templates.extend( + ( + model_info.get("native_chat_template"), + (model_info.get("chat_template_info") or {}).get("template"), + ) + ) + if any( + detect_reasoning_channel_markers_from_template(template, tools = tools) is not None + for template in templates + ): + flags["supports_reasoning"] = True + flags["reasoning_always_on"] = True + logger.info("safetensors: model always reasons (native channel markers)") + except Exception: + logger.debug("safetensors_native_reasoning_marker_check_failed", exc_info = True) # 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); @@ -1938,9 +1991,9 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: ) if ( flags.get("supports_tools") - and chat_template - and not any(m in chat_template for m in _PARSER_MARKERS) - and not _BARE_JSON_NAME_MARKER_RE.search(chat_template) + and isinstance(feature_template, str) + and not any(m in feature_template for m in _PARSER_MARKERS) + and not _BARE_JSON_NAME_MARKER_RE.search(feature_template) ): logger.info( "safetensors: template advertises tools but uses an " @@ -3968,6 +4021,7 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre native_grant_backed = False model_log_label = request.model_path + gguf_load_stack = ExitStack() try: # Validate user pass-through args up front so a managed-flag collision # returns 400 before any model work. @@ -4187,15 +4241,9 @@ 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 (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" - ) - await asyncio.to_thread( - unsloth_backend.unload_model, unsloth_backend.active_model_name - ) + if config.gguf_hf_repo: + from core.inference.llama_cpp import gguf_load_in_flight + gguf_load_stack.enter_context(gguf_load_in_flight(config.gguf_hf_repo)) # Inherit llama_extra_args from the previous load when the request # omits the field (the chat-settings Apply path doesn't round-trip @@ -4275,6 +4323,38 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre extra_llama_args, ) + # Block cache writes that would race the download manager. This runs + # after pass-through argument inheritance so a carried --no-mmproj + # changes the companion requirement exactly as it does for the load. + if config.gguf_hf_repo: + from core.inference.llama_cpp import _hub_download_blocks_gguf_load + if await asyncio.to_thread( + _hub_download_blocks_gguf_load, + config.gguf_hf_repo, + config.gguf_variant, + require_mmproj = bool( + config.is_vision and not extra_args_disable_mmproj(extra_llama_args) + ), + hf_token = request.hf_token, + ): + raise HTTPException( + status_code = 409, + detail = ( + f"'{model_log_label}' is currently being downloaded " + "by the download manager. Wait for the download to " + "finish (or cancel it), then load the model." + ), + ) + + # Unload any active Unsloth model only after every hub conflict check. + if unsloth_backend.active_model_name: + logger.info( + f"Unloading Unsloth model '{unsloth_backend.active_model_name}' before loading GGUF" + ) + await asyncio.to_thread( + unsloth_backend.unload_model, unsloth_backend.active_model_name + ) + # Route to HF or local mode based on config. Run in a thread so the # event loop stays free for progress polling and other requests # during the (potentially long) GGUF download + llama-server start. @@ -4639,6 +4719,8 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre logger.error(f"Error loading model: {e}", exc_info = True) msg = _maybe_unsupported_message(redacted_msg) raise HTTPException(status_code = 500, detail = f"Failed to load model: {msg}") + finally: + gguf_load_stack.close() def _requires_trust_remote_code_for_model( @@ -5392,6 +5474,10 @@ async def generate_stream( if chunk is _DONE: completed = True break + if isinstance(chunk, GenStreamError): + yield f"data: {json.dumps({'error': _friendly_gen_stream_error(chunk)})}\n\n" + yield "data: [DONE]\n\n" + return yield f"data: {json.dumps({'content': chunk})}\n\n" if completed: yield "data: [DONE]\n\n" @@ -5405,6 +5491,7 @@ async def generate_stream( backend.reset_generation_state() logger.error(f"Error during generation: {e}", exc_info = True) yield f"data: {json.dumps({'error': _friendly_error(e)})}\n\n" + yield "data: [DONE]\n\n" finally: await _stop_local_disconnect_cancel_watcher(disconnect_watcher) if not completed and not cancel_event.is_set(): @@ -7028,6 +7115,13 @@ async def openai_chat_completions( chunk_text = await asyncio.to_thread(next, gen, _DONE) if chunk_text is _DONE: break + if isinstance(chunk_text, GenStreamError): + _msg = _friendly_gen_stream_error(chunk_text) + api_monitor.fail(monitor_id, _msg) + yield _openai_stream_error_sse( + {"error": {"message": _msg, "type": "server_error"}} + ) + return if chunk_text: api_monitor.append_reply(monitor_id, chunk_text) yield _chat_content_chunk( @@ -7043,8 +7137,11 @@ async def openai_chat_completions( raise except Exception as e: logger.error(f"Error during audio input streaming: {e}", exc_info = True) - api_monitor.fail(monitor_id, _friendly_error(e)) - yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n" + _msg = _friendly_error(e) + api_monitor.fail(monitor_id, _msg) + yield _openai_stream_error_sse( + {"error": {"message": _msg, "type": "server_error"}} + ) finally: await _stop_local_disconnect_cancel_watcher(disconnect_watcher) _tracker.__exit__(None, None, None) @@ -7061,7 +7158,15 @@ async def openai_chat_completions( ) else: try: - full_text = "".join(audio_input_generate()) + full_text = "" + for chunk_text in audio_input_generate(): + if isinstance(chunk_text, GenStreamError): + _msg = _friendly_gen_stream_error(chunk_text) + api_monitor.fail(monitor_id, _msg) + raise HTTPException(status_code = 500, detail = _msg) + full_text += chunk_text + except HTTPException: + raise except Exception as e: api_monitor.fail(monitor_id, _friendly_error(e)) raise @@ -8605,19 +8710,33 @@ async def openai_chat_completions( # Classify capability flags from the loaded template. _sf_model_info = backend.models.get(backend.active_model_name, {}) _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") + # Named templates may expose native reasoning only in their ``tool_use`` + # branch. Use a truthy placeholder for Studio-managed tools, whose concrete + # schemas are selected below, and the request schemas for client passthrough. + _sf_server_tool_intent = bool( + _effective_enable_tools(payload) or _explicit_studio_tool_loop_requested(payload) ) - # 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, + _sf_template_tools = payload.tools if payload.tool_choice != "none" else None + if not _sf_template_tools and _sf_server_tool_intent: + _sf_template_tools = ({},) + + def _sf_response_protocol(tools = None): + features = _detect_safetensors_features(backend, _sf_tpl, tools = tools) + parse_think = bool( + features.get("supports_reasoning") or features.get("reasoning_always_on") + ) + reasoning_prefilled = _sf_reasoning_prefill_mode( + features, + payload.enable_thinking, + _sf_tpl, + reasoning_effort = payload.reasoning_effort, + ) + return features, parse_think, reasoning_prefilled + + # GGUF parity: split canonical output into reasoning_content. The + # selected template branch must match whether this request renders tools. + _sf_features, _sf_parse_think, _sf_reasoning_prefilled = _sf_response_protocol( + _sf_template_tools ) def _new_sf_reasoning_extractor(): @@ -8767,6 +8886,7 @@ async def openai_chat_completions( permission_mode = payload.permission_mode, use_adapter = payload.use_adapter, stats_holder = _sf_stats_holder, + reasoning_prefilled = _sf_reasoning_prefilled, ) _sf_tool_sentinel = object() @@ -8826,6 +8946,18 @@ async def openai_chat_completions( _sf_next_task = None if event is _sf_tool_sentinel: break + if isinstance(event, GenStreamError): + backend.reset_generation_state() + _msg = _friendly_gen_stream_error(event) + api_monitor.fail(monitor_id, _msg) + yield _openai_stream_error_sse( + {"error": {"message": _msg, "type": "server_error"}} + ) + return + if not isinstance(event, dict): + raise RuntimeError( + f"Invalid safetensors tool event: {type(event).__name__}" + ) if event["type"] == "heartbeat": # Tool-execution wrapper heartbeat -> SSE keepalive. @@ -8913,6 +9045,11 @@ async def openai_chat_completions( backend.reset_generation_state() api_monitor.finish(monitor_id, "cancelled") raise + except GenStreamErrorRaised as exc: + backend.reset_generation_state() + _msg = _friendly_gen_stream_error(exc) + api_monitor.fail(monitor_id, _msg) + yield _openai_stream_error_sse({"error": {"message": _msg, "type": "server_error"}}) except Exception: backend.reset_generation_state() # Generic wire message; full trace stays in the log (CWE-209: @@ -8962,6 +9099,15 @@ async def openai_chat_completions( for event in gen: if cancel_event.is_set(): break + if isinstance(event, GenStreamError): + raise HTTPException( + status_code = 500, + detail = _friendly_gen_stream_error(event), + ) + if not isinstance(event, dict): + raise RuntimeError( + f"Invalid safetensors tool event: {type(event).__name__}" + ) if event.get("type") == "content": full_text = _strip_tool_xml_for_display( event.get("text", ""), @@ -9002,6 +9148,15 @@ async def openai_chat_completions( backend.reset_generation_state() api_monitor.finish(monitor_id, "cancelled") raise + except GenStreamErrorRaised as exc: + backend.reset_generation_state() + _msg = _friendly_gen_stream_error(exc) + api_monitor.fail(monitor_id, _msg) + raise HTTPException(status_code = 500, detail = _msg) + except HTTPException as exc: + backend.reset_generation_state() + api_monitor.fail(monitor_id, str(exc.detail)) + raise except Exception: backend.reset_generation_state() # CWE-209: generic detail; full trace in log. @@ -9088,6 +9243,12 @@ async def openai_chat_completions( else: gen_kwargs["tools"] = payload.tools + # The potential tool context above is needed before server/client routing is + # known. This standard path now has the exact schemas that will be rendered, + # so resolve reasoning parsing again to keep empty registries, forced-tool + # misses, and tool_choice="none" on the marker-free template branch. + _, _sf_parse_think, _sf_reasoning_prefilled = _sf_response_protocol(gen_kwargs.get("tools")) + # Request-scoped usage/timings receptacle (filled at gen_done). stats_holder: dict = {} @@ -9168,6 +9329,14 @@ async def openai_chat_completions( _next_task = None if cumulative is _DONE: break + if isinstance(cumulative, GenStreamError): + backend.reset_generation_state() + _msg = _friendly_gen_stream_error(cumulative) + api_monitor.fail(monitor_id, _msg) + yield _openai_stream_error_sse( + {"error": {"message": _msg, "type": "server_error"}} + ) + return if await request.is_disconnected(): cancel_event.set() backend.reset_generation_state() @@ -9317,6 +9486,11 @@ async def openai_chat_completions( try: full_text = "" for token in generate(): + if isinstance(token, GenStreamError): + backend.reset_generation_state() + _msg = _friendly_gen_stream_error(token) + api_monitor.fail(monitor_id, _msg) + raise HTTPException(status_code = 500, detail = _msg) full_text = token # Split prefilled reasoning (GGUF parity); also covers MLX via @@ -9415,6 +9589,8 @@ async def openai_chat_completions( api_monitor.finish(monitor_id) return _model_json_response(response) + except HTTPException: + raise except Exception as e: backend.reset_generation_state() logger.error(f"Error during OpenAI completion: {e}", exc_info = True) diff --git a/studio/backend/run.py b/studio/backend/run.py index 56b9c78343..4f105b53b1 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -1354,11 +1354,23 @@ def run_server( if secure: os.environ["UNSLOTH_SECURE"] = "1" - import nest_asyncio - - nest_asyncio.apply() - import asyncio + + # nest_asyncio is for Colab/IPython, where the main thread already runs a loop + # the blocking waits below would collide with. Apply it only with a loop running + # (a plain CLI start has nothing to nest) and only on Python <= 3.13: on 3.14+ + # its global Task patch leaves asyncio.current_task() None (tracking moved into + # C), which also breaks the background uvicorn loop and 500s every request. It + # is archived upstream, so no 3.14 fix is coming; skip it there. + if sys.version_info < (3, 14): + try: + asyncio.get_running_loop() + except RuntimeError: + pass + else: + import nest_asyncio + nest_asyncio.apply() + from threading import Thread, Event import uvicorn diff --git a/studio/backend/tests/test_deepseek_v4_thinking_effort.py b/studio/backend/tests/test_deepseek_v4_thinking_effort.py index 19808ad0d7..0d60d9b5ec 100644 --- a/studio/backend/tests/test_deepseek_v4_thinking_effort.py +++ b/studio/backend/tests/test_deepseek_v4_thinking_effort.py @@ -16,7 +16,6 @@ from __future__ import annotations import sys from pathlib import Path -from types import SimpleNamespace import pytest @@ -128,15 +127,13 @@ 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 {} + shim = object.__new__(LlamaCppBackend) + shim._supports_reasoning = flags["supports_reasoning"] + shim._reasoning_always_on = flags["reasoning_always_on"] + shim._reasoning_style = flags["reasoning_style"] + shim._reasoning_effort_levels = flags["reasoning_effort_levels"] + shim._supports_preserve_thinking = flags["supports_preserve_thinking"] + return shim._request_reasoning_kwargs(enable_thinking, reasoning_effort, None) or {} def _flags(): diff --git a/studio/backend/tests/test_gguf_load_cache_reuse.py b/studio/backend/tests/test_gguf_load_cache_reuse.py new file mode 100644 index 0000000000..15d91cd324 --- /dev/null +++ b/studio/backend/tests/test_gguf_load_cache_reuse.py @@ -0,0 +1,740 @@ +# 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 cached GGUF reuse and load/download exclusion. + +No GPU, network, or subprocesses are required. +""" + +from __future__ import annotations + +import asyncio +import sys +import threading +import types as _types +from pathlib import Path +from unittest.mock import patch + +import pytest + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Stub optional dependencies before importing the modules under test. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +_structlog_stub = _types.ModuleType("structlog") +sys.modules.setdefault("structlog", _structlog_stub) + +try: + import httpx # noqa: F401 +except ImportError: + _httpx_stub = _types.ModuleType("httpx") + for _exc_name in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", + "HTTPError", + "RequestError", + "HTTPStatusError", + ): + setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) + _httpx_stub.Response = type("Response", (), {}) + _httpx_stub.Request = type("Request", (), {}) + + class _FakeTimeout: + def __init__(self, *a, **kw): + pass + + _httpx_stub.Timeout = _FakeTimeout + _httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **kw: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, + ) + sys.modules.setdefault("httpx", _httpx_stub) + + +from huggingface_hub import constants as hf_constants + +from core.inference.llama_cpp import ( + LlamaCppBackend, + cached_gguf_for_load, + gguf_load_in_flight, + hf_gguf_load_in_flight, +) + + +REPO = "unsloth/gemma-test-GGUF" +VARIANT = "UD-Q4_K_XL" +MAIN = f"gemma-test-{VARIANT}.gguf" + + +def _build_cache( + root: Path, + repo_id: str, + files: dict[str, int], + *, + snapshot_sha: str = "a" * 40, +) -> Path: + """Create ``$root/models--/snapshots//`` for each entry.""" + repo_dir = root / f"models--{repo_id.replace('/', '--')}" + (repo_dir / "blobs").mkdir(parents = True, exist_ok = True) + snap = repo_dir / "snapshots" / snapshot_sha + snap.mkdir(parents = True, exist_ok = True) + for rel, size in files.items(): + full = snap / rel + full.parent.mkdir(parents = True, exist_ok = True) + full.write_bytes(b"\0" * size) + return snap + + +@pytest.fixture +def hf_cache(tmp_path, monkeypatch): + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + return tmp_path + + +def _fail_download(*_args, **_kwargs): + raise AssertionError("must reuse the cached GGUF instead of downloading") + + +def _fail_get_paths_info(*_args, **_kwargs): + raise AssertionError("cached reuse must return before the sizing preflight") + + +class TestLoadReusesCachedCopy: + def test_online_reuse_after_revision_bump(self, hf_cache): + """A new repo revision does not replace a complete cached model.""" + backend = LlamaCppBackend() + snap = _build_cache(hf_cache, REPO, {MAIN: 4}) + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", _fail_get_paths_info), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(snap / MAIN) + + def test_reuse_size_check_uses_cached_snapshot_revision(self, hf_cache): + """Current-revision size changes do not invalidate an older complete copy.""" + backend = LlamaCppBackend() + snap = _build_cache(hf_cache, REPO, {MAIN: 4}) + revisions: list[str | None] = [] + + def fake_get_paths_info( + _repo, + paths, + *, + revision = None, + token = None, + ): + revisions.append(revision) + size = 4 if revision == snap.name else 8 + return [_types.SimpleNamespace(path = path, size = size) for path in paths] + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(snap / MAIN) + assert revisions == [snap.name] + + def test_reuse_when_cached_revision_vanished_from_hub(self, hf_cache): + """The Hub answers an unknown revision with an empty result, not an error.""" + backend = LlamaCppBackend() + snap = _build_cache(hf_cache, REPO, {MAIN: 4}) + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", lambda *_a, **_k: []), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(snap / MAIN) + + def test_truncated_cached_file_is_not_reused(self, hf_cache): + backend = LlamaCppBackend() + _build_cache(hf_cache, REPO, {MAIN: 4}) + downloaded: list[str] = [] + + def fake_get_paths_info( + _repo, + paths, + *, + revision = None, + token = None, + ): + return [_types.SimpleNamespace(path = path, size = 8) for path in paths] + + 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: [MAIN]), + 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 = VARIANT) + + assert downloaded == [MAIN] + assert out == f"/fake/{REPO}/{MAIN}" + + def test_truncated_cached_split_shard_is_not_reused(self, hf_cache): + backend = LlamaCppBackend() + shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf" + shard2 = f"gemma-test-{VARIANT}-00002-of-00002.gguf" + _build_cache(hf_cache, REPO, {shard1: 8, shard2: 4}) + downloaded: list[str] = [] + + def fake_get_paths_info( + _repo, + paths, + *, + revision = None, + token = None, + ): + return [_types.SimpleNamespace(path = path, size = 8) for path in paths] + + 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: [shard1, shard2]), + 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 = VARIANT) + + assert downloaded == [shard1, shard2] + assert out == f"/fake/{REPO}/{shard1}" + + def test_online_reuse_when_reupload_renamed_the_file(self, hf_cache): + """A renamed variant still reuses its cached file.""" + backend = LlamaCppBackend() + old_name = f"gemma-test-old-{VARIANT}.gguf" + snap = _build_cache(hf_cache, REPO, {old_name: 4}) + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", _fail_get_paths_info), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(snap / old_name) + + def test_downloads_when_nothing_cached(self, hf_cache): + backend = LlamaCppBackend() + downloaded: list[str] = [] + + def fake_download( + repo_id, + filename, + token = None, + **_kwargs, + ): + downloaded.append(filename) + return f"/fake/{repo_id}/{filename}" + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 1) for p in paths if p is not None] + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + 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 = VARIANT) + + assert downloaded == [MAIN] + assert out == f"/fake/{REPO}/{MAIN}" + + def test_force_redownloads_despite_cache(self, hf_cache): + """A forced download ignores a complete cached copy.""" + backend = LlamaCppBackend() + _build_cache(hf_cache, REPO, {MAIN: 4}) + downloaded: list[str] = [] + + def fake_download( + repo_id, + filename, + token = None, + **kwargs, + ): + assert kwargs.get("force_download") is True + downloaded.append(filename) + return f"/fake/{repo_id}/{filename}" + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 1) for p in paths if p is not None] + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + 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 = VARIANT, force = True) + + assert downloaded == [MAIN] + assert out == f"/fake/{REPO}/{MAIN}" + + def test_split_reused_only_when_colocated(self, hf_cache): + backend = LlamaCppBackend() + shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf" + shard2 = f"gemma-test-{VARIANT}-00002-of-00002.gguf" + snap = _build_cache(hf_cache, REPO, {shard1: 4, shard2: 4}) + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [shard1, shard2]), + patch("huggingface_hub.get_paths_info", _fail_get_paths_info), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(snap / shard1) + + def test_partial_split_set_downloads(self, hf_cache): + """A partial split set is not reused.""" + backend = LlamaCppBackend() + shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf" + shard2 = f"gemma-test-{VARIANT}-00002-of-00002.gguf" + _build_cache(hf_cache, REPO, {shard1: 4}) + downloaded: list[str] = [] + + def fake_download( + repo_id, + filename, + token = None, + **_kwargs, + ): + downloaded.append(filename) + return f"/fake/{repo_id}/{filename}" + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p is not None] + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [shard1, shard2]), + 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 = VARIANT) + + assert downloaded == [shard1, shard2] + assert out == f"/fake/{REPO}/{shard1}" + + def test_reuse_prefers_newest_snapshot_after_update(self, hf_cache): + """Loads prefer the newest complete snapshot.""" + import os + + backend = LlamaCppBackend() + old_snap = _build_cache(hf_cache, REPO, {MAIN: 4}, snapshot_sha = "a" * 40) + new_snap = _build_cache(hf_cache, REPO, {MAIN: 6}, snapshot_sha = "b" * 40) + os.utime(old_snap, (1_000_000, 1_000_000)) + os.utime(new_snap, (2_000_000, 2_000_000)) + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", _fail_get_paths_info), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(new_snap / MAIN) + + def test_low_disk_fallback_reuses_cached_copy(self, hf_cache): + backend = LlamaCppBackend() + fallback = "gemma-test-Q2_K.gguf" + snap = _build_cache(hf_cache, REPO, {fallback: 4}) + + def fake_get_paths_info( + _repo, + paths, + *, + revision = None, + token = None, + ): + size = 4 if revision == snap.name else 100 + return [_types.SimpleNamespace(path = path, size = size) for path in paths] + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("shutil.disk_usage", lambda *_a, **_k: _types.SimpleNamespace(free = 10)), + patch.object( + backend, + "_find_smallest_fitting_variant", + lambda *_a, **_k: (fallback, 4, []), + ), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(snap / fallback) + + def test_companion_prefers_main_snapshot_sibling(self, hf_cache): + """A cached mmproj is reused from the main model's snapshot.""" + backend = LlamaCppBackend() + snap = _build_cache(hf_cache, REPO, {MAIN: 4, "mmproj-F16.gguf": 2}) + + def _fail_list(*_args, **_kwargs): + raise AssertionError("snapshot sibling must resolve without a repo listing") + + with patch("huggingface_hub.list_repo_files", _fail_list): + out = backend._download_mmproj(hf_repo = REPO, near_path = str(snap / MAIN)) + + assert out == str(snap / "mmproj-F16.gguf") + + def test_companion_finds_snapshot_through_hf_symlink(self, hf_cache): + backend = LlamaCppBackend() + snap = _build_cache(hf_cache, REPO, {}) + blobs = snap.parent.parent / "blobs" + main_blob = blobs / "main" + mmproj_blob = blobs / "mmproj" + main_blob.write_bytes(b"main") + mmproj_blob.write_bytes(b"mmproj") + try: + (snap / MAIN).symlink_to(main_blob) + (snap / "mmproj-F16.gguf").symlink_to(mmproj_blob) + except OSError as exc: + pytest.skip(f"symlinks unavailable: {exc}") + + with patch("huggingface_hub.list_repo_files", _fail_download): + out = backend._download_mmproj(hf_repo = REPO, near_path = str(snap / MAIN)) + + assert out == str(snap / "mmproj-F16.gguf") + + def test_companion_does_not_download_during_hub_job(self, hf_cache): + backend = LlamaCppBackend() + snap = _build_cache(hf_cache, REPO, {MAIN: 4}) + registry = _types.SimpleNamespace(active_job_refs = lambda _repo: [object()]) + + with ( + patch("huggingface_hub.list_repo_files", _fail_download), + patch("hub.utils.download_registry.get_models_registry", lambda: registry), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_mmproj(hf_repo = REPO, near_path = str(snap / MAIN)) + + assert out is None + + +class TestCachedGgufForLoadProbe: + def test_complete_copy_found(self, hf_cache): + snap = _build_cache(hf_cache, REPO, {MAIN: 4}) + assert cached_gguf_for_load(REPO, VARIANT) == str(snap / MAIN) + + def test_absent_copy_is_none(self, hf_cache): + assert cached_gguf_for_load(REPO, VARIANT) is None + + def test_partial_split_is_none(self, hf_cache): + shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf" + _build_cache(hf_cache, REPO, {shard1: 4}) + assert cached_gguf_for_load(REPO, VARIANT) is None + + def test_partial_new_snapshot_does_not_hide_complete_split(self, hf_cache): + import os + + shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf" + shard2 = f"gemma-test-{VARIANT}-00002-of-00002.gguf" + old = _build_cache( + hf_cache, + REPO, + {shard1: 4, shard2: 4}, + snapshot_sha = "a" * 40, + ) + new = _build_cache(hf_cache, REPO, {shard1: 4}, snapshot_sha = "b" * 40) + os.utime(old, (1_000_000, 1_000_000)) + os.utime(new, (2_000_000, 2_000_000)) + + assert cached_gguf_for_load(REPO, VARIANT) == str(old / shard1) + + def test_split_requires_every_declared_shard(self, hf_cache): + shard1 = f"gemma-test-{VARIANT}-00001-of-00003.gguf" + shard2 = f"gemma-test-{VARIANT}-00002-of-00003.gguf" + _build_cache(hf_cache, REPO, {shard1: 4, shard2: 4}) + + assert cached_gguf_for_load(REPO, VARIANT) is None + + def test_required_mmproj_must_share_main_snapshot(self, hf_cache): + snap = _build_cache(hf_cache, REPO, {MAIN: 4}) + assert cached_gguf_for_load(REPO, VARIANT) == str(snap / MAIN) + assert cached_gguf_for_load(REPO, VARIANT, require_mmproj = True) is None + + (snap / "mmproj-F16.gguf").write_bytes(b"mmproj") + assert cached_gguf_for_load(REPO, VARIANT, require_mmproj = True) == str(snap / MAIN) + + def test_required_mmproj_scans_past_newer_main_only_snapshot(self, hf_cache): + import os + + old = _build_cache( + hf_cache, + REPO, + {MAIN: 4, "mmproj-F16.gguf": 2}, + snapshot_sha = "a" * 40, + ) + new = _build_cache(hf_cache, REPO, {MAIN: 4}, snapshot_sha = "b" * 40) + os.utime(old, (1_000_000, 1_000_000)) + os.utime(new, (2_000_000, 2_000_000)) + + assert cached_gguf_for_load(REPO, VARIANT, require_mmproj = True) == str(old / MAIN) + + +class TestLoadHubDownloadExclusion: + def test_in_flight_marker_counts_and_normalizes_case(self): + assert not hf_gguf_load_in_flight(REPO) + with gguf_load_in_flight(REPO): + assert hf_gguf_load_in_flight(REPO.upper()) + with gguf_load_in_flight(REPO.lower()): + assert hf_gguf_load_in_flight(REPO) + assert hf_gguf_load_in_flight(REPO) + assert not hf_gguf_load_in_flight(REPO) + + def test_marker_noops_for_local_loads(self): + with gguf_load_in_flight(None): + assert not hf_gguf_load_in_flight("") + + def test_marker_cleared_on_exception(self): + with pytest.raises(RuntimeError): + with gguf_load_in_flight(REPO): + raise RuntimeError("boom") + assert not hf_gguf_load_in_flight(REPO) + + def test_hub_download_refused_while_load_in_flight(self): + from fastapi import HTTPException + + from hub.schemas.downloads import DownloadModelRequest + from hub.services.models import downloads as dl + + body = DownloadModelRequest(repo_id = REPO, gguf_variant = VARIANT) + with ( + patch.object(dl, "resolve_cached_repo_id_case", lambda repo_id, repo_type: repo_id), + gguf_load_in_flight(REPO), + ): + with pytest.raises(HTTPException) as exc_info: + asyncio.run(dl.download_model_response(body)) + + assert exc_info.value.status_code == 409 + assert "load" in exc_info.value.detail.lower() + + def test_hub_download_rechecks_marker_before_claim(self): + from fastapi import HTTPException + + from hub.schemas.downloads import DownloadModelRequest + from hub.services.models import downloads as dl + + scope = None + + def mark_load(*_args, **_kwargs): + nonlocal scope + if scope is None: + scope = gguf_load_in_flight(REPO) + scope.__enter__() + return frozenset() + + class _Registry: + def claim(self, *_args, admission_check, **_kwargs): + assert admission_check() is False + return False, "admission_blocked" + + def current_generation(self, _key): + return 0 + + registry = _Registry() + body = DownloadModelRequest(repo_id = REPO, gguf_variant = VARIANT) + try: + with ( + patch.object(dl, "resolve_cached_repo_id_case", lambda repo_id, repo_type: repo_id), + patch.object(dl.gguf_variants, "gguf_variant_blob_hashes", mark_load), + patch.object(dl, "_registry", registry), + ): + with pytest.raises(HTTPException) as exc_info: + asyncio.run(dl.download_model_response(body)) + finally: + if scope is not None: + scope.__exit__(None, None, None) + + assert exc_info.value.status_code == 409 + + def test_registry_admission_check_prevents_claim(self): + from hub.utils.download_registry import DownloadRegistry, TRANSPORT_HTTP + + registry = DownloadRegistry() + claimed, state = registry.claim( + f"{REPO}::{VARIANT}", + TRANSPORT_HTTP, + repo_type = "model", + repo_id = REPO, + variant = VARIANT, + admission_check = lambda: False, + ) + + assert claimed is False + assert state == "admission_blocked" + assert registry.active_jobs(REPO) == {} + + def test_same_variant_job_stays_visible_during_retry_handoff(self): + from hub.utils.download_registry import DownloadRegistry, TRANSPORT_XET + from core.inference.llama_cpp import _hub_download_blocks_gguf_load + + registry = DownloadRegistry() + key = f"{REPO}::{VARIANT}" + claimed, _ = registry.claim( + key, + TRANSPORT_XET, + repo_type = "model", + repo_id = REPO, + variant = VARIANT, + ) + assert claimed is True + assert registry.has_active_variant(REPO, VARIANT.lower()) is True + + registry.release_active_slot(key) + + assert registry.active_jobs(REPO) == {} + assert registry.active_job_refs(REPO) + assert registry.has_active_variant(REPO, VARIANT) is True + with ( + patch("hub.utils.download_registry.get_models_registry", lambda: registry), + patch( + "core.inference.llama_cpp.cached_gguf_for_load", + side_effect = AssertionError("same-variant jobs must block before cache reuse"), + ), + ): + assert _hub_download_blocks_gguf_load(REPO, VARIANT) is True + + registry.set_job(key, "complete") + assert registry.has_active_variant(REPO, VARIANT) is False + + def test_other_variant_job_still_allows_complete_cached_load(self): + from core.inference.llama_cpp import _hub_download_blocks_gguf_load + from hub.utils.download_registry import DownloadRegistry, TRANSPORT_HTTP + + registry = DownloadRegistry() + registry.claim( + f"{REPO}::Q8_0", + TRANSPORT_HTTP, + repo_type = "model", + repo_id = REPO, + variant = "Q8_0", + ) + with ( + patch("hub.utils.download_registry.get_models_registry", lambda: registry), + patch( + "core.inference.llama_cpp.cached_gguf_for_load", + return_value = "/cached/model.gguf", + ) as cached_probe, + ): + assert _hub_download_blocks_gguf_load(REPO, VARIANT) is False + + cached_probe.assert_called_once_with( + REPO, + VARIANT, + require_mmproj = False, + verify_sizes = True, + hf_token = None, + ) + + def test_cancelled_request_keeps_marker_until_load_thread_finishes(self): + from core.inference.llama_cpp import _with_gguf_load_marker + + started = threading.Event() + release = threading.Event() + finished = threading.Event() + + class FakeBackend: + @_with_gguf_load_marker + def load_model(self, *, hf_repo): + started.set() + release.wait(timeout = 2) + finished.set() + return True + + async def scenario(): + with patch( + "core.inference.llama_cpp._hub_download_blocks_gguf_load", + return_value = False, + ): + task = asyncio.create_task( + asyncio.to_thread(FakeBackend().load_model, hf_repo = REPO) + ) + assert await asyncio.to_thread(started.wait, 1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert hf_gguf_load_in_flight(REPO) + + release.set() + assert await asyncio.to_thread(finished.wait, 1) + for _ in range(100): + if not hf_gguf_load_in_flight(REPO): + break + await asyncio.sleep(0.001) + assert not hf_gguf_load_in_flight(REPO) + + asyncio.run(scenario()) + + def test_load_marker_precedes_hub_guard_and_unload(self): + source = (Path(__file__).resolve().parent.parent / "routes" / "inference.py").read_text() + gguf_branch = source[source.index("if config.is_gguf:") :] + + assert ( + gguf_branch.index("enter_context(gguf_load_in_flight") + < gguf_branch.index("if request.llama_extra_args is None") + < gguf_branch.index("_hub_download_blocks_gguf_load") + < gguf_branch.index("unsloth_backend.unload_model") + ) + llama_source = ( + Path(__file__).resolve().parent.parent / "core" / "inference" / "llama_cpp.py" + ).read_text() + assert "@_with_gguf_load_marker\n def load_model(" in llama_source diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index c3161c5714..bd2c008589 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -1061,6 +1061,80 @@ def test_same_turn_duplicate_web_search_is_internal_noop(monkeypatch): ] +def test_same_turn_duplicate_does_not_drop_later_parallel_call(monkeypatch): + # One batch: search(a), search(a) [duplicate], search(b). The duplicate is an + # internal no-op, but the distinct search(b) after it must still run, and the + # no-op nudge must land after the tool results rather than splitting them. + batch = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_a1", + "type": "function", + "function": {"name": "web_search", "arguments": json.dumps({"query": "a"})}, + }, + { + "index": 1, + "id": "call_a2", + "type": "function", + "function": {"name": "web_search", "arguments": json.dumps({"query": "a"})}, + }, + { + "index": 2, + "id": "call_b", + "type": "function", + "function": {"name": "web_search", "arguments": json.dumps({"query": "b"})}, + }, + ] + } + ), + _done(), + ] + final_stream = [_sse({"content": "Final answer."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [batch, final_stream], payloads) + + calls: list[dict] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append(arguments) + return "search-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 = 3, + ) + ) + + # Both distinct calls ran; the duplicate did not (old `break` dropped search(b)). + assert calls == [{"query": "a"}, {"query": "b"}] + assert [e.get("tool_call_id") for e in events if e.get("type") == "tool_end"] == [ + "call_a1", + "call_b", + ] + + # The next generation's conversation must be well-formed: the assistant lists + # only the executed calls (no orphan for the duplicate), the two tool results + # follow contiguously, and the no-op nudge lands after them, never between. + conv = payloads[1]["messages"] + asst = next(m for m in conv if m["role"] == "assistant" and m.get("tool_calls")) + assert [tc.get("id") for tc in asst["tool_calls"]] == ["call_a1", "call_b"] + after = conv[conv.index(asst) + 1 :] + assert [m["role"] for m in after[:2]] == ["tool", "tool"] + assert [m.get("tool_call_id") for m in after[:2]] == ["call_a1", "call_b"] + assert after[2]["role"] == "user" # deferred duplicate nudge, after the results + assert after[2]["content"].startswith( + "One earlier request to call tool 'web_search' in this batch was not executed" + ) + assert "previous tool request" not in after[2]["content"].lower() + + def test_same_turn_repeated_render_html_does_not_emit_second_provisional_start(monkeypatch): same_turn_render_calls = [ _sse( diff --git a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py index 3c21f41701..d0213f6079 100644 --- a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py +++ b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py @@ -373,6 +373,7 @@ def test_kill_orphaned_servers_returns_count(): with ( patch.dict(sys.modules, {"psutil": fake_psutil}), patch.dict(os.environ, {"LLAMA_SERVER_PATH": fake_path}), + patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)), ): n = LlamaCppBackend._kill_orphaned_servers() assert n == 1, "only the Studio-owned orphan should be counted" @@ -384,11 +385,53 @@ def test_kill_orphaned_servers_returns_count(): with ( patch.dict(sys.modules, {"psutil": fake_psutil}), patch.dict(os.environ, {"LLAMA_SERVER_PATH": fake_path}), + patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)), ): assert LlamaCppBackend._kill_orphaned_servers() == 0 assert killed == [] +def test_kill_orphaned_servers_spares_live_parent(): + """A Studio-owned llama-server whose parent is still running is not an + orphan (a live Studio or the user's shell owns it) and must never be + killed; only the true orphan (parent gone) is reaped.""" + import os + + mypid = os.getpid() + fake_path = "/tmp/unsloth-test-llama/llama-server" + killed: list[int] = [] + + class _FakeProc: + def __init__(self, pid, name, exe): + self.info = {"pid": pid, "name": name, "exe": exe} + + def kill(self): + killed.append(self.info["pid"]) + + live_parent = _FakeProc(mypid + 1, "llama-server", fake_path) + true_orphan = _FakeProc(mypid + 2, "llama-server", fake_path) + + fake_psutil = _types.ModuleType("psutil") + fake_psutil.NoSuchProcess = type("NoSuchProcess", (Exception,), {}) + fake_psutil.AccessDenied = type("AccessDenied", (Exception,), {}) + fake_psutil.ZombieProcess = type("ZombieProcess", (Exception,), {}) + fake_psutil.process_iter = lambda attrs = None: [live_parent, true_orphan] + + with ( + patch.dict(sys.modules, {"psutil": fake_psutil}), + patch.dict(os.environ, {"LLAMA_SERVER_PATH": fake_path}), + patch.object(LlamaCppBackend, "_reap_recorded_pid", staticmethod(lambda: 0)), + patch.object( + LlamaCppBackend, + "_pid_parent_is_alive", + staticmethod(lambda pid: pid == mypid + 1), + ), + ): + n = LlamaCppBackend._kill_orphaned_servers() + assert n == 1, "only the true orphan should be reaped" + assert killed == [mypid + 2], "the live-parent server must be spared" + + def test_startup_reaper_arms_settle_timestamp(): """__init__ arms ``_last_kill_monotonic`` when the startup reaper kills an orphan (so the first load_model waits for VRAM to settle), and leaves the diff --git a/studio/backend/tests/test_mcp_server.py b/studio/backend/tests/test_mcp_server.py new file mode 100644 index 0000000000..71792605ae --- /dev/null +++ b/studio/backend/tests/test_mcp_server.py @@ -0,0 +1,290 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import asyncio +import sys +import types + +import pytest + +from mcp_server import BearerTokenMiddleware, _clamp, _dump, create_studio_mcp + + +def _get_tool(name): + tools = asyncio.run(create_studio_mcp().list_tools()) + return {tool.name: tool for tool in tools}[name] + + +def test_studio_mcp_registers_control_plane_tools(): + tools = asyncio.run(create_studio_mcp().list_tools()) + + assert {tool.name for tool in tools} == { + "studio_status", + "list_local_models", + "get_training_status", + "start_training", + "stop_training", + "list_training_runs", + "validate_recipe", + "get_recipe_job_status", + "get_recipe_job_dataset", + "load_checkpoint", + "export_gguf", + } + + +def test_dump_serializes_pydantic_values(): + class Response: + def model_dump(self, *, mode): + assert mode == "json" + return {"ok": True} + + assert _dump(Response()) == {"ok": True} + assert _dump({"already": "json"}) == {"already": "json"} + + +def test_bearer_token_middleware_rejects_wrong_token(): + events = [] + + async def app(scope, receive, send): + events.append("app") + + async def send(message): + events.append(message) + + middleware = BearerTokenMiddleware(app, "secret") + asyncio.run( + middleware( + {"type": "http", "headers": [(b"authorization", b"Bearer wrong")]}, + None, + send, + ) + ) + + assert events[0]["status"] == 401 + assert "app" not in events + + +def test_bearer_token_middleware_closes_unauthorized_websocket(): + events = [] + + async def app(scope, receive, send): + events.append("app") + + async def send(message): + events.append(message) + + middleware = BearerTokenMiddleware(app, "secret") + asyncio.run( + middleware( + {"type": "websocket", "headers": []}, + None, + send, + ) + ) + + assert events == [{"type": "websocket.close", "code": 4401}] + + +def test_bearer_token_middleware_rejects_non_ascii_authorization(): + # A non-ASCII bearer value must produce a clean 401, not a 500. Comparing on + # bytes avoids the str hmac.compare_digest TypeError on non-ASCII input. + events = [] + + async def app(scope, receive, send): + events.append("app") + + async def send(message): + events.append(message) + + middleware = BearerTokenMiddleware(app, "secret") + asyncio.run( + middleware( + {"type": "http", "headers": [(b"authorization", b"Bearer \xff\xff")]}, + None, + send, + ) + ) + + assert events[0]["status"] == 401 + assert "app" not in events + + +def test_bearer_token_middleware_accepts_correct_token(): + events = [] + + async def app(scope, receive, send): + events.append("app") + + async def send(message): + events.append(message) + + middleware = BearerTokenMiddleware(app, "secret") + asyncio.run( + middleware( + {"type": "http", "headers": [(b"authorization", b"Bearer secret")]}, + None, + send, + ) + ) + + assert events == ["app"] + + +def test_bearer_token_middleware_requires_non_empty_token(): + async def app(scope, receive, send): + pass + + for bad in ("", " "): + with pytest.raises(ValueError): + BearerTokenMiddleware(app, bad) + + +def test_bearer_token_middleware_rejects_non_ascii_token(): + async def app(scope, receive, send): + pass + + # non-ASCII tokens cannot be transmitted in an HTTP header by a standard + # client, so they are rejected at construction instead of locking out. + for bad in ("töken", "\U0001f600"): + with pytest.raises(ValueError): + BearerTokenMiddleware(app, bad) + + +def test_bearer_token_middleware_passes_through_non_http_scopes(): + events = [] + + async def app(scope, receive, send): + events.append("app") + + async def send(message): + events.append(message) + + middleware = BearerTokenMiddleware(app, "secret") + asyncio.run(middleware({"type": "lifespan"}, None, send)) + + assert events == ["app"] + + +def test_clamp_restricts_to_inclusive_bounds(): + assert _clamp(5, 1, 200) == 5 + assert _clamp(-10, 1, 200) == 1 + assert _clamp(10_000, 1, 200) == 200 + assert _clamp(0, 1, 500) == 1 + assert _clamp(1_000, 1, 500) == 500 + + +def test_export_and_checkpoint_tools_expose_forwarded_fields(): + export_props = set(_get_tool("export_gguf").parameters["properties"]) + assert {"hf_token", "imatrix", "imatrix_path"} <= export_props + + checkpoint_props = set(_get_tool("load_checkpoint").parameters["properties"]) + assert {"hf_token", "approved_remote_code_fingerprint"} <= checkpoint_props + + +def _stub_module(monkeypatch, name, **attrs): + module = types.ModuleType(name) + for key, value in attrs.items(): + setattr(module, key, value) + if "." in name: + module.__path__ = [] # mark package-like so submodule imports resolve + monkeypatch.setitem(sys.modules, name, module) + return module + + +def test_export_gguf_forwards_hf_token_and_imatrix(monkeypatch): + captured = {} + + class FakeExportGGUFRequest: + def __init__(self, **kwargs): + captured.update(kwargs) + + async def fake_export(request, current_subject): + return {"current_subject": current_subject} + + _stub_module(monkeypatch, "models", ExportGGUFRequest = FakeExportGGUFRequest) + _stub_module(monkeypatch, "routes") + _stub_module(monkeypatch, "routes.export", export_gguf = fake_export) + + tool = _get_tool("export_gguf") + result = asyncio.run( + tool.fn( + save_directory = "/tmp/out", + quantization_method = ["Q4_K_M", "Q8_0"], + push_to_hub = True, + repo_id = "me/model", + hf_token = "hf_secret", + imatrix = True, + imatrix_path = "/tmp/imatrix.dat", + ) + ) + + assert captured["hf_token"] == "hf_secret" + assert captured["imatrix"] is True + assert captured["imatrix_path"] == "/tmp/imatrix.dat" + assert captured["quantization_method"] == ["Q4_K_M", "Q8_0"] + assert result["current_subject"] == "mcp" + + +def test_load_checkpoint_forwards_token_and_fingerprint(monkeypatch): + captured = {} + + class FakeLoadCheckpointRequest: + def __init__(self, **kwargs): + captured.update(kwargs) + + async def fake_load(request, current_subject): + return {"current_subject": current_subject} + + _stub_module(monkeypatch, "models", LoadCheckpointRequest = FakeLoadCheckpointRequest) + _stub_module(monkeypatch, "routes") + _stub_module(monkeypatch, "routes.export", load_checkpoint = fake_load) + + tool = _get_tool("load_checkpoint") + asyncio.run( + tool.fn( + checkpoint_path = "/tmp/ckpt", + approved_remote_code_fingerprint = "sha256:abc", + hf_token = "hf_secret", + ) + ) + + assert captured["hf_token"] == "hf_secret" + assert captured["approved_remote_code_fingerprint"] == "sha256:abc" + + +def test_list_training_runs_clamps_pagination(monkeypatch): + captured = {} + + async def fake_list_runs(limit, offset, current_subject): + captured["limit"] = limit + captured["offset"] = offset + return {"ok": True} + + _stub_module(monkeypatch, "routes") + _stub_module(monkeypatch, "routes.training_history", list_training_runs = fake_list_runs) + + tool = _get_tool("list_training_runs") + asyncio.run(tool.fn(limit = 10_000, offset = -5)) + + assert captured["limit"] == 200 + assert captured["offset"] == 0 + + +def test_get_recipe_job_dataset_clamps_pagination(monkeypatch): + captured = {} + + def fake_job_dataset(job_id, limit, offset): + captured["limit"] = limit + captured["offset"] = offset + return {"ok": True} + + _stub_module(monkeypatch, "routes") + _stub_module(monkeypatch, "routes.data_recipe") + _stub_module(monkeypatch, "routes.data_recipe.jobs", job_dataset = fake_job_dataset) + + tool = _get_tool("get_recipe_job_dataset") # this tool is synchronous + tool.fn(job_id = "job-1", limit = -1, offset = -9) + + assert captured["limit"] == 1 + assert captured["offset"] == 0 diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index 29fbb45158..fa50cd84d6 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -523,3 +523,169 @@ def test_mlx_generate_text_forwards_kwargs_into_template_helper(monkeypatch): assert render["kwargs"]["enable_thinking"] is True assert render["kwargs"]["reasoning_effort"] == "medium" assert render["kwargs"]["preserve_thinking"] is True + + +def test_mlx_text_normalizes_native_reasoning_and_close_releases_lock(monkeypatch): + _install_fake_mlx(monkeypatch) + from core.inference.mlx_inference import MLXInferenceBackend + + monkeypatch.setattr( + "core.inference.chat_template_helpers.apply_chat_template_for_generation", + lambda *_args, **_kwargs: "prompt", + raising = True, + ) + monkeypatch.setattr( + "core.inference.chat_template_helpers.render_with_native_template_fallback", + lambda formatted_prompt, **_kwargs: SimpleNamespace( + prompt = formatted_prompt, + reasoning_channel_markers = ("<|channel>thought\n", ""), + ), + raising = True, + ) + + mlx_lm_pkg = types.ModuleType("mlx_lm") + mlx_lm_sample = types.ModuleType("mlx_lm.sample_utils") + mlx_lm_sample.make_sampler = lambda **_kw: object() + mlx_lm_sample.make_logits_processors = lambda **_kw: None + + class _Resp: + def __init__(self, text, tok): + self.text = text + self.token = tok + + def _stream_generate(_model, _tokenizer, **_kw): + yield _Resp("<|channel>thought\n", 10) + yield _Resp("r", 11) + yield _Resp("", 12) + yield _Resp("a", 13) + + mlx_lm_pkg.stream_generate = _stream_generate + monkeypatch.setitem(sys.modules, "mlx_lm", mlx_lm_pkg) + monkeypatch.setitem(sys.modules, "mlx_lm.sample_utils", mlx_lm_sample) + + backend = MLXInferenceBackend() + backend._model = object() + backend._tokenizer = SimpleNamespace(all_special_tokens = []) + backend._is_vlm = False + + assert list( + backend.generate_chat_response( + messages = [{"role": "user", "content": "ping"}], + max_new_tokens = 4, + ) + ) == ["", "r", "r", "ra"] + + gen = backend.generate_chat_response( + messages = [{"role": "user", "content": "ping"}], + max_new_tokens = 4, + ) + assert next(gen) == "" + assert backend._generation_lock.locked() + gen.close() + assert not backend._generation_lock.locked() + + +def test_mlx_text_native_metadata_preserves_prefilled_think_snapshots(monkeypatch): + _install_fake_mlx(monkeypatch) + from core.inference.mlx_inference import MLXInferenceBackend + + monkeypatch.setattr( + "core.inference.chat_template_helpers.apply_chat_template_for_generation", + lambda *_args, **_kwargs: "prompt\n", + raising = True, + ) + monkeypatch.setattr( + "core.inference.chat_template_helpers.render_with_native_template_fallback", + lambda formatted_prompt, **_kwargs: SimpleNamespace( + prompt = formatted_prompt, + reasoning_channel_markers = ("<|channel>thought", ""), + ), + raising = True, + ) + + mlx_lm_pkg = types.ModuleType("mlx_lm") + mlx_lm_sample = types.ModuleType("mlx_lm.sample_utils") + mlx_lm_sample.make_sampler = lambda **_kw: object() + mlx_lm_sample.make_logits_processors = lambda **_kw: None + + class _Resp: + def __init__(self, text, tok): + self.text = text + self.token = tok + + def _stream_generate(_model, _tokenizer, **_kw): + yield _Resp("reason", 10) + yield _Resp("", 11) + yield _Resp("answer", 12) + + mlx_lm_pkg.stream_generate = _stream_generate + monkeypatch.setitem(sys.modules, "mlx_lm", mlx_lm_pkg) + monkeypatch.setitem(sys.modules, "mlx_lm.sample_utils", mlx_lm_sample) + + backend = MLXInferenceBackend() + backend._model = object() + backend._tokenizer = SimpleNamespace(all_special_tokens = []) + backend._is_vlm = False + + snapshots = list( + backend.generate_chat_response( + messages = [{"role": "user", "content": "ping"}], + max_new_tokens = 3, + ) + ) + assert snapshots == [ + "\n", + "\nreason", + "\nreason", + "\nreasonanswer", + ] + assert all(current.startswith(previous) for previous, current in zip(snapshots, snapshots[1:])) + + +def test_mlx_vlm_normalizes_native_reasoning_channels(monkeypatch): + _install_fake_mlx(monkeypatch) + from core.inference.mlx_inference import MLXInferenceBackend + + monkeypatch.setattr( + "core.inference.chat_template_helpers.apply_chat_template_for_generation", + lambda *_args, **_kwargs: "prompt", + raising = True, + ) + + mlx_vlm_pkg = types.ModuleType("mlx_vlm") + + class _Resp: + def __init__(self, text, tok): + self.text = text + self.token = tok + + def _stream_generate(_model, _processor, _prompt, _images, **_kw): + yield _Resp("<|channel>thought\n", 10) + yield _Resp("vision", 11) + yield _Resp("", 12) + yield _Resp(" answer", 13) + + mlx_vlm_pkg.stream_generate = _stream_generate + monkeypatch.setitem(sys.modules, "mlx_vlm", mlx_vlm_pkg) + + backend = MLXInferenceBackend() + backend._model = SimpleNamespace(config = SimpleNamespace()) + backend._processor = SimpleNamespace( + chat_template = "<|channel>thought\n...", + all_special_tokens = [], + apply_chat_template = lambda *_args, **_kwargs: "prompt", + ) + backend._is_vlm = True + + assert list( + backend.generate_chat_response( + messages = [{"role": "user", "content": "describe"}], + image = object(), + max_new_tokens = 4, + ) + ) == [ + "", + "vision", + "vision", + "vision answer", + ] diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index 86e528ae67..02230632b6 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -328,6 +328,7 @@ def test_download_mtp_prefers_root_over_new_scheme_copies(monkeypatch): pick, label, cancel_event = None, + near_path = None, ): captured["pick"] = pick return None @@ -428,6 +429,32 @@ def test_download_mtp_reuse_follows_snapshot_order_offline(tmp_path, monkeypatch assert got is not None and Path(got).parent.parent.name == "newest" +def test_download_mtp_prefers_main_snapshot_offline(tmp_path, monkeypatch): + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + snapshots = tmp_path / "models--unsloth--gemma" / "snapshots" + old = snapshots / "old" + new = snapshots / "new" + old.mkdir(parents = True) + new.mkdir(parents = True) + main = old / "gemma-UD-Q4_K_XL.gguf" + old_drafter = old / "mtp-gemma.gguf" + new_drafter = new / "mtp-gemma.gguf" + main.write_bytes(b"main") + old_drafter.write_bytes(b"old") + new_drafter.write_bytes(b"new") + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda _repo: [new, old]) + + got = LlamaCppBackend()._download_mtp( + hf_repo = "unsloth/gemma-GGUF", + near_path = str(main), + ) + + assert got == str(old_drafter) + + 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). @@ -447,6 +474,7 @@ def test_download_mtp_online_skips_cache_reuse(tmp_path, monkeypatch): pick, label, cancel_event = None, + near_path = None, ): reached["hit"] = True return None diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index e24e2ca451..0b9a1e704f 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -244,9 +244,7 @@ class TestGgufVariantFileResolution: 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. + # Keep coverage for offline reuse; online reuse is tested separately. monkeypatch.setenv("HF_HUB_OFFLINE", "1") backend = LlamaCppBackend() repo = "unsloth/vision-GGUF" @@ -292,8 +290,7 @@ class TestGgufVariantFileResolution: 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. + # Keep coverage for case-insensitive offline cache lookup. monkeypatch.setenv("HF_HUB_OFFLINE", "1") backend = LlamaCppBackend() canonical_repo = "unsloth/gemma-4-E2B-it-GGUF" @@ -348,45 +345,26 @@ class TestGgufVariantFileResolution: 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. + def test_download_online_reuses_complete_cached_snapshot(self, monkeypatch, hf_cache): + # Loads reuse complete cached models across repo revisions. 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] = [] + snap = _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 fake_download( - repo_id, - filename, - token = None, - **kwargs, - ): - downloaded.append(filename) - return f"/fresh/{filename}" + def fail_download(*_args, **_kwargs): + raise AssertionError("must 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", fake_download), + 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 downloaded == ["model-UD-Q4_K_XL.gguf"] - assert out == "/fresh/model-UD-Q4_K_XL.gguf" + assert out == str(snap / "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 diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 5300a48557..8725b28ac8 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -1054,7 +1054,7 @@ class TestChatCompletionRequestToolFields: monkeypatch.setattr( inference_route, "_detect_safetensors_features", - lambda backend, chat_template: {"supports_tools": True}, + lambda backend, chat_template, tools = None: {"supports_tools": True}, ) monitor = ApiMonitor(max_entries = 3) monkeypatch.setattr(inference_route, "api_monitor", monitor) @@ -6620,6 +6620,29 @@ class TestApiMonitorAudioInput: assert entry["reply"] == "hello world" assert monitor.active_count() == 0 + def failing_chunks(): + yield "partial" + raise RuntimeError("generation failed") + + self._patch_audio_backend(monkeypatch, failing_chunks()) + error_monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", error_monitor) + error_response = await openai_chat_completions( + payload, + request = request, + current_subject = "test", + ) + error_chunks = [ + chunk.decode() if isinstance(chunk, bytes) else chunk + async for chunk in error_response.body_iterator + ] + + assert '"type": "server_error"' in error_chunks[-1] + assert error_chunks[-1].endswith("data: [DONE]\n\n") + [error_entry] = error_monitor.snapshot() + assert error_entry["status"] == "error" + assert error_monitor.active_count() == 0 + asyncio.run(_run()) def test_non_gguf_tts_auto_route_records_monitor(self, monkeypatch): diff --git a/studio/backend/tests/test_rag_parsing.py b/studio/backend/tests/test_rag_parsing.py index 14ab0efe2e..3e259f6bd0 100644 --- a/studio/backend/tests/test_rag_parsing.py +++ b/studio/backend/tests/test_rag_parsing.py @@ -54,6 +54,56 @@ def test_pdf_markdown_off_uses_plain_text(tmp_path, monkeypatch): assert "#" not in text and "|" not in text # plain text path emits no Markdown markup +def test_pdf_bytes_use_same_extraction_path(tmp_path, monkeypatch): + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", False) + pdf = tmp_path / "table.pdf" + _table_pdf(pdf) + from_file = parsers.parse(str(pdf)) + from_bytes, total_pages = parsers.parse_pdf_bytes(pdf.read_bytes()) + assert [page.text for page in from_bytes] == [page.text for page in from_file] + assert total_pages == len(from_file) + + +def test_pdf_bytes_limit_pages_before_extraction(monkeypatch): + import pymupdf + + from core.rag import config, parsers + + monkeypatch.setattr(config, "PDF_MARKDOWN", False) + doc = pymupdf.open() + for marker in ("page one", "page two", "page three"): + page = doc.new_page() + page.insert_text((40, 40), marker) + data = doc.tobytes() + doc.close() + + pages, total_pages = parsers.parse_pdf_bytes(data, max_pages = 2) + assert len(pages) == 2 + assert "page two" in pages[-1].text + assert total_pages == 3 # full count, not the 2 extracted + + +def test_pdf_markdown_receives_page_limit(monkeypatch): + from core.rag import parsers + + captured = {} + + class _FakePymupdf4llm: + @staticmethod + def to_markdown(doc, **kwargs): + captured.update(kwargs) + return [{"text": "page"} for _ in kwargs["pages"]] + + class _Doc: + page_count = 100 + + monkeypatch.setitem(__import__("sys").modules, "pymupdf4llm", _FakePymupdf4llm) + assert parsers._pdf_markdown(_Doc(), range(2)) == ["page", "page"] + assert captured == {"page_chunks": True, "show_progress": False, "pages": [0, 1]} + + def test_pdf_markdown_passes_only_supported_legacy_kwargs(monkeypatch): # The pinned PyMuPDF4LLM legacy path ignores unknown kwargs; do not pass the # newer layout-only OCR knobs or Markdown extraction silently loses policy control. diff --git a/studio/backend/tests/test_safetensors_capability_advertise.py b/studio/backend/tests/test_safetensors_capability_advertise.py index 0ed670ac01..bd3d8d16b9 100644 --- a/studio/backend/tests/test_safetensors_capability_advertise.py +++ b/studio/backend/tests/test_safetensors_capability_advertise.py @@ -417,6 +417,59 @@ def test_detect_safetensors_features_gemma_native_tool_call_keeps_tools_on(): assert flags["supports_tools"] is True +def test_detect_safetensors_features_gemma_native_reasoning_is_parseable_not_prefilled(): + """Native Gemma channels are normalized to , then split by the route.""" + from routes.inference import _detect_safetensors_features, _sf_reasoning_prefill_mode + + tpl_with_gemma_native = "{% if add_generation_prompt %}<|channel>thought\n{% endif %}" + backend = SimpleNamespace( + active_model_name = "unsloth/gemma-4-E2B-it", + models = { + "unsloth/gemma-4-E2B-it": { + "native_chat_template": tpl_with_gemma_native, + "chat_template_info": {"template": "override has no native markers"}, + } + }, + ) + flags = _detect_safetensors_features(backend, "override has no native markers") + missing_arg_flags = _detect_safetensors_features(backend, None) + + assert flags["supports_reasoning"] is True + assert flags["reasoning_always_on"] is True + assert missing_arg_flags["supports_reasoning"] is True + assert _sf_reasoning_prefill_mode(flags, None, tpl_with_gemma_native) is False + + +def test_detect_safetensors_features_selects_native_reasoning_from_tool_template(): + """Request tools select a marker-bearing named template without affecting default chat.""" + from routes.inference import _detect_safetensors_features + + named_template = { + "default": "plain default template", + "tool_use": "{% if tools %}<|channel>thought\n{% endif %}", + } + backend = SimpleNamespace( + active_model_name = "custom/named-native-reasoning", + models = { + "custom/named-native-reasoning": { + "native_chat_template": named_template, + "chat_template_info": {"template": "{% if tools %}{% endif %}"}, + } + }, + ) + + default_flags = _detect_safetensors_features(backend, "plain override") + tool_flags = _detect_safetensors_features( + backend, + "plain override", + tools = [{"type": "function"}], + ) + + assert default_flags["supports_reasoning"] is False + assert tool_flags["supports_reasoning"] is True + assert tool_flags["reasoning_always_on"] is True + + # Qwen3.5 family pin: the live GGUF + safetensors templates both wrap tool # calls as ``\n...``. Faithful slice so the # classifier never silently regresses for this family. diff --git a/studio/backend/tests/test_safetensors_reasoning_stream.py b/studio/backend/tests/test_safetensors_reasoning_stream.py index 4e708139b7..af5a05d266 100644 --- a/studio/backend/tests/test_safetensors_reasoning_stream.py +++ b/studio/backend/tests/test_safetensors_reasoning_stream.py @@ -215,3 +215,135 @@ def test_s6_reasoning_effort_none_disables_prefill_for_enable_thinking_effort(): swallowed = _replay_sf_reasoning_stream(events, prefilled = True) assert swallowed["visible"] == "" assert swallowed["reasoning"] == "The capital of France is Paris." + + +def test_native_reasoning_streamer_selected_and_errors_raise(): + import threading + import pytest + + torch = pytest.importorskip("torch") + inf = pytest.importorskip("core.inference.inference") + + class Batch(dict): + def to(self, _device): + return self + + class Tok: + chat_template = "<|channel>thought\n..." + all_special_tokens = [] + eos_token_id = 1 + pad_token_id = None + pieces = {10: "<|channel>thought\n", 11: "r", 12: "", 13: "a"} + + def __call__(self, *_args, **_kwargs): + return Batch({"input_ids": torch.zeros((1, 1), dtype = torch.long)}) + + def decode(self, ids, **_kwargs): + return "".join(self.pieces.get(int(token_id), "") for token_id in ids) + + class Model: + device = "cpu" + generation_config = type("Cfg", (), {"eos_token_id": 1})() + config = generation_config + + def __init__(self, fail = False): + self.fail = fail + self.kwargs = None + + def generate(self, **kwargs): + self.kwargs = kwargs + streamer = kwargs["streamer"] + streamer.put(torch.zeros((1, 1), dtype = torch.long)) + for token_id in [10, 11, 12, 13]: + streamer.put(torch.tensor([token_id])) + if self.fail: + raise RuntimeError("boom") + + backend = inf.InferenceBackend.__new__(inf.InferenceBackend) + backend.active_model_name = "gemma-test" + backend._generation_lock = threading.Lock() + backend.models = {"gemma-test": {"model": Model(), "tokenizer": Tok()}} + + assert list(backend.generate_stream("prompt", max_new_tokens = 4))[-1] == "ra" + + backend.models["gemma-test"]["model"] = Model(fail = True) + + with pytest.raises(inf._GenerationThreadError, match = "boom"): + list(backend.generate_stream("prompt", max_new_tokens = 4)) + + +def test_text_only_vlm_fallback_resolves_native_markers_off(): + import threading + import pytest + + torch = pytest.importorskip("torch") + inf = pytest.importorskip("core.inference.inference") + + class Batch(dict): + def to(self, _device): + return self + + class Tokenizer: + all_special_tokens = [] + eos_token_id = 1 + pad_token_id = None + + def __call__(self, *_args, **_kwargs): + return Batch({"input_ids": torch.zeros((1, 1), dtype = torch.long)}) + + class Processor: + chat_template = "<|channel>thought\n..." + tokenizer = Tokenizer() + + class Model: + device = "cpu" + generation_config = type("Cfg", (), {"eos_token_id": 1})() + config = generation_config + + def generate(self, **_kwargs): + return None + + class EmptyStreamer: + def __next__(self): + raise StopIteration + + def end(self): + return None + + captured = {} + backend = inf.InferenceBackend.__new__(inf.InferenceBackend) + backend.active_model_name = "vision-test" + backend._generation_lock = threading.Lock() + backend.models = { + "vision-test": { + "model": Model(), + "processor": Processor(), + "tokenizer": Processor(), + } + } + backend.format_chat_prompt = lambda *_args, **_kwargs: "manual text-only prompt" + + def make_streamer(*_args, **kwargs): + captured.update(kwargs) + return EmptyStreamer() + + backend._make_text_streamer = make_streamer + + assert ( + list( + backend._generate_vision_response( + messages = [{"role": "user", "content": "hello"}], + system_prompt = "", + image = None, + temperature = 0.7, + top_p = 0.9, + top_k = 40, + min_p = 0.0, + max_new_tokens = 1, + repetition_penalty = 1.0, + ) + ) + == [] + ) + assert captured["reasoning_channel_markers"] is None + assert captured["reasoning_channel_markers_resolved"] is True diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index eae1a75161..915f82ac8e 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -2843,6 +2843,61 @@ class TestLoopBehaviour: ] assert len(duplicate_nudges) == 1 + def test_same_turn_duplicate_does_not_drop_later_parallel_call(self): + # Turn 1 runs search(x). Turn 2's batch is [search(x) duplicate, python]: + # the duplicate is a no-op, but python after it must still run, and the + # no-op nudge must land after python's result rather than splitting it. + captured_messages: list[list[dict]] = [] + 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_messages.append([dict(m) for m in messages]) + chunks = next(turns) + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + exec_fn = FakeExecuteTool(["search-x", "py-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 = 4, + ) + ) + + # Turn-1 search and turn-2 python both ran; the turn-2 duplicate search did not. + assert exec_fn.calls == [ + ("web_search", {"query": "x"}), + ("python", {"code": "print(1)"}), + ] + + conv = captured_messages[-1] + turn2 = [m for m in conv if m.get("role") == "assistant" and m.get("tool_calls")][-1] + assert [tc["function"]["name"] for tc in turn2["tool_calls"]] == ["python"] + after = conv[conv.index(turn2) + 1 :] + assert after[0]["role"] == "tool" and after[0]["content"] == "py-result" + assert after[1]["role"] == "user" # deferred duplicate nudge, after the result + assert after[1]["content"].startswith( + "One earlier request to call tool 'web_search' in this batch was not executed" + ) + assert "previous tool request" not in after[1]["content"].lower() + def test_duplicate_tool_call_internal_noop_allows_distinct_followup_tool(self): captured_messages: list[list[dict]] = [] captured_tool_names: list[list[str]] = [] @@ -3150,6 +3205,196 @@ class TestLoopBehaviour: 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_reasoning_intent_does_not_reprompt_a_visible_answer(self): + generations = 0 + + def _gen(_messages, active_tools = None): + nonlocal generations + generations += 1 + yield ( + "Let me prepare the requested summary carefully." + "This is the final visible answer." + ) + + exec_fn = FakeExecuteTool([]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "summarize this"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + nudge_tool_calls = True, + ) + ) + + assert generations == 1 + assert exec_fn.calls == [] + contents = [e["text"] for e in events if e["type"] == "content"] + assert contents[-1].endswith("This is the final visible answer.") + + def test_prefilled_reasoning_intent_does_not_reprompt_a_visible_answer(self): + generations = 0 + + def _gen(_messages, active_tools = None): + nonlocal generations + generations += 1 + yield "Let me prepare the requested summary carefully.This is the final visible answer." + + exec_fn = FakeExecuteTool([]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "summarize this"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + nudge_tool_calls = True, + reasoning_prefilled = True, + ) + ) + + assert generations == 1 + assert exec_fn.calls == [] + contents = [e["text"] for e in events if e["type"] == "content"] + assert contents[-1].endswith("This is the final visible answer.") + + def test_prefilled_reasoning_with_reemitted_think_does_not_reprompt(self): + generations = 0 + + def _gen(_messages, active_tools = None): + nonlocal generations + generations += 1 + yield ( + "Let me prepare the requested summary carefully." + "more private planningThis is the final visible answer." + ) + + exec_fn = FakeExecuteTool([]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "summarize this"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + nudge_tool_calls = True, + reasoning_prefilled = True, + ) + ) + + assert generations == 1 + assert exec_fn.calls == [] + contents = [e["text"] for e in events if e["type"] == "content"] + assert contents[-1].endswith("This is the final visible answer.") + + def test_prefilled_reasoning_with_later_think_does_not_reprompt(self): + generations = 0 + + def _gen(_messages, active_tools = None): + nonlocal generations + generations += 1 + yield ( + "private prefilled planning" + "Let me prepare the requested summary carefully." + "This is the final visible answer." + ) + + exec_fn = FakeExecuteTool([]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "summarize this"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + nudge_tool_calls = True, + reasoning_prefilled = True, + ) + ) + + assert generations == 1 + assert exec_fn.calls == [] + contents = [e["text"] for e in events if e["type"] == "content"] + assert contents[-1].endswith("This is the final visible answer.") + + def test_reasoning_only_intent_still_reprompts_and_uses_a_tool(self): + loop, exec_fn = _make_loop( + turns = [ + ["Let me search for that."], + ['{"name":"web_search","arguments":{"query":"cats"}}'], + ["Here is the answer."], + ], + exec_results = ["result"], + nudge_tool_calls = True, + ) + + events = _collect_events(loop) + + assert exec_fn.calls == [("web_search", {"query": "cats"})] + contents = [e["text"] for e in events if e["type"] == "content"] + assert contents[-1] == "Here is the answer." + + def test_prefilled_no_close_reasoning_intent_still_reprompts(self): + loop, exec_fn = _make_loop( + turns = [ + ["I need more context.Let me search for that."], + ['{"name":"web_search","arguments":{"query":"cats"}}'], + ["Here is the answer."], + ], + exec_results = ["result"], + nudge_tool_calls = True, + reasoning_prefilled = True, + ) + + events = _collect_events(loop) + + assert exec_fn.calls == [("web_search", {"query": "cats"})] + contents = [e["text"] for e in events if e["type"] == "content"] + assert contents[-1] == "Here is the answer." + + def test_prefilled_reasoning_prefix_is_kept_for_reasoning_only_reprompt(self): + loop, exec_fn = _make_loop( + turns = [ + ["Let me search for that.checking details"], + ['{"name":"web_search","arguments":{"query":"cats"}}'], + ["Here is the answer."], + ], + exec_results = ["result"], + nudge_tool_calls = True, + reasoning_prefilled = True, + ) + + events = _collect_events(loop) + + assert exec_fn.calls == [("web_search", {"query": "cats"})] + contents = [e["text"] for e in events if e["type"] == "content"] + assert contents[-1] == "Here is the answer." + + def test_reprompt_history_uses_visible_intent_text(self): + captured: list[list[dict]] = [] + + def _gen(messages, active_tools = None): + captured.append([dict(message) for message in messages]) + if len(captured) == 1: + yield "private planning detailsLet me search for that." + elif len(captured) == 2: + yield '{"name":"web_search","arguments":{"query":"cats"}}' + else: + yield "Here is the answer." + + exec_fn = FakeExecuteTool(["result"]) + events = _collect_events( + run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "find cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + nudge_tool_calls = True, + ) + ) + + assert exec_fn.calls == [("web_search", {"query": "cats"})] + assert captured[1][1] == {"role": "assistant", "content": "Let me search for that."} + contents = [e["text"] for e in events if e["type"] == "content"] + assert contents[-1] == "Here is the answer." + def test_intent_signal_triggers_reprompt(self): # Turn 1: intent signal, no tool call. # Turn 2 (re-prompt): proper tool call -> executes. diff --git a/studio/backend/tests/test_sf_client_tools_passthrough.py b/studio/backend/tests/test_sf_client_tools_passthrough.py index 01905b712c..f91eec9817 100644 --- a/studio/backend/tests/test_sf_client_tools_passthrough.py +++ b/studio/backend/tests/test_sf_client_tools_passthrough.py @@ -177,6 +177,15 @@ def _sse_objects(chunks): # ── Non-streaming ───────────────────────────────────────────────── +def test_non_reasoning_backend_keeps_literal_think_tags(monkeypatch): + backend = _ScriptedBackend(_fixed("show example tags")) + response = _call(_request(stream = False), monkeypatch, backend, supports_tools = False) + + message = _json_body(response)["choices"][0]["message"] + assert message["content"] == "show example tags" + assert message["reasoning_content"] is None + + def test_xml_healed_to_tool_calls_non_streaming(monkeypatch): backend = _ScriptedBackend(_fixed(_CALL_XML)) payload = _request(tools = [LOOKUP_TOOL], stream = False) @@ -485,6 +494,52 @@ def test_streaming_no_tools_verbatim(monkeypatch): assert finishes == ["stop"] +def test_streaming_gen_stream_error_is_not_model_text(monkeypatch): + from core.inference.orchestrator import GenStreamError + + class _ErrorAfterPartial(_ScriptedBackend): + def __init__(self): + super().__init__(_fixed()) + + def generate_chat_response(self, **_kwargs): + yield "partial" + yield GenStreamError("Error: /tmp/secret traceback") + + backend = _ErrorAfterPartial() + payload = _request(stream = True) + response = _call(payload, monkeypatch, backend, supports_tools = False) + chunks = _collect_sse(response) + objs = _sse_objects(chunks) + + deltas = [o.get("choices", [{}])[0].get("delta", {}) for o in objs if o.get("choices")] + assert any("partial" in json.dumps(delta) for delta in deltas) + assert not any("/tmp/secret" in json.dumps(delta) for delta in deltas) + errors = [o["error"]["message"] for o in objs if "error" in o] + assert errors == ["An internal error occurred."] + assert any( + "data: [DONE]" in (chunk.decode() if isinstance(chunk, bytes) else chunk) + for chunk in chunks + ) + + +def test_server_tool_streaming_invalid_event_is_error(monkeypatch): + class _InvalidEventBackend(_ScriptedBackend): + def __init__(self): + super().__init__(_fixed()) + + def generate_chat_completion_with_tools(self, **_kwargs): + yield {"type": "content", "text": "partial"} + yield "not-an-event" + + backend = _InvalidEventBackend() + payload = _request(tools = [LOOKUP_TOOL], enable_tools = True, stream = True) + response = _call(payload, monkeypatch, backend) + objs = _sse_objects(_collect_sse(response)) + + errors = [o["error"]["message"] for o in objs if "error" in o] + assert errors == ["An internal error occurred."] + + 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)) diff --git a/studio/backend/tests/test_think_prefill_reemit.py b/studio/backend/tests/test_think_prefill_reemit.py index 300ff92776..346399c3b2 100644 --- a/studio/backend/tests/test_think_prefill_reemit.py +++ b/studio/backend/tests/test_think_prefill_reemit.py @@ -2,7 +2,7 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """ -Unit tests for detect_think_prefill. +Unit tests for local reasoning-stream helpers. Reasoning templates (Qwen3.6-style) end the generation prompt with an open ``\\n`` so the model starts reasoning immediately. skip_prompt @@ -16,7 +16,13 @@ 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 +from core.inference.chat_template_helpers import ( + ReasoningChannelNormalizer, + detect_reasoning_channel_markers, + detect_reasoning_channel_markers_from_model_info, + detect_think_prefill, + render_with_native_template_fallback, +) QWEN_PROMPT = "<|im_start|>user\nHi!<|im_end|>\n<|im_start|>assistant\n" @@ -87,3 +93,170 @@ def test_guard_emits_when_think_not_special(): 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" + + +def test_gemma_channel_detection_uses_active_template_not_token_metadata(): + class TemplateTokenizer: + chat_template = {"default": "...<|channel>thought\\n{{ eoc_token }}"} + + class NamedTemplateTokenizer: + chat_template = { + "default": "plain assistant template", + "tool_use": "...<|channel>thought\\n{{ eoc_token }}", + } + + class TokenMetadataOnly: + chat_template = None + soc_token = "<|channel>" + eoc_token = "" + + class NamedTemplateProcessor: + chat_template = { + "default": "plain processor default", + "tool_use": "<|channel>thought\nprocessor tool template", + } + tokenizer = TokenMetadataOnly() + + def apply_chat_template(self, *_args, **_kwargs): + raise NotImplementedError + + expected = ("<|channel>thought", "") + assert detect_reasoning_channel_markers(TemplateTokenizer()) == expected + assert detect_reasoning_channel_markers(NamedTemplateTokenizer()) is None + assert ( + detect_reasoning_channel_markers( + NamedTemplateTokenizer(), tools = [{"function": {"name": "web_search"}}] + ) + == expected + ) + assert detect_reasoning_channel_markers(NamedTemplateTokenizer(), tools = []) is None + assert ( + detect_reasoning_channel_markers( + NamedTemplateProcessor(), tools = [{"function": {"name": "web_search"}}] + ) + is None + ) + assert detect_reasoning_channel_markers(TokenMetadataOnly()) is None + + +def test_gemma_channel_detection_tries_no_argument_getter_fallback(): + class FallbackTokenizer: + chat_template = "plain fallback template" + + def get_chat_template(self, **kwargs): + if kwargs: + raise ValueError("tools are not supported") + return "...<|channel>thought\n" + + assert detect_reasoning_channel_markers( + FallbackTokenizer(), tools = [{"function": {"name": "web_search"}}] + ) == ("<|channel>thought", "") + + +def test_native_template_fallback_returns_selected_reasoning_metadata(): + from types import SimpleNamespace + + messages = [{"role": "user", "content": "hi"}] + tools = [{"type": "function", "function": {"name": "web_search"}}] + + def render(tokenizer, msgs, *, tools, **_kw): + body = "".join(message["content"] for message in msgs) + suffix = "|TOOLS" if tools else "" + return body + suffix if tokenizer.chat_template == "NATIVE <|channel>thought\n" else body + + result = render_with_native_template_fallback( + formatted_prompt = "hi", + tokenizer = SimpleNamespace(chat_template = "OVERRIDE"), + model_info = { + "native_chat_template": "NATIVE <|channel>thought\n", + "tokenizer": SimpleNamespace(chat_template = "OVERRIDE"), + }, + active_model_name = "gemma-test", + messages = messages, + tools = tools, + apply_fn = render, + return_metadata = True, + ) + + assert result.prompt == "hi|TOOLS" + assert result.reasoning_channel_markers == ("<|channel>thought", "") + + +def test_cached_native_template_metadata_recovers_reasoning_markers_without_tools(): + from types import SimpleNamespace + + model_info = {"chat_template_info": {"template": "native <|channel>thought\n"}} + + assert detect_reasoning_channel_markers_from_model_info( + SimpleNamespace(chat_template = "override has no native markers"), + model_info, + tools = None, + ) == ("<|channel>thought", "") + result = render_with_native_template_fallback( + formatted_prompt = "prompt from override", + tokenizer = SimpleNamespace(chat_template = "override has no native markers"), + model_info = model_info, + active_model_name = "gemma-test", + messages = [{"role": "user", "content": "hi"}], + tools = None, + return_metadata = True, + ) + assert result.prompt == "prompt from override" + assert result.reasoning_channel_markers == ("<|channel>thought", "") + + +def test_cached_native_markers_do_not_describe_live_tool_template(): + from types import SimpleNamespace + + tools = [{"type": "function", "function": {"name": "web_search"}}] + + class LiveTokenizer: + chat_template = "live tool template without native markers" + + def render(_tokenizer, _messages, *, tools, **_kwargs): + return "prompt with tools" if tools else "prompt without tools" + + result = render_with_native_template_fallback( + formatted_prompt = "prompt with tools", + tokenizer = LiveTokenizer(), + model_info = { + "chat_template_info": {"template": "native <|channel>thought\n"}, + "tokenizer": SimpleNamespace(), + }, + active_model_name = "gemma-test", + messages = [{"role": "user", "content": "hi"}], + tools = tools, + apply_fn = render, + return_metadata = True, + ) + + assert result.prompt == "prompt with tools" + assert result.reasoning_channel_markers is None + + +def test_gemma_channel_normalization_is_prefix_monotonic_and_preserves_tools(): + parser = ReasoningChannelNormalizer("<|channel>thought", "") + output = "" + snapshots = [] + for chunk in ( + "<|chan", + "nel>thought", + "\nReason", + "<|tool_call>web_search", + ): + delta = parser.feed(chunk) + if delta: + output += delta + snapshots.append(output) + + assert snapshots == [ + "", + "Reason", + "Reason<|tool_call>web_search", + ] + assert snapshots[1].startswith(snapshots[0]) + compact = ReasoningChannelNormalizer("<|channel>thought", "") + assert compact.feed("<|channel>thoughtanswer") + compact.finish() == ( + "answer" + ) diff --git a/studio/backend/tests/test_tool_loop_controller.py b/studio/backend/tests/test_tool_loop_controller.py index 0e8ae798af..496c30ac13 100644 --- a/studio/backend/tests/test_tool_loop_controller.py +++ b/studio/backend/tests/test_tool_loop_controller.py @@ -13,6 +13,7 @@ if _BACKEND_DIR not in sys.path: from core.inference.tool_loop_controller import ( ToolLoopController, + append_deferred_nudges, canonical_tool_call_key, coerce_tool_arguments, status_for_tool, @@ -21,6 +22,22 @@ from core.inference.tool_loop_controller import ( ) +def test_append_deferred_nudges_merges_deduped_into_one_message(): + conversation = [{"role": "assistant", "tool_calls": [1]}, {"role": "tool", "content": "r"}] + nudges = [ + {"role": "user", "content": "duplicate"}, + {"role": "user", "content": "duplicate"}, # dropped: same content + {"role": "user", "content": "disabled foo"}, + ] + append_deferred_nudges(conversation, nudges) + # One user message, after the results, with distinct contents joined. + assert conversation[2:] == [{"role": "user", "content": "duplicate\n\ndisabled foo"}] + # Empty is a no-op. + before = list(conversation) + append_deferred_nudges(conversation, []) + assert conversation == before + + def _tool(name: str) -> dict: return {"type": "function", "function": {"name": name}} @@ -111,6 +128,10 @@ def test_successful_duplicate_is_internal_noop_and_keeps_remaining_tools(): assert not duplicate.should_execute assert not duplicate.emit_visible_events duplicate_nudge = completion.model_message()["content"] + assert duplicate_nudge.startswith( + "One earlier request to call tool 'web_search' in this batch was not executed" + ) + assert "previous tool request" not in duplicate_nudge.lower() assert "already completed successfully" in duplicate_nudge assert "different enabled tool" in duplicate_nudge assert completion.model_message()["role"] == "user" @@ -165,7 +186,12 @@ def test_empty_enabled_tool_list_blocks_all_tool_calls(): assert decision.action == "disabled" assert not decision.emit_visible_events assert completion.model_message()["role"] == "user" - assert "not enabled" in completion.model_message()["content"] + disabled_nudge = completion.model_message()["content"] + assert disabled_nudge.startswith( + "One earlier request to call tool 'web_search' in this batch was not executed" + ) + assert "previous tool request" not in disabled_nudge.lower() + assert "not enabled" in disabled_nudge assert controller.force_final_answer assert controller.active_tools() == [] diff --git a/studio/backend/tests/test_web_fetch_binary_guard.py b/studio/backend/tests/test_web_fetch_binary_guard.py index 3041ed5c34..10db953913 100644 --- a/studio/backend/tests/test_web_fetch_binary_guard.py +++ b/studio/backend/tests/test_web_fetch_binary_guard.py @@ -59,6 +59,18 @@ def _fetch_with(monkeypatch, body: bytes, content_type: str | None) -> str: return tools._fetch_page_text("https://example.com/thing", timeout = 5) +def _pdf_bytes(*page_texts: str) -> bytes: + pymupdf = pytest.importorskip("pymupdf") + doc = pymupdf.open() + for text in page_texts: + page = doc.new_page() + if text: + page.insert_textbox(pymupdf.Rect(40, 40, 550, 750), text, fontsize = 11) + data = doc.tobytes() + doc.close() + return data + + @pytest.mark.parametrize( "content_type,expected", [ @@ -90,10 +102,119 @@ def test_is_text_candidate_content_type(content_type, expected): assert tools._is_text_candidate_content_type(content_type) is expected -def test_pdf_rejected_by_content_type(monkeypatch): - out = _fetch_with(monkeypatch, b"%PDF-1.7\n\xff\xd8\xff\x00\x89PNG" * 200, "application/pdf") - assert "�" not in out - assert "non-text content" in out and "application/pdf" in out +@pytest.mark.parametrize( + "content_type", + ["application/pdf", "application/octet-stream", "text/html", "text/plain", None], +) +def test_pdf_text_extracted(monkeypatch, content_type): + out = _fetch_with( + monkeypatch, + _pdf_bytes("First page marker", "Second page marker"), + content_type, + ) + assert "## Page 1\n\nFirst page marker" in out + assert "## Page 2" in out and "Second page marker" in out + assert "binary content" not in out and "non-text content" not in out + + +@pytest.mark.parametrize("content_type", ["application/pdf", "text/plain"]) +def test_malformed_pdf_returns_safe_placeholder(monkeypatch, content_type): + out = _fetch_with(monkeypatch, b"%PDF-1.7\nnot a complete PDF", content_type) + assert out == "(PDF content could not be read as text)" + + +def test_pdf_without_text_layer_reported(monkeypatch): + out = _fetch_with(monkeypatch, _pdf_bytes(""), "application/pdf") + assert out == "(PDF contains no extractable text)" + + +def test_encrypted_pdf_returns_safe_placeholder(monkeypatch): + pymupdf = pytest.importorskip("pymupdf") + doc = pymupdf.open() + doc.new_page().insert_text((40, 40), "private text") + data = doc.tobytes( + encryption = pymupdf.PDF_ENCRYPT_AES_256, + owner_pw = "owner", + user_pw = "secret", + ) + doc.close() + out = _fetch_with(monkeypatch, data, "application/pdf") + assert out == "(PDF content could not be read as text)" + + +def test_pdf_download_limit_enforced(monkeypatch): + monkeypatch.setattr(tools, "_MAX_PDF_FETCH_BYTES", 256) + out = _fetch_with(monkeypatch, _pdf_bytes("Readable but oversized"), "application/pdf") + assert out == "(PDF content exceeds the download limit; not readable as text)" + + +def test_mislabeled_pdf_is_read_past_text_download_cap(monkeypatch): + body = _pdf_bytes("Cross-reference data was fetched") + monkeypatch.setattr(tools, "_MAX_FETCH_BYTES", 128) + monkeypatch.setattr(tools, "_MAX_PDF_FETCH_BYTES", len(body) + 100) + out = _fetch_with(monkeypatch, body, "text/plain") + assert "Cross-reference data was fetched" in out + + +def test_pdf_extraction_caps_pages_and_intermediate_text(monkeypatch): + from core.rag.parsers import Page + + seen = {} + + def fake_parse(data, *, max_pages = None): + seen["max_pages"] = max_pages + pages = [Page(text = "x" * 1000, page_number = i, char_count = 1000) for i in range(1, 51)] + return pages, 60 # document actually has more pages than the cap + + monkeypatch.setattr("core.rag.parsers.parse_pdf_bytes", fake_parse) + text = tools._extract_pdf_text(b"unused") + assert seen["max_pages"] == tools._MAX_WEB_PDF_PAGES + assert len(text) <= tools._MAX_PAGE_CHARS + assert "text limited to 16,000 characters" in text + assert "page processing capped at 50 pages" in text + + +def test_pdf_exactly_at_page_cap_not_marked_capped(monkeypatch): + from core.rag.parsers import Page + + # Exactly _MAX_WEB_PDF_PAGES pages are fully read, so no "capped" marker. + monkeypatch.setattr( + "core.rag.parsers.parse_pdf_bytes", + lambda data, *, max_pages = None: ( + [Page(text = "short", page_number = i, char_count = 5) for i in range(1, 51)], + 50, + ), + ) + text = tools._extract_pdf_text(b"unused") + assert "page processing capped" not in text + assert "## Page 50\n\nshort" in text + + +def test_pdf_page_cap_does_not_claim_later_pages_are_textless(monkeypatch): + from core.rag.parsers import Page + monkeypatch.setattr( + "core.rag.parsers.parse_pdf_bytes", + lambda data, *, max_pages = None: ( + [Page(text = "", page_number = i, char_count = 0) for i in range(1, 51)], + 60, + ), + ) + assert tools._extract_pdf_text(b"unused") == ( + "(PDF contains no extractable text in the first 50 pages)" + ) + + +def test_pdf_result_discarded_after_fetch_deadline(monkeypatch): + clock = {"time": 1000.0} + monkeypatch.setattr(tools.time, "monotonic", lambda: clock["time"]) + + def slow_extract(data): + clock["time"] += 10.0 + return "late PDF text" + + monkeypatch.setattr(tools, "_extract_pdf_text", slow_extract) + out = _fetch_with(monkeypatch, _pdf_bytes("Readable text"), "application/pdf") + assert out == "Failed to fetch URL: timed out." def test_text_octet_stream_kept_after_sniffing(monkeypatch): @@ -154,7 +275,6 @@ def test_valid_utf8_binary_caught_by_control_chars(monkeypatch): @pytest.mark.parametrize( "magic", [ - b"%PDF-", b"PK\x03\x04", b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1", b"\x1f\x8b", @@ -180,8 +300,8 @@ def test_text_labeled_binary_caught_by_magic(monkeypatch, magic): b"\t\xef\xbb\xbf ", ], ) -def test_pdf_magic_after_harmless_prefix(monkeypatch, prefix): - body = prefix + b"%PDF-1.7\n" + b"1 0 obj<>endobj\n" * 100 +def test_binary_magic_after_harmless_prefix(monkeypatch, prefix): + body = prefix + b"\x1f\x8b" + b" printable text-heavy body" * 100 out = _fetch_with(monkeypatch, body, "text/plain") assert "binary content" in out @@ -243,10 +363,10 @@ def test_html_page_unaffected(monkeypatch): def test_content_type_sanitized_in_message(monkeypatch): # Do not echo obs-folded header content into the model response. - out = _fetch_with(monkeypatch, b"\x00\x01\x02" * 500, "application/pdf\r\n data: injected") + out = _fetch_with(monkeypatch, b"PK\x03\x04" * 500, "application/zip\r\n data: injected") assert "\n" not in out and "\r" not in out assert "injected" not in out - assert "application/pdf" in out + assert "application/zip" in out @pytest.mark.parametrize( diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 8d6c919ebd..117ad7b780 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -49,6 +49,11 @@ logger = get_logger(__name__) # and spawn workers copy os.environ. setdefault so an explicit user override wins. os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID") +# Studio workers can import MLX without importing unsloth first, so mirror the +# package bootstrap here. Keep an explicit user value authoritative. +if platform.system() == "Darwin" and platform.machine() == "arm64": + os.environ.setdefault("AGX_RELAX_CDM_CTXSTORE_TIMEOUT", "1") + # ========== Device Enum ========== diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index 25d88aeeea..b6e218793a 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -1551,7 +1551,7 @@ export function AppSidebar() {
-
+ {/* min-w-0 so long names truncate instead of overflowing; + pr on the button reserves room for the settings cog */} +
{displayTitle} Unsloth
- {/* settings cog (replaces the up/down chevron) */} - + {/* settings cog; sibling of the trigger (buttons cannot nest), + overlaid on the row's right edge, opens settings directly */} + diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 57213f07cb..62b8af6e3a 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -1326,7 +1326,9 @@ const ThreadWelcome: FC<{ useEffect(() => { // Prefer the nickname; otherwise first name only. Blank falls back to none. - const name = nickname.trim() || (displayName.trim().split(/\s+/)[0] ?? ""); + const raw = nickname.trim() || (displayName.trim().split(/\s+/)[0] ?? ""); + // Cap very long names so the greeting stays on one line. + const name = raw.length > 20 ? `${raw.slice(0, 20)}…` : raw; setWelcome(buildWelcome(new Date().getHours(), name)); }, [displayName, nickname]); @@ -2707,6 +2709,7 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({ const setCodeToolsEnabled = useChatRuntimeStore((s) => s.setCodeToolsEnabled); const artifactsEnabled = useChatRuntimeStore((s) => s.artifactsEnabled); const setArtifactsEnabled = useChatRuntimeStore((s) => s.setArtifactsEnabled); + const showCanvasMenuItem = useChatRuntimeStore((s) => s.showCanvasMenuItem); const mcpEnabledForChat = useChatRuntimeStore((s) => s.mcpEnabledForChat); const setMcpEnabledForChat = useChatRuntimeStore( (s) => s.setMcpEnabledForChat, @@ -2957,7 +2960,8 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({ ), - canvas: ( + // Hidden by default; enabled from Settings > Chat > Canvas. + canvas: showCanvasMenuItem ? ( setArtifactsEnabled(!artifactsEnabled)} @@ -2968,7 +2972,7 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({ ) : null} - ), + ) : null, bypassPermissions: , projects: ( diff --git a/studio/frontend/src/components/ui/dropdown-menu.tsx b/studio/frontend/src/components/ui/dropdown-menu.tsx index fe0270a63e..ea6dbbb6e7 100644 --- a/studio/frontend/src/components/ui/dropdown-menu.tsx +++ b/studio/frontend/src/components/ui/dropdown-menu.tsx @@ -2,10 +2,11 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { DropdownMenu as DropdownMenuPrimitive } from "radix-ui"; -import type * as React from "react"; +import * as React from "react"; import { Tick02Icon } from "@/lib/tick-icon"; import { ChevronRightStandardIcon } from "@/lib/chevron-icons"; +import { useIsMobile } from "@/hooks/use-mobile"; import { cn } from "@/lib/utils"; import { HugeiconsIcon } from "@hugeicons/react"; @@ -206,6 +207,14 @@ function DropdownMenuShortcut({ ); } +function assignRef(ref: React.Ref | undefined, value: T | null) { + if (typeof ref === "function") { + ref(value); + } else if (ref) { + ref.current = value; + } +} + function DropdownMenuSub({ ...props }: React.ComponentProps) { @@ -242,17 +251,63 @@ function DropdownMenuSubTrigger({ function DropdownMenuSubContent({ className, + sideOffset, + style, + ref, ...props }: React.ComponentProps) { + const isMobile = useIsMobile(); + const [contentWidth, setContentWidth] = React.useState(0); + const resizeObserverRef = React.useRef(null); + const composedRef = React.useCallback( + ( + element: React.ComponentRef< + typeof DropdownMenuPrimitive.SubContent + > | null, + ) => { + resizeObserverRef.current?.disconnect(); + resizeObserverRef.current = null; + assignRef(ref, element); + if (!element) return; + + const updateContentWidth = () => { + setContentWidth(element.offsetWidth); + }; + updateContentWidth(); + + if (typeof ResizeObserver !== "undefined") { + resizeObserverRef.current = new ResizeObserver(updateContentWidth); + resizeObserverRef.current.observe(element); + } + }, + [ref], + ); + + React.useEffect( + () => () => { + resizeObserverRef.current?.disconnect(); + }, + [], + ); + + const compactSideOffset = + isMobile && contentWidth > 0 ? -contentWidth : sideOffset; return ( // Portaled like DropdownMenuContent: rendered inline, the fixed popper // wrapper is a descendant of the parent menu's scroll container, so any // transform there turns on overflow clipping and hides the submenu. s.artifactsEnabled); const setArtifactsEnabled = useChatRuntimeStore((s) => s.setArtifactsEnabled); + const showCanvasMenuItem = useChatRuntimeStore((s) => s.showCanvasMenuItem); const permissionMode = useChatRuntimeStore((s) => s.permissionMode); const mcpEnabledForChat = useChatRuntimeStore((s) => s.mcpEnabledForChat); const setMcpEnabledForChat = useChatRuntimeStore( @@ -1411,7 +1412,8 @@ export function SharedComposer({ ), - canvas: ( + // Hidden by default; enabled from Settings > Chat > Canvas. + canvas: showCanvasMenuItem ? ( setArtifactsEnabled(!artifactsEnabled)} @@ -1422,7 +1424,7 @@ export function SharedComposer({ ) : null} - ), + ) : null, bypassPermissions: , projects: ( 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 2a81fe6ea8..af99458349 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -2,7 +2,11 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import type { RememberedLoadSettings } from "@/components/assistant-ui/model-selector/remembered-load-settings"; -import { cancelStagedModelDownload } from "@/features/hub"; +import { + cancelStagedModelDownload, + mirrorHfTokenInto, + useHfTokenStore, +} from "@/features/hub"; import { toast } from "@/lib/toast"; import { create } from "zustand"; import { isExternalModelId, parseExternalModelId } from "../external-providers"; @@ -23,14 +27,15 @@ import { savePersistedChatSettingsPatch, } from "../utils/chat-settings-storage"; import { useExternalProvidersStore } from "./external-providers-store"; +import { PLUS_MENU_PINS_STORAGE_KEY } from "./plus-menu-prefs-store"; -const HF_TOKEN_KEY = "unsloth_hf_token"; -const HF_TOKEN_CHANGED_EVENT = "unsloth:hf-token-changed"; export const CHAT_REASONING_ENABLED_KEY = "unsloth_chat_reasoning_enabled"; export const CHAT_TOOLS_ENABLED_KEY = "unsloth_chat_tools_enabled"; export const CHAT_CODE_TOOLS_ENABLED_KEY = "unsloth_chat_code_tools_enabled"; export const CHAT_IMAGE_TOOLS_ENABLED_KEY = "unsloth_chat_image_tools_enabled"; export const CHAT_ARTIFACTS_ENABLED_KEY = "unsloth_chat_artifacts_enabled"; +export const CHAT_SHOW_CANVAS_MENU_ITEM_KEY = + "unsloth_chat_show_canvas_menu_item"; export const CHAT_COLLAPSE_HTML_ARTIFACTS_KEY = "unsloth_chat_collapse_html_artifacts"; export const CHAT_ALLOW_ARTIFACT_NETWORK_ACCESS_KEY = @@ -357,6 +362,24 @@ function saveBool(key: string, value: boolean): void { } } +// The visibility flag shipped after the menu pins, so when it is absent, +// profiles that had explicitly pinned Canvas keep it visible. +function loadShowCanvasMenuItem(): boolean { + const stored = loadOptionalBool(CHAT_SHOW_CANVAS_MENU_ITEM_KEY); + if (stored !== null) return stored; + if (!canUseStorage()) return false; + try { + const raw = localStorage.getItem(PLUS_MENU_PINS_STORAGE_KEY); + if (raw === null) return false; + const parsed = JSON.parse(raw) as { + state?: { pins?: { canvas?: boolean } }; + }; + return parsed.state?.pins?.canvas === true; + } catch { + return false; + } +} + /** * "full" is intentionally not restorable: it disables the sandbox and every * confirmation gate, so it must be re-enabled (through the warning dialog) @@ -474,17 +497,6 @@ export function saveSpeculativeType(value: string | null): void { } } -function notifyHfTokenChanged(value: string): void { - if (!canUseStorage()) return; - try { - window.dispatchEvent( - new CustomEvent(HF_TOKEN_CHANGED_EVENT, { detail: value }), - ); - } catch { - // ignore - } -} - /** A local model staged for a deferred load (see `pendingSelection`). Shape is * a subset of the load hook's `SelectedModelInput`, structurally assignable. */ export type PendingModelSelection = { @@ -642,6 +654,8 @@ type ChatRuntimeStore = { codeToolsEnabled: boolean; imageToolsEnabled: boolean; artifactsEnabled: boolean; + // Whether the Canvas toggle is offered in the composer + menu (hidden by default). + showCanvasMenuItem: boolean; collapseHtmlArtifacts: boolean; allowArtifactNetworkAccess: boolean; mcpEnabledForChat: boolean; @@ -818,6 +832,7 @@ type ChatRuntimeStore = { enabled: boolean, options?: { persist?: boolean }, ) => void; + setShowCanvasMenuItem: (enabled: boolean) => void; setCollapseHtmlArtifacts: (enabled: boolean) => void; setAllowArtifactNetworkAccess: (enabled: boolean) => void; setMcpEnabledForChat: (enabled: boolean) => void; @@ -1120,7 +1135,7 @@ export const useChatRuntimeStore = create((set, get) => ({ runningByThreadId: {}, cancelByThreadId: {}, autoTitle: false, - hfToken: loadString(HF_TOKEN_KEY, ""), + hfToken: useHfTokenStore.getState().token, modelsError: null, lastModelLoadError: null, activeGgufVariant: null, @@ -1147,6 +1162,7 @@ export const useChatRuntimeStore = create((set, get) => ({ codeToolsEnabled: loadBool(CHAT_CODE_TOOLS_ENABLED_KEY, false), imageToolsEnabled: loadBool(CHAT_IMAGE_TOOLS_ENABLED_KEY, false), artifactsEnabled: loadBool(CHAT_ARTIFACTS_ENABLED_KEY, false), + showCanvasMenuItem: loadShowCanvasMenuItem(), collapseHtmlArtifacts: loadBool(CHAT_COLLAPSE_HTML_ARTIFACTS_KEY, false), allowArtifactNetworkAccess: loadBool( CHAT_ALLOW_ARTIFACT_NETWORK_ACCESS_KEY, @@ -1324,11 +1340,7 @@ export const useChatRuntimeStore = create((set, get) => ({ setScalarSettingVersion("autoTitle", autoTitle, state.autoTitle); return { autoTitle }; }), - setHfToken: (hfToken) => { - saveString(HF_TOKEN_KEY, hfToken); - set({ hfToken }); - notifyHfTokenChanged(hfToken); - }, + setHfToken: (hfToken) => useHfTokenStore.getState().setToken(hfToken), setModelsError: (modelsError) => set({ modelsError }), setLastModelLoadError: (lastModelLoadError) => set({ lastModelLoadError }), setCheckpoint: (modelId, ggufVariant) => @@ -1504,6 +1516,11 @@ export const useChatRuntimeStore = create((set, get) => ({ } return { artifactsEnabled }; }), + setShowCanvasMenuItem: (showCanvasMenuItem) => + set(() => { + saveBool(CHAT_SHOW_CANVAS_MENU_ITEM_KEY, showCanvasMenuItem); + return { showCanvasMenuItem }; + }), setCollapseHtmlArtifacts: (collapseHtmlArtifacts) => set((state) => { saveBool(CHAT_COLLAPSE_HTML_ARTIFACTS_KEY, collapseHtmlArtifacts); @@ -1807,6 +1824,12 @@ export const useChatRuntimeStore = create((set, get) => ({ setContextUsage: (contextUsage) => set({ contextUsage }), })); +// Mirror token edits made through the shared store (e.g. Studio's field). +const unsubscribeHfTokenMirror = mirrorHfTokenInto(useChatRuntimeStore); +if (import.meta.hot) { + import.meta.hot.dispose(unsubscribeHfTokenMirror); +} + export function resolveSpeculativeSettingsForLoad({ usePersistedPreference = false, }: { diff --git a/studio/frontend/src/features/chat/stores/plus-menu-prefs-store.ts b/studio/frontend/src/features/chat/stores/plus-menu-prefs-store.ts index 822c66eeb5..960649a901 100644 --- a/studio/frontend/src/features/chat/stores/plus-menu-prefs-store.ts +++ b/studio/frontend/src/features/chat/stores/plus-menu-prefs-store.ts @@ -44,6 +44,8 @@ const DEFAULT_PINS: Record = { bypassPermissions: false, }; +export const PLUS_MENU_PINS_STORAGE_KEY = "unsloth_plus_menu_pins"; + export interface PlusMenuPrefsState { pins: Record; setPin: (id: PlusMenuItemId, value: boolean) => void; @@ -72,7 +74,7 @@ export const usePlusMenuPrefsStore = create()( })), }), { - name: "unsloth_plus_menu_pins", + name: PLUS_MENU_PINS_STORAGE_KEY, // Backfill any ids added in a later release so persisted state from an // older version still resolves every menu item. merge: (persisted, current) => { diff --git a/studio/frontend/src/features/hub/index.ts b/studio/frontend/src/features/hub/index.ts index ddcef146b0..3515f6ca76 100644 --- a/studio/frontend/src/features/hub/index.ts +++ b/studio/frontend/src/features/hub/index.ts @@ -2,3 +2,8 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 export { cancelStagedModelDownload } from "./download-manager"; +export { + getHfToken, + mirrorHfTokenInto, + useHfTokenStore, +} from "./stores/hf-token-store"; diff --git a/studio/frontend/src/features/hub/stores/hf-token-store.ts b/studio/frontend/src/features/hub/stores/hf-token-store.ts index b1e2560f02..499ba9f644 100644 --- a/studio/frontend/src/features/hub/stores/hf-token-store.ts +++ b/studio/frontend/src/features/hub/stores/hf-token-store.ts @@ -5,11 +5,9 @@ import { create } from "zustand"; import { bumpInventoryVersion } from "./inventory-events"; const HF_TOKEN_KEY = "unsloth_hf_token"; -const HF_TOKEN_CHANGED_EVENT = "unsloth:hf-token-changed"; const LEGACY_TRAINING_KEY = "unsloth_training_config_v1"; let storageSyncStarted = false; let storageSyncListener: ((event: StorageEvent) => void) | null = null; -let tokenChangedListener: ((event: Event) => void) | null = null; function canUseStorage(): boolean { return typeof window !== "undefined"; @@ -63,10 +61,6 @@ function stopStorageSync(): void { window.removeEventListener("storage", storageSyncListener); storageSyncListener = null; } - if (tokenChangedListener !== null) { - window.removeEventListener(HF_TOKEN_CHANGED_EVENT, tokenChangedListener); - tokenChangedListener = null; - } storageSyncStarted = false; } @@ -100,10 +94,6 @@ export const useHfTokenStore = create((set) => { applyToken(event.newValue ?? "", false); }; window.addEventListener("storage", storageSyncListener); - tokenChangedListener = (event) => { - applyToken((event as CustomEvent).detail ?? "", false); - }; - window.addEventListener(HF_TOKEN_CHANGED_EVENT, tokenChangedListener); } return { @@ -117,6 +107,21 @@ export function getHfToken(): string { return useHfTokenStore.getState().token; } +// Keep a plain zustand store's `hfToken` field in sync with the shared token: +// seed the current value, then mirror later edits. Returns the unsubscribe so +// callers can wire it to HMR disposal. +export function mirrorHfTokenInto(store: { + getState: () => T; + setState: (partial: Partial) => void; +}): () => void { + store.setState({ hfToken: getHfToken() } as Partial); + return useHfTokenStore.subscribe((state) => { + if (store.getState().hfToken !== state.token) { + store.setState({ hfToken: state.token } as Partial); + } + }); +} + // HF's JS client throws on a non-empty token that isn't `hf_...` instead of // browsing anonymously, so treat anything malformed as no token. export function hfApiToken( diff --git a/studio/frontend/src/features/settings/components/api-monitor-console.tsx b/studio/frontend/src/features/settings/components/api-monitor-console.tsx index ab67d4e006..ef4daa1349 100644 --- a/studio/frontend/src/features/settings/components/api-monitor-console.tsx +++ b/studio/frontend/src/features/settings/components/api-monitor-console.tsx @@ -111,7 +111,7 @@ function MonitorEntry({ const reply = replyText || (entry.status === "running" ? "Waiting..." : "No reply"); return ( -
+