From a8ff8673df640710e58fabfb226df565d9751afe Mon Sep 17 00:00:00 2001 From: Ritwij Aryan Parmar <88580521+RitwijParmar@users.noreply.github.com> Date: Fri, 17 Jul 2026 19:50:39 -0400 Subject: [PATCH] feat(studio): expose an opt-in MCP control plane (#7191) * feat(studio): expose opt-in MCP control plane * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Harden Studio MCP tools: byte-safe auth, page clamping, forward export/checkpoint fields Follow-up hardening on the opt-in MCP control plane. All changes are additive and backwards compatible. - BearerTokenMiddleware now compares the Authorization header on raw bytes. A non-ASCII bearer value previously reached str-based hmac.compare_digest, which raises TypeError and surfaced as a 500 instead of a clean 401. The constructor also rejects an empty or whitespace-only token so an empty token can never match an empty "Bearer " header. - MCP tools call the route functions directly, which skips FastAPI Query validation. list_training_runs and get_recipe_job_dataset now clamp limit and offset to the same bounds the HTTP routes enforce (a negative SQLite LIMIT otherwise means "no limit"). - export_gguf forwards hf_token (the backend rejects a Hub upload without it), accepts a list of quantization methods, and exposes imatrix / imatrix_path so the IQ low-bit quants are reachable. - load_checkpoint forwards hf_token and approved_remote_code_fingerprint so gated checkpoints and the remote-code approval retry work. Its docstring is corrected: the export backend coexists with training and inference rather than freeing GPU work. - start_training passes via_api_key=False explicitly instead of relying on the unfilled Depends default. Tests: add coverage for non-ASCII and empty-token auth, the correct-token pass through, non-http scope pass through, pagination clamping, and the forwarded export/checkpoint fields. * Harden Studio MCP: cap /mcp request bodies, reject unusable tokens, fix docs Follow-up hardening from a full review pass. All changes are additive and backwards compatible. - Add "/mcp" to _BODY_PROTECTED_PREFIXES so MaxBodyMiddleware enforces the same request-body cap it already applies to every other write endpoint (/api/train, /api/export, /api/data-recipe, ...). The MCP endpoint accepts authenticated POST tool-call bodies; without this an authenticated client could send an unbounded body. The middleware only buffers the request body (not the SSE response), so streaming is unaffected, and the 500MB default cap never affects a real JSON-RPC tool call (verified live). - Reject a non-ASCII UNSLOTH_STUDIO_MCP_TOKEN at construction. HTTP header values are ASCII, so a non-ASCII token cannot be sent by a standard client and would silently lock out the endpoint; fail fast instead. - MCP.md: document the canonical /mcp/ endpoint and note that /mcp redirects to it, so clients that do not follow redirected POSTs still connect. Tests: add non-ASCII token rejection coverage. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten MCP server comments --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen --- studio/MCP.md | 34 +++ studio/backend/main.py | 17 ++ studio/backend/mcp_server.py | 259 +++++++++++++++++++++ studio/backend/tests/test_mcp_server.py | 290 ++++++++++++++++++++++++ 4 files changed, 600 insertions(+) create mode 100644 studio/MCP.md create mode 100644 studio/backend/mcp_server.py create mode 100644 studio/backend/tests/test_mcp_server.py diff --git a/studio/MCP.md b/studio/MCP.md new file mode 100644 index 0000000000..91b39fcc69 --- /dev/null +++ b/studio/MCP.md @@ -0,0 +1,34 @@ +# Unsloth Studio MCP server + +Studio can expose a local MCP server so an MCP client can inspect models and +GPU state, validate recipes, start or stop training, inspect recipe output, and +export a loaded model. + +The server is disabled by default. Enable it for a local Studio process with: + +```bash +UNSLOTH_STUDIO_ENABLE_MCP=1 \ +UNSLOTH_STUDIO_MCP_TOKEN='use-a-local-secret' \ +unsloth studio +``` + +The endpoint is `http://127.0.0.1:8888/mcp/` when Studio uses its default port +(a request to `/mcp` redirects to the canonical `/mcp/`). Use the actual Studio +port when it is configured differently. + +The high-impact tools are: + +- `studio_status` and `list_local_models` for discovery +- `get_training_status`, `start_training`, `stop_training`, and `list_training_runs` +- `validate_recipe`, `get_recipe_job_status`, and `get_recipe_job_dataset` +- `load_checkpoint` and `export_gguf` + +`start_training` accepts the same fields as the Studio `TrainingStartRequest`. +The request is validated by the existing Pydantic model before a subprocess is +started. Export paths use the existing Studio validation as well. + +The endpoint always requires `UNSLOTH_STUDIO_MCP_TOKEN` and checks an exact +Bearer token for both HTTP and WebSocket connections. Keep it on localhost +unless the deployment has an authenticated reverse proxy. The MCP endpoint is +intentionally opt-in because tools can consume GPU memory, write model +artifacts, and stop active work. \ No newline at end of file diff --git a/studio/backend/main.py b/studio/backend/main.py index e64048dc00..6e16dc00ca 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -612,6 +612,22 @@ app = FastAPI( lifespan = lifespan, ) +# The MCP surface is opt-in because it can start GPU jobs and write model +# artifacts. Mount it only when explicitly enabled by the Studio process. +if os.environ.get("UNSLOTH_STUDIO_ENABLE_MCP") == "1": + from fastmcp.utilities.lifespan import combine_lifespans + + from mcp_server import BearerTokenMiddleware, create_studio_mcp + + _studio_mcp_app = create_studio_mcp().http_app(path = "/") + _studio_mcp_lifespan = _studio_mcp_app.lifespan + _mcp_token = os.environ.get("UNSLOTH_STUDIO_MCP_TOKEN") + if not _mcp_token: + raise RuntimeError("UNSLOTH_STUDIO_MCP_TOKEN is required when MCP is enabled") + _studio_mcp_app = BearerTokenMiddleware(_studio_mcp_app, _mcp_token) + app.router.lifespan_context = combine_lifespans(lifespan, _studio_mcp_lifespan) + app.mount("/mcp", _studio_mcp_app) + from loggers.config import LogConfig from loggers.handlers import LoggingMiddleware @@ -752,6 +768,7 @@ _BODY_PROTECTED_PREFIXES = ( "/api/settings", "/api/train", "/api/export", + "/mcp", ) _DATASET_UPLOAD_PASSTHROUGH_PREFIX = "/api/datasets/upload" _DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX = ( diff --git a/studio/backend/mcp_server.py b/studio/backend/mcp_server.py new file mode 100644 index 0000000000..f837f46425 --- /dev/null +++ b/studio/backend/mcp_server.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Curated MCP tools for driving an Unsloth Studio instance. + +The MCP surface deliberately wraps the existing Studio services instead of +duplicating training or export logic. It is opt-in because several tools can +start GPU work or write model artifacts. +""" + +from __future__ import annotations + +import hmac +from typing import Any + +from fastmcp import FastMCP + + +class BearerTokenMiddleware: + """Require an exact bearer token when Studio MCP is exposed remotely.""" + + def __init__(self, app: Any, token: str) -> None: + if not token or not token.strip(): + raise ValueError("Studio MCP bearer token must be a non-empty value") + if not token.isascii(): + # A non-ASCII token cannot be sent in an HTTP header; reject it here. + raise ValueError("Studio MCP bearer token must contain ASCII characters only") + self.app = app + # Compare on raw header bytes: str hmac.compare_digest raises on non-ASCII + # input, which would surface as a 500 instead of a clean 401. + self.expected = token.encode("utf-8") + + async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None: + scope_type = scope.get("type") + if scope_type not in ("http", "websocket"): + await self.app(scope, receive, send) + return + + headers = dict(scope.get("headers", [])) + raw_auth = headers.get(b"authorization", b"") + scheme, _, supplied = raw_auth.partition(b" ") + if scheme.lower() != b"bearer" or not hmac.compare_digest(supplied, self.expected): + await _send_unauthorized(send, scope_type) + return + + await self.app(scope, receive, send) + + +async def _send_unauthorized(send: Any, scope_type: str) -> None: + if scope_type == "websocket": + await send({"type": "websocket.close", "code": 4401}) + return + + await send( + { + "type": "http.response.start", + "status": 401, + "headers": [(b"content-type", b"application/json"), (b"www-authenticate", b"Bearer")], + } + ) + await send( + { + "type": "http.response.body", + "body": b'{"detail":"MCP bearer token required"}', + } + ) + + +def _dump(value: Any) -> Any: + """Convert Pydantic responses to plain JSON values for MCP clients.""" + if hasattr(value, "model_dump"): + return value.model_dump(mode = "json") + return value + + +def _clamp(value: int, low: int, high: int) -> int: + """Clamp an MCP-supplied integer into an inclusive range. + + MCP tools call the Studio route functions directly, which skips FastAPI's + Query(ge=, le=) validation, so we re-apply the same bounds here. + """ + return max(low, min(value, high)) + + +def create_studio_mcp() -> FastMCP: + """Create the Studio MCP server and register the high-value tools.""" + mcp = FastMCP( + "Unsloth Studio", + instructions = ( + "Use read tools to inspect the local Studio state before starting GPU work. " + "Training and export tools can consume substantial VRAM and write files. " + "Never expose tokens or local paths from tool results unless the user asks." + ), + ) + + @mcp.tool + async def studio_status() -> dict[str, Any]: + """Return the current training, export, inference, and GPU state.""" + from routes.export import get_export_status + from routes.inference import get_status as get_inference_status + from routes.training import get_training_status + + from utils.hardware import get_gpu_utilization + + training, export, inference = await _gather_status( + get_training_status(current_subject = "mcp"), + get_export_status(current_subject = "mcp"), + get_inference_status(current_subject = "mcp"), + ) + return { + "training": _dump(training), + "export": _dump(export), + "inference": _dump(inference), + "hardware": get_gpu_utilization(), + } + + @mcp.tool + async def list_local_models(models_dir: str = "./models") -> dict[str, Any]: + """List local and cached models available to Studio.""" + from routes.models import list_local_models as list_models + return _dump(await list_models(models_dir = models_dir, current_subject = "mcp")) + + @mcp.tool + async def get_training_status() -> dict[str, Any]: + """Read the active training job, phase, progress, and recent metrics.""" + from routes.training import get_training_status as get_status + return _dump(await get_status(current_subject = "mcp")) + + @mcp.tool + async def start_training(config: dict[str, Any]) -> dict[str, Any]: + """Start a validated Studio training job from a TrainingStartRequest-shaped object. + + The config is validated by the same Pydantic model used by the Studio UI. + Call get_training_status first and do not start work while another job runs. + """ + from models import TrainingStartRequest + from routes.training import start_training as start + + request = TrainingStartRequest.model_validate(config) + # Pass via_api_key explicitly (a direct call leaves it a Depends object). + # MCP drives Studio like the UI session, so it coexists and frees VRAM. + return _dump(await start(request, current_subject = "mcp", via_api_key = False)) + + @mcp.tool + async def stop_training(save: bool = True) -> dict[str, Any]: + """Ask the active training process to stop at its next safe checkpoint.""" + from routes.training import TrainingStopRequest, stop_training as stop + return _dump(await stop(TrainingStopRequest(save = save), current_subject = "mcp")) + + @mcp.tool + async def list_training_runs(limit: int = 50, offset: int = 0) -> dict[str, Any]: + """List completed and stopped training runs, newest first.""" + from routes.training_history import list_training_runs as list_runs + + # Clamp here (direct call skips Query bounds); a negative LIMIT = no limit. + limit = _clamp(limit, 1, 200) + offset = max(0, offset) + return _dump(await list_runs(limit = limit, offset = offset, current_subject = "mcp")) + + @mcp.tool + def validate_recipe(recipe: dict[str, Any]) -> dict[str, Any]: + """Validate a Data Recipe with the same validator used by Studio.""" + from models.data_recipe import RecipePayload + from routes.data_recipe.validate import validate + + return _dump(validate(RecipePayload(recipe = recipe))) + + @mcp.tool + def get_recipe_job_status(job_id: str) -> dict[str, Any]: + """Read the status of a Data Recipe job.""" + from routes.data_recipe.jobs import job_status + return _dump(job_status(job_id)) + + @mcp.tool + def get_recipe_job_dataset( + job_id: str, + limit: int = 20, + offset: int = 0, + ) -> dict[str, Any]: + """Read a bounded page of generated Data Recipe rows.""" + from routes.data_recipe.jobs import job_dataset + + # Clamp here (direct call skips FastAPI's Query bounds). + limit = _clamp(limit, 1, 500) + offset = max(0, offset) + return _dump(job_dataset(job_id, limit = limit, offset = offset)) + + @mcp.tool + async def load_checkpoint( + checkpoint_path: str, + max_seq_length: int = 2048, + load_in_4bit: bool = True, + trust_remote_code: bool = False, + approved_remote_code_fingerprint: str | None = None, + hf_token: str | None = None, + ) -> dict[str, Any]: + """Load a checkpoint into the export backend. + + Export runs in its own subprocess and coexists with training and + inference; it does not unload them, so a load can fail with a clear + out-of-memory error if the GPU is already full. Pass hf_token to load a + gated checkpoint, and approved_remote_code_fingerprint to retry a + trust_remote_code load that was blocked pending review. + """ + from models import LoadCheckpointRequest + from routes.export import load_checkpoint as load + + request = LoadCheckpointRequest( + checkpoint_path = checkpoint_path, + max_seq_length = max_seq_length, + load_in_4bit = load_in_4bit, + trust_remote_code = trust_remote_code, + approved_remote_code_fingerprint = approved_remote_code_fingerprint, + hf_token = hf_token, + ) + return _dump(await load(request, current_subject = "mcp")) + + @mcp.tool + async def export_gguf( + save_directory: str, + quantization_method: str | list[str] = "Q4_K_M", + push_to_hub: bool = False, + repo_id: str | None = None, + hf_token: str | None = None, + imatrix: bool = False, + imatrix_path: str | None = None, + ) -> dict[str, Any]: + """Export the loaded model to GGUF using Studio's existing path validation. + + quantization_method may be a single method or a list to produce several + GGUFs from one load. Pass hf_token when push_to_hub is set (the backend + rejects a Hub upload without it). Set imatrix (or imatrix_path) for the + IQ low-bit quants that require an importance matrix. + """ + from models import ExportGGUFRequest + from routes.export import export_gguf as export + + request = ExportGGUFRequest( + save_directory = save_directory, + quantization_method = quantization_method, + push_to_hub = push_to_hub, + repo_id = repo_id, + hf_token = hf_token, + imatrix = imatrix, + imatrix_path = imatrix_path, + ) + return _dump(await export(request, current_subject = "mcp")) + + return mcp + + +async def _gather_status(*coroutines: Any) -> tuple[Any, ...]: + """Gather independent status calls without letting one optional backend fail all state.""" + import asyncio + + results = await asyncio.gather(*coroutines, return_exceptions = True) + return tuple( + {"error": str(result)} if isinstance(result, Exception) else result for result in results + ) diff --git a/studio/backend/tests/test_mcp_server.py b/studio/backend/tests/test_mcp_server.py new file mode 100644 index 0000000000..71792605ae --- /dev/null +++ b/studio/backend/tests/test_mcp_server.py @@ -0,0 +1,290 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import asyncio +import sys +import types + +import pytest + +from mcp_server import BearerTokenMiddleware, _clamp, _dump, create_studio_mcp + + +def _get_tool(name): + tools = asyncio.run(create_studio_mcp().list_tools()) + return {tool.name: tool for tool in tools}[name] + + +def test_studio_mcp_registers_control_plane_tools(): + tools = asyncio.run(create_studio_mcp().list_tools()) + + assert {tool.name for tool in tools} == { + "studio_status", + "list_local_models", + "get_training_status", + "start_training", + "stop_training", + "list_training_runs", + "validate_recipe", + "get_recipe_job_status", + "get_recipe_job_dataset", + "load_checkpoint", + "export_gguf", + } + + +def test_dump_serializes_pydantic_values(): + class Response: + def model_dump(self, *, mode): + assert mode == "json" + return {"ok": True} + + assert _dump(Response()) == {"ok": True} + assert _dump({"already": "json"}) == {"already": "json"} + + +def test_bearer_token_middleware_rejects_wrong_token(): + events = [] + + async def app(scope, receive, send): + events.append("app") + + async def send(message): + events.append(message) + + middleware = BearerTokenMiddleware(app, "secret") + asyncio.run( + middleware( + {"type": "http", "headers": [(b"authorization", b"Bearer wrong")]}, + None, + send, + ) + ) + + assert events[0]["status"] == 401 + assert "app" not in events + + +def test_bearer_token_middleware_closes_unauthorized_websocket(): + events = [] + + async def app(scope, receive, send): + events.append("app") + + async def send(message): + events.append(message) + + middleware = BearerTokenMiddleware(app, "secret") + asyncio.run( + middleware( + {"type": "websocket", "headers": []}, + None, + send, + ) + ) + + assert events == [{"type": "websocket.close", "code": 4401}] + + +def test_bearer_token_middleware_rejects_non_ascii_authorization(): + # A non-ASCII bearer value must produce a clean 401, not a 500. Comparing on + # bytes avoids the str hmac.compare_digest TypeError on non-ASCII input. + events = [] + + async def app(scope, receive, send): + events.append("app") + + async def send(message): + events.append(message) + + middleware = BearerTokenMiddleware(app, "secret") + asyncio.run( + middleware( + {"type": "http", "headers": [(b"authorization", b"Bearer \xff\xff")]}, + None, + send, + ) + ) + + assert events[0]["status"] == 401 + assert "app" not in events + + +def test_bearer_token_middleware_accepts_correct_token(): + events = [] + + async def app(scope, receive, send): + events.append("app") + + async def send(message): + events.append(message) + + middleware = BearerTokenMiddleware(app, "secret") + asyncio.run( + middleware( + {"type": "http", "headers": [(b"authorization", b"Bearer secret")]}, + None, + send, + ) + ) + + assert events == ["app"] + + +def test_bearer_token_middleware_requires_non_empty_token(): + async def app(scope, receive, send): + pass + + for bad in ("", " "): + with pytest.raises(ValueError): + BearerTokenMiddleware(app, bad) + + +def test_bearer_token_middleware_rejects_non_ascii_token(): + async def app(scope, receive, send): + pass + + # non-ASCII tokens cannot be transmitted in an HTTP header by a standard + # client, so they are rejected at construction instead of locking out. + for bad in ("töken", "\U0001f600"): + with pytest.raises(ValueError): + BearerTokenMiddleware(app, bad) + + +def test_bearer_token_middleware_passes_through_non_http_scopes(): + events = [] + + async def app(scope, receive, send): + events.append("app") + + async def send(message): + events.append(message) + + middleware = BearerTokenMiddleware(app, "secret") + asyncio.run(middleware({"type": "lifespan"}, None, send)) + + assert events == ["app"] + + +def test_clamp_restricts_to_inclusive_bounds(): + assert _clamp(5, 1, 200) == 5 + assert _clamp(-10, 1, 200) == 1 + assert _clamp(10_000, 1, 200) == 200 + assert _clamp(0, 1, 500) == 1 + assert _clamp(1_000, 1, 500) == 500 + + +def test_export_and_checkpoint_tools_expose_forwarded_fields(): + export_props = set(_get_tool("export_gguf").parameters["properties"]) + assert {"hf_token", "imatrix", "imatrix_path"} <= export_props + + checkpoint_props = set(_get_tool("load_checkpoint").parameters["properties"]) + assert {"hf_token", "approved_remote_code_fingerprint"} <= checkpoint_props + + +def _stub_module(monkeypatch, name, **attrs): + module = types.ModuleType(name) + for key, value in attrs.items(): + setattr(module, key, value) + if "." in name: + module.__path__ = [] # mark package-like so submodule imports resolve + monkeypatch.setitem(sys.modules, name, module) + return module + + +def test_export_gguf_forwards_hf_token_and_imatrix(monkeypatch): + captured = {} + + class FakeExportGGUFRequest: + def __init__(self, **kwargs): + captured.update(kwargs) + + async def fake_export(request, current_subject): + return {"current_subject": current_subject} + + _stub_module(monkeypatch, "models", ExportGGUFRequest = FakeExportGGUFRequest) + _stub_module(monkeypatch, "routes") + _stub_module(monkeypatch, "routes.export", export_gguf = fake_export) + + tool = _get_tool("export_gguf") + result = asyncio.run( + tool.fn( + save_directory = "/tmp/out", + quantization_method = ["Q4_K_M", "Q8_0"], + push_to_hub = True, + repo_id = "me/model", + hf_token = "hf_secret", + imatrix = True, + imatrix_path = "/tmp/imatrix.dat", + ) + ) + + assert captured["hf_token"] == "hf_secret" + assert captured["imatrix"] is True + assert captured["imatrix_path"] == "/tmp/imatrix.dat" + assert captured["quantization_method"] == ["Q4_K_M", "Q8_0"] + assert result["current_subject"] == "mcp" + + +def test_load_checkpoint_forwards_token_and_fingerprint(monkeypatch): + captured = {} + + class FakeLoadCheckpointRequest: + def __init__(self, **kwargs): + captured.update(kwargs) + + async def fake_load(request, current_subject): + return {"current_subject": current_subject} + + _stub_module(monkeypatch, "models", LoadCheckpointRequest = FakeLoadCheckpointRequest) + _stub_module(monkeypatch, "routes") + _stub_module(monkeypatch, "routes.export", load_checkpoint = fake_load) + + tool = _get_tool("load_checkpoint") + asyncio.run( + tool.fn( + checkpoint_path = "/tmp/ckpt", + approved_remote_code_fingerprint = "sha256:abc", + hf_token = "hf_secret", + ) + ) + + assert captured["hf_token"] == "hf_secret" + assert captured["approved_remote_code_fingerprint"] == "sha256:abc" + + +def test_list_training_runs_clamps_pagination(monkeypatch): + captured = {} + + async def fake_list_runs(limit, offset, current_subject): + captured["limit"] = limit + captured["offset"] = offset + return {"ok": True} + + _stub_module(monkeypatch, "routes") + _stub_module(monkeypatch, "routes.training_history", list_training_runs = fake_list_runs) + + tool = _get_tool("list_training_runs") + asyncio.run(tool.fn(limit = 10_000, offset = -5)) + + assert captured["limit"] == 200 + assert captured["offset"] == 0 + + +def test_get_recipe_job_dataset_clamps_pagination(monkeypatch): + captured = {} + + def fake_job_dataset(job_id, limit, offset): + captured["limit"] = limit + captured["offset"] = offset + return {"ok": True} + + _stub_module(monkeypatch, "routes") + _stub_module(monkeypatch, "routes.data_recipe") + _stub_module(monkeypatch, "routes.data_recipe.jobs", job_dataset = fake_job_dataset) + + tool = _get_tool("get_recipe_job_dataset") # this tool is synchronous + tool.fn(job_id = "job-1", limit = -1, offset = -9) + + assert captured["limit"] == 1 + assert captured["offset"] == 0