unsloth/studio/backend/main.py
Lee Jackson 61ed4cac51
Studio: persist chat history in backend storage (#5272)
* feat: Persist chat history in backend storage

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address chat tombstone batching review

* fix: update desktop auth routes stub

* chat db settings storage

* chat db settings routes

* chat db settings client

* chat db settings store

* chat db settings wiring

* chat db history storage

* chat db settings migration

* chat db settings fallback

* chat db container metadata

* chat db legacy migration fixes

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* chat ci auth background reads

* chat auth storage fixes

* chat migration final fixes

* chat export batch message lookup

* chat history review fixes

* chat prune sync fix

* chat settings hydration retry

* gate settings persistence

* Scope chat-history rows by subject; fix hijack, clear-confirm, hydrate race

Backend storage and routes:
- chat_threads / chat_messages / chat_settings carry a NOT NULL subject
  column with composite PRIMARY KEY (id, subject). Two authenticated
  identities can no longer see or wipe each other's data.
- Pre-existing rows on an existing studio.db migrate under sentinel
  subject __legacy_unscoped__ via rename + rebuild + copy; single-user
  installs see no behavior change.
- ON CONFLICT(id, subject) DO UPDATE ... WHERE chat_messages.thread_id =
  excluded.thread_id refuses cross-thread re-parenting via upsert.
  upsert_chat_message + sync_chat_messages now raise
  ChatMessageThreadMismatch which the routes map to HTTP 409.
- replace_thread_messages rejects body messages whose threadId does not
  match the URL thread (HTTP 400) instead of silently rewriting them.
- DELETE /api/chat requires ?confirm=true, returns row count, logs the
  subject and count.
- upsert_chat_settings_merge does read + deep-merge + write inside a
  single BEGIN IMMEDIATE so concurrent writers no longer drop each
  other's updates. The route delegates to this helper.
- New POST /api/chat/messages:batch returns {thread_id -> messages[]}
  for many threads in one HTTP call. Subject-scoped. Unknown ids return
  empty lists instead of 404 so the sidebar/search caller can rebuild
  atomically.

Frontend:
- chat-runtime-store: hydrate-failure catch sets settingsHydrated:true
  so a transient backend blip no longer permanently disables
  persistence. setParams bumps inferenceParamMutationVersions
  unconditionally so a slow hydration response cannot clobber a
  pre-hydrate user edit. saveSettingsPatch replaces the serial chain
  with a debounced pendingPatch + deep merge; flush on beforeunload.
- chat-history-storage: clearStoredChats returns ClearStoredChatsResult
  distinguishing backend / legacy / both outcomes.
  listStoredChatThreadsWithMessages uses the batched fetch (one HTTP
  call) instead of Promise.all per-thread; legacy Dexie fallback only
  fires when the batch result is empty.
- chat-api: batchListChatMessages with graceful 404 / 405 fallback to
  per-thread listChatMessages for older servers.
- chat-thread-tombstones: store {id, deletedAt} tuples with 90-day GC
  and a 5000-entry cap so localStorage stays bounded. Back-compat reads
  pre-fix plain strings. Adds removeChatThreadTombstones (rollback) and
  clearAllChatThreadTombstones (post-legacy-purge clean-up).
- use-chat-sidebar-items: deleteChatItem tombstones synchronously
  BEFORE the backend round-trip and rolls back on failure (restores
  pre-PR optimistic UX). 300 ms trailing debounce on
  CHAT_HISTORY_UPDATED_EVENT plus requestSeq guard so stream-time event
  bursts produce at most one fetch per quiet window.

Tests:
- studio/backend/tests/pr5272_sim/ adds 64 regression tests covering
  schema migration from pre-fix shape, subject scoping, cross-thread
  hijack, bulk-replace mismatch, clear-confirm, concurrent settings,
  unicode + 2MB content + SQL-injection-safe binding, chunking
  boundary at 900 and 901 ids, batched endpoint (multi-subject + 1200
  ids + per-thread order), and grep contracts for the frontend patches.
  test_chat_history_storage.py updated to pass subject.

Verified locally on Linux + macOS + Windows GitHub Actions runners
(staging fork): 64 pass + 2 from the PR's own backend test on all
three OSes.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Drop subject scoping and clear-confirm gate (Studio is single-user)

Per maintainer feedback: subject scoping, cross-thread message hijack
guard, and DELETE /api/chat ?confirm=true gate are unnecessary because
Studio is intentionally single-user (the client already shows a confirm
dialog before clear-all).

This commit reverts those backend changes and keeps only the
non-multi-user pieces from the earlier fix commit:

- studio_db.py: restored to pre-fix shape; adds upsert_chat_settings_merge
  which does atomic read + deep-merge + write under BEGIN IMMEDIATE so
  two concurrent slider drags cannot drop one another's updates.
- routes/chat_history.py: restored; put_settings now calls the atomic
  merge instead of doing the read-merge-write across three separate
  connections. Adds POST /api/chat/messages:batch to collapse the
  sidebar/search rebuild from N round-trips to 1.
- frontend/api/chat-api.ts: align batchListChatMessages request and
  response keys with the backend (threadIds / messagesByThreadId).
- tests/test_chat_history_storage.py: add atomic-merge concurrency test,
  deep-merge nested-key test, and 901-id chunking-boundary test.
- Drop the pr5272_sim test directory (those tests covered the reverted
  subject-scoping/hijack/confirm behavior).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix sidebar delete crash, keepalive on settings beforeunload flush, search rebuild race

Two correctness bugs and one perf race surfaced by a fresh code review of
the prior fix commit:

- chat-api.ts: notifyChatHistoryUpdated was declared as a non-exported
  function, but use-chat-sidebar-items.ts imports it. The import would
  fail tsc with TS2305 and at runtime the optimistic-delete and
  delete-failure rollback paths would both throw.
- chat-runtime-store.ts + chat-settings-api.ts + chat-settings-storage.ts:
  the beforeunload settings flush is now actually keepalive. Without it
  the browser cancels the in-flight PUT on tab close, so the last slider
  drag is silently dropped (which is exactly the case the
  debounce+beforeunload combination was meant to protect against).
- use-chat-search-index.ts: rebuilds now coalesce with a 300ms trailing
  debounce and discard out-of-order responses via a requestSeq guard.
  Matches the sibling pattern in use-chat-sidebar-items.ts so two rapid
  CHAT_HISTORY_UPDATED_EVENTs (run-start + run-end save during a turn)
  cannot land with stale data winning.
- chat-thread-tombstones.ts: drop dead clearAllChatThreadTombstones with
  no call sites; Dexie is never wiped so the function has no use.

* fix(studio): protect chat persistence writes

* fix(studio): align chat history clear semantics

* fix(studio): show partial chat clear feedback

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix(studio): preserve chat persistence fallbacks

* fix(studio): harden chat thread persistence checks

* Preserve chat message timestamps

* Gate chat stream on history save

* Make chat thread backfill best effort

* Avoid chat message 404 probe

* Tighten chat legacy fallbacks

* chat: server-side ledger so legacy Dexie import is recoverable

The boolean localStorage sentinel
(unsloth_chat_legacy_imported_to_studio_db) made importLegacyChatsIfNeeded
non-recoverable: deleting studio.db while the browser keeps the flag
silently hides every legacy Dexie thread from the sidebar (verified by
the 3-GPU validation probe; matches the third review comment on PR
#5272). Same trap fires for browser-profile sync to a fresh machine
and any other path that wipes studio.db while keeping IndexedDB.

Source of truth moves into studio.db itself via a new
chat_legacy_import_log table keyed by legacy thread id. The ledger
disappears together with studio.db, so the next launch re-runs the
import from whatever Dexie still holds. localStorage stays as a
per-session perf hint only.

Performance, all bounded by the three new fast-paths before any
backend work:

  A) localStorage hint says "imported earlier in this session" -- 0
     network, ~0 ms. Covers the warm sidebar mount.

  B) indexedDB.databases() reports no "unsloth-chat" DB -- 0 network,
     ~1 ms. Covers every new user who never had the old browser-only
     Studio (the common case after launch).

  C) db.threads.count() + db.messages.count() are both 0 -- 0 network,
     ~5 ms. Covers returning users who migrated long ago and Dexie was
     never repopulated.

