Studio: require signed capability tokens for /p preview links (#6666)
* Studio: require signed capability tokens for /p preview links The public /p preview routes added in #6486 run model load and chat generation as the admin user with no authentication. The only gate is the preview ref, a deterministic outputs-root path (run or run/checkpoint) that is guessable rather than secret. On a network-reachable Studio (--secure tunnel or -H 0.0.0.0), an unauthenticated caller who guesses a ref can consume GPU and probe a private fine-tuned checkpoint. Make the share link an unguessable, revocable capability: - Sign the canonical ref with a dedicated server-side secret (HMAC-SHA256, stored in app_secrets, independent of the JWT/login secret). - Require a valid token on every /p chat, models, and page request before resolving a checkpoint or loading a model; missing or invalid tokens get a generic 404 so the surface never confirms a ref exists. - Accept the token via ?k= (browser link and preview page) or Authorization: Bearer (OpenAI-compatible clients). - Rotate the secret to revoke every outstanding link (POST /api/settings/preview-links/rotate). - Clamp preview generation (max_tokens/max_completion_tokens <= 1024, n = 1) and set Referrer-Policy: no-referrer on the page so the token is not leaked via Referer. Training history hands the authenticated owner the signed token, and the copy-link button builds /p/{ref}?k={sig}. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: honor a lower caller token limit in the preview clamp Codex review: when only the legacy max_tokens was sent, the clamp left max_completion_tokens at the 1024 default, and _effective_max_tokens prefers max_completion_tokens, so a request like max_tokens=16 could still generate up to 1024 tokens. Derive one effective limit (max_completion_tokens wins, else the legacy max_tokens) and pin both fields to it so a caller's lower limit is kept. * Studio: add preview kill switch, rate limit, and revoke-links UI Follow-ups to the /p preview capability work: - Public-sharing kill switch: a persisted setting (default on) gates the public /p surface. When off, every preview request 404s even with a valid token, and the owner UI stops offering share links. GET/PUT /api/settings/preview-sharing; enforced in _verify_or_404. - Per-IP rate limit on the preview chat route: a coarse in-process sliding-window limiter (20 req/min/IP) returns 429 + Retry-After before the GPU lock is taken. Client IP honors X-Forwarded-For only when UNSLOTH_STUDIO_TRUST_FORWARDED is set, matching the login limiter's trust model. - Settings UI: a "Preview sharing" section with the public-sharing toggle and a "Revoke all preview links" button (confirm dialog) that rotates the secret. Tests cover the kill switch (404 when off), the 429 path, the sliding window, client-IP trust behavior, and the setting default. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: fix preview-fields sharing arg and refresh sigs after revoke Codex review: - P1: get_training_run_detail and update_training_run called _preview_fields with only output_dir after it gained a required sharing_on parameter, raising a 500 TypeError once get_run succeeded. Pass get_preview_sharing_enabled() at both sites; add a detail-endpoint regression test. - P2: after rotating the preview secret from settings, the history grid still held stale preview_sig values, so a freshly copied link would 404. Emit emitTrainingRunsChanged() after a successful revoke so the grid refetches freshly signed refs. * Studio: harden preview sharing controls (Codex review) - Fail closed: a read failure on the preview-sharing kill switch now returns False instead of defaulting to enabled, so an unavailable settings DB can't reopen the public surface. A missing key still defaults to enabled. - Per-IP rate limit behind the managed Cloudflare tunnel: client_ip now honors CF-Connecting-IP when the socket peer is loopback, so tunneled visitors are keyed by their real IP instead of collapsing onto the local cloudflared peer. - GET /p no longer mints key/share_url when sharing is disabled; it returns sharing_enabled=false so clients don't distribute links that 404. - Settings UI: toggling public sharing emits the training-runs-changed event so the history grid shows/hides Copy preview link without a manual refresh. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: harden preview rate limiter and IP keying (Opus review) From a two-agent review of the PR: - Rate limiter no longer evicts an active bucket when the table is full: a flood of distinct keys could otherwise cycle out a throttled bucket and reset its counter. Evict only aged-out buckets; if the table is full of live clients, fail closed (deny the new key) instead. - client_ip keys on the rightmost (proxy-appended) X-Forwarded-For hop when the trust env is set; the leftmost is client-spoofable. Documented the append/overwrite-proxy assumption. - _verify_or_404 checks the capability token before the kill-switch DB read, so unauthenticated /p spam can't be used as an unbounded settings-DB sink and the response is identical regardless of the sharing on/off state. Tests: nested run/checkpoint happy path + wrong-ref rejection, the eviction fail-closed behavior, and route-level coverage for the rotate / preview-sharing settings endpoints. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
ed5e2a1590
commit
80d3434d61
20 changed files with 1256 additions and 35 deletions
|
|
@ -283,6 +283,11 @@
|
|||
</div>
|
||||
<script>
|
||||
const base = location.pathname.replace(/\/+$/, "");
|
||||
// The capability token rides in ?k=; location.pathname drops it, so carry it
|
||||
// onto the chat request explicitly. Not stored or logged.
|
||||
const k = new URLSearchParams(location.search).get("k");
|
||||
const chatUrl =
|
||||
base + "/v1/chat/completions" + (k ? "?k=" + encodeURIComponent(k) : "");
|
||||
const log = document.getElementById("log"),
|
||||
thread = document.getElementById("thread"),
|
||||
welcome = document.getElementById("welcome");
|
||||
|
|
@ -328,7 +333,7 @@
|
|||
out.innerHTML = '<span class="dots"><i></i><i></i><i></i></span>';
|
||||
let acc = "";
|
||||
try {
|
||||
const r = await fetch(base + "/v1/chat/completions", {
|
||||
const r = await fetch(chatUrl, {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({
|
||||
|
|
|
|||
|
|
@ -270,6 +270,63 @@ def compute_identity_proof(nonce: bytes, host: str, port: int) -> str:
|
|||
return hmac.new(get_or_create_identity_secret(), msg, hashlib.sha256).hexdigest()
|
||||
|
||||
|
||||
# Capability secret for public ``/p`` preview share links. HMAC(secret, ref)
|
||||
# turns the deterministic preview ref into an unguessable bearer capability, so a
|
||||
# guessed run/checkpoint name can't reach inference. Dedicated (not the per-user
|
||||
# JWT secret) so rotating it revokes every shared link without touching logins.
|
||||
_PREVIEW_LINK_SECRET_DB_KEY = "preview_link_secret"
|
||||
_preview_link_secret_cache: Optional[bytes] = None
|
||||
|
||||
|
||||
def get_or_create_preview_link_secret() -> bytes:
|
||||
"""Return the preview-link signing secret (hex 32-byte row in app_secrets), creating it once."""
|
||||
global _preview_link_secret_cache
|
||||
if _preview_link_secret_cache is not None:
|
||||
return _preview_link_secret_cache
|
||||
|
||||
conn = get_connection()
|
||||
try:
|
||||
row = conn.execute(
|
||||
"SELECT value FROM app_secrets WHERE key = ?",
|
||||
(_PREVIEW_LINK_SECRET_DB_KEY,),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
conn.execute(
|
||||
"INSERT OR IGNORE INTO app_secrets (key, value) VALUES (?, ?)",
|
||||
(_PREVIEW_LINK_SECRET_DB_KEY, secrets.token_hex(32)),
|
||||
)
|
||||
conn.commit()
|
||||
row = conn.execute(
|
||||
"SELECT value FROM app_secrets WHERE key = ?",
|
||||
(_PREVIEW_LINK_SECRET_DB_KEY,),
|
||||
).fetchone()
|
||||
secret = bytes.fromhex(row["value"])
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
_preview_link_secret_cache = secret
|
||||
return secret
|
||||
|
||||
|
||||
def rotate_preview_link_secret() -> bytes:
|
||||
"""Rotate the preview-link secret, immediately revoking every outstanding ``/p`` share link."""
|
||||
global _preview_link_secret_cache
|
||||
new_secret_hex = secrets.token_hex(32)
|
||||
conn = get_connection()
|
||||
try:
|
||||
conn.execute(
|
||||
"INSERT OR REPLACE INTO app_secrets (key, value) VALUES (?, ?)",
|
||||
(_PREVIEW_LINK_SECRET_DB_KEY, new_secret_hex),
|
||||
)
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
secret = bytes.fromhex(new_secret_hex)
|
||||
_preview_link_secret_cache = secret
|
||||
return secret
|
||||
|
||||
|
||||
_API_KEY_PBKDF2_ITERATIONS = 100_000
|
||||
DESKTOP_SECRET_PREFIX = "desktop-"
|
||||
_DESKTOP_SECRET_HASH_KEY = "desktop_secret_hash"
|
||||
|
|
|
|||
|
|
@ -603,6 +603,9 @@ class TrainingRunSummary(BaseModel):
|
|||
resumed_later: bool = False
|
||||
has_preview_model: bool = False
|
||||
preview_ref: Optional[str] = None
|
||||
# HMAC capability token for the `/p/{preview_ref}` share link; None when not
|
||||
# previewable. The frontend appends it as `?k=` so a guessed ref can't be used.
|
||||
preview_sig: Optional[str] = None
|
||||
|
||||
|
||||
class TrainingRunUpdateRequest(BaseModel):
|
||||
|
|
|
|||
|
|
@ -19,17 +19,68 @@ from auth.storage import DEFAULT_ADMIN_USERNAME
|
|||
from models.inference import ChatCompletionRequest, LoadRequest
|
||||
from routes.inference import load_model, openai_chat_completions
|
||||
from state.tool_policy import tools_force_disabled
|
||||
from utils.client_ip import client_ip
|
||||
from utils.models.checkpoints import list_preview_targets, resolve_preview_checkpoint
|
||||
from utils.preview_rate_limit import check_rate_limit
|
||||
from utils.preview_sharing_settings import get_preview_sharing_enabled
|
||||
from utils.preview_token import sign_preview_ref, verify_preview_ref
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
# Public (no key); resolve_preview_checkpoint pins `run` under outputs_root.
|
||||
# One model loads at a time, so serialize load+generate across previews.
|
||||
# A shared preview link is a public bearer capability; cap per-request generation
|
||||
# so a single call can't tie up the (serialized) preview GPU indefinitely.
|
||||
_PREVIEW_MAX_OUTPUT_TOKENS = 1024
|
||||
|
||||
# Capability-gated (signed ref required); resolve_preview_checkpoint pins `run`
|
||||
# under outputs_root. One model loads at a time, so serialize load+generate.
|
||||
_preview_lock = asyncio.Lock()
|
||||
|
||||
|
||||
def _extract_token(request: Request) -> str | None:
|
||||
"""Capability token from the ``?k=`` query (browser link + preview page) or an
|
||||
``Authorization: Bearer`` header (OpenAI-compatible clients using it as api_key)."""
|
||||
token = request.query_params.get("k")
|
||||
if token:
|
||||
return token
|
||||
header = request.headers.get("authorization", "")
|
||||
if header[:7].lower() == "bearer ":
|
||||
return header[7:].strip() or None
|
||||
return None
|
||||
|
||||
|
||||
def _verify_or_404(run: str, checkpoint: str | None, request: Request) -> None:
|
||||
"""Require a valid preview capability BEFORE any checkpoint resolve / model load.
|
||||
|
||||
Missing or invalid tokens get a generic 404 -- identical to a non-existent ref --
|
||||
so the public surface never confirms whether a run/checkpoint exists. When an
|
||||
admin has switched public sharing off, every public request 404s regardless of
|
||||
token.
|
||||
|
||||
Verify the (cheap, no-I/O) capability first: an unauthenticated caller with a
|
||||
bad/missing token is rejected without the kill-switch DB read, so spamming
|
||||
``/p/...`` can't be used as an unbounded settings-DB sink, and the response is
|
||||
identical whether or not sharing is enabled (no on/off oracle).
|
||||
"""
|
||||
ref = run if not checkpoint else f"{run}/{checkpoint}"
|
||||
if not verify_preview_ref(ref, _extract_token(request)):
|
||||
raise HTTPException(status_code = 404, detail = "Not found")
|
||||
if not get_preview_sharing_enabled():
|
||||
raise HTTPException(status_code = 404, detail = "Not found")
|
||||
|
||||
|
||||
def _enforce_rate_limit(request: Request) -> None:
|
||||
"""Throttle the GPU-backed preview chat per client IP (429 on exceed)."""
|
||||
retry_after = check_rate_limit(client_ip(request))
|
||||
if retry_after:
|
||||
raise HTTPException(
|
||||
status_code = 429,
|
||||
detail = "Too many preview requests. Please slow down.",
|
||||
headers = {"Retry-After": str(retry_after)},
|
||||
)
|
||||
|
||||
|
||||
def _resolve_or_4xx(run: str, checkpoint: str | None):
|
||||
try:
|
||||
return resolve_preview_checkpoint(run, checkpoint)
|
||||
|
|
@ -49,6 +100,21 @@ def _sanitize_preview_payload(
|
|||
# Normalize use_adapter (never trust the caller): pin True for LoRA, None for
|
||||
# merged. _apply_adapter_state mutates the shared model without restoring, so an
|
||||
# unpinned `false` would persist to later visitors who omit the field.
|
||||
#
|
||||
# Cap generation cost on this public, GPU-backed surface. Derive one effective
|
||||
# limit (mirroring _effective_max_tokens: max_completion_tokens wins, else the
|
||||
# legacy max_tokens) and pin BOTH fields to it, so a caller's lower limit is
|
||||
# honored and neither field can exceed the ceiling.
|
||||
requested = (
|
||||
payload.max_completion_tokens
|
||||
if payload.max_completion_tokens is not None
|
||||
else payload.max_tokens
|
||||
)
|
||||
capped_max_tokens = (
|
||||
min(requested, _PREVIEW_MAX_OUTPUT_TOKENS)
|
||||
if requested is not None
|
||||
else _PREVIEW_MAX_OUTPUT_TOKENS
|
||||
)
|
||||
return payload.model_copy(
|
||||
update = {
|
||||
"tools": None,
|
||||
|
|
@ -67,6 +133,9 @@ def _sanitize_preview_payload(
|
|||
"encrypted_api_key": None,
|
||||
"provider_base_url": None,
|
||||
"use_adapter": True if is_lora else None,
|
||||
"max_tokens": capped_max_tokens,
|
||||
"max_completion_tokens": capped_max_tokens,
|
||||
"n": 1,
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -105,15 +174,30 @@ async def _serve_chat(
|
|||
@router.get("")
|
||||
async def list_previews(request: Request, current_subject: str = Depends(get_current_subject)):
|
||||
base = str(request.base_url)
|
||||
sharing_on = get_preview_sharing_enabled()
|
||||
previews = []
|
||||
for target in list_preview_targets():
|
||||
ref = quote(target["ref"], safe = "/")
|
||||
previews.append({**target, "url": f"{base}p/{ref}/v1"})
|
||||
return {"object": "list", "data": previews}
|
||||
# Mint the capability for the authenticated owner: ``key`` for OpenAI
|
||||
# clients (Bearer / api_key), ``share_url`` for the browser link. When
|
||||
# public sharing is off, every public /p request 404s, so don't hand out
|
||||
# dead credentials -- omit the capability and signal the disabled state.
|
||||
token = sign_preview_ref(target["ref"]) if sharing_on else None
|
||||
previews.append(
|
||||
{
|
||||
**target,
|
||||
"url": f"{base}p/{ref}/v1",
|
||||
"key": token,
|
||||
"share_url": f"{base}p/{ref}?k={token}" if token else None,
|
||||
}
|
||||
)
|
||||
return {"object": "list", "data": previews, "sharing_enabled": sharing_on}
|
||||
|
||||
|
||||
@router.post("/{run}/v1/chat/completions")
|
||||
async def preview_chat_latest(run: str, payload: ChatCompletionRequest, request: Request):
|
||||
_verify_or_404(run, None, request)
|
||||
_enforce_rate_limit(request)
|
||||
return await _serve_chat(run, None, payload, request)
|
||||
|
||||
|
||||
|
|
@ -121,6 +205,8 @@ async def preview_chat_latest(run: str, payload: ChatCompletionRequest, request:
|
|||
async def preview_chat_checkpoint(
|
||||
run: str, checkpoint: str, payload: ChatCompletionRequest, request: Request
|
||||
):
|
||||
_verify_or_404(run, checkpoint, request)
|
||||
_enforce_rate_limit(request)
|
||||
return await _serve_chat(run, checkpoint, payload, request)
|
||||
|
||||
|
||||
|
|
@ -140,13 +226,17 @@ def _models_response(run: str, checkpoint: str | None):
|
|||
}
|
||||
|
||||
|
||||
# The models/page GET routes only stat the checkpoint dir (no GPU), so they are
|
||||
# token-gated but not rate-limited; only the GPU-backed chat path is throttled.
|
||||
@router.get("/{run}/v1/models")
|
||||
async def preview_models_latest(run: str):
|
||||
async def preview_models_latest(run: str, request: Request):
|
||||
_verify_or_404(run, None, request)
|
||||
return _models_response(run, None)
|
||||
|
||||
|
||||
@router.get("/{run}/{checkpoint}/v1/models")
|
||||
async def preview_models_checkpoint(run: str, checkpoint: str):
|
||||
async def preview_models_checkpoint(run: str, checkpoint: str, request: Request):
|
||||
_verify_or_404(run, checkpoint, request)
|
||||
return _models_response(run, checkpoint)
|
||||
|
||||
|
||||
|
|
@ -183,14 +273,24 @@ def _preview_page(run: str, checkpoint: str | None) -> HTMLResponse:
|
|||
_resolve_or_4xx(run, checkpoint)
|
||||
title = run if not checkpoint else f"{run}/{checkpoint}"
|
||||
page = _PREVIEW_PAGE_HTML.replace("__TITLE__", html.escape(title))
|
||||
return HTMLResponse(page, headers = {"Content-Security-Policy": _PREVIEW_PAGE_CSP})
|
||||
# no-referrer: the capability token rides in the query string, so keep it out
|
||||
# of the Referer header on any outbound navigation.
|
||||
return HTMLResponse(
|
||||
page,
|
||||
headers = {
|
||||
"Content-Security-Policy": _PREVIEW_PAGE_CSP,
|
||||
"Referrer-Policy": "no-referrer",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@router.get("/{run}", response_class = HTMLResponse)
|
||||
async def preview_page_latest(run: str):
|
||||
async def preview_page_latest(run: str, request: Request):
|
||||
_verify_or_404(run, None, request)
|
||||
return _preview_page(run, None)
|
||||
|
||||
|
||||
@router.get("/{run}/{checkpoint}", response_class = HTMLResponse)
|
||||
async def preview_page_checkpoint(run: str, checkpoint: str):
|
||||
async def preview_page_checkpoint(run: str, checkpoint: str, request: Request):
|
||||
_verify_or_404(run, checkpoint, request)
|
||||
return _preview_page(run, checkpoint)
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from fastapi import APIRouter, Depends
|
|||
from pydantic import BaseModel, ConfigDict, Field, field_validator
|
||||
|
||||
from auth.authentication import get_current_subject
|
||||
from auth.storage import rotate_preview_link_secret
|
||||
from loggers import get_logger
|
||||
from utils.utils import safe_error_detail, log_and_http_error
|
||||
from utils.personalization_settings import (
|
||||
|
|
@ -31,6 +32,11 @@ from utils.helper_precache_settings import (
|
|||
helper_model_disabled_by_env,
|
||||
set_helper_precache_enabled,
|
||||
)
|
||||
from utils.preview_sharing_settings import (
|
||||
DEFAULT_PREVIEW_SHARING_ENABLED,
|
||||
get_preview_sharing_enabled,
|
||||
set_preview_sharing_enabled,
|
||||
)
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
|
@ -122,6 +128,55 @@ def update_helper_precache(
|
|||
return _helper_precache_response(enabled)
|
||||
|
||||
|
||||
class PreviewLinkRotateResponse(BaseModel):
|
||||
rotated: bool = True
|
||||
|
||||
|
||||
@router.post("/preview-links/rotate", response_model = PreviewLinkRotateResponse)
|
||||
def rotate_preview_links(
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
) -> PreviewLinkRotateResponse:
|
||||
"""Rotate the preview-link signing secret, revoking every previously shared `/p` link."""
|
||||
rotate_preview_link_secret()
|
||||
logger.info("settings.preview_links_rotated subject=%s", current_subject)
|
||||
return PreviewLinkRotateResponse(rotated = True)
|
||||
|
||||
|
||||
class PreviewSharingPayload(BaseModel):
|
||||
enabled: bool
|
||||
|
||||
|
||||
class PreviewSharingResponse(BaseModel):
|
||||
enabled: bool
|
||||
default_enabled: bool = DEFAULT_PREVIEW_SHARING_ENABLED
|
||||
|
||||
|
||||
@router.get("/preview-sharing", response_model = PreviewSharingResponse)
|
||||
def get_preview_sharing(
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
) -> PreviewSharingResponse:
|
||||
return PreviewSharingResponse(enabled = get_preview_sharing_enabled())
|
||||
|
||||
|
||||
@router.put("/preview-sharing", response_model = PreviewSharingResponse)
|
||||
def update_preview_sharing(
|
||||
payload: PreviewSharingPayload, current_subject: str = Depends(get_current_subject)
|
||||
) -> PreviewSharingResponse:
|
||||
"""Enable/disable the public `/p` preview surface. When off, links 404 even with a token."""
|
||||
try:
|
||||
enabled = set_preview_sharing_enabled(payload.enabled)
|
||||
except ValueError as exc:
|
||||
raise log_and_http_error(
|
||||
exc,
|
||||
400,
|
||||
safe_error_detail(exc, fallback = "Invalid preview sharing setting."),
|
||||
event = "settings.update_preview_sharing_failed",
|
||||
log = logger,
|
||||
) from exc
|
||||
logger.info("settings.preview_sharing_updated subject=%s enabled=%s", current_subject, enabled)
|
||||
return PreviewSharingResponse(enabled = enabled)
|
||||
|
||||
|
||||
def _is_bundled_avatar_url(value: str) -> bool:
|
||||
parsed = urlsplit(value)
|
||||
if parsed.scheme or parsed.netloc:
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ Training history API routes — browse, view, and delete past training runs.
|
|||
"""
|
||||
|
||||
import json
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
from loggers import get_logger
|
||||
|
|
@ -28,12 +29,30 @@ from storage.studio_db import (
|
|||
update_run_display_name,
|
||||
)
|
||||
from utils.models.checkpoints import has_preview_model, preview_ref
|
||||
from utils.preview_sharing_settings import get_preview_sharing_enabled
|
||||
from utils.preview_token import sign_preview_ref
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
||||
def _preview_fields(output_dir: Optional[str], sharing_on: bool) -> dict:
|
||||
"""Previewability + the signed `/p` share ref for a run's output dir.
|
||||
|
||||
The signature is what makes the share link a capability: these routes are
|
||||
authenticated, so only the run's owner ever receives it. When public sharing
|
||||
is switched off, omit the signature so the UI hides the copy-link affordance
|
||||
(and the link would 404 anyway). ``sharing_on`` is resolved once per request.
|
||||
"""
|
||||
ref = preview_ref(output_dir)
|
||||
return {
|
||||
"has_preview_model": has_preview_model(output_dir),
|
||||
"preview_ref": ref,
|
||||
"preview_sig": sign_preview_ref(ref) if (ref and sharing_on) else None,
|
||||
}
|
||||
|
||||
|
||||
@router.get("/runs", response_model = TrainingRunListResponse)
|
||||
async def list_training_runs(
|
||||
limit: int = Query(50, ge = 1, le = 200),
|
||||
|
|
@ -42,14 +61,14 @@ async def list_training_runs(
|
|||
):
|
||||
"""List training runs, newest first."""
|
||||
result = list_runs(limit = limit, offset = offset)
|
||||
sharing_on = get_preview_sharing_enabled()
|
||||
return TrainingRunListResponse(
|
||||
runs = [
|
||||
TrainingRunSummary(
|
||||
**{
|
||||
**r,
|
||||
"can_resume": can_resume_run(r),
|
||||
"has_preview_model": has_preview_model(r.get("output_dir")),
|
||||
"preview_ref": preview_ref(r.get("output_dir")),
|
||||
**_preview_fields(r.get("output_dir"), sharing_on),
|
||||
}
|
||||
)
|
||||
for r in result["runs"]
|
||||
|
|
@ -78,8 +97,7 @@ async def get_training_run_detail(run_id: str, current_subject: str = Depends(ge
|
|||
**{
|
||||
**{k: v for k, v in run.items() if k != "config_json"},
|
||||
"can_resume": can_resume_run(run),
|
||||
"has_preview_model": has_preview_model(run.get("output_dir")),
|
||||
"preview_ref": preview_ref(run.get("output_dir")),
|
||||
**_preview_fields(run.get("output_dir"), get_preview_sharing_enabled()),
|
||||
}
|
||||
),
|
||||
config = config,
|
||||
|
|
@ -111,8 +129,7 @@ async def update_training_run(
|
|||
**{
|
||||
**{k: v for k, v in refreshed.items() if k != "config_json"},
|
||||
"can_resume": can_resume_run(refreshed),
|
||||
"has_preview_model": has_preview_model(refreshed.get("output_dir")),
|
||||
"preview_ref": preview_ref(refreshed.get("output_dir")),
|
||||
**_preview_fields(refreshed.get("output_dir"), get_preview_sharing_enabled()),
|
||||
}
|
||||
)
|
||||
|
||||
|
|
|
|||
143
studio/backend/tests/test_preview_followups.py
Normal file
143
studio/backend/tests/test_preview_followups.py
Normal file
|
|
@ -0,0 +1,143 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Unit coverage for the preview follow-ups: rate limiter, client IP, kill switch."""
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import types as _types
|
||||
|
||||
|
||||
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
||||
if _BACKEND_DIR not in sys.path:
|
||||
sys.path.insert(0, _BACKEND_DIR)
|
||||
|
||||
_loggers_stub = _types.ModuleType("loggers")
|
||||
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
|
||||
sys.modules.setdefault("loggers", _loggers_stub)
|
||||
|
||||
import utils.preview_rate_limit as rl
|
||||
from utils.client_ip import client_ip
|
||||
from utils.preview_sharing_settings import (
|
||||
DEFAULT_PREVIEW_SHARING_ENABLED,
|
||||
_coerce_bool,
|
||||
get_preview_sharing_enabled,
|
||||
)
|
||||
|
||||
|
||||
# ── Rate limiter ─────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_rate_limit_per_key(monkeypatch):
|
||||
monkeypatch.setattr(rl, "_MAX_REQUESTS", 3)
|
||||
rl.reset()
|
||||
assert rl.check_rate_limit("ip1") == 0
|
||||
assert rl.check_rate_limit("ip1") == 0
|
||||
assert rl.check_rate_limit("ip1") == 0
|
||||
# 4th request over the ceiling -> positive retry-after seconds.
|
||||
assert rl.check_rate_limit("ip1") > 0
|
||||
# A different client is unaffected.
|
||||
assert rl.check_rate_limit("ip2") == 0
|
||||
|
||||
|
||||
def test_rate_limit_window_rolls_off(monkeypatch):
|
||||
monkeypatch.setattr(rl, "_MAX_REQUESTS", 1)
|
||||
monkeypatch.setattr(rl, "_WINDOW_SECONDS", 10.0)
|
||||
rl.reset()
|
||||
t = {"now": 1000.0}
|
||||
monkeypatch.setattr(rl.time, "monotonic", lambda: t["now"])
|
||||
assert rl.check_rate_limit("ip") == 0
|
||||
assert rl.check_rate_limit("ip") > 0 # immediately over
|
||||
t["now"] += 11.0 # window elapsed
|
||||
assert rl.check_rate_limit("ip") == 0
|
||||
|
||||
|
||||
def test_rate_limit_eviction_does_not_reset_active_bucket(monkeypatch):
|
||||
# A flood of distinct keys must not cycle the table and clear a live limit.
|
||||
monkeypatch.setattr(rl, "_MAX_REQUESTS", 1)
|
||||
monkeypatch.setattr(rl, "_MAX_BUCKETS", 2)
|
||||
rl.reset()
|
||||
assert rl.check_rate_limit("a") == 0
|
||||
assert rl.check_rate_limit("a") > 0 # 'a' throttled (active)
|
||||
assert rl.check_rate_limit("b") == 0
|
||||
assert rl.check_rate_limit("b") > 0 # 'b' throttled; table now full of actives
|
||||
# A new key can't evict an active bucket -> denied (fail closed)...
|
||||
assert rl.check_rate_limit("c") > 0
|
||||
# ...and the flood did not reset 'a'.
|
||||
assert rl.check_rate_limit("a") > 0
|
||||
|
||||
|
||||
# ── Client IP ────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class _Req:
|
||||
def __init__(
|
||||
self,
|
||||
host,
|
||||
headers = None,
|
||||
):
|
||||
self.client = _types.SimpleNamespace(host = host) if host else None
|
||||
self.headers = headers or {}
|
||||
|
||||
|
||||
def test_client_ip_uses_socket_peer_by_default(monkeypatch):
|
||||
monkeypatch.delenv("UNSLOTH_STUDIO_TRUST_FORWARDED", raising = False)
|
||||
# Forwarded header is ignored unless the operator opts in.
|
||||
req = _Req("203.0.113.9", {"x-forwarded-for": "198.51.100.7"})
|
||||
assert client_ip(req) == "203.0.113.9"
|
||||
assert client_ip(None) == "_unknown"
|
||||
|
||||
|
||||
def test_client_ip_uses_rightmost_forwarded_when_trusted(monkeypatch):
|
||||
monkeypatch.setenv("UNSLOTH_STUDIO_TRUST_FORWARDED", "1")
|
||||
# Leftmost is client-spoofable; the trusted proxy appends the real peer on the
|
||||
# right, so the rightmost hop is the one we key on.
|
||||
req = _Req("127.0.0.1", {"x-forwarded-for": "1.2.3.4, 198.51.100.7"})
|
||||
assert client_ip(req) == "198.51.100.7"
|
||||
|
||||
|
||||
def test_client_ip_uses_cf_connecting_ip_on_loopback(monkeypatch):
|
||||
# Managed Cloudflare tunnel terminates at loopback; key by the real visitor.
|
||||
monkeypatch.delenv("UNSLOTH_STUDIO_TRUST_FORWARDED", raising = False)
|
||||
req = _Req("127.0.0.1", {"cf-connecting-ip": "198.51.100.7"})
|
||||
assert client_ip(req) == "198.51.100.7"
|
||||
|
||||
|
||||
def test_client_ip_ignores_cf_header_from_non_loopback(monkeypatch):
|
||||
# A direct (non-loopback) caller can't spoof CF-Connecting-IP to skew the limit.
|
||||
monkeypatch.delenv("UNSLOTH_STUDIO_TRUST_FORWARDED", raising = False)
|
||||
req = _Req("203.0.113.9", {"cf-connecting-ip": "198.51.100.7"})
|
||||
assert client_ip(req) == "203.0.113.9"
|
||||
|
||||
|
||||
def test_client_ip_loopback_without_cf_returns_peer(monkeypatch):
|
||||
monkeypatch.delenv("UNSLOTH_STUDIO_TRUST_FORWARDED", raising = False)
|
||||
assert client_ip(_Req("127.0.0.1")) == "127.0.0.1"
|
||||
|
||||
|
||||
# ── Kill-switch setting ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_sharing_defaults_enabled_and_coerces():
|
||||
assert DEFAULT_PREVIEW_SHARING_ENABLED is True
|
||||
assert _coerce_bool("off") is False
|
||||
assert _coerce_bool("on") is True
|
||||
assert _coerce_bool(True) is True
|
||||
assert _coerce_bool("nonsense") is None
|
||||
|
||||
|
||||
def test_sharing_missing_key_defaults_enabled(monkeypatch):
|
||||
import storage.studio_db as sdb
|
||||
monkeypatch.setattr(sdb, "get_app_setting", lambda key, fallback = None: None)
|
||||
assert get_preview_sharing_enabled() is True
|
||||
|
||||
|
||||
def test_sharing_read_error_fails_closed(monkeypatch):
|
||||
# A transient settings-DB failure must not reopen the public surface.
|
||||
import storage.studio_db as sdb
|
||||
|
||||
def _boom(*args, **kwargs):
|
||||
raise RuntimeError("settings db unavailable")
|
||||
|
||||
monkeypatch.setattr(sdb, "get_app_setting", _boom)
|
||||
assert get_preview_sharing_enabled() is False
|
||||
|
|
@ -5,10 +5,12 @@
|
|||
|
||||
Exercises the route layer with a real ``preview_router`` while stubbing the
|
||||
expensive model calls (``load_model`` / ``openai_chat_completions``). Covers the
|
||||
public-surface guarantees: path-traversal rejection, request sanitization
|
||||
(tools / provider routing / use_adapter), asset-path containment, the page CSP
|
||||
header + HTML escaping, and that the preview lock is held until a streaming
|
||||
response is fully drained.
|
||||
public-surface guarantees: HMAC capability gating (a valid ``?k=`` token or
|
||||
Bearer credential is required; missing/invalid/wrong-ref tokens 404 before any
|
||||
model load), path-traversal rejection, request sanitization (tools / provider
|
||||
routing / use_adapter / generation clamp), asset-path containment, the page CSP
|
||||
+ no-referrer headers and HTML escaping, and that the preview lock is held until
|
||||
a streaming response is fully drained.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
|
@ -34,9 +36,23 @@ from fastapi.responses import StreamingResponse
|
|||
from fastapi.testclient import TestClient
|
||||
|
||||
import routes.preview as preview
|
||||
import utils.preview_token as preview_token
|
||||
from models.inference import ChatCompletionRequest
|
||||
|
||||
|
||||
# A fixed secret keeps signing deterministic and avoids touching auth.db.
|
||||
_TEST_SECRET = b"unit-test-preview-secret-0123456789"
|
||||
|
||||
|
||||
def _use_test_secret(monkeypatch) -> None:
|
||||
monkeypatch.setattr(preview_token, "get_or_create_preview_link_secret", lambda: _TEST_SECRET)
|
||||
|
||||
|
||||
def _sig(ref: str) -> str:
|
||||
"""Valid capability token for ``ref`` under the patched test secret."""
|
||||
return preview_token.sign_preview_ref(ref)
|
||||
|
||||
|
||||
def _make_run(outputs: Path, name: str = "demorun") -> Path:
|
||||
run = outputs / name
|
||||
run.mkdir(parents = True)
|
||||
|
|
@ -59,6 +75,14 @@ def client(tmp_path, monkeypatch, captured):
|
|||
outputs = tmp_path / "outputs"
|
||||
_make_run(outputs)
|
||||
|
||||
_use_test_secret(monkeypatch)
|
||||
|
||||
# Public sharing on by default; reset the per-IP rate buckets each test.
|
||||
monkeypatch.setattr(preview, "get_preview_sharing_enabled", lambda: True)
|
||||
import utils.preview_rate_limit as _rl
|
||||
|
||||
_rl.reset()
|
||||
|
||||
# resolve_preview_checkpoint -> resolve_output_dir -> outputs_root().
|
||||
from utils.paths import storage_roots as _sr
|
||||
|
||||
|
|
@ -86,18 +110,21 @@ def client(tmp_path, monkeypatch, captured):
|
|||
|
||||
|
||||
def test_page_renders_with_csp(client):
|
||||
r = client.get("/p/demorun")
|
||||
r = client.get(f"/p/demorun?k={_sig('demorun')}")
|
||||
assert r.status_code == 200
|
||||
assert "text/html" in r.headers["content-type"]
|
||||
csp = r.headers.get("content-security-policy", "")
|
||||
assert "default-src 'self'" in csp
|
||||
assert "base-uri 'none'" in csp
|
||||
# Token rides in the query string; keep it out of the Referer header.
|
||||
assert r.headers.get("referrer-policy") == "no-referrer"
|
||||
|
||||
|
||||
def test_page_escapes_title(tmp_path, monkeypatch, captured):
|
||||
outputs = tmp_path / "outputs"
|
||||
# Run dir name carries an HTML-special char; the page must escape it.
|
||||
_make_run(outputs, name = "a<b")
|
||||
_use_test_secret(monkeypatch)
|
||||
from utils.paths import storage_roots as _sr
|
||||
|
||||
monkeypatch.setattr(_sr, "outputs_root", lambda: outputs)
|
||||
|
|
@ -106,14 +133,15 @@ def test_page_escapes_title(tmp_path, monkeypatch, captured):
|
|||
app.include_router(preview.router, prefix = "/p")
|
||||
c = TestClient(app, raise_server_exceptions = False)
|
||||
|
||||
r = c.get("/p/a%3Cb")
|
||||
# Sign the decoded canonical ref ("a<b"), not the %-encoded path segment.
|
||||
r = c.get(f"/p/a%3Cb?k={_sig('a<b')}")
|
||||
assert r.status_code == 200
|
||||
assert "a<b" not in r.text
|
||||
assert "a<b" in r.text
|
||||
|
||||
|
||||
def test_models_endpoint_shape(client):
|
||||
r = client.get("/p/demorun/v1/models")
|
||||
r = client.get(f"/p/demorun/v1/models?k={_sig('demorun')}")
|
||||
assert r.status_code == 200
|
||||
body = r.json()
|
||||
assert body["object"] == "list"
|
||||
|
|
@ -131,6 +159,25 @@ def test_list_previews_builds_urls(client, monkeypatch):
|
|||
assert r.status_code == 200
|
||||
data = r.json()["data"]
|
||||
assert data[0]["url"].endswith("/p/demorun/v1")
|
||||
# The listing hands the authenticated owner a usable capability.
|
||||
assert data[0]["key"] == _sig("demorun")
|
||||
assert data[0]["share_url"].endswith(f"/p/demorun?k={_sig('demorun')}")
|
||||
|
||||
|
||||
def test_list_previews_omits_capability_when_sharing_disabled(client, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
preview,
|
||||
"list_preview_targets",
|
||||
lambda: [{"ref": "demorun", "is_latest": True}],
|
||||
)
|
||||
monkeypatch.setattr(preview, "get_preview_sharing_enabled", lambda: False)
|
||||
r = client.get("/p")
|
||||
assert r.status_code == 200
|
||||
body = r.json()
|
||||
# Don't hand out credentials that 404; signal the disabled state instead.
|
||||
assert body["sharing_enabled"] is False
|
||||
assert body["data"][0]["key"] is None
|
||||
assert body["data"][0]["share_url"] is None
|
||||
|
||||
|
||||
# ── Path traversal / containment ────────────────────────────────────────────
|
||||
|
|
@ -179,7 +226,7 @@ def test_asset_path_contained(client, asset):
|
|||
|
||||
def test_chat_payload_sanitized(client, captured):
|
||||
r = client.post(
|
||||
"/p/demorun/v1/chat/completions",
|
||||
f"/p/demorun/v1/chat/completions?k={_sig('demorun')}",
|
||||
json = {
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"tools": [{"type": "function", "function": {"name": "rm", "parameters": {}}}],
|
||||
|
|
@ -217,6 +264,10 @@ def test_chat_payload_sanitized(client, captured):
|
|||
assert p.external_model is None
|
||||
# Adapter pinned on for LoRA: a caller can't flip the shared backend to base.
|
||||
assert p.use_adapter is True
|
||||
# Generation cost capped on this public surface (no override sent -> ceiling).
|
||||
assert p.max_tokens == preview._PREVIEW_MAX_OUTPUT_TOKENS
|
||||
assert p.max_completion_tokens == preview._PREVIEW_MAX_OUTPUT_TOKENS
|
||||
assert p.n == 1
|
||||
# Loads the resolved checkpoint dir, not an attacker-supplied path.
|
||||
assert captured["load_path"].endswith("demorun")
|
||||
|
||||
|
|
@ -228,6 +279,7 @@ def test_merged_checkpoint_strips_use_adapter(tmp_path, monkeypatch, captured):
|
|||
merged.mkdir(parents = True)
|
||||
(merged / "config.json").write_text(json.dumps({"_name_or_path": "some/base"}))
|
||||
|
||||
_use_test_secret(monkeypatch)
|
||||
from utils.paths import storage_roots as _sr
|
||||
|
||||
monkeypatch.setattr(_sr, "outputs_root", lambda: outputs)
|
||||
|
|
@ -246,7 +298,7 @@ def test_merged_checkpoint_strips_use_adapter(tmp_path, monkeypatch, captured):
|
|||
app.include_router(preview.router, prefix = "/p")
|
||||
c = TestClient(app, raise_server_exceptions = False)
|
||||
r = c.post(
|
||||
"/p/mergedrun/v1/chat/completions",
|
||||
f"/p/mergedrun/v1/chat/completions?k={_sig('mergedrun')}",
|
||||
json = {"messages": [{"role": "user", "content": "hi"}], "use_adapter": False},
|
||||
)
|
||||
assert r.status_code == 200
|
||||
|
|
@ -291,3 +343,154 @@ def test_streaming_holds_lock_until_drained(tmp_path, monkeypatch, captured):
|
|||
chunks = asyncio.run(_run())
|
||||
assert any(b"[DONE]" in c for c in chunks)
|
||||
assert not preview._preview_lock.locked()
|
||||
|
||||
|
||||
# ── Capability gating ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_chat_without_token_404_and_no_load(client, captured):
|
||||
r = client.post(
|
||||
"/p/demorun/v1/chat/completions",
|
||||
json = {"messages": [{"role": "user", "content": "hi"}]},
|
||||
)
|
||||
assert r.status_code == 404
|
||||
# Verified before any model work: nothing loaded, nothing generated.
|
||||
assert "load_path" not in captured
|
||||
assert "payload" not in captured
|
||||
|
||||
|
||||
def test_chat_with_invalid_token_404(client, captured):
|
||||
r = client.post(
|
||||
"/p/demorun/v1/chat/completions?k=not-a-valid-token",
|
||||
json = {"messages": [{"role": "user", "content": "hi"}]},
|
||||
)
|
||||
assert r.status_code == 404
|
||||
assert "load_path" not in captured
|
||||
|
||||
|
||||
def test_token_for_other_ref_rejected(client, captured):
|
||||
# A capability minted for a different ref must not unlock demorun.
|
||||
r = client.post(
|
||||
f"/p/demorun/v1/chat/completions?k={_sig('otherrun')}",
|
||||
json = {"messages": [{"role": "user", "content": "hi"}]},
|
||||
)
|
||||
assert r.status_code == 404
|
||||
assert "load_path" not in captured
|
||||
|
||||
|
||||
def test_models_without_token_404(client):
|
||||
assert client.get("/p/demorun/v1/models").status_code == 404
|
||||
|
||||
|
||||
def test_page_without_token_404(client):
|
||||
assert client.get("/p/demorun").status_code == 404
|
||||
|
||||
|
||||
def test_checkpoint_route_with_valid_sig(client, captured):
|
||||
# Nested ref: the signed/verified/resolved canonical ref is "run/checkpoint".
|
||||
sig = _sig("demorun/checkpoint-1")
|
||||
r = client.post(
|
||||
f"/p/demorun/checkpoint-1/v1/chat/completions?k={sig}",
|
||||
json = {"messages": [{"role": "user", "content": "hi"}]},
|
||||
)
|
||||
assert r.status_code == 200
|
||||
assert captured["load_path"].endswith("checkpoint-1")
|
||||
|
||||
|
||||
def test_checkpoint_token_does_not_unlock_bare_run(client, captured):
|
||||
# A token minted for the nested checkpoint must not unlock the run ref.
|
||||
r = client.post(
|
||||
f"/p/demorun/v1/chat/completions?k={_sig('demorun/checkpoint-1')}",
|
||||
json = {"messages": [{"role": "user", "content": "hi"}]},
|
||||
)
|
||||
assert r.status_code == 404
|
||||
assert "load_path" not in captured
|
||||
|
||||
|
||||
def test_bearer_token_accepted(client, captured):
|
||||
# OpenAI-compatible clients pass the capability as the api_key (Bearer header).
|
||||
r = client.post(
|
||||
"/p/demorun/v1/chat/completions",
|
||||
headers = {"Authorization": f"Bearer {_sig('demorun')}"},
|
||||
json = {"messages": [{"role": "user", "content": "hi"}]},
|
||||
)
|
||||
assert r.status_code == 200
|
||||
assert captured["load_path"].endswith("demorun")
|
||||
|
||||
|
||||
def test_generation_clamp_caps_overrides(client, captured):
|
||||
r = client.post(
|
||||
f"/p/demorun/v1/chat/completions?k={_sig('demorun')}",
|
||||
json = {
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"max_tokens": 999999,
|
||||
"max_completion_tokens": 888888,
|
||||
"n": 64,
|
||||
},
|
||||
)
|
||||
assert r.status_code == 200
|
||||
p = captured["payload"]
|
||||
assert p.max_tokens == preview._PREVIEW_MAX_OUTPUT_TOKENS
|
||||
assert p.max_completion_tokens == preview._PREVIEW_MAX_OUTPUT_TOKENS
|
||||
assert p.n == 1
|
||||
|
||||
|
||||
def test_generation_clamp_honors_lower_legacy_max_tokens(client, captured):
|
||||
# A caller asking for fewer tokens via the legacy field must not be bumped up
|
||||
# to the ceiling: _effective_max_tokens prefers max_completion_tokens, so both
|
||||
# fields have to carry the lower value.
|
||||
r = client.post(
|
||||
f"/p/demorun/v1/chat/completions?k={_sig('demorun')}",
|
||||
json = {"messages": [{"role": "user", "content": "hi"}], "max_tokens": 16},
|
||||
)
|
||||
assert r.status_code == 200
|
||||
p = captured["payload"]
|
||||
assert p.max_tokens == 16
|
||||
assert p.max_completion_tokens == 16
|
||||
|
||||
|
||||
def test_generation_clamp_honors_lower_completion_tokens(client, captured):
|
||||
r = client.post(
|
||||
f"/p/demorun/v1/chat/completions?k={_sig('demorun')}",
|
||||
json = {"messages": [{"role": "user", "content": "hi"}], "max_completion_tokens": 32},
|
||||
)
|
||||
assert r.status_code == 200
|
||||
p = captured["payload"]
|
||||
assert p.max_tokens == 32
|
||||
assert p.max_completion_tokens == 32
|
||||
|
||||
|
||||
# ── Public-sharing kill switch ───────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_chat_blocked_when_sharing_disabled(client, monkeypatch, captured):
|
||||
# Admin turned public sharing off: even a valid token 404s, with no model load.
|
||||
monkeypatch.setattr(preview, "get_preview_sharing_enabled", lambda: False)
|
||||
r = client.post(
|
||||
f"/p/demorun/v1/chat/completions?k={_sig('demorun')}",
|
||||
json = {"messages": [{"role": "user", "content": "hi"}]},
|
||||
)
|
||||
assert r.status_code == 404
|
||||
assert "load_path" not in captured
|
||||
|
||||
|
||||
def test_page_blocked_when_sharing_disabled(client, monkeypatch):
|
||||
monkeypatch.setattr(preview, "get_preview_sharing_enabled", lambda: False)
|
||||
assert client.get(f"/p/demorun?k={_sig('demorun')}").status_code == 404
|
||||
|
||||
|
||||
# ── Rate limiting ────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_chat_rate_limited_returns_429(client, monkeypatch):
|
||||
import utils.preview_rate_limit as rl
|
||||
|
||||
monkeypatch.setattr(rl, "_MAX_REQUESTS", 2)
|
||||
rl.reset()
|
||||
url = f"/p/demorun/v1/chat/completions?k={_sig('demorun')}"
|
||||
body = {"messages": [{"role": "user", "content": "hi"}]}
|
||||
assert client.post(url, json = body).status_code == 200
|
||||
assert client.post(url, json = body).status_code == 200
|
||||
r = client.post(url, json = body)
|
||||
assert r.status_code == 429
|
||||
assert r.headers.get("retry-after")
|
||||
|
|
|
|||
77
studio/backend/tests/test_preview_sharing_settings.py
Normal file
77
studio/backend/tests/test_preview_sharing_settings.py
Normal file
|
|
@ -0,0 +1,77 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Route-level tests for the preview settings endpoints (rotate + sharing toggle)."""
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import types as _types
|
||||
|
||||
|
||||
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
||||
if _BACKEND_DIR not in sys.path:
|
||||
sys.path.insert(0, _BACKEND_DIR)
|
||||
|
||||
_loggers_stub = _types.ModuleType("loggers")
|
||||
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
|
||||
sys.modules.setdefault("loggers", _loggers_stub)
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import routes.settings as settings
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(monkeypatch):
|
||||
# Stub the persistence helpers so the endpoints don't touch the real DBs.
|
||||
calls: dict = {"enabled": True}
|
||||
|
||||
def _set(value):
|
||||
calls["set"] = bool(value)
|
||||
calls["enabled"] = bool(value)
|
||||
return bool(value)
|
||||
|
||||
monkeypatch.setattr(settings, "get_preview_sharing_enabled", lambda: calls["enabled"])
|
||||
monkeypatch.setattr(settings, "set_preview_sharing_enabled", _set)
|
||||
monkeypatch.setattr(
|
||||
settings, "rotate_preview_link_secret", lambda: calls.__setitem__("rotated", True)
|
||||
)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(settings.router)
|
||||
app.dependency_overrides[settings.get_current_subject] = lambda: "admin"
|
||||
return TestClient(app, raise_server_exceptions = False), calls
|
||||
|
||||
|
||||
def test_rotate_preview_links(client):
|
||||
c, calls = client
|
||||
r = c.post("/preview-links/rotate")
|
||||
assert r.status_code == 200
|
||||
assert r.json() == {"rotated": True}
|
||||
assert calls.get("rotated") is True
|
||||
|
||||
|
||||
def test_get_preview_sharing(client):
|
||||
c, _ = client
|
||||
r = c.get("/preview-sharing")
|
||||
assert r.status_code == 200
|
||||
body = r.json()
|
||||
assert body["enabled"] is True
|
||||
assert "default_enabled" in body
|
||||
|
||||
|
||||
def test_put_preview_sharing_disables(client):
|
||||
c, calls = client
|
||||
r = c.put("/preview-sharing", json = {"enabled": False})
|
||||
assert r.status_code == 200
|
||||
assert r.json()["enabled"] is False
|
||||
assert calls["set"] is False
|
||||
|
||||
|
||||
def test_put_preview_sharing_rejects_non_bool(client):
|
||||
# Pydantic rejects a non-bool body (422) before the handler runs.
|
||||
c, _ = client
|
||||
r = c.put("/preview-sharing", json = {"enabled": "maybe"})
|
||||
assert r.status_code == 422
|
||||
75
studio/backend/tests/test_preview_token.py
Normal file
75
studio/backend/tests/test_preview_token.py
Normal file
|
|
@ -0,0 +1,75 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Unit + rotation coverage for `/p` preview capability tokens.
|
||||
|
||||
The token turns a guessable preview ref into an unguessable bearer capability:
|
||||
it must round-trip for the ref it was signed for, reject tampering / wrong refs,
|
||||
and stop verifying once the signing secret is rotated (link revocation).
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import types as _types
|
||||
|
||||
|
||||
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
||||
if _BACKEND_DIR not in sys.path:
|
||||
sys.path.insert(0, _BACKEND_DIR)
|
||||
|
||||
# Mirror the other preview tests: avoid the heavy real `loggers` handlers.
|
||||
_loggers_stub = _types.ModuleType("loggers")
|
||||
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
|
||||
sys.modules.setdefault("loggers", _loggers_stub)
|
||||
|
||||
import auth.storage as storage
|
||||
import utils.preview_token as preview_token
|
||||
|
||||
|
||||
_S1 = b"secret-one-aaaaaaaaaaaaaaaaaaaaaaaa"
|
||||
_S2 = b"secret-two-bbbbbbbbbbbbbbbbbbbbbbbb"
|
||||
|
||||
|
||||
def test_sign_verify_roundtrip(monkeypatch):
|
||||
monkeypatch.setattr(preview_token, "get_or_create_preview_link_secret", lambda: _S1)
|
||||
token = preview_token.sign_preview_ref("run/checkpoint-1")
|
||||
assert preview_token.verify_preview_ref("run/checkpoint-1", token)
|
||||
# URL-safe, unpadded, and high-entropy (SHA-256 -> 43 base64url chars).
|
||||
assert "=" not in token and "/" not in token and "+" not in token
|
||||
assert len(token) >= 40
|
||||
|
||||
|
||||
def test_missing_or_tampered_token_rejected(monkeypatch):
|
||||
monkeypatch.setattr(preview_token, "get_or_create_preview_link_secret", lambda: _S1)
|
||||
token = preview_token.sign_preview_ref("demorun")
|
||||
assert not preview_token.verify_preview_ref("demorun", None)
|
||||
assert not preview_token.verify_preview_ref("demorun", "")
|
||||
flipped = token[:-1] + ("A" if token[-1] != "A" else "B")
|
||||
assert not preview_token.verify_preview_ref("demorun", flipped)
|
||||
# A token minted for one ref does not unlock another.
|
||||
assert not preview_token.verify_preview_ref("otherrun", token)
|
||||
# A non-ASCII token is invalid, not a crash (the route would 500 otherwise).
|
||||
assert not preview_token.verify_preview_ref("demorun", "tøken-é")
|
||||
|
||||
|
||||
def test_secret_change_invalidates_token(monkeypatch):
|
||||
monkeypatch.setattr(preview_token, "get_or_create_preview_link_secret", lambda: _S1)
|
||||
token = preview_token.sign_preview_ref("demorun")
|
||||
monkeypatch.setattr(preview_token, "get_or_create_preview_link_secret", lambda: _S2)
|
||||
assert not preview_token.verify_preview_ref("demorun", token)
|
||||
|
||||
|
||||
def test_rotation_revokes_links(tmp_path, monkeypatch):
|
||||
# Exercise the real storage helpers against a throwaway auth.db.
|
||||
monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db")
|
||||
monkeypatch.setattr(storage, "_preview_link_secret_cache", None)
|
||||
|
||||
token = preview_token.sign_preview_ref("demorun")
|
||||
assert preview_token.verify_preview_ref("demorun", token)
|
||||
# Secret persists across calls (the link keeps working until rotated).
|
||||
assert preview_token.verify_preview_ref("demorun", token)
|
||||
|
||||
storage.rotate_preview_link_secret()
|
||||
# Old shared link is revoked; a freshly minted one works.
|
||||
assert not preview_token.verify_preview_ref("demorun", token)
|
||||
assert preview_token.verify_preview_ref("demorun", preview_token.sign_preview_ref("demorun"))
|
||||
|
|
@ -90,6 +90,23 @@ def test_update_run_whitespace_clears_display_name(monkeypatch: pytest.MonkeyPat
|
|||
assert result.display_name is None
|
||||
|
||||
|
||||
def test_get_run_detail_includes_preview_fields(monkeypatch: pytest.MonkeyPatch):
|
||||
# Regression: detail/update must pass the sharing flag into _preview_fields;
|
||||
# a missing arg used to surface as a 500 TypeError after get_run succeeded.
|
||||
monkeypatch.setattr(training_history, "get_run", lambda run_id: dict(BASE_RUN))
|
||||
monkeypatch.setattr(training_history, "get_run_metrics", lambda run_id: {})
|
||||
monkeypatch.setattr(training_history, "can_resume_run", lambda run: False)
|
||||
monkeypatch.setattr(training_history, "get_preview_sharing_enabled", lambda: True)
|
||||
|
||||
detail = asyncio.run(
|
||||
training_history.get_training_run_detail("run-1", current_subject = "test-user")
|
||||
)
|
||||
|
||||
assert detail.run.id == "run-1"
|
||||
# Not a previewable dir, so no signed ref - but the field is built without error.
|
||||
assert detail.run.preview_sig is None
|
||||
|
||||
|
||||
def test_update_run_rejects_unknown_fields():
|
||||
with pytest.raises(ValidationError):
|
||||
TrainingRunUpdateRequest.model_validate({"unknown": "value"})
|
||||
|
|
|
|||
71
studio/backend/utils/client_ip.py
Normal file
71
studio/backend/utils/client_ip.py
Normal file
|
|
@ -0,0 +1,71 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Resolve the caller's IP for rate limiting.
|
||||
|
||||
Trust model, in order:
|
||||
1. If the operator opts in via ``UNSLOTH_STUDIO_TRUST_FORWARDED`` (Studio behind
|
||||
their own reverse proxy), honor the *rightmost* ``X-Forwarded-For`` hop -- the
|
||||
one the trusted proxy appended. The leftmost entry is client-controlled and
|
||||
spoofable, so this assumes a proxy that appends (or overwrites) the header;
|
||||
only enable the env var behind such a proxy.
|
||||
2. If the socket peer is loopback, honor ``CF-Connecting-IP``. Studio's managed
|
||||
Cloudflare tunnel terminates at 127.0.0.1, so every tunneled visitor would
|
||||
otherwise collapse onto the same socket peer (the local cloudflared process)
|
||||
and share one rate-limit bucket. ``CF-Connecting-IP`` is set by Cloudflare's
|
||||
edge and can't be forged by a tunneled client.
|
||||
3. Otherwise the socket peer, so a direct LAN caller can't spoof a header to
|
||||
dodge a per-IP limit.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ipaddress
|
||||
import os
|
||||
|
||||
_TRUST_FORWARDED_ENV = "UNSLOTH_STUDIO_TRUST_FORWARDED"
|
||||
|
||||
|
||||
def _trust_forwarded_for() -> bool:
|
||||
return os.environ.get(_TRUST_FORWARDED_ENV, "").strip().lower() in {"1", "true", "yes"}
|
||||
|
||||
|
||||
def _is_loopback(host: str | None) -> bool:
|
||||
try:
|
||||
return bool(host) and ipaddress.ip_address(host).is_loopback
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def _normalize_addr(value: str | None) -> str | None:
|
||||
"""Parse an ``X-Forwarded-For`` entry into a bare, validated IP (strip port/brackets)."""
|
||||
raw = (value or "").strip().strip('"')
|
||||
if not raw:
|
||||
return None
|
||||
if raw.startswith("["): # [ipv6]:port
|
||||
raw = raw[1:].split("]", 1)[0]
|
||||
elif raw.count(":") == 1: # ipv4:port
|
||||
raw = raw.split(":", 1)[0]
|
||||
try:
|
||||
return ipaddress.ip_address(raw).compressed
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def client_ip(request) -> str:
|
||||
"""Best-effort client IP, or ``"_unknown"`` when it can't be determined."""
|
||||
if request is None:
|
||||
return "_unknown"
|
||||
peer = request.client.host if request.client else None
|
||||
if _trust_forwarded_for():
|
||||
# Rightmost hop = what the trusted proxy saw; the leftmost is spoofable.
|
||||
xff = request.headers.get("x-forwarded-for", "")
|
||||
if xff:
|
||||
normalized = _normalize_addr(xff.rsplit(",", 1)[-1])
|
||||
if normalized:
|
||||
return normalized
|
||||
if _is_loopback(peer):
|
||||
cf = _normalize_addr(request.headers.get("cf-connecting-ip"))
|
||||
if cf:
|
||||
return cf
|
||||
return peer or "_unknown"
|
||||
67
studio/backend/utils/preview_rate_limit.py
Normal file
67
studio/backend/utils/preview_rate_limit.py
Normal file
|
|
@ -0,0 +1,67 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Coarse per-IP sliding-window rate limit for the public ``/p`` preview chat.
|
||||
|
||||
A signed link stops ref guessing, but anyone with a link can still drive GPU
|
||||
generation. This bounds sustained abuse from a single source. In-process and
|
||||
single-worker only (like the login limiter in ``routes/auth.py``); Studio runs as
|
||||
one uvicorn process, so a shared store isn't needed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
from collections import deque
|
||||
|
||||
# Window / ceiling for preview chat-completions per client IP.
|
||||
_WINDOW_SECONDS = 60.0
|
||||
_MAX_REQUESTS = 20
|
||||
# Bound memory on a public surface (many distinct IPs).
|
||||
_MAX_BUCKETS = 4096
|
||||
|
||||
_buckets: dict[str, deque] = {}
|
||||
_lock = threading.Lock()
|
||||
|
||||
|
||||
def _prune(bucket: deque, now: float) -> None:
|
||||
while bucket and now - bucket[0] > _WINDOW_SECONDS:
|
||||
bucket.popleft()
|
||||
|
||||
|
||||
def _evict_aged(now: float) -> None:
|
||||
"""Drop only buckets that have fully aged out. Never evict an active bucket:
|
||||
evicting a throttled key would reset its counter, so a flood of distinct keys
|
||||
could cycle the table and clear a victim's (or its own) limit."""
|
||||
for key in list(_buckets.keys()):
|
||||
_prune(_buckets[key], now)
|
||||
if not _buckets[key]:
|
||||
del _buckets[key]
|
||||
|
||||
|
||||
def check_rate_limit(key: str) -> int:
|
||||
"""Record a hit for ``key``; return seconds-to-wait if over the limit, else 0."""
|
||||
now = time.monotonic()
|
||||
with _lock:
|
||||
bucket = _buckets.get(key)
|
||||
if bucket is None:
|
||||
if len(_buckets) >= _MAX_BUCKETS:
|
||||
_evict_aged(now)
|
||||
if len(_buckets) >= _MAX_BUCKETS:
|
||||
# Table is full of currently-active clients. Fail closed: deny the
|
||||
# new key rather than evict a live bucket (which would hand out a
|
||||
# rate-limit reset). Pathological only (>= _MAX_BUCKETS live IPs).
|
||||
return max(1, int(_WINDOW_SECONDS))
|
||||
bucket = _buckets[key] = deque()
|
||||
_prune(bucket, now)
|
||||
if len(bucket) >= _MAX_REQUESTS:
|
||||
return max(1, int(_WINDOW_SECONDS - (now - bucket[0])) + 1)
|
||||
bucket.append(now)
|
||||
return 0
|
||||
|
||||
|
||||
def reset() -> None:
|
||||
"""Clear all buckets (test isolation)."""
|
||||
with _lock:
|
||||
_buckets.clear()
|
||||
55
studio/backend/utils/preview_sharing_settings.py
Normal file
55
studio/backend/utils/preview_sharing_settings.py
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Persisted kill switch for public ``/p`` preview link sharing."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
PREVIEW_SHARING_SETTING_KEY = "preview_public_sharing_enabled"
|
||||
# Default on: signed share links work out of the box (current behavior). An admin
|
||||
# can flip this off to take the public ``/p`` surface offline entirely - links
|
||||
# then 404 even with a valid token, leaving preview to the authenticated app.
|
||||
DEFAULT_PREVIEW_SHARING_ENABLED = True
|
||||
|
||||
|
||||
def _coerce_bool(value: Any) -> bool | None:
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
if isinstance(value, str):
|
||||
normalized = value.strip().lower()
|
||||
if normalized in {"1", "true", "yes", "on"}:
|
||||
return True
|
||||
if normalized in {"0", "false", "no", "off", ""}:
|
||||
return False
|
||||
return None
|
||||
|
||||
|
||||
def get_preview_sharing_enabled() -> bool:
|
||||
"""Read the persisted public-preview-sharing preference.
|
||||
|
||||
A *missing* setting defaults to enabled so the feature keeps working as
|
||||
before unless an admin explicitly turns it off. A *read failure* (e.g. a
|
||||
transient SQLite/permission error) fails closed -- this is a kill switch, so
|
||||
an unreadable settings DB must not silently reopen the public surface.
|
||||
"""
|
||||
try:
|
||||
from storage.studio_db import get_app_setting
|
||||
stored = get_app_setting(PREVIEW_SHARING_SETTING_KEY, None)
|
||||
except Exception:
|
||||
return False
|
||||
parsed = _coerce_bool(stored)
|
||||
return parsed if parsed is not None else DEFAULT_PREVIEW_SHARING_ENABLED
|
||||
|
||||
|
||||
def set_preview_sharing_enabled(value: Any) -> bool:
|
||||
"""Persist whether public ``/p`` preview links are accepted."""
|
||||
parsed = _coerce_bool(value)
|
||||
if parsed is None:
|
||||
raise ValueError("Public preview sharing must be true or false.")
|
||||
|
||||
from storage.studio_db import upsert_app_settings
|
||||
|
||||
upsert_app_settings({PREVIEW_SHARING_SETTING_KEY: parsed})
|
||||
return parsed
|
||||
52
studio/backend/utils/preview_token.py
Normal file
52
studio/backend/utils/preview_token.py
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""HMAC capability tokens for public ``/p`` preview share links.
|
||||
|
||||
The preview ref (``run`` or ``run/checkpoint``) is a deterministic, guessable
|
||||
outputs-root path, so it can't gate access on its own. We sign the canonical ref
|
||||
with a dedicated server-side secret and require the resulting token on every
|
||||
public preview request: guessing a ref no longer grants access, and rotating the
|
||||
secret (``auth.storage.rotate_preview_link_secret``) revokes every link at once.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import hmac
|
||||
from typing import Optional
|
||||
|
||||
from auth.storage import get_or_create_preview_link_secret
|
||||
|
||||
# Versioned so the token format can evolve without silently honoring old shapes.
|
||||
_PREVIEW_TOKEN_VERSION = "v1"
|
||||
|
||||
|
||||
def _canonical_payload(ref: str) -> bytes:
|
||||
# Sign the canonical ref only (never host/path) so links stay portable across
|
||||
# localhost / LAN IP / tunnel host changes.
|
||||
return f"preview:{_PREVIEW_TOKEN_VERSION}:{ref}".encode("utf-8")
|
||||
|
||||
|
||||
def sign_preview_ref(ref: str) -> str:
|
||||
"""Return the URL-safe HMAC capability token for a canonical preview ref."""
|
||||
mac = hmac.new(
|
||||
get_or_create_preview_link_secret(),
|
||||
_canonical_payload(ref),
|
||||
hashlib.sha256,
|
||||
).digest()
|
||||
return base64.urlsafe_b64encode(mac).rstrip(b"=").decode("ascii")
|
||||
|
||||
|
||||
def verify_preview_ref(ref: str, token: Optional[str]) -> bool:
|
||||
"""Constant-time check that ``token`` is a valid capability for ``ref``."""
|
||||
if not token:
|
||||
return False
|
||||
# Compare as bytes: a non-ASCII token (e.g. a %-encoded query value) would make
|
||||
# hmac.compare_digest on two str raise TypeError -> treat it as simply invalid.
|
||||
try:
|
||||
provided = token.encode("ascii")
|
||||
except UnicodeEncodeError:
|
||||
return False
|
||||
return hmac.compare_digest(sign_preview_ref(ref).encode("ascii"), provided)
|
||||
62
studio/frontend/src/features/settings/api/preview-sharing.ts
Normal file
62
studio/frontend/src/features/settings/api/preview-sharing.ts
Normal file
|
|
@ -0,0 +1,62 @@
|
|||
// SPDX-License-Identifier: AGPL-3.0-only
|
||||
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
import { authFetch } from "@/features/auth";
|
||||
import { readFastApiError } from "@/lib/format-fastapi-error";
|
||||
|
||||
export type PreviewSharingSettings = {
|
||||
enabled: boolean;
|
||||
defaultEnabled: boolean;
|
||||
};
|
||||
|
||||
type ApiPreviewSharingSettings = {
|
||||
enabled: boolean;
|
||||
// biome-ignore lint/style/useNamingConvention: API schema
|
||||
default_enabled: boolean;
|
||||
};
|
||||
|
||||
function fromApi(settings: ApiPreviewSharingSettings): PreviewSharingSettings {
|
||||
return {
|
||||
enabled: settings.enabled,
|
||||
defaultEnabled: settings.default_enabled,
|
||||
};
|
||||
}
|
||||
|
||||
export async function loadPreviewSharing(): Promise<PreviewSharingSettings> {
|
||||
const res = await authFetch("/api/settings/preview-sharing");
|
||||
if (!res.ok) {
|
||||
throw new Error(
|
||||
await readFastApiError(res, "Failed to load preview sharing settings"),
|
||||
);
|
||||
}
|
||||
return fromApi(await res.json());
|
||||
}
|
||||
|
||||
export async function updatePreviewSharing(
|
||||
enabled: boolean,
|
||||
): Promise<PreviewSharingSettings> {
|
||||
const res = await authFetch("/api/settings/preview-sharing", {
|
||||
method: "PUT",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ enabled }),
|
||||
});
|
||||
if (!res.ok) {
|
||||
throw new Error(
|
||||
await readFastApiError(res, "Failed to update preview sharing settings"),
|
||||
);
|
||||
}
|
||||
return fromApi(await res.json());
|
||||
}
|
||||
|
||||
// Rotate the server-side signing secret, invalidating every previously shared
|
||||
// /p preview link in one step. Newly copied links keep working.
|
||||
export async function rotatePreviewLinks(): Promise<void> {
|
||||
const res = await authFetch("/api/settings/preview-links/rotate", {
|
||||
method: "POST",
|
||||
});
|
||||
if (!res.ok) {
|
||||
throw new Error(
|
||||
await readFastApiError(res, "Failed to revoke preview links"),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -13,18 +13,18 @@ import {
|
|||
import { Input } from "@/components/ui/input";
|
||||
import { Switch } from "@/components/ui/switch";
|
||||
import { usePlatformStore } from "@/config/env";
|
||||
import { isTauri } from "@/lib/api-base";
|
||||
import { openModelsDir } from "@/features/native-intents";
|
||||
import { copyToClipboard } from "@/lib/copy-to-clipboard";
|
||||
import { toast } from "@/lib/toast";
|
||||
import { loadModelsFolder, type ModelsFolder } from "../api/models-folder";
|
||||
import { resetOnboardingDone } from "@/features/auth";
|
||||
import { useChatRuntimeStore } from "@/features/chat";
|
||||
import { openModelsDir } from "@/features/native-intents";
|
||||
import { emitTrainingRunsChanged } from "@/features/training";
|
||||
import {
|
||||
setShowLlamaUpdateBanner,
|
||||
useShowLlamaUpdateBanner,
|
||||
} from "@/hooks/use-llama-update-pref";
|
||||
import { LOCALE_STORAGE_KEY, useT } from "@/i18n";
|
||||
import { isTauri } from "@/lib/api-base";
|
||||
import { copyToClipboard } from "@/lib/copy-to-clipboard";
|
||||
import { toast } from "@/lib/toast";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { useNavigate, useRouterState } from "@tanstack/react-router";
|
||||
import { Check, Eye, EyeOff } from "lucide-react";
|
||||
|
|
@ -34,6 +34,13 @@ import {
|
|||
loadHelperPrecacheSettings,
|
||||
updateHelperPrecacheSettings,
|
||||
} from "../api/helper-precache";
|
||||
import { type ModelsFolder, loadModelsFolder } from "../api/models-folder";
|
||||
import {
|
||||
type PreviewSharingSettings,
|
||||
loadPreviewSharing,
|
||||
rotatePreviewLinks,
|
||||
updatePreviewSharing,
|
||||
} from "../api/preview-sharing";
|
||||
import {
|
||||
DEFAULT_UPLOAD_LIMIT_MB,
|
||||
type UploadLimitSettings,
|
||||
|
|
@ -146,6 +153,14 @@ export function GeneralTab() {
|
|||
null,
|
||||
);
|
||||
const [isSavingHelperPrecache, setIsSavingHelperPrecache] = useState(false);
|
||||
const [previewSharing, setPreviewSharing] =
|
||||
useState<PreviewSharingSettings | null>(null);
|
||||
const [previewSharingError, setPreviewSharingError] = useState<string | null>(
|
||||
null,
|
||||
);
|
||||
const [isSavingPreviewSharing, setIsSavingPreviewSharing] = useState(false);
|
||||
const [revokePreviewOpen, setRevokePreviewOpen] = useState(false);
|
||||
const [isRevokingPreview, setIsRevokingPreview] = useState(false);
|
||||
const [modelsFolder, setModelsFolder] = useState<ModelsFolder | null>(null);
|
||||
|
||||
const draftRef = useRef(draftToken);
|
||||
|
|
@ -220,6 +235,27 @@ export function GeneralTab() {
|
|||
};
|
||||
}, [t]);
|
||||
|
||||
useEffect(() => {
|
||||
let cancelled = false;
|
||||
void loadPreviewSharing()
|
||||
.then((settings) => {
|
||||
if (cancelled) return;
|
||||
setPreviewSharing(settings);
|
||||
setPreviewSharingError(null);
|
||||
})
|
||||
.catch((error) => {
|
||||
if (cancelled) return;
|
||||
setPreviewSharingError(
|
||||
error instanceof Error
|
||||
? error.message
|
||||
: t("settings.general.previewSharing.loadError"),
|
||||
);
|
||||
});
|
||||
return () => {
|
||||
cancelled = true;
|
||||
};
|
||||
}, [t]);
|
||||
|
||||
useEffect(() => {
|
||||
let cancelled = false;
|
||||
void loadModelsFolder()
|
||||
|
|
@ -274,6 +310,44 @@ export function GeneralTab() {
|
|||
}
|
||||
};
|
||||
|
||||
const savePreviewSharing = async (enabled: boolean) => {
|
||||
setIsSavingPreviewSharing(true);
|
||||
setPreviewSharingError(null);
|
||||
try {
|
||||
const settings = await updatePreviewSharing(enabled);
|
||||
setPreviewSharing(settings);
|
||||
// Toggling sharing changes whether /api/train/runs returns preview_sig, so
|
||||
// refresh the history grid (hide/show the Copy preview link buttons).
|
||||
emitTrainingRunsChanged();
|
||||
} catch (error) {
|
||||
setPreviewSharingError(
|
||||
error instanceof Error
|
||||
? error.message
|
||||
: t("settings.general.previewSharing.saveError"),
|
||||
);
|
||||
} finally {
|
||||
setIsSavingPreviewSharing(false);
|
||||
}
|
||||
};
|
||||
|
||||
const revokePreviewLinks = async () => {
|
||||
setIsRevokingPreview(true);
|
||||
try {
|
||||
await rotatePreviewLinks();
|
||||
// The secret rotated, so any preview_sig the history grid still holds is
|
||||
// now stale. Refresh so copied links use freshly minted signatures.
|
||||
emitTrainingRunsChanged();
|
||||
setRevokePreviewOpen(false);
|
||||
toast.success(t("settings.general.previewSharing.revoked"));
|
||||
} catch (error) {
|
||||
toast.error(t("settings.general.previewSharing.revokeError"), {
|
||||
description: error instanceof Error ? error.message : undefined,
|
||||
});
|
||||
} finally {
|
||||
setIsRevokingPreview(false);
|
||||
}
|
||||
};
|
||||
|
||||
const saveUploadLimit = async () => {
|
||||
const parsed = Number(draftUploadLimit);
|
||||
if (!Number.isInteger(parsed)) {
|
||||
|
|
@ -454,6 +528,42 @@ export function GeneralTab() {
|
|||
</SettingsRow>
|
||||
</SettingsSection>
|
||||
|
||||
<SettingsSection
|
||||
title={t("settings.general.previewSharing.sectionTitle")}
|
||||
>
|
||||
<SettingsRow
|
||||
label={t("settings.general.previewSharing.enableLabel")}
|
||||
description={t("settings.general.previewSharing.enableDescription")}
|
||||
>
|
||||
<div className="flex flex-col items-end gap-1">
|
||||
<Switch
|
||||
checked={previewSharing?.enabled ?? false}
|
||||
disabled={!previewSharing || isSavingPreviewSharing}
|
||||
onCheckedChange={(enabled) => void savePreviewSharing(enabled)}
|
||||
/>
|
||||
{previewSharingError ? (
|
||||
<span className="max-w-[260px] text-right text-xs text-destructive">
|
||||
{previewSharingError}
|
||||
</span>
|
||||
) : null}
|
||||
</div>
|
||||
</SettingsRow>
|
||||
<SettingsRow
|
||||
destructive={true}
|
||||
label={t("settings.general.previewSharing.revokeLabel")}
|
||||
description={t("settings.general.previewSharing.revokeDescription")}
|
||||
>
|
||||
<Button
|
||||
variant="outline"
|
||||
size="sm"
|
||||
onClick={() => setRevokePreviewOpen(true)}
|
||||
className="text-destructive hover:text-destructive hover:border-destructive/60"
|
||||
>
|
||||
{t("settings.general.previewSharing.revokeAction")}
|
||||
</Button>
|
||||
</SettingsRow>
|
||||
</SettingsSection>
|
||||
|
||||
<SettingsSection title={t("settings.general.uploads.sectionTitle")}>
|
||||
<SettingsRow
|
||||
label={t("settings.general.uploads.maxUploadSize")}
|
||||
|
|
@ -561,6 +671,36 @@ export function GeneralTab() {
|
|||
</DialogFooter>
|
||||
</DialogContent>
|
||||
</Dialog>
|
||||
|
||||
<Dialog open={revokePreviewOpen} onOpenChange={setRevokePreviewOpen}>
|
||||
<DialogContent className="max-w-md">
|
||||
<DialogHeader>
|
||||
<DialogTitle>
|
||||
{t("settings.general.previewSharing.revokeConfirmTitle")}
|
||||
</DialogTitle>
|
||||
<DialogDescription>
|
||||
{t("settings.general.previewSharing.revokeConfirmDescription")}
|
||||
</DialogDescription>
|
||||
</DialogHeader>
|
||||
<DialogFooter>
|
||||
<Button
|
||||
variant="outline"
|
||||
onClick={() => setRevokePreviewOpen(false)}
|
||||
>
|
||||
{t("common.cancel")}
|
||||
</Button>
|
||||
<Button
|
||||
onClick={() => void revokePreviewLinks()}
|
||||
disabled={isRevokingPreview}
|
||||
className="bg-destructive hover:bg-destructive/90 text-destructive-foreground"
|
||||
>
|
||||
{isRevokingPreview
|
||||
? t("settings.general.previewSharing.revoking")
|
||||
: t("settings.general.previewSharing.revokeConfirmAction")}
|
||||
</Button>
|
||||
</DialogFooter>
|
||||
</DialogContent>
|
||||
</Dialog>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -387,8 +387,9 @@ export function HistoryCardGrid({
|
|||
const isRunning = run.status === "running";
|
||||
const canResume = run.can_resume && !wasContinued;
|
||||
const isResuming = resumeTarget === run.id;
|
||||
// Backend /p ref, gated on previewability + route-expressible depth.
|
||||
const canCopyPreview = !!run.preview_ref;
|
||||
// Backend /p ref + its capability token. Both are required: the link
|
||||
// is useless (404s) without the signature, so don't offer to copy it.
|
||||
const canCopyPreview = !!run.preview_ref && !!run.preview_sig;
|
||||
return (
|
||||
<div
|
||||
role="button"
|
||||
|
|
@ -456,7 +457,9 @@ export function HistoryCardGrid({
|
|||
serverUrl ??
|
||||
window.location.origin
|
||||
).replace(/\/+$/, "");
|
||||
const url = `${base}/p/${ref}`;
|
||||
// The signature is a bearer capability carried as ?k=; the
|
||||
// recipient's page forwards it on its chat requests.
|
||||
const url = `${base}/p/${ref}?k=${encodeURIComponent(run.preview_sig ?? "")}`;
|
||||
const ok = await copyToClipboard(url);
|
||||
toast[ok ? "success" : "error"](
|
||||
t(
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ export interface TrainingRunSummary {
|
|||
resumed_later: boolean;
|
||||
has_preview_model: boolean;
|
||||
preview_ref: string | null;
|
||||
preview_sig: string | null;
|
||||
duration_seconds: number | null;
|
||||
error_message: string | null;
|
||||
loss_sparkline: number[] | null;
|
||||
|
|
|
|||
|
|
@ -142,6 +142,25 @@ export const en = {
|
|||
loadError: "Failed to load Helper LLM settings.",
|
||||
saveError: "Failed to save Helper LLM settings.",
|
||||
},
|
||||
previewSharing: {
|
||||
sectionTitle: "Preview sharing",
|
||||
enableLabel: "Public preview links",
|
||||
enableDescription:
|
||||
"Let anyone with a signed link chat with a finished model, no login required. Turn off to take the public preview surface offline; shared links stop working.",
|
||||
loadError: "Failed to load preview sharing settings.",
|
||||
saveError: "Failed to save preview sharing settings.",
|
||||
revokeLabel: "Revoke all preview links",
|
||||
revokeDescription:
|
||||
"Rotate the signing secret so every link you've shared stops working. Newly copied links keep working.",
|
||||
revokeAction: "Revoke links",
|
||||
revoking: "Revoking...",
|
||||
revokeConfirmTitle: "Revoke all preview links?",
|
||||
revokeConfirmDescription:
|
||||
"Every preview link you've shared will stop working immediately. This can't be undone.",
|
||||
revokeConfirmAction: "Revoke all links",
|
||||
revoked: "All preview links revoked",
|
||||
revokeError: "Couldn't revoke preview links",
|
||||
},
|
||||
notifications: {
|
||||
sectionTitle: "Notifications",
|
||||
showLlamaUpdates: "llama.cpp update notifications",
|
||||
|
|
@ -161,8 +180,7 @@ export const en = {
|
|||
storage: {
|
||||
sectionTitle: "Storage",
|
||||
modelsFolder: "Models folder",
|
||||
modelsFolderDescription:
|
||||
"Where downloaded models are stored.",
|
||||
modelsFolderDescription: "Where downloaded models are stored.",
|
||||
openAction: "Open",
|
||||
copyAction: "Copy path",
|
||||
copied: "Path copied",
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue