diff --git a/studio/backend/hub/routes/__init__.py b/studio/backend/hub/routes/__init__.py index e9579635b0..7c5cfb9b3c 100644 --- a/studio/backend/hub/routes/__init__.py +++ b/studio/backend/hub/routes/__init__.py @@ -5,8 +5,10 @@ from hub.routes.inventory import router as inventory_router from hub.routes.datasets import router as datasets_router +from hub.routes.token import router as token_router __all__ = [ "inventory_router", "datasets_router", + "token_router", ] diff --git a/studio/backend/hub/routes/token.py b/studio/backend/hub/routes/token.py new file mode 100644 index 0000000000..1b7ad733a2 --- /dev/null +++ b/studio/backend/hub/routes/token.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Hugging Face token validation endpoint.""" + +from __future__ import annotations + +import asyncio +from typing import Literal, Optional + +from fastapi import APIRouter, Depends, Request +from pydantic import BaseModel + +from auth.authentication import get_current_subject +from hub.dependencies import get_hf_token +from utils.client_ip import client_ip +from utils.hf_token_validation import validate_hf_token + + +router = APIRouter() + + +class HfTokenValidationResponse(BaseModel): + status: Literal["missing", "valid", "invalid", "rate_limited", "unavailable"] + retry_after_seconds: Optional[int] = None + + +@router.post("/token/validate", response_model = HfTokenValidationResponse) +async def validate_token( + request: Request, + hf_token: Optional[str] = Depends(get_hf_token), + current_subject: str = Depends(get_current_subject), +): + if not hf_token: + return HfTokenValidationResponse(status = "missing") + result = await asyncio.to_thread( + validate_hf_token, + hf_token, + rate_key = f"{current_subject}:{client_ip(request)}", + ) + return HfTokenValidationResponse( + status = result.status, + retry_after_seconds = result.retry_after_seconds, + ) diff --git a/studio/backend/main.py b/studio/backend/main.py index 0ffc4489a1..f686e29bf5 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -312,6 +312,7 @@ from routes.preview import router as preview_router from hub.routes import ( inventory_router as hub_inventory_router, datasets_router as hub_datasets_router, + token_router as hub_token_router, ) from hub.schemas.downloads import TransportCapabilities from hub.utils.download_registry import ( @@ -993,6 +994,7 @@ 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"]) +app.include_router(hub_token_router, prefix = "/api/hub", 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. diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 9c08ea4b79..afd942e9a5 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1778,7 +1778,7 @@ from core.inference.providers import get_base_url from core.inference.external_provider import ExternalProviderClient from core.inference.chat_templates import resolve_effective_chat_template_override from storage import providers_db -from utils.utils import safe_error_detail, log_and_http_error +from utils.utils import is_hf_authentication_error, safe_error_detail, log_and_http_error import io import base64 @@ -5244,6 +5244,14 @@ async def validate_model( raise HTTPException(status_code = 400, detail = str(e)) except Exception as e: redacted_msg = redact_native_paths(str(e)) + if is_hf_authentication_error(e): + raise HTTPException( + status_code = 400, + detail = ( + "Hugging Face authentication failed. Check or clear the token " + "in Settings, and confirm access to this gated repository." + ), + ) if _is_unsupported_nvfp4_inference_error(redacted_msg): logger.warning( "NVFP4 inference is not supported yet while validating '%s'", diff --git a/studio/backend/tests/test_hf_token_validation.py b/studio/backend/tests/test_hf_token_validation.py new file mode 100644 index 0000000000..31b30fc37d --- /dev/null +++ b/studio/backend/tests/test_hf_token_validation.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Focused coverage for cached, rate-limited HF token validation.""" + +from __future__ import annotations + +from pathlib import Path +import sys + +import httpx +import pytest + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +import utils.hf_token_validation as validation + + +@pytest.fixture(autouse = True) +def _reset_validation_state(): + validation.reset_hf_token_validation_state() + yield + validation.reset_hf_token_validation_state() + + +def test_cached_token_does_not_spend_another_attempt(monkeypatch): + calls = [] + + def _check(token): + calls.append(token) + return validation.TokenValidationResult(status = "valid") + + monkeypatch.setattr(validation, "_check_remote", _check) + first = validation.validate_hf_token("hf_valid", rate_key = "user:ip") + second = validation.validate_hf_token("hf_valid", rate_key = "user:ip") + + assert first.status == second.status == "valid" + assert calls == ["hf_valid"] + + +def test_three_uncached_attempts_per_hour(monkeypatch): + monkeypatch.setattr( + validation, + "_check_remote", + lambda _token: validation.TokenValidationResult(status = "invalid"), + ) + + for index in range(3): + result = validation.validate_hf_token(f"hf_bad_{index}", rate_key = "user:ip") + assert result.status == "invalid" + + limited = validation.validate_hf_token("hf_bad_4", rate_key = "user:ip") + assert limited.status == "rate_limited" + assert limited.retry_after_seconds is not None + assert limited.retry_after_seconds > 0 + + other_user = validation.validate_hf_token("hf_other", rate_key = "other:ip") + assert other_user.status == "invalid" + + +def test_window_rolls_forward(monkeypatch): + clock = {"now": 100.0} + monkeypatch.setattr(validation.time, "monotonic", lambda: clock["now"]) + monkeypatch.setattr(validation, "_MAX_ATTEMPTS", 1) + monkeypatch.setattr(validation, "_WINDOW_SECONDS", 10.0) + monkeypatch.setattr( + validation, + "_check_remote", + lambda _token: validation.TokenValidationResult(status = "invalid"), + ) + + assert validation.validate_hf_token("hf_a", rate_key = "user:ip").status == "invalid" + assert validation.validate_hf_token("hf_b", rate_key = "user:ip").status == "rate_limited" + clock["now"] += 11.0 + assert validation.validate_hf_token("hf_b", rate_key = "user:ip").status == "invalid" + + +@pytest.mark.parametrize( + ("status_code", "expected"), + [(200, "valid"), (401, "invalid"), (429, "rate_limited"), (500, "unavailable")], +) +def test_remote_status_classification(monkeypatch, status_code, expected): + response = httpx.Response( + status_code, + request = httpx.Request("GET", "https://huggingface.co/api/whoami-v2"), + headers = {"Retry-After": "42"} if status_code == 429 else None, + ) + + class _Session: + def get(self, url, *, headers, timeout): + assert url == "https://huggingface.co/api/whoami-v2" + assert headers["authorization"] == "Bearer hf_test" + assert timeout == validation._REMOTE_TIMEOUT_SECONDS + return response + + monkeypatch.setattr(validation, "get_session", lambda: _Session()) + result = validation._check_remote("hf_test") + assert result.status == expected + if status_code == 429: + assert result.retry_after_seconds == 42 + + +def test_wrapped_http_401_is_invalid(monkeypatch): + response = httpx.Response( + 401, + request = httpx.Request("GET", "https://huggingface.co/api/whoami-v2"), + ) + + class _Session: + def get(self, _url, **_kwargs): + error = RuntimeError("Invalid user token.") + error.response = response + raise error + + monkeypatch.setattr(validation, "get_session", lambda: _Session()) + assert validation._check_remote("hf_test").status == "invalid" + + +def test_remote_timeout_is_bounded_and_unavailable(monkeypatch): + class _Session: + def get(self, _url, *, headers, timeout): + assert headers["authorization"] == "Bearer hf_test" + assert timeout == validation._REMOTE_TIMEOUT_SECONDS + raise TimeoutError("timed out") + + monkeypatch.setattr(validation, "get_session", lambda: _Session()) + assert validation._check_remote("hf_test").status == "unavailable" + + +def test_raw_token_is_not_retained(monkeypatch): + monkeypatch.setattr( + validation, + "_check_remote", + lambda _token: validation.TokenValidationResult(status = "valid"), + ) + token = "hf_do_not_store_this_value" + validation.validate_hf_token(token, rate_key = "user:ip") + + assert token not in repr(validation._cache) + assert token not in repr(validation._attempts) + + +def test_unexpected_remote_exception_releases_singleflight(monkeypatch): + calls = 0 + monkeypatch.setattr(validation, "_INFLIGHT_WAIT_SECONDS", 0.0) + + def _check(_token): + nonlocal calls + calls += 1 + if calls == 1: + raise RuntimeError("unexpected failure") + return validation.TokenValidationResult(status = "valid") + + monkeypatch.setattr(validation, "_check_remote", _check) + + with pytest.raises(RuntimeError, match = "unexpected failure"): + validation.validate_hf_token("hf_test", rate_key = "user:ip") + + result = validation.validate_hf_token("hf_test", rate_key = "user:ip") + assert result.status == "valid" + assert calls == 2 + assert validation._inflight == {} diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index 64a3c62156..741f19c67a 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -38,7 +38,7 @@ from utils.hardware import ( DeviceType, ) import utils.hardware.hardware as _hw_module -from utils.utils import format_error_message +from utils.utils import format_error_message, is_hf_authentication_error # ========== Helpers ========== @@ -439,6 +439,20 @@ class TestFormatErrorMessage: msg = format_error_message(err, "any/model") assert "invalid" in msg.lower() + def test_hf_authentication_error_follows_wrapped_401(self): + response = type("Response", (), {"status_code": 401})() + auth_error = Exception("request failed") + auth_error.response = response + wrapper = RuntimeError("model validation failed") + wrapper.__cause__ = auth_error + assert is_hf_authentication_error(wrapper) is True + + def test_hf_authentication_error_does_not_treat_429_as_invalid(self): + response = type("Response", (), {"status_code": 429})() + rate_error = Exception("too many requests") + rate_error.response = response + assert is_hf_authentication_error(rate_error) is False + # --- OOM on CUDA --- @needs_torch diff --git a/studio/backend/utils/hf_token_validation.py b/studio/backend/utils/hf_token_validation.py new file mode 100644 index 0000000000..7247c6e756 --- /dev/null +++ b/studio/backend/utils/hf_token_validation.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Cached, rate-limited Hugging Face token validation.""" + +from __future__ import annotations + +import hashlib +import threading +import time +from collections import deque +from dataclasses import dataclass +from typing import Literal + +from huggingface_hub import HfApi +from huggingface_hub.utils import build_hf_headers, get_session + + +TokenValidationStatus = Literal["valid", "invalid", "rate_limited", "unavailable"] + + +@dataclass(frozen = True) +class TokenValidationResult: + status: TokenValidationStatus + retry_after_seconds: int | None = None + + +_WINDOW_SECONDS = 3600.0 +_MAX_ATTEMPTS = 3 +_CACHE_TTL_SECONDS = 3600.0 +_TEMPORARY_CACHE_TTL_SECONDS = 15.0 +_MAX_BUCKETS = 4096 +_MAX_CACHE_ENTRIES = 4096 +_INFLIGHT_WAIT_SECONDS = 30.0 +_REMOTE_TIMEOUT_SECONDS = 10.0 + +_attempts: dict[str, deque[float]] = {} +_cache: dict[str, tuple[float, TokenValidationResult]] = {} +_inflight: dict[str, threading.Event] = {} +_lock = threading.Lock() + + +def _fingerprint(token: str) -> str: + return hashlib.sha256(token.encode("utf-8")).hexdigest() + + +def _prune_attempts(bucket: deque[float], now: float) -> None: + while bucket and now - bucket[0] >= _WINDOW_SECONDS: + bucket.popleft() + + +def _prune_locked(now: float) -> None: + for key in list(_attempts): + bucket = _attempts[key] + _prune_attempts(bucket, now) + if not bucket: + del _attempts[key] + for key, (expires_at, _result) in list(_cache.items()): + if expires_at <= now: + del _cache[key] + + +def _cached_locked(fingerprint: str, now: float) -> TokenValidationResult | None: + cached = _cache.get(fingerprint) + if cached is None: + return None + expires_at, result = cached + if expires_at <= now: + del _cache[fingerprint] + return None + return result + + +def _retry_after(bucket: deque[float], now: float) -> int: + return max(1, int(_WINDOW_SECONDS - (now - bucket[0])) + 1) + + +def _reserve_attempt_locked(rate_key: str, now: float) -> TokenValidationResult | None: + bucket = _attempts.get(rate_key) + if bucket is None: + if len(_attempts) >= _MAX_BUCKETS: + _prune_locked(now) + if len(_attempts) >= _MAX_BUCKETS: + return TokenValidationResult( + status = "rate_limited", + retry_after_seconds = max(1, int(_WINDOW_SECONDS)), + ) + bucket = _attempts[rate_key] = deque() + _prune_attempts(bucket, now) + if len(bucket) >= _MAX_ATTEMPTS: + return TokenValidationResult( + status = "rate_limited", + retry_after_seconds = _retry_after(bucket, now), + ) + bucket.append(now) + return None + + +def _http_status(response: object | None) -> int | None: + status = getattr(response, "status_code", None) + try: + return int(status) if status is not None else None + except (TypeError, ValueError): + return None + + +def _remote_retry_after(response: object | None) -> int | None: + headers = getattr(response, "headers", None) + if not headers: + return None + raw = headers.get("Retry-After") + try: + return max(1, int(float(raw))) if raw is not None else None + except (TypeError, ValueError): + return None + + +def _classify_response(response: object | None) -> TokenValidationResult: + status = _http_status(response) + if status is not None and 200 <= status < 300: + return TokenValidationResult(status = "valid") + if status == 401: + return TokenValidationResult(status = "invalid") + if status == 429: + return TokenValidationResult( + status = "rate_limited", + retry_after_seconds = _remote_retry_after(response), + ) + return TokenValidationResult(status = "unavailable") + + +def _check_remote(token: str) -> TokenValidationResult: + api = HfApi() + try: + # HfApi.whoami has no timeout parameter in the pinned Hub client. + # Use its session and headers against the same whoami endpoint. + response = get_session().get( + f"{api.endpoint}/api/whoami-v2", + headers = build_hf_headers(token = token), + timeout = _REMOTE_TIMEOUT_SECONDS, + ) + except Exception as exc: + # huggingface-hub 0.36.x can wrap a 401 as requests.HTTPError. + return _classify_response(getattr(exc, "response", None)) + return _classify_response(response) + + +def validate_hf_token(token: str, *, rate_key: str) -> TokenValidationResult: + """Validate ``token`` without retaining it, sharing results across callers. + + Cached checks do not consume the caller's three-per-hour network budget. A + single-flight event also prevents simultaneously mounted UI surfaces from + sending duplicate ``whoami`` requests for the same token. + """ + normalized = token.strip() + if not normalized: + return TokenValidationResult(status = "invalid") + token_fingerprint = _fingerprint(normalized) + owner_event: threading.Event | None = None + + try: + while True: + now = time.monotonic() + with _lock: + cached = _cached_locked(token_fingerprint, now) + if cached is not None: + return cached + waiting = _inflight.get(token_fingerprint) + if waiting is None: + limited = _reserve_attempt_locked(rate_key, now) + if limited is not None: + return limited + owner_event = threading.Event() + _inflight[token_fingerprint] = owner_event + break + if not waiting.wait(_INFLIGHT_WAIT_SECONDS): + return TokenValidationResult(status = "unavailable") + + result = _check_remote(normalized) + now = time.monotonic() + ttl = ( + _CACHE_TTL_SECONDS + if result.status in ("valid", "invalid") + else max(_TEMPORARY_CACHE_TTL_SECONDS, float(result.retry_after_seconds or 0)) + ) + with _lock: + if len(_cache) >= _MAX_CACHE_ENTRIES: + _prune_locked(now) + if len(_cache) < _MAX_CACHE_ENTRIES: + _cache[token_fingerprint] = (now + ttl, result) + return result + finally: + if owner_event is not None: + with _lock: + event = _inflight.get(token_fingerprint) + if event is owner_event: + _inflight.pop(token_fingerprint, None) + event.set() + + +def reset_hf_token_validation_state() -> None: + """Clear process state for test isolation.""" + with _lock: + for event in _inflight.values(): + event.set() + _inflight.clear() + _attempts.clear() + _cache.clear() diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index 3818253ac9..31f5f31bee 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -123,6 +123,26 @@ def without_hf_auth(): os.environ.pop("HF_HUB_DISABLE_IMPLICIT_TOKEN", None) +def is_hf_authentication_error(error: Exception) -> bool: + """Return whether an exception chain contains a definitive HF auth failure.""" + seen: set[int] = set() + current: BaseException | None = error + while current is not None and id(current) not in seen: + seen.add(id(current)) + response = getattr(current, "response", None) + status = getattr(response, "status_code", None) + try: + if status is not None and int(status) == 401: + return True + except (TypeError, ValueError): + pass + message = str(current).lower() + if "invalid user token" in message or "invalid hf token" in message: + return True + current = current.__cause__ or current.__context__ + return False + + def format_error_message(error: Exception, model_name: str) -> str: """ Format a user-friendly error message for common load issues. diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index ba56ce7525..e23892e020 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -16,6 +16,7 @@ import { type ChatSearch, } from "@/features/chat"; import { RemoteCodeConsentDialog } from "@/features/security"; +import { HfTokenWarningDialog } from "@/features/hf-auth"; import { TransformersUpgradeDialog } from "@/features/transformers-upgrade"; import { useTrainingUnloadGuard } from "@/features/training"; import { useExportRuntimeLifecycle } from "@/features/export"; @@ -230,6 +231,7 @@ function RootLayout() { {!isAuthFlowRoute && } + {hideNavbar ? ( diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index d9d0493a03..475bd0f801 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -2,6 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { authFetch } from "@/features/auth"; +import { prepareHfTokenForUse } from "@/features/hf-auth"; // These helpers are deliberately API-layer-only and are not part of their // features' React-facing public barrels. // eslint-disable-next-line no-restricted-imports @@ -108,11 +109,14 @@ export async function getApiMonitorEntry(id: string): Promise { export async function loadModel( payload: LoadModelRequest, ): Promise { + const preparedToken = await prepareHfTokenForUse(payload.hf_token); + if (!preparedToken.proceed) throw new Error("Model load cancelled."); const response = await authFetch("/api/inference/load", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ ...payload, + hf_token: preparedToken.token, native_path_lease: payload.nativePathLease ?? null, nativePathLease: undefined, }), @@ -123,13 +127,15 @@ export async function loadModel( export async function validateModel( payload: LoadModelRequest, ): Promise { + const preparedToken = await prepareHfTokenForUse(payload.hf_token); + if (!preparedToken.proceed) throw new Error("Model load cancelled."); const response = await authFetch("/api/inference/validate", { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ model_path: payload.model_path, native_path_lease: payload.nativePathLease ?? null, - hf_token: payload.hf_token, + hf_token: preparedToken.token, gguf_variant: payload.gguf_variant ?? null, // Intended load settings so validate's preflight matches the follow-up // /load. Default placement is sized against the selected GPUs. diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index 4cd7958fb8..51e9f209dd 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -44,6 +44,7 @@ import { import { usePlatformStore } from "@/config/env"; import { useHubModelSearch } from "@/features/hub/hooks/use-hub-model-search"; import { confirmRemoteCodeIfNeeded } from "@/features/security"; +import { prepareHfTokenForUse } from "@/features/hf-auth"; import { GuidedTour, useGuidedTourController } from "@/features/tour"; import { type LocalModelInfo, @@ -724,11 +725,17 @@ export function ExportPage() { const checkpointPath = selectedCp?.path ?? null; const pushToHub = destination === "hub"; + const preparedToken = await prepareHfTokenForUse(hfToken, { + allowAnonymous: !pushToHub, + }); + if (!preparedToken.proceed) return; + const actionHfToken = preparedToken.token ?? ""; + const repoId = pushToHub && hfUsername && modelName ? `${hfUsername}/${modelName}` : undefined; - const token = pushToHub && hfToken ? hfToken : undefined; + const token = pushToHub && actionHfToken ? actionHfToken : undefined; // The GGUF method with the LoRA target reuses the LoRA-adapter export path. const effectiveMethod: ExportMethod = ggufAsLora ? "lora" : exportMethod; const emitLoraGguf = @@ -747,7 +754,7 @@ export function ExportPage() { if (sourceMode !== "checkpoint") { const remoteCodeOk = await confirmRemoteCodeIfNeeded({ modelName: source, - hfToken: hfToken || null, + hfToken: actionHfToken || null, // An HF source can need trust_remote_code via its YAML default with no // auto_map to review; signal it so a YAML-only model does not export // with it false. @@ -767,7 +774,7 @@ export function ExportPage() { modelSource, trustRemoteCode, approvedRemoteCodeFingerprint, - loadToken: hfToken || null, + loadToken: actionHfToken || null, exportMethod: effectiveMethod, isAdapter: adapterExport, quantLevels, diff --git a/studio/frontend/src/features/hf-auth/api.ts b/studio/frontend/src/features/hf-auth/api.ts new file mode 100644 index 0000000000..ab4b049566 --- /dev/null +++ b/studio/frontend/src/features/hf-auth/api.ts @@ -0,0 +1,44 @@ +// 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"; +// This header helper is API-layer-only and is not part of the feature's +// React-facing public barrel. +// eslint-disable-next-line no-restricted-imports +import { hubTokenHeader } from "@/features/hub/lib/hub-token-header"; + +export type HfTokenValidationStatus = + | "missing" + | "valid" + | "invalid" + | "rate_limited" + | "unavailable"; + +export interface HfTokenValidationResult { + status: HfTokenValidationStatus; + retryAfterSeconds: number | null; +} + +export async function validateHfToken( + token: string | null | undefined, +): Promise { + const normalized = token?.trim() ?? ""; + if (!normalized) { + return { status: "missing", retryAfterSeconds: null }; + } + const response = await authFetch("/api/hub/token/validate", { + method: "POST", + headers: hubTokenHeader(normalized), + }); + if (!response.ok) { + return { status: "unavailable", retryAfterSeconds: null }; + } + const body = (await response.json()) as { + status?: HfTokenValidationStatus; + retry_after_seconds?: number | null; + }; + return { + status: body.status ?? "unavailable", + retryAfterSeconds: body.retry_after_seconds ?? null, + }; +} diff --git a/studio/frontend/src/features/hf-auth/confirm-token.ts b/studio/frontend/src/features/hf-auth/confirm-token.ts new file mode 100644 index 0000000000..e1f6e6c705 --- /dev/null +++ b/studio/frontend/src/features/hf-auth/confirm-token.ts @@ -0,0 +1,63 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +// These stores are used outside React and are not part of their features' +// React-facing public barrels. +// eslint-disable-next-line no-restricted-imports +import { useHfTokenStore } from "@/features/hub/stores/hf-token-store"; +// eslint-disable-next-line no-restricted-imports +import { useSettingsDialogStore } from "@/features/settings/stores/settings-dialog-store"; +import { validateHfToken } from "./api"; +import { useHfTokenWarningStore } from "./store"; + +export interface PreparedHfToken { + proceed: boolean; + token: string | null; +} + +interface PrepareHfTokenOptions { + allowAnonymous?: boolean; +} + +// A caller can retain the pre-dialog payload while the shared store is cleared. +// Remember that one-session choice so a follow-up /load does not prompt again +// after its preceding /validate already continued anonymously. +const anonymousForSession = new Set(); + +export async function prepareHfTokenForUse( + token: string | null | undefined, + options: PrepareHfTokenOptions = {}, +): Promise { + const normalized = token?.trim() ?? ""; + if (!normalized) return { proceed: true, token: null }; + const allowAnonymous = options.allowAnonymous ?? true; + if (allowAnonymous && anonymousForSession.has(normalized)) { + return { proceed: true, token: null }; + } + + let validation; + try { + validation = await validateHfToken(normalized); + } catch { + // Validation is advisory. Let the real operation retain its own error. + return { proceed: true, token: normalized }; + } + if (validation.status !== "invalid") { + // A connectivity failure or rate limit cannot prove that a token is bad. + // Let the real operation proceed and retain its repository-specific error. + return { proceed: true, token: normalized }; + } + + const decision = await useHfTokenWarningStore + .getState() + .requestDecision(allowAnonymous); + if (decision === "anonymous") { + anonymousForSession.add(normalized); + useHfTokenStore.getState().clearToken(); + return { proceed: true, token: null }; + } + if (decision === "replace") { + useSettingsDialogStore.getState().openDialog("general"); + } + return { proceed: false, token: normalized }; +} diff --git a/studio/frontend/src/features/hf-auth/hf-token-warning-dialog.tsx b/studio/frontend/src/features/hf-auth/hf-token-warning-dialog.tsx new file mode 100644 index 0000000000..4370c39693 --- /dev/null +++ b/studio/frontend/src/features/hf-auth/hf-token-warning-dialog.tsx @@ -0,0 +1,66 @@ +// 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 { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alert-dialog"; +import { Button } from "@/components/ui/button"; +import { AlertTriangle } from "lucide-react"; +import { useHfTokenWarningStore } from "./store"; + +export function HfTokenWarningDialog() { + const open = useHfTokenWarningStore((state) => state.open); + const allowAnonymous = useHfTokenWarningStore( + (state) => state.allowAnonymous, + ); + const resolve = useHfTokenWarningStore((state) => state.resolve); + + return ( + { + if (!next) resolve("cancel"); + }} + > + + +
+
+ +
+
+ Hugging Face token is invalid + + {allowAnonymous + ? "Hugging Face rejected the saved token. Replace it to access private or gated repositories, or continue without it for public and fully downloaded models." + : "Hugging Face rejected the saved token. Replace it before uploading to the Hub."} + +
+
+ + + resolve("cancel")}> + Cancel + +
+ {allowAnonymous ? ( + + ) : null} + resolve("replace")}> + Replace token + +
+ + + + ); +} diff --git a/studio/frontend/src/features/hf-auth/index.ts b/studio/frontend/src/features/hf-auth/index.ts new file mode 100644 index 0000000000..e8bcb48193 --- /dev/null +++ b/studio/frontend/src/features/hf-auth/index.ts @@ -0,0 +1,10 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +export { validateHfToken } from "./api"; +export type { + HfTokenValidationResult, + HfTokenValidationStatus, +} from "./api"; +export { prepareHfTokenForUse } from "./confirm-token"; +export { HfTokenWarningDialog } from "./hf-token-warning-dialog"; diff --git a/studio/frontend/src/features/hf-auth/store.ts b/studio/frontend/src/features/hf-auth/store.ts new file mode 100644 index 0000000000..faa2543a8a --- /dev/null +++ b/studio/frontend/src/features/hf-auth/store.ts @@ -0,0 +1,33 @@ +// 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 { create } from "zustand"; + +export type HfTokenWarningDecision = "anonymous" | "replace" | "cancel"; +type Resolver = (decision: HfTokenWarningDecision) => void; + +let pendingResolver: Resolver | null = null; + +interface HfTokenWarningStore { + open: boolean; + allowAnonymous: boolean; + requestDecision: (allowAnonymous: boolean) => Promise; + resolve: (decision: HfTokenWarningDecision) => void; +} + +export const useHfTokenWarningStore = create((set) => ({ + open: false, + allowAnonymous: true, + requestDecision: (allowAnonymous) => + new Promise((resolve) => { + pendingResolver?.("cancel"); + pendingResolver = resolve; + set({ open: true, allowAnonymous }); + }), + resolve: (decision) => { + const resolver = pendingResolver; + pendingResolver = null; + set({ open: false, allowAnonymous: true }); + resolver?.(decision); + }, +})); diff --git a/studio/frontend/src/features/settings/tabs/general-tab.tsx b/studio/frontend/src/features/settings/tabs/general-tab.tsx index 73a33fd155..d5a85c01c6 100644 --- a/studio/frontend/src/features/settings/tabs/general-tab.tsx +++ b/studio/frontend/src/features/settings/tabs/general-tab.tsx @@ -21,6 +21,7 @@ import { setShowLlamaUpdateBanner, useShowLlamaUpdateBanner, } from "@/hooks/use-llama-update-pref"; +import { useHfTokenValidation } from "@/hooks"; import { LOCALE_STORAGE_KEY, useT } from "@/i18n"; import { isTauri } from "@/lib/api-base"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; @@ -216,11 +217,18 @@ export function GeneralTab() { if (trimmed !== hfToken) setHfToken(trimmed); }; + const clearHfToken = () => { + draftRef.current = ""; + setDraftToken(""); + setHfToken(""); + }; + // Show an "accepted" tick once a non-empty token has been committed to the // store and the field still matches it (i.e. not mid-edit). Gives the user // feedback that a pasted token was saved. const tokenSaved = draftToken.trim().length > 0 && draftToken.trim() === (hfToken ?? ""); + const tokenValidation = useHfTokenValidation(hfToken ?? ""); useEffect(() => { let cancelled = false; @@ -498,46 +506,70 @@ export function GeneralTab() { label={t("settings.general.huggingFaceToken")} description={t("settings.general.huggingFaceTokenDescription")} > -
- setDraftToken(e.target.value)} - onBlur={commitToken} - className={cn( - "h-8 w-full font-mono text-xs", - tokenSaved ? "pr-14" : "pr-8", - )} - /> - {tokenSaved ? ( - // Decorative: pointer-events-none lets clicks reach the input - // underneath so the field still focuses anywhere. - +
+
+ setDraftToken(e.target.value)} + onBlur={commitToken} + className={cn( + "h-8 w-full font-mono text-xs", + tokenSaved ? "pr-14" : "pr-8", + )} + /> + {tokenSaved ? ( + // Decorative: pointer-events-none lets clicks reach the input + // underneath so the field still focuses anywhere. + + + + ) : null} + +
+ +
+ {tokenValidation.isChecking ? ( +

Checking token…

+ ) : tokenValidation.error ? ( +

+ {tokenValidation.error} +

) : null} -
{/* The desktop app authenticates via desktop auto-auth with a generated diff --git a/studio/frontend/src/features/training/api/train-api.ts b/studio/frontend/src/features/training/api/train-api.ts index 609e4f8591..c5f09c8bd0 100644 --- a/studio/frontend/src/features/training/api/train-api.ts +++ b/studio/frontend/src/features/training/api/train-api.ts @@ -2,6 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { authFetch } from "@/features/auth"; +import { prepareHfTokenForUse } from "@/features/hf-auth"; import { readFastApiError } from "@/lib/format-fastapi-error"; import type { TrainingStartRequest, @@ -30,10 +31,12 @@ async function parseJson(response: Response): Promise { export async function startTraining( payload: TrainingStartRequest, ): Promise { + const preparedToken = await prepareHfTokenForUse(payload.hf_token); + if (!preparedToken.proceed) throw new Error("Training start cancelled."); const response = await authFetch("/api/train/start", { method: "POST", headers: { "Content-Type": "application/json" }, - body: JSON.stringify(payload), + body: JSON.stringify({ ...payload, hf_token: preparedToken.token }), }); return parseJson(response); } diff --git a/studio/frontend/src/features/training/hooks/use-training-actions.ts b/studio/frontend/src/features/training/hooks/use-training-actions.ts index 2f04656c23..0d32e066d5 100644 --- a/studio/frontend/src/features/training/hooks/use-training-actions.ts +++ b/studio/frontend/src/features/training/hooks/use-training-actions.ts @@ -2,6 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { primeNativeNotificationPermission } from "@/lib/native-notifications"; +import { prepareHfTokenForUse } from "@/features/hf-auth"; import { confirmRemoteCodeIfNeeded } from "@/features/security"; import { useCallback } from "react"; import { toast } from "@/lib/toast"; @@ -43,7 +44,7 @@ export function useTrainingActions() { const startError = useTrainingRuntimeStore((state) => state.startError); const startTrainingRun = useCallback(async (): Promise => { - const config = useTrainingConfigStore.getState(); + let config = useTrainingConfigStore.getState(); const runtimeStore = useTrainingRuntimeStore.getState(); const dialogStore = useDatasetPreviewDialogStore.getState(); @@ -54,6 +55,13 @@ export function useTrainingActions() { return false; } + const preparedToken = await prepareHfTokenForUse(config.hfToken); + if (!preparedToken.proceed) return false; + if ((preparedToken.token ?? "") !== config.hfToken) { + config.setHfToken(preparedToken.token ?? ""); + config = useTrainingConfigStore.getState(); + } + primeNativeNotificationPermission().catch(() => undefined); runtimeStore.setStartResources( @@ -226,6 +234,13 @@ export function useTrainingActions() { resume_from_checkpoint: outputDir, } as TrainingStartRequest; + const preparedToken = await prepareHfTokenForUse(payload.hf_token); + if (!preparedToken.proceed) { + runtimeStore.setStarting(false); + return false; + } + payload.hf_token = preparedToken.token; + runtimeStore.setStartResources( payload.model_name, payload.hf_dataset, diff --git a/studio/frontend/src/hooks/use-hf-token-validation.ts b/studio/frontend/src/hooks/use-hf-token-validation.ts index 11f56a2151..e0ed02cca2 100644 --- a/studio/frontend/src/hooks/use-hf-token-validation.ts +++ b/studio/frontend/src/hooks/use-hf-token-validation.ts @@ -1,8 +1,8 @@ // 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 { whoAmI } from "@huggingface/hub"; -import { useCallback, useEffect, useRef, useState } from "react"; +import { validateHfToken } from "@/features/hf-auth"; +import { useEffect, useRef, useState } from "react"; import { useDebouncedValue } from "./use-debounced-value"; export interface HfTokenValidationState { @@ -17,6 +17,20 @@ const INITIAL: HfTokenValidationState = { isChecking: false, }; +interface CompletedValidation extends HfTokenValidationState { + token: string; +} + +const NO_COMPLETED_VALIDATION: CompletedValidation = { + ...INITIAL, + token: "", +}; + +// Current user access tokens contain 34 characters after the hf_ prefix. +// Action-time validation still accepts legacy shapes without spending quota +// on every intermediate value typed into a live form field. +const COMPLETE_HF_TOKEN = /^hf_[A-Za-z0-9]{34}$/; + /** * Validates the HF token via the whoami-v2 API, debounced to avoid excessive * requests while typing. isValid is null until checked. @@ -26,39 +40,73 @@ export function useHfTokenValidation(token: string): HfTokenValidationState { token.trim().replace(/^["']+|["']+$/g, ""), 500, ); - const [state, setState] = useState(INITIAL); + const [completed, setCompleted] = useState( + NO_COMPLETED_VALIDATION, + ); const versionRef = useRef(0); - - const runCheck = useCallback(async (t: string) => { - if (!t) { - setState({ isValid: null, error: null, isChecking: false }); - return; - } - - const v = ++versionRef.current; - setState((prev) => ({ ...prev, isChecking: true, error: null })); - - try { - await whoAmI({ accessToken: t }); - if (versionRef.current !== v) return; - setState({ isValid: true, error: null, isChecking: false }); - } catch { - if (versionRef.current !== v) return; - setState({ - isValid: false, - error: "invalid or expired token", - isChecking: false, - }); - } - }, []); + const shouldValidate = COMPLETE_HF_TOKEN.test(debouncedToken); useEffect(() => { - if (!debouncedToken) { - setState(INITIAL); + if (!shouldValidate) { + versionRef.current += 1; return; } - runCheck(debouncedToken); - }, [debouncedToken, runCheck]); + const version = ++versionRef.current; + void validateHfToken(debouncedToken).then( + (result) => { + if (versionRef.current !== version) return; + if (result.status === "valid") { + setCompleted({ + token: debouncedToken, + isValid: true, + error: null, + isChecking: false, + }); + } else if (result.status === "invalid") { + setCompleted({ + token: debouncedToken, + isValid: false, + error: "invalid or expired token", + isChecking: false, + }); + } else if (result.status === "rate_limited") { + const wait = result.retryAfterSeconds + ? ` Try again in about ${Math.ceil(result.retryAfterSeconds / 60)} minute(s).` + : " Try again later."; + setCompleted({ + token: debouncedToken, + isValid: null, + error: `Token verification is rate limited.${wait}`, + isChecking: false, + }); + } else { + setCompleted({ + token: debouncedToken, + isValid: null, + error: "Could not verify the token. Check your connection and try again.", + isChecking: false, + }); + } + }, + () => { + if (versionRef.current !== version) return; + setCompleted({ + token: debouncedToken, + isValid: null, + error: "Could not verify the token. Check your connection and try again.", + isChecking: false, + }); + }, + ); + }, [debouncedToken, shouldValidate]); - return state; + if (!shouldValidate) return INITIAL; + if (completed.token !== debouncedToken) { + return { isValid: null, error: null, isChecking: true }; + } + return { + isValid: completed.isValid, + error: completed.error, + isChecking: false, + }; }