mirror of
https://github.com/unslothai/unsloth.git
synced 2026-10-19 17:06:43 +00:00
43ceef8c04
## Summary
Agents that drive Unsloth Studio over MCP couldn't tell one failure from
another, saw cancelled work reported as failed, got hints naming
parameters the MCP tools don't have, and sat through long calls with no
progress. MCP errors now start with a machine-readable tag, cancelled
jobs read as cancelled, long calls report a heartbeat, and `chat`,
`embed` and `cancel` take the options agents were missing.
## Changes
### Error tags (`studio_mcp/errors.py`)
A failure an agent can branch on now starts with `[<code> k=v]`. Plain
refusals stay untagged.
- `[gpu_busy retry_after=12] GPU busy: ... Retry after 12 s.` for all
three gpu_busy shapes.
- `[multi_model_always kind=chat] ... Call unload_model(kind='llm') to
free the GPU; retrying won't help.` Detected from the
`X-Unsloth-Refusal: multi-model-always` header, the `multi_model_always`
code, or the 409 copy. No Retry-After.
- `[active_generations running=N] ...` with a hint to wait or cancel
your own chat by `cancel_id`. The route's `force_cancel_active` advice
is stripped, since the MCP tools don't expose it.
- `[memory_estimate] ... Pass allow_oversized=true to try anyway.` for
the 400 memory refusal.
- `[not_loaded kind=image]` (or chat, video, audio, stt) for "No model
loaded" answers, with the kind taken from the message or the route path.
- A route's `detail.code` or /v1 `error.code` becomes the tag verbatim.
The /v1 code used to be appended as a `(code)` suffix.
- Deferred (padded 200) errors now read headers the route copied into
`_deferred_error`, so the multi-model refusal is tagged there too.
- llama-server startup failures keep the error lines (up to 8) instead
of the full startup log, so the cause fits under the 2000-character cap.
### Jobs and status
- `get_job` on a cancelled video returns `status: "cancelled"` and
`error_code: "cancelled"` instead of `failed`.
- The export job listing returns `format`, `phase`, `output` and `error`
for each job, newest first (`ExportJobSummary`).
- `VideoJobRef` adds `progress_percent`, the same number `JobStatus`
uses.
- `studio_status` training: `loss` is `None` at step 0 instead of the
0.0 placeholder, and `elapsed_seconds` is reported. `/api/train/status`
details now carry `eta_seconds` and `elapsed_seconds`.
- A stop that lands during the final save no longer reads "Training
completed!". It reads "Training cancelled." or "Training stopped."
(`core/training/training.py`).
### Tools
- `chat` takes `enable_thinking`, `reasoning_effort` and
`response_format`. Any `finish_reason: "length"` now adds a note: "The
reply was cut off at max_tokens" when there is text, the existing
reasoning note when there isn't.
- `chat`, `generate_audio` and `transcribe` report a progress heartbeat
every 2 s (`studio_mcp/progress.py`). The tool's own progress, such as
`transcribe`'s download fraction, goes into the heartbeat message so MCP
progress only increases.
- `embed(query, top_k)` ranks `texts` by cosine similarity to `query`
and returns the `top_k` best (default 10) as `matches` with index, score
and text, without vectors. `top_k` without `query` is refused, and
`query` allows at most 2047 texts.
- `cancel(kind="video", id=...)` stops only that video. `POST
/api/inference/video/generate/cancel` accepts an optional `{"video_id":
"..."}` body and answers 409 "That video is not rendering." when the id
is not the running render, and 422 for an empty `video_id`. With no body
it behaves as before.
- `cancel(kind="image" | "diffusion_training")` now refuses an `id`
instead of ignoring it.
### Docs and copy
- MCP `INSTRUCTIONS` are rewritten one sentence per line: the one-GPU
rule, the heartbeat, polling intervals, the training flow, cancel by id
and the error tags.
- `studio/MCP.md` documents the Settings > API "Allow agent connections"
switch, sets `tool_timeout_sec = 300` for Codex, and adds a "Long calls"
paragraph (`MCP_TOOL_TIMEOUT=300000` for Claude Code). The `chat`,
`embed` and `cancel` rows are updated.
- The agent access copy in Settings > API
(`settings.apiKeys.mcp.description`, `enableDescription`,
`exportKeyHint`) says "Unsloth Studio" and "Unsloth Studio API key" in
all 14 locales. "API key" matches `UNSLOTH_API_KEY` and the page's
existing coding-agents hint; the token list above it still says "Access
tokens", which this PR leaves alone. Values only, no key changes.
MCP clients that parse error text will see the tag prefix and the /v1
code moved from the end to the tag. `EmbedResult.embeddings` is now
optional, but it is still returned whenever `query` is not set.
## Testing
Run on this head (`studio/backend`, Unsloth Studio venv, CPU):
- `pytest tests/test_mcp_*.py tests/test_video_cancel_by_id.py
tests/test_video_routes.py tests/test_training_status_eta_and_stop.py
tests/test_training_status_terminal.py
tests/test_training_stop_watchdog.py tests/test_job_accounts.py
tests/test_openai_videos_route.py`: 1823 passed, 4 failed. The same 4
fail on main e795149d99 in this environment: 3 in `test_mcp_images.py`
(transformers vs huggingface-hub 1.33 import) and
`test_video_download_plan_judges_a_quantized_reference_pick_per_partition`
(PyAV not installed).
- `tests/test_training_*.py`: the only failures are the same
huggingface-hub import errors, also on main.
- The new tests fail against main's sources: 38 tag/hint cases in
`test_mcp_errors.py`, `test_mcp_tool_contract.py` and
`test_mcp_tools_loading.py`, plus the training eta/elapsed and
stop-during-save cases. The negative controls pass on both.
- `ruff check` (0.6.9) is clean, `scripts/run_ruff_format.py` was
applied, and `npm run i18n:check:strict` passes.
New test files:
- `tests/test_mcp_progress.py`: heartbeat, the progress relay, strictly
increasing progress, ctx-free schemas.
- `tests/test_training_status_eta_and_stop.py`: eta/elapsed in
`/api/train/status`, the stop-during-save and drafter-phase messages,
and a no-stop run that still reads completed.
- `tests/test_video_cancel_by_id.py`: a wrong id gets 409 and the render
keeps going, the matching id cancels, no body behaves as before, bad
bodies and an empty id get 422.
UI: before/after screenshots of Settings > API > Agent access (English
and German) are in the evidence comment below.
## Notes for related work
- `studio_mcp/errors.py` keeps the `multi_model_always` hint on one
constant (`MULTI_MODEL_HINT`) and has an empty `CODE_HINTS` map, so
later work can add `unload_model(kind="all")` or audio-code hints
without touching the detection.
- The INSTRUCTIONS text marks the load, one-GPU, audio, polling,
training and galleries lines with comments, so a later PR can add its
sentence after the matching line.
<!-- ui-evidence:start -->
### UI before / after
Settings > API > Agent access (MCP) with agent connections on: the
section description, the switch description and the UNSLOTH_API_KEY
setup hint now say "Unsloth API key" instead of "access token". BEFORE
is the merge base, AFTER is this PR, two separate installs.
<!-- ui-evidence:end -->
---------
Co-authored-by: Daniel Han <23090290+danielhanchen@users.noreply.github.com>
Co-authored-by: danielhanchen <danielhanchen@gmail.com>
71 lines
3.1 KiB
Python
71 lines
3.1 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Unsloth Studio's MCP server at /mcp: curated tools that reach Unsloth Studio through its own routes as the agent's API key. Off until the owner turns it on, because several tools start GPU work or write files."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from fastmcp import FastMCP
|
|
|
|
from studio_mcp import prompts
|
|
from studio_mcp.tools import (
|
|
audio,
|
|
cancel,
|
|
data,
|
|
export,
|
|
images,
|
|
jobs,
|
|
library,
|
|
models,
|
|
status,
|
|
text,
|
|
training,
|
|
video,
|
|
)
|
|
|
|
INSTRUCTIONS = " ".join(
|
|
(
|
|
"Unsloth Studio runs local AI models for chat, embed, images, audio, transcribe, video, datasets, fine-tune and export.",
|
|
"Call studio_status first to see what is loaded and running, and list_models to find models.",
|
|
"Load a model with load_model before chat, generate_image, generate_audio or generate_video; embed and transcribe fall back to Unsloth Studio's own embedding and speech-to-text models.",
|
|
"Unsloth Studio may unload one model to fit another on the GPU, so load a model, generate, then unload_model before you switch kinds; speech-to-text models stay loaded beside them.",
|
|
"load_model lists what it unloaded in evicted, so check studio_status before reusing a model you loaded earlier.",
|
|
"generate_audio can speak, clone a voice, edit or convert speech, make music (a song or sfx) and separate stems, with a tts or audio model loaded.",
|
|
"chat, generate_audio and transcribe wait for their result and report progress every 2 s.",
|
|
"Video, recipe, export and training jobs return at once: poll get_job (video, recipe, export) or studio_status (training) every 15 to 30 s.",
|
|
"Stop work with cancel: a chat by the cancel_id you passed when starting it, a video by its id.",
|
|
"To fine-tune, find a dataset with datasets, check the config with start_training(validate_only=true), start it, then export_model the finished run.",
|
|
"Errors start with a tag such as [gpu_busy retry_after=12] or [not_loaded kind=image]; the hint after the message says what to do.",
|
|
"File paths in inputs work only when you run on the Unsloth Studio computer; otherwise send data or an Unsloth Studio id.",
|
|
"Generated images, audio and videos are saved to the Unsloth Studio galleries and returned with their id and URL.",
|
|
)
|
|
)
|
|
|
|
|
|
def create_studio_mcp() -> FastMCP:
|
|
"""Create the Unsloth MCP server and register the high-value tools."""
|
|
mcp = FastMCP("Unsloth Studio", instructions = INSTRUCTIONS)
|
|
|
|
for module in (
|
|
status,
|
|
models,
|
|
text,
|
|
images,
|
|
audio,
|
|
video,
|
|
jobs,
|
|
data,
|
|
training,
|
|
export,
|
|
library,
|
|
cancel,
|
|
):
|
|
for tool, annotations, *schema in module.TOOLS:
|
|
mcp.tool(
|
|
tool, annotations = annotations, **({"output_schema": schema[0]} if schema else {})
|
|
)
|
|
|
|
for prompt, title in prompts.PROMPTS:
|
|
mcp.prompt(prompt, title = title)
|
|
|
|
return mcp
|