Only when all three miss does the code talk to the backend
(GET /api/chat/import-ledger -> diff vs Dexie -> existing import path
-> POST /api/chat/import-ledger to record what was just imported).
Per-thread tracking is enough because Dexie is read-only after this
PR; a thread's message set does not grow.

Backend deployments that predate the import-ledger routes are
handled transparently: the client treats 404/405 as an empty ledger
and re-runs the (idempotent via UPSERT) import on next launch.

Changes:
- storage/studio_db.py: new chat_legacy_import_log table (WITHOUT
  ROWID, PK on legacy_thread_id) + list_chat_legacy_import_log() +
  record_chat_legacy_import_log() (idempotent batch UPSERT).
- routes/chat_history.py: GET + POST /api/chat/import-ledger with the
  obvious request/response models.
- frontend api/chat-api.ts: listChatImportLedger() (returns a Set for
  O(1) diff) + recordChatImportLedger(), both with 404/405 fallback.
- frontend utils/chat-history-storage.ts: importLegacyChatsIfNeeded
  gains three fast-paths, ledger fetch on the slow path, and writes
  the ledger after a successful import. The localStorage helper is
  unchanged on the surface; it just stops being authoritative.
- tests: 5 new test_legacy_import_log_* cases (empty default, record
  + list round-trip, idempotency, input dedup, empty/null ignore).
  All 9 pre-existing tests still pass.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Make the legacy-import recovery actually recoverable

