ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
unsloth/studio/MCP.md
Ritwij Aryan Parmar a8ff8673df
 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 <unslothshared@gmail.com>
2026-07-17 16:50:39 -07:00

1.4 KiB
Raw Permalink Blame History

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:

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.