ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(studio): expose an opt-in MCP control plane (#7191)

Browse source 
* 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 <unslothshared@gmail.com>
This commit is contained in:
Ritwij Aryan Parmar 2026-07-17 19:50:39 -04:00 • committed by GitHub
commit a8ff8673df
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 600 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

34
studio/MCP.md Normal file
View file

@ -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.

17
studio/backend/main.py
View file

@ -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 = (

259
studio/backend/mcp_server.py Normal file
View file

@ -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
)

290
studio/backend/tests/test_mcp_server.py Normal file
View file

@ -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