Skip to content

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

Merged
danielhanchen merged 6 commits into
unslothai:mainfrom
RitwijParmar:codex/studio-mcp-server
Jul 17, 2026
Merged

feat(studio): expose an opt-in MCP control plane#7191
danielhanchen merged 6 commits into
unslothai:mainfrom
RitwijParmar:codex/studio-mcp-server

Conversation

@RitwijParmar

Copy link
Copy Markdown
Contributor

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=1 and 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_TOKEN is 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_compile for the changed Python files
  • FastMCP smoke test covering all 11 registered tools
  • ruff check on changed files
  • focused bearer-token and tool-registration tests added in studio/backend/tests/test_mcp_server.py
  • git diff --check

@gemini-code-assist gemini-code-assist Bot left a comment

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.

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.

Comment on lines +26 to +38
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)

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

Comment thread studio/backend/mcp_server.py Outdated
Comment on lines +41 to +54
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"}',
}
)

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

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

Suggested change
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"}',
}
)

Comment on lines +84 to +88
training, export, inference = await _gather_status(
get_training_status(),
get_export_status(),
get_inference_status(),
)

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

Comment on lines +104 to +108
"""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

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

Comment on lines +120 to +121

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

return _dump(await stop(TrainingStopRequest(save = save)))

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

)
return _dump(await load(request))

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

)
return _dump(await export(request))

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

Comment thread studio/backend/main.py Outdated
Comment on lines +556 to +577
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

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.

medium

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.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 Codex Review

_bin = LlamaCppBackend._find_llama_server_binary()
_caps = LlamaCppBackend.probe_server_capabilities(_bin)
app.state.llama_cpp_capabilities = _caps
_freshness = check_prebuilt_freshness(_bin)
app.state.llama_cpp_freshness = _freshness

P2 Badge Keep llama.cpp probes off the startup path

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

Comment thread studio/backend/main.py Outdated
Comment on lines 1123 to 1124
if _is_same_origin_request(request):
content, nonce = _inject_bootstrap(content, app)

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P1 Badge 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 👍 / 👎.

Comment thread studio/backend/main.py Outdated

from mcp_server import BearerTokenMiddleware, create_studio_mcp

_studio_mcp_app = create_studio_mcp().http_app()

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 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 👍 / 👎.

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.

Satisfied on the current head: the app is mounted at /mcp via http_app(path="/"), matching MCP.md.

Comment thread studio/backend/main.py Outdated
Comment on lines +481 to +482
if _mcp_token := os.environ.get("UNSLOTH_STUDIO_MCP_TOKEN"):
_studio_mcp_app = BearerTokenMiddleware(_studio_mcp_app, _mcp_token)

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

P1 Badge 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 👍 / 👎.

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.

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.

Comment thread studio/backend/main.py
Comment on lines 815 to 816
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"])

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 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 👍 / 👎.

Comment thread studio/backend/main.py Outdated
Comment on lines +124 to +125
# rocm bnb DLL present, or HIP_PATH/ROCM_PATH set (DLL unparsable -> "72")
if _found_rocm_bnb or _hip_env:

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 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 👍 / 👎.

Comment thread studio/backend/main.py Outdated
Comment on lines 981 to 984
return {
"gpu": get_gpu_summary(),
"versions": get_package_versions(),
# Export capability + torch-aware reason; the Export UI grays out with the message.
**export_capability(),
}

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 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 👍 / 👎.

Comment thread studio/backend/main.py Outdated
Comment on lines +943 to +947
visibility_info = get_backend_visible_gpu_info()
gpu_info = {
"available": visibility_info["available"],
"devices": visibility_info["devices"],
}

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 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 👍 / 👎.

Comment thread studio/backend/main.py Outdated
app.state.llama_cpp_capabilities = None
app.state.llama_cpp_freshness = None
_start_llama_cpp_probes_if_enabled(app)
_start_helper_precache_if_enabled()

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 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 👍 / 👎.

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.

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.

Comment thread studio/backend/main.py Outdated
clear_unsloth_compiled_cache,
_hw_module,
)
await asyncio.to_thread(terminate_hub_downloads)

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 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 👍 / 👎.

Comment thread studio/backend/mcp_server.py Outdated
Comment on lines +170 to +174
request = LoadCheckpointRequest(
checkpoint_path=checkpoint_path,
max_seq_length=max_seq_length,
load_in_4bit=load_in_4bit,
trust_remote_code=trust_remote_code,

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 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 👍 / 👎.

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: load_checkpoint now accepts and forwards approved_remote_code_fingerprint and hf_token into LoadCheckpointRequest.

@RitwijParmar
RitwijParmar force-pushed the codex/studio-mcp-server branch 3 times, most recently from d147c5c to 7abeff6 Compare July 16, 2026 21:55
@RitwijParmar
RitwijParmar force-pushed the codex/studio-mcp-server branch from 823d490 to fe64345 Compare July 16, 2026 21:57

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

💡 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".

Comment on lines +192 to +196
request = ExportGGUFRequest(
save_directory = save_directory,
quantization_method = quantization_method,
push_to_hub = push_to_hub,
repo_id = repo_id,

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.

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

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.

Comment on lines +182 to +187
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]:

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.

…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.
danielhanchen added a commit to danielhanchen/unsloth-staging-2 that referenced this pull request Jul 17, 2026
danielhanchen and others added 2 commits July 17, 2026 12:50
…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.
danielhanchen added a commit to danielhanchen/unsloth-staging-2 that referenced this pull request Jul 17, 2026
@danielhanchen

Copy link
Copy Markdown
Member

@codex review

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. Another round soon, please!

Reviewed commit: b83044414f

ℹ️ 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".

@danielhanchen

Copy link
Copy Markdown
Member

@codex review

1 similar comment
@danielhanchen

Copy link
Copy Markdown
Member

@codex review

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. Keep it up!

Reviewed commit: 53e9e97f34

ℹ️ 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".

@danielhanchen

Copy link
Copy Markdown
Member

Thanks!

@danielhanchen
danielhanchen merged commit a8ff867 into unslothai:main Jul 17, 2026
46 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature] Standalone MCP server for Unsloth Studio

2 participants