unsloth/studio/backend/routes/data_recipe/mcp.py
Daniel Han 4b6a733fc3
Studio: harden stdio MCP gating and fix transport edge cases (#5892)
* Studio: harden stdio MCP gating and fix transport edge cases

- Gate the Data Recipe stdio path behind UNSLOTH_STUDIO_ALLOW_STDIO_MCP so a hosted deployment cannot spawn local processes through recipes
- Enforce the gate inside _client() so the transport sink cannot spawn when disabled
- keep_alive=False so stdio probes/calls do not leave orphan subprocesses
- Force OAuth off for stdio servers on create and update
- Drop stored headers when a server switches transport type
- Reject a command whose first token is a URL scheme
- Add MCP gate and improvement tests

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* tests: skip Data Recipe stdio tests when data_designer is absent

The data_designer plugin is only installed in the Studio test job, so guard
the two build_mcp_providers tests with importorskip so the core matrix skips
them instead of failing on ModuleNotFoundError.

---------

Co-authored-by: Daniel Han <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-05-31 04:38:52 -07:00

92 lines
3 KiB
Python

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""MCP helper endpoints for data recipe."""
from __future__ import annotations
from collections import defaultdict
from fastapi import APIRouter
from core.data_recipe.service import build_mcp_providers
from models.data_recipe import (
McpToolsListRequest,
McpToolsListResponse,
McpToolsProviderResult,
)
router = APIRouter()
@router.post("/mcp/tools", response_model = McpToolsListResponse)
def list_mcp_tools(payload: McpToolsListRequest) -> McpToolsListResponse:
try:
from data_designer.engine.mcp import io as mcp_io
except ImportError as exc:
return McpToolsListResponse(
providers = [
McpToolsProviderResult(
name = "",
error = f"MCP dependencies unavailable: {exc}",
)
]
)
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(
McpToolsProviderResult(
name = provider_name,
error = "Unsupported MCP provider config.",
)
)
continue
provider = built[0]
try:
tools = mcp_io.list_tools(provider, timeout_sec = payload.timeout_sec)
tool_names = sorted(
{tool.name for tool in tools if getattr(tool, "name", "")}
)
for tool_name in tool_names:
tool_to_providers[tool_name].append(provider.name)
providers.append(
McpToolsProviderResult(
name = provider.name,
tools = tool_names,
)
)
except Exception as exc:
providers.append(
McpToolsProviderResult(
name = provider.name or provider_name,
error = str(exc).strip() or "Failed to load tools.",
)
)
duplicate_tools = {
tool_name: provider_names
for tool_name, provider_names in sorted(tool_to_providers.items())
if len(provider_names) > 1
}
return McpToolsListResponse(
providers = providers,
duplicate_tools = duplicate_tools,
)