* 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>
790 lines
27 KiB
Python
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
|