diff --git a/README.md b/README.md index 948d84a789..562c35ff1a 100644 --- a/README.md +++ b/README.md @@ -72,10 +72,13 @@ Unsloth Studio (Beta) works on **Windows, Linux, WSL** and **macOS**. ```bash curl -fsSL https://unsloth.ai/install.sh | sh ``` +Use the same command to update. + #### Windows: ```powershell irm https://unsloth.ai/install.ps1 | iex ``` +Use the same command to update. #### Launch ```bash @@ -83,9 +86,6 @@ unsloth studio -p 8888 ``` For cloud or global access, add `-H 0.0.0.0`. By default, Unsloth is accessible only locally. -#### Update -To update, use the same install commands above or use `unsloth studio update`. - #### Docker Use our [Docker image](https://hub.docker.com/r/unsloth/unsloth) ```unsloth/unsloth``` container. Run: ```bash @@ -171,7 +171,9 @@ unsloth studio -p 8888 ``` Then to update : ```bash -unsloth studio update +cd unsloth && git pull +./install.sh --local +unsloth studio -p 8888 ``` #### Developer installs: Windows PowerShell: @@ -184,7 +186,9 @@ unsloth studio -p 8888 ``` Then to update : ```bash -unsloth studio update +cd unsloth && git pull +./install.sh --local +unsloth studio -p 8888 ``` #### Nightly: MacOS, Linux, WSL: diff --git a/install.ps1 b/install.ps1 index cab66f5ae1..d3942bf2fc 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1566,7 +1566,7 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.8" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.9" unsloth-zoo } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below @@ -1580,7 +1580,7 @@ shell.Run cmd, 0, False } } } else { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.8" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.9" unsloth-zoo } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -1627,7 +1627,7 @@ shell.Run cmd, 0, False if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.5.8" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.5.9" unsloth-zoo } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic } @@ -1639,7 +1639,7 @@ shell.Run cmd, 0, False } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.5.8" unsloth-zoo } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.5.9" unsloth-zoo } } else { $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -1667,7 +1667,7 @@ shell.Run cmd, 0, False Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.5.8" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.5.9" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) diff --git a/install.sh b/install.sh index f0af60e2d6..89c506bc9d 100755 --- a/install.sh +++ b/install.sh @@ -1530,17 +1530,51 @@ if [ -x "$VENV_DIR/bin/python" ]; then : > "$VENV_DIR/.unsloth-studio-owned" 2>/dev/null || true fi -# Guard against Python 3.13.8 torch import bug on Apple Silicon -# (skip when the user explicitly chose a version via --python) +# Guard against two independent Apple Silicon venv problems, in order: +# 1. uv may create the venv from a cached x86_64 (Rosetta) Python when a +# same-version x86_64 build is already cached (often because uv itself +# is an x86_64 build). That venv reports x86_64 to wheel resolvers, and +# PyTorch ships no macOS wheels on the CPU index for any architecture, +# so the torch install can never resolve. Recreate it with an +# arch-explicit arm64 CPython. +# 2. Python 3.13.8 has a known torch import bug. +# The two are independent: a venv may be x86_64 and, once recreated, still +# land on 3.13.8. So we re-inspect the interpreter between the checks instead +# of chaining them with elif, guaranteeing both invariants hold on whatever +# venv we end up with. Skip both when the user explicitly chose an interpreter +# via --python. if [ -z "$_USER_PYTHON" ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then - _PY_VER=$("$VENV_DIR/bin/python" -c \ - "import sys; print('{}.{}.{}'.format(*sys.version_info[:3]))" 2>/dev/null || echo "") + _inspect_venv() { + "$VENV_DIR/bin/python" -c \ + "import platform, sys; print(platform.machine(), '{}.{}.{}'.format(*sys.version_info[:3]))" \ + 2>/dev/null || echo " " + } + _info=$(_inspect_venv) + _VENV_ARCH=${_info%% *} + _PY_VER=${_info##* } + + if [ "$_VENV_ARCH" = "x86_64" ]; then + echo " WARNING: venv was created with an x86_64 (Rosetta) Python on Apple Silicon." + echo " Recreating venv with native arm64 Python ${PYTHON_VERSION}..." + rm -rf "$VENV_DIR" + run_install_cmd "recreate venv (arm64)" uv venv "$VENV_DIR" \ + --python "cpython-${PYTHON_VERSION}-macos-aarch64-none" + if [ -x "$VENV_DIR/bin/python" ]; then + : > "$VENV_DIR/.unsloth-studio-owned" 2>/dev/null || true + fi + # Re-inspect: the recreated arm64 venv may still be 3.13.8. + _info=$(_inspect_venv) + _VENV_ARCH=${_info%% *} + _PY_VER=${_info##* } + fi + if [ "$_PY_VER" = "3.13.8" ]; then echo " WARNING: Python 3.13.8 has a known torch import bug." echo " Recreating venv with Python 3.12..." rm -rf "$VENV_DIR" PYTHON_VERSION="3.12" - run_install_cmd "recreate venv" uv venv "$VENV_DIR" --python "$PYTHON_VERSION" + run_install_cmd "recreate venv" uv venv "$VENV_DIR" \ + --python "cpython-${PYTHON_VERSION}-macos-aarch64-none" if [ -x "$VENV_DIR/bin/python" ]; then : > "$VENV_DIR/.unsloth-studio-owned" 2>/dev/null || true fi @@ -2049,7 +2083,7 @@ if [ "$_MIGRATED" = true ]; then # to prevent transitive torch resolution. run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.5.8" unsloth-zoo + "unsloth>=2026.5.9" unsloth-zoo # Resolve pydantic WITH deps so pip pins pydantic-core to the # matching version (no-torch-runtime.txt below is --no-deps). # All transitive deps are torch-free. @@ -2062,7 +2096,7 @@ if [ "$_MIGRATED" = true ]; then else run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.5.8" unsloth-zoo + "unsloth>=2026.5.9" unsloth-zoo fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2266,7 +2300,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.5.8" unsloth-zoo + "unsloth>=2026.5.9" unsloth-zoo # Same pydantic-with-deps trick as the migrated branch. run_install_cmd "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic @@ -2284,7 +2318,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.5.8" unsloth-zoo + --upgrade-package unsloth "unsloth>=2026.5.9" unsloth-zoo substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -2316,7 +2350,7 @@ else tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.5.8" --torch-backend=auto + run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.5.9" --torch-backend=auto substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." diff --git a/pyproject.toml b/pyproject.toml index aef88d90f5..acc65f12cf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,7 +69,7 @@ triton = [ ] huggingfacenotorch = [ - "unsloth_zoo>=2026.5.4", + "unsloth_zoo>=2026.5.5", "wheel>=0.42.0", "packaging", "numpy", @@ -90,7 +90,7 @@ huggingfacenotorch = [ ] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.5.4", + "unsloth_zoo>=2026.5.5", "torchvision", "unsloth[triton]", ] @@ -580,7 +580,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2026.5.4", + "unsloth_zoo>=2026.5.5", "packaging", "tyro", "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0", diff --git a/studio/backend/core/data_recipe/service.py b/studio/backend/core/data_recipe/service.py index b4ec0ccd94..85b567885f 100644 --- a/studio/backend/core/data_recipe/service.py +++ b/studio/backend/core/data_recipe/service.py @@ -176,12 +176,21 @@ def build_mcp_providers( ) -> list: from data_designer.config.mcp import LocalStdioMCPProvider, MCPProvider # pyright: ignore[reportMissingImports] + # Same gate as the chat MCP path: stdio providers spawn a local subprocess, + # so only build them when this host allows it (desktop / explicit opt-in). + # Skip them otherwise so a recipe carried onto a hosted host cannot spawn. + from core.inference.mcp_client import stdio_mcp_enabled + + stdio_allowed = stdio_mcp_enabled() + providers: list[MCPProvider | LocalStdioMCPProvider] = [] for provider in recipe.get("mcp_providers", []): if not isinstance(provider, dict): continue provider_type = provider.get("provider_type") if provider_type == "stdio": + if not stdio_allowed: + continue env = provider.get("env") if not isinstance(env, dict): env = {} diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index c8250bd4f5..38ba635627 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -2591,6 +2591,105 @@ class LlamaCppBackend: # ── Lifecycle ───────────────────────────────────────────────── + # GGUF ``general.architecture`` values for diffusion / image models. + # llama.cpp proper has no such architectures, so loading one as a chat + # model dies with "unknown model architecture: ''". These match + # the patched stable-diffusion.cpp / ComfyUI-GGUF enums (LLM_ARCH_FLUX, + # LLM_ARCH_QWEN_IMAGE, ...). Unsloth publishes FLUX and Qwen-Image GGUFs + # under https://huggingface.co/collections/unsloth/unsloth-diffusion-ggufs. + # Matched exactly (not as a substring) so a chat arch merely containing a + # short token like "wan"/"sd1" (e.g. "taiwan") is not misrouted to Images. + _DIFFUSION_ARCHES = frozenset( + ( + "qwen_image", + "flux", + "sd1", + "sdxl", + "sd3", + "aura", + "hidream", + "cosmos", + "ltxv", + "hyvid", + "wan", + "lumina2", + ) + ) + + @staticmethod + def _classify_llama_start_failure( + output: str, + gguf_path: Optional[str], + model_identifier: Optional[str], + ) -> str: + """Explain *why* llama-server failed to start, from its output. + + Several distinct failures all otherwise collapse into the same + opaque "invalid GGUF or out of memory" message. The worst case is + a diffusion / image GGUF (FLUX, Qwen-Image, ...) loaded as a chat + model: the file is perfectly valid and there is plenty of memory, + but llama.cpp has no such architecture, so the user is told to free + memory that was never the problem (issue #5842). Pick the most + specific message the captured output supports. + """ + lowered = (output or "").lower() + + # Detect Ollama source up front so the arch branch can keep the + # Ollama hint instead of the generic "unsupported arch" message. + gguf = gguf_path or "" + is_ollama = ( + ".studio_links" in gguf + or os.sep + "ollama_links" + os.sep in gguf + or os.sep + ".cache" + os.sep + "ollama" + os.sep in gguf + or (model_identifier or "").startswith("ollama/") + ) + + # "unknown model architecture: ''": diffusion -> Images page, + # Ollama -> Ollama hint, else a precise "unsupported" message. Exact + # match so chat archs are never misrouted. + arch_match = re.search(r"unknown model architecture:\s*'([^']+)'", lowered) + if arch_match: + arch = arch_match.group(1) + if arch in LlamaCppBackend._DIFFUSION_ARCHES: + return ( + f"'{arch}' is a diffusion (image-generation) GGUF, which " + "llama-server cannot run as a chat/completion model. Use " + "Studio's Images page to generate with local diffusion " + "GGUFs such as FLUX and Qwen-Image." + ) + if is_ollama: + return ( + "Some Ollama models do not work with llama.cpp. Try a " + "different model, or use this model directly through " + "Ollama instead." + ) + return ( + f"llama.cpp does not support this GGUF's model architecture " + f"('{arch}'). The file is valid, but this model type cannot " + "be run with llama-server." + ) + + # Other Ollama compat failures that do not name an arch. Only when + # the output shows a GGUF compat issue, not OOM / missing binaries. + if is_ollama: + gguf_compat_hints = ( + "key not found", + "unknown model architecture", + "failed to load model", + ) + if any(h in lowered for h in gguf_compat_hints): + return ( + "Some Ollama models do not work with llama.cpp. Try a " + "different model, or use this model directly through " + "Ollama instead." + ) + + # Fallback: genuinely unknown failure (OOM, missing binary, ...). + return ( + "llama-server failed to start. " + "Check that the GGUF file is valid and you have enough memory." + ) + def load_model( self, *, @@ -3383,31 +3482,12 @@ class LlamaCppBackend: # Wait for llama-server to become healthy if not self._wait_for_health(timeout = 600.0): self._kill_process() - _gguf = gguf_path or "" - _is_ollama = ( - ".studio_links" in _gguf - or os.sep + "ollama_links" + os.sep in _gguf - or os.sep + ".cache" + os.sep + "ollama" + os.sep in _gguf - or (self._model_identifier or "").startswith("ollama/") - ) - # Only show the Ollama-specific message when the server - # output indicates a GGUF compatibility issue, not for - # unrelated failures like OOM or missing binaries. - if _is_ollama: - _output = "\n".join(self._stdout_lines[-50:]).lower() - _gguf_compat_hints = ( - "key not found", - "unknown model architecture", - "failed to load model", - ) - if any(h in _output for h in _gguf_compat_hints): - raise RuntimeError( - "Some Ollama models do not work with llama.cpp. " - "Try a different model, or use this model directly through Ollama instead." - ) raise RuntimeError( - "llama-server failed to start. " - "Check that the GGUF file is valid and you have enough memory." + self._classify_llama_start_failure( + "\n".join(self._stdout_lines[-50:]), + gguf_path, + self._model_identifier, + ) ) self._healthy = True diff --git a/studio/backend/core/inference/mcp_client.py b/studio/backend/core/inference/mcp_client.py index a5e614899d..2ed1a630dc 100644 --- a/studio/backend/core/inference/mcp_client.py +++ b/studio/backend/core/inference/mcp_client.py @@ -5,6 +5,9 @@ from __future__ import annotations import asyncio import json +import os +import shlex +import sys from typing import Any, Optional from loggers import get_logger @@ -16,7 +19,55 @@ MCP_TOOL_PREFIX = "mcp__" _oauth_token_store = None +def is_stdio(address: str) -> bool: + """A non-HTTP address is a local stdio command, e.g. + 'npx -y @modelcontextprotocol/server-filesystem /path'.""" + return not address.strip().lower().startswith(("http://", "https://")) + + +def parse_stdio_command(address: str) -> list[str]: + """Split a stdio command line into argv. Shared by route validation and the + transport so both agree on quoting (notably Windows backslash paths).""" + posix = sys.platform != "win32" + parts = shlex.split(address, posix = posix) + if not posix: + # posix=False keeps backslash paths intact but also keeps the surrounding + # quotes on a token. Strip a matched pair so the argv reaches the + # subprocess clean ('"C:\\Program Files\\node"' -> C:\\Program Files\\node). + parts = [ + p[1:-1] if len(p) >= 2 and p[0] == p[-1] and p[0] in "\"'" else p + for p in parts + ] + return parts + + +def stdio_mcp_enabled() -> bool: + """stdio MCP servers spawn local processes as the backend user (and bypass + the python/terminal sandbox), so they are only allowed when the backend + host is the user's own machine. The Tauri desktop app sets + UNSLOTH_STUDIO_ALLOW_STDIO_MCP=1 (see main.py); advanced localhost / + self-hosted users can opt in with the same variable. It stays off for + Colab and any network (0.0.0.0) bind.""" + return os.environ.get("UNSLOTH_STUDIO_ALLOW_STDIO_MCP") == "1" + + +# Probe timeouts for discovering a server's tool list. OAuth needs minutes for +# first-connect/expired-token browser sign-in; stdio allows for first-run +# package download (e.g. `npx -y ...`); HTTP fails fast. +_HTTP_PROBE_TIMEOUT = 8.0 +_OAUTH_PROBE_TIMEOUT = 305.0 +_STDIO_PROBE_TIMEOUT = 60.0 + + +def probe_timeout(address: str, use_oauth: bool) -> float: + if use_oauth: + return _OAUTH_PROBE_TIMEOUT + return _STDIO_PROBE_TIMEOUT if is_stdio(address) else _HTTP_PROBE_TIMEOUT + + def parse_server_headers(server: dict) -> Optional[dict]: + """Parsed headers_json. For stdio servers this dict is the process + environment instead of HTTP headers (see _client).""" raw = server.get("headers_json") if not raw: return None @@ -63,6 +114,28 @@ async def clear_oauth_tokens_async(url: str) -> None: def _client(url: str, headers: Optional[dict], use_oauth: bool = False): from fastmcp import Client + + if is_stdio(url): + # Belt-and-suspenders: never spawn unless stdio is enabled on this host. + if not stdio_mcp_enabled(): + raise PermissionError("stdio MCP servers are disabled on this host") + from fastmcp.client.transports import StdioTransport + + parts = parse_stdio_command(url) + if not parts: + raise ValueError(f"Empty stdio command: {url!r}") + # env vars ride the headers field (merged over the SDK's safe default env). + # keep_alive=False tears the subprocess down on exit, so a one-shot + # probe/tool call never leaves an orphan process. + return Client( + StdioTransport( + command = parts[0], + args = parts[1:], + env = headers or None, + keep_alive = False, + ) + ) + from fastmcp.client.transports import SSETransport, StreamableHttpTransport from fastmcp.mcp_config import infer_transport_type_from_url diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 9572a2169a..baf1236456 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -28,8 +28,11 @@ import urllib.request from core.inference.mcp_client import ( MCP_TOOL_PREFIX, call_tool_sync, + is_stdio, list_tools_async, parse_server_headers, + probe_timeout, + stdio_mcp_enabled, ) from storage import mcp_servers_db @@ -568,17 +571,19 @@ def _mcp_specs_for_server(server: dict, mcp_tools: list[dict]) -> list[dict]: async def get_enabled_mcp_tools() -> list[dict]: servers = [s for s in mcp_servers_db.list_servers() if s.get("is_enabled")] + # Never spawn stdio servers when stdio is disabled on this host (e.g. a DB + # carried over from a desktop install onto a Colab / network deployment). + if not stdio_mcp_enabled(): + servers = [s for s in servers if not is_stdio(s["url"])] if not servers: return [] - # OAuth probes need minutes for first-connect/expired-token browser - # sign-in; non-OAuth probes fail fast. Matches routes/mcp_servers.py. results = await asyncio.gather( *( list_tools_async( url = s["url"], headers = parse_server_headers(s), - timeout = 305.0 if s.get("use_oauth") else 8.0, + timeout = probe_timeout(s["url"], bool(s.get("use_oauth"))), use_oauth = bool(s.get("use_oauth")), ) for s in servers @@ -630,6 +635,8 @@ def execute_tool( return f"Error: MCP server '{server_id}' not found" if not server.get("is_enabled"): return f"Error: MCP server '{server_id}' is disabled" + if is_stdio(server["url"]) and not stdio_mcp_enabled(): + return f"Error: stdio MCP server '{server_id}' is disabled on this host" return call_tool_sync( url = server["url"], headers = parse_server_headers(server), diff --git a/studio/backend/main.py b/studio/backend/main.py index 6b8ac438c0..b6cf58a02c 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -297,6 +297,11 @@ def _load_desktop_owner() -> dict[str, str] | None: _DESKTOP_OWNER = _load_desktop_owner() +# The Tauri desktop app runs the backend on the owner's own machine, so local +# stdio MCP servers are safe there. setdefault lets an explicit "0" opt out. +if _DESKTOP_OWNER: + os.environ.setdefault("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1") + def _desktop_owner() -> dict[str, str] | None: return _DESKTOP_OWNER diff --git a/studio/backend/routes/data_recipe/mcp.py b/studio/backend/routes/data_recipe/mcp.py index 1f5c0f34e0..7184934ce9 100644 --- a/studio/backend/routes/data_recipe/mcp.py +++ b/studio/backend/routes/data_recipe/mcp.py @@ -36,8 +36,18 @@ def list_mcp_tools(payload: McpToolsListRequest) -> McpToolsListResponse: providers: list[McpToolsProviderResult] = [] tool_to_providers: dict[str, list[str]] = defaultdict(list) + from core.inference.mcp_client import stdio_mcp_enabled + for provider_payload in payload.mcp_providers: provider_name = str(provider_payload.get("name", "")).strip() + if provider_payload.get("provider_type") == "stdio" and not stdio_mcp_enabled(): + providers.append( + McpToolsProviderResult( + name = provider_name, + error = "Local (stdio) MCP servers are disabled on this host.", + ) + ) + continue built = build_mcp_providers({"mcp_providers": [provider_payload]}) if len(built) != 1: providers.append( diff --git a/studio/backend/routes/mcp_servers.py b/studio/backend/routes/mcp_servers.py index a7501d1691..6c63bc20ca 100644 --- a/studio/backend/routes/mcp_servers.py +++ b/studio/backend/routes/mcp_servers.py @@ -11,8 +11,12 @@ from fastapi import APIRouter, Depends, HTTPException from auth.authentication import get_current_subject from core.inference.mcp_client import ( clear_oauth_tokens_async, + is_stdio, list_tools_async, parse_server_headers, + parse_stdio_command, + probe_timeout, + stdio_mcp_enabled, ) from models.mcp_servers import ( McpServerCreate, @@ -28,16 +32,30 @@ logger = structlog.get_logger(__name__) router = APIRouter() -_PROBE_TIMEOUT_SECONDS = 8.0 -# When OAuth probes need to open a browser, wait long enough for the user to -# sign in. Matches fastmcp's default OAuth callback_timeout (300 s) + slack. -_OAUTH_PROBE_TIMEOUT_SECONDS = 305.0 - - def _validate_url(url: str) -> str: trimmed = (url or "").strip() if not trimmed: raise HTTPException(status_code = 400, detail = "url must not be empty") + # When stdio is enabled on this host, a non-HTTP value is a local command. + # Reuse this field so stdio servers ride the existing CRUD/storage with no + # schema change. When stdio is disabled the value falls through to the + # http-only validation below, so non-HTTP input is just a bad URL (400). + if stdio_mcp_enabled() and is_stdio(trimmed): + try: + parts = parse_stdio_command(trimmed) + except ValueError as exc: + raise HTTPException(status_code = 400, detail = f"Invalid command: {exc}") + if not parts or not parts[0].strip(): + raise HTTPException(status_code = 400, detail = "command must not be empty") + if "://" in parts[0]: + # A URL-scheme first token is a mistyped URL, not a command. Reject + # it cleanly instead of exec-ing it (mirrors the frontend check). + raise HTTPException( + status_code = 400, + detail = "Enter an http(s):// URL, or a local command whose " + "first token is an executable (not a URL).", + ) + return trimmed parsed = urlparse(trimmed) if parsed.scheme not in ("http", "https"): raise HTTPException( @@ -91,6 +109,9 @@ async def create_mcp_server( raise HTTPException(status_code = 400, detail = "display_name must not be empty") url = _validate_url(payload.url) headers = _normalize_headers(payload.headers) + # OAuth is HTTP-only; force it off for stdio commands so a stale flag can't + # push the probe onto the 305s OAuth timeout. Backend is the enforcer. + use_oauth = payload.use_oauth and not is_stdio(url) server_id = uuid.uuid4().hex[:16] mcp_servers_db.create_server( @@ -99,7 +120,7 @@ async def create_mcp_server( url = url, headers_json = json.dumps(headers) if headers else None, is_enabled = payload.is_enabled, - use_oauth = payload.use_oauth, + use_oauth = use_oauth, ) return _row_to_response(mcp_servers_db.get_server(server_id)) @@ -132,6 +153,9 @@ def _changes_from_payload(payload: McpServerUpdate) -> dict: status_code = 400, detail = "use_oauth must be true or false" ) changes["use_oauth"] = payload.use_oauth + # stdio is OAuth-less: drop a stale OAuth flag when switching to a command. + if "url" in changes and is_stdio(changes["url"]): + changes["use_oauth"] = False return changes @@ -147,6 +171,15 @@ async def update_mcp_server( changes = _changes_from_payload(payload) if not changes: raise HTTPException(status_code = 400, detail = "No fields to update") + # headers == HTTP headers (remote) or env vars (stdio). On a transport-type + # switch with no new headers, drop the old ones so env secrets are not + # re-sent as HTTP headers (or vice versa). + if ( + "url" in changes + and is_stdio(changes["url"]) != is_stdio(old["url"]) + and "headers_json" not in changes + ): + changes["headers_json"] = None # Clear persisted OAuth tokens when the URL changes or OAuth is # disabled; fastmcp keys tokens by URL and would otherwise let a # re-pointed server silently inherit the old account's credentials. @@ -180,15 +213,19 @@ async def refresh_mcp_server_tools( server = mcp_servers_db.get_server(server_id) if not server: raise HTTPException(status_code = 404, detail = "MCP server not found") + # Refresh uses the stored address, so re-check the stdio gate here too: a + # stdio row from a desktop DB must not spawn on a hosted/network host. + if is_stdio(server["url"]) and not stdio_mcp_enabled(): + raise HTTPException( + status_code = 400, detail = "stdio MCP servers are disabled on this host" + ) use_oauth = bool(server.get("use_oauth")) try: tools = await list_tools_async( url = server["url"], headers = parse_server_headers(server), - timeout = _OAUTH_PROBE_TIMEOUT_SECONDS - if use_oauth - else _PROBE_TIMEOUT_SECONDS, + timeout = probe_timeout(server["url"], use_oauth), use_oauth = use_oauth, ) except Exception as exc: # noqa: BLE001 — surface transport+timeout errors to UI @@ -212,9 +249,7 @@ async def test_mcp_server( tools = await list_tools_async( url = url, headers = headers, - timeout = _OAUTH_PROBE_TIMEOUT_SECONDS - if payload.use_oauth - else _PROBE_TIMEOUT_SECONDS, + timeout = probe_timeout(url, payload.use_oauth), use_oauth = payload.use_oauth, ) except Exception as exc: # noqa: BLE001 diff --git a/studio/backend/tests/test_llama_cpp_start_failure_classification.py b/studio/backend/tests/test_llama_cpp_start_failure_classification.py new file mode 100644 index 0000000000..e647ff2c7d --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_start_failure_classification.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for LlamaCppBackend._classify_llama_start_failure. + +When llama-server exits before becoming healthy, load_model turns its +captured stdout/stderr into a user-facing reason. A diffusion / image +GGUF (FLUX, Qwen-Image, ...) is a valid file with plenty of memory, so +the generic "invalid file or out of memory" message is actively +misleading (issue #5842). These tests pin the classification. +""" + +from __future__ import annotations + +import sys +import types as _types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Match the stubbing pattern in sibling tests so the module imports in a +# lightweight env without fastapi. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) +# Give the structlog stub a real get_logger: a bare ModuleType poisons +# sys.modules for later tests that call structlog.get_logger at import time. +_structlog_stub = _types.ModuleType("structlog") +_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger( + "structlog" +) +sys.modules.setdefault("structlog", _structlog_stub) +if not hasattr(sys.modules["structlog"], "get_logger"): + sys.modules["structlog"].get_logger = _structlog_stub.get_logger + +from core.inference.llama_cpp import LlamaCppBackend # noqa: E402 + +_classify = LlamaCppBackend._classify_llama_start_failure + +# Real llama-server failure lines (lower-cased downstream anyway). +_QWEN_IMAGE_OUT = ( + "load_model: loading model 'qwen-image-edit-2511-Q4_K_M.gguf'\n" + "llama_model_load: error loading model: unknown model architecture: 'qwen_image'\n" + "llama_model_load_from_file_impl: failed to load model" +) +_OOM_OUT = ( + "ggml_backend_cuda_buffer_type_alloc_buffer: allocating 12000.00 MiB on " + "device 0: cudaMalloc failed: out of memory" +) + + +class TestDiffusionArchitectures: + def test_qwen_image_routes_to_images_page(self): + msg = _classify(_QWEN_IMAGE_OUT, "/models/qwen-image.gguf", "local/qwen-image") + assert "diffusion" in msg.lower() + assert "Images page" in msg + assert "qwen_image" in msg + # Must NOT keep blaming memory / file validity. + assert "out of memory" not in msg.lower() + assert "enough memory" not in msg.lower() + + # Parametrize over the production set so new arches are auto-covered. + @pytest.mark.parametrize("arch", sorted(LlamaCppBackend._DIFFUSION_ARCHES)) + def test_every_diffusion_arch_is_recognised(self, arch): + out = f"error loading model: unknown model architecture: '{arch}'" + msg = _classify(out, f"/models/{arch}.gguf", f"local/{arch}") + assert "diffusion" in msg.lower() + assert "Images page" in msg + assert arch in msg + + +class TestUnsupportedNonDiffusionArchitecture: + def test_unknown_llm_arch_says_unsupported_not_oom(self): + out = "error loading model: unknown model architecture: 'some_new_llm'" + msg = _classify(out, "/models/x.gguf", "local/x") + assert "some_new_llm" in msg + assert "architecture" in msg.lower() + # Specific, not the misleading memory message. + assert "enough memory" not in msg.lower() + assert "diffusion" not in msg.lower() + + # Exact match: a chat arch merely containing a diffusion token (wan, + # sd1, flux, ...) must not be routed to the Images page. + @pytest.mark.parametrize( + "arch", + [ + "taiwan", # contains "wan" + "swan_llm", # contains "wan" + "fluxion", # contains "flux" + "sd1234", # contains "sd1" + "sd3_chat", # contains "sd3" + "aura2_text", # contains "aura" + "cosmos_reason", # contains "cosmos" + "qwen_image_text", # contains "qwen_image" + ], + ) + def test_arch_containing_diffusion_token_is_not_misrouted(self, arch): + out = f"error loading model: unknown model architecture: '{arch}'" + msg = _classify(out, f"/models/{arch}.gguf", f"local/{arch}") + assert arch in msg + assert "does not support" in msg.lower() + assert "diffusion" not in msg.lower() + assert "Images page" not in msg + + +class TestOllamaAndFallback: + _OLLAMA_GGUF = ( + f"/home/u/.ollama{__import__('os').sep}ollama_links" + f"{__import__('os').sep}m.gguf" + ) + + def test_ollama_compat_message_still_works(self): + out = "llama_model_load: error loading model: key not found" + msg = _classify(out, self._OLLAMA_GGUF, "ollama/llama3") + assert "Ollama" in msg + + def test_ollama_unknown_arch_keeps_ollama_guidance(self): + # Ollama + non-diffusion unknown arch keeps the Ollama hint, not the + # generic llama.cpp "unsupported" message. + out = "error loading model: unknown model architecture: 'some_new_llm'" + msg = _classify(out, self._OLLAMA_GGUF, "ollama/some-new") + assert "Ollama" in msg + assert "directly through Ollama" in msg + assert "does not support" not in msg.lower() + + def test_ollama_diffusion_arch_still_routes_to_images(self): + # Diffusion routing wins over the Ollama hint. + out = "error loading model: unknown model architecture: 'flux'" + msg = _classify(out, self._OLLAMA_GGUF, "ollama/flux") + assert "diffusion" in msg.lower() + assert "Images page" in msg + + def test_generic_oom_keeps_memory_message(self): + msg = _classify(_OOM_OUT, "/models/big.gguf", "local/big") + assert "enough memory" in msg.lower() + assert "diffusion" not in msg.lower() + + def test_empty_output_is_safe(self): + msg = _classify("", None, None) + assert "llama-server failed to start" in msg diff --git a/studio/backend/tests/test_mcp_stdio_improvements.py b/studio/backend/tests/test_mcp_stdio_improvements.py new file mode 100644 index 0000000000..e980e6a057 --- /dev/null +++ b/studio/backend/tests/test_mcp_stdio_improvements.py @@ -0,0 +1,236 @@ +"""Tests for the proposed PR #5863 improvements. + +Covers: _client() self-gating + keep_alive, OAuth normalised off for stdio +(create + update), env/header dropped on a transport-type switch, and the +backend rejecting a command whose first token is a URL scheme. + +Run from studio/backend: python -m pytest tests/test_mcp_stdio_improvements.py -q +""" + +import asyncio + +import pytest +from fastapi import HTTPException + +from core.inference import mcp_client +from storage import mcp_servers_db + + +def _reset_db(tmp_path, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(mcp_servers_db, "_schema_ready", False) + + +def _enable(monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1") + + +def _disable(monkeypatch): + monkeypatch.delenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", raising = False) + + +# ── P1: _client() self-gates the stdio sink ───────────────────────── + + +def test_client_refuses_stdio_when_disabled(monkeypatch): + _disable(monkeypatch) + with pytest.raises(PermissionError): + mcp_client._client("npx -y server /tmp", None) + + +def test_client_builds_stdio_when_enabled_without_spawning(monkeypatch): + _enable(monkeypatch) + # Constructing the Client must not spawn the subprocess (spawn happens on + # __aenter__); we only assert it builds. + client = mcp_client._client("npx -y server /tmp", {"K": "v"}) + assert client is not None + + +def test_client_http_unaffected_by_gate(monkeypatch): + _disable(monkeypatch) + assert mcp_client._client("https://example.com/mcp", None) is not None + + +# ── P3: OAuth normalised off for stdio (create + update) ──────────── + + +def test_create_forces_oauth_off_for_stdio(tmp_path, monkeypatch): + import routes.mcp_servers as routes_mcp + from models.mcp_servers import McpServerCreate + + _reset_db(tmp_path, monkeypatch) + _enable(monkeypatch) + resp = asyncio.run( + routes_mcp.create_mcp_server( + McpServerCreate( + display_name = "FS", url = "npx -y server /tmp", use_oauth = True + ), + current_subject = "u", + ) + ) + assert resp.use_oauth is False + assert mcp_servers_db.get_server(resp.id)["use_oauth"] == 0 + + +def test_create_keeps_oauth_for_http(tmp_path, monkeypatch): + import routes.mcp_servers as routes_mcp + from models.mcp_servers import McpServerCreate + + _reset_db(tmp_path, monkeypatch) + _enable(monkeypatch) + resp = asyncio.run( + routes_mcp.create_mcp_server( + McpServerCreate(display_name = "GH", url = "https://gh/mcp", use_oauth = True), + current_subject = "u", + ) + ) + assert resp.use_oauth is True + + +def test_update_url_to_stdio_clears_oauth(tmp_path, monkeypatch): + import routes.mcp_servers as routes_mcp + from models.mcp_servers import McpServerUpdate + + _reset_db(tmp_path, monkeypatch) + _enable(monkeypatch) + monkeypatch.setattr(mcp_client, "_oauth_token_store", None) + monkeypatch.setattr( + routes_mcp, "clear_oauth_tokens_async", lambda *a, **k: asyncio.sleep(0) + ) + mcp_servers_db.create_server( + id = "s1", display_name = "A", url = "https://a/mcp", use_oauth = True + ) + resp = asyncio.run( + routes_mcp.update_mcp_server( + "s1", McpServerUpdate(url = "npx -y server /tmp"), current_subject = "u" + ) + ) + assert resp.use_oauth is False + + +# ── P4: env/headers dropped on a transport-type switch ────────────── + + +def test_switch_stdio_to_http_drops_env(tmp_path, monkeypatch): + import routes.mcp_servers as routes_mcp + from models.mcp_servers import McpServerUpdate + + _reset_db(tmp_path, monkeypatch) + _enable(monkeypatch) + mcp_servers_db.create_server( + id = "s1", + display_name = "A", + url = "npx server", + headers_json = '{"API_KEY": "secret"}', + ) + resp = asyncio.run( + routes_mcp.update_mcp_server( + "s1", McpServerUpdate(url = "https://remote/mcp"), current_subject = "u" + ) + ) + # the stdio env must NOT survive as HTTP headers on the remote endpoint + assert resp.headers == {} + assert mcp_servers_db.get_server("s1")["headers_json"] is None + + +def test_switch_keeps_explicitly_supplied_headers(tmp_path, monkeypatch): + import routes.mcp_servers as routes_mcp + from models.mcp_servers import McpServerUpdate + + _reset_db(tmp_path, monkeypatch) + _enable(monkeypatch) + mcp_servers_db.create_server( + id = "s1", + display_name = "A", + url = "npx server", + headers_json = '{"API_KEY": "secret"}', + ) + resp = asyncio.run( + routes_mcp.update_mcp_server( + "s1", + McpServerUpdate( + url = "https://remote/mcp", headers = {"Authorization": "Bearer new"} + ), + current_subject = "u", + ) + ) + assert resp.headers == {"Authorization": "Bearer new"} + + +def test_same_transport_edit_keeps_headers(tmp_path, monkeypatch): + import routes.mcp_servers as routes_mcp + from models.mcp_servers import McpServerUpdate + + _reset_db(tmp_path, monkeypatch) + _enable(monkeypatch) + mcp_servers_db.create_server( + id = "s1", + display_name = "A", + url = "npx server", + headers_json = '{"API_KEY": "secret"}', + ) + # editing only the display name (still stdio) must not wipe env vars + resp = asyncio.run( + routes_mcp.update_mcp_server( + "s1", McpServerUpdate(display_name = "B"), current_subject = "u" + ) + ) + assert resp.headers == {"API_KEY": "secret"} + + +# ── P5: reject a command whose first token is a URL scheme ─────────── + + +def test_validate_url_rejects_url_scheme_command_when_enabled(monkeypatch): + from routes.mcp_servers import _validate_url + + _enable(monkeypatch) + for bad in ["ftp://host/x", "file:///etc/passwd", "ws://h/y"]: + with pytest.raises(HTTPException) as exc: + _validate_url(bad) + assert exc.value.status_code == 400 + + +def test_validate_url_allows_url_in_argument(monkeypatch): + from routes.mcp_servers import _validate_url + + _enable(monkeypatch) + # :// inside an ARGUMENT (not the first token) is still a valid command + assert _validate_url("npx server --url https://x/mcp") == ( + "npx server --url https://x/mcp" + ) + + +# ── P6: Data Recipe stdio path obeys the same host gate ───────────── +# build_mcp_providers needs the data_designer plugin, which is only installed in +# the Studio test job; skip there rather than fail the core matrix. + +_STDIO_RECIPE = { + "mcp_providers": [ + { + "provider_type": "stdio", + "name": "fs", + "command": "npx", + "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"], + "env": {}, + } + ] +} + + +def test_data_recipe_skips_stdio_when_disabled(monkeypatch): + pytest.importorskip("data_designer") + _disable(monkeypatch) + from core.data_recipe.service import build_mcp_providers + + # gate off -> the stdio provider is dropped (no subprocess can be spawned) + assert build_mcp_providers(_STDIO_RECIPE) == [] + + +def test_data_recipe_builds_stdio_when_enabled(monkeypatch): + pytest.importorskip("data_designer") + _enable(monkeypatch) + from core.data_recipe.service import build_mcp_providers + + built = build_mcp_providers(_STDIO_RECIPE) + assert len(built) == 1 # constructed (not spawned) only when enabled diff --git a/studio/backend/tests/test_mcp_stdio_pr5863.py b/studio/backend/tests/test_mcp_stdio_pr5863.py new file mode 100644 index 0000000000..d9a9510378 --- /dev/null +++ b/studio/backend/tests/test_mcp_stdio_pr5863.py @@ -0,0 +1,367 @@ +"""Verification tests for PR #5863 (stdio MCP server support). + +Covers the pure helpers (is_stdio / parse_stdio_command / stdio_mcp_enabled / +probe_timeout), the route-level _validate_url gate, and - most importantly - +that the UNSLOTH_STUDIO_ALLOW_STDIO_MCP gate blocks the stdio transport at all +five enforcement points (create, update, test, refresh, discovery, execute) +when disabled, and reaches it when enabled. The transport (_client) is stubbed +so no real subprocess is spawned; a recorder asserts whether it was reached. + +Run from studio/backend: python -m pytest tests/test_mcp_stdio_pr5863.py -q +""" + +import sys + +import pytest +from fastapi import HTTPException + +from core.inference import mcp_client +from storage import mcp_servers_db + + +def _reset_db(tmp_path, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.setattr(mcp_servers_db, "_schema_ready", False) + + +def _enable(monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1") + + +def _disable(monkeypatch): + monkeypatch.delenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", raising = False) + + +# ── transport stub + recorder ─────────────────────────────────────── + + +class _FakeTool: + def __init__(self, name): + self._name = name + + def model_dump(self, exclude_none = True): + return {"name": self._name, "description": f"{self._name} tool"} + + +class _Block: + def __init__(self, text): + self.type = "text" + self.text = text + + +class _FakeResult: + is_error = False + + def __init__(self, text): + self.content = [_Block(text)] + + +class _RecordingClient: + """Stands in for fastmcp.Client; records that the transport was opened.""" + + def __init__(self, url, headers, use_oauth, recorder): + recorder.append({"url": url, "headers": headers, "use_oauth": use_oauth}) + + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def list_tools(self): + return [_FakeTool("list_directory"), _FakeTool("write_file")] + + async def call_tool(self, name, args): + return _FakeResult(f"called {name}") + + +@pytest.fixture +def transport(monkeypatch): + """Patch mcp_client._client with a recorder. Returns the recorder list; + empty == the stdio transport was never reached.""" + recorder = [] + monkeypatch.setattr( + mcp_client, + "_client", + lambda url, headers, use_oauth = False: _RecordingClient( + url, headers, use_oauth, recorder + ), + ) + return recorder + + +# ── 1. is_stdio ───────────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "addr", + [ + "http://localhost:8000/mcp", + "https://example.com/mcp", + " https://example.com/mcp ", + "HTTPS://EXAMPLE.COM/mcp", + ], +) +def test_is_stdio_false_for_http(addr): + assert mcp_client.is_stdio(addr) is False + + +@pytest.mark.parametrize( + "addr", + [ + "npx -y @modelcontextprotocol/server-filesystem /tmp", + "python -m some.module", + "uvx some-server --flag", + "/usr/local/bin/my-server", + ], +) +def test_is_stdio_true_for_commands(addr): + assert mcp_client.is_stdio(addr) is True + + +# ── 2. parse_stdio_command ────────────────────────────────────────── + + +def test_parse_basic_argv(): + assert mcp_client.parse_stdio_command( + "npx -y @modelcontextprotocol/server-filesystem /tmp" + ) == ["npx", "-y", "@modelcontextprotocol/server-filesystem", "/tmp"] + + +def test_parse_keeps_url_argument_as_one_command(): + # gemini "high": a :// inside an ARGUMENT must not break the command. + assert mcp_client.parse_stdio_command( + "npx server --endpoint https://example.com/mcp" + ) == ["npx", "server", "--endpoint", "https://example.com/mcp"] + + +def test_parse_quoted_arg(): + assert mcp_client.parse_stdio_command('python -m mod --name "a b"') == [ + "python", + "-m", + "mod", + "--name", + "a b", + ] + + +def test_parse_empty_returns_empty_list(): + assert mcp_client.parse_stdio_command(" ") == [] + + +def test_parse_unclosed_quote_raises_valueerror(): + with pytest.raises(ValueError): + mcp_client.parse_stdio_command('npx "unclosed') + + +def test_parse_windows_strips_wrapping_quotes(monkeypatch): + # gemini "medium": posix=False keeps backslash paths but also the wrapping + # quotes; the PR strips a matched pair so argv[0] reaches the OS clean. + monkeypatch.setattr(sys, "platform", "win32") + parts = mcp_client.parse_stdio_command( + r'"C:\Program Files\node\node.exe" server.js' + ) + assert parts[0] == r"C:\Program Files\node\node.exe" + assert parts[1] == "server.js" + + +# ── 3. stdio_mcp_enabled ──────────────────────────────────────────── + + +@pytest.mark.parametrize("val", ["0", "false", "true", "", " 1 ", "yes", "2"]) +def test_stdio_disabled_for_non_exact_one(monkeypatch, val): + monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", val) + assert mcp_client.stdio_mcp_enabled() is False + + +def test_stdio_enabled_only_for_exact_one(monkeypatch): + _disable(monkeypatch) + assert mcp_client.stdio_mcp_enabled() is False + monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1") + assert mcp_client.stdio_mcp_enabled() is True + + +# ── 4. probe_timeout ──────────────────────────────────────────────── + + +def test_probe_timeout_matrix(): + assert mcp_client.probe_timeout("https://x/mcp", False) == 8.0 + assert mcp_client.probe_timeout("https://x/mcp", True) == 305.0 + assert mcp_client.probe_timeout("npx server", False) == 60.0 + # oauth wins regardless of address kind (documented behaviour) + assert mcp_client.probe_timeout("npx server", True) == 305.0 + + +# ── 5. _validate_url gate ─────────────────────────────────────────── + + +def test_validate_url_gate_off_rejects_stdio(monkeypatch): + _disable(monkeypatch) + from routes.mcp_servers import _validate_url + + assert _validate_url("https://example.com/mcp") == "https://example.com/mcp" + for bad in ["npx server", "python -m mod", "ftp://host"]: + with pytest.raises(HTTPException) as exc: + _validate_url(bad) + assert exc.value.status_code == 400 + + +def test_validate_url_gate_on_accepts_stdio(monkeypatch): + _enable(monkeypatch) + from routes.mcp_servers import _validate_url + + assert _validate_url("npx -y server /tmp") == "npx -y server /tmp" + # http still works when stdio is on + assert _validate_url("https://x/mcp") == "https://x/mcp" + # url-bearing argument accepted as a command + assert _validate_url("npx server --url https://x/mcp") == ( + "npx server --url https://x/mcp" + ) + # empty / unparseable still rejected + for bad in [" ", '"unclosed']: + with pytest.raises(HTTPException) as exc: + _validate_url(bad) + assert exc.value.status_code == 400 + + +# ── 6. gate enforcement at every spawn path (mocked transport) ────── + + +def test_create_route_gate(tmp_path, monkeypatch, transport): + import asyncio + + from models.mcp_servers import McpServerCreate + import routes.mcp_servers as routes_mcp + + _reset_db(tmp_path, monkeypatch) + payload = McpServerCreate(display_name = "FS", url = "npx -y server /tmp") + + _disable(monkeypatch) + with pytest.raises(HTTPException) as exc: + asyncio.run(routes_mcp.create_mcp_server(payload, current_subject = "u")) + assert exc.value.status_code == 400 + + _enable(monkeypatch) + resp = asyncio.run(routes_mcp.create_mcp_server(payload, current_subject = "u")) + assert resp.url == "npx -y server /tmp" + + +def test_update_http_to_stdio_blocked_when_off(tmp_path, monkeypatch): + import asyncio + + from models.mcp_servers import McpServerUpdate + import routes.mcp_servers as routes_mcp + + _reset_db(tmp_path, monkeypatch) + _disable(monkeypatch) + mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://a/mcp") + # editing url -> stdio command must 400 (http->stdio edit bypass closed) + with pytest.raises(HTTPException) as exc: + asyncio.run( + routes_mcp.update_mcp_server( + "s1", McpServerUpdate(url = "npx server"), current_subject = "u" + ) + ) + assert exc.value.status_code == 400 + + +def test_test_route_gate(tmp_path, monkeypatch, transport): + import asyncio + + from models.mcp_servers import McpServerTestRequest + import routes.mcp_servers as routes_mcp + + _reset_db(tmp_path, monkeypatch) + req = McpServerTestRequest(url = "npx -y server /tmp") + + _disable(monkeypatch) + with pytest.raises(HTTPException) as exc: + asyncio.run(routes_mcp.test_mcp_server(req, current_subject = "u")) + assert exc.value.status_code == 400 + assert transport == [] # transport never opened + + _enable(monkeypatch) + res = asyncio.run(routes_mcp.test_mcp_server(req, current_subject = "u")) + assert res.ok and res.tool_count == 2 + assert len(transport) == 1 + + +def test_refresh_route_gate(tmp_path, monkeypatch, transport): + import asyncio + + import routes.mcp_servers as routes_mcp + + _reset_db(tmp_path, monkeypatch) + # a stdio row as if carried over from a desktop DB + mcp_servers_db.create_server(id = "stdio1", display_name = "FS", url = "npx server") + + _disable(monkeypatch) + with pytest.raises(HTTPException) as exc: + asyncio.run(routes_mcp.refresh_mcp_server_tools("stdio1", current_subject = "u")) + assert exc.value.status_code == 400 + assert transport == [] + + _enable(monkeypatch) + res = asyncio.run( + routes_mcp.refresh_mcp_server_tools("stdio1", current_subject = "u") + ) + assert res.ok and res.tool_count == 2 + assert len(transport) == 1 + + +def test_discovery_gate(tmp_path, monkeypatch, transport): + import asyncio + + from core.inference.tools import get_enabled_mcp_tools + + _reset_db(tmp_path, monkeypatch) + mcp_servers_db.create_server( + id = "stdio1", display_name = "FS", url = "npx server", is_enabled = True + ) + + _disable(monkeypatch) + assert asyncio.run(get_enabled_mcp_tools()) == [] + assert transport == [] # filtered out before any probe + + _enable(monkeypatch) + specs = asyncio.run(get_enabled_mcp_tools()) + assert len(specs) == 2 + assert len(transport) == 1 + + +def test_execute_gate(tmp_path, monkeypatch, transport): + from core.inference.tools import execute_tool + + _reset_db(tmp_path, monkeypatch) + mcp_servers_db.create_server( + id = "stdio1", display_name = "FS", url = "npx server", is_enabled = True + ) + + _disable(monkeypatch) + out = execute_tool("mcp__stdio1__list_directory", {"path": "/tmp"}) + assert "disabled on this host" in out + assert transport == [] + + _enable(monkeypatch) + out = execute_tool("mcp__stdio1__list_directory", {"path": "/tmp"}) + assert out == "called list_directory" + assert len(transport) == 1 + + +# ── 7. env vars ride headers_json as the subprocess env ───────────── + + +def test_stdio_env_passed_through(tmp_path, monkeypatch, transport): + from core.inference.tools import execute_tool + + _reset_db(tmp_path, monkeypatch) + _enable(monkeypatch) + mcp_servers_db.create_server( + id = "stdio1", + display_name = "FS", + url = "npx server", + headers_json = '{"API_KEY": "sk-test"}', + is_enabled = True, + ) + execute_tool("mcp__stdio1__list_directory", {}) + assert transport[-1]["headers"] == {"API_KEY": "sk-test"} diff --git a/studio/frontend/src/components/assistant-ui/sources.tsx b/studio/frontend/src/components/assistant-ui/sources.tsx index 140b61f932..5593369bf9 100644 --- a/studio/frontend/src/components/assistant-ui/sources.tsx +++ b/studio/frontend/src/components/assistant-ui/sources.tsx @@ -263,20 +263,34 @@ const SourcesGroup: FC = () => { return (
- {/* Hidden measurement container — renders all badges to measure row positions */} + {/* Hidden measurement container. Renders all badges off-screen so we + can read each child's offsetTop and decide how many fit in two + rows. Wrapped in an absolute, h-0, overflow-hidden box so the + measurement pills do NOT contribute to the viewport's scrollable + overflow region. Without this clip, every hidden source row + adds ~30px to scrollHeight, producing a phantom empty scroll + area below the message: visible to users as unbounded blank + space below the assistant action bar. The inner div still + flex-wraps its children for measurement; offsetTop reads + correctly because the wrapper is positioned (absolute) and the + children's offsetTop is measured relative to it. */}
- {sources.map((source) => ( - - - - {source.title || extractDomain(source.url)} - - - ))} +
+ {sources.map((source) => ( + + + + {source.title || extractDomain(source.url)} + + + ))} +
{/* Visible container */} diff --git a/studio/frontend/src/features/chat/chat-mcp-servers-dialog.tsx b/studio/frontend/src/features/chat/chat-mcp-servers-dialog.tsx index 35b5aca64c..b377e7a88e 100644 --- a/studio/frontend/src/features/chat/chat-mcp-servers-dialog.tsx +++ b/studio/frontend/src/features/chat/chat-mcp-servers-dialog.tsx @@ -66,23 +66,40 @@ function headersToObject(rows: HeaderRow[]): Record | undefined return Object.keys(out).length > 0 ? out : undefined; } -function isValidUrl(url: string): boolean { - const trimmed = url.trim(); +// A non-HTTP address is a local stdio command. Case-insensitive to match the +// backend's is_stdio(), so all layers split http-vs-command identically. +function isHttpAddress(value: string): boolean { + const trimmed = value.trim().toLowerCase(); + return trimmed.startsWith("http://") || trimmed.startsWith("https://"); +} + +function isValidAddress(value: string): boolean { + const trimmed = value.trim(); if (!trimmed) return false; - try { - const parsed = new URL(trimmed); - return parsed.protocol === "http:" || parsed.protocol === "https:"; - } catch { - return false; + if (isHttpAddress(trimmed)) { + try { + const parsed = new URL(trimmed); + return parsed.protocol === "http:" || parsed.protocol === "https:"; + } catch { + return false; + } } + // Anything else is treated as a local command (stdio); the backend gates + // whether stdio servers are allowed on this host. Reject other URL schemes + // only when the command itself is a URL; "://" is fine inside an argument + // (e.g. a database connection string passed to the server). + return !trimmed.split(/\s+/)[0].includes("://"); } function HeadersEditor({ rows, onChange, + stdio, }: { rows: HeaderRow[]; onChange: (rows: HeaderRow[]) => void; + // stdio servers reuse this editor for environment variables instead of headers. + stdio: boolean; }) { const update = (id: string, patch: Partial) => onChange(rows.map((row) => (row.id === id ? { ...row, ...patch } : row))); @@ -91,19 +108,41 @@ function HeadersEditor({ const remove = (id: string) => onChange(rows.filter((row) => row.id !== id)); + const copy = stdio + ? { + label: "Environment variables", + add: "Add variable", + keyPlaceholder: "Variable name", + valuePlaceholder: "Variable value", + remove: "Remove variable", + } + : { + label: "Custom headers", + add: "Add header", + keyPlaceholder: "Header name", + valuePlaceholder: "Header value", + remove: "Remove header", + }; + return ( <>
- +
{rows.length === 0 ? (
- Optional. Add an Authorization header here for servers - that require auth. + {stdio ? ( + "Optional. Environment variables passed to the server process." + ) : ( + <> + Optional. Add an Authorization header here for servers + that require auth. + + )}
) : (
@@ -111,12 +150,12 @@ function HeadersEditor({
update(row.id, { key: e.target.value })} /> update(row.id, { value: e.target.value })} /> @@ -199,8 +238,8 @@ export function ChatMcpServersDialog({ async function testConnection() { const trimmedUrl = form.url.trim(); - if (!isValidUrl(trimmedUrl)) { - toast.error("Enter a valid http:// or https:// URL first"); + if (!isValidAddress(trimmedUrl)) { + toast.error("Enter an http(s):// URL or a local command first"); return; } setTesting(true); @@ -236,11 +275,11 @@ export function ChatMcpServersDialog({ return; } if (!trimmedUrl) { - toast.error("URL is required"); + toast.error("URL or command is required"); return; } - if (!isValidUrl(trimmedUrl)) { - toast.error("URL must start with http:// or https://"); + if (!isValidAddress(trimmedUrl)) { + toast.error("Enter an http(s):// URL or a local command"); return; } setSaving(true); @@ -331,6 +370,9 @@ export function ChatMcpServersDialog({ } const showForm = view.kind !== "list"; + // A local stdio command uses env vars, not headers or OAuth. + const addressIsCommand = + form.url.trim() !== "" && !isHttpAddress(form.url); return ( @@ -338,7 +380,7 @@ export function ChatMcpServersDialog({ MCP Servers - Register remote MCP servers. + Register remote (HTTP) or local (stdio command) MCP servers. @@ -356,40 +398,47 @@ export function ChatMcpServersDialog({ />
- + setForm((prev) => ({ ...prev, url: e.target.value })) } - placeholder="https://example.com/mcp" + placeholder="https://example.com/mcp or npx -y @modelcontextprotocol/server-filesystem /tmp" /> + + An http(s) URL for a remote server, or a local command to run an + stdio server (desktop app only). +
-
-
- - - For servers that require browser-based authentication - (GitHub, Linear, etc.). A browser window will open on first - connect. - + {!addressIsCommand && ( +
+
+ + + For servers that require browser-based authentication + (GitHub, Linear, etc.). A browser window will open on first + connect. + +
+ + setForm((prev) => ({ ...prev, useOauth })) + } + />
- - setForm((prev) => ({ ...prev, useOauth })) - } - /> -
+ )} setForm((prev) => ({ ...prev, headers }))} + stdio={addressIsCommand} />
@@ -415,7 +464,7 @@ export function ChatMcpServersDialog({
) : ( -
+