From ff00fdd155910801f35ec5a6b616cff8cdd726e8 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Sun, 31 May 2026 05:54:46 -0300 Subject: [PATCH] Studio: add stdio MCP server support (#5863) * Studio: add stdio MCP server support * Fix stdio command validation and Windows quoting --- studio/backend/core/inference/mcp_client.py | 64 +++++++++ studio/backend/core/inference/tools.py | 13 +- studio/backend/main.py | 5 + studio/backend/routes/mcp_servers.py | 36 +++-- .../features/chat/chat-mcp-servers-dialog.tsx | 129 ++++++++++++------ 5 files changed, 192 insertions(+), 55 deletions(-) diff --git a/studio/backend/core/inference/mcp_client.py b/studio/backend/core/inference/mcp_client.py index a5e614899d..1a38bc8c49 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,19 @@ 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): + from fastmcp.client.transports import StdioTransport + + parts = parse_stdio_command(url) + if not parts: + raise ValueError(f"Empty stdio command: {url!r}") + # stdio env vars ride the (HTTP-only) headers field. The MCP SDK merges + # them over its default safe env (PATH etc.), so pass them through as-is. + return Client( + StdioTransport(command = parts[0], args = parts[1:], env = headers or None) + ) + 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/mcp_servers.py b/studio/backend/routes/mcp_servers.py index a7501d1691..04b776776a 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,22 @@ 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") + return trimmed parsed = urlparse(trimmed) if parsed.scheme not in ("http", "https"): raise HTTPException( @@ -180,15 +190,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 +226,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/frontend/src/features/chat/chat-mcp-servers-dialog.tsx b/studio/frontend/src/features/chat/chat-mcp-servers-dialog.tsx index 35b5aca64c..dfb50b5bc2 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} />