The previous commit added a server-side ledger to make Dexie -> studio.db
import recoverable after a studio.db wipe, but the localStorage perf hint
still short-circuited the import gate before the ledger was ever consulted.
After a wipe, the hint stayed "true" and the bulk re-import never ran -- the
ledger sat empty and only the per-thread lazy materialize-on-continue path
restored data.

Changes:

- Remove the localStorage short-circuit from importLegacyChatsIfNeeded so
  the ledger is checked on every fresh tab. legacyChatImportPromise keeps
  the per-session cache; the hint now only matters for the listing paths.
- Batch the slow path: one db.messages.where().anyOf().toArray() and one
  batchListChatMessages() instead of 2N round-trips. At 1k threads this
  drops a multi-second blocking import to a single request pair.
- recordChatImportLedger returns {accepted, inserted, supported}. The
  localStorage hint is only flipped when supported is true, so old
  backends (404 / 405 / 501) no longer permanently poison recovery.
- Ledger backfill: threads already present in chat_threads but missing
  from the ledger now get added too, so old-FE-then-new-FE deployments
  don't redo the diff every launch.
- Backend response field renamed recorded -> {accepted, inserted}.
  accepted is the deduped non-empty input count; inserted is the rows
  actually new (via INSERT ... RETURNING). Bounded by Field(max_length=
  10_000) on the request payload.
- Storage helpers renamed: chat_legacy_import_log -> chat_legacy_imports,
  record_* -> upsert_* to match the existing noun/verb conventions.
- DEXIE_DB_NAME exported from db.ts; duplicate constant in
  chat-history-storage.ts removed.
- 3 new route-level tests for /api/chat/import-ledger covering the
  round-trip, the (accepted, inserted) split, and the 10k payload cap.

All 18 chat-history tests pass.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: shine1i <wasimysdev@gmail.com>
Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-05-22 06:18:05 -07:00

