From e5cf956601a9271ddb2441441c734baeb5be3bc4 Mon Sep 17 00:00:00 2001 From: Nilay <118994073+NilayYadav@users.noreply.github.com> Date: Wed, 24 Jun 2026 19:01:53 +0530 Subject: [PATCH] Studio: shareable per-checkpoint preview links (#6486) * checkpoint preview endpoint * harden new preview endpoints * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * address review * Studio preview: pin adapter, guard streaming submit, robust copy-link Harden the public per-checkpoint preview surface: - Pin use_adapter=True in the preview payload sanitizer. Otherwise an unauthenticated /p caller can POST use_adapter=false, which calls disable_adapter_layers() on the shared in-memory model without restoring it; since load_model skips reloads for the same checkpoint, every later visitor (the page never sends the field) keeps getting base-model output instead of the fine-tuned checkpoint. Forcing it on also re-enables a previously disabled adapter and no-ops on merged checkpoints. - Ignore preview-page submits while a response is streaming. The send button was disabled but the Enter handler still called requestSubmit(), so a second request could start before the first reply landed in msgs and reorder the chat history. Both the keydown and submit handlers now honor the disabled button. - Keep the cloudflare-URL polling loop alive across transient startup fetch errors instead of letting one rejection halt it. - Build the copy-link from a backend preview_ref (output dir relative to outputs_root, gated on previewability and the two-segment /p route limit) so a nested output dir no longer copies a basename-only link that 404s. Expose preview_ref on training run summaries. Add route-level security tests (path traversal, payload sanitization, asset containment, CSP header, HTML title escaping, streaming lock held until drained) and preview_ref unit tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio preview: Safari-safe submit and adapter pin only for LoRA Follow-ups from cross-browser and route simulations: - Preview page: send the message from a shared send() helper called by both the form submit and the Enter key, instead of form.requestSubmit(). The latter throws on Safari < 16 and older iOS, which broke Enter-to-send there. Verified across Chromium, Firefox and WebKit with Playwright. - Only pin use_adapter=True when the resolved checkpoint is a LoRA adapter (adapter_config.json present); for a merged checkpoint strip it to None. A merged model has no adapter to toggle, so forcing it on only produced a per-request "not a PeftModel" warning. The cross-request base-model contamination fix still holds for LoRA previews. Add a merged-checkpoint test asserting use_adapter is stripped to None. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio preview: trim verbose comments Tighten comments across the preview routes, page, checkpoint helpers, and tests to short single-line notes; drop ones that just restate the code. No behavior change (verified comment/docstring-only with comment_tools.py check). * Harden preview routes for PR #6486 - Return a generic 400 detail on a rejected preview path so the public /p route never echoes the absolute install path (the real reason is logged server-side instead). - Strip confirm_tool_calls, session_id and rag_scope in the preview payload sanitizer so the public surface stays inert regardless of the tool gate. - Use Path.is_relative_to for the asset containment check, matching the rest of the codebase. - Add img-src 'self' and font-src 'self' to the preview page CSP. - Preview page: on a mid-stream error keep the streamed text, flag the break, and restore the prompt so the user can retry; drop the unused --font-sans var. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: Daniel Han --- pyproject.toml | 1 + studio/backend/assets/preview_page.html | 398 ++++++++++++++++++ studio/backend/main.py | 3 + studio/backend/models/training.py | 2 + studio/backend/routes/preview.py | 196 +++++++++ studio/backend/routes/training_history.py | 17 +- studio/backend/run.py | 5 +- studio/backend/state/tool_policy.py | 21 +- studio/backend/tests/test_preview.py | 134 ++++++ studio/backend/tests/test_preview_routes.py | 293 +++++++++++++ studio/backend/utils/api_errors.py | 17 +- studio/backend/utils/models/checkpoints.py | 61 +++ .../src/features/studio/history-card-grid.tsx | 63 ++- .../src/features/training/types/history.ts | 2 + studio/frontend/src/i18n/locales/en.ts | 3 + 15 files changed, 1206 insertions(+), 10 deletions(-) create mode 100644 studio/backend/assets/preview_page.html create mode 100644 studio/backend/routes/preview.py create mode 100644 studio/backend/tests/test_preview.py create mode 100644 studio/backend/tests/test_preview_routes.py diff --git a/pyproject.toml b/pyproject.toml index bf4b118d68..13c421d8ea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,6 +57,7 @@ studio = [ "backend/requirements/**/*", "backend/plugins/**/*", "backend/assets/**/*.jinja", + "backend/assets/**/*.html", "backend/core/data_recipe/oxc-validator/*.json", "backend/core/data_recipe/oxc-validator/*.mjs", ] diff --git a/studio/backend/assets/preview_page.html b/studio/backend/assets/preview_page.html new file mode 100644 index 0000000000..272f7ac89b --- /dev/null +++ b/studio/backend/assets/preview_page.html @@ -0,0 +1,398 @@ + + + + + + __TITLE__ - Unsloth + + + +
+ Unsloth__TITLE__ +
+
+
+

Chat with your model

+

Fine-tuned with Unsloth

+
+
+
+
+
+
+ + +
+
Served by Unsloth Studio
+
+
+ + + diff --git a/studio/backend/main.py b/studio/backend/main.py index a56bd46c4b..a8e81b68e6 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -282,6 +282,7 @@ from routes import ( training_router, ) from routes.llama import router as llama_router +from routes.preview import router as preview_router from hub.routes import ( inventory_router as hub_inventory_router, datasets_router as hub_datasets_router, @@ -672,6 +673,7 @@ from utils.upload_limits import ( # noqa: E402 _BODY_PROTECTED_PREFIXES = ( "/v1/chat/completions", "/v1/completions", + "/p/", "/api/inference", "/api/data-recipe", "/api/datasets", @@ -885,6 +887,7 @@ app.include_router(inference_studio_router, prefix = "/api/inference", tags = [" # OpenAI-compatible: mount the inference router at /v1 for external tools. app.include_router(inference_router, prefix = "/v1", tags = ["openai-compat"]) +app.include_router(preview_router, prefix = "/p", tags = ["preview"]) app.include_router(providers_router, prefix = "/api/providers", tags = ["providers"]) app.include_router(settings_router, prefix = "/api/settings", tags = ["settings"]) app.include_router(mcp_servers_router, prefix = "/api/mcp/servers", tags = ["mcp"]) diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 670d5f911d..8b6fb36471 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -601,6 +601,8 @@ class TrainingRunSummary(BaseModel): loss_sparkline: Optional[List[float]] = None can_resume: bool = False resumed_later: bool = False + has_preview_model: bool = False + preview_ref: Optional[str] = None class TrainingRunUpdateRequest(BaseModel): diff --git a/studio/backend/routes/preview.py b/studio/backend/routes/preview.py new file mode 100644 index 0000000000..d5247a2bcf --- /dev/null +++ b/studio/backend/routes/preview.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Per-checkpoint preview endpoints: /p/{run}[/{checkpoint}]/v1/...""" + +from __future__ import annotations + +import asyncio +import html +from pathlib import Path +from urllib.parse import quote + +from fastapi import APIRouter, Depends, HTTPException, Request +from fastapi.responses import FileResponse, HTMLResponse, StreamingResponse +from loggers import get_logger + +from auth.authentication import get_current_subject +from auth.storage import DEFAULT_ADMIN_USERNAME +from models.inference import ChatCompletionRequest, LoadRequest +from routes.inference import load_model, openai_chat_completions +from state.tool_policy import tools_force_disabled +from utils.models.checkpoints import list_preview_targets, resolve_preview_checkpoint + +logger = get_logger(__name__) + +router = APIRouter() + +# Public (no key); resolve_preview_checkpoint pins `run` under outputs_root. +# One model loads at a time, so serialize load+generate across previews. +_preview_lock = asyncio.Lock() + + +def _resolve_or_4xx(run: str, checkpoint: str | None): + try: + return resolve_preview_checkpoint(run, checkpoint) + except ValueError as exc: + # Detail can carry the absolute install path on a symlink escape; log it, + # return a generic message on this public route. + logger.warning("preview path rejected: %s", exc) + raise HTTPException(status_code = 400, detail = "Invalid run or checkpoint") + except FileNotFoundError as exc: + raise HTTPException(status_code = 404, detail = str(exc)) + + +def _sanitize_preview_payload( + payload: ChatCompletionRequest, is_lora: bool +) -> ChatCompletionRequest: + # Public surface: strip tools/MCP + provider routing (no host code / open proxy). + # Normalize use_adapter (never trust the caller): pin True for LoRA, None for + # merged. _apply_adapter_state mutates the shared model without restoring, so an + # unpinned `false` would persist to later visitors who omit the field. + return payload.model_copy( + update = { + "tools": None, + "enable_tools": False, + "enabled_tools": None, + "mcp_enabled": False, + "bypass_permissions": False, + "confirm_tool_calls": False, + "session_id": None, + "rag_scope": None, + "openai_code_exec_container_id": None, + "anthropic_code_exec_container_id": None, + "provider_id": None, + "provider_type": None, + "external_model": None, + "encrypted_api_key": None, + "provider_base_url": None, + "use_adapter": True if is_lora else None, + } + ) + + +async def _unlock_after(body_iterator): + # Hold the lock until the stream drains so another checkpoint can't swap mid-stream. + try: + async for chunk in body_iterator: + yield chunk + finally: + _preview_lock.release() + + +async def _serve_chat( + run: str, checkpoint: str | None, payload: ChatCompletionRequest, request: Request +): + path = _resolve_or_4xx(run, checkpoint) + is_lora = (path / "adapter_config.json").exists() + payload = _sanitize_preview_payload(payload, is_lora) + await _preview_lock.acquire() + keep_locked = False + try: + await load_model(LoadRequest(model_path = str(path)), request, DEFAULT_ADMIN_USERNAME) + # Beats a process-wide `--enable-tools` (enable_tools=False alone wouldn't). + with tools_force_disabled(): + response = await openai_chat_completions(payload, request, DEFAULT_ADMIN_USERNAME) + if isinstance(response, StreamingResponse): + response.body_iterator = _unlock_after(response.body_iterator) + keep_locked = True + return response + finally: + if not keep_locked: + _preview_lock.release() + + +@router.get("") +async def list_previews(request: Request, current_subject: str = Depends(get_current_subject)): + base = str(request.base_url) + previews = [] + for target in list_preview_targets(): + ref = quote(target["ref"], safe = "/") + previews.append({**target, "url": f"{base}p/{ref}/v1"}) + return {"object": "list", "data": previews} + + +@router.post("/{run}/v1/chat/completions") +async def preview_chat_latest(run: str, payload: ChatCompletionRequest, request: Request): + return await _serve_chat(run, None, payload, request) + + +@router.post("/{run}/{checkpoint}/v1/chat/completions") +async def preview_chat_checkpoint( + run: str, checkpoint: str, payload: ChatCompletionRequest, request: Request +): + return await _serve_chat(run, checkpoint, payload, request) + + +def _models_response(run: str, checkpoint: str | None): + path = _resolve_or_4xx(run, checkpoint) + model_id = run if not checkpoint else f"{run}/{checkpoint}" + return { + "object": "list", + "data": [ + { + "id": model_id, + "object": "model", + "created": int(path.stat().st_mtime), + "owned_by": "unsloth-studio", + } + ], + } + + +@router.get("/{run}/v1/models") +async def preview_models_latest(run: str): + return _models_response(run, None) + + +@router.get("/{run}/{checkpoint}/v1/models") +async def preview_models_checkpoint(run: str, checkpoint: str): + return _models_response(run, checkpoint) + + +# Serve logo/fonts here too: the SPA static mount is absent in --api-only (Tauri). +_FRONTEND_DIST = (Path(__file__).resolve().parents[2] / "frontend" / "dist").resolve() +_PREVIEW_ASSET_MEDIA_TYPES = { + ".png": "image/png", + ".woff": "font/woff", + ".woff2": "font/woff2", +} + + +@router.get("/_assets/{asset_path:path}") +async def preview_asset(asset_path: str): + target = (_FRONTEND_DIST / asset_path).resolve() + media_type = _PREVIEW_ASSET_MEDIA_TYPES.get(target.suffix.lower()) + if media_type is None or not target.is_relative_to(_FRONTEND_DIST) or not target.is_file(): + raise HTTPException(status_code = 404, detail = "Not found") + return FileResponse(target, media_type = media_type) + + +# Self-contained public page; only the title is interpolated. +_PREVIEW_PAGE_HTML = ( + Path(__file__).resolve().parent.parent / "assets" / "preview_page.html" +).read_text(encoding = "utf-8") + +_PREVIEW_PAGE_CSP = ( + "default-src 'self'; script-src 'unsafe-inline'; style-src 'unsafe-inline'; " + "img-src 'self'; font-src 'self'; connect-src 'self'; base-uri 'none'" +) + + +def _preview_page(run: str, checkpoint: str | None) -> HTMLResponse: + _resolve_or_4xx(run, checkpoint) + title = run if not checkpoint else f"{run}/{checkpoint}" + page = _PREVIEW_PAGE_HTML.replace("__TITLE__", html.escape(title)) + return HTMLResponse(page, headers = {"Content-Security-Policy": _PREVIEW_PAGE_CSP}) + + +@router.get("/{run}", response_class = HTMLResponse) +async def preview_page_latest(run: str): + return _preview_page(run, None) + + +@router.get("/{run}/{checkpoint}", response_class = HTMLResponse) +async def preview_page_checkpoint(run: str, checkpoint: str): + return _preview_page(run, checkpoint) diff --git a/studio/backend/routes/training_history.py b/studio/backend/routes/training_history.py index 1560c72767..a64d2a938e 100644 --- a/studio/backend/routes/training_history.py +++ b/studio/backend/routes/training_history.py @@ -27,6 +27,7 @@ from storage.studio_db import ( list_runs, update_run_display_name, ) +from utils.models.checkpoints import has_preview_model, preview_ref logger = get_logger(__name__) @@ -42,7 +43,17 @@ async def list_training_runs( """List training runs, newest first.""" result = list_runs(limit = limit, offset = offset) return TrainingRunListResponse( - runs = [TrainingRunSummary(**{**r, "can_resume": can_resume_run(r)}) for r in result["runs"]], + runs = [ + TrainingRunSummary( + **{ + **r, + "can_resume": can_resume_run(r), + "has_preview_model": has_preview_model(r.get("output_dir")), + "preview_ref": preview_ref(r.get("output_dir")), + } + ) + for r in result["runs"] + ], total = result["total"], ) @@ -67,6 +78,8 @@ async def get_training_run_detail(run_id: str, current_subject: str = Depends(ge **{ **{k: v for k, v in run.items() if k != "config_json"}, "can_resume": can_resume_run(run), + "has_preview_model": has_preview_model(run.get("output_dir")), + "preview_ref": preview_ref(run.get("output_dir")), } ), config = config, @@ -98,6 +111,8 @@ async def update_training_run( **{ **{k: v for k, v in refreshed.items() if k != "config_json"}, "can_resume": can_resume_run(refreshed), + "has_preview_model": has_preview_model(refreshed.get("output_dir")), + "preview_ref": preview_ref(refreshed.get("output_dir")), } ) diff --git a/studio/backend/run.py b/studio/backend/run.py index 481fd623f3..9cb7868949 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -1095,7 +1095,10 @@ def run_server( app.state.server_port = port if port and port > 0 else None # Direct (non-tunnel) base for the API panel; resolve 0.0.0.0 to the LAN IP. if port and port > 0: - _direct_host = _resolve_external_ip() if host == "0.0.0.0" else host + _direct_host = _resolve_external_ip() if host in ("0.0.0.0", "::") else host + # Bracket IPv6 literals so the URL is valid (http://[2405:...]:port). + if ":" in _direct_host and not _direct_host.startswith("["): + _direct_host = f"[{_direct_host}]" app.state.server_url = f"http://{_direct_host}:{port}" else: app.state.server_url = None diff --git a/studio/backend/state/tool_policy.py b/studio/backend/state/tool_policy.py index 9b0fc7d6cb..e0792321f9 100644 --- a/studio/backend/state/tool_policy.py +++ b/studio/backend/state/tool_policy.py @@ -10,15 +10,34 @@ Set by `unsloth run` at startup; consulted by the inference route gates. False -> CLI forced tools off for every request. """ -from typing import Optional +import contextvars +from contextlib import contextmanager +from typing import Iterator, Optional _tool_policy: Optional[bool] = None +# Per-request hard-off so public surfaces refuse tools even under a CLI `--enable-tools`. +_force_disabled: contextvars.ContextVar[bool] = contextvars.ContextVar( + "tool_policy_force_disabled", default = False +) + def get_tool_policy() -> Optional[bool]: + if _force_disabled.get(): + return False return _tool_policy +@contextmanager +def tools_force_disabled() -> Iterator[None]: + """Hard-disable server-side tools for the current async context.""" + token = _force_disabled.set(True) + try: + yield + finally: + _force_disabled.reset(token) + + def set_tool_policy(value: Optional[bool]) -> None: if value is not None and not isinstance(value, bool): raise TypeError(f"tool_policy must be Optional[bool], got {type(value).__name__}") diff --git a/studio/backend/tests/test_preview.py b/studio/backend/tests/test_preview.py new file mode 100644 index 0000000000..e131f99951 --- /dev/null +++ b/studio/backend/tests/test_preview.py @@ -0,0 +1,134 @@ +# 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 json +from pathlib import Path +import sys +import types as _types + +import pytest + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +from utils.models.checkpoints import ( + list_preview_targets, + preview_ref, + resolve_preview_checkpoint, +) + + +def _make_run(outputs: Path) -> tuple[Path, Path]: + run = outputs / "unsloth_SmolLM-135M_1775412608" + run.mkdir(parents = True) + (run / "adapter_config.json").write_text( + json.dumps({"base_model_name_or_path": "HuggingFaceTB/SmolLM-135M"}) + ) + ckpt = run / "checkpoint-60" + ckpt.mkdir() + (ckpt / "adapter_config.json").write_text( + json.dumps({"base_model_name_or_path": "HuggingFaceTB/SmolLM-135M"}) + ) + return run, ckpt + + +def _point_outputs_root_at(monkeypatch, outputs: Path) -> None: + from utils.paths import storage_roots as _sr + from utils.models import checkpoints as _ckpt + + monkeypatch.setattr(_sr, "outputs_root", lambda: outputs) + # checkpoints imported outputs_root by name; patch that alias too (preview_ref uses it). + monkeypatch.setattr(_ckpt, "outputs_root", lambda: outputs) + + +def test_resolve_main_adapter_and_checkpoint(tmp_path: Path, monkeypatch): + outputs = tmp_path / "outputs" + run, ckpt = _make_run(outputs) + _point_outputs_root_at(monkeypatch, outputs) + + assert resolve_preview_checkpoint(run.name) == run + assert resolve_preview_checkpoint(run.name, "checkpoint-60") == ckpt + + +def test_resolve_missing_raises_not_found(tmp_path: Path, monkeypatch): + outputs = tmp_path / "outputs" + _make_run(outputs) + _point_outputs_root_at(monkeypatch, outputs) + + with pytest.raises(FileNotFoundError): + resolve_preview_checkpoint("does-not-exist") + (outputs / "empty").mkdir() + with pytest.raises(FileNotFoundError): + resolve_preview_checkpoint("empty") + + +def test_resolve_rejects_traversal(tmp_path: Path, monkeypatch): + outputs = tmp_path / "outputs" + _make_run(outputs) + _point_outputs_root_at(monkeypatch, outputs) + + with pytest.raises(ValueError): + resolve_preview_checkpoint("..", "etc") + + +def test_list_preview_targets_flattens_with_latest_flag(tmp_path: Path, monkeypatch): + outputs = tmp_path / "outputs" + run, _ = _make_run(outputs) + _point_outputs_root_at(monkeypatch, outputs) + + targets = list_preview_targets(str(outputs)) + by_ref = {t["ref"]: t for t in targets} + + assert by_ref[run.name]["is_latest"] is True + assert by_ref[run.name]["checkpoint"] is None + assert by_ref[f"{run.name}/checkpoint-60"]["is_latest"] is False + assert by_ref[f"{run.name}/checkpoint-60"]["checkpoint"] == "checkpoint-60" + assert all(t["base_model"] == "HuggingFaceTB/SmolLM-135M" for t in targets) + + +def test_preview_ref_flat_run_is_basename(tmp_path: Path, monkeypatch): + outputs = tmp_path / "outputs" + run, _ = _make_run(outputs) + _point_outputs_root_at(monkeypatch, outputs) + + assert preview_ref(str(run)) == run.name + + +def test_preview_ref_preserves_one_level_nesting(tmp_path: Path, monkeypatch): + outputs = tmp_path / "outputs" + _point_outputs_root_at(monkeypatch, outputs) + nested = outputs / "experiments" / "run1" + nested.mkdir(parents = True) + (nested / "adapter_config.json").write_text("{}") + + # /p route supports run/checkpoint, so a single level of nesting survives. + assert preview_ref(str(nested)) == "experiments/run1" + + +def test_preview_ref_none_for_unpreviewable_or_too_deep(tmp_path: Path, monkeypatch): + outputs = tmp_path / "outputs" + _point_outputs_root_at(monkeypatch, outputs) + + # Missing / no model artifact -> not previewable. + assert preview_ref(None) is None + empty = outputs / "empty" + empty.mkdir(parents = True) + assert preview_ref(str(empty)) is None + + # Too deep for the two-segment /p route -> no dead link. + deep = outputs / "a" / "b" / "run" + deep.mkdir(parents = True) + (deep / "adapter_config.json").write_text("{}") + assert preview_ref(str(deep)) is None + + # Outside outputs_root -> None. + outside = tmp_path / "elsewhere" + outside.mkdir() + (outside / "adapter_config.json").write_text("{}") + assert preview_ref(str(outside)) is None diff --git a/studio/backend/tests/test_preview_routes.py b/studio/backend/tests/test_preview_routes.py new file mode 100644 index 0000000000..d6edd4ef4d --- /dev/null +++ b/studio/backend/tests/test_preview_routes.py @@ -0,0 +1,293 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Security smoke for the public /p preview routes. + +Exercises the route layer with a real ``preview_router`` while stubbing the +expensive model calls (``load_model`` / ``openai_chat_completions``). Covers the +public-surface guarantees: path-traversal rejection, request sanitization +(tools / provider routing / use_adapter), asset-path containment, the page CSP +header + HTML escaping, and that the preview lock is held until a streaming +response is fully drained. +""" + +import asyncio +import json +from pathlib import Path +import sys +import types as _types + +import pytest + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Mirror test_preview.py: the real `loggers` package pulls in heavy handlers. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +from fastapi import FastAPI +from fastapi.responses import StreamingResponse +from fastapi.testclient import TestClient + +import routes.preview as preview +from models.inference import ChatCompletionRequest + + +def _make_run(outputs: Path, name: str = "demorun") -> Path: + run = outputs / name + run.mkdir(parents = True) + (run / "adapter_config.json").write_text( + json.dumps({"base_model_name_or_path": "HuggingFaceTB/SmolLM-135M"}) + ) + ckpt = run / "checkpoint-1" + ckpt.mkdir() + (ckpt / "adapter_config.json").write_text("{}") + return run + + +@pytest.fixture +def captured(): + return {} + + +@pytest.fixture +def client(tmp_path, monkeypatch, captured): + outputs = tmp_path / "outputs" + _make_run(outputs) + + # resolve_preview_checkpoint -> resolve_output_dir -> outputs_root(). + from utils.paths import storage_roots as _sr + + monkeypatch.setattr(_sr, "outputs_root", lambda: outputs) + + async def _fake_load_model(load_req, request, subject): + captured["load_path"] = load_req.model_path + return None + + async def _fake_chat(payload, request, subject): + captured["payload"] = payload + return {"ok": True} + + monkeypatch.setattr(preview, "load_model", _fake_load_model) + monkeypatch.setattr(preview, "openai_chat_completions", _fake_chat) + + app = FastAPI() + app.include_router(preview.router, prefix = "/p") + app.dependency_overrides[preview.get_current_subject] = lambda: "admin" + # raise_server_exceptions=False so a 5xx surfaces as a response, not a throw. + return TestClient(app, raise_server_exceptions = False) + + +# ── Page rendering ──────────────────────────────────────────────────────── + + +def test_page_renders_with_csp(client): + r = client.get("/p/demorun") + assert r.status_code == 200 + assert "text/html" in r.headers["content-type"] + csp = r.headers.get("content-security-policy", "") + assert "default-src 'self'" in csp + assert "base-uri 'none'" in csp + + +def test_page_escapes_title(tmp_path, monkeypatch, captured): + outputs = tmp_path / "outputs" + # Run dir name carries an HTML-special char; the page must escape it. + _make_run(outputs, name = "a None. + outputs = tmp_path / "outputs" + merged = outputs / "mergedrun" + merged.mkdir(parents = True) + (merged / "config.json").write_text(json.dumps({"_name_or_path": "some/base"})) + + from utils.paths import storage_roots as _sr + + monkeypatch.setattr(_sr, "outputs_root", lambda: outputs) + + async def _fake_load(load_req, request, subject): + return None + + async def _fake_chat(payload, request, subject): + captured["payload"] = payload + return {"ok": True} + + monkeypatch.setattr(preview, "load_model", _fake_load) + monkeypatch.setattr(preview, "openai_chat_completions", _fake_chat) + + app = FastAPI() + app.include_router(preview.router, prefix = "/p") + c = TestClient(app, raise_server_exceptions = False) + r = c.post( + "/p/mergedrun/v1/chat/completions", + json = {"messages": [{"role": "user", "content": "hi"}], "use_adapter": False}, + ) + assert r.status_code == 200 + assert captured["payload"].use_adapter is None + + +# ── Streaming lock lifetime ────────────────────────────────────────────────── + + +def test_streaming_holds_lock_until_drained(tmp_path, monkeypatch, captured): + outputs = tmp_path / "outputs" + _make_run(outputs) + from utils.paths import storage_roots as _sr + + monkeypatch.setattr(_sr, "outputs_root", lambda: outputs) + + async def _fake_load_model(load_req, request, subject): + return None + + async def _gen(): + yield b"data: {}\n\n" + yield b"data: [DONE]\n\n" + + async def _fake_chat(payload, request, subject): + return StreamingResponse(_gen()) + + monkeypatch.setattr(preview, "load_model", _fake_load_model) + monkeypatch.setattr(preview, "openai_chat_completions", _fake_chat) + + async def _run(): + assert not preview._preview_lock.locked() + payload = ChatCompletionRequest(messages = [{"role": "user", "content": "hi"}]) + resp = await preview._serve_chat("demorun", None, payload, request = None) + # Lock must still be held: a second checkpoint must not swap the backend + # mid-stream. + assert preview._preview_lock.locked() + chunks = [c async for c in resp.body_iterator] + # Released only after the stream fully drains. + assert not preview._preview_lock.locked() + return chunks + + chunks = asyncio.run(_run()) + assert any(b"[DONE]" in c for c in chunks) + assert not preview._preview_lock.locked() diff --git a/studio/backend/utils/api_errors.py b/studio/backend/utils/api_errors.py index b1c55b61b9..cae8daf287 100644 --- a/studio/backend/utils/api_errors.py +++ b/studio/backend/utils/api_errors.py @@ -125,6 +125,12 @@ def is_anthropic_path(path: str) -> bool: return path.startswith("/v1/messages") +def wants_api_error_envelope(path: str) -> bool: + """True for the OpenAI/Anthropic-compatible surfaces: the ``/v1/*`` mount and + the preview ``/p/[/]/v1/*`` mount.""" + return path.startswith("/v1/") or (path.startswith("/p/") and "/v1/" in path) + + def error_body_for_path( path, message, @@ -183,15 +189,16 @@ def _summarize_validation_errors(errors) -> tuple: def install_api_error_handlers(app) -> None: """Register validation + HTTPException handlers that emit ``/v1/*`` envelopes. - Both handlers are global but only transform responses for paths starting with - ``/v1/``. Non-``/v1/`` paths reproduce FastAPI's default ``{"detail": ...}`` - behavior exactly so the Studio frontend keeps working. + Both handlers are global but only transform responses for OpenAI/Anthropic- + compatible surfaces (see :func:`wants_api_error_envelope`: the ``/v1/*`` mount + and the preview ``/p/.../v1/*`` mount). Every other path reproduces FastAPI's + default ``{"detail": ...}`` behavior exactly so the Studio frontend keeps working. """ @app.exception_handler(RequestValidationError) async def _handle_validation_error(request, exc): path = request.url.path - if path.startswith("/v1/"): + if wants_api_error_envelope(path): summary, param = _summarize_validation_errors(exc.errors()) return JSONResponse( status_code = 400, @@ -211,7 +218,7 @@ def install_api_error_handlers(app) -> None: # default http_exception_handler, which returns a bodiless Response. if not is_body_allowed_for_status_code(exc.status_code): return Response(status_code = exc.status_code, headers = headers) - if path.startswith("/v1/"): + if wants_api_error_envelope(path): detail = exc.detail # Already a fully-formed envelope: pass through untouched. if isinstance(detail, dict) and ("error" in detail or detail.get("type") == "error"): diff --git a/studio/backend/utils/models/checkpoints.py b/studio/backend/utils/models/checkpoints.py index 63f599d0df..d174f6677b 100644 --- a/studio/backend/utils/models/checkpoints.py +++ b/studio/backend/utils/models/checkpoints.py @@ -159,3 +159,64 @@ def scan_checkpoints( except Exception as e: logger.error(f"Error scanning checkpoints: {e}") return [] + + +def _is_model_dir(path: Path) -> bool: + return (path / "config.json").exists() or (path / "adapter_config.json").exists() + + +def has_preview_model(output_dir: Optional[str]) -> bool: + """True when ``output_dir`` holds a previewable root model (what ``/p/{run}`` + resolves). A cancelled run keeps ``output_dir`` but saves no root adapter.""" + if not output_dir: + return False + path = Path(output_dir) + return path.is_dir() and _is_model_dir(path) + + +def preview_ref(output_dir: Optional[str]) -> Optional[str]: + """``/p`` ref (``run`` or ``run/checkpoint``) relative to outputs_root, or None. + + Posix-joined so a nested output dir keeps a working link instead of collapsing + to its basename. None when not previewable, outside outputs_root, or deeper than + the two path segments the ``/p`` route matches (so the UI omits a dead link). + """ + if not has_preview_model(output_dir): + return None + try: + rel = Path(output_dir).resolve().relative_to(outputs_root().resolve()) + except (ValueError, OSError): + return None + parts = rel.parts + if not parts or len(parts) > 2: + return None + return "/".join(parts) + + +def resolve_preview_checkpoint(run: str, checkpoint: Optional[str] = None) -> Path: + relative = run if not checkpoint else f"{run}/{checkpoint}" + path = resolve_output_dir(relative) + if not path.is_dir() or not _is_model_dir(path): + raise FileNotFoundError( + f"No trained checkpoint at '{relative}'. Check the run/checkpoint name (see GET /p)." + ) + return path + + +def list_preview_targets(outputs_dir: str = str(outputs_root())) -> List[dict]: + targets: List[dict] = [] + for run_name, checkpoints, metadata in scan_checkpoints(outputs_dir): + for display_name, path, loss in checkpoints: + is_latest = display_name == run_name + checkpoint = None if is_latest else Path(path).name + targets.append( + { + "run": run_name, + "checkpoint": checkpoint, + "ref": run_name if is_latest else f"{run_name}/{checkpoint}", + "is_latest": is_latest, + "loss": loss, + "base_model": metadata.get("base_model"), + } + ) + return targets diff --git a/studio/frontend/src/features/studio/history-card-grid.tsx b/studio/frontend/src/features/studio/history-card-grid.tsx index dbe7fff01b..75ef1c2d85 100644 --- a/studio/frontend/src/features/studio/history-card-grid.tsx +++ b/studio/frontend/src/features/studio/history-card-grid.tsx @@ -24,7 +24,10 @@ import { useTrainingRuntimeStore, } from "@/features/training"; import { formatDuration } from "@/features/studio/sections/progress-section-lib"; +import { fetchDeviceType, usePlatformStore } from "@/config/env"; import { cn } from "@/lib/utils"; +import { copyToClipboard } from "@/lib/copy-to-clipboard"; +import { toast } from "@/lib/toast"; import { Delete02Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { type ReactElement, useCallback, useEffect, useRef, useState } from "react"; @@ -194,6 +197,28 @@ export function HistoryCardGrid({ const [manualFetchInFlight, setManualFetchInFlight] = useState(false); const { resumeTrainingRunFromHistory } = useTrainingActions(); const isStarting = useTrainingRuntimeStore((state) => state.isStarting); + // Copy-link base: Cloudflare tunnel > LAN host:port > origin. The tunnel + // registers shortly after startup, so poll (bounded) until it shows. + const cloudflareUrl = usePlatformStore((s) => s.cloudflareUrl); + const serverUrl = usePlatformStore((s) => s.serverUrl); + useEffect(() => { + if (cloudflareUrl) return; + let cancelled = false; + void (async () => { + for (let attempt = 0; attempt < 12 && !cancelled; attempt++) { + try { + await fetchDeviceType({ force: true }); + } catch { + // Ignore startup blips; copy-link falls back to serverUrl/origin. + } + if (cancelled || usePlatformStore.getState().cloudflareUrl) return; + await new Promise((r) => setTimeout(r, 2500)); + } + })(); + return () => { + cancelled = true; + }; + }, [cloudflareUrl]); const userControllerRef = useRef(null); const pollControllerRef = useRef(null); @@ -362,6 +387,8 @@ export function HistoryCardGrid({ const isRunning = run.status === "running"; const canResume = run.can_resume && !wasContinued; const isResuming = resumeTarget === run.id; + // Backend /p ref, gated on previewability + route-expressible depth. + const canCopyPreview = !!run.preview_ref; return (
onSelectRun(run.id)} onKeyDown={(e) => { @@ -411,6 +438,38 @@ export function HistoryCardGrid({ {isResuming ? t("studio.history.resuming") : t("studio.history.resumeTraining")} )} + {canCopyPreview && ( + + )}

{run.loss_sparkline && run.loss_sparkline.length >= 2 && ( -
+