mirror of
https://github.com/unslothai/unsloth.git
synced 2026-10-19 00:49:01 +00:00
Studio: one-command MCP setup, recipe prompts, a model catalog and tool annotations (#14077)
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>
This commit is contained in:
Etherl
GitHub
parent
df9a01c7a1
commit
4bed2fa246
21 changed files with 2527 additions and 26 deletions
|
|
@ -132,6 +132,38 @@ The tools changed:
|
|||
|
||||
### Set up an agent
|
||||
|
||||
#### One command
|
||||
|
||||
`unsloth mcp add` writes the agent's MCP entry for you. It supports
|
||||
`claude-code`, `codex`, `cursor` and `generic`:
|
||||
|
||||
```bash
|
||||
unsloth mcp add claude-code
|
||||
unsloth mcp add codex --url http://192.168.1.20:8888
|
||||
unsloth mcp add cursor --dry-run
|
||||
```
|
||||
|
||||
- The entry names the `UNSLOTH_API_KEY` environment variable, never the key
|
||||
itself, and the command never prints a key. Pick another variable with `--env-var`.
|
||||
- It finds a running Unsloth Studio on this computer, or takes `--url`.
|
||||
- `--dry-run` prints the entry it would write and writes nothing. Running it
|
||||
again changes nothing.
|
||||
- A config file it creates is readable by you only (mode 0600).
|
||||
- `codex` also sets `tool_timeout_sec = 300`.
|
||||
- `generic` prints a JSON entry to paste into any other agent's config.
|
||||
|
||||
Then check the connection:
|
||||
|
||||
```bash
|
||||
unsloth doctor --mcp
|
||||
```
|
||||
|
||||
It checks that Unsloth Studio answers, that your key works, and that the MCP
|
||||
server is on. It also warns when an agent's tool timeout is under 240 s, because
|
||||
long tools such as `load_model` can wait that long.
|
||||
|
||||
#### By hand
|
||||
|
||||
Each one reads the key from the `UNSLOTH_API_KEY` environment variable, so set
|
||||
it before you start the agent. Replace `http://127.0.0.1:8888` with your own
|
||||
address.
|
||||
|
|
@ -139,13 +171,13 @@ address.
|
|||
**Claude Code** (macOS or Linux):
|
||||
|
||||
```bash
|
||||
claude mcp add --transport http unsloth-studio http://127.0.0.1:8888/mcp/ --header "Authorization: Bearer $UNSLOTH_API_KEY"
|
||||
claude mcp add --transport http --scope user unsloth-studio http://127.0.0.1:8888/mcp/ --header 'Authorization: Bearer ${UNSLOTH_API_KEY}'
|
||||
```
|
||||
|
||||
**Claude Code** (Windows PowerShell):
|
||||
|
||||
```powershell
|
||||
claude mcp add --transport http unsloth-studio http://127.0.0.1:8888/mcp/ --header "Authorization: Bearer $env:UNSLOTH_API_KEY"
|
||||
claude mcp add --transport http --scope user unsloth-studio http://127.0.0.1:8888/mcp/ --header 'Authorization: Bearer ${UNSLOTH_API_KEY}'
|
||||
```
|
||||
|
||||
**OpenAI Codex**:
|
||||
|
|
@ -235,7 +267,7 @@ api_key_format = "Bearer {token}"
|
|||
| Tool | What it does |
|
||||
|---|---|
|
||||
| `studio_status` | What is loaded, loading and running: chat, image, video, speech-to-text and embedding models, training, export and GPUs. |
|
||||
| `list_models` | Models Unsloth Studio can serve, by kind or by Audio page `workflow`. With `model`, its training defaults. |
|
||||
| `list_models` | Models Unsloth Studio can serve, by kind or by Audio page `workflow`. With `model`, its training defaults. With `source="catalog"`, the image, video and audio models it can download, with their size and whether they fit on the GPU now. |
|
||||
| `load_model` | Load a chat, image, video or speech-to-text model, downloading it first if needed. Reports progress. |
|
||||
| `unload_model` | Unload a model. |
|
||||
| `chat` | Ask the loaded model, with optional images. Returns a `cancel_id`. |
|
||||
|
|
@ -299,6 +331,13 @@ Tool results never contain paths on the Unsloth Studio computer. Checkpoints are
|
|||
Unsloth Studio's exports folder, and `export_model` takes only a relative
|
||||
`save_directory`.
|
||||
|
||||
### Prompts
|
||||
|
||||
The server also offers two prompts, which Claude Code shows as slash commands:
|
||||
`make_short_film` and `build_dataset_and_train`. Each one reads what is loaded
|
||||
and which catalog models fit on the GPU right now, then lays out the tool calls
|
||||
for the job.
|
||||
|
||||
### Known limits
|
||||
|
||||
- An API key cannot unload a chat model that was loaded from a local folder.
|
||||
|
|
|
|||
78
studio/backend/core/inference/media_catalog.py
Normal file
78
studio/backend/core/inference/media_catalog.py
Normal file
|
|
@ -0,0 +1,78 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""The image and video models Unsloth Studio can offer, built from the family registries.
|
||||
|
||||
Pure and torch-free. Rows carry no path, so account filtering keys on ``id``; the route filters
|
||||
per request and this payload is never filtered in place.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import List, Literal, Optional
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from .diffusion_families import _FAMILIES as IMAGE_FAMILIES
|
||||
from .diffusion_families import family_selectable, mirror_repo
|
||||
from .diffusion_inference_info import family_inference_infos
|
||||
from .video_families import pipeline_available_video_families
|
||||
|
||||
|
||||
class MediaCatalogModel(BaseModel):
|
||||
"""One image or video model: a family's base repo, or a video family's GGUF repo."""
|
||||
|
||||
id: str
|
||||
kind: Literal["image", "video"]
|
||||
family: str
|
||||
display_name: str
|
||||
# Decimal bytes of the bf16 weights resident at once; None when the registry has no size.
|
||||
size_bytes: Optional[int] = None
|
||||
# A reduced-memory artifact (a GGUF repo) rather than the full-precision base.
|
||||
low_vram: bool = False
|
||||
|
||||
|
||||
class MediaCatalogResponse(BaseModel):
|
||||
models: List[MediaCatalogModel] = []
|
||||
|
||||
|
||||
def _gb_to_bytes(gb: Optional[float]) -> Optional[int]:
|
||||
return None if gb is None else round(gb * 1e9)
|
||||
|
||||
|
||||
def _row(
|
||||
kind: str,
|
||||
family: str,
|
||||
repo: str,
|
||||
size_gb: Optional[float],
|
||||
low_vram: bool = False,
|
||||
):
|
||||
# The unsloth mirror when there is one, as the model picker shows it; gated vendor repos stay out.
|
||||
repo_id = mirror_repo(repo) or repo
|
||||
return MediaCatalogModel(
|
||||
id = repo_id,
|
||||
kind = kind,
|
||||
family = family,
|
||||
display_name = repo_id.rsplit("/", 1)[-1],
|
||||
size_bytes = _gb_to_bytes(size_gb),
|
||||
low_vram = low_vram,
|
||||
)
|
||||
|
||||
|
||||
def media_catalog_payload(device: Optional[str] = None) -> MediaCatalogResponse:
|
||||
"""Every selectable image family, then every video family ``device`` can run and its GGUF repo, in registry order."""
|
||||
bf16_gb = {
|
||||
info["family"]: info["estimated_resident_gb"].get("bf16")
|
||||
for info in family_inference_infos()
|
||||
}
|
||||
rows = [
|
||||
_row("image", fam.name, fam.base_repo, bf16_gb.get(fam.name))
|
||||
for fam in IMAGE_FAMILIES
|
||||
if family_selectable(fam)
|
||||
]
|
||||
for fam in pipeline_available_video_families(device = device):
|
||||
sizes = fam.bf16_components_gb
|
||||
rows.append(_row("video", fam.name, fam.base_repo, sum(sizes) if sizes else None))
|
||||
if fam.gguf_repo:
|
||||
rows.append(_row("video", fam.name, fam.gguf_repo, None, low_vram = True))
|
||||
return MediaCatalogResponse(models = rows)
|
||||
|
|
@ -7,6 +7,7 @@ from __future__ import annotations
|
|||
|
||||
from fastmcp import FastMCP
|
||||
|
||||
from studio_mcp import prompts
|
||||
from studio_mcp.tools import (
|
||||
audio,
|
||||
cancel,
|
||||
|
|
@ -60,4 +61,7 @@ def create_studio_mcp() -> FastMCP:
|
|||
tool, annotations = annotations, **({"output_schema": schema[0]} if schema else {})
|
||||
)
|
||||
|
||||
for prompt, title in prompts.PROMPTS:
|
||||
mcp.prompt(prompt, title = title)
|
||||
|
||||
return mcp
|
||||
|
|
|
|||
|
|
@ -51387,6 +51387,33 @@ async def audio_catalog(current_subject: str = Depends(get_current_subject)):
|
|||
return payload.model_copy(update = {"models": [AudioCatalogModel(**r) for r in rows]})
|
||||
|
||||
|
||||
@studio_router.get("/media/catalog")
|
||||
async def media_catalog(current_subject: str = Depends(get_current_subject)):
|
||||
"""The image and video models Unsloth Studio can offer, filtered to what this account may see."""
|
||||
from core.inference.diffusion_device import resolve_diffusion_device_target
|
||||
from core.inference.media_catalog import MediaCatalogResponse, media_catalog_payload
|
||||
|
||||
def device() -> Optional[str]:
|
||||
# The Video page hides families the device can't run (Apple Silicon); so does the catalog.
|
||||
try:
|
||||
return resolve_diffusion_device_target().device
|
||||
except Exception: # noqa: BLE001 - no device read: list every family
|
||||
return None
|
||||
|
||||
try:
|
||||
payload = await asyncio.to_thread(lambda: media_catalog_payload(device()))
|
||||
rows = await asyncio.to_thread(
|
||||
account_access.filter_model_rows, [m.model_dump() for m in payload.models]
|
||||
)
|
||||
except Exception: # noqa: BLE001 - no exception text reaches the client
|
||||
logger.debug("media catalog failed", exc_info = True)
|
||||
raise HTTPException(
|
||||
status_code = 503,
|
||||
detail = "Unsloth Studio could not list image and video models right now.",
|
||||
) from None
|
||||
return MediaCatalogResponse(models = rows)
|
||||
|
||||
|
||||
@router.get("/audio/voices", response_model = AudioVoiceListResponse)
|
||||
async def list_audio_voices(request: Request, current_subject: str = Depends(get_current_subject)):
|
||||
from core.inference import audio_voices
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ from fastmcp.exceptions import ToolError
|
|||
from studio_mcp.caller import Caller
|
||||
from studio_mcp.errors import tool_error
|
||||
from studio_mcp.outputs import LoadResult, UnloadResult
|
||||
from studio_mcp.tools import as_dict, number, opt_text, route_json, strings
|
||||
from studio_mcp.tools import as_dict, number, opt_text, route_json, strings, try_json
|
||||
|
||||
POLL_INTERVAL_S = 2.0
|
||||
|
||||
|
|
@ -306,6 +306,28 @@ def stt_engine(status: dict, model: str) -> str:
|
|||
return "transformers"
|
||||
|
||||
|
||||
async def catalog_stt_engine(caller: Caller, status: dict, model: str) -> str:
|
||||
engine = stt_engine(status, model)
|
||||
known = any(
|
||||
isinstance(state, dict)
|
||||
and model in {*strings(state.get("models")), *strings(state.get("downloaded_models"))}
|
||||
for state in status.values()
|
||||
)
|
||||
if known:
|
||||
return engine
|
||||
# list_models(source="catalog") offers audio.cpp ASR rows no status list names until downloaded.
|
||||
rows = as_dict(await try_json(caller, "/api/inference/audio/catalog")).get("models")
|
||||
for row in rows if isinstance(rows, list) else []:
|
||||
if (
|
||||
isinstance(row, dict)
|
||||
and row.get("task") == "asr"
|
||||
and not row.get("unsupported_reason")
|
||||
and str(row.get("id") or "").lower() == model.split(":", 1)[0].lower()
|
||||
):
|
||||
return "audiocpp"
|
||||
return engine
|
||||
|
||||
|
||||
async def download_stt(
|
||||
caller: Caller,
|
||||
ctx: Optional[Context],
|
||||
|
|
@ -359,7 +381,7 @@ async def load_stt(
|
|||
hf_token: Optional[str],
|
||||
) -> LoadResult:
|
||||
status = await stt_status(caller, model)
|
||||
engine = stt_engine(status, model)
|
||||
engine = await catalog_stt_engine(caller, status, model)
|
||||
state = as_dict(status.get(engine))
|
||||
# Load refuses a model that is not on disk, so download first and load only once that finished.
|
||||
# The route folds a variant into "model:variant"; another cached variant does not count.
|
||||
|
|
|
|||
|
|
@ -182,6 +182,12 @@ class ModelEntry(ToolOutput):
|
|||
display_name: Optional[RouteText] = None
|
||||
quant: Optional[RouteText] = None
|
||||
context_length: Optional[int] = None
|
||||
# Catalog rows only: decimal bytes, a reduced-memory artifact, whether it fits one GPU's free memory now.
|
||||
size_bytes: Optional[int] = None
|
||||
low_vram: Optional[bool] = None
|
||||
fits_now: Optional[bool] = None
|
||||
# "catalog" for list_models(source="catalog") rows.
|
||||
source: Optional[RouteText] = None
|
||||
audio_workflows: Optional[list[RouteText]] = None
|
||||
|
||||
|
||||
|
|
|
|||
201
studio/backend/studio_mcp/prompts.py
Normal file
201
studio/backend/studio_mcp/prompts.py
Normal file
|
|
@ -0,0 +1,201 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Recipe prompts for agents: each one reads what Unsloth Studio has right now and spells out the tool calls for a whole job."""
|
||||
|
||||
# No ``from __future__ import annotations`` here: fastmcp tells agents to send any argument whose
|
||||
# annotation is not the ``str`` type itself as a "JSON string", and a postponed annotation is the text "str".
|
||||
|
||||
import asyncio
|
||||
from typing import Any, Optional
|
||||
|
||||
from studio_mcp.tools import models, status
|
||||
|
||||
MAX_CHARS = 4000
|
||||
# What the agent asked for is clipped so the recipe always fits.
|
||||
MAX_ASK_CHARS = 800
|
||||
CANDIDATES = 5
|
||||
# The Audio page a film needs from each audio kind; the catalog also holds detectors and separators.
|
||||
FILM_PAGES = {"audio": "music", "tts": "speak"}
|
||||
# GpuDevice reports GiB; sizes and fits_now use decimal GB.
|
||||
GIB_TO_GB = 1.073741824
|
||||
FALLBACK = (
|
||||
"Unsloth Studio could not be read from here: call studio_status and "
|
||||
'list_models(source="catalog") first.'
|
||||
)
|
||||
|
||||
|
||||
def _gb(value: float) -> str:
|
||||
return f"{value:.1f} GB"
|
||||
|
||||
|
||||
def _gpu_line(studio: Any) -> str:
|
||||
devices = []
|
||||
for device in studio.hardware.devices:
|
||||
if device.vram_total_gb is None or device.vram_used_gb is None:
|
||||
continue
|
||||
free = max(0.0, device.vram_total_gb - device.vram_used_gb) * GIB_TO_GB
|
||||
name = f"GPU {device.index}" if device.index is not None else "GPU"
|
||||
devices.append(f"{name} {_gb(free)} free of {_gb(device.vram_total_gb * GIB_TO_GB)}")
|
||||
if not devices:
|
||||
return "Unsloth Studio reports no GPU memory right now; check studio_status before loading."
|
||||
return "Unsloth Studio free GPU memory: " + "; ".join(devices) + "."
|
||||
|
||||
|
||||
def _loaded_line(studio: Any) -> str:
|
||||
loaded = [model.id for model in studio.chat.loaded]
|
||||
loaded += [slot.model for slot in (studio.image, studio.video) if slot.loaded and slot.model]
|
||||
if studio.stt.model:
|
||||
loaded.append(studio.stt.model)
|
||||
if not loaded:
|
||||
return "Unsloth Studio has no model loaded."
|
||||
return "Unsloth Studio has loaded: " + ", ".join(loaded) + "."
|
||||
|
||||
|
||||
def _catalog_line(
|
||||
kind: str,
|
||||
rows: Any,
|
||||
page: Optional[str] = None,
|
||||
) -> str:
|
||||
fitting = [
|
||||
row
|
||||
for row in rows.models
|
||||
if getattr(row, "fits_now", None) is True
|
||||
and (page is None or page in (getattr(row, "audio_workflows", None) or []))
|
||||
]
|
||||
fitting.sort(key = lambda row: getattr(row, "size_bytes", None) or 0)
|
||||
if not fitting:
|
||||
return f"- {kind}: no Unsloth Studio catalog model fits the free GPU memory now; unload a model or try a low_vram row."
|
||||
shown = []
|
||||
for row in fitting[:CANDIDATES]:
|
||||
size = getattr(row, "size_bytes", None)
|
||||
shown.append(f"{row.id} ({_gb(size / 1e9)})" if size else row.id)
|
||||
label = f"{kind} ({page})" if page else kind
|
||||
return f"- {label}: " + ", ".join(shown)
|
||||
|
||||
|
||||
def _local_line(kind: str, rows: Any) -> str:
|
||||
ids = [row.id + (" (loaded)" if row.loaded else "") for row in rows.models[:CANDIDATES]]
|
||||
if not ids:
|
||||
return f"- {kind}: none downloaded in Unsloth Studio yet; start_training downloads a Hugging Face id."
|
||||
return f"- {kind}: " + ", ".join(ids)
|
||||
|
||||
|
||||
async def _facts(
|
||||
kinds: tuple[str, ...],
|
||||
local: bool = False,
|
||||
pages: Optional[dict[str, str]] = None,
|
||||
) -> str:
|
||||
"""The GPU, loaded models and model candidates as text; never raises, so a prompt renders without HTTP."""
|
||||
|
||||
async def read(kind: str) -> Any:
|
||||
if local:
|
||||
return await models.list_models(kind = kind)
|
||||
return await models.list_models(kind = kind, source = "catalog")
|
||||
|
||||
results = await asyncio.gather(
|
||||
status.studio_status(), *(read(kind) for kind in kinds), return_exceptions = True
|
||||
)
|
||||
studio, listings = results[0], results[1:]
|
||||
if all(isinstance(result, BaseException) for result in results):
|
||||
return FALLBACK
|
||||
lines = (
|
||||
[FALLBACK]
|
||||
if isinstance(studio, BaseException)
|
||||
else [_gpu_line(studio), _loaded_line(studio)]
|
||||
)
|
||||
title = "Downloaded models" if local else "Catalog models that fit now, smallest first"
|
||||
lines.append(f"{title} (Unsloth Studio, at most {CANDIDATES} each):")
|
||||
for kind, rows in zip(kinds, listings):
|
||||
if isinstance(rows, BaseException):
|
||||
lines.append(f'- {kind}: unknown; call list_models(kind="{kind}") yourself.')
|
||||
else:
|
||||
page = (pages or {}).get(kind)
|
||||
lines.append(_local_line(kind, rows) if local else _catalog_line(kind, rows, page))
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _clip(text: str) -> str:
|
||||
text = text.strip()
|
||||
return text if len(text) <= MAX_ASK_CHARS else text[:MAX_ASK_CHARS] + "…"
|
||||
|
||||
|
||||
def _capped(text: str) -> str:
|
||||
return text if len(text) <= MAX_CHARS else text[: MAX_CHARS - 1] + "…"
|
||||
|
||||
|
||||
async def make_short_film(idea: str, length_seconds: str = "20") -> str:
|
||||
"""Plan and render a short film with Unsloth Studio: video shots, then music or a voice-over.
|
||||
|
||||
Args:
|
||||
idea: What the film is about.
|
||||
length_seconds: The total length to aim for, in seconds.
|
||||
"""
|
||||
facts = await _facts(("video", "audio", "tts"), pages = FILM_PAGES)
|
||||
length = length_seconds.strip() or "20"
|
||||
return _capped(
|
||||
f"Make a short film of about {length} seconds with Unsloth Studio.\n"
|
||||
f"Idea: {_clip(idea)}\n\n"
|
||||
f"What Unsloth Studio has now:\n{facts}\n\n"
|
||||
"Steps:\n"
|
||||
"1. Call studio_status to see the GPUs and what Unsloth Studio has loaded.\n"
|
||||
'2. Call list_models(source="catalog", kind="video") and pick a row with fits_now true, '
|
||||
"or a low_vram row on a small GPU.\n"
|
||||
'3. Call load_model(model=<id>, kind="video"); Unsloth Studio downloads it first if needed.\n'
|
||||
f"4. Write a shot list: Unsloth Studio video models make clips of a few seconds, so plan several "
|
||||
f"short shots that add up to about {length} seconds.\n"
|
||||
"5. Call generate_video once per shot with a self-contained prompt, keeping characters, style "
|
||||
"and lighting the same in every shot.\n"
|
||||
'6. Poll get_job(kind="video", id=<job id>) until each shot is completed and keep its URL.\n'
|
||||
'7. Call unload_model(kind="video") before loading an audio model, so Unsloth Studio has room.\n'
|
||||
'8. For music, load a list_models(source="catalog", kind="audio", workflow="music") row and call '
|
||||
'generate_audio(workflow="music", mode="song", instrumental=true, duration_s=<length>). '
|
||||
'Not every music model makes an instrumental song: if Unsloth Studio refuses mode="song", the '
|
||||
'model makes sound effects only, so call mode="sfx" with text describing the sound or pick '
|
||||
"another row; if it says the model always sings, set instrumental=false and pass lyrics. "
|
||||
'For a voice-over, load a list_models(source="catalog", kind="tts", workflow="speak") row and '
|
||||
'call generate_audio(workflow="speak", text=<script>).\n'
|
||||
"9. Unload each model when you are done with it, then download the clips and audio from their "
|
||||
"Unsloth Studio URLs and stitch them yourself, for example with ffmpeg.\n"
|
||||
)
|
||||
|
||||
|
||||
async def build_dataset_and_train(goal: str, base_model: str = "") -> str:
|
||||
"""Build a training dataset with Unsloth Studio, fine-tune a model on it and export the result.
|
||||
|
||||
Args:
|
||||
goal: What the fine-tuned model should do.
|
||||
base_model: A Hugging Face model id to start from; leave empty to pick one.
|
||||
"""
|
||||
facts = await _facts(("llm",), local = True)
|
||||
base = base_model.strip()
|
||||
pick = (
|
||||
f'Start from {base}; list_models(model="{base}") returns Unsloth Studio\'s training defaults for it.'
|
||||
if base
|
||||
else "Pick a small base model: a downloaded one below, or a Hugging Face id such as unsloth/Qwen3-0.6B."
|
||||
)
|
||||
return _capped(
|
||||
"Build a dataset and fine-tune a model with Unsloth Studio.\n"
|
||||
f"Goal: {_clip(goal)}\n\n"
|
||||
f"What Unsloth Studio has now:\n{facts}\n\n"
|
||||
"Steps:\n"
|
||||
f"1. {pick}\n"
|
||||
'2. Find data with datasets(action="list"), or datasets(action="download", repo_id=<Hub id>), and '
|
||||
'check it with datasets(action="check_format", name=<dataset>).\n'
|
||||
"3. To make data instead, write a Data Recipe, check it with run_recipe(recipe=<recipe>, "
|
||||
'mode="validate"), try mode="preview", then run mode="full" and follow '
|
||||
'get_job(kind="recipe") until Unsloth Studio has saved the dataset.\n'
|
||||
"4. Build the training config and check it with start_training(config=<config>, "
|
||||
"validate_only=true), which starts nothing in Unsloth Studio.\n"
|
||||
"5. Call start_training(config=<config>) and poll studio_status until Unsloth Studio "
|
||||
"reports the run finished.\n"
|
||||
"6. Find the run with list_training_runs(include_checkpoints=true).\n"
|
||||
'7. Call export_model(checkpoint=<run>, format="gguf", save_directory=<folder>) and follow '
|
||||
'get_job(kind="export") until Unsloth Studio has written it.\n'
|
||||
)
|
||||
|
||||
|
||||
PROMPTS = (
|
||||
(make_short_film, "Make a short film"),
|
||||
(build_dataset_and_train, "Build a dataset and train"),
|
||||
)
|
||||
|
|
@ -211,7 +211,7 @@ async def _download_then_retry(
|
|||
model = _default_stt(status)
|
||||
if model is None:
|
||||
raise ToolError("Unsloth Studio did not name a default speech-to-text model to download")
|
||||
engine = loading.stt_engine(status, model)
|
||||
engine = await loading.catalog_stt_engine(caller, status, model)
|
||||
await loading.download_stt(caller, ctx, model = model, engine = engine, hf_token = None)
|
||||
|
||||
|
||||
|
|
@ -250,7 +250,7 @@ async def _stt_model(caller: Caller, model: Optional[str]) -> tuple[str, str]:
|
|||
model = _default_stt(status)
|
||||
if model is None:
|
||||
raise ToolError("Unsloth Studio did not name a speech-to-text model; pass model.")
|
||||
return model, loading.stt_engine(status, model)
|
||||
return model, await loading.catalog_stt_engine(caller, status, model)
|
||||
|
||||
|
||||
async def _transcribe_source(
|
||||
|
|
@ -386,7 +386,7 @@ async def _resident_stt(
|
|||
payload = await try_json(caller, loading.STT_STATUS, params = {"model": model} if model else None)
|
||||
status = payload if isinstance(payload, dict) else {}
|
||||
if model:
|
||||
return model, loading.stt_engine(status, model)
|
||||
return model, await loading.catalog_stt_engine(caller, status, model)
|
||||
for engine in loading.STT_ENGINES:
|
||||
state = status.get(engine)
|
||||
if isinstance(state, dict) and opt_text(state.get("loaded_model")):
|
||||
|
|
|
|||
|
|
@ -5,15 +5,29 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Literal, Optional
|
||||
import asyncio
|
||||
from typing import Annotated, Any, Literal, Optional
|
||||
from urllib.parse import quote
|
||||
|
||||
from fastmcp import Context
|
||||
from fastmcp.exceptions import ToolError
|
||||
from pydantic import Field
|
||||
|
||||
from studio_mcp import loading
|
||||
from studio_mcp.caller import current_caller
|
||||
from studio_mcp.outputs import LoadResult, ModelEntry, ModelList, UnloadResult
|
||||
from studio_mcp.tools import DESTRUCTIVE, READ_ONLY, WRITES, as_dict, opt_text, route_json
|
||||
from studio_mcp.outputs import HardwareSlot, LoadResult, ModelEntry, ModelList, UnloadResult
|
||||
from studio_mcp.tools import (
|
||||
DESTRUCTIVE,
|
||||
READ_ONLY,
|
||||
WRITES,
|
||||
as_dict,
|
||||
integer,
|
||||
opt_text,
|
||||
route_json,
|
||||
strings,
|
||||
try_json,
|
||||
)
|
||||
from studio_mcp.tools.status import SLOT_ROUTES, _hardware
|
||||
|
||||
# "tts" and "audio" are what list_models reports for audio models; they load in the llm slot.
|
||||
ModelKind = Literal["llm", "image", "video", "stt", "tts", "audio"]
|
||||
|
|
@ -50,6 +64,9 @@ TRAINING_DEFAULT_KEYS = {
|
|||
}
|
||||
|
||||
|
||||
CATALOG_FIELDS = ("size_bytes", "low_vram", "fits_now", "source")
|
||||
|
||||
|
||||
def kind_of(entry: dict) -> str:
|
||||
task = entry.get("task")
|
||||
if task is None:
|
||||
|
|
@ -60,7 +77,8 @@ def kind_of(entry: dict) -> str:
|
|||
def _entry(entry: dict) -> Optional[ModelEntry]:
|
||||
if not opt_text(entry.get("id")):
|
||||
return None
|
||||
return ModelEntry.from_route(entry, kind = kind_of(entry))
|
||||
# The catalog fields stay null on local rows, whatever the route sends.
|
||||
return ModelEntry.from_route(entry, kind = kind_of(entry), **dict.fromkeys(CATALOG_FIELDS))
|
||||
|
||||
|
||||
def _training_defaults(details: Any) -> Optional[dict[str, Any]]:
|
||||
|
|
@ -81,13 +99,143 @@ def _same_model_id(listed: str, asked: str) -> bool:
|
|||
return listed.lower() == asked.split(":", 1)[0].lower() or listed.lower() == asked.lower()
|
||||
|
||||
|
||||
# The audio catalog's task -> the kind list_models reports and load_model takes.
|
||||
AUDIO_KIND = {"tts": "tts", "music": "audio", "sep": "audio", "asr": "stt", "": "audio"}
|
||||
MEDIA_KINDS = ("image", "video")
|
||||
AUDIO_KINDS = ("tts", "audio", "stt")
|
||||
# GpuDevice memory is GiB; catalog sizes are decimal GB.
|
||||
GIB_TO_GB = 1.073741824
|
||||
|
||||
|
||||
def free_gb(hardware: Optional[HardwareSlot]) -> Optional[float]:
|
||||
# One load lands on one GPU, so the most free memory on any one device counts, never the sum.
|
||||
free = [
|
||||
(device.vram_total_gb - device.vram_used_gb) * GIB_TO_GB
|
||||
for device in (hardware.devices if hardware else [])
|
||||
if device.vram_total_gb is not None and device.vram_used_gb is not None
|
||||
]
|
||||
return max(free) if free else None
|
||||
|
||||
|
||||
def fits(size_bytes: Optional[int], free: Optional[float]) -> Optional[bool]:
|
||||
if size_bytes is None or free is None:
|
||||
return None
|
||||
return size_bytes / 1e9 <= free
|
||||
|
||||
|
||||
async def _catalog_part(path: str, label: str, others: Optional[str]) -> list:
|
||||
try:
|
||||
listing = await route_json("GET", path)
|
||||
except ToolError as error:
|
||||
hint = f" Pass kind={others} to list the others." if others else ""
|
||||
message = f"Unsloth Studio could not list the {label} catalog: {error}{hint}"
|
||||
raise ToolError(message) from None
|
||||
rows = as_dict(listing).get("models")
|
||||
return [row for row in rows if isinstance(row, dict)] if isinstance(rows, list) else []
|
||||
|
||||
|
||||
async def _catalog(
|
||||
kind: Optional[str], loaded_only: bool, workflow: Optional[str]
|
||||
) -> list[ModelEntry]:
|
||||
if loaded_only:
|
||||
raise ToolError("loaded_only lists resident models; use the default source.")
|
||||
if kind == "llm":
|
||||
return []
|
||||
want_media = kind in (None, *MEDIA_KINDS) and not workflow
|
||||
want_audio = kind in (None, *AUDIO_KINDS)
|
||||
both = want_media and want_audio
|
||||
|
||||
async def nothing() -> list:
|
||||
return []
|
||||
|
||||
media, audio, hardware = await asyncio.gather(
|
||||
_catalog_part("/api/inference/media/catalog", "image and video", '"tts"' if both else None)
|
||||
if want_media
|
||||
else nothing(),
|
||||
_catalog_part("/api/inference/audio/catalog", "audio", '"image"' if both else None)
|
||||
if want_audio
|
||||
else nothing(),
|
||||
try_json(current_caller(), SLOT_ROUTES["hardware"]),
|
||||
)
|
||||
free = free_gb(_hardware(hardware) if isinstance(hardware, dict) else None)
|
||||
entries = []
|
||||
for row in media:
|
||||
if not opt_text(row.get("id")) or row.get("kind") not in MEDIA_KINDS:
|
||||
continue
|
||||
size = integer(row.get("size_bytes"))
|
||||
low_vram = row.get("low_vram") if isinstance(row.get("low_vram"), bool) else None
|
||||
entries.append(
|
||||
ModelEntry(
|
||||
id = row["id"],
|
||||
kind = row["kind"],
|
||||
display_name = opt_text(row.get("display_name")),
|
||||
size_bytes = size,
|
||||
low_vram = low_vram,
|
||||
fits_now = fits(size, free),
|
||||
source = "catalog",
|
||||
)
|
||||
)
|
||||
for row in audio:
|
||||
# A row the catalog refuses (unsupported family, no Audio page) can't be loaded.
|
||||
if not opt_text(row.get("id")) or row.get("unsupported_reason"):
|
||||
continue
|
||||
size = integer(row.get("size_bytes"))
|
||||
entries.append(
|
||||
ModelEntry(
|
||||
id = row["id"],
|
||||
kind = AUDIO_KIND.get(row.get("task") or "", "audio"),
|
||||
display_name = opt_text(row.get("display_name")),
|
||||
audio_workflows = strings(row.get("workflows")),
|
||||
size_bytes = size,
|
||||
fits_now = fits(size, free),
|
||||
source = "catalog",
|
||||
)
|
||||
)
|
||||
return entries
|
||||
|
||||
|
||||
async def _training_defaults_for(model: Optional[str]) -> Optional[dict[str, Any]]:
|
||||
if not model:
|
||||
return None
|
||||
details = await route_json(
|
||||
"GET", "/api/models/config/" + quote(model, safe = "/:"), hub_header = True
|
||||
)
|
||||
return _training_defaults(details)
|
||||
|
||||
|
||||
async def list_models(
|
||||
kind: Optional[ModelKind] = None,
|
||||
loaded_only: bool = False,
|
||||
model: Optional[str] = None,
|
||||
workflow: Optional[AudioWorkflow] = None,
|
||||
source: Annotated[
|
||||
Literal["local", "catalog"],
|
||||
Field(
|
||||
description = '"local" lists downloaded and loaded models; "catalog" lists image, video '
|
||||
"and audio models Unsloth Studio can download."
|
||||
),
|
||||
] = "local",
|
||||
) -> ModelList:
|
||||
"""Models this Unsloth Studio can serve: loaded ones and those already downloaded. Each entry has the id to pass to load_model, its kind and whether it is loaded. ``kind`` is read from the model's task; a model with no task is reported as "llm", which also covers embedding models and unloaded text-to-speech models Unsloth Studio cannot classify yet, so pass the kind explicitly to load_model. ``loaded_only`` returns just the resident models. ``model`` narrows the list to that model (empty when it is not downloaded yet) and adds Unsloth Studio's training defaults for it, which also works for a Hugging Face repo id that is not downloaded (send X-Unsloth-HF-Token for a gated repo). ``workflow`` keeps only audio models that run that Unsloth Studio Audio page (their ``audio_workflows``), such as the ids to pass to generate_audio as ``model``."""
|
||||
"""Models this Unsloth Studio can serve: loaded ones and those already downloaded.
|
||||
Each entry has the id to pass to load_model, its kind and whether it is loaded.
|
||||
``kind`` is read from the model's task; a model with no task is reported as "llm", which also covers embedding models and unloaded text-to-speech models Unsloth Studio cannot classify yet, so pass the kind explicitly to load_model.
|
||||
``loaded_only`` returns just the resident models.
|
||||
``model`` narrows the list to that model (empty when it is not downloaded yet) and adds Unsloth Studio's training defaults for it, which also works for a Hugging Face repo id that is not downloaded (send X-Unsloth-HF-Token for a gated repo).
|
||||
``workflow`` keeps only audio models that run that Unsloth Studio Audio page (their ``audio_workflows``), such as the ids to pass to generate_audio as ``model``.
|
||||
With source="catalog" it lists image, video and audio models Unsloth Studio can download, with size_bytes (decimal), low_vram and fits_now (whether it fits in the free memory of one GPU right now; null when unknown).
|
||||
LLM rows are not in the catalog: list downloaded LLMs with the default source, or load a Hugging Face id.
|
||||
Load a catalog row with load_model(model=id, kind=its kind).
|
||||
"""
|
||||
if source == "catalog":
|
||||
models = await _catalog(kind, loaded_only, workflow)
|
||||
models = [
|
||||
entry
|
||||
for entry in models
|
||||
if (not kind or entry.kind == kind)
|
||||
and (not model or _same_model_id(entry.id, model))
|
||||
and (not workflow or workflow in (entry.audio_workflows or []))
|
||||
]
|
||||
return ModelList(models = models, training_defaults = await _training_defaults_for(model))
|
||||
if loaded_only and kind == "llm" and not workflow:
|
||||
# Resident chat models without a scan of every model folder; it never lists image, video or STT.
|
||||
listing = await route_json("GET", "/api/inference/loaded-models")
|
||||
|
|
@ -104,13 +252,7 @@ async def list_models(
|
|||
if workflow and workflow not in (entry.audio_workflows or []):
|
||||
continue
|
||||
models.append(entry)
|
||||
defaults = None
|
||||
if model:
|
||||
details = await route_json(
|
||||
"GET", "/api/models/config/" + quote(model, safe = "/:"), hub_header = True
|
||||
)
|
||||
defaults = _training_defaults(details)
|
||||
return ModelList(models = models, training_defaults = defaults)
|
||||
return ModelList(models = models, training_defaults = await _training_defaults_for(model))
|
||||
|
||||
|
||||
async def load_model(
|
||||
|
|
|
|||
|
|
@ -68,7 +68,10 @@ EXPORT_PAYLOADS = {
|
|||
# tool name -> (fake Studio routes, tool arguments), or a list of them
|
||||
CASES: dict[str, Any] = {
|
||||
"studio_status": (STATUS_PAYLOADS, {}),
|
||||
"list_models": (MODELS_PAYLOADS, {"model": "unsloth/Qwen3-0.6B"}),
|
||||
"list_models": [
|
||||
(MODELS_PAYLOADS, {"model": "unsloth/Qwen3-0.6B"}),
|
||||
(MODELS_PAYLOADS, {"source": "catalog"}),
|
||||
],
|
||||
"load_model": (LOADING_PAYLOADS, {"model": "unsloth/Llama-3.2-1B-Instruct-GGUF"}),
|
||||
"unload_model": (LOADING_PAYLOADS, {}),
|
||||
"chat": (TEXT_PAYLOADS, {"prompt": "hi"}),
|
||||
|
|
|
|||
243
studio/backend/tests/test_mcp_prompts.py
Normal file
243
studio/backend/tests/test_mcp_prompts.py
Normal file
|
|
@ -0,0 +1,243 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import re
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from mcp_server import create_studio_mcp
|
||||
from studio_mcp import prompts
|
||||
from studio_mcp.outputs import (
|
||||
ChatModel,
|
||||
ChatSlot,
|
||||
GpuDevice,
|
||||
HardwareSlot,
|
||||
ModelEntry,
|
||||
ModelList,
|
||||
StudioStatus,
|
||||
)
|
||||
from studio_mcp.tools import models, status
|
||||
from .mcp_harness import MCP_HEADERS, TEST_TOKEN, fake_studio, served
|
||||
|
||||
ARGUMENTS = {
|
||||
"make_short_film": [("idea", True), ("length_seconds", False)],
|
||||
"build_dataset_and_train": [("goal", True), ("base_model", False)],
|
||||
}
|
||||
TITLES = {
|
||||
"make_short_film": "Make a short film",
|
||||
"build_dataset_and_train": "Build a dataset and train",
|
||||
}
|
||||
# 24 GiB card with 4 GiB used: 20 GiB free is 21.5 decimal GB.
|
||||
STATUS = StudioStatus(
|
||||
chat = ChatSlot(loaded = [ChatModel(id = "unsloth/Qwen3-0.6B")]),
|
||||
hardware = HardwareSlot(
|
||||
available = True, devices = [GpuDevice(index = 0, vram_used_gb = 4.0, vram_total_gb = 24.0)]
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _row(
|
||||
id,
|
||||
kind,
|
||||
size_gb,
|
||||
fits_now,
|
||||
pages = None,
|
||||
):
|
||||
# Catalog fields come from list_models(source="catalog"); duck-typed so the test does not pin ModelEntry.
|
||||
return SimpleNamespace(
|
||||
id = id,
|
||||
kind = kind,
|
||||
loaded = False,
|
||||
size_bytes = int(size_gb * 1e9),
|
||||
fits_now = fits_now,
|
||||
audio_workflows = pages,
|
||||
)
|
||||
|
||||
|
||||
CATALOG = {
|
||||
"video": [
|
||||
_row("unsloth/Wan-big", "video", 60.0, False),
|
||||
_row("unsloth/LTX-small", "video", 12.0, True),
|
||||
_row("unsloth/Wan-mid", "video", 18.5, True),
|
||||
],
|
||||
"audio": [
|
||||
_row("unsloth/wake-word", "audio", 0.1, True, []),
|
||||
_row("unsloth/ACE-Step", "audio", 8.0, True, ["music"]),
|
||||
],
|
||||
"tts": [_row(f"unsloth/tts-{i}", "tts", 1.0 + i, True, ["speak"]) for i in range(7)]
|
||||
+ [_row("unsloth/voice-clone-only", "tts", 0.5, True, ["clone"])],
|
||||
}
|
||||
LLMS = [
|
||||
ModelEntry(id = "unsloth/Qwen3-0.6B", kind = "llm", loaded = True),
|
||||
ModelEntry(id = "unsloth/gemma-3-1b-it", kind = "llm"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def stubbed(monkeypatch):
|
||||
calls = []
|
||||
|
||||
async def fake_status():
|
||||
return STATUS
|
||||
|
||||
async def fake_list(
|
||||
kind = None,
|
||||
source = "local",
|
||||
**_,
|
||||
):
|
||||
calls.append((kind, source))
|
||||
if source == "catalog":
|
||||
return SimpleNamespace(models = CATALOG.get(kind, []))
|
||||
return ModelList(models = [m for m in LLMS if m.kind == kind])
|
||||
|
||||
monkeypatch.setattr(status, "studio_status", fake_status)
|
||||
monkeypatch.setattr(models, "list_models", fake_list)
|
||||
return calls
|
||||
|
||||
|
||||
def _render(name, args):
|
||||
result = asyncio.run(create_studio_mcp().render_prompt(name, args))
|
||||
return "".join(message.content.text for message in result.messages)
|
||||
|
||||
|
||||
def _tool_names():
|
||||
return {tool.name for tool in asyncio.run(create_studio_mcp().list_tools())}
|
||||
|
||||
|
||||
def test_prompts_are_listed_with_titles_and_plain_arguments():
|
||||
listed = {p.name: p for p in asyncio.run(create_studio_mcp().list_prompts())}
|
||||
assert set(listed) == set(ARGUMENTS)
|
||||
for name, prompt in listed.items():
|
||||
assert prompt.title == TITLES[name]
|
||||
assert "Unsloth Studio" in prompt.description
|
||||
assert [(a.name, a.required) for a in prompt.arguments] == ARGUMENTS[name]
|
||||
for argument in prompt.arguments:
|
||||
assert argument.description, (name, argument.name)
|
||||
assert "JSON string" not in argument.description, (name, argument.name)
|
||||
|
||||
|
||||
def test_film_prompt_shows_fitting_models_and_free_gpu(stubbed):
|
||||
text = _render("make_short_film", {"idea": "a fox at dawn", "length_seconds": "20"})
|
||||
assert len(text) <= prompts.MAX_CHARS
|
||||
assert "a fox at dawn" in text and "about 20 seconds" in text
|
||||
assert "GPU 0 21.5 GB free of 25.8 GB" in text
|
||||
assert "unsloth/Qwen3-0.6B" in text
|
||||
# Fitting rows only, smallest first, at most five per kind.
|
||||
assert "unsloth/Wan-big" not in text
|
||||
assert text.index("unsloth/LTX-small (12.0 GB)") < text.index("unsloth/Wan-mid (18.5 GB)")
|
||||
# Audio rows are only the ones a film can use: music for the bed, speak for the voice-over.
|
||||
assert "- audio (music): unsloth/ACE-Step (8.0 GB)" in text
|
||||
assert "unsloth/wake-word" not in text and "unsloth/voice-clone-only" not in text
|
||||
assert 'kind="audio", workflow="music")' in text and 'kind="tts", workflow="speak")' in text
|
||||
assert "unsloth/tts-4" in text and "unsloth/tts-5" not in text
|
||||
assert sorted(stubbed) == [("audio", "catalog"), ("tts", "catalog"), ("video", "catalog")]
|
||||
|
||||
|
||||
def test_train_prompt_lists_downloaded_llms(stubbed):
|
||||
text = _render(
|
||||
"build_dataset_and_train",
|
||||
{"goal": "answer support tickets", "base_model": "unsloth/gemma-3-1b-it"},
|
||||
)
|
||||
assert "answer support tickets" in text
|
||||
assert "unsloth/Qwen3-0.6B (loaded)" in text and "unsloth/gemma-3-1b-it" in text
|
||||
assert 'list_models(model="unsloth/gemma-3-1b-it")' in text
|
||||
assert "21.5 GB free" in text
|
||||
# There are no LLM catalog rows, so the train prompt reads the downloaded list.
|
||||
assert stubbed == [("llm", "local")]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"name, args", [("make_short_film", {"idea": "x"}), ("build_dataset_and_train", {"goal": "x"})]
|
||||
)
|
||||
def test_prompts_fall_back_without_live_data(monkeypatch, name, args):
|
||||
async def boom(*_, **__):
|
||||
raise RuntimeError("No active HTTP request found.")
|
||||
|
||||
monkeypatch.setattr(status, "studio_status", boom)
|
||||
monkeypatch.setattr(models, "list_models", boom)
|
||||
text = _render(name, args)
|
||||
assert prompts.FALLBACK in text
|
||||
assert "Steps:" in text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"name, args", [("make_short_film", {"idea": "x"}), ("build_dataset_and_train", {"goal": "x"})]
|
||||
)
|
||||
def test_prompts_render_in_process_without_stubs(name, args):
|
||||
# No HTTP request means no caller: the real tools raise and the prompt still renders.
|
||||
assert prompts.FALLBACK in _render(name, args)
|
||||
|
||||
|
||||
def test_one_failed_listing_does_not_hide_the_rest(stubbed, monkeypatch):
|
||||
real = models.list_models
|
||||
|
||||
async def no_audio(kind = None, **kwargs):
|
||||
if kind == "audio":
|
||||
raise RuntimeError("down")
|
||||
return await real(kind = kind, **kwargs)
|
||||
|
||||
monkeypatch.setattr(models, "list_models", no_audio)
|
||||
text = _render("make_short_film", {"idea": "x"})
|
||||
assert '- audio: unknown; call list_models(kind="audio") yourself.' in text
|
||||
assert "unsloth/LTX-small" in text and "21.5 GB free" in text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"name, args", [("make_short_film", {"idea": "x"}), ("build_dataset_and_train", {"goal": "x"})]
|
||||
)
|
||||
def test_rendered_calls_name_registered_tools(stubbed, name, args):
|
||||
tools = _tool_names()
|
||||
called = set(re.findall(r"\b([a-z_]+)\(", _render(name, args)))
|
||||
assert called and called <= tools, called - tools
|
||||
|
||||
|
||||
def test_long_input_stays_under_the_cap(stubbed):
|
||||
text = _render("make_short_film", {"idea": "x" * 10_000, "length_seconds": "20"})
|
||||
assert len(text) <= prompts.MAX_CHARS
|
||||
assert "Steps:" in text
|
||||
|
||||
|
||||
def test_film_soundtrack_call_fits_models_without_instrumental_songs(stubbed):
|
||||
# The music page also lists sound-effect-only models (Stable Audio SFX, ControlFoley) and ones that
|
||||
# always sing (HeartMuLa, MiniMax Music 3); the recipe says how to call those instead of only song.
|
||||
text = _render("make_short_film", {"idea": "x" * 10_000, "length_seconds": "20"})
|
||||
assert 'refuses mode="song"' in text and 'mode="sfx"' in text
|
||||
assert "always sings, set instrumental=false and pass lyrics" in text
|
||||
# Still whole under the cap with the longest idea.
|
||||
assert "stitch them yourself" in text
|
||||
|
||||
|
||||
def test_every_mention_says_unsloth_studio(stubbed):
|
||||
for name, args in (
|
||||
("make_short_film", {"idea": "x"}),
|
||||
("build_dataset_and_train", {"goal": "x"}),
|
||||
):
|
||||
text = _render(name, args)
|
||||
assert not re.search(r"(?<!Unsloth )Studio", text), name
|
||||
|
||||
|
||||
def _rpc(http, method, params):
|
||||
body = {"jsonrpc": "2.0", "id": 1, "method": method, "params": params}
|
||||
headers = {**MCP_HEADERS, "Authorization": f"Bearer {TEST_TOKEN}"}
|
||||
response = http.post("/mcp/", json = body, headers = headers)
|
||||
assert response.status_code == 200, response.text
|
||||
data = [json.loads(line[5:]) for line in response.text.splitlines() if line.startswith("data:")]
|
||||
return data[-1]["result"]
|
||||
|
||||
|
||||
def test_prompts_over_http(stubbed, monkeypatch):
|
||||
with TestClient(served(create_studio_mcp(), fake_studio(), monkeypatch = monkeypatch)) as http:
|
||||
listed = {p["name"]: p for p in _rpc(http, "prompts/list", {})["prompts"]}
|
||||
got = _rpc(http, "prompts/get", {"name": "make_short_film", "arguments": {"idea": "a fox"}})
|
||||
assert {name: p["title"] for name, p in listed.items()} == TITLES
|
||||
for name, prompt in listed.items():
|
||||
assert [(a["name"], a.get("required", False)) for a in prompt["arguments"]] == ARGUMENTS[
|
||||
name
|
||||
]
|
||||
assert "JSON string" not in json.dumps(listed)
|
||||
text = got["messages"][0]["content"]["text"]
|
||||
assert "unsloth/LTX-small (12.0 GB)" in text and "GPU 0 21.5 GB free" in text
|
||||
|
|
@ -50,7 +50,7 @@ SCHEMA_FACTS = [
|
|||
(
|
||||
"list_models",
|
||||
"parameters/properties",
|
||||
lambda props: set(props) == {"kind", "loaded_only", "model", "workflow"},
|
||||
lambda props: {"kind", "loaded_only", "model", "workflow"} <= set(props),
|
||||
),
|
||||
(
|
||||
"load_model",
|
||||
|
|
|
|||
|
|
@ -492,6 +492,48 @@ def test_a_shared_whisper_id_runs_on_the_engine_holding_it(monkeypatch, audio):
|
|||
assert bodies(studio, "/api/inference/audio/transcribe/source")[0]["engine"] == "gguf"
|
||||
|
||||
|
||||
CITRINET = "audio-cpp/audio.cpp-gguf/Citrinet-ASR-GGUF"
|
||||
ASR_CATALOG = ("GET", "/api/inference/audio/catalog")
|
||||
|
||||
|
||||
def test_a_catalog_asr_row_no_status_names_runs_on_audiocpp(monkeypatch):
|
||||
# list_models(source="catalog") offers it; stt/status names it only once downloaded.
|
||||
routes = {
|
||||
STT_STATUS_ROUTE: {"transformers": stt_state(), "audiocpp": stt_state()},
|
||||
ASR_CATALOG: {"models": [{"id": CITRINET, "task": "asr"}]},
|
||||
SOURCE: _ndjson(COMPLETE),
|
||||
}
|
||||
args = {"audio": {"clip_id": "clip-7"}, "model": CITRINET}
|
||||
result, studio = run_tool(monkeypatch, routes, "transcribe", args)
|
||||
assert result["isError"] is False
|
||||
assert bodies(studio, "/api/inference/audio/transcribe/source")[0]["engine"] == "audiocpp"
|
||||
|
||||
|
||||
def test_a_refused_catalog_asr_row_is_downloaded_on_audiocpp(monkeypatch, fast_polls):
|
||||
missing = f"STT model '{CITRINET}' is not downloaded yet."
|
||||
idle = {"transformers": stt_state(), "audiocpp": stt_state()}
|
||||
done = {
|
||||
"transformers": stt_state(),
|
||||
"audiocpp": stt_state(
|
||||
downloaded = [CITRINET],
|
||||
download = {"downloading": False, "completed_download_ids": ["d1"]},
|
||||
),
|
||||
}
|
||||
routes = {
|
||||
TRANSCRIPTIONS: sequence(
|
||||
openai_error(missing, 409, type = "conflict_error"), JSONResponse(TRANSCRIPT)
|
||||
),
|
||||
STT_STATUS_ROUTE: sequence(idle, idle, done),
|
||||
ASR_CATALOG: {"models": [{"id": CITRINET, "task": "asr"}]},
|
||||
STT_DOWNLOAD: {"downloading": True, "download_id": "d1"},
|
||||
}
|
||||
result, studio = run_tool(monkeypatch, routes, "transcribe", _memo(model = CITRINET))
|
||||
assert result["isError"] is False
|
||||
assert bodies(studio, "/api/inference/audio/stt/download") == [
|
||||
{"model": CITRINET, "engine": "audiocpp"}
|
||||
]
|
||||
|
||||
|
||||
def test_audio_over_25_mib_is_uploaded_then_transcribed(monkeypatch):
|
||||
from studio_mcp.tools import audio
|
||||
|
||||
|
|
|
|||
|
|
@ -343,7 +343,12 @@ def test_media_unload_takes_no_body(monkeypatch, kind, statuses, model):
|
|||
assert unloads == ([b""] if model else [])
|
||||
|
||||
|
||||
def _stt_studio(statuses, *, download = None):
|
||||
def _stt_studio(
|
||||
statuses,
|
||||
*,
|
||||
download = None,
|
||||
extra = None,
|
||||
):
|
||||
return fake_studio(
|
||||
{
|
||||
("GET", f"{STT}/status"): sequence(*statuses),
|
||||
|
|
@ -351,6 +356,7 @@ def _stt_studio(statuses, *, download = None):
|
|||
or {"downloading": True, "model": WHISPER, "download_id": "d1"},
|
||||
("POST", f"{STT}/load"): {"loaded_model": WHISPER, "device": "cuda"},
|
||||
("POST", f"{STT}/unload"): {"loaded_model": None, "device": None},
|
||||
**(extra or {}),
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -416,6 +422,29 @@ def test_a_requested_variant_is_downloaded_even_when_another_is_cached(monkeypat
|
|||
assert bodies(studio, f"{STT}/load") == expected
|
||||
|
||||
|
||||
def test_a_catalog_asr_row_downloads_and_loads_on_audiocpp(monkeypatch, fast_polls):
|
||||
# list_models(source="catalog") offers it; no status list names it before the download.
|
||||
citrinet = "audio-cpp/audio.cpp-gguf/Citrinet-ASR-GGUF"
|
||||
idle = {**stt_state(), "models": []}
|
||||
done = {"downloading": False, "completed_download_ids": ["d1"]}
|
||||
studio = _stt_studio(
|
||||
[
|
||||
{"transformers": stt_state(), "audiocpp": idle},
|
||||
{"transformers": stt_state(), "audiocpp": {**idle, "download": done}},
|
||||
],
|
||||
extra = {
|
||||
("GET", "/api/inference/audio/catalog"): {"models": [{"id": citrinet, "task": "asr"}]}
|
||||
},
|
||||
)
|
||||
result, _studio = run_tool(
|
||||
monkeypatch, studio, "load_model", {"model": citrinet, "kind": "stt"}
|
||||
)
|
||||
assert result["isError"] is False
|
||||
expected = [{"model": citrinet, "engine": "audiocpp"}]
|
||||
assert bodies(studio, f"{STT}/download") == expected
|
||||
assert bodies(studio, f"{STT}/load") == expected
|
||||
|
||||
|
||||
def test_a_failed_download_never_loads(monkeypatch, fast_polls):
|
||||
failed = {"downloading": False, "error": "401 gated repo"}
|
||||
studio = _stt_studio(
|
||||
|
|
|
|||
|
|
@ -1,7 +1,10 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
from .mcp_harness import run_tool
|
||||
|
||||
|
|
@ -57,10 +60,46 @@ CONFIG = {
|
|||
},
|
||||
}
|
||||
|
||||
MEDIA_ROUTE = "/api/inference/media/catalog"
|
||||
AUDIO_ROUTE = "/api/inference/audio/catalog"
|
||||
HARDWARE_ROUTE = "/api/train/hardware"
|
||||
TI2V = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
|
||||
TI2V_GGUF = "unsloth/Wan2.2-TI2V-5B-GGUF"
|
||||
MEDIA = {
|
||||
"models": [
|
||||
{"id": "unsloth/Z-Image-Turbo", "kind": "image", "size_bytes": 20500000000},
|
||||
{"id": "unsloth/stable-diffusion-xl-base-1.0", "kind": "image", "family": "sdxl"},
|
||||
{"id": TI2V, "kind": "video", "display_name": "TI2V", "size_bytes": 24200000000},
|
||||
{"id": TI2V_GGUF, "kind": "video", "family": "wan2.2-ti2v-5b", "low_vram": True},
|
||||
]
|
||||
}
|
||||
# Rows as C2's /audio/catalog serves them, one of each task.
|
||||
AUDIO = {
|
||||
"source": "hub",
|
||||
"models": [
|
||||
{"id": "audio-cpp/audio.cpp-gguf/Kokoro-82M-GGUF", "task": "tts", "workflows": ["speak"]},
|
||||
{"id": "Jmica/Drama-3-GGUF", "task": "tts", "workflows": ["speak", "clone"]},
|
||||
{"id": "audio-cpp/audio.cpp-gguf/ACE-Step-GGUF", "task": "music", "workflows": ["music"]},
|
||||
{"id": "audio-cpp/audio.cpp-gguf/Whisper-GGUF", "task": "asr", "size_bytes": 900000000},
|
||||
{"id": "audio-cpp/audio.cpp-gguf/Experimental-GGUF", "task": "", "workflows": []},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def gpus(*sizes_gib, used = 0.0):
|
||||
devices = [
|
||||
{"index": i, "vram_total_gb": gib, "vram_used_gb": used} for i, gib in enumerate(sizes_gib)
|
||||
]
|
||||
return {"available": True, "backend": "cuda", "devices": devices}
|
||||
|
||||
|
||||
PAYLOADS = {
|
||||
("GET", "/v1/models"): CATALOG,
|
||||
("GET", "/api/inference/loaded-models"): LOADED,
|
||||
("GET", "/api/models/config/unsloth/Qwen3-0.6B"): CONFIG,
|
||||
("GET", MEDIA_ROUTE): MEDIA,
|
||||
("GET", AUDIO_ROUTE): AUDIO,
|
||||
("GET", HARDWARE_ROUTE): gpus(24),
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -102,6 +141,10 @@ def test_entries_carry_only_the_allowlisted_fields(monkeypatch):
|
|||
"display_name": "Llama 3.2 1B",
|
||||
"quant": "Q4_K_M",
|
||||
"context_length": 8192,
|
||||
"size_bytes": None,
|
||||
"low_vram": None,
|
||||
"fits_now": None,
|
||||
"source": None,
|
||||
"audio_workflows": None,
|
||||
}
|
||||
assert models["unsloth/orpheus-3b"]["audio_workflows"] == ["speak", "clone"]
|
||||
|
|
@ -174,3 +217,131 @@ def test_the_config_request_carries_the_hub_token_only_when_given(monkeypatch, t
|
|||
calls = {path: h for _m, path, h, _b in studio.state.calls}
|
||||
assert calls["/api/models/config/unsloth/Qwen3-0.6B"].get("x-unsloth-hf-token") == token
|
||||
assert "x-unsloth-hf-token" not in calls["/v1/models"]
|
||||
|
||||
|
||||
def _catalog(
|
||||
monkeypatch,
|
||||
args = None,
|
||||
routes = PAYLOADS,
|
||||
):
|
||||
result, studio = _list(monkeypatch, {"source": "catalog", **(args or {})}, routes = routes)
|
||||
assert result["isError"] is False, result
|
||||
return result["structuredContent"]["models"], [path for _m, path, _q in studio.state.queries]
|
||||
|
||||
|
||||
def test_catalog_lists_media_rows_then_audio_rows_one_to_one(monkeypatch):
|
||||
models, paths = _catalog(monkeypatch)
|
||||
media, audio = models[: len(MEDIA["models"])], models[len(MEDIA["models"]) :]
|
||||
assert [m["id"] for m in media] == [row["id"] for row in MEDIA["models"]]
|
||||
assert [m["kind"] for m in media] == ["image", "image", "video", "video"]
|
||||
assert [m["low_vram"] for m in media] == [None, None, None, True]
|
||||
assert [(m["id"], m["audio_workflows"]) for m in audio] == [
|
||||
(row["id"], row.get("workflows", [])) for row in AUDIO["models"]
|
||||
]
|
||||
assert [m["kind"] for m in audio] == ["tts", "tts", "audio", "stt", "audio"]
|
||||
assert {m["source"] for m in models} == {"catalog"}
|
||||
assert models[2]["display_name"] == "TI2V" and models[2]["size_bytes"] == 24200000000
|
||||
assert sorted(paths) == sorted([MEDIA_ROUTE, AUDIO_ROUTE, HARDWARE_ROUTE])
|
||||
|
||||
|
||||
# A 24 GiB card holds 25.77 decimal GB; one load lands on one GPU, so two cards never add up.
|
||||
@pytest.mark.parametrize(
|
||||
"hardware,size,fits",
|
||||
[
|
||||
(gpus(24), 24200000000, True),
|
||||
(gpus(12), 24200000000, False),
|
||||
(gpus(24), 25500000000, True),
|
||||
(gpus(24), 25900000000, False),
|
||||
(gpus(12, 12), 20000000000, False),
|
||||
(gpus(12, 24), 20000000000, True),
|
||||
(gpus(80, used = 70.0), 20000000000, False),
|
||||
({"available": False, "devices": []}, 20000000000, None),
|
||||
({"devices": [{"index": 0, "vram_total_gb": 24}]}, 20000000000, None),
|
||||
(gpus(24), None, None),
|
||||
],
|
||||
)
|
||||
def test_fits_now_reads_the_most_free_memory_on_one_gpu(monkeypatch, hardware, size, fits):
|
||||
row = {"id": TI2V, "kind": "video", "size_bytes": size}
|
||||
routes = {**PAYLOADS, ("GET", MEDIA_ROUTE): {"models": [row]}}
|
||||
routes[("GET", HARDWARE_ROUTE)] = hardware
|
||||
models, _paths = _catalog(monkeypatch, {"kind": "video"}, routes = routes)
|
||||
assert [m["fits_now"] for m in models] == [fits]
|
||||
|
||||
|
||||
def test_fits_now_is_unknown_when_the_hardware_route_fails(monkeypatch):
|
||||
routes = {**PAYLOADS, ("GET", HARDWARE_ROUTE): JSONResponse({"detail": "x"}, status_code = 500)}
|
||||
models, _paths = _catalog(monkeypatch, routes = routes)
|
||||
assert {m["fits_now"] for m in models} == {None}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"args,ids,paths",
|
||||
[
|
||||
({"kind": "video"}, [TI2V, TI2V_GGUF], [MEDIA_ROUTE]),
|
||||
({"kind": "stt"}, ["audio-cpp/audio.cpp-gguf/Whisper-GGUF"], [AUDIO_ROUTE]),
|
||||
({"workflow": "clone"}, ["Jmica/Drama-3-GGUF"], [AUDIO_ROUTE]),
|
||||
({"kind": "llm"}, [], []),
|
||||
],
|
||||
)
|
||||
def test_catalog_fetches_only_what_the_filters_need(monkeypatch, args, ids, paths):
|
||||
models, called = _catalog(monkeypatch, args)
|
||||
assert [m["id"] for m in models] == ids
|
||||
assert [p for p in called if p != HARDWARE_ROUTE] == paths
|
||||
assert (HARDWARE_ROUTE in called) == bool(paths)
|
||||
|
||||
|
||||
def test_catalog_model_filter_keeps_training_defaults(monkeypatch):
|
||||
routes = {**PAYLOADS, ("GET", f"/api/models/config/{TI2V_GGUF.lower()}"): CONFIG}
|
||||
args = {"source": "catalog", "model": TI2V_GGUF.lower()}
|
||||
result, _studio = _list(monkeypatch, args, routes = routes)
|
||||
assert [m["id"] for m in result["structuredContent"]["models"]] == [TI2V_GGUF]
|
||||
assert result["structuredContent"]["training_defaults"]["lora_r"] == 16
|
||||
|
||||
|
||||
def test_catalog_refuses_loaded_only(monkeypatch):
|
||||
result, studio = _list(monkeypatch, {"source": "catalog", "loaded_only": True})
|
||||
assert result["isError"] is True
|
||||
assert "use the default source" in result["content"][0]["text"]
|
||||
assert studio.state.queries == []
|
||||
|
||||
|
||||
def test_a_failed_catalog_part_names_itself_instead_of_a_partial_list(monkeypatch):
|
||||
routes = {**PAYLOADS, ("GET", AUDIO_ROUTE): JSONResponse({"detail": "down"}, status_code = 503)}
|
||||
result, _studio = _list(monkeypatch, {"source": "catalog"}, routes = routes)
|
||||
assert result["isError"] is True
|
||||
text = result["content"][0]["text"]
|
||||
assert "could not list the audio catalog" in text and 'Pass kind="image"' in text
|
||||
|
||||
|
||||
# Negative control: the default source reads only /v1/models, and the new fields stay null.
|
||||
def test_the_default_source_is_unchanged(monkeypatch):
|
||||
result, studio = _list(monkeypatch)
|
||||
explicit, _studio = _list(monkeypatch, {"source": "local"})
|
||||
models = result["structuredContent"]["models"]
|
||||
assert explicit["structuredContent"] == result["structuredContent"]
|
||||
assert [m["id"] for m in models] == [row["id"] for row in CATALOG["data"]]
|
||||
assert [path for _m, path, _q in studio.state.queries] == ["/v1/models"]
|
||||
new = ("size_bytes", "low_vram", "fits_now", "source")
|
||||
assert {m[key] for m in models for key in new} == {None}
|
||||
|
||||
|
||||
def test_source_defaults_to_local_in_the_schema():
|
||||
from mcp_server import create_studio_mcp
|
||||
|
||||
tools = {tool.name: tool for tool in asyncio.run(create_studio_mcp().list_tools())}
|
||||
source = tools["list_models"].parameters["properties"]["source"]
|
||||
assert source["default"] == "local" and source["enum"] == ["local", "catalog"]
|
||||
|
||||
|
||||
def test_catalog_skips_audio_rows_the_catalog_refuses(monkeypatch):
|
||||
refused = {
|
||||
"id": "audio-cpp/audio.cpp-gguf/Unknown-GGUF",
|
||||
"task": "",
|
||||
"workflows": [],
|
||||
"size_bytes": 1000,
|
||||
"unsupported_reason": "Unsloth Studio does not support this model family yet.",
|
||||
}
|
||||
routes = {**PAYLOADS, ("GET", AUDIO_ROUTE): {"models": [*AUDIO["models"], refused]}}
|
||||
models, _paths = _catalog(monkeypatch, {"kind": "audio"}, routes = routes)
|
||||
assert refused["id"] not in [m["id"] for m in models]
|
||||
assert "audio-cpp/audio.cpp-gguf/ACE-Step-GGUF" in [m["id"] for m in models]
|
||||
|
|
|
|||
155
studio/backend/tests/test_media_catalog.py
Normal file
155
studio/backend/tests/test_media_catalog.py
Normal file
|
|
@ -0,0 +1,155 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from auth.authentication import get_current_subject
|
||||
from core.inference import diffusion_families, media_catalog, video_families
|
||||
from core.inference.diffusion_inference_info import family_inference_infos
|
||||
from hub.services.models import account_access
|
||||
from utils.api_errors import install_api_error_handlers
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def selectable(monkeypatch):
|
||||
# Selectability needs diffusers; the registry decides the rows here.
|
||||
hidden = set()
|
||||
|
||||
def pick(fam):
|
||||
return fam.name not in hidden
|
||||
|
||||
monkeypatch.setattr(diffusion_families, "family_selectable", pick)
|
||||
monkeypatch.setattr(media_catalog, "family_selectable", pick)
|
||||
return hidden
|
||||
|
||||
|
||||
def _rows():
|
||||
return [row.model_dump() for row in media_catalog.media_catalog_payload().models]
|
||||
|
||||
|
||||
def test_every_selectable_family_is_a_row_sized_from_its_registry(selectable):
|
||||
rows = _rows()
|
||||
bf16 = {i["family"]: i["estimated_resident_gb"]["bf16"] for i in family_inference_infos()}
|
||||
images = [row for row in rows if row["kind"] == "image"]
|
||||
assert [row["family"] for row in images] == [f.name for f in diffusion_families._FAMILIES]
|
||||
for row in images:
|
||||
expected = bf16.get(row["family"])
|
||||
assert row["size_bytes"] == (None if expected is None else round(expected * 1e9))
|
||||
assert row["low_vram"] is False
|
||||
videos = [row for row in rows if row["kind"] == "video" and not row["low_vram"]]
|
||||
assert [row["family"] for row in videos] == [f.name for f in video_families._FAMILIES]
|
||||
sizes = {f.name: f.bf16_components_gb for f in video_families._FAMILIES}
|
||||
for row in videos:
|
||||
assert row["size_bytes"] == round(sum(sizes[row["family"]]) * 1e9)
|
||||
assert {"id", "kind", "family", "display_name", "size_bytes", "low_vram"} == set(rows[0])
|
||||
# Images first, then videos, each in registry order.
|
||||
kinds = [row["kind"] for row in rows]
|
||||
assert kinds == sorted(kinds)
|
||||
|
||||
|
||||
def test_video_gguf_repos_are_low_vram_rows_without_a_size(selectable):
|
||||
ggufs = [row for row in _rows() if row["low_vram"]]
|
||||
expected = [(f.gguf_repo, f.name) for f in video_families._FAMILIES if f.gguf_repo]
|
||||
assert [(row["id"], row["family"]) for row in ggufs] == expected
|
||||
assert {row["size_bytes"] for row in ggufs} == {None}
|
||||
assert {row["kind"] for row in ggufs} == {"video"}
|
||||
|
||||
|
||||
def test_mirror_ids_resolve_back_to_their_family(selectable):
|
||||
rows = _rows()
|
||||
by_id = {row["id"]: row for row in rows}
|
||||
assert by_id["unsloth/FLUX.1-schnell"]["family"] == "flux.1"
|
||||
assert by_id["unsloth/Wan2.2-TI2V-5B-GGUF"]["display_name"] == "Wan2.2-TI2V-5B-GGUF"
|
||||
for row in rows:
|
||||
detect = diffusion_families.detect_family
|
||||
if row["kind"] == "video":
|
||||
detect = video_families.detect_video_family
|
||||
assert detect(row["id"]).name == row["family"], row["id"]
|
||||
|
||||
|
||||
def test_apple_silicon_drops_the_video_families_it_cannot_run(selectable):
|
||||
everywhere = {(r["family"], r["id"]) for r in _rows() if r["kind"] == "video"}
|
||||
on_mps = media_catalog.media_catalog_payload(device = "mps").models
|
||||
dropped = {fam.name for fam in video_families._FAMILIES if fam.modular_workflow}
|
||||
assert dropped and {(r.family, r.id) for r in on_mps if r.kind == "video"} == {
|
||||
row for row in everywhere if row[0] not in dropped
|
||||
}
|
||||
|
||||
|
||||
def test_the_route_asks_for_the_device(client, monkeypatch):
|
||||
from core.inference import diffusion_device
|
||||
|
||||
seen = []
|
||||
target = type("Target", (), {"device": "mps"})()
|
||||
monkeypatch.setattr(diffusion_device, "resolve_diffusion_device_target", lambda: target)
|
||||
monkeypatch.setattr(
|
||||
media_catalog, "media_catalog_payload", lambda device = None: seen.append(device) or PAYLOAD
|
||||
)
|
||||
assert client.get("/api/inference/media/catalog").status_code == 200
|
||||
assert seen == ["mps"]
|
||||
|
||||
|
||||
def test_an_unselectable_family_is_left_out(selectable):
|
||||
selectable.update({"flux.1", "wan2.2-ti2v-5b"})
|
||||
families = {row["family"] for row in _rows()}
|
||||
assert "flux.1" not in families and "wan2.2-ti2v-5b" not in families
|
||||
assert "z-image" in families and "ltx-2" in families
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client():
|
||||
from routes import inference
|
||||
|
||||
app = FastAPI()
|
||||
install_api_error_handlers(app)
|
||||
app.dependency_overrides[get_current_subject] = lambda: "tester"
|
||||
app.include_router(inference.router, prefix = "/api/inference")
|
||||
app.include_router(inference.studio_router, prefix = "/api/inference")
|
||||
app.include_router(inference.router, prefix = "/v1")
|
||||
yield TestClient(app)
|
||||
|
||||
|
||||
PAYLOAD = media_catalog.MediaCatalogResponse(
|
||||
models = [
|
||||
{"id": "unsloth/Z-Image-Turbo", "kind": "image", "family": "z-image", "display_name": "Z"},
|
||||
{"id": "unsloth/FLUX.2-dev", "kind": "image", "family": "flux.2-dev", "display_name": "F"},
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_the_route_serves_every_row_to_the_owner(client, monkeypatch):
|
||||
monkeypatch.setattr(media_catalog, "media_catalog_payload", lambda device = None: PAYLOAD)
|
||||
response = client.get("/api/inference/media/catalog")
|
||||
assert response.status_code == 200
|
||||
assert media_catalog.MediaCatalogResponse.model_validate(response.json()) == PAYLOAD
|
||||
assert client.get("/v1/media/catalog").status_code == 404
|
||||
|
||||
|
||||
def test_a_managed_account_sees_only_the_rows_it_may_use(client, monkeypatch):
|
||||
monkeypatch.setattr(media_catalog, "media_catalog_payload", lambda device = None: PAYLOAD)
|
||||
seen = []
|
||||
monkeypatch.setattr(account_access, "managed_account", lambda: True)
|
||||
monkeypatch.setattr(account_access, "model_grants", lambda: set())
|
||||
monkeypatch.setattr(account_access, "_warm_public_repos", lambda *a: None)
|
||||
monkeypatch.setattr(
|
||||
account_access,
|
||||
"model_visible",
|
||||
lambda reference, **_k: seen.append(reference) or "FLUX" not in reference,
|
||||
)
|
||||
response = client.get("/api/inference/media/catalog")
|
||||
assert [row["id"] for row in response.json()["models"]] == ["unsloth/Z-Image-Turbo"]
|
||||
assert seen == ["unsloth/Z-Image-Turbo", "unsloth/FLUX.2-dev"]
|
||||
assert len(PAYLOAD.models) == 2
|
||||
|
||||
|
||||
def test_a_builder_failure_answers_503_without_its_text(client, monkeypatch):
|
||||
def build(device = None):
|
||||
raise RuntimeError("/home/me/.cache/huggingface exploded")
|
||||
|
||||
monkeypatch.setattr(media_catalog, "media_catalog_payload", build)
|
||||
response = client.get("/api/inference/media/catalog")
|
||||
assert response.status_code == 503
|
||||
assert "Unsloth Studio could not list image and video models right now." in response.text
|
||||
assert "exploded" not in response.text
|
||||
|
|
@ -101,6 +101,7 @@ else:
|
|||
from unsloth_cli.commands.export import export, list_checkpoints
|
||||
from unsloth_cli.commands.eval import evaluate as eval_command
|
||||
from unsloth_cli.commands.doctor import doctor
|
||||
from unsloth_cli.commands.mcp import mcp_app
|
||||
from unsloth_cli.commands.studio import (
|
||||
run as studio_run,
|
||||
studio_app,
|
||||
|
|
@ -243,6 +244,9 @@ if not _windows_studio_mutation_entry:
|
|||
help = "Start a coding agent (Claude, Codex, OpenClaw, OpenCode, Hermes, Pi, dsh, Vibe) "
|
||||
"against Unsloth.",
|
||||
)
|
||||
app.add_typer(
|
||||
mcp_app, name = "mcp", help = "Connect coding agents to the Unsloth Studio MCP server."
|
||||
)
|
||||
# backwards-compatible hidden alias: `unsloth connect` routes to `unsloth start`.
|
||||
app.add_typer(
|
||||
start_app,
|
||||
|
|
|
|||
|
|
@ -475,11 +475,23 @@ def _load_pynvml():
|
|||
|
||||
|
||||
def doctor(
|
||||
ctx: typer.Context,
|
||||
gpus: bool = typer.Option(
|
||||
True,
|
||||
"--gpus/--no-gpus",
|
||||
help = "Check GPUs, interconnect and a multi-GPU all_reduce (the default check).",
|
||||
),
|
||||
mcp: bool = typer.Option(
|
||||
False,
|
||||
"--mcp",
|
||||
help = "Check the Unsloth Studio MCP server, the API key and agent timeouts. "
|
||||
"Skips the GPU checks unless --gpus is also given.",
|
||||
),
|
||||
env_var: str = typer.Option(
|
||||
"UNSLOTH_API_KEY",
|
||||
"--env-var",
|
||||
help = "With --mcp: the environment variable that holds the Unsloth Studio API key.",
|
||||
),
|
||||
timeout: float = typer.Option(
|
||||
60.0,
|
||||
"--timeout",
|
||||
|
|
@ -487,9 +499,27 @@ def doctor(
|
|||
),
|
||||
):
|
||||
"""Preflight checks for this machine. Exit status 1 means a problem was found."""
|
||||
if not gpus:
|
||||
typer.echo("Nothing to check.")
|
||||
raise typer.Exit()
|
||||
# Plain --mcp skips the GPU checks; an explicit --gpus keeps them.
|
||||
# By name: newer typer ships its own copy of click's ParameterSource.
|
||||
explicit = getattr(ctx.get_parameter_source("gpus"), "name", "") != "DEFAULT"
|
||||
run_gpus = gpus if explicit else not mcp
|
||||
code = 0
|
||||
if mcp:
|
||||
from unsloth_cli.commands.mcp import mcp_doctor
|
||||
code = mcp_doctor(env_var)
|
||||
if not run_gpus:
|
||||
if not mcp:
|
||||
typer.echo("Nothing to check.")
|
||||
raise typer.Exit(code)
|
||||
if mcp:
|
||||
typer.echo("")
|
||||
try:
|
||||
_gpu_checks(timeout)
|
||||
except typer.Exit as e:
|
||||
raise typer.Exit(max(e.exit_code, code))
|
||||
|
||||
|
||||
def _gpu_checks(timeout: float):
|
||||
try:
|
||||
import torch
|
||||
except ImportError:
|
||||
|
|
|
|||
581
unsloth_cli/commands/mcp.py
Normal file
581
unsloth_cli/commands/mcp.py
Normal file
|
|
@ -0,0 +1,581 @@
|
|||
# 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 mcp add <agent>` and the checks behind `unsloth doctor --mcp`.
|
||||
|
||||
Agent configs only name the environment variable that holds the API key. `mcp add`
|
||||
never reads the key; doctor sends it to the server to test it and never prints it.
|
||||
"""
|
||||
|
||||
import copy
|
||||
import difflib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shlex
|
||||
import subprocess
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Tuple
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import typer
|
||||
|
||||
from unsloth_cli._inference import (
|
||||
_USER_AGENT,
|
||||
find_studio_server,
|
||||
is_loopback_url,
|
||||
urlopen_no_redirect,
|
||||
)
|
||||
from unsloth_cli.commands.start import (
|
||||
_codex_source_home,
|
||||
_fail,
|
||||
_load_app_config,
|
||||
_read_app_file,
|
||||
_which_with_install_dirs,
|
||||
_write_app_file,
|
||||
)
|
||||
|
||||
mcp_app = typer.Typer(
|
||||
help = "Connect coding agents to the Unsloth Studio MCP server.",
|
||||
no_args_is_help = True,
|
||||
context_settings = {"help_option_names": ["-h", "--help"]},
|
||||
)
|
||||
|
||||
_DEFAULT_BASE = "http://127.0.0.1:8888"
|
||||
_ENV_VAR_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
|
||||
_NAME_RE = re.compile(r"^[A-Za-z0-9_-]+$")
|
||||
_TABLE_LINE_RE = re.compile(r"^[ \t]*\[\[?[^\[\]\n]+\]\]?[ \t]*(#[^\n]*)?\r?\n?$")
|
||||
# The longest an Unsloth Studio MCP tool waits before answering (the server's wait cap).
|
||||
# The backend can't be imported from the CLI, so it is repeated here.
|
||||
_MCP_WAIT_CAP_S = 240
|
||||
_CODEX_DEFAULT_TIMEOUT_S = 60
|
||||
_CLAUDE_ADD = ["mcp", "add", "--transport", "http", "--scope", "user"]
|
||||
_INITIALIZE = (
|
||||
b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18",'
|
||||
b'"capabilities":{},"clientInfo":{"name":"unsloth-doctor","version":"1"}}}'
|
||||
)
|
||||
|
||||
|
||||
class Agent(str, Enum):
|
||||
claude_code = "claude-code"
|
||||
codex = "codex"
|
||||
cursor = "cursor"
|
||||
generic = "generic"
|
||||
|
||||
|
||||
def mcp_url(base: str) -> str:
|
||||
base = base.strip().rstrip("/")
|
||||
if base.endswith("/mcp"):
|
||||
base = base[: -len("/mcp")]
|
||||
return base + "/mcp/"
|
||||
|
||||
|
||||
def _bearer(env_var: str, style: str = "${%s}") -> str:
|
||||
return "Bearer " + style % env_var
|
||||
|
||||
|
||||
def claude_entry(url: str, env_var: str) -> dict:
|
||||
return {"type": "http", "url": url, "headers": {"Authorization": _bearer(env_var)}}
|
||||
|
||||
|
||||
def _with_server(existing: dict, name: str, entry: dict) -> dict:
|
||||
data = copy.deepcopy(existing)
|
||||
servers = data.get("mcpServers")
|
||||
if not isinstance(servers, dict):
|
||||
servers = data["mcpServers"] = {}
|
||||
servers[name] = entry
|
||||
return data
|
||||
|
||||
|
||||
def cursor_merge(existing: dict, name: str, url: str, env_var: str) -> dict:
|
||||
entry = {"url": url, "headers": {"Authorization": _bearer(env_var, "${env:%s}")}}
|
||||
return _with_server(existing, name, entry)
|
||||
|
||||
|
||||
def generic_snippet(name: str, url: str, env_var: str) -> dict:
|
||||
return {"mcpServers": {name: claude_entry(url, env_var)}}
|
||||
|
||||
|
||||
def _codex_header(name: str):
|
||||
n = re.escape(name)
|
||||
return re.compile(rf"""^\s*\[\s*mcp_servers\s*\.\s*(?:"{n}"|'{n}'|{n})\s*[\].]""")
|
||||
|
||||
|
||||
def _codex_block(name: str, url: str, env_var: str, timeout: int) -> str:
|
||||
return (
|
||||
f"[mcp_servers.{name}]\n"
|
||||
f"url = {json.dumps(url)}\n"
|
||||
f"bearer_token_env_var = {json.dumps(env_var)}\n"
|
||||
f"tool_timeout_sec = {int(timeout)}\n"
|
||||
)
|
||||
|
||||
|
||||
def _comment_tail(chunk: str) -> str:
|
||||
# Blank and comment lines at the end of a table usually describe the next one: keep them.
|
||||
lines = chunk.splitlines(True)
|
||||
end = len(lines)
|
||||
while end and (not lines[end - 1].strip() or lines[end - 1].lstrip().startswith("#")):
|
||||
end -= 1
|
||||
return "".join(lines[end:])
|
||||
|
||||
|
||||
def codex_merge(existing_text: str, name: str, url: str, env_var: str, timeout: int) -> str:
|
||||
"""Replace the [mcp_servers.NAME] table (and its subtables) in place, or append it."""
|
||||
header = _codex_header(name)
|
||||
block = _codex_block(name, url, env_var, timeout)
|
||||
out, placed = [], False
|
||||
for chunk in re.split(r"(?m)^(?=[ \t]*\[)", existing_text):
|
||||
if not header.match(chunk):
|
||||
out.append(chunk)
|
||||
elif not placed:
|
||||
out.append(block + _comment_tail(chunk))
|
||||
placed = True
|
||||
else:
|
||||
tail = _comment_tail(chunk)
|
||||
out.append(tail if "#" in tail else "")
|
||||
text = "".join(out)
|
||||
if not placed:
|
||||
if text.strip():
|
||||
text = text.rstrip("\n") + "\n\n"
|
||||
else:
|
||||
text = ""
|
||||
text += block
|
||||
return text
|
||||
|
||||
|
||||
def _toml_loads():
|
||||
try:
|
||||
import tomllib # novermin
|
||||
except ImportError:
|
||||
try:
|
||||
import tomli as tomllib
|
||||
except ImportError:
|
||||
return None
|
||||
return tomllib.loads
|
||||
|
||||
|
||||
def codex_check(old_text: str, new_text: str, name: str, expected: dict) -> None:
|
||||
"""Raise ValueError unless the splice changed exactly mcp_servers.NAME."""
|
||||
loads = _toml_loads()
|
||||
if loads is None:
|
||||
# No TOML parser (Python < 3.11 without tomli): refuse any other mention of NAME, and
|
||||
# anything that makes the line-based splice unsafe: a multi-line string can hold a
|
||||
# line that looks like a table header, and a line opening with "[" can be array data.
|
||||
block = _codex_block(name, *expected.values())
|
||||
if name in new_text.replace(block, ""):
|
||||
raise ValueError(f"{name} is also defined in another form")
|
||||
if '"""' in old_text or "'''" in old_text:
|
||||
raise ValueError("it has multi-line strings and no TOML parser is installed")
|
||||
for chunk in re.split(r"(?m)^(?=[ \t]*\[)", old_text)[1:]:
|
||||
if not _TABLE_LINE_RE.match(chunk.splitlines(True)[0]):
|
||||
raise ValueError("a line starting with '[' is not a table header")
|
||||
return
|
||||
try:
|
||||
old, new = loads(old_text), loads(new_text)
|
||||
except Exception as e:
|
||||
raise ValueError(f"not valid TOML after the edit ({e})")
|
||||
if (new.get("mcp_servers") or {}).get(name) != expected:
|
||||
raise ValueError(f"mcp_servers.{name} did not come out as expected")
|
||||
|
||||
def without(doc):
|
||||
doc = dict(doc)
|
||||
servers = doc.get("mcp_servers")
|
||||
if isinstance(servers, dict):
|
||||
servers = {k: v for k, v in servers.items() if k != name}
|
||||
if servers:
|
||||
doc["mcp_servers"] = servers
|
||||
else:
|
||||
doc.pop("mcp_servers")
|
||||
return doc
|
||||
|
||||
if without(old) != without(new):
|
||||
raise ValueError("other settings would change")
|
||||
|
||||
|
||||
def _redact(text: str) -> str:
|
||||
# Last line of defence for text we don't parse, such as claude's own error output.
|
||||
text = re.sub(r"(?i)(bearer\s+)(?!\$\{)\S+", r"\1***", text)
|
||||
return re.sub(r"sk-unsloth-\S+", "sk-unsloth-***", text)
|
||||
|
||||
|
||||
def _diff(path: Path, old_text: Optional[str], new_text: str) -> str:
|
||||
return "".join(
|
||||
difflib.unified_diff(
|
||||
(old_text or "").splitlines(True),
|
||||
new_text.splitlines(True),
|
||||
fromfile = str(path) if old_text is not None else "/dev/null",
|
||||
tofile = str(path),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _show_entry(path: Path, name: str, entry_text: str, replaces: bool) -> None:
|
||||
# Only the new entry, which names an env var: an old one may hold a key pasted in by hand.
|
||||
typer.echo(_diff(path, None, entry_text), nl = False)
|
||||
if replaces:
|
||||
typer.echo(f"This replaces the current {name} entry; its old values are not shown.")
|
||||
|
||||
|
||||
def write_config(
|
||||
path: Path,
|
||||
old_text: Optional[str],
|
||||
new_text: str,
|
||||
*,
|
||||
dry_run: bool,
|
||||
shown: Tuple[str, str, bool],
|
||||
) -> str:
|
||||
"""Write ``new_text``; a dry run prints only ``shown``: (name, new entry text, replaces)."""
|
||||
if old_text == new_text:
|
||||
return "unchanged"
|
||||
if dry_run:
|
||||
_show_entry(path, *shown)
|
||||
return "dry-run"
|
||||
mode = None if old_text is None else path.stat().st_mode & 0o777
|
||||
_write_app_file(path, new_text)
|
||||
if mode is not None and os.name != "nt":
|
||||
os.chmod(path, mode)
|
||||
return "written"
|
||||
|
||||
|
||||
def _read_json_config(path: Path) -> Tuple[Optional[str], dict]:
|
||||
text = _read_app_file(path)
|
||||
try:
|
||||
data = _load_app_config(path, text)
|
||||
except Exception:
|
||||
data = None
|
||||
if not isinstance(data, dict):
|
||||
_fail(
|
||||
f"Couldn't parse {path} as a JSON object, so it was left as is. "
|
||||
"Fix it and re-run, or add the entry by hand (see `unsloth mcp add generic`)."
|
||||
)
|
||||
if not isinstance(data.get("mcpServers", {}), dict):
|
||||
_fail(f'"mcpServers" in {path} is not an object, so the file was left as is.')
|
||||
return text, data
|
||||
|
||||
|
||||
def _json_text(data: dict) -> str:
|
||||
return json.dumps(data, indent = 2, ensure_ascii = False) + "\n"
|
||||
|
||||
|
||||
def _merge_json_file(
|
||||
path: Path, name: str, data: dict, text: Optional[str], new: dict, dry_run: bool
|
||||
) -> str:
|
||||
if text is not None and new == data:
|
||||
return "unchanged"
|
||||
replaces = name in data.get("mcpServers", {})
|
||||
shown = (name, _json_text({name: new["mcpServers"][name]}), replaces)
|
||||
return write_config(path, text, _json_text(new), dry_run = dry_run, shown = shown)
|
||||
|
||||
|
||||
def claude_config_path() -> Path:
|
||||
configured = os.environ.get("CLAUDE_CONFIG_DIR")
|
||||
return (
|
||||
Path(configured).expanduser() / ".claude.json"
|
||||
if configured
|
||||
else Path.home() / ".claude.json"
|
||||
)
|
||||
|
||||
|
||||
def _add_claude(name, url, env_var, dry_run) -> Tuple[str, Path]:
|
||||
path = claude_config_path()
|
||||
text, data = _read_json_config(path)
|
||||
entry = claude_entry(url, env_var)
|
||||
servers = data.get("mcpServers") if isinstance(data.get("mcpServers"), dict) else {}
|
||||
current = servers.get(name)
|
||||
if current == entry:
|
||||
return "unchanged", path
|
||||
exe = _which_with_install_dirs("claude")
|
||||
if exe is None:
|
||||
typer.echo("Claude Code was not found on PATH, so its config file is edited directly.")
|
||||
new = _with_server(data, name, entry)
|
||||
return _merge_json_file(path, name, data, text, new, dry_run), path
|
||||
commands = []
|
||||
if current is not None:
|
||||
commands.append([exe, "mcp", "remove", "--scope", "user", name])
|
||||
commands.append(
|
||||
[exe, *_CLAUDE_ADD, name, url, "--header", f"Authorization: {_bearer(env_var)}"]
|
||||
)
|
||||
if dry_run:
|
||||
_show_entry(path, name, _json_text({name: entry}), current is not None)
|
||||
for argv in commands:
|
||||
typer.echo("Would run: " + shlex.join(argv))
|
||||
return "dry-run", path
|
||||
for argv in commands:
|
||||
removed = current is not None and argv is commands[-1]
|
||||
_run_claude(
|
||||
argv, removed and f"The old {name} entry was already removed; re-run this command."
|
||||
)
|
||||
if _read_json_config(path)[1].get("mcpServers", {}).get(name) != entry:
|
||||
typer.echo(f"Check the result with `claude mcp get {name}`: it is not what was asked.")
|
||||
return "written", path
|
||||
|
||||
|
||||
def _run_claude(argv: List[str], after: Optional[str] = None) -> None:
|
||||
run = dict(stdin = subprocess.DEVNULL, capture_output = True, text = True, timeout = 60)
|
||||
note = f"\n{after}" if after else ""
|
||||
try:
|
||||
done = subprocess.run(argv, **run)
|
||||
except (OSError, subprocess.TimeoutExpired) as e:
|
||||
_fail(f"`claude mcp {argv[2]}` failed: {e}{note}")
|
||||
if done.returncode != 0:
|
||||
output = _redact((done.stderr or done.stdout).strip())
|
||||
_fail(f"`claude mcp {argv[2]}` failed:\n{output}{note}")
|
||||
|
||||
|
||||
def _add_codex(name, url, env_var, timeout, dry_run) -> Tuple[str, Path]:
|
||||
path = _codex_source_home() / "config.toml"
|
||||
text = _read_app_file(path)
|
||||
old = text or ""
|
||||
new = codex_merge(old, name, url, env_var, timeout)
|
||||
expected = {"url": url, "bearer_token_env_var": env_var, "tool_timeout_sec": int(timeout)}
|
||||
try:
|
||||
codex_check(old, new, name, expected)
|
||||
except ValueError as e:
|
||||
_fail(
|
||||
f"Couldn't update {path} safely ({e}), so it was left as is. Add this table by hand:\n\n"
|
||||
+ _codex_block(name, url, env_var, timeout)
|
||||
)
|
||||
chunks = re.split(r"(?m)^(?=[ \t]*\[)", old)
|
||||
replaces = any(_codex_header(name).match(chunk) for chunk in chunks)
|
||||
shown = (name, _codex_block(name, url, env_var, timeout), replaces)
|
||||
return write_config(path, text, new, dry_run = dry_run, shown = shown), path
|
||||
|
||||
|
||||
def _add_json(path: Path, name: str, merge, dry_run) -> Tuple[str, Path]:
|
||||
text, data = _read_json_config(path)
|
||||
return _merge_json_file(path, name, data, text, merge(data), dry_run), path
|
||||
|
||||
|
||||
def _resolve_base(url: Optional[str]) -> str:
|
||||
return url or os.environ.get("UNSLOTH_STUDIO_URL") or find_studio_server() or _DEFAULT_BASE
|
||||
|
||||
|
||||
@mcp_app.command("add")
|
||||
def add(
|
||||
agent: Agent = typer.Argument(..., help = "The agent to set up."),
|
||||
url: Optional[str] = typer.Option(None, "--url", help = "Unsloth Studio address."),
|
||||
env_var: str = typer.Option("UNSLOTH_API_KEY", "--env-var", help = "Env var with the key."),
|
||||
name: str = typer.Option("unsloth-studio", "--name", help = "Server name in the config."),
|
||||
timeout: int = typer.Option(300, "--timeout", min = 240, help = "Codex tool timeout (s)."),
|
||||
output: Optional[Path] = typer.Option(None, "--output", help = "generic: merge into this file."),
|
||||
dry_run: bool = typer.Option(False, "--dry-run", help = "Show the change; write nothing."),
|
||||
) -> None:
|
||||
"""Add the Unsloth Studio MCP server to an agent's config. The API key stays in an env var."""
|
||||
if not _ENV_VAR_RE.match(env_var):
|
||||
_fail(f"--env-var {env_var!r} is not a valid environment variable name.")
|
||||
if not _NAME_RE.match(name):
|
||||
_fail("--name may only use letters, digits, '-' and '_'.")
|
||||
base = _resolve_base(url)
|
||||
parsed = urlparse(base.strip())
|
||||
if parsed.scheme not in ("http", "https") or not parsed.netloc:
|
||||
_fail(f"{base!r} is not an http(s) URL.")
|
||||
target = mcp_url(base)
|
||||
|
||||
if agent is Agent.generic and output is None:
|
||||
typer.echo(_json_text(generic_snippet(name, target, env_var)), nl = False)
|
||||
typer.echo(f"Replace ${{{env_var}}} with your agent's syntax for environment variables.")
|
||||
state = None
|
||||
elif agent is Agent.claude_code:
|
||||
state, path = _add_claude(name, target, env_var, dry_run)
|
||||
elif agent is Agent.codex:
|
||||
state, path = _add_codex(name, target, env_var, timeout, dry_run)
|
||||
elif agent is Agent.cursor:
|
||||
path = Path.home() / ".cursor" / "mcp.json"
|
||||
state, _ = _add_json(path, name, lambda d: cursor_merge(d, name, target, env_var), dry_run)
|
||||
else:
|
||||
if output.suffix.lower() not in (".json", ".jsonc"):
|
||||
_fail("--output must be a .json file; for another format, paste the printed entry.")
|
||||
entry = claude_entry(target, env_var)
|
||||
state, _ = _add_json(output, name, lambda d: _with_server(d, name, entry), dry_run)
|
||||
path = output
|
||||
|
||||
if state == "written":
|
||||
typer.echo(f"Unsloth Studio MCP set up for {agent.value} in {path}")
|
||||
elif state == "unchanged":
|
||||
typer.echo(f"Unsloth Studio MCP for {agent.value} is already set up in {path} (unchanged).")
|
||||
elif state == "dry-run":
|
||||
typer.echo("Dry run: nothing was written.")
|
||||
typer.echo(f"Set {env_var} to an Unsloth Studio API key (Settings → API) where it starts.")
|
||||
if not os.environ.get(env_var):
|
||||
typer.echo(f"{env_var} is not set in this shell.")
|
||||
if agent is Agent.claude_code:
|
||||
typer.echo("For long jobs, also set MCP_TOOL_TIMEOUT=300000 where Claude Code starts.")
|
||||
|
||||
|
||||
# --- unsloth doctor --mcp ---
|
||||
|
||||
|
||||
@dataclass
|
||||
class McpFacts:
|
||||
env_var: str = "UNSLOTH_API_KEY"
|
||||
base: Optional[str] = None
|
||||
key: str = "missing" # missing | unsent | valid | invalid | error
|
||||
key_detail: str = ""
|
||||
mcp: str = "unknown" # on | off | absent | refused | error | unknown
|
||||
mcp_detail: str = ""
|
||||
timeouts: List[Tuple[str, float, str]] = field(default_factory = list) # (label, s, fix)
|
||||
|
||||
|
||||
def _probe(url: str, key: Optional[str], timeout: float, data: Optional[bytes]):
|
||||
"""(status, detail); status is None when nothing answered. POSTs when data is given."""
|
||||
headers = {
|
||||
"Accept": "application/json, text/event-stream",
|
||||
"Content-Type": "application/json",
|
||||
"User-Agent": _USER_AGENT,
|
||||
}
|
||||
if key:
|
||||
headers["Authorization"] = f"Bearer {key}"
|
||||
try:
|
||||
with urlopen_no_redirect(urllib.request.Request(url, data, headers), timeout) as response:
|
||||
return response.status, ""
|
||||
except urllib.error.HTTPError as e:
|
||||
try:
|
||||
detail = str(json.loads(e.read(65536) or b"{}").get("detail") or e.reason)
|
||||
except Exception:
|
||||
detail = str(e.reason)
|
||||
return e.code, _scrub(detail, key)
|
||||
except Exception as e:
|
||||
# http.client quotes a header value it refuses (a key with a newline) in its error.
|
||||
return None, _scrub(str(getattr(e, "reason", None) or e), key)
|
||||
|
||||
|
||||
def _scrub(detail: str, key: Optional[str]) -> str:
|
||||
for secret in (key, (key or "").strip(), repr(key or "")[1:-1]):
|
||||
if secret:
|
||||
detail = detail.replace(secret, "***")
|
||||
return _redact(detail)
|
||||
|
||||
|
||||
def _ours(name: str, key: str, entry, urls: Tuple[str, ...]) -> bool:
|
||||
if not isinstance(entry, dict):
|
||||
return False
|
||||
url = str(entry.get("url", ""))
|
||||
return key == name or (url.startswith(("http://", "https://")) and mcp_url(url) in urls)
|
||||
|
||||
|
||||
def _codex_timeouts(name: str, urls: Tuple[str, ...] = ()) -> List[Tuple[str, float, str]]:
|
||||
loads = _toml_loads()
|
||||
try:
|
||||
text = (_codex_source_home() / "config.toml").read_text(encoding = "utf-8")
|
||||
servers = loads(text).get("mcp_servers") or {}
|
||||
except (Exception, typer.Exit):
|
||||
return []
|
||||
fix = "Run `unsloth mcp add codex` to set tool_timeout_sec = 300."
|
||||
found = [
|
||||
(key, entry.get("tool_timeout_sec", _CODEX_DEFAULT_TIMEOUT_S))
|
||||
for key, entry in servers.items()
|
||||
if _ours(name, key, entry, urls)
|
||||
]
|
||||
return [(f"Codex {key}", s, fix) for key, s in found if isinstance(s, (int, float))]
|
||||
|
||||
|
||||
def _claude_timeouts(name: str, urls: Tuple[str, ...] = ()) -> List[Tuple[str, float, str]]:
|
||||
root = os.environ.get("CLAUDE_CONFIG_DIR")
|
||||
settings = (Path(root).expanduser() if root else Path.home() / ".claude") / "settings.json"
|
||||
try:
|
||||
servers = json.loads(claude_config_path().read_text(encoding = "utf-8"))["mcpServers"]
|
||||
if not any(_ours(name, k, v, urls) for k, v in servers.items()):
|
||||
return []
|
||||
value = os.environ.get("MCP_TOOL_TIMEOUT")
|
||||
if value is None:
|
||||
value = json.loads(settings.read_text(encoding = "utf-8"))["env"]["MCP_TOOL_TIMEOUT"]
|
||||
return [("Claude Code", float(value) / 1000, "Set MCP_TOOL_TIMEOUT=300000.")]
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
def gather_mcp_facts(env_var: str = "UNSLOTH_API_KEY", timeout: float = 5.0) -> McpFacts:
|
||||
facts = McpFacts(env_var = env_var)
|
||||
facts.base = find_studio_server(timeout = min(timeout, 3.0))
|
||||
urls = (mcp_url(facts.base),) if facts.base else ()
|
||||
facts.timeouts = _codex_timeouts("unsloth-studio", urls) + _claude_timeouts(
|
||||
"unsloth-studio", urls
|
||||
)
|
||||
if facts.base is None:
|
||||
return facts
|
||||
key = os.environ.get(env_var) or None
|
||||
# Like the other auto flows, the key only goes to loopback or an explicit UNSLOTH_STUDIO_URL.
|
||||
if key and not (is_loopback_url(facts.base) or os.environ.get("UNSLOTH_STUDIO_URL")):
|
||||
facts.key, key = "unsent", None
|
||||
elif key:
|
||||
status, detail = _probe(f"{facts.base}/api/inference/loaded-models", key, timeout, None)
|
||||
if status is not None and 200 <= status < 300:
|
||||
facts.key = "valid"
|
||||
else:
|
||||
facts.key = "invalid" if status in (401, 403) else "error"
|
||||
facts.key_detail, key = f"{status or ''} {detail}".strip(), None
|
||||
status, detail = _probe(f"{facts.base}/mcp/", key, timeout, _INITIALIZE)
|
||||
if status == 404:
|
||||
facts.mcp = "off"
|
||||
elif status == 405: # an Unsloth Studio from before MCP: its page catch-all takes GET only
|
||||
facts.mcp = "absent"
|
||||
elif status == 401:
|
||||
facts.mcp, facts.mcp_detail = "refused", detail
|
||||
elif status is not None and 200 <= status < 300:
|
||||
facts.mcp = "on"
|
||||
else:
|
||||
facts.mcp, facts.mcp_detail = "error", f"{status or ''} {detail}".strip()
|
||||
return facts
|
||||
|
||||
|
||||
def mcp_verdict(facts: McpFacts) -> Tuple[bool, List[str], List[str]]:
|
||||
"""(healthy, hints, advice). Advice never makes the verdict fail."""
|
||||
var, hints = facts.env_var, []
|
||||
key_hints = {
|
||||
"missing": f"Set {var} to an Unsloth Studio API key (Settings → API).",
|
||||
"unsent": f"{facts.base} is not on this machine; set UNSLOTH_STUDIO_URL to test the key.",
|
||||
"invalid": f"Unsloth Studio refused the key in {var}. Create one in Settings → API.",
|
||||
"error": f"Unsloth Studio failed while checking the key: {facts.key_detail}",
|
||||
}
|
||||
mcp_hints = {
|
||||
"off": "Turn on Settings → API → Allow agent connections, or start Unsloth Studio "
|
||||
"with UNSLOTH_STUDIO_ENABLE_MCP=1.",
|
||||
"absent": "This Unsloth Studio has no MCP server. Update it: unsloth studio update",
|
||||
"refused": f"The MCP server refused the key: {facts.mcp_detail}",
|
||||
"error": f"The MCP server did not answer properly: {facts.mcp_detail}",
|
||||
}
|
||||
if facts.base is None:
|
||||
where = os.environ.get("UNSLOTH_STUDIO_URL") or _DEFAULT_BASE
|
||||
hints.append(
|
||||
f"Unsloth Studio is not reachable at {where}. Start it, or set UNSLOTH_STUDIO_URL."
|
||||
)
|
||||
else:
|
||||
hints += [key_hints[facts.key]] if facts.key in key_hints else []
|
||||
# A missing or refused key already explains a refused MCP call.
|
||||
if facts.mcp in mcp_hints and not (facts.mcp == "refused" and facts.key != "valid"):
|
||||
hints.append(mcp_hints[facts.mcp])
|
||||
advice = [
|
||||
f"{label} gives a tool {s:g} s; Unsloth Studio tools can wait {_MCP_WAIT_CAP_S} s. {fix}"
|
||||
for label, s, fix in facts.timeouts
|
||||
if s < _MCP_WAIT_CAP_S
|
||||
]
|
||||
return not hints, hints, advice
|
||||
|
||||
|
||||
def mcp_doctor(env_var: str = "UNSLOTH_API_KEY", timeout: float = 5.0) -> int:
|
||||
facts = gather_mcp_facts(env_var, timeout)
|
||||
key = "set" if os.environ.get(env_var) else "not set"
|
||||
checked = {
|
||||
"valid": "accepted",
|
||||
"invalid": "refused",
|
||||
"error": "not checked",
|
||||
"unsent": "not sent",
|
||||
}
|
||||
if facts.base and facts.key in checked:
|
||||
key += f", {checked[facts.key]}"
|
||||
typer.echo("=== Unsloth Studio MCP ===")
|
||||
typer.echo(f" server: {facts.base or 'not reachable'}")
|
||||
typer.echo(f" API key ({env_var}): {key}")
|
||||
if facts.base:
|
||||
typer.echo(f" MCP endpoint: {facts.mcp}")
|
||||
for label, seconds, _ in facts.timeouts:
|
||||
typer.echo(f" {label} tool timeout: {seconds:g} s")
|
||||
healthy, hints, advice = mcp_verdict(facts)
|
||||
typer.echo(f"\n=== Verdict: {'OK' if healthy else 'PROBLEMS FOUND'} ===")
|
||||
for line in hints + advice or ["No problems found."]:
|
||||
typer.echo(f"- {line}")
|
||||
return 0 if healthy else 1
|
||||
269
unsloth_cli/tests/test_doctor_mcp.py
Normal file
269
unsloth_cli/tests/test_doctor_mcp.py
Normal file
|
|
@ -0,0 +1,269 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import json
|
||||
import sys
|
||||
import urllib.error
|
||||
from pathlib import Path
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
if str(_REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(_REPO_ROOT))
|
||||
|
||||
import pytest
|
||||
import typer
|
||||
from typer.testing import CliRunner
|
||||
|
||||
import unsloth_cli
|
||||
import unsloth_cli.commands.doctor as doctor_mod
|
||||
import unsloth_cli.commands.mcp as mcp_mod
|
||||
|
||||
CANARY = "sk-unsloth-LEAKCANARY0123456789"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def env(tmp_path, monkeypatch):
|
||||
for var, sub in (("HOME", "home"), ("CLAUDE_CONFIG_DIR", "claude"), ("CODEX_HOME", "codex")):
|
||||
(tmp_path / sub).mkdir()
|
||||
monkeypatch.setenv(var, str(tmp_path / sub))
|
||||
for var in ("UNSLOTH_STUDIO_URL", "MCP_TOOL_TIMEOUT", "WSL_DISTRO_NAME"):
|
||||
monkeypatch.delenv(var, raising = False)
|
||||
monkeypatch.setenv("UNSLOTH_API_KEY", CANARY)
|
||||
return tmp_path
|
||||
|
||||
|
||||
def _facts(**kw):
|
||||
base = dict(base = "http://127.0.0.1:8888", key = "valid", mcp = "on")
|
||||
base.update(kw)
|
||||
return mcp_mod.McpFacts(**base)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"facts, needle",
|
||||
[
|
||||
(_facts(base = None, key = "missing", mcp = "unknown"), "not reachable"),
|
||||
(_facts(key = "missing", mcp = "refused"), "Set UNSLOTH_API_KEY"),
|
||||
(_facts(key = "invalid", mcp = "refused"), "refused the key in UNSLOTH_API_KEY"),
|
||||
(_facts(key = "error", key_detail = "500 boom"), "500 boom"),
|
||||
(_facts(mcp = "off"), "Allow agent connections"),
|
||||
(_facts(mcp = "absent"), "unsloth studio update"),
|
||||
(_facts(mcp = "refused", mcp_detail = "Workflow keys cannot use"), "Workflow keys cannot use"),
|
||||
(_facts(mcp = "error", mcp_detail = "timed out"), "timed out"),
|
||||
(_facts(key = "unsent"), "UNSLOTH_STUDIO_URL"),
|
||||
],
|
||||
)
|
||||
def test_verdict_problems(facts, needle):
|
||||
healthy, hints, _ = mcp_mod.mcp_verdict(facts)
|
||||
assert not healthy
|
||||
assert any(needle in h for h in hints), hints
|
||||
|
||||
|
||||
def test_verdict_ok_with_timeout_advice():
|
||||
facts = _facts(
|
||||
timeouts = [
|
||||
("Codex unsloth-studio", 60, "Set tool_timeout_sec = 300."),
|
||||
("Codex other", 300, "x"),
|
||||
("Claude Code", 60.0, "Set MCP_TOOL_TIMEOUT=300000."),
|
||||
]
|
||||
)
|
||||
healthy, hints, advice = mcp_mod.mcp_verdict(facts)
|
||||
assert healthy and hints == []
|
||||
assert len(advice) == 2
|
||||
assert "60 s" in advice[0] and "tool_timeout_sec = 300" in advice[0]
|
||||
assert "MCP_TOOL_TIMEOUT=300000" in advice[1]
|
||||
|
||||
|
||||
class _Resp:
|
||||
status = 200
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *a):
|
||||
return False
|
||||
|
||||
|
||||
def _stub_http(
|
||||
monkeypatch,
|
||||
answers,
|
||||
base = "http://127.0.0.1:8888",
|
||||
):
|
||||
"""answers: path suffix -> status. Records (url, Authorization header)."""
|
||||
seen = []
|
||||
monkeypatch.setattr(mcp_mod, "find_studio_server", lambda *a, **k: base)
|
||||
|
||||
def opener(request, timeout):
|
||||
seen.append((request.full_url, request.get_header("Authorization")))
|
||||
for suffix, status in answers.items():
|
||||
if request.full_url.endswith(suffix) and status != 200:
|
||||
body = io.BytesIO(json.dumps({"detail": f"nope {status}"}).encode())
|
||||
raise urllib.error.HTTPError(request.full_url, status, "x", {}, body)
|
||||
return _Resp()
|
||||
|
||||
monkeypatch.setattr(mcp_mod, "urlopen_no_redirect", opener)
|
||||
return seen
|
||||
|
||||
|
||||
def test_gather_sends_the_key_to_loopback(env, monkeypatch):
|
||||
seen = _stub_http(monkeypatch, {"/loaded-models": 200, "/mcp/": 200})
|
||||
facts = mcp_mod.gather_mcp_facts()
|
||||
assert (facts.key, facts.mcp) == ("valid", "on")
|
||||
assert [auth for _, auth in seen] == [f"Bearer {CANARY}"] * 2
|
||||
|
||||
|
||||
def test_gather_reports_off_and_refused(env, monkeypatch):
|
||||
seen = _stub_http(monkeypatch, {"/loaded-models": 401, "/mcp/": 404})
|
||||
facts = mcp_mod.gather_mcp_facts()
|
||||
assert (facts.key, facts.mcp) == ("invalid", "off")
|
||||
# A refused key is not sent again.
|
||||
assert seen[1][1] is None
|
||||
_stub_http(monkeypatch, {"/loaded-models": 200, "/mcp/": 401})
|
||||
facts = mcp_mod.gather_mcp_facts()
|
||||
assert (facts.mcp, facts.mcp_detail) == ("refused", "nope 401")
|
||||
|
||||
|
||||
def test_gather_never_sends_the_key_to_a_non_loopback_base(env, monkeypatch):
|
||||
seen = _stub_http(monkeypatch, {}, base = "http://10.0.0.5:8888")
|
||||
facts = mcp_mod.gather_mcp_facts()
|
||||
assert facts.key == "unsent"
|
||||
assert seen and all(auth is None for _, auth in seen)
|
||||
monkeypatch.setenv("UNSLOTH_STUDIO_URL", "http://10.0.0.5:8888")
|
||||
seen = _stub_http(monkeypatch, {}, base = "http://10.0.0.5:8888")
|
||||
assert mcp_mod.gather_mcp_facts().key == "valid"
|
||||
assert seen[0][1] == f"Bearer {CANARY}"
|
||||
|
||||
|
||||
def test_gather_without_a_server_or_key(env, monkeypatch):
|
||||
monkeypatch.setattr(mcp_mod, "find_studio_server", lambda *a, **k: None)
|
||||
assert mcp_mod.gather_mcp_facts().base is None
|
||||
monkeypatch.delenv("UNSLOTH_API_KEY")
|
||||
seen = _stub_http(monkeypatch, {"/mcp/": 401})
|
||||
facts = mcp_mod.gather_mcp_facts()
|
||||
assert (facts.key, facts.mcp) == ("missing", "refused")
|
||||
assert [url for url, _ in seen] == ["http://127.0.0.1:8888/mcp/"]
|
||||
|
||||
|
||||
def test_timeouts_read_from_codex_and_claude(env, monkeypatch):
|
||||
if mcp_mod._toml_loads() is None:
|
||||
pytest.skip("needs a TOML parser")
|
||||
(env / "codex" / "config.toml").write_text(
|
||||
'[mcp_servers.unsloth-studio]\nurl = "http://127.0.0.1:8888/mcp/"\n\n'
|
||||
'[mcp_servers.other]\nurl = "http://h/mcp/"\ntool_timeout_sec = 400\n\n'
|
||||
'[mcp_servers.unrelated]\ncommand = "x"\n'
|
||||
)
|
||||
# Another server is ours only when its URL is this Unsloth Studio's, not any /mcp/ URL.
|
||||
found = mcp_mod._codex_timeouts("unsloth-studio")
|
||||
assert [(label, s) for label, s, _ in found] == [("Codex unsloth-studio", 60)]
|
||||
found = mcp_mod._codex_timeouts("unsloth-studio", ("http://h/mcp/",))
|
||||
assert [(label, s) for label, s, _ in found] == [
|
||||
("Codex unsloth-studio", 60),
|
||||
("Codex other", 400),
|
||||
]
|
||||
assert mcp_mod._claude_timeouts("unsloth-studio") == []
|
||||
(env / "claude" / ".claude.json").write_text(
|
||||
json.dumps({"mcpServers": {"unsloth-studio": {"url": "http://h/mcp/"}}})
|
||||
)
|
||||
(env / "claude" / "settings.json").write_text(
|
||||
json.dumps({"env": {"MCP_TOOL_TIMEOUT": "60000"}})
|
||||
)
|
||||
assert mcp_mod._claude_timeouts("unsloth-studio")[0][:2] == ("Claude Code", 60.0)
|
||||
monkeypatch.setenv("MCP_TOOL_TIMEOUT", "300000")
|
||||
assert mcp_mod._claude_timeouts("unsloth-studio")[0][:2] == ("Claude Code", 300.0)
|
||||
|
||||
|
||||
def _cli():
|
||||
cli = typer.Typer()
|
||||
cli.command()(doctor_mod.doctor)
|
||||
return cli
|
||||
|
||||
|
||||
def test_doctor_mcp_skips_gpus_and_exits_1_on_problems(env, monkeypatch):
|
||||
_stub_http(monkeypatch, {"/loaded-models": 200, "/mcp/": 404})
|
||||
monkeypatch.setattr(doctor_mod, "_gpu_checks", lambda t: pytest.fail("GPU checks ran"))
|
||||
result = CliRunner().invoke(_cli(), ["--mcp"])
|
||||
assert result.exit_code == 1, result.output
|
||||
assert "=== Unsloth Studio MCP ===" in result.output
|
||||
assert "API key (UNSLOTH_API_KEY): set, accepted" in result.output
|
||||
assert "Verdict: PROBLEMS FOUND" in result.output
|
||||
assert "LEAKCANARY" not in result.output and "sk-unsloth" not in result.output
|
||||
|
||||
|
||||
def test_doctor_mcp_ok_with_advice(env, monkeypatch):
|
||||
_stub_http(monkeypatch, {})
|
||||
(env / "codex" / "config.toml").write_text('[mcp_servers.unsloth-studio]\nurl = "u"\n')
|
||||
result = CliRunner().invoke(_cli(), ["--mcp", "--no-gpus"])
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "Codex unsloth-studio tool timeout: 60 s" in result.output
|
||||
assert "Verdict: OK" in result.output and "tool_timeout_sec = 300" in result.output
|
||||
assert "LEAKCANARY" not in result.output
|
||||
|
||||
|
||||
def test_doctor_mcp_checks_the_named_env_var(env, monkeypatch):
|
||||
_stub_http(monkeypatch, {})
|
||||
monkeypatch.delenv("UNSLOTH_API_KEY", raising = False)
|
||||
monkeypatch.setenv("MY_STUDIO_KEY", "sk-unsloth-LEAKCANARY0123456789")
|
||||
result = CliRunner().invoke(_cli(), ["--mcp", "--env-var", "MY_STUDIO_KEY"])
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "API key (MY_STUDIO_KEY): set, accepted" in result.output
|
||||
assert "LEAKCANARY" not in result.output
|
||||
|
||||
|
||||
@pytest.mark.parametrize("gpu_code, mcp_code", [(0, 1), (1, 0), (0, 0)])
|
||||
def test_doctor_mcp_with_gpus_runs_both_and_keeps_the_worst_code(monkeypatch, gpu_code, mcp_code):
|
||||
ran = []
|
||||
monkeypatch.setattr(mcp_mod, "mcp_doctor", lambda env_var: ran.append("mcp") or mcp_code)
|
||||
|
||||
def gpu_checks(timeout):
|
||||
ran.append("gpus")
|
||||
raise typer.Exit(gpu_code)
|
||||
|
||||
monkeypatch.setattr(doctor_mod, "_gpu_checks", gpu_checks)
|
||||
result = CliRunner().invoke(_cli(), ["--mcp", "--gpus"])
|
||||
assert ran == ["mcp", "gpus"]
|
||||
assert result.exit_code == max(gpu_code, mcp_code)
|
||||
|
||||
|
||||
def test_plain_doctor_is_unchanged(monkeypatch):
|
||||
monkeypatch.setattr(mcp_mod, "mcp_doctor", lambda: pytest.fail("MCP checks ran"))
|
||||
result = CliRunner().invoke(_cli(), ["--no-gpus"])
|
||||
assert result.exit_code == 0 and result.output == "Nothing to check.\n"
|
||||
ran = []
|
||||
monkeypatch.setattr(doctor_mod, "_gpu_checks", lambda t: ran.append(t))
|
||||
assert CliRunner().invoke(_cli(), []).exit_code == 0
|
||||
assert ran == [60.0]
|
||||
help_text = CliRunner().invoke(unsloth_cli.app, ["doctor", "--help"]).output
|
||||
assert "--gpus" in help_text and "--mcp" in help_text
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", [400, 403, 406, 421])
|
||||
def test_gather_reports_an_unhandled_4xx_from_mcp_as_an_error(env, monkeypatch, status):
|
||||
_stub_http(monkeypatch, {"/loaded-models": 200, "/mcp/": status})
|
||||
facts = mcp_mod.gather_mcp_facts()
|
||||
assert facts.mcp == "error" and str(status) in facts.mcp_detail
|
||||
assert not mcp_mod.mcp_verdict(facts)[0]
|
||||
|
||||
|
||||
def test_probe_errors_never_quote_the_key(env, monkeypatch):
|
||||
# http.client quotes a refused header value, so a key with a newline lands in the error.
|
||||
monkeypatch.setenv("UNSLOTH_API_KEY", CANARY + "\r\n")
|
||||
monkeypatch.setattr(mcp_mod, "find_studio_server", lambda *a, **k: "http://127.0.0.1:9")
|
||||
facts = mcp_mod.gather_mcp_facts(timeout = 1.0)
|
||||
assert facts.key == "error"
|
||||
assert "LEAKCANARY" not in facts.key_detail, facts.key_detail
|
||||
assert all("LEAKCANARY" not in line for line in mcp_mod.mcp_verdict(facts)[1])
|
||||
|
||||
|
||||
def test_probes_send_the_cli_user_agent(env, monkeypatch):
|
||||
agents = []
|
||||
monkeypatch.setattr(mcp_mod, "find_studio_server", lambda *a, **k: "http://127.0.0.1:8888")
|
||||
|
||||
def opener(request, timeout):
|
||||
agents.append(request.get_header("User-agent"))
|
||||
return _Resp()
|
||||
|
||||
monkeypatch.setattr(mcp_mod, "urlopen_no_redirect", opener)
|
||||
mcp_mod.gather_mcp_facts()
|
||||
assert agents == ["unsloth-cli"] * 2
|
||||
455
unsloth_cli/tests/test_mcp_add.py
Normal file
455
unsloth_cli/tests/test_mcp_add.py
Normal file
|
|
@ -0,0 +1,455 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
if str(_REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(_REPO_ROOT))
|
||||
|
||||
import pytest
|
||||
from typer.testing import CliRunner
|
||||
|
||||
import unsloth_cli
|
||||
import unsloth_cli.commands.mcp as mcp_mod
|
||||
|
||||
CANARY = "sk-unsloth-LEAKCANARY0123456789"
|
||||
URL = "http://127.0.0.1:18799/mcp/"
|
||||
POSIX = pytest.mark.skipif(os.name == "nt", reason = "POSIX file modes")
|
||||
tomllib = pytest.importorskip("tomllib") if sys.version_info >= (3, 11) else None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def home(tmp_path, monkeypatch):
|
||||
for var, sub in (("HOME", "home"), ("CLAUDE_CONFIG_DIR", "claude"), ("CODEX_HOME", "codex")):
|
||||
(tmp_path / sub).mkdir()
|
||||
monkeypatch.setenv(var, str(tmp_path / sub))
|
||||
monkeypatch.delenv("UNSLOTH_STUDIO_URL", raising = False)
|
||||
monkeypatch.delenv("WSL_DISTRO_NAME", raising = False)
|
||||
monkeypatch.setenv("UNSLOTH_API_KEY", CANARY)
|
||||
monkeypatch.setattr(mcp_mod, "find_studio_server", lambda *a, **k: None)
|
||||
monkeypatch.setattr(mcp_mod, "_which_with_install_dirs", lambda name: None)
|
||||
return tmp_path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_claude(home, monkeypatch):
|
||||
"""A `claude` binary that records argv and edits .claude.json the way the real one does."""
|
||||
calls = []
|
||||
path = home / "claude" / ".claude.json"
|
||||
|
||||
def run(argv, **kwargs):
|
||||
calls.append(list(argv))
|
||||
data = json.loads(path.read_text()) if path.exists() else {}
|
||||
servers = data.setdefault("mcpServers", {})
|
||||
if argv[2] == "remove":
|
||||
servers.pop(argv[-1], None)
|
||||
else:
|
||||
name, url, header = argv[7], argv[8], argv[10]
|
||||
key, value = header.split(": ", 1)
|
||||
servers[name] = {"type": "http", "url": url, "headers": {key: value}}
|
||||
path.write_text(json.dumps(data))
|
||||
return mcp_mod.subprocess.CompletedProcess(argv, 0, "", "")
|
||||
|
||||
monkeypatch.setattr(mcp_mod, "_which_with_install_dirs", lambda name: "/opt/bin/claude")
|
||||
monkeypatch.setattr(mcp_mod.subprocess, "run", run)
|
||||
return calls
|
||||
|
||||
|
||||
def _add(*args):
|
||||
return CliRunner().invoke(unsloth_cli.app, ["mcp", "add", *args])
|
||||
|
||||
|
||||
def _mode(path):
|
||||
return path.stat().st_mode & 0o777
|
||||
|
||||
|
||||
# Pure renderers
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"base",
|
||||
["http://h:1", "http://h:1/", "http://h:1/mcp", "http://h:1/mcp/", " http://h:1// "],
|
||||
)
|
||||
def test_mcp_url_normalises_to_trailing_slash(base):
|
||||
assert mcp_mod.mcp_url(base) == "http://h:1/mcp/"
|
||||
|
||||
|
||||
def test_entries_reference_the_env_var_only():
|
||||
assert mcp_mod.claude_entry(URL, "MY_KEY") == {
|
||||
"type": "http",
|
||||
"url": URL,
|
||||
"headers": {"Authorization": "Bearer ${MY_KEY}"},
|
||||
}
|
||||
merged = mcp_mod.cursor_merge({"mcpServers": {"o": {}}, "k": 1}, "s", URL, "V")
|
||||
assert merged["k"] == 1 and merged["mcpServers"]["o"] == {}
|
||||
assert merged["mcpServers"]["s"] == {
|
||||
"url": URL,
|
||||
"headers": {"Authorization": "Bearer ${env:V}"},
|
||||
}
|
||||
snippet = mcp_mod.generic_snippet("s", URL, "V")["mcpServers"]["s"]
|
||||
assert snippet["headers"] == {"Authorization": "Bearer ${V}"}
|
||||
|
||||
|
||||
EXPECTED = {"url": URL, "bearer_token_env_var": "V", "tool_timeout_sec": 300}
|
||||
EXISTING = (
|
||||
'model = "o3"\n\n'
|
||||
"[mcp_servers.other]\n"
|
||||
'command = "x"\n\n'
|
||||
"[mcp_servers.unsloth-studio]\n"
|
||||
'url = "http://old/mcp/"\n\n'
|
||||
"[mcp_servers.unsloth-studio.env]\n"
|
||||
'A = "1"\n\n'
|
||||
'[profiles."p"]\n'
|
||||
"x = 1\n"
|
||||
)
|
||||
|
||||
|
||||
def test_codex_merge_replaces_in_place_and_drops_subtables():
|
||||
new = mcp_mod.codex_merge(EXISTING, "unsloth-studio", URL, "UNSLOTH_API_KEY", 300)
|
||||
assert new == (
|
||||
'model = "o3"\n\n'
|
||||
"[mcp_servers.other]\n"
|
||||
'command = "x"\n\n'
|
||||
"[mcp_servers.unsloth-studio]\n"
|
||||
f'url = "{URL}"\n'
|
||||
'bearer_token_env_var = "UNSLOTH_API_KEY"\n'
|
||||
"tool_timeout_sec = 300\n\n"
|
||||
'[profiles."p"]\n'
|
||||
"x = 1\n"
|
||||
)
|
||||
assert mcp_mod.codex_merge(new, "unsloth-studio", URL, "UNSLOTH_API_KEY", 300) == new
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"header", ['[mcp_servers."unsloth-studio"]', "[ mcp_servers . 'unsloth-studio' ]"]
|
||||
)
|
||||
def test_codex_merge_matches_quoted_headers(header):
|
||||
new = mcp_mod.codex_merge(f'{header}\nurl = "x"\n', "unsloth-studio", URL, "V", 300)
|
||||
assert new.count("unsloth-studio") == 1
|
||||
|
||||
|
||||
def test_codex_merge_appends_and_keeps_other_tables_byte_identical():
|
||||
old = '# mine\n[mcp_servers.unsloth-studio-2]\nurl = "y"\n'
|
||||
new = mcp_mod.codex_merge(old, "unsloth-studio", URL, "V", 300)
|
||||
assert new.startswith(old + "\n[mcp_servers.unsloth-studio]\n")
|
||||
assert mcp_mod.codex_merge("", "s", URL, "V", 300).startswith("[mcp_servers.s]\n")
|
||||
|
||||
|
||||
@pytest.mark.skipif(tomllib is None, reason = "needs tomllib")
|
||||
@pytest.mark.parametrize(
|
||||
"old",
|
||||
[
|
||||
'mcp_servers = { unsloth-studio = { url = "x" } }\n',
|
||||
'mcp_servers.unsloth-studio.url = "x"\n',
|
||||
"[mcp_servers]\nunsloth-studio = { url = 'x' }\n",
|
||||
"[broken\n",
|
||||
],
|
||||
)
|
||||
def test_codex_check_refuses_other_forms_and_invalid_toml(old):
|
||||
new = mcp_mod.codex_merge(old, "unsloth-studio", URL, "V", 300)
|
||||
with pytest.raises(ValueError):
|
||||
mcp_mod.codex_check(old, new, "unsloth-studio", EXPECTED)
|
||||
|
||||
|
||||
def test_codex_check_without_a_toml_parser_refuses_an_inline_table(monkeypatch):
|
||||
monkeypatch.setattr(mcp_mod, "_toml_loads", lambda: None)
|
||||
for old, ok in (
|
||||
('mcp_servers = { unsloth-studio = { url = "x" } }\n', False),
|
||||
(EXISTING, True),
|
||||
):
|
||||
new = mcp_mod.codex_merge(old, "unsloth-studio", URL, "V", 300)
|
||||
try:
|
||||
mcp_mod.codex_check(old, new, "unsloth-studio", EXPECTED)
|
||||
assert ok
|
||||
except ValueError:
|
||||
assert not ok
|
||||
|
||||
|
||||
# The command
|
||||
|
||||
|
||||
def test_mcp_is_registered_on_the_cli():
|
||||
result = CliRunner().invoke(unsloth_cli.app, ["mcp", "add", "--help"])
|
||||
assert result.exit_code == 0, result.output
|
||||
for agent in ("claude-code", "codex", "cursor", "generic"):
|
||||
assert agent in result.output
|
||||
assert "--api-key" not in result.output
|
||||
|
||||
|
||||
@POSIX
|
||||
@pytest.mark.skipif(tomllib is None, reason = "needs tomllib")
|
||||
def test_codex_writes_once_then_reports_unchanged(home):
|
||||
path = home / "codex" / "config.toml"
|
||||
result = _add("codex", "--url", "http://127.0.0.1:18799")
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "Unsloth Studio MCP set up for codex" in result.output
|
||||
assert _mode(path) == 0o600
|
||||
assert tomllib.loads(path.read_text())["mcp_servers"]["unsloth-studio"] == {
|
||||
"url": URL,
|
||||
"bearer_token_env_var": "UNSLOTH_API_KEY",
|
||||
"tool_timeout_sec": 300,
|
||||
}
|
||||
before = path.stat().st_mtime_ns
|
||||
result = _add("codex", "--url", "http://127.0.0.1:18799")
|
||||
assert result.exit_code == 0 and "unchanged" in result.output
|
||||
assert path.stat().st_mtime_ns == before
|
||||
|
||||
|
||||
@POSIX
|
||||
@pytest.mark.parametrize("mode", [0o644, 0o400])
|
||||
def test_existing_file_keeps_its_mode(home, mode):
|
||||
path = home / "codex" / "config.toml"
|
||||
path.write_text(EXISTING)
|
||||
path.chmod(mode)
|
||||
result = _add("codex", "--url", URL, "--timeout", "600")
|
||||
assert result.exit_code == 0, result.output
|
||||
assert _mode(path) == mode
|
||||
assert "tool_timeout_sec = 600" in path.read_text()
|
||||
assert "[mcp_servers.unsloth-studio.env]" not in path.read_text()
|
||||
|
||||
|
||||
def test_dry_run_prints_a_diff_and_writes_nothing(home, fake_claude):
|
||||
path = home / "codex" / "config.toml"
|
||||
path.write_text(EXISTING)
|
||||
stat = path.stat()
|
||||
result = _add("codex", "--url", URL, "--dry-run")
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "+bearer_token_env_var" in result.output and "Dry run" in result.output
|
||||
assert path.read_text() == EXISTING and path.stat().st_mtime_ns == stat.st_mtime_ns
|
||||
result = _add("cursor", "--url", URL, "--dry-run")
|
||||
assert result.exit_code == 0 and not (home / "home" / ".cursor").exists()
|
||||
result = _add("claude-code", "--url", URL, "--dry-run")
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "Would run: /opt/bin/claude mcp add" in result.output
|
||||
assert fake_claude == [] and not (home / "claude" / ".claude.json").exists()
|
||||
|
||||
|
||||
def test_codex_refuses_an_inline_table_and_leaves_the_file(home):
|
||||
path = home / "codex" / "config.toml"
|
||||
old = 'mcp_servers = { unsloth-studio = { url = "x" } }\n'
|
||||
path.write_text(old)
|
||||
result = _add("codex", "--url", URL)
|
||||
assert result.exit_code == 1
|
||||
assert "[mcp_servers.unsloth-studio]" in result.output
|
||||
assert path.read_text() == old
|
||||
|
||||
|
||||
@POSIX
|
||||
def test_symlinked_config_is_refused(home):
|
||||
target = home / "real.json"
|
||||
target.write_text("{}")
|
||||
(home / "home" / ".cursor").mkdir()
|
||||
(home / "home" / ".cursor" / "mcp.json").symlink_to(target)
|
||||
result = _add("cursor", "--url", URL)
|
||||
assert result.exit_code == 1 and "symlink" in result.output
|
||||
assert target.read_text() == "{}"
|
||||
|
||||
|
||||
@POSIX
|
||||
def test_cursor_reads_jsonc_and_keeps_other_servers(home):
|
||||
path = home / "home" / ".cursor" / "mcp.json"
|
||||
path.parent.mkdir()
|
||||
path.write_text('{\n // mine\n "mcpServers": {"other": {"url": "x"},},\n}\n')
|
||||
result = _add("cursor", "--url", URL)
|
||||
assert result.exit_code == 0, result.output
|
||||
data = json.loads(path.read_text())
|
||||
assert data["mcpServers"]["other"] == {"url": "x"}
|
||||
assert data["mcpServers"]["unsloth-studio"] == {
|
||||
"url": URL,
|
||||
"headers": {"Authorization": "Bearer ${env:UNSLOTH_API_KEY}"},
|
||||
}
|
||||
assert "unchanged" in _add("cursor", "--url", URL).output
|
||||
|
||||
|
||||
@pytest.mark.parametrize("text", ["[1, 2]", "{nope"])
|
||||
def test_unparseable_json_is_refused(home, text):
|
||||
path = home / "home" / ".cursor" / "mcp.json"
|
||||
path.parent.mkdir()
|
||||
path.write_text(text)
|
||||
result = _add("cursor", "--url", URL)
|
||||
assert result.exit_code == 1 and "Couldn't parse" in result.output
|
||||
assert path.read_text() == text
|
||||
|
||||
|
||||
@POSIX
|
||||
def test_claude_without_the_binary_writes_the_json_directly(home):
|
||||
path = home / "claude" / ".claude.json"
|
||||
result = _add("claude-code", "--url", "http://127.0.0.1:18799")
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "Claude Code was not found" in result.output
|
||||
assert json.loads(path.read_text())["mcpServers"]["unsloth-studio"] == {
|
||||
"type": "http",
|
||||
"url": URL,
|
||||
"headers": {"Authorization": "Bearer ${UNSLOTH_API_KEY}"},
|
||||
}
|
||||
assert _mode(path) == 0o600
|
||||
assert "unchanged" in _add("claude-code", "--url", URL).output
|
||||
|
||||
|
||||
def test_claude_shells_out_with_list_args(home, fake_claude):
|
||||
result = _add("claude-code", "--url", URL, "--env-var", "MY_KEY")
|
||||
assert result.exit_code == 0, result.output
|
||||
argv = "/opt/bin/claude mcp add --transport http --scope user unsloth-studio".split()
|
||||
assert fake_claude == [argv + [URL, "--header", "Authorization: Bearer ${MY_KEY}"]]
|
||||
assert "Check the result" not in result.output
|
||||
fake_claude.clear()
|
||||
assert "unchanged" in _add("claude-code", "--url", URL, "--env-var", "MY_KEY").output
|
||||
assert fake_claude == []
|
||||
result = _add("claude-code", "--url", "http://127.0.0.1:9/", "--env-var", "MY_KEY")
|
||||
assert result.exit_code == 0, result.output
|
||||
assert [argv[2] for argv in fake_claude] == ["remove", "add"]
|
||||
assert fake_claude[0][-3:] == ["--scope", "user", "unsloth-studio"]
|
||||
|
||||
|
||||
def test_claude_failure_and_unparseable_config_exit_1(home, fake_claude, monkeypatch):
|
||||
path = home / "claude" / ".claude.json"
|
||||
path.write_text("{broken")
|
||||
assert _add("claude-code", "--url", URL).exit_code == 1
|
||||
assert fake_claude == [] and path.read_text() == "{broken"
|
||||
path.unlink()
|
||||
failed = lambda argv, **k: mcp_mod.subprocess.CompletedProcess(argv, 1, "", "boom")
|
||||
monkeypatch.setattr(mcp_mod.subprocess, "run", failed)
|
||||
result = _add("claude-code", "--url", URL)
|
||||
assert result.exit_code == 1 and "boom" in result.output
|
||||
|
||||
|
||||
def test_generic_prints_or_merges(home):
|
||||
result = _add("generic", "--url", URL)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert '"Authorization": "Bearer ${UNSLOTH_API_KEY}"' in result.output
|
||||
out = home / "agent.json"
|
||||
out.write_text('{"mcpServers": {"a": {}}, "x": 1}')
|
||||
assert _add("generic", "--url", URL, "--output", str(out)).exit_code == 0
|
||||
data = json.loads(out.read_text())
|
||||
assert data["x"] == 1 and set(data["mcpServers"]) == {"a", "unsloth-studio"}
|
||||
|
||||
|
||||
def test_codex_merge_keeps_comments_about_the_next_table():
|
||||
old = '[mcp_servers.unsloth-studio]\nurl = "old"\n\n# the work proxy\n[mcp_servers.work]\nurl = "w"\n'
|
||||
new = mcp_mod.codex_merge(old, "unsloth-studio", URL, "K", 300)
|
||||
assert "\n\n# the work proxy\n[mcp_servers.work]\n" in new and 'url = "old"' not in new
|
||||
|
||||
|
||||
def test_dry_run_never_prints_a_key_already_in_a_config(home, fake_claude):
|
||||
# Older hand-written entries may hold a literal key in any spelling; other servers hold their own.
|
||||
literal = {
|
||||
"url": "u",
|
||||
"headers": {"Authorization": "bearer sk-OLDKEY1", "X-Api-Key": "OLDKEY2"},
|
||||
}
|
||||
other = {"command": "gh-mcp", "env": {"GITHUB_TOKEN": "ghp_OLDKEY3"}}
|
||||
servers = {"mcpServers": {"unsloth-studio": literal, "github": other}}
|
||||
cursor = home / "home" / ".cursor" / "mcp.json"
|
||||
cursor.parent.mkdir()
|
||||
cursor.write_text(json.dumps(servers, indent = 4))
|
||||
(home / "claude" / ".claude.json").write_text(json.dumps(servers))
|
||||
(home / "codex" / "config.toml").write_text(
|
||||
'[mcp_servers.github]\nbearer_token = "ghp_OLDKEY3"\n\n'
|
||||
'[mcp_servers.unsloth-studio]\nurl = "u"\nbearer_token = """sk-OLDKEY1"""\n'
|
||||
'http_headers = { "X-Api-Key" = "OLDKEY2" }\n'
|
||||
)
|
||||
for agent in ("cursor", "claude-code", "codex"):
|
||||
result = _add(agent, "--url", URL, "--dry-run")
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "OLDKEY" not in result.output, (agent, result.output)
|
||||
assert "its old values are not shown" in result.output, agent
|
||||
assert "${UNSLOTH_API_KEY}" in result.output or "UNSLOTH_API_KEY" in result.output
|
||||
assert fake_claude == []
|
||||
|
||||
|
||||
def test_a_non_object_mcp_servers_is_left_alone(home):
|
||||
path = home / "home" / ".cursor" / "mcp.json"
|
||||
path.parent.mkdir()
|
||||
path.write_text('{"mcpServers": ["mine"]}')
|
||||
result = _add("cursor", "--url", URL)
|
||||
assert result.exit_code == 1 and "not an object" in result.output
|
||||
assert path.read_text() == '{"mcpServers": ["mine"]}'
|
||||
|
||||
|
||||
def test_generic_output_must_be_json(home):
|
||||
out = home / "agent.toml"
|
||||
out.write_text("[x]\n")
|
||||
result = _add("generic", "--url", URL, "--output", str(out))
|
||||
assert result.exit_code == 1 and ".json" in result.output
|
||||
assert out.read_text() == "[x]\n"
|
||||
|
||||
|
||||
def test_a_failed_add_after_remove_says_the_entry_is_gone(home, fake_claude, monkeypatch):
|
||||
(home / "claude" / ".claude.json").write_text(
|
||||
json.dumps({"mcpServers": {"unsloth-studio": {"url": "old"}}})
|
||||
)
|
||||
fake = mcp_mod.subprocess.run
|
||||
|
||||
def run(argv, **kwargs):
|
||||
if argv[2] == "add":
|
||||
return mcp_mod.subprocess.CompletedProcess(argv, 1, "", "add broke")
|
||||
return fake(argv, **kwargs)
|
||||
|
||||
monkeypatch.setattr(mcp_mod.subprocess, "run", run)
|
||||
result = _add("claude-code", "--url", URL)
|
||||
assert result.exit_code == 1 and "add broke" in result.output
|
||||
assert "already removed" in result.output
|
||||
|
||||
|
||||
BAD = ["--env-var=BAD-NAME", "--env-var=$(id)", "--name=a b", "--name=x]\n[y", "--url=file:///e"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("arg", BAD)
|
||||
def test_bad_input_is_refused(home, arg):
|
||||
result = _add("codex", "--url", URL, arg)
|
||||
assert result.exit_code == 1, result.output
|
||||
assert not (home / "codex" / "config.toml").exists()
|
||||
|
||||
|
||||
def test_timeout_below_the_wait_cap_is_rejected(home):
|
||||
assert _add("codex", "--url", URL, "--timeout", "120").exit_code == 2
|
||||
|
||||
|
||||
def test_url_falls_back_to_env_then_discovery_then_default(home, monkeypatch):
|
||||
assert mcp_mod._resolve_base(None) == "http://127.0.0.1:8888"
|
||||
monkeypatch.setattr(mcp_mod, "find_studio_server", lambda *a, **k: "http://127.0.0.1:9000")
|
||||
assert mcp_mod._resolve_base(None) == "http://127.0.0.1:9000"
|
||||
monkeypatch.setenv("UNSLOTH_STUDIO_URL", "http://box:1")
|
||||
assert mcp_mod._resolve_base(None) == "http://box:1"
|
||||
assert mcp_mod._resolve_base("http://x:2") == "http://x:2"
|
||||
|
||||
|
||||
def test_the_key_never_leaks(home, fake_claude, monkeypatch):
|
||||
assert "is not set" not in _add("generic", "--url", URL).output
|
||||
outputs = []
|
||||
for agent in ("claude-code", "codex", "cursor", "generic"):
|
||||
for extra in (["--dry-run"], []):
|
||||
result = _add(agent, "--url", URL, *extra)
|
||||
assert result.exit_code == 0, result.output
|
||||
outputs.append(result.output)
|
||||
outputs.append(_add("generic", "--url", URL, "--output", str(home / "g.json")).output)
|
||||
assert fake_claude, "the claude stub should have run"
|
||||
assert all("LEAKCANARY" not in " ".join(argv) for argv in fake_claude)
|
||||
assert all("LEAKCANARY" not in text for text in outputs)
|
||||
for path in home.rglob("*"):
|
||||
if path.is_file():
|
||||
assert "LEAKCANARY" not in path.read_text(), path
|
||||
monkeypatch.delenv("UNSLOTH_API_KEY")
|
||||
assert "UNSLOTH_API_KEY is not set" in _add("cursor", "--url", URL).output
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"old",
|
||||
[
|
||||
# A header inside a multi-line string: the splice would cut the string and the keys after it.
|
||||
'model = "o3"\nnotes = """\n[mcp_servers.unsloth-studio]\nkeep me\n"""\nsandbox = "x"\n',
|
||||
# A nested array line inside the table would be left behind as broken TOML.
|
||||
'[mcp_servers.unsloth-studio]\nurl = "x"\nm = [\n [1, 2],\n]\n',
|
||||
],
|
||||
)
|
||||
def test_codex_check_without_a_toml_parser_refuses_unsafe_splices(monkeypatch, old):
|
||||
monkeypatch.setattr(mcp_mod, "_toml_loads", lambda: None)
|
||||
new = mcp_mod.codex_merge(old, "unsloth-studio", URL, "V", 300)
|
||||
assert new != old
|
||||
with pytest.raises(ValueError):
|
||||
mcp_mod.codex_check(old, new, "unsloth-studio", EXPECTED)
|
||||
Loading…
Add table
Add a link
Reference in a new issue