feat(studio): expose an opt-in MCP control plane#7191
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces an opt-in Model Context Protocol (MCP) server for Unsloth Studio, allowing clients to inspect GPU state, manage training runs, and export models. It also refactors the backend startup/shutdown lifespan and middleware. The review feedback highlights several critical issues: the MCP bearer token middleware must be updated to authenticate WebSocket connections to prevent bypasses, and several MCP tools must explicitly pass current_subject="mcp" to avoid passing unresolved FastAPI Depends objects downstream. Additionally, it is recommended to revert SecurityHeadersMiddleware to a pure ASGI implementation to prevent performance and cancellation issues with streaming endpoints.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| 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) |
There was a problem hiding this comment.
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.
| 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) |
There was a problem hiding this comment.
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) -> None: | ||
| 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"}', | ||
| } | ||
| ) |
There was a problem hiding this comment.
Update _send_unauthorized to accept the scope_type parameter and handle WebSocket closure appropriately (by sending a websocket.close message instead of http.response.start which would crash the ASGI server).
| async def _send_unauthorized(send: Any) -> None: | |
| 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"}', | |
| } | |
| ) | |
| 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"}', | |
| } | |
| ) |
| training, export, inference = await _gather_status( | ||
| get_training_status(), | ||
| get_export_status(), | ||
| get_inference_status(), | ||
| ) |
There was a problem hiding this comment.
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.
| 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"), | |
| ) |
There was a problem hiding this comment.
Already handled on the current head: the three status calls pass current_subject="mcp".
| """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 |
There was a problem hiding this comment.
Pass current_subject="mcp" to get_status. Otherwise, it will default to the Depends object.
| """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")) |
|
|
||
| @mcp.tool |
There was a problem hiding this comment.
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.
| @mcp.tool | |
| request = TrainingStartRequest.model_validate(config) | |
| return _dump(await start(request, current_subject="mcp", via_api_key=False)) |
There was a problem hiding this comment.
Fixed in 2ac311d: start_training now passes current_subject="mcp" and via_api_key=False explicitly (mcp_server.py).
| return _dump(await stop(TrainingStopRequest(save = save))) | ||
|
|
||
| @mcp.tool | ||
| async def list_training_runs(limit: int = 50, offset: int = 0) -> dict[str, Any]: |
There was a problem hiding this comment.
| ) | ||
| return _dump(await load(request)) | ||
|
|
||
| @mcp.tool |
| ) | ||
| return _dump(await export(request)) | ||
|
|
||
| return mcp |
| class SecurityHeadersMiddleware(BaseHTTPMiddleware): | ||
| """Set baseline security headers; splice per-response inline-script nonces into CSP.""" | ||
|
|
||
| async def dispatch(self, request: _StarletteRequest, call_next): | ||
| response = await call_next(request) | ||
| # Strip the internal nonce hand-off header so it never reaches the client | ||
| nonce = response.headers.get(_CSP_SCRIPT_NONCE_HEADER) | ||
| if nonce is not None: | ||
| del response.headers[_CSP_SCRIPT_NONCE_HEADER] | ||
| response.headers.setdefault("Content-Security-Policy", _build_csp(nonce)) | ||
| # Omit X-Frame-Options in Colab — CSP frame-ancestors handles it, and | ||
| # DENY would block serve_kernel_port_as_iframe regardless of CSP. | ||
| if not _IS_COLAB and request.url.path != _ARTIFACT_PREVIEW_FRAME_PATH: | ||
| response.headers.setdefault("X-Frame-Options", "DENY") | ||
| response.headers.setdefault("X-Content-Type-Options", "nosniff") | ||
| response.headers.setdefault("Referrer-Policy", "no-referrer") | ||
| response.headers.setdefault( | ||
| "Permissions-Policy", | ||
| "camera=(), microphone=(self), geolocation=()", | ||
| ) | ||
| response.headers["server"] = "unsloth-studio" | ||
| return response |
There was a problem hiding this comment.
Changing SecurityHeadersMiddleware from a pure ASGI middleware to BaseHTTPMiddleware re-introduces the issues with streaming responses being wrapped in AnyIO streams. As noted in the original docstring, pure ASGI is preferred for streaming endpoints (like training progress, export logs, and chat completions) to ensure proper cancellation propagation and avoid AnyIO overhead. Consider reverting to the pure ASGI implementation.
There was a problem hiding this comment.
💡 Codex Review
unsloth/studio/backend/main.py
Lines 391 to 395 in 81d1295
These probes are now run inline during FastAPI startup, including check_prebuilt_freshness(), which may perform a blocking GitHub freshness lookup on a cold or expired cache and previously was skipped entirely when UNSLOTH_DISABLE_UPDATE_CHECK=1. On offline or slow-network machines this can delay Application startup complete for the duration of the network/probe timeout even though the result is only advisory.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if _is_same_origin_request(request): | ||
| content, nonce = _inject_bootstrap(content, app) |
There was a problem hiding this comment.
Keep bootstrap credentials off public HTML
When Studio is bound to a LAN/public address or served through a proxy, a normal top-level GET usually has no Origin, and _is_same_origin_request() treats that as trusted; this now injects window.__UNSLOTH_BOOTSTRAP__ into the page for any remote visitor while the default admin still requires a password change. The previous loopback/proxy checks and suppress_bootstrap_injection guard prevented handing the bootstrap password to non-local clients, so this reopens the default-admin credential leak on exposed launches.
Useful? React with 👍 / 👎.
|
|
||
| from mcp_server import BearerTokenMiddleware, create_studio_mcp | ||
|
|
||
| _studio_mcp_app = create_studio_mcp().http_app() |
There was a problem hiding this comment.
Mount the MCP app at the documented endpoint
With FastMCP, http_app() already serves the streamable HTTP endpoint at /mcp by default; mounting that sub-app at /mcp makes the actual URL /mcp/mcp/ (the FastMCP mounting docs show an app created with path='/mcp' and mounted under /analytics is reached at /analytics/mcp/: https://gofastmcp.com/integrations/fastapi). In the documented setup clients configured for http://127.0.0.1:8888/mcp will miss the server, so pass path='/' to http_app() or mount the default app at root.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Satisfied on the current head: the app is mounted at /mcp via http_app(path="/"), matching MCP.md.
| if _mcp_token := os.environ.get("UNSLOTH_STUDIO_MCP_TOKEN"): | ||
| _studio_mcp_app = BearerTokenMiddleware(_studio_mcp_app, _mcp_token) |
There was a problem hiding this comment.
Require authentication before exposing MCP controls
If an operator enables UNSLOTH_STUDIO_ENABLE_MCP=1 on a Studio process that is reachable beyond loopback but forgets to set UNSLOTH_STUDIO_MCP_TOKEN, this branch mounts the control plane with no bearer middleware at all. Those tools can start/stop training and export artifacts, bypassing the auth dependencies used by the normal /api/* routes, so a network client can control GPU work in that deployment; fail closed for non-local exposure or require a token whenever MCP is enabled remotely.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Satisfied on the current head: BearerTokenMiddleware wraps the MCP app and startup raises RuntimeError when UNSLOTH_STUDIO_MCP_TOKEN is unset, so the endpoint is never exposed without a token.
| app.include_router(inference_router, prefix = "/v1", tags = ["openai-compat"]) | ||
| app.include_router(preview_router, prefix = "/p", tags = ["preview"]) | ||
| app.include_router(providers_router, prefix = "/api/providers", tags = ["providers"]) |
There was a problem hiding this comment.
Re-register the public preview router
The /p/... preview route is no longer included even though routes/preview.py still defines the signed public preview API and the history UI still copies /p/${preview_ref}?k=... links. Those links now fall through to the SPA catch-all or 404 instead of serving the preview page/OpenAI-compatible endpoints, so users who enable public preview sharing get unusable share links.
Useful? React with 👍 / 👎.
| # rocm bnb DLL present, or HIP_PATH/ROCM_PATH set (DLL unparsable -> "72") | ||
| if _found_rocm_bnb or _hip_env: |
There was a problem hiding this comment.
Avoid forcing ROCm bitsandbytes without a ROCm wheel
On Windows machines that have HIP_PATH/ROCM_PATH from the AMD HIP SDK but are using a CUDA/CPU bitsandbytes wheel, _hip_env now sets BNB_ROCM_VERSION even when no libbitsandbytes_rocm*.dll was found. That was previously gated on the actual ROCm DLL because the environment override makes bitsandbytes look for a ROCm backend and fail at import on non-ROCm wheels, so CUDA/CPU hosts with HIP tooling installed can lose startup/training imports.
Useful? React with 👍 / 👎.
| return { | ||
| "gpu": get_gpu_summary(), | ||
| "versions": get_package_versions(), | ||
| # Export capability + torch-aware reason; the Export UI grays out with the message. | ||
| **export_capability(), | ||
| } |
There was a problem hiding this comment.
Preserve detailed hardware response fields
The frontend still calls /api/system/hardware?include_details=true and expects the detailed response to include gpus, llama_cpp, and export_supported/unsupported reason fields. This endpoint now ignores include_details and returns only the summary GPU/version block, so the About page loses llama.cpp and multi-GPU details and the Export UI no longer pre-disables export with the hardware-specific unsupported reason on CPU or missing-MLX hosts.
Useful? React with 👍 / 👎.
| visibility_info = get_backend_visible_gpu_info() | ||
| gpu_info = { | ||
| "available": visibility_info["available"], | ||
| "devices": visibility_info["devices"], | ||
| } |
There was a problem hiding this comment.
Preserve live VRAM metrics in /api/system
The response now copies only get_backend_visible_gpu_info(), which reports visible devices and total memory but not the live vram_used_gb, vram_free_gb, or utilization fields that /api/system consumers still sum for the floating monitor and Resources tab. On GPU hosts those live monitors will show VRAM as unused/unknown even while training or inference is active; keep merging the utilization probe into these devices as before.
Useful? React with 👍 / 👎.
| app.state.llama_cpp_capabilities = None | ||
| app.state.llama_cpp_freshness = None | ||
| _start_llama_cpp_probes_if_enabled(app) | ||
| _start_helper_precache_if_enabled() |
There was a problem hiding this comment.
Reconcile orphaned RAG ingestion jobs on startup
After a crash during knowledge-base ingestion, storage.rag_db.reconcile_orphaned_ingestion_jobs() is what marks unfinished jobs/documents failed or completed and removes partial chunks. This startup path now goes from training-run cleanup straight to helper precache, so stale running ingestion rows can remain stuck as processing after restart and block users from cleanly retrying that document.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
This targets an earlier, force-pushed-away revision. The current main.py change is a single +16/-0 additive hunk (a conditional MCP mount); reconcile_orphaned_ingestion_jobs and the other referenced functions are unchanged from main. The same applies to the other main.py comments in this review.
| clear_unsloth_compiled_cache, | ||
| _hw_module, | ||
| ) | ||
| await asyncio.to_thread(terminate_hub_downloads) |
There was a problem hiding this comment.
Keep shutdown cleanup resilient to dead executors
During abrupt shutdowns, such as a Windows console close or uvicorn teardown after the default executor has already shut down, asyncio.to_thread(terminate_hub_downloads) can raise before DEVICE is reset or the compiled cache is cleared. The removed run_lifespan_shutdown() helper caught executor submission/body failures, ran inline when needed, and continued with later cleanup, so this can reintroduce noisy shutdown failures and skipped cleanup.
Useful? React with 👍 / 👎.
| request = LoadCheckpointRequest( | ||
| checkpoint_path=checkpoint_path, | ||
| max_seq_length=max_seq_length, | ||
| load_in_4bit=load_in_4bit, | ||
| trust_remote_code=trust_remote_code, |
There was a problem hiding this comment.
Forward approval and HF token fields when loading checkpoints
The MCP load_checkpoint wrapper drops LoadCheckpointRequest.approved_remote_code_fingerprint and hf_token. For gated/private checkpoints or LoRA base models the export route cannot scan/load with the user's token, and for approvable trust_remote_code models an MCP client has no way to retry with the fingerprint returned by the security scan, so those checkpoints are unloadable via MCP even though the normal Studio route supports them.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 2ac311d: load_checkpoint now accepts and forwards approved_remote_code_fingerprint and hf_token into LoadCheckpointRequest.
d147c5c to
7abeff6
Compare
823d490 to
fe64345
Compare
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a5a2c2206b
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| request = ExportGGUFRequest( | ||
| save_directory = save_directory, | ||
| quantization_method = quantization_method, | ||
| push_to_hub = push_to_hub, | ||
| repo_id = repo_id, |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
Fixed in 2ac311d: export_gguf now forwards hf_token into ExportGGUFRequest so push_to_hub can authenticate.
| ) -> dict[str, Any]: | ||
| """Read a bounded page of generated Data Recipe rows.""" | ||
| from routes.data_recipe.jobs import job_dataset | ||
| return _dump(job_dataset(job_id, limit = limit, offset = offset)) |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
Fixed in 2ac311d: get_recipe_job_dataset now clamps limit to [1,500] and floors offset at 0, matching the route's Query bounds.
| async def export_gguf( | ||
| save_directory: str, | ||
| quantization_method: str = "Q4_K_M", | ||
| push_to_hub: bool = False, | ||
| repo_id: str | None = None, | ||
| ) -> dict[str, Any]: |
There was a problem hiding this comment.
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 👍 / 👎.
There was a problem hiding this comment.
Fixed in 2ac311d: export_gguf now exposes and forwards imatrix and imatrix_path, so the IQ low-bit quants are reachable.
…t/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.
…ix 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.
for more information, see https://pre-commit.ci
|
@codex review |
|
Codex Review: Didn't find any major issues. Another round soon, please! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
@codex review |
1 similar comment
|
@codex review |
|
Codex Review: Didn't find any major issues. Keep it up! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
Thanks! |
Summary
Closes #6941.
Unsloth Studio already has the training, recipe, model, and export services. This adds a standalone MCP surface that lets an MCP client inspect and operate those services through the same validated route models instead of duplicating Studio logic.
The endpoint is opt-in with
UNSLOTH_STUDIO_ENABLE_MCP=1and is mounted at/mcp. The write-capable surface includes training start/stop, checkpoint loading, GGUF export, and recipe validation/status inspection. Read tools cover Studio status, local models, training history, and generated recipe rows.When
UNSLOTH_STUDIO_MCP_TOKENis set the endpoint requires an exact bearer token. Without that setting the endpoint remains intended for a local Studio process only.Verification
python3 -m py_compilefor the changed Python filesruff checkon changed filesstudio/backend/tests/test_mcp_server.pygit diff --check