Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

Appearance settings
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 34 additions & 0 deletions 34 studio/MCP.md
Original file line number Diff line number Diff line change
@@ -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 changes: 17 additions & 0 deletions 17 studio/backend/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -612,6 +612,22 @@ async def lifespan(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

Expand Down Expand Up @@ -752,6 +768,7 @@ async def send_wrapper(message):
"/api/settings",
"/api/train",
"/api/export",
"/mcp",
)
_DATASET_UPLOAD_PASSTHROUGH_PREFIX = "/api/datasets/upload"
_DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX = (
Expand Down
259 changes: 259 additions & 0 deletions 259 studio/backend/mcp_server.py
Original file line number Diff line number Diff line change
@@ -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)
Comment on lines +33 to +46

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

security-critical critical

The middleware currently only enforces the bearer token check for "http" scope types. If the MCP server is accessed via WebSockets (scope type "websocket"), the authentication check is completely bypassed, allowing unauthorized clients to access the control plane. Enforce the token check for both "http" and "websocket" scopes, and pass the scope type to _send_unauthorized to handle rejection properly.

Suggested change
async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None:
if scope.get("type") != "http":
await self.app(scope, receive, send)
return
headers = dict(scope.get("headers", []))
raw_auth = headers.get(b"authorization", b"").decode("latin-1")
scheme, _, supplied = raw_auth.partition(" ")
if scheme.lower() != "bearer" or not hmac.compare_digest(supplied, self.token):
await _send_unauthorized(send)
return
await self.app(scope, receive, send)
async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None:
if scope.get("type") not in ("http", "websocket"):
await self.app(scope, receive, send)
return
headers = dict(scope.get("headers", []))
raw_auth = headers.get(b"authorization", b"").decode("latin-1")
scheme, _, supplied = raw_auth.partition(" ")
if scheme.lower() != "bearer" or not hmac.compare_digest(supplied, self.token):
await _send_unauthorized(send, scope.get("type", "http"))
return
await self.app(scope, receive, send)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Already handled on the current head: the middleware gates both http and websocket scopes and closes an unauthorized websocket with code 4401 (mcp_server.py).



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"),
)
Comment on lines +105 to +109

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Pass current_subject="mcp" to the status route functions. Otherwise, they will default to the Depends(get_current_subject) dependency injection object, which can cause unexpected behavior or errors downstream.

Suggested change
training, export, inference = await _gather_status(
get_training_status(),
get_export_status(),
get_inference_status(),
)
training, export, inference = await _gather_status(
get_training_status(current_subject="mcp"),
get_export_status(current_subject="mcp"),
get_inference_status(current_subject="mcp"),
)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Already handled on the current head: the three status calls pass 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
Comment on lines +125 to +129

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Pass current_subject="mcp" to get_status. Otherwise, it will default to the Depends object.

Suggested change
"""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())
@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"))

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
Comment on lines +143 to +144

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Pass current_subject="mcp" and via_api_key=False to start_training. Otherwise, they will default to Depends objects, which will be passed to the training backend and cause serialization or runtime errors.

Suggested change
@mcp.tool
request = TrainingStartRequest.model_validate(config)
return _dump(await start(request, current_subject="mcp", via_api_key=False))

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 2ac311d: start_training now passes current_subject="mcp" and via_api_key=False explicitly (mcp_server.py).

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]:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Pass current_subject="mcp" to stop_training. Otherwise, it will default to the Depends object.

Suggested change
async def list_training_runs(limit: int = 50, offset: int = 0) -> dict[str, Any]:
return _dump(await stop(TrainingStopRequest(save=save), current_subject="mcp"))

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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Clamp MCP recipe dataset pages

When an MCP caller supplies a very large limit, this direct call bypasses the FastAPI Query(..., le=500) validation on the underlying recipe dataset endpoint and forwards the raw value into the job manager. For large generated datasets that can materialize far more than the intended bounded preview into memory and into the MCP response, so the wrapper should enforce the same 1–500 range before calling the route.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 2ac311d: get_recipe_job_dataset now clamps limit to [1,500] and floors offset at 0, matching the route's Query bounds.


@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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Pass current_subject="mcp" to load_checkpoint. Otherwise, it will default to the Depends object.

Suggested change
@mcp.tool
return _dump(await load(request, current_subject="mcp"))

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]:
Comment on lines +219 to +227

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Allow imatrix for IQ GGUF exports

When an MCP client requests an IQ low-bit quant such as iq2_xxs or iq4_xs, this wrapper has no way to set imatrix/imatrix_path, so ExportGGUFRequest always uses the defaults and the route forwards imatrix_file=None. The existing UI explicitly forces imatrix for those quant levels because llama.cpp rejects them without it (studio/frontend/src/features/export/export-page.tsx lines 257-261), so those supported GGUF exports fail through MCP.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 2ac311d: export_gguf now exposes and forwards imatrix and imatrix_path, so the IQ low-bit quants are reachable.

"""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,
Comment on lines +238 to +242

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Forward the GGUF Hugging Face token

When an MCP client sets push_to_hub=True, this wrapper has no parameter for hf_token and constructs ExportGGUFRequest without it. The normal /api/export/export/gguf path forwards request.hf_token, and the export backend rejects Hub uploads unless both repo_id and hf_token are present, so GGUF Hub uploads through MCP always fail even though the Studio route supports them.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 2ac311d: export_gguf now forwards hf_token into ExportGGUFRequest so push_to_hub can authenticate.

hf_token = hf_token,
imatrix = imatrix,
imatrix_path = imatrix_path,
)
return _dump(await export(request, current_subject = "mcp"))

return mcp

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Pass current_subject="mcp" to export_gguf. Otherwise, it will default to the Depends object.

Suggested change
return mcp
return _dump(await export(request, current_subject="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
)
Loading
Loading
Morty Proxy This is a proxified and sanitized view of the page, visit original site.