790 lines
27 KiB
Python

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Main FastAPI application for Unsloth UI Backend
"""
import os
import sys
from pathlib import Path as _Path
# Suppress annoying C-level dependency warnings globally
os.environ["PYTHONWARNINGS"] = "ignore"
# Ensure backend dir is on sys.path so _platform_compat is importable when
# main.py is launched directly (e.g. `uvicorn main:app`).
_backend_dir = str(_Path(__file__).parent)
if _backend_dir not in sys.path:
sys.path.insert(0, _backend_dir)
# Fix for Anaconda/conda-forge Python: seed platform._sys_version_cache before
# any library imports that trigger attrs -> rich -> structlog -> platform crash.
# See: https://github.com/python/cpython/issues/102396
import _platform_compat # noqa: F401
# Direct `uvicorn main:app` launches bypass run.py, so re-export here too
# (mirrors run.py). Required BEFORE the unsloth-zoo import below, since
# its LLAMA_CPP_DEFAULT_DIR binding is import-time.
from utils.paths.storage_roots import studio_root as _studio_root
try:
_LEGACY_STUDIO_ROOT = (_Path.home() / ".unsloth" / "studio").resolve()
except (OSError, ValueError):
_LEGACY_STUDIO_ROOT = _Path.home() / ".unsloth" / "studio"
try:
_STUDIO_ROOT_RESOLVED = _studio_root().resolve()
except (OSError, ValueError):
_STUDIO_ROOT_RESOLVED = _studio_root()
if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT:
if not os.environ.get("UNSLOTH_STUDIO_HOME"):
os.environ["UNSLOTH_STUDIO_HOME"] = str(_STUDIO_ROOT_RESOLVED)
if not os.environ.get("UNSLOTH_LLAMA_CPP_PATH"):
os.environ["UNSLOTH_LLAMA_CPP_PATH"] = str(_STUDIO_ROOT_RESOLVED / "llama.cpp")
import hashlib
import mimetypes
import re as _re
import shutil
import warnings
from contextlib import asynccontextmanager
from importlib.metadata import PackageNotFoundError, version as package_version
_STUDIO_INSTALL_ID_RE = _re.compile(r"^[0-9a-f]{64}$")
def _read_studio_install_id() -> str:
"""Per-install opaque id written by install.sh / install.ps1 at
$STUDIO_HOME/share/studio_install_id. Returns "" when the file is
absent (pre-PR install, fresh tree never run through the installer)
or contains anything other than a 64-char lowercase-hex token --
in which case /api/health emits "" and the launcher's _check_health
falls back to the existing "no baked id, accept any healthy
Unsloth backend" path. This intentionally replaces a previous
sha256(resolved_install_path) so the field carries no install-path
information for callers reaching /api/health (relevant when Studio
is run with -H 0.0.0.0)."""
try:
token = (
(_STUDIO_ROOT_RESOLVED / "share" / "studio_install_id").read_text().strip()
)
except (OSError, ValueError):
return ""
return token if _STUDIO_INSTALL_ID_RE.fullmatch(token) else ""
_STUDIO_ROOT_ID_CACHE: str = _read_studio_install_id()
def _studio_root_id() -> str:
"""Same-install discriminator for /api/health: a per-install opaque
token written once by the installer and read once at module import.
Empty when no installer-written token is present; the launcher
contract treats "" as "no baked id, accept any healthy backend"."""
return _STUDIO_ROOT_ID_CACHE
# Fix broken Windows registry MIME types. Some Windows installs map .js to
# "text/plain" in the registry (HKCR\.js\Content Type). Python's mimetypes
# module reads from the registry, and FastAPI/Starlette's StaticFiles uses
# mimetypes.guess_type() to set Content-Type headers. Browsers enforce strict
# MIME checking for ES module scripts (<script type="module">) and will refuse
# to execute .js files served as text/plain — resulting in a blank page.
# Calling add_type() *before* StaticFiles is instantiated ensures the correct
# types are used regardless of the OS registry.
if sys.platform == "win32":
mimetypes.add_type("application/javascript", ".js")
mimetypes.add_type("text/css", ".css")
# Suppress annoying dependency warnings in production
if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
warnings.filterwarnings("ignore")
# Alternatively, you can be more specific:
# warnings.filterwarnings("ignore", category=DeprecationWarning)
# warnings.filterwarnings("ignore", module="triton.*")
from fastapi import Depends, FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, HTMLResponse, Response
from pathlib import Path
from datetime import datetime
# Import routers
from routes import (
auth_router,
chat_history_router,
data_recipe_router,
datasets_router,
export_router,
inference_router,
inference_studio_router,
models_router,
providers_router,
training_history_router,
training_router,
)
from auth import storage
from auth.authentication import get_current_subject
from utils.hardware import (
detect_hardware,
get_device,
DeviceType,
get_backend_visible_gpu_info,
)
import utils.hardware.hardware as _hw_module
from utils.cache_cleanup import clear_unsloth_compiled_cache
from utils.native_path_leases import native_path_leases_supported
from utils.update_status import (
get_studio_install_source_status,
get_studio_update_status,
)
from utils.studio_version import get_studio_version
def get_unsloth_version() -> str:
try:
return package_version("unsloth")
except PackageNotFoundError:
pass
version_file = (
_Path(__file__).resolve().parents[2] / "unsloth" / "models" / "_utils.py"
)
try:
for line in version_file.read_text(encoding = "utf-8").splitlines():
if line.startswith("__version__ = "):
return line.split("=", 1)[1].strip().strip('"').strip("'")
except OSError:
pass
return "dev"
UNSLOTH_VERSION = get_unsloth_version()
STUDIO_VERSION = get_studio_version()
def _load_desktop_owner() -> dict[str, str] | None:
token = os.environ.pop("UNSLOTH_STUDIO_DESKTOP_OWNER_TOKEN", "")
kind = os.environ.pop("UNSLOTH_STUDIO_DESKTOP_OWNER_KIND", "")
if kind != "tauri" or not token:
return None
return {
"kind": "tauri",
"token_sha256": hashlib.sha256(token.encode("utf-8")).hexdigest(),
}
_DESKTOP_OWNER = _load_desktop_owner()
def _desktop_owner() -> dict[str, str] | None:
return _DESKTOP_OWNER
@asynccontextmanager
async def lifespan(app: FastAPI):
"""Startup: detect hardware, seed default admin if needed. Shutdown: clean up compiled cache."""
# Clean up any stale compiled cache from previous runs
clear_unsloth_compiled_cache()
# Remove stale .venv_overlay from previous versions — no longer used.
# Version switching now uses .venv_t5/ (pre-installed by setup.sh).
overlay_dir = Path(__file__).resolve().parent.parent.parent / ".venv_overlay"
if overlay_dir.is_dir():
shutil.rmtree(overlay_dir, ignore_errors = True)
# Detect hardware first — sets DEVICE global used everywhere
detect_hardware()
# llama.cpp probes: capability (MTP support) + freshness (release age).
# Both cached; freshness has a 24h disk TTL.
try:
from core.inference.llama_cpp import LlamaCppBackend
from utils.llama_cpp_freshness import (
check_prebuilt_freshness,
format_stale_warning,
)
_bin = LlamaCppBackend._find_llama_server_binary()
_caps = LlamaCppBackend.probe_server_capabilities(_bin)
app.state.llama_cpp_capabilities = _caps
_freshness = check_prebuilt_freshness(_bin)
app.state.llama_cpp_freshness = _freshness
import structlog as _structlog
_log = _structlog.get_logger(__name__)
if _caps.get("found") and not _caps.get("supports_mtp"):
_msg = (
"llama.cpp prebuilt lacks MTP support "
"(--spec-type mtp/draft-mtp). Run `unsloth studio update`. "
"MTP GGUFs will load without speculative decoding."
)
_log.warning(_msg)
print(f"WARNING: {_msg}", flush = True)
if _freshness.get("stale"):
_msg = format_stale_warning(_freshness)
_log.warning(_msg)
print(f"WARNING: {_msg}", flush = True)
except Exception as _probe_exc:
import structlog as _structlog
_structlog.get_logger(__name__).debug(
"llama.cpp startup probes failed: %s", _probe_exc
)
from storage.studio_db import cleanup_orphaned_runs
try:
cleanup_orphaned_runs()
except Exception as exc:
import structlog
structlog.get_logger(__name__).warning(
"cleanup_orphaned_runs failed at startup: %s", exc
)
# Pre-cache the helper GGUF model for LLM-assisted dataset detection.
# Runs in a background thread so it doesn't block server startup.
import threading
def _precache():
try:
from utils.datasets.llm_assist import precache_helper_gguf
precache_helper_gguf()
except Exception:
pass # non-critical
threading.Thread(target = _precache, daemon = True).start()
# Initialize RSA key pair for API key encryption (external providers)
from core.inference.key_exchange import init_key_pair
init_key_pair()
if storage.ensure_default_admin():
bootstrap_pw = storage.get_bootstrap_password()
app.state.bootstrap_password = bootstrap_pw
bootstrap_path = storage.DB_PATH.parent / ".bootstrap_password"
print("\n" + "=" * 60)
print("DEFAULT ADMIN ACCOUNT CREATED")
print(f" username: {storage.DEFAULT_ADMIN_USERNAME}")
print(f" password saved to: {bootstrap_path}")
print(" Open the Studio UI to sign in and change it.")
print("=" * 60 + "\n")
else:
app.state.bootstrap_password = storage.get_bootstrap_password()
yield
# Cleanup
_hw_module.DEVICE = None
clear_unsloth_compiled_cache()
# Create FastAPI app
app = FastAPI(
title = "Unsloth UI Backend",
version = UNSLOTH_VERSION,
description = "Backend API for Unsloth UI - Training and Model Management",
lifespan = lifespan,
)
# Initialize structured logging
from loggers.config import LogConfig
from loggers.handlers import LoggingMiddleware
logger = LogConfig.setup_logging(
service_name = "unsloth-studio-backend",
env = os.getenv("ENVIRONMENT_TYPE", "production"),
)
app.add_middleware(LoggingMiddleware)
# Citation favicons load from www.google.com/s2/favicons; *.gstatic.com is
# kept for legacy web-search faviconV2 paths. Everything else is same-origin.
from starlette.middleware.base import BaseHTTPMiddleware # noqa: E402
from starlette.requests import Request as _StarletteRequest # noqa: E402
_CSP_SCRIPT_NONCE_HEADER = "x-internal-script-nonce"
def _build_csp(script_nonce: "str | None" = None) -> str:
script_src = "script-src 'self'"
if script_nonce:
script_src += f" 'nonce-{script_nonce}'"
return (
"default-src 'self'; "
"img-src 'self' data: blob: https://t0.gstatic.com "
"https://t1.gstatic.com https://t2.gstatic.com "
"https://t3.gstatic.com https://www.google.com; "
"connect-src 'self' https://huggingface.co https://datasets-server.huggingface.co; "
"style-src 'self' 'unsafe-inline'; "
f"{script_src}; "
"font-src 'self' data:; "
"frame-ancestors 'none'; "
"form-action 'self'; "
"base-uri 'self'"
)
class SecurityHeadersMiddleware(BaseHTTPMiddleware):
"""Set baseline security headers; splice per-response inline-script nonces into CSP."""
async def dispatch(self, request: _StarletteRequest, call_next):
response = await call_next(request)
# Strip the internal nonce hand-off header so it never reaches the client.
nonce = response.headers.get(_CSP_SCRIPT_NONCE_HEADER)
if nonce is not None:
del response.headers[_CSP_SCRIPT_NONCE_HEADER]
response.headers.setdefault("Content-Security-Policy", _build_csp(nonce))
response.headers.setdefault("X-Frame-Options", "DENY")
response.headers.setdefault("X-Content-Type-Options", "nosniff")
response.headers.setdefault("Referrer-Policy", "no-referrer")
response.headers.setdefault(
"Permissions-Policy",
"camera=(), microphone=(), geolocation=(), interest-cohort=()",
)
response.headers["server"] = "unsloth-studio"
return response
app.add_middleware(SecurityHeadersMiddleware)
# Cap upload body on protected POSTs; default 500 MB, env-tunable.
import json as _json_for_413 # noqa: E402
_MAX_BODY_BYTES = int(os.environ.get("UNSLOTH_STUDIO_MAX_BODY_MB", "500")) * 1024 * 1024
_BODY_PROTECTED_PREFIXES = (
"/v1/chat/completions",
"/v1/completions",
"/api/inference",
"/api/data-recipe",
"/api/datasets",
"/api/chat",
"/api/train",
"/api/export",
)
async def _send_413(send, total_bytes: int) -> None:
payload = _json_for_413.dumps(
{
"detail": (
f"Request body too large "
f"({total_bytes:,} bytes; max {_MAX_BODY_BYTES:,})."
)
},
).encode("utf-8")
await send(
{
"type": "http.response.start",
"status": 413,
"headers": [
(b"content-type", b"application/json"),
(b"content-length", str(len(payload)).encode("ascii")),
],
}
)
await send({"type": "http.response.body", "body": payload, "more_body": False})
class MaxBodyMiddleware:
"""Reject oversized bodies on protected POST/PUT/PATCH; raw ASGI so chunked uploads cannot bypass the cap."""
def __init__(self, app, max_bytes: int, protected_prefixes: tuple):
self.app = app
self.max_bytes = max_bytes
self.protected_prefixes = protected_prefixes
async def __call__(self, scope, receive, send):
if scope["type"] != "http":
await self.app(scope, receive, send)
return
method = scope.get("method", "").upper()
path = scope.get("path", "")
if method not in ("POST", "PUT", "PATCH") or not any(
path.startswith(p) for p in self.protected_prefixes
):
await self.app(scope, receive, send)
return
declared = None
for name, value in scope.get("headers", []):
if name == b"content-length":
try:
declared = int(value.decode("latin-1"))
except (ValueError, UnicodeDecodeError):
declared = None
break
if declared is not None and declared > self.max_bytes:
await _send_413(send, declared)
return
chunks: list = []
total = 0
while True:
msg = await receive()
mtype = msg.get("type")
if mtype == "http.disconnect":
return
if mtype != "http.request":
# Mid-stream unexpected frame: forwarding would corrupt downstream.
return
body = msg.get("body", b"") or b""
if body:
total += len(body)
if total > self.max_bytes:
await _send_413(send, total)
return
chunks.append(body)
if not msg.get("more_body", False):
break
replayed = {"sent": False}
async def replay_receive():
if not replayed["sent"]:
replayed["sent"] = True
return {
"type": "http.request",
"body": b"".join(chunks),
"more_body": False,
}
# After replay, fall through so http.disconnect still propagates.
return await receive()
await self.app(scope, replay_receive, send)
app.add_middleware(
MaxBodyMiddleware,
max_bytes = _MAX_BODY_BYTES,
protected_prefixes = _BODY_PROTECTED_PREFIXES,
)
from starlette.responses import RedirectResponse as _RedirectResponse # noqa: E402
@app.get("/recipes", include_in_schema = False)
@app.get("/recipes/{rest:path}", include_in_schema = False)
async def _recipes_redirect(rest: str = ""):
target = "/data-recipes" + (("/" + rest) if rest else "")
return _RedirectResponse(url = target, status_code = 308)
# CORS middleware
_api_only = os.environ.get("UNSLOTH_API_ONLY") == "1"
_cors_origins = ["*"]
if _api_only:
_cors_origins = [
"tauri://localhost", # Linux/macOS Tauri webview
"http://tauri.localhost", # Windows Tauri webview
"http://localhost", # dev fallback
"http://localhost:5173", # Tauri dev/Vite
"http://127.0.0.1:5173", # Tauri dev/Vite fallback
]
_cors_origin_regex = None
else:
_cors_origin_regex = None
app.add_middleware(
CORSMiddleware,
allow_origins = _cors_origins,
allow_origin_regex = _cors_origin_regex,
allow_credentials = True,
allow_methods = ["*"],
allow_headers = ["*"],
)
# ============ Register API Routes ============
# Register routers
app.include_router(auth_router, prefix = "/api/auth", tags = ["auth"])
app.include_router(training_router, prefix = "/api/train", tags = ["training"])
app.include_router(models_router, prefix = "/api/models", tags = ["models"])
app.include_router(chat_history_router, prefix = "/api/chat", tags = ["chat"])
app.include_router(inference_router, prefix = "/api/inference", tags = ["inference"])
# Studio-only inference endpoints (cancel, etc.) are intentionally NOT
# exposed on the /v1 OpenAI-compat prefix below.
app.include_router(inference_studio_router, prefix = "/api/inference", tags = ["inference"])
# OpenAI-compatible endpoints: mount the same inference router at /v1
# so external tools (Open WebUI, SillyTavern, etc.) can use the
# standard /v1/chat/completions path.
app.include_router(inference_router, prefix = "/v1", tags = ["openai-compat"])
app.include_router(providers_router, prefix = "/api/providers", tags = ["providers"])
app.include_router(datasets_router, prefix = "/api/datasets", tags = ["datasets"])
app.include_router(data_recipe_router, prefix = "/api/data-recipe", tags = ["data-recipe"])
app.include_router(export_router, prefix = "/api/export", tags = ["export"])
app.include_router(
training_history_router, prefix = "/api/train", tags = ["training-history"]
)
# ============ Health and System Endpoints ============
@app.get("/api/health")
async def health_check(request: Request):
"""Liveness plus launcher capability bits; install fingerprint gated on a valid bearer.
Unauthenticated callers (Tauri watchdog, frontend bootstrap polls) need
``service`` / ``studio_root_id`` / ``chat_only`` / ``desktop_*`` / ``native_path_leases_supported``
to (a) re-adopt a sibling backend across restarts and (b) gate UI surfaces
before any token is available. None of those leak install path or version.
``version`` / ``studio_version`` / ``device_type`` still require a bearer
because they fingerprint the host.
"""
base = {
"status": "healthy",
"timestamp": datetime.now().isoformat(),
"service": "Unsloth UI Backend",
"chat_only": _hw_module.CHAT_ONLY,
"desktop_protocol_version": 1,
"desktop_manageability_version": 1,
"supports_desktop_auth": True,
"supports_desktop_backend_ownership": True,
# Opaque per-install id; launchers reject sibling Studios on the same port.
"studio_root_id": _studio_root_id(),
"native_path_leases_supported": native_path_leases_supported(),
**({"desktop_owner": owner} if (owner := _desktop_owner()) else {}),
}
auth = request.headers.get("authorization", "")
if not auth.lower().startswith("bearer "):
return base
try:
from auth.authentication import get_current_subject as _gcs
from fastapi.security import HTTPAuthorizationCredentials
creds = HTTPAuthorizationCredentials(
scheme = "Bearer", credentials = auth.split(" ", 1)[1]
)
# Must await: a bare coroutine is truthy and would skip the auth check.
subject = await _gcs(creds)
except HTTPException:
return base
except Exception:
return base
if not subject:
return base
platform_map = {"darwin": "mac", "win32": "windows", "linux": "linux"}
device_type = platform_map.get(sys.platform, sys.platform)
return {
**base,
"version": UNSLOTH_VERSION,
"studio_version": STUDIO_VERSION,
"device_type": device_type,
}
@app.get("/api/studio/install-source")
def studio_install_source(_current_subject: str = Depends(get_current_subject)):
"""Return source-aware install metadata without remote update checks."""
return get_studio_install_source_status(UNSLOTH_VERSION)
@app.get("/api/studio/update-status")
def studio_update_status(_current_subject: str = Depends(get_current_subject)):
"""Return source-aware manual update status for browser-served Studio."""
return get_studio_update_status(UNSLOTH_VERSION)
@app.post("/api/shutdown")
async def shutdown_server(
request: Request,
current_subject: str = Depends(get_current_subject),
):
"""Gracefully shut down the Unsloth Studio server.
Called by the frontend quit dialog so users can stop the server from the UI
without needing to use the CLI or kill the process manually.
"""
import asyncio
async def _delayed_shutdown():
await asyncio.sleep(0.2) # Let the HTTP response return first
trigger = getattr(request.app.state, "trigger_shutdown", None)
if trigger is not None:
trigger()
else:
# Fallback when not launched via run_server() (e.g. direct uvicorn)
import signal
import os
os.kill(os.getpid(), signal.SIGTERM)
request.app.state._shutdown_task = asyncio.create_task(_delayed_shutdown())
return {"status": "shutting_down"}
@app.get("/api/system")
async def get_system_info(
current_subject: str = Depends(get_current_subject),
):
"""Get system information.
Gated behind auth: the response includes platform, Python version,
GPU name, memory total, and ML package set -- enough to fingerprint
a host. Studio's chat-only-mode design assumes only the local user
reaches /api/system; in -H 0.0.0.0 / Colab / Tauri-relayed setups
that assumption breaks unless we require a bearer.
"""
import platform
import psutil
from utils.hardware import get_device
from utils.hardware.hardware import _backend_label
visibility_info = get_backend_visible_gpu_info()
gpu_info = {
"available": visibility_info["available"],
"devices": visibility_info["devices"],
}
# CPU & Memory
memory = psutil.virtual_memory()
return {
"platform": platform.platform(),
"python_version": platform.python_version(),
# Use the centralized _backend_label helper so the /api/system
# endpoint reports "rocm" on AMD hosts instead of "cuda", matching
# the /api/hardware and /api/gpu-visibility endpoints.
"device_backend": _backend_label(get_device()),
"cpu_count": psutil.cpu_count(),
"memory": {
"total_gb": round(memory.total / 1e9, 2),
"available_gb": round(memory.available / 1e9, 2),
"percent_used": memory.percent,
},
"gpu": gpu_info,
}
@app.get("/api/system/gpu-visibility")
async def get_gpu_visibility(
current_subject: str = Depends(get_current_subject),
):
return get_backend_visible_gpu_info()
@app.get("/api/system/hardware")
async def get_hardware_info(
current_subject: str = Depends(get_current_subject),
):
"""Return GPU name, total VRAM, and key ML package versions.
Gated behind auth alongside /api/system -- same fingerprinting
concern. /api/system/gpu-visibility is also auth-gated already.
"""
from utils.hardware import get_gpu_summary, get_package_versions
return {
"gpu": get_gpu_summary(),
"versions": get_package_versions(),
}
# ============ Serve Frontend (Optional) ============
def _strip_crossorigin(html_bytes: bytes) -> bytes:
"""Remove ``crossorigin`` attributes from script/link tags.
Vite adds ``crossorigin`` by default which forces CORS mode on font
subresource loads. When Studio is served over plain HTTP, Firefox
HTTPS-Only Mode does not exempt CORS font requests -- causing all
@font-face downloads to fail silently. Stripping the attribute
makes them regular same-origin fetches that work on any protocol.
"""
import re as _re
html = html_bytes.decode("utf-8")
html = _re.sub(r'\s+crossorigin(?:="[^"]*")?', "", html)
return html.encode("utf-8")
def _inject_bootstrap(html_bytes: bytes, app: FastAPI):
"""Inject bootstrap credentials when password change is pending.
Returns ``(html_bytes, script_nonce_or_None)``. Callers must forward
the nonce via ``_CSP_SCRIPT_NONCE_HEADER`` so the inline script is
not blocked by CSP.
"""
import json as _json
import secrets as _secrets
if not storage.requires_password_change(storage.DEFAULT_ADMIN_USERNAME):
return html_bytes, None
bootstrap_pw = getattr(app.state, "bootstrap_password", None)
if not bootstrap_pw:
return html_bytes, None
payload = _json.dumps(
{
"username": storage.DEFAULT_ADMIN_USERNAME,
"password": bootstrap_pw,
}
)
nonce = _secrets.token_urlsafe(16)
tag = f'<script nonce="{nonce}">window.__UNSLOTH_BOOTSTRAP__={payload}</script>'
html = html_bytes.decode("utf-8")
html = html.replace("</head>", f"{tag}</head>", 1)
return html.encode("utf-8"), nonce
def setup_frontend(app: FastAPI, build_path: Path):
"""Mount frontend static files (optional)"""
if not build_path.exists():
return False
# Mount assets
assets_dir = build_path / "assets"
if assets_dir.exists():
app.mount("/assets", StaticFiles(directory = assets_dir), name = "assets")
def _build_index_response() -> Response:
content = (build_path / "index.html").read_bytes()
content = _strip_crossorigin(content)
content, nonce = _inject_bootstrap(content, app)
headers = {"Cache-Control": "no-cache, no-store, must-revalidate"}
if nonce:
headers[_CSP_SCRIPT_NONCE_HEADER] = nonce
return Response(
content = content,
media_type = "text/html",
headers = headers,
)
@app.get("/")
async def serve_root():
return _build_index_response()
@app.get("/{full_path:path}")
async def serve_frontend(full_path: str):
if full_path in {"api", "v1"} or full_path.startswith(("api/", "v1/")):
return {"error": "API endpoint not found"}
file_path = (build_path / full_path).resolve()
# Block path traversal — ensure resolved path stays inside build_path
if not file_path.is_relative_to(build_path.resolve()):
return Response(status_code = 403)
if file_path.is_file():
return FileResponse(file_path)
# Serve index.html as bytes — avoids Content-Length mismatch
return _build_index_response()
return True