Studio: add stdio MCP server support (#5863)
* Studio: add stdio MCP server support * Fix stdio command validation and Windows quoting
This commit is contained in:
oobabooga
GitHub
parent
f213663d5b
commit
ff00fdd155
5 changed files with 192 additions and 55 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -66,23 +66,40 @@ function headersToObject(rows: HeaderRow[]): Record<string, string> | 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<HeaderRow>) =>
|
||||
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 (
|
||||
<>
|
||||
<div className="flex items-center justify-between">
|
||||
<Label className="text-sm">Custom headers</Label>
|
||||
<Label className="text-sm">{copy.label}</Label>
|
||||
<Button type="button" variant="ghost" size="sm" onClick={add}>
|
||||
<HugeiconsIcon icon={PlusSignIcon} size={14} />
|
||||
Add header
|
||||
{copy.add}
|
||||
</Button>
|
||||
</div>
|
||||
{rows.length === 0 ? (
|
||||
<div className="text-xs text-muted-foreground">
|
||||
Optional. Add an <code>Authorization</code> header here for servers
|
||||
that require auth.
|
||||
{stdio ? (
|
||||
"Optional. Environment variables passed to the server process."
|
||||
) : (
|
||||
<>
|
||||
Optional. Add an <code>Authorization</code> header here for servers
|
||||
that require auth.
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
) : (
|
||||
<div className="flex flex-col gap-2">
|
||||
|
|
@ -111,12 +150,12 @@ function HeadersEditor({
|
|||
<div key={row.id} className="flex items-center gap-2">
|
||||
<Input
|
||||
value={row.key}
|
||||
placeholder="Header name"
|
||||
placeholder={copy.keyPlaceholder}
|
||||
onChange={(e) => update(row.id, { key: e.target.value })}
|
||||
/>
|
||||
<Input
|
||||
value={row.value}
|
||||
placeholder="Header value"
|
||||
placeholder={copy.valuePlaceholder}
|
||||
onChange={(e) => update(row.id, { value: e.target.value })}
|
||||
/>
|
||||
<Button
|
||||
|
|
@ -124,7 +163,7 @@ function HeadersEditor({
|
|||
variant="ghost"
|
||||
size="icon"
|
||||
onClick={() => remove(row.id)}
|
||||
aria-label="Remove header"
|
||||
aria-label={copy.remove}
|
||||
>
|
||||
<HugeiconsIcon icon={Delete02Icon} size={14} />
|
||||
</Button>
|
||||
|
|
@ -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 (
|
||||
<Dialog open={open} onOpenChange={onOpenChange}>
|
||||
|
|
@ -338,7 +380,7 @@ export function ChatMcpServersDialog({
|
|||
<DialogHeader>
|
||||
<DialogTitle>MCP Servers</DialogTitle>
|
||||
<DialogDescription>
|
||||
Register remote MCP servers.
|
||||
Register remote (HTTP) or local (stdio command) MCP servers.
|
||||
</DialogDescription>
|
||||
</DialogHeader>
|
||||
|
||||
|
|
@ -356,40 +398,47 @@ export function ChatMcpServersDialog({
|
|||
/>
|
||||
</div>
|
||||
<div className="grid gap-2">
|
||||
<Label htmlFor="mcp-url">URL</Label>
|
||||
<Label htmlFor="mcp-url">URL or command</Label>
|
||||
<Input
|
||||
id="mcp-url"
|
||||
value={form.url}
|
||||
onChange={(e) =>
|
||||
setForm((prev) => ({ ...prev, url: e.target.value }))
|
||||
}
|
||||
placeholder="https://example.com/mcp"
|
||||
placeholder="https://example.com/mcp or npx -y @modelcontextprotocol/server-filesystem /tmp"
|
||||
/>
|
||||
<span className="text-xs text-muted-foreground">
|
||||
An http(s) URL for a remote server, or a local command to run an
|
||||
stdio server (desktop app only).
|
||||
</span>
|
||||
</div>
|
||||
|
||||
<div className="flex items-start justify-between gap-3">
|
||||
<div className="flex flex-col gap-0.5">
|
||||
<Label className="text-sm" htmlFor="mcp-oauth">
|
||||
Use OAuth sign-in
|
||||
</Label>
|
||||
<span className="text-xs text-muted-foreground">
|
||||
For servers that require browser-based authentication
|
||||
(GitHub, Linear, etc.). A browser window will open on first
|
||||
connect.
|
||||
</span>
|
||||
{!addressIsCommand && (
|
||||
<div className="flex items-start justify-between gap-3">
|
||||
<div className="flex flex-col gap-0.5">
|
||||
<Label className="text-sm" htmlFor="mcp-oauth">
|
||||
Use OAuth sign-in
|
||||
</Label>
|
||||
<span className="text-xs text-muted-foreground">
|
||||
For servers that require browser-based authentication
|
||||
(GitHub, Linear, etc.). A browser window will open on first
|
||||
connect.
|
||||
</span>
|
||||
</div>
|
||||
<Switch
|
||||
id="mcp-oauth"
|
||||
checked={form.useOauth}
|
||||
onCheckedChange={(useOauth) =>
|
||||
setForm((prev) => ({ ...prev, useOauth }))
|
||||
}
|
||||
/>
|
||||
</div>
|
||||
<Switch
|
||||
id="mcp-oauth"
|
||||
checked={form.useOauth}
|
||||
onCheckedChange={(useOauth) =>
|
||||
setForm((prev) => ({ ...prev, useOauth }))
|
||||
}
|
||||
/>
|
||||
</div>
|
||||
)}
|
||||
|
||||
<HeadersEditor
|
||||
rows={form.headers}
|
||||
onChange={(headers) => setForm((prev) => ({ ...prev, headers }))}
|
||||
stdio={addressIsCommand}
|
||||
/>
|
||||
|
||||
<div className="flex items-center justify-between gap-2 pt-2">
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue