mirror of
https://github.com/unslothai/unsloth.git
synced 2026-10-19 00:49:01 +00:00
4bed2fa246
Includes #14072 (Studio: audio.cpp models land on the right pages); review this PR's own commits from 4d58b31e05 on. `git diff e1d6db49bc HEAD` (the #14072 head) shows only this PR's changes; the merge commit just brings #14072 up to date. ## Summary Connecting an agent to the Unsloth Studio MCP server meant copying a per-agent command from `MCP.md`, and the documented Claude Code command expanded `$UNSLOTH_API_KEY` in the shell, so the key itself was written into `~/.claude.json`. Agents also had no way to see which image, video or audio models they could download, or whether one would fit on the GPU. This adds a setup command, a connection check, a download catalog in `list_models`, and two recipe prompts. ## Changes `unsloth mcp add {claude-code,codex,cursor,generic}` writes the agent's MCP entry. The entry names an environment variable (`UNSLOTH_API_KEY` by default, `--env-var` to change it) in the `Authorization` header, and the command never reads, prints or writes the key. - It uses `--url`, then `UNSLOTH_STUDIO_URL`, then a running Unsloth Studio on this machine, then `http://127.0.0.1:8888`. - Claude Code goes through `claude mcp add --scope user` when `claude` is on PATH, otherwise it edits `~/.claude.json` directly. Codex gets a `[mcp_servers.unsloth-studio]` table with `tool_timeout_sec = 300`, and the edit is refused if any other TOML setting would change. Cursor gets `~/.cursor/mcp.json` with `${env:UNSLOTH_API_KEY}`. `generic` prints a JSON snippet, or merges it into `--output`. - `--dry-run` prints just the entry it would write and writes nothing. It never prints the old entry it would replace, since one written by hand may hold a key. Re-running with the same arguments reports "unchanged". New config files are created with mode 0600, and existing files keep their mode. `unsloth doctor --mcp [--env-var NAME]` checks that Unsloth Studio answers, that the key is accepted, and that `/mcp/` is on (it tells "turned off" apart from "this Unsloth Studio predates MCP"). It warns, without failing, when a configured Codex or Claude Code tool timeout is under 240 s. Plain `--mcp` skips the GPU checks; `--mcp --gpus` runs both. `list_models(source="catalog")` lists image, video and audio models Unsloth Studio can download. Audio rows are forwarded from the audio.cpp catalog route in #14072 (`/api/inference/audio/catalog`), rather than importing its builder, so they stay filtered per account like every other MCP call. Image and video rows come from a new `GET /api/inference/media/catalog`, built from the diffusion and video family registries, filtered to what the calling account may load, and limited to the video families this device can run (as on the Video page). Catalog rows carry `size_bytes` (decimal bytes), `low_vram` (a video family's GGUF repo) and `fits_now`, which compares the size against the free memory of the single GPU with the most free memory, or is null when either value is unknown. The default `source="local"` output is unchanged. LLM rows in `list_models(source="catalog")` are a possible follow-up pending a maintainer decision. For now the docstring tells agents to list downloaded LLMs with the default source or load a Hugging Face id. The MCP server now offers two prompts, which Claude Code shows as slash commands: - `make_short_film(idea, length_seconds)` reads free GPU memory, loaded models and the video, music and voice catalog rows that fit now, then lists the calls from `load_model` through `generate_video`, `get_job` and `generate_audio`. - `build_dataset_and_train(goal, base_model)` reads the downloaded LLMs and lists the dataset, Data Recipe, `start_training` and `export_model` calls. Both fall back to plain steps when Unsloth Studio can't be read, and the rendered text is capped at 4000 characters. `MCP.md` gets a "One command" section for `mcp add` and `doctor --mcp`, the catalog in the `list_models` row, and a Prompts section. The hand-written Claude Code commands now single-quote `${UNSLOTH_API_KEY}` and use `--scope user`, so the key stays out of `~/.claude.json`. Tool annotations (titles, `idempotentHint`, `destructiveHint` on `load_model`) will land as a final commit after the other MCP PRs merge. ## Testing New tests: `unsloth_cli/tests/test_mcp_add.py`, `unsloth_cli/tests/test_doctor_mcp.py`, `studio/backend/tests/test_mcp_prompts.py`, `studio/backend/tests/test_media_catalog.py`, and catalog cases in `studio/backend/tests/test_mcp_tools_models.py`. `test_mcp_no_host_paths.py` also runs `list_models(source="catalog")`, and the `list_models` schema row in `test_mcp_server.py` is now a superset check, so it accepts `source`. - `unsloth_cli/tests/`: 2084 passed, 4 skipped. The 1 failure, `test_studio_update_verify.py::test_a_no_torch_install_keeps_that_mode_in_the_reinstall`, fails the same way on the #14072 head. - `studio/backend/tests/test_mcp_*.py` plus the media and audio.cpp catalog tests: 1708 passed, 3 skipped (fastmcp 3.4.8). The MCP, catalog and prompt files also pass on fastmcp 3.2.4. `FastMCP.prompt(title=)` exists back to the 3.0.2 floor. - Leak tests: a dry run over configs holding literal keys (any `Bearer` spelling, other headers, `args`, TOML `bearer_token` in any quoting, multi-line strings, another server's token) prints none of them, across 56 probe cases. Another test sets `UNSLOTH_API_KEY` to a canary and checks stdout, every written file and the `claude` argv for all four agents, with and without `--dry-run`. - `fits_now` is covered for a stubbed 24 GB and 12 GB card (Wan2.2-TI2V-5B, 24.2 GB, fits only the first), the GiB/GB boundary, and two GPUs (the max counts, not the sum). - Negative controls: `unsloth doctor --no-gpus` output is byte-identical to the #14072 head, and `doctor --help` only gains the `--mcp` line. Dumping every MCP tool's schema, description and annotations before and after shows only `list_models` changed. Against a running Unsloth Studio (no model loaded): - `unsloth doctor --mcp` with a valid key printed "API key: set, accepted", "MCP endpoint: on" and "Verdict: OK", exit 0. With a wrong key it printed "set, refused" and exit 1. The key appeared in neither output. - `mcp add claude-code`, `codex` and `cursor`, each run twice in a sandboxed home, wrote once and then reported "unchanged". `--dry-run` printed the diff and wrote nothing, and every created file was mode 600. `claude mcp get` and `codex mcp get` read the entries back with the env var name, not a key. - Claude Code 2.1.302 expands `${UNSLOTH_API_KEY}` in a user-scope header when it connects, so the entry `mcp add claude-code` writes authenticates. - A no-context Claude Code agent with only this MCP server, asked to plan a 20 s film with music and a voice-over for one 24 GB GPU, called `studio_status`, then `list_models(source="catalog")` for video, audio and TTS. It picked `unsloth/Wan2.2-TI2V-5B-GGUF` (a `low_vram` row), Stable-Audio-3-Medium-GGUF (3.6 GB) and Qwen3-TTS-12Hz-1.7B-VoiceDesign-GGUF (2.8 GB). That took 6 tool calls with no errors. The full catalog is 130 rows (33 KB) and returns in 0.09 s. - `prompts/list` and `prompts/get` for both prompts work over HTTP with live GPU and catalog data. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han <23090290+danielhanchen@users.noreply.github.com> |
||
|---|---|---|
| .. | ||
| __init__.py | ||
| _password_prompt.py | ||
| chat.py | ||
| doctor.py | ||
| eval.py | ||
| export.py | ||
| inference.py | ||
| mcp.py | ||
| start.py | ||
| studio.py | ||
| train.py | ||