b91116caccThe studio backend lazily imports unsloth_zoo submodules (export's llama_cpp, hardware's vllm_utils, mlx, chat_templates). Each imports the top-level unsloth_zoo, which raises 'Please install Unsloth via pip install unsloth' unless UNSLOTH_IS_PRESENT is set (normally by import unsloth). Only hardware.py set it per-site; the rest crash on a clean install, and it surfaced on Windows. Set it once at backend startup (as unsloth does on import) so every site, and the DiffusionGemma runner, works across platforms. Follow-up to #6259 (which covered only the shim subprocess env).
1156 lines
43 KiB
Python
1156 lines
43 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
|
|
import threading
|
|
from pathlib import Path as _Path
|
|
import asyncio
|
|
from dataclasses import asdict
|
|
|
|
# Suppress C-level dependency warnings globally
|
|
os.environ["PYTHONWARNINGS"] = "ignore"
|
|
|
|
# ── Windows AMD ROCm DLL injection ──────────────────────────────────────────
|
|
# Python 3.8+ ignores PATH for extension modules; register ROCm bin dirs with
|
|
# os.add_dll_directory() so amdhip64.dll etc. are found before any torch import.
|
|
if sys.platform == "win32":
|
|
# Retained at module scope; os.add_dll_directory returns a handle that
|
|
# removes the search-path entry when garbage collected.
|
|
_ROCM_DLL_HANDLES: list = []
|
|
|
|
def _add_rocm_dll_dirs() -> None:
|
|
candidates = []
|
|
# 1. HIP_PATH / ROCM_PATH set by the AMD HIP SDK installer
|
|
for _var in ("HIP_PATH", "ROCM_PATH"):
|
|
_val = os.environ.get(_var)
|
|
if _val:
|
|
candidates.append(os.path.join(_val, "bin"))
|
|
# 2. AMD installer: C:\Program Files\AMD\ROCm\<ver>\bin, newest first.
|
|
_default_root = os.path.join(
|
|
os.environ.get("ProgramFiles", r"C:\Program Files"), "AMD", "ROCm"
|
|
)
|
|
|
|
def _ver_key(name: str) -> tuple:
|
|
# Numeric tuple key so "10.0" sorts after "7.0"; non-numeric chunks fall back to string
|
|
parts = []
|
|
for chunk in name.split("."):
|
|
try:
|
|
parts.append((0, int(chunk)))
|
|
except ValueError:
|
|
parts.append((1, chunk))
|
|
return tuple(parts)
|
|
|
|
try:
|
|
if os.path.isdir(_default_root):
|
|
for _ver in sorted(os.listdir(_default_root), key = _ver_key, reverse = True):
|
|
_bin = os.path.join(_default_root, _ver, "bin")
|
|
if os.path.isdir(_bin):
|
|
candidates.append(_bin)
|
|
except OSError:
|
|
pass
|
|
for _d in candidates:
|
|
if os.path.isdir(_d):
|
|
try:
|
|
_ROCM_DLL_HANDLES.append(os.add_dll_directory(_d))
|
|
except (OSError, AttributeError):
|
|
pass
|
|
|
|
_add_rocm_dll_dirs()
|
|
del _add_rocm_dll_dirs
|
|
|
|
# ── Windows AMD ROCm: make hipInfo.exe resolvable for subprocess probes ──
|
|
# bitsandbytes' get_rocm_gpu_arch() runs `hipinfo.exe` via PATH at import
|
|
# time; the AMD torch wheel ships it in the venv Scripts dir, which is on
|
|
# PATH only when the venv is activated -- Studio launches python directly.
|
|
# Without this, every bitsandbytes import logs a scary (but harmless)
|
|
# "Could not detect ROCm GPU architecture: [WinError 2]" ERROR + WARNING.
|
|
# Gated on the file existing: only AMD ROCm wheels ship hipInfo.exe, so
|
|
# NVIDIA/CPU hosts are untouched. os.add_dll_directory above does not help
|
|
# here -- subprocess PATH resolution ignores DLL search directories.
|
|
_scripts_dir = os.path.dirname(sys.executable)
|
|
if os.path.isfile(os.path.join(_scripts_dir, "hipInfo.exe")):
|
|
import shutil as _shutil
|
|
if not _shutil.which("hipinfo.exe"):
|
|
os.environ["PATH"] = _scripts_dir + os.pathsep + os.environ.get("PATH", "")
|
|
del _shutil
|
|
del _scripts_dir
|
|
|
|
# ── Windows AMD ROCm: set BNB_ROCM_VERSION before any bitsandbytes import ─
|
|
# bitsandbytes derives the rocm<ver>.dll name from torch.version.hip, but the
|
|
# wheel ships rocm72.dll, so the server crashes ("Configured ROCm binary not
|
|
# found") without this. Detect the shipped DLL (mirrors worker.py); gate on
|
|
# the rocm bnb DLL rather than torch.version.hip to avoid importing torch on
|
|
# every Windows host.
|
|
# Values seeded by the installer's sitecustomize.py are redetectable
|
|
# defaults; explicit caller values remain authoritative.
|
|
if (
|
|
"BNB_ROCM_VERSION" not in os.environ
|
|
or os.environ.get("UNSLOTH_BNB_ROCM_VERSION_SOURCE") == "sitecustomize"
|
|
):
|
|
import glob as _glob
|
|
import logging as _logging
|
|
|
|
_bnb_rocm_ver = None
|
|
_found_rocm_bnb = False
|
|
try:
|
|
import importlib.util as _ilu
|
|
_bnb_spec = _ilu.find_spec("bitsandbytes")
|
|
# submodule_search_locations (not spec.origin) handles editable installs
|
|
if _bnb_spec and _bnb_spec.submodule_search_locations:
|
|
import re as _re_bnb
|
|
|
|
_all_vers_main: list[str] = []
|
|
for _pkg_dir in _bnb_spec.submodule_search_locations:
|
|
for _dll in _glob.glob(os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll")):
|
|
_found_rocm_bnb = True
|
|
_km = _re_bnb.search(
|
|
r"libbitsandbytes_rocm(\d+)\.dll", os.path.basename(_dll)
|
|
)
|
|
if _km:
|
|
_all_vers_main.append(_km.group(1))
|
|
if _all_vers_main:
|
|
_bnb_rocm_ver = max(_all_vers_main, key = lambda v: int(v))
|
|
except Exception as _e:
|
|
_logging.getLogger(__name__).warning(
|
|
"Windows ROCm: BNB DLL detection failed (%s); leaving BNB_ROCM_VERSION as is",
|
|
_e,
|
|
)
|
|
# Only when a ROCm bnb DLL actually exists: HIP_PATH/ROCM_PATH alone
|
|
# (HIP SDK on a CUDA/CPU box) must not force a ROCm backend onto a
|
|
# non-ROCm bitsandbytes, which raises at import. DLL unparsable -> "72".
|
|
if _found_rocm_bnb:
|
|
_bnb_rocm_ver_final = _bnb_rocm_ver or os.environ.get("BNB_ROCM_VERSION") or "72"
|
|
os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver_final
|
|
os.environ["UNSLOTH_BNB_ROCM_VERSION_SOURCE"] = "detected"
|
|
_logging.getLogger(__name__).info(
|
|
"Windows ROCm: set BNB_ROCM_VERSION=%s (from installed BNB wheel)",
|
|
_bnb_rocm_ver_final,
|
|
)
|
|
|
|
# ── WSL AMD Strix Halo (gfx1151): enable ROCDXG before any torch import ──────
|
|
# In WSL the AMD GPU is reached via the ROCDXG bridge (librocdxg.so over
|
|
# /dev/dxg), which HSA loads only when HSA_ENABLE_DXG_DETECTION=1 is set BEFORE
|
|
# torch touches the GPU. A worker launched outside a login shell (e.g.
|
|
# `wsl.exe -d Ubuntu-24.04 python ...`) misses the installer's persisted env
|
|
# and silently falls back to CPU. Set it here, gated to no-op unless BOTH
|
|
# /dev/dxg AND librocdxg.so exist -- native Linux ROCm, NVIDIA, macOS and
|
|
# Windows are unaffected.
|
|
elif sys.platform.startswith("linux") and "HSA_ENABLE_DXG_DETECTION" not in os.environ:
|
|
try:
|
|
if os.path.exists("/dev/dxg") and any(
|
|
os.path.exists(os.path.join(_p, "librocdxg.so"))
|
|
for _p in ("/opt/rocm/lib", "/opt/rocm/lib64")
|
|
):
|
|
os.environ["HSA_ENABLE_DXG_DETECTION"] = "1"
|
|
import logging as _logging
|
|
_logging.getLogger(__name__).info(
|
|
"WSL ROCm: set HSA_ENABLE_DXG_DETECTION=1 (librocdxg bridge present)"
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
# Put backend dir 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)
|
|
|
|
# `uvicorn main:app` bypasses run.py; seed thread caps here too.
|
|
from utils.cpu_threads import configure_cpu_threads
|
|
|
|
try:
|
|
configure_cpu_threads()
|
|
except ValueError as exc:
|
|
_raw = os.environ.get("UNSLOTH_CPU_THREADS")
|
|
raise SystemExit(f"Error: Invalid UNSLOTH_CPU_THREADS value {_raw!r}: {exc}") from None
|
|
|
|
# Anaconda/conda-forge Python: seed platform._sys_version_cache before any
|
|
# library import triggers 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, whose
|
|
# 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")
|
|
|
|
# The studio bundles unsloth_zoo; declare unsloth present (as `import unsloth`
|
|
# does) so its lazy submodule imports (export, hardware, mlx) and the
|
|
# DiffusionGemma runner never trip the install guard on a clean install.
|
|
os.environ.setdefault("UNSLOTH_IS_PRESENT", "1")
|
|
|
|
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
|
|
from typing import Optional
|
|
from urllib.parse import urlparse
|
|
|
|
|
|
_STUDIO_INSTALL_ID_RE = _re.compile(r"^[0-9a-f]{64}$")
|
|
|
|
|
|
def _read_studio_install_id() -> str:
|
|
"""Per-install opaque id at $STUDIO_HOME/share/studio_install_id.
|
|
|
|
Returns "" when absent or not a 64-char lowercase-hex token; then
|
|
/api/health emits "" and the launcher accepts any healthy backend.
|
|
Carries no install-path info (matters when Studio runs -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 (cached at import).
|
|
|
|
Empty when no installer token is present; the launcher treats "" as
|
|
"accept any healthy backend"."""
|
|
return _STUDIO_ROOT_ID_CACHE
|
|
|
|
|
|
# Fix broken Windows registry MIME types: some installs map .js to text/plain,
|
|
# which mimetypes (hence StaticFiles) inherits and browsers reject for ES
|
|
# modules. add_type() before StaticFiles forces correct types.
|
|
if sys.platform == "win32":
|
|
mimetypes.add_type("application/javascript", ".js")
|
|
mimetypes.add_type("text/css", ".css")
|
|
|
|
# Suppress dependency warnings in production
|
|
if os.getenv("ENVIRONMENT_TYPE", "production") == "production":
|
|
warnings.filterwarnings("ignore")
|
|
# Or 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
|
|
|
|
from routes import (
|
|
auth_router,
|
|
chat_history_router,
|
|
data_recipe_router,
|
|
datasets_router,
|
|
export_router,
|
|
inference_router,
|
|
inference_studio_router,
|
|
mcp_servers_router,
|
|
models_router,
|
|
providers_router,
|
|
rag_router,
|
|
training_history_router,
|
|
training_router,
|
|
)
|
|
from routes.llama import router as llama_router
|
|
from hub.routes import (
|
|
inventory_router as hub_inventory_router,
|
|
datasets_router as hub_datasets_router,
|
|
)
|
|
from hub.schemas.downloads import TransportCapabilities
|
|
from hub.utils.download_registry import (
|
|
get_download_transport_capabilities,
|
|
reap_orphan_workers as reap_hub_orphan_workers,
|
|
terminate_active_downloads as terminate_hub_downloads,
|
|
)
|
|
from routes.settings import router as settings_router
|
|
from routes.prompts import router as prompts_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
|
|
from utils.api_errors import install_api_error_handlers
|
|
|
|
|
|
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()
|
|
|
|
# The Tauri desktop app runs the backend on the owner's own machine, so local
|
|
# stdio MCP servers are safe there. setdefault lets an explicit "0" opt out.
|
|
if _DESKTOP_OWNER:
|
|
os.environ.setdefault("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1")
|
|
|
|
|
|
def _desktop_owner() -> dict[str, str] | None:
|
|
return _DESKTOP_OWNER
|
|
|
|
|
|
def _start_helper_precache_if_enabled() -> None:
|
|
"""Start optional Helper LLM GGUF pre-cache only after explicit opt-in."""
|
|
try:
|
|
from utils.helper_precache_settings import should_preload_helper_on_startup
|
|
if not should_preload_helper_on_startup():
|
|
return
|
|
except Exception:
|
|
return
|
|
|
|
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, name = "helper-gguf-precache").start()
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
"""Startup: detect hardware, seed default admin if needed. Shutdown: clean up compiled cache."""
|
|
clear_unsloth_compiled_cache()
|
|
|
|
# Remove stale .venv_overlay from old versions; switching now uses .venv_t5/.
|
|
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 the DEVICE global used everywhere.
|
|
detect_hardware()
|
|
|
|
# Reap download workers orphaned by a previous crash before new downloads start.
|
|
reap_hub_orphan_workers()
|
|
|
|
# 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)
|
|
|
|
_start_helper_precache_if_enabled()
|
|
|
|
# Warm the RAG embedder so the first upload skips the cold load. Non-fatal.
|
|
def _warm_rag_embedder():
|
|
try:
|
|
from storage import rag_db
|
|
|
|
if not rag_db.RAG_AVAILABLE:
|
|
return
|
|
from core.rag import embeddings
|
|
|
|
embeddings.warm()
|
|
except Exception:
|
|
pass
|
|
|
|
threading.Thread(target = _warm_rag_embedder, 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
|
|
await asyncio.to_thread(terminate_hub_downloads)
|
|
_hw_module.DEVICE = None
|
|
clear_unsloth_compiled_cache()
|
|
|
|
|
|
app = FastAPI(
|
|
title = "Unsloth UI Backend",
|
|
version = UNSLOTH_VERSION,
|
|
description = "Backend API for Unsloth UI - Training and Model Management",
|
|
lifespan = lifespan,
|
|
)
|
|
|
|
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)
|
|
|
|
|
|
# img/media-src allow any https origin so HF model-card assets render (mirrors
|
|
# tauri.conf.json); scripts/frames/connect-src stay same-origin + HF.
|
|
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"
|
|
_ARTIFACT_PREVIEW_FRAME_PATH = "/api/inference/artifact-preview-frame"
|
|
|
|
|
|
# /content is Colab's working directory — more reliable than env vars, which
|
|
# aren't always set depending on Colab runtime version.
|
|
import importlib.util as _importlib_util
|
|
|
|
_IS_COLAB = os.path.isdir("/content") and (
|
|
bool(os.environ.get("COLAB_BACKEND_URL"))
|
|
or bool(os.environ.get("COLAB_JUPYTER_IP"))
|
|
or _importlib_util.find_spec("google.colab") is not None
|
|
)
|
|
|
|
|
|
def _build_csp(script_nonce: "str | None" = None) -> str:
|
|
script_src = "script-src 'self'"
|
|
if script_nonce:
|
|
script_src += f" 'nonce-{script_nonce}'"
|
|
# Colab parent frames span multi-level *.prod.colab.dev subdomains (CSP
|
|
# wildcards match one level only) and null-origin iframes; use '*' since
|
|
# Colab is already a sandboxed single-user environment.
|
|
frame_ancestors = "*" if _IS_COLAB else "'none'"
|
|
|
|
# In Colab, the kernel/output scaffolding injects scripts and fetch/WS from
|
|
# *.prod.colab.dev and *.googleusercontent.com, so widen script-src and
|
|
# connect-src for those. Scripts still use a nonce, not 'unsafe-inline'.
|
|
if _IS_COLAB:
|
|
script_src += " https://*.prod.colab.dev https://*.googleusercontent.com"
|
|
connect_src = (
|
|
"'self' blob: data: "
|
|
"https://huggingface.co https://datasets-server.huggingface.co "
|
|
"https://*.prod.colab.dev wss://*.prod.colab.dev "
|
|
"https://*.googleusercontent.com wss://*.googleusercontent.com"
|
|
)
|
|
else:
|
|
connect_src = "'self' https://huggingface.co https://datasets-server.huggingface.co"
|
|
|
|
return (
|
|
"default-src 'self'; "
|
|
"img-src 'self' data: blob: https:; "
|
|
"media-src 'self' data: blob: https:; "
|
|
f"connect-src {connect_src}; "
|
|
"style-src 'self' 'unsafe-inline'; "
|
|
f"{script_src}; "
|
|
"font-src 'self' data:; "
|
|
"frame-src 'self'; "
|
|
f"frame-ancestors {frame_ancestors}; "
|
|
"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))
|
|
# Omit X-Frame-Options in Colab — CSP frame-ancestors handles it, and
|
|
# DENY would block serve_kernel_port_as_iframe regardless of CSP.
|
|
if not _IS_COLAB and request.url.path != _ARTIFACT_PREVIEW_FRAME_PATH:
|
|
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=(self), geolocation=()",
|
|
)
|
|
response.headers["server"] = "unsloth-studio"
|
|
return response
|
|
|
|
|
|
app.add_middleware(SecurityHeadersMiddleware)
|
|
|
|
|
|
# Cap request bodies on protected POSTs. Upload routes get explicit multipart
|
|
# headroom; non-upload routes keep the default body cap.
|
|
import json as _json_for_413 # noqa: E402
|
|
from utils.upload_limits import ( # noqa: E402
|
|
UNSTRUCTURED_RECIPE_UPLOAD_MAX_BYTES,
|
|
default_request_body_limit_bytes,
|
|
upload_request_limit_bytes,
|
|
)
|
|
|
|
_BODY_PROTECTED_PREFIXES = (
|
|
"/v1/chat/completions",
|
|
"/v1/completions",
|
|
"/api/inference",
|
|
"/api/data-recipe",
|
|
"/api/datasets",
|
|
"/api/hub",
|
|
"/api/chat",
|
|
"/api/settings",
|
|
"/api/train",
|
|
"/api/export",
|
|
)
|
|
_DATASET_UPLOAD_PASSTHROUGH_PREFIX = "/api/datasets/upload"
|
|
_DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX = (
|
|
"/api/data-recipe/seed/upload-unstructured-file"
|
|
)
|
|
_BODY_UPLOAD_PASSTHROUGH_PREFIXES = (
|
|
_DATASET_UPLOAD_PASSTHROUGH_PREFIX,
|
|
_DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX,
|
|
)
|
|
|
|
|
|
def _get_upload_passthrough_request_max_bytes(path: str) -> int:
|
|
if path.startswith(_DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX):
|
|
return upload_request_limit_bytes(UNSTRUCTURED_RECIPE_UPLOAD_MAX_BYTES)
|
|
if path.startswith(_DATASET_UPLOAD_PASSTHROUGH_PREFIX):
|
|
return upload_request_limit_bytes()
|
|
return default_request_body_limit_bytes()
|
|
|
|
|
|
async def _send_411(send) -> None:
|
|
payload = _json_for_413.dumps(
|
|
{"detail": "Content-Length required for upload requests."},
|
|
).encode("utf-8")
|
|
await send(
|
|
{
|
|
"type": "http.response.start",
|
|
"status": 411,
|
|
"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})
|
|
|
|
|
|
async def _send_413(send, total_bytes: int, max_bytes: int) -> None:
|
|
payload = _json_for_413.dumps(
|
|
{"detail": (f"Request body too large ({total_bytes:,} bytes; max {max_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_getter,
|
|
protected_prefixes: tuple,
|
|
upload_passthrough_prefixes: tuple = (),
|
|
upload_passthrough_max_bytes_getter = None,
|
|
):
|
|
self.app = app
|
|
self.max_bytes_getter = max_bytes_getter
|
|
self.protected_prefixes = protected_prefixes
|
|
self.upload_passthrough_prefixes = upload_passthrough_prefixes
|
|
self.upload_passthrough_max_bytes_getter = upload_passthrough_max_bytes_getter
|
|
|
|
def _upload_passthrough_max_bytes(self, path: str) -> int:
|
|
if self.upload_passthrough_max_bytes_getter is None:
|
|
return int(self.max_bytes_getter())
|
|
try:
|
|
return int(self.upload_passthrough_max_bytes_getter(path))
|
|
except TypeError:
|
|
try:
|
|
return int(self.upload_passthrough_max_bytes_getter())
|
|
except Exception:
|
|
return int(self.max_bytes_getter())
|
|
except Exception:
|
|
return int(self.max_bytes_getter())
|
|
|
|
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
|
|
|
|
max_bytes = int(self.max_bytes_getter())
|
|
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 any(path.startswith(p) for p in self.upload_passthrough_prefixes):
|
|
upload_max_bytes = self._upload_passthrough_max_bytes(path)
|
|
if declared is None:
|
|
await _send_411(send)
|
|
return
|
|
if declared > upload_max_bytes:
|
|
await _send_413(send, declared, upload_max_bytes)
|
|
return
|
|
await self.app(scope, receive, send)
|
|
return
|
|
|
|
if declared is not None and declared > max_bytes:
|
|
await _send_413(send, declared, max_bytes)
|
|
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 > max_bytes:
|
|
await _send_413(send, total, max_bytes)
|
|
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_getter = default_request_body_limit_bytes,
|
|
protected_prefixes = _BODY_PROTECTED_PREFIXES,
|
|
upload_passthrough_prefixes = _BODY_UPLOAD_PASSTHROUGH_PREFIXES,
|
|
upload_passthrough_max_bytes_getter = _get_upload_passthrough_request_max_bytes,
|
|
)
|
|
|
|
|
|
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)
|
|
|
|
|
|
_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 NOT exposed on the /v1
|
|
# OpenAI-compat prefix below.
|
|
app.include_router(inference_studio_router, prefix = "/api/inference", tags = ["inference"])
|
|
|
|
# OpenAI-compatible: mount the inference router at /v1 for external tools.
|
|
app.include_router(inference_router, prefix = "/v1", tags = ["openai-compat"])
|
|
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"])
|
|
app.include_router(prompts_router, prefix = "/api/prompts", tags = ["prompts"])
|
|
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(llama_router, prefix = "/api/llama", tags = ["llama"])
|
|
app.include_router(export_router, prefix = "/api/export", tags = ["export"])
|
|
app.include_router(rag_router, prefix = "/api/rag", tags = ["rag"])
|
|
app.include_router(training_history_router, prefix = "/api/train", tags = ["training-history"])
|
|
app.include_router(hub_inventory_router, prefix = "/api/hub", tags = ["hub"])
|
|
app.include_router(hub_datasets_router, prefix = "/api/hub/datasets", tags = ["hub"])
|
|
|
|
# Re-wrap client-error responses on the /v1/* surface into OpenAI/Anthropic
|
|
# error envelopes; non-/v1 paths keep FastAPI's default {"detail": ...} shape.
|
|
install_api_error_handlers(app)
|
|
|
|
|
|
# ============ Health and System Endpoints ============
|
|
|
|
|
|
@app.get("/api/health")
|
|
async def health_check(request: Request):
|
|
"""Liveness plus launcher capability bits; host fingerprint gated on a bearer.
|
|
|
|
Unauthenticated callers get non-sensitive fields (service, studio_root_id,
|
|
chat_only, desktop_*, native_path_leases_supported) to re-adopt a sibling
|
|
backend and gate UI before a token exists. version / studio_version /
|
|
device_type require a bearer since 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.get(
|
|
"/api/studio/download-transport-capabilities",
|
|
response_model = TransportCapabilities,
|
|
)
|
|
def studio_download_transport_capabilities(_current_subject: str = Depends(get_current_subject)):
|
|
return asdict(get_download_transport_capabilities())
|
|
|
|
|
|
@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 the CLI or killing the process manually.
|
|
"""
|
|
|
|
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.
|
|
|
|
Auth-gated: the response (platform, Python/GPU, memory, ML packages) can
|
|
fingerprint a host, which matters in -H 0.0.0.0 / Colab / Tauri-relayed
|
|
setups where remote callers can reach /api/system.
|
|
"""
|
|
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(),
|
|
# _backend_label so /api/system reports "rocm" (not "cuda") on AMD,
|
|
# matching /api/hardware and /api/gpu-visibility.
|
|
"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.
|
|
"""
|
|
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's default ``crossorigin`` forces CORS mode on font loads, which
|
|
Firefox HTTPS-Only Mode breaks over plain HTTP; stripping it makes them
|
|
same-origin fetches that work on any protocol.
|
|
"""
|
|
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 forward the nonce
|
|
via ``_CSP_SCRIPT_NONCE_HEADER`` so CSP allows the inline script.
|
|
"""
|
|
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
|
|
|
|
|
|
_DEFAULT_PORTS = {"http": 80, "https": 443, "ws": 80, "wss": 443}
|
|
|
|
|
|
def _canonical_origin(scheme: str, netloc: str) -> Optional[tuple[str, str, int]]:
|
|
"""Canonicalise an Origin to ``(scheme, host, port)`` for equality.
|
|
Browsers strip default ports (RFC 6454 sec 6.1) and scheme/host are
|
|
case-insensitive (RFC 3986), so a bare string compare misclassifies
|
|
same-origin requests as cross-origin. Returns ``None`` on unparseable input
|
|
so callers fall to the safer cross-origin default.
|
|
"""
|
|
scheme = (scheme or "").strip().lower()
|
|
if not scheme or not netloc:
|
|
return None
|
|
# Strip userinfo (RFC 3986); Origin never carries credentials.
|
|
if "@" in netloc:
|
|
netloc = netloc.rsplit("@", 1)[1]
|
|
# IPv6 hosts use brackets (RFC 3986 sec 3.2.2): ``[::1]:8902``. Bare
|
|
# ``partition(":")`` mis-parses these, breaking ``unsloth studio -H ::1``.
|
|
if netloc.startswith("["):
|
|
close = netloc.find("]")
|
|
if close == -1:
|
|
return None
|
|
host = netloc[1:close]
|
|
rest = netloc[close + 1 :]
|
|
if rest.startswith(":"):
|
|
port_str = rest[1:]
|
|
elif rest == "":
|
|
port_str = ""
|
|
else:
|
|
return None
|
|
else:
|
|
host, _, port_str = netloc.partition(":")
|
|
host = host.strip().lower()
|
|
if not host:
|
|
return None
|
|
if port_str:
|
|
try:
|
|
port = int(port_str)
|
|
except ValueError:
|
|
return None
|
|
else:
|
|
port = _DEFAULT_PORTS.get(scheme, 0)
|
|
return (scheme, host, port)
|
|
|
|
|
|
def _is_same_origin_request(request: Request) -> bool:
|
|
"""True when Origin is missing or matches request's scheme://host:port.
|
|
|
|
Missing Origin counts as same-origin (top-level GETs omit it). Both sides
|
|
are canonicalised via :func:`_canonical_origin`; callers must emit
|
|
``Vary: Origin``.
|
|
"""
|
|
origin = request.headers.get("origin")
|
|
if origin is None:
|
|
# Missing header: top-level same-document GETs omit Origin.
|
|
return True
|
|
# Empty string is not a valid serialised origin (RFC 6454 sec 6.1).
|
|
if not origin:
|
|
return False
|
|
# "null" token (sandboxed iframes, file:// pages) is never same-origin.
|
|
if origin == "null":
|
|
return False
|
|
# ``urlparse`` raises ``ValueError`` on malformed IPv6 brackets; swallow
|
|
# so a garbage Origin doesn't 500 the SPA handler.
|
|
try:
|
|
parsed = urlparse(origin)
|
|
except ValueError:
|
|
return False
|
|
origin_canon = _canonical_origin(parsed.scheme, parsed.netloc)
|
|
if origin_canon is None:
|
|
return False
|
|
try:
|
|
self_canon = _canonical_origin(request.url.scheme, request.url.netloc)
|
|
except ValueError:
|
|
return False
|
|
if self_canon is None:
|
|
return False
|
|
return origin_canon == self_canon
|
|
|
|
|
|
def setup_frontend(app: FastAPI, build_path: Path):
|
|
"""Mount frontend static files (optional)"""
|
|
if not build_path.exists():
|
|
return False
|
|
|
|
assets_dir = build_path / "assets"
|
|
if assets_dir.exists():
|
|
app.mount("/assets", StaticFiles(directory = assets_dir), name = "assets")
|
|
|
|
def _build_index_response(request: Request) -> Response:
|
|
content = (build_path / "index.html").read_bytes()
|
|
content = _strip_crossorigin(content)
|
|
# Bootstrap pw is same-origin only; Vary: Origin keeps caches honest.
|
|
if _is_same_origin_request(request):
|
|
content, nonce = _inject_bootstrap(content, app)
|
|
else:
|
|
nonce = None
|
|
headers = {
|
|
"Cache-Control": "no-cache, no-store, must-revalidate",
|
|
"Vary": "Origin",
|
|
}
|
|
if nonce:
|
|
headers[_CSP_SCRIPT_NONCE_HEADER] = nonce
|
|
return Response(
|
|
content = content,
|
|
media_type = "text/html",
|
|
headers = headers,
|
|
)
|
|
|
|
@app.get("/")
|
|
async def serve_root(request: Request):
|
|
return _build_index_response(request)
|
|
|
|
@app.get("/{full_path:path}")
|
|
async def serve_frontend(request: Request, full_path: str):
|
|
# Unknown API paths: raise a real 404 so the api_errors handlers can
|
|
# render the correct envelope for /v1/* (and {"detail":...} for /api/*).
|
|
# This handler only sees paths NOT matched by a real route. The full
|
|
# request path is "/" + full_path.
|
|
if full_path in {"api", "v1"} or full_path.startswith(("api/", "v1/")):
|
|
raise HTTPException(status_code = 404, detail = "API endpoint not found")
|
|
|
|
file_path = (build_path / full_path).resolve()
|
|
|
|
# Block path traversal — resolved path must stay 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(request)
|
|
|
|
return True
|