diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 3ae974448e..f290028a37 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -19,6 +19,10 @@ from pydantic import ( ) +# Redefined from picker/schemas.py so this core schema needn't import the picker package. +MAX_CHAT_TEMPLATE_BYTES = 65_536 + + class LoadRequest(BaseModel): """Request to load a model for inference""" @@ -54,8 +58,16 @@ class LoadRequest(BaseModel): @field_validator("chat_template_override") @classmethod def normalize_blank_chat_template_override(cls, value: Optional[str]) -> Optional[str]: - if value is not None and value.strip() == "": + if value is None: return None + # Reject on character count first (a lower bound on the UTF-8 byte length) + # so an oversized template is rejected before allocating the encoded bytes. + if len(value) > MAX_CHAT_TEMPLATE_BYTES: + raise ValueError(f"Chat template exceeds the {MAX_CHAT_TEMPLATE_BYTES}-byte limit.") + if value.strip() == "": + return None + if len(value.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES: + raise ValueError(f"Chat template exceeds the {MAX_CHAT_TEMPLATE_BYTES}-byte limit.") return value cache_type_kv: Optional[str] = Field( diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index b8526c75e7..a6320fb3fa 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -1797,10 +1797,14 @@ def _get_model_size_bytes(model_name: str, hf_token: Optional[str] = None) -> Op async def get_model_config( model_name: str, hf_token: Optional[str] = Query(None), + hf_token_header: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ): """Get configuration for a specific model (wraps load_model_defaults).""" try: + # Prefer the header token; keep the query param as a fallback for older + # clients that predate the header (as /gguf-variants does). + hf_token = _normalize_hf_token(hf_token_header) or _normalize_hf_token(hf_token) if not is_local_path(model_name): resolved = resolve_cached_repo_id_case(model_name) if resolved != model_name: @@ -2518,6 +2522,7 @@ async def get_lora_base_model(lora_path: str, current_subject: str = Depends(get async def check_vision_model( model_name: str, hf_token: Optional[str] = Query(None), + hf_token_header: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ): """ @@ -2526,6 +2531,7 @@ async def check_vision_model( This endpoint wraps the backend is_vision_model function. """ try: + hf_token = _normalize_hf_token(hf_token_header) or _normalize_hf_token(hf_token) logger.info(f"Checking if vision model: {model_name}") # Authenticate so a gated/private VLM classifies correctly (else 404 -> non-vision). is_vision = is_vision_model(model_name, hf_token = hf_token) @@ -2550,6 +2556,7 @@ async def check_vision_model( async def check_embedding_model( model_name: str, hf_token: Optional[str] = Query(None), + hf_token_header: Optional[str] = Depends(get_hf_token), current_subject: str = Depends(get_current_subject), ): """ @@ -2558,6 +2565,7 @@ async def check_embedding_model( This endpoint wraps the backend is_embedding_model function. """ try: + hf_token = _normalize_hf_token(hf_token_header) or _normalize_hf_token(hf_token) logger.info(f"Checking if embedding model: {model_name}") is_embedding = is_embedding_model(model_name, hf_token = hf_token) diff --git a/studio/frontend/src/features/model-picker/model-config/per-model-config.ts b/studio/frontend/src/features/model-picker/model-config/per-model-config.ts index 53a376e8b4..f1b4140718 100644 --- a/studio/frontend/src/features/model-picker/model-config/per-model-config.ts +++ b/studio/frontend/src/features/model-picker/model-config/per-model-config.ts @@ -139,8 +139,7 @@ function deleteOldestEvictableEntry( protectedKeys?: ReadonlySet, ): { key: string; value: StoredMap[string] } | null { for (const key of Object.keys(map)) { - // Never evict a future-schema entry an older client cannot interpret, - // matching the save/delete guards. + // Never evict a future-schema entry an older client can't interpret. if ( protectedKeys?.has(key) || storedConfigVersion(map[key]) > STORAGE_SCHEMA_VERSION @@ -263,6 +262,10 @@ function migrateLegacyLoadSettingsOnce(): void { return; } const map = readMapRaw(); + // Snapshot the user's existing entries before merging: legacy settings are + // imported only into spare budget and must never evict a config the user + // already saved in the current schema. + const existingKeys = new Set(Object.keys(map)); const migratedKeys = mergeLegacyEntries( map, legacy as Record, @@ -271,11 +274,11 @@ function migrateLegacyLoadSettingsOnce(): void { localStorage.setItem(LEGACY_MIGRATION_FLAG, "1"); return; } - // Protect the just-migrated entries during eviction. If the budget cannot - // fit them (e.g. storage is full of future-schema records an older client - // cannot evict), leave the flag unset so migration retries once space frees - // up rather than marking it complete and dropping the migrated config. - if (!enforceStorageBudget(map, new Set(migratedKeys))) { + // Protect existing entries so only the just-migrated legacy entries are + // evicted when over budget: importing old settings never discards a newer + // config, and since existing entries already fit the budget this cannot + // deadlock. On failure the flag stays unset so migration retries. + if (!enforceStorageBudget(map, existingKeys)) { return; } if (writeMap(map)) { @@ -489,8 +492,7 @@ function loadPerModelConfig( if (!key) { return null; } - // Never apply a future-schema record an older client cannot interpret, - // matching the save/delete/evict guards. + // Never apply a future-schema record an older client can't interpret. if (storedConfigVersion(map[key]) > STORAGE_SCHEMA_VERSION) { return null; } @@ -548,8 +550,7 @@ export function deletePerModelConfig( ggufVariant?: string | null, ): boolean { const map = readMap(); - // Mirror savePerModelConfig: never let an older client destroy a - // future-schema entry it cannot interpret. + // Never let an older client destroy a future-schema entry it can't interpret. if (hasFutureConfigForModelVariant(map, modelId, ggufVariant)) { return false; } diff --git a/studio/frontend/src/features/training/api/models-api.ts b/studio/frontend/src/features/training/api/models-api.ts index e5fd8b043d..2ccc49f9e5 100644 --- a/studio/frontend/src/features/training/api/models-api.ts +++ b/studio/frontend/src/features/training/api/models-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 { hubTokenHeader } from "@/features/hub"; interface VisionCheckResponse { model_name: string; @@ -98,10 +99,10 @@ export async function checkVisionModel( hfToken?: string | null, ): Promise { const encoded = encodeURIComponent(modelName); - const query = hfToken?.trim() ? `?hf_token=${encodeURIComponent(hfToken.trim())}` : ""; - const response = await authFetch(`/api/models/check-vision/${encoded}${query}`); + const response = await authFetch(`/api/models/check-vision/${encoded}`, { + headers: hubTokenHeader(hfToken?.trim() || null), + }); if (!response.ok) { - // If the check fails (e.g. network error), default to non-vision return false; } const data = (await response.json()) as VisionCheckResponse; @@ -114,10 +115,10 @@ export async function checkEmbeddingModel( hfToken?: string | null, ): Promise { const encoded = encodeURIComponent(modelName); - const query = hfToken?.trim() ? `?hf_token=${encodeURIComponent(hfToken.trim())}` : ""; - const response = await authFetch(`/api/models/check-embedding/${encoded}${query}`); + const response = await authFetch(`/api/models/check-embedding/${encoded}`, { + headers: hubTokenHeader(hfToken?.trim() || null), + }); if (!response.ok) { - // If the check fails (e.g. network error), default to non-embedding return false; } const data = (await response.json()) as EmbeddingCheckResponse; @@ -130,8 +131,10 @@ export async function getModelConfig( hfToken?: string, ): Promise { const encoded = encodeURIComponent(modelName); - const params = hfToken ? `?hf_token=${encodeURIComponent(hfToken)}` : ""; - const response = await authFetch(`/api/models/config/${encoded}${params}`, { signal }); + const response = await authFetch(`/api/models/config/${encoded}`, { + headers: hubTokenHeader(hfToken), + signal, + }); if (!response.ok) { throw new Error(`Failed to fetch model config (${response.status})`); }