diff --git a/studio/backend/hub/utils/gguf.py b/studio/backend/hub/utils/gguf.py index 4eb8180628..acd650bf42 100644 --- a/studio/backend/hub/utils/gguf.py +++ b/studio/backend/hub/utils/gguf.py @@ -353,6 +353,20 @@ def list_partial_gguf_variants_from_state( return variants, has_vision +def resolve_local_gguf_path(repo_id: str, gguf_variant: Optional[str]) -> Optional[str]: + """Absolute path to the (shard-1) GGUF file for ``repo_id`` + ``gguf_variant`` + if it is already downloaded in the HF cache, else ``None``. Read-only — never + triggers a download. Lets callers read header metadata before a load.""" + for snapshot in iter_hf_cache_snapshots(repo_id): + variants, _ = list_local_gguf_variants(str(snapshot)) + for variant in variants: + if gguf_variant is None or variant.quant == gguf_variant: + candidate = snapshot / variant.filename + if candidate.is_file(): + return str(candidate) + return None + + def list_gguf_variants( repo_id: str, hf_token: Optional[str] = None ) -> tuple[list[GgufVariantInfo], bool, Optional[list]]: diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 82949870da..fdc3bb25c6 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -125,6 +125,11 @@ class ValidateModelRequest(BaseModel): gguf_variant: Optional[str] = Field( None, description = "GGUF quantization variant (e.g. 'Q4_K_M')" ) + include_context_length: bool = Field( + False, + description = "Also read the native context length from the local GGUF header. " + "Opt-in so the normal load preflight doesn't pay for a cache scan it doesn't need.", + ) class ValidateModelResponse(BaseModel): @@ -144,6 +149,11 @@ class ValidateModelResponse(BaseModel): False, description = "Whether the model defaults require trust_remote_code to be enabled for loading.", ) + context_length: Optional[int] = Field( + None, + description = "Native training context length, read from the GGUF header when the file " + "is already downloaded locally; None for non-GGUF, gated, or not-yet-downloaded models.", + ) class GenerateRequest(BaseModel): diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 48b13d31c9..71cee33b24 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -2536,6 +2536,33 @@ async def validate_model( detail = f"Invalid model identifier: {model_log_label}", ) + is_gguf = getattr(config, "is_gguf", False) + # Native context length, read from the local GGUF header when present. + # Lets the staged ("Load on selection" off) flow populate the context + # slider before the GPU load; None until the file is downloaded. + context_length: Optional[int] = None + if request.include_context_length and is_gguf: + from hub.utils.gguf import resolve_local_gguf_path + from utils.models.gguf_metadata import read_gguf_context_length + + # Best-effort: a header-read failure must never fail validation of an + # otherwise-valid model (the outer except turns it into a 400). + try: + if native_grant_backed: + # model_identifier is the resolved canonical .gguf path. + local_gguf = model_identifier + else: + # Local folder / exported GGUFs already have their file + # resolved on the config (gguf_file is None for HF repos, so + # those fall back to the HF-cache lookup). + local_gguf = config.gguf_file or resolve_local_gguf_path( + model_identifier, request.gguf_variant + ) + if local_gguf: + context_length = read_gguf_context_length(local_gguf) + except Exception as e: + logger.debug("Context-length probe failed for %s: %s", model_log_label, e) + return ValidateModelResponse( valid = True, message = "Model identifier is valid.", @@ -2543,12 +2570,13 @@ async def validate_model( display_name = model_log_label if native_grant_backed else getattr(config, "display_name", config.identifier), - is_gguf = getattr(config, "is_gguf", False), + is_gguf = is_gguf, is_lora = getattr(config, "is_lora", False), is_vision = getattr(config, "is_vision", False), requires_trust_remote_code = bool( load_inference_config(config.identifier).get("trust_remote_code", False) ), + context_length = context_length, ) except HTTPException: diff --git a/studio/backend/tests/test_gguf_metadata.py b/studio/backend/tests/test_gguf_metadata.py index d3d4387720..a5be07f8e3 100644 --- a/studio/backend/tests/test_gguf_metadata.py +++ b/studio/backend/tests/test_gguf_metadata.py @@ -13,6 +13,7 @@ from typing import Iterable, Mapping from utils.models.gguf_metadata import ( is_mmproj_by_metadata, pairing_score, + read_gguf_context_length, read_gguf_general_metadata, read_mmproj_audio_capability, ) @@ -21,6 +22,7 @@ from utils.models.gguf_metadata import ( _GGUF_MAGIC = 0x46554747 _VTYPE_STRING = 8 _VTYPE_UINT32 = 4 +_VTYPE_UINT64 = 10 _VTYPE_ARRAY = 9 _VTYPE_BOOL = 7 @@ -38,6 +40,10 @@ def _enc_kv_uint32(key: str, value: int) -> bytes: return _enc_string(key) + struct.pack(" bytes: + return _enc_string(key) + struct.pack(" bytes: return _enc_string(key) + struct.pack(" Path: """Minimal GGUF: header + KV body, no tensors.""" extra_uint32 = extra_uint32 or {} + extra_uint64 = extra_uint64 or {} extra_string_arrays = extra_string_arrays or {} extra_bools = extra_bools or {} kv_count = ( - len(general_strings) + len(extra_uint32) + len(extra_string_arrays) + len(extra_bools) + len(general_strings) + + len(extra_uint32) + + len(extra_uint64) + + len(extra_string_arrays) + + len(extra_bools) ) body = b"" for k, v in general_strings.items(): body += _enc_kv_string(k, v) for k, v in extra_uint32.items(): body += _enc_kv_uint32(k, v) + for k, v in extra_uint64.items(): + body += _enc_kv_uint64(k, v) for k, v in extra_string_arrays.items(): body += _enc_kv_string_array(k, v) for k, v in extra_bools.items(): @@ -100,6 +114,66 @@ def test_returns_none_for_non_gguf(tmp_path: Path): assert read_gguf_general_metadata(str(p)) is None +def test_context_length_none_for_missing_file(tmp_path: Path): + assert read_gguf_context_length(str(tmp_path / "nope.gguf")) is None + + +def test_context_length_none_for_non_gguf(tmp_path: Path): + p = tmp_path / "garbage.gguf" + p.write_bytes(b"not a gguf file at all, just bytes") + assert read_gguf_context_length(str(p)) is None + + +def test_context_length_read_from_arch_namespaced_key(tmp_path: Path): + p = _write_synthetic_gguf( + tmp_path / "model.gguf", + {"general.architecture": "llama"}, + extra_uint32 = {"llama.context_length": 4096, "llama.block_count": 32}, + ) + assert read_gguf_context_length(str(p)) == 4096 + + +def test_context_length_none_when_absent(tmp_path: Path): + # Architecture present but no .context_length key. + p = _write_synthetic_gguf( + tmp_path / "model.gguf", + {"general.architecture": "llama"}, + extra_uint32 = {"llama.block_count": 32}, + ) + assert read_gguf_context_length(str(p)) is None + + +def test_context_length_ignores_foreign_arch_key(tmp_path: Path): + # A context_length under a different arch namespace must not match. + p = _write_synthetic_gguf( + tmp_path / "model.gguf", + {"general.architecture": "llama"}, + extra_uint32 = {"qwen2.context_length": 8192}, + ) + assert read_gguf_context_length(str(p)) is None + + +def test_context_length_read_from_uint64(tmp_path: Path): + # Some models store context_length as a uint64 (vtype 10). + p = _write_synthetic_gguf( + tmp_path / "model.gguf", + {"general.architecture": "qwen3"}, + extra_uint64 = {"qwen3.context_length": 262144}, + ) + assert read_gguf_context_length(str(p)) == 262144 + + +def test_context_length_zero_treated_as_absent(tmp_path: Path): + # A zero/garbage ceiling must read as None so the UI can't build a slider + # with max < min. + p = _write_synthetic_gguf( + tmp_path / "model.gguf", + {"general.architecture": "llama"}, + extra_uint32 = {"llama.context_length": 0}, + ) + assert read_gguf_context_length(str(p)) is None + + def test_extracts_general_string_fields(tmp_path: Path): p = _write_synthetic_gguf( tmp_path / "model.gguf", diff --git a/studio/backend/utils/models/gguf_metadata.py b/studio/backend/utils/models/gguf_metadata.py index a2912cc843..c24ec28e1d 100644 --- a/studio/backend/utils/models/gguf_metadata.py +++ b/studio/backend/utils/models/gguf_metadata.py @@ -50,6 +50,10 @@ _CACHE_MAX_ENTRIES = 4096 # keyed by (file cache key, wanted key). None = key absent / file unreadable. _BOOL_CACHE: Dict[Tuple[_CacheKey, str], Optional[bool]] = {} +# Native training context length (``{arch}.context_length``). None = absent / +# unreadable. Lets the UI show the real context ceiling before a model loads. +_CONTEXT_CACHE: Dict[_CacheKey, Optional[int]] = {} + def _cache_key(path: str) -> Optional[_CacheKey]: try: @@ -138,6 +142,92 @@ def _parse_gguf_header(path: str) -> Optional[Dict[str, str]]: return out +def read_gguf_context_length(path: str) -> Optional[int]: + """Return the GGUF's native training context length (``{arch}.context_length``), + or ``None`` if missing/unreadable/not a GGUF. Cached by (path, mtime, size). + Lets the UI populate the context slider before the model is loaded.""" + key = _cache_key(path) + if key is None: + return None + with _CACHE_LOCK: + if key in _CONTEXT_CACHE: + return _CONTEXT_CACHE[key] + result = _parse_gguf_context_length(path) + with _CACHE_LOCK: + while len(_CONTEXT_CACHE) >= _CACHE_MAX_ENTRIES: + try: + _CONTEXT_CACHE.pop(next(iter(_CONTEXT_CACHE))) + except StopIteration: + break + _CONTEXT_CACHE[key] = result + return result + + +def _parse_gguf_context_length(path: str) -> Optional[int]: + # The context key is architecture-namespaced (``llama.context_length`` etc.), + # so we learn the key only after reading ``general.architecture``. GGUF writes + # general.* before arch.* keys, matching the loader's own parser. + ctx_key: Optional[str] = None + try: + with open(path, "rb") as f: + head = f.read(24) + if len(head) < 24: + return None + magic, _version, _tcount, kv_count = struct.unpack(" 1 << 20: # 1 MB sanity bound + break + kbytes = f.read(klen) + if len(kbytes) < klen: + break + key = kbytes.decode("utf-8", "replace") + vt_bytes = f.read(4) + if len(vt_bytes) < 4: + break + vtype = struct.unpack(" 1 << 22: # 4 MB sanity bound + break + sbytes = f.read(slen) + if len(sbytes) < slen: + break + ctx_key = f"{sbytes.decode('utf-8', 'replace')}.context_length" + elif ctx_key is not None and key == ctx_key and vtype in (4, 10): + width = 4 if vtype == 4 else 8 + n_bytes = f.read(width) + if len(n_bytes) < width: + break + value = struct.unpack(" 0 else None + else: + if not _skip_gguf_value(f, vtype): + break + except (struct.error, UnicodeDecodeError): + break + except OSError as e: + logger.debug(f"read_gguf_context_length: cannot open {path}: {e}") + return None + except Exception as e: + logger.debug(f"read_gguf_context_length: parse failure on {path}: {e}") + return None + return None + + # Strings (8) and arrays (9) are handled inline. _FIXED_VTYPE_SIZES: Dict[int, int] = { 0: 1, # uint8 diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index 4792e78c4d..7fb9380037 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -119,6 +119,7 @@ function RootLayout() { chatRuntime.setActiveThreadId(null); chatRuntime.setActiveProjectId(null); chatRuntime.setIncognito(false); + if (chatRuntime.pendingSelection) chatRuntime.abandonStagedModel(); void navigate({ to: "/chat", search: { new: crypto.randomUUID() }, @@ -135,6 +136,7 @@ function RootLayout() { chatRuntime.setActiveProjectId(null); chatRuntime.setActiveThreadId(null); chatRuntime.setIncognito(false); + if (chatRuntime.pendingSelection) chatRuntime.abandonStagedModel(); }, [isChatRoute]); return ( diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index 80090bfd63..d17892e1bd 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -8,6 +8,8 @@ import { PopoverContent, PopoverTrigger, } from "@/components/ui/popover"; +import { InfoHint } from "@/components/ui/info-hint"; +import { Switch } from "@/components/ui/switch"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { usePlatformStore } from "@/config/env"; import { isCustomProviderType } from "@/features/chat/external-providers"; @@ -108,6 +110,11 @@ interface ModelSelectorProps { activeGgufVariant?: string | null; onValueChange?: (value: string, meta: ModelSelectorChangeMeta) => void; onEject?: () => void; + /** When provided, renders a persisted "Load on selection" toggle in the + * popover. Off → picking a model stages it for a deferred, configured load + * instead of loading immediately. */ + loadOnSelection?: boolean; + onLoadOnSelectionChange?: (value: boolean) => void; onFoldersChange?: () => void; onPickLocalModel?: () => void | Promise; onModelsChange?: (deletedModel?: DeletedModelRef) => void; @@ -210,6 +217,8 @@ function ModelSelectorContent({ onPickLocalModel, onModelsChange, deleteDisabled, + loadOnSelection, + onLoadOnSelectionChange, className, dataTour, }: { @@ -223,6 +232,8 @@ function ModelSelectorContent({ onPickLocalModel?: () => void; onModelsChange?: (deletedModel?: DeletedModelRef) => void; deleteDisabled?: boolean; + loadOnSelection?: boolean; + onLoadOnSelectionChange?: (value: boolean) => void; className?: string; dataTour?: string; }) { @@ -378,6 +389,37 @@ function ModelSelectorContent({ ) : null} + {onLoadOnSelectionChange ? ( +
+
+
+
+ Load on selection + +
+
+ On: load the model + immediately after selection. +
+
+ Off: configure options + first, then click Load model. +
+
+ +
+ + Local GGUF models only + +
+ +
+
+ ) : null} ); } @@ -395,6 +437,8 @@ export function ModelSelector({ onPickLocalModel, onModelsChange, deleteDisabled, + loadOnSelection, + onLoadOnSelectionChange, variant = "outline", size = "default", className, @@ -513,6 +557,8 @@ export function ModelSelector({ onPickLocalModel={onPickLocalModel ? handlePickLocalModel : undefined} onModelsChange={onModelsChange} deleteDisabled={deleteDisabled} + loadOnSelection={loadOnSelection} + onLoadOnSelectionChange={onLoadOnSelectionChange} className={contentClassName} dataTour={contentDataTour} /> diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 3165057021..6f11ec96f7 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -1357,15 +1357,14 @@ export function HubModelPicker({ <> LM Studio {lmStudioModels.map((m) => { + const isGgufFile = m.path.toLowerCase().endsWith(".gguf"); const isGguf = isGgufRepo(m.id) || isGgufRepo(m.display_name); const optionKey = makeModelOptionKey("lm-studio", m.id); return (
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/types.ts b/studio/frontend/src/components/assistant-ui/model-selector/types.ts index 4cc5d779ce..6a86e4f7ed 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/types.ts +++ b/studio/frontend/src/components/assistant-ui/model-selector/types.ts @@ -31,6 +31,9 @@ export interface ModelSelectorChangeMeta { ggufVariant?: string; isDownloaded?: boolean; expectedBytes?: number; + /** Direct local .gguf file picked without a variant (custom folder / LM + * Studio). Marks it as a GGUF source for the deferred-load staging flow. */ + isGguf?: boolean; } export interface DeletedModelRef { diff --git a/studio/frontend/src/components/ui/info-hint.tsx b/studio/frontend/src/components/ui/info-hint.tsx new file mode 100644 index 0000000000..4433478647 --- /dev/null +++ b/studio/frontend/src/components/ui/info-hint.tsx @@ -0,0 +1,43 @@ +// 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 type { ReactNode } from "react"; + +import { + Tooltip, + TooltipContent, + TooltipTrigger, +} from "@/components/ui/tooltip"; +import { InformationCircleIcon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; + +/** Small "i" affordance that reveals a styled tooltip on hover/focus. The + * standard inline help control across the settings UI. */ +export function InfoHint({ children }: { children: ReactNode }) { + return ( + + + + + + {children} + + + ); +} diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 2d2fa81012..f3235fa447 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 { consumeNativePathToken } from "@/features/native-intents/api"; import { formatFastApiDetail } from "@/lib/format-fastapi-error"; import type { MessageRecord, @@ -115,6 +116,45 @@ export async function validateModel( return parseJsonOrThrow(response); } +/** + * Read a GGUF's native context length from its local header (no GPU load, no + * download). Returns null when the file isn't downloaded yet, the model isn't a + * GGUF, or it's gated. For a native (drag-drop / picked) file, pass + * `nativePathToken` so the backend reads the granted local path. Used by the + * deferred-load staging flow to fill the context slider before the single load. + */ +export async function fetchGgufContextLength(payload: { + model_path: string; + gguf_variant?: string | null; + hf_token?: string | null; + nativePathToken?: string | null; +}): Promise { + let nativePathLease: string | null = null; + if (payload.nativePathToken) { + try { + nativePathLease = ( + await consumeNativePathToken(payload.nativePathToken, "validate-model") + ).nativePathLease; + } catch { + // Lease expired / revoked: degrade to no context (the load can re-mint). + return null; + } + } + const response = await authFetch("/api/inference/validate", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + model_path: payload.model_path, + gguf_variant: payload.gguf_variant ?? null, + hf_token: payload.hf_token ?? null, + native_path_lease: nativePathLease, + include_context_length: true, + }), + }); + const res = await parseJsonOrThrow(response); + return res.context_length ?? null; +} + export async function unloadModel(payload: UnloadModelRequest): Promise { const response = await authFetch("/api/inference/unload", { method: "POST", diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 3b518f83bf..5ed834fb9e 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -62,7 +62,10 @@ import { parseExternalModelId, } from "./external-providers"; import { useChatModelRuntime } from "./hooks/use-chat-model-runtime"; +import type { SelectedModelInput } from "./hooks/use-chat-model-runtime"; import { useChatProjects } from "./hooks/use-chat-projects"; +import { useStagedModelPreparation } from "./hooks/use-staged-model-preparation"; +import { useLatestRef } from "@/features/hub/hooks/use-latest-ref"; import { type SidebarItem, useChatSidebarItems, @@ -93,6 +96,7 @@ import { CHAT_IMAGE_TOOLS_ENABLED_KEY, CHAT_TOOLS_ENABLED_KEY, CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, + hasGgufSource, loadOptionalBool, useChatRuntimeStore, } from "./stores/chat-runtime-store"; @@ -1024,6 +1028,20 @@ export function ChatPage(): ReactElement { const settingsOpen = useChatRuntimeStore((s) => s.settingsPanelOpen); const setSettingsOpen = useChatRuntimeStore((s) => s.setSettingsPanelOpen); + const loadOnSelection = useChatRuntimeStore((s) => s.loadOnSelection); + const setLoadOnSelection = useChatRuntimeStore((s) => s.setLoadOnSelection); + // Deferred-load staging: downloads a staged GGUF (if needed) and reads its + // header context so the sheet can show the context slider before the load. + const stagedDownload = useStagedModelPreparation(); + // Abandon a staged pick: the store action cancels its in-flight download and + // reverts the edited knobs, so nothing lingers after the user walks away. + const abandonStaged = useCallback(() => { + useChatRuntimeStore.getState().abandonStagedModel(); + }, []); + // Tracks whether the chat page is still mounted, so a staged-load failure that + // resolves after the user left chat doesn't resurrect the abandoned pick. + const mountedRef = useRef(true); + useEffect(() => () => void (mountedRef.current = false), []); const incognito = useChatRuntimeStore((s) => s.incognito); const setIncognito = useChatRuntimeStore((s) => s.setIncognito); const incognitoLabel = incognito @@ -1503,12 +1521,64 @@ export function ChatPage(): ReactElement { closeArtifactSurface(); }, [activeThreadId, closeArtifactSurface, selectedArtifact, view]); + // Abandon a staged (not-yet-loaded) pick when the chat context actually + // changes — switching threads, leaving single view, or starting a new chat / + // project — so a stale Load button can't resurface in a different context. + // New Chat keeps activeThreadId null and only bumps the `new` search nonce, so + // the key includes the route identity, not just the thread. Mirrors the + // incognito reset pattern. (Route exit is handled in __root.tsx, which runs + // after this unmounts.) Clear only on a real change, never on mount: staging + // from the Hub sets pendingSelection then navigates here, and clearing on + // mount would wipe it. Comparing the previous context (rather than a first-run + // flag) is also safe under StrictMode's double-invoke and component remounts. + const chatContextKey = `${view.mode}|${activeThreadId ?? ""}|${search.new ?? ""}|${search.project ?? ""}`; + const chatContextKeyRef = useLatestRef(chatContextKey); + const prevChatContextRef = useRef(null); + useEffect(() => { + const prev = prevChatContextRef.current; + prevChatContextRef.current = chatContextKey; + if (prev === null || prev === chatContextKey) return; + abandonStaged(); + }, [chatContextKey, abandonStaged]); + const hasActiveModel = Boolean(inferenceParams.checkpoint); + // Load immediately, or — when "Load on selection" is off — stage the pick so + // its load options can be set first. Shared by the main selector, native + // drag-drop/picker, and the dropped-file chip (the Hub stages via the store). + const stageOrLoad = useCallback( + async (selection: SelectedModelInput) => { + const store = useChatRuntimeStore.getState(); + // Only GGUF picks have pre-load options worth staging. Non-GGUF models + // (and the toggle-on case) load immediately, so e.g. a trust_remote_code + // approval surfaces through the normal load path. + if (store.loadOnSelection || !hasGgufSource(selection)) { + // Abandon any staged GGUF first so its edited knobs (e.g. a custom + // context length) don't leak into this immediate load -- resolveLoad + // reads customContextLength before checking the target is GGUF. + abandonStaged(); + await selectModel(selection); + return; + } + // Tear down any existing staged pick first so its in-flight download is + // cancelled, not left running after we rebind to the new pick. + abandonStaged(); + store.stageModel({ + id: selection.id, + isLora: selection.isLora, + ggufVariant: selection.ggufVariant, + isDownloaded: selection.isDownloaded, + expectedBytes: selection.expectedBytes, + nativePathToken: selection.nativePathToken, + isGguf: selection.isGguf, + }); + }, + [abandonStaged, selectModel], + ); const loadNativeModelIntent = useCallback( async (intent: NativeIntent, loadingDescription: string) => { const label = intent.path.displayLabel || intent.displayLabel || "Local GGUF model"; - await selectModel({ + await stageOrLoad({ id: label, nativePathToken: intent.path.token, isDownloaded: true, @@ -1518,7 +1588,7 @@ export function ChatPage(): ReactElement { }); useNativeIntentStore.getState().clearModelIntent(intent.id); }, - [selectModel], + [stageOrLoad], ); const handleNativeModelDropAutoLoad = useCallback( (intent: NativeIntent) => @@ -1567,6 +1637,7 @@ export function ChatPage(): ReactElement { ggufVariant?: string; isDownloaded?: boolean; expectedBytes?: number; + isGguf?: boolean; }, ) => { const store = useChatRuntimeStore.getState(); @@ -1579,6 +1650,9 @@ export function ChatPage(): ReactElement { ) return; if (meta?.source === "external" || isExternalModelId(value)) { + // Switching to an external model abandons any staged local pick: cancel + // its download too (setCheckpoint below only clears the pending + knobs). + abandonStaged(); const selectedExternal = parseExternalModelId(value); const selectedProvider = selectedExternal ? externalProvidersForChat.find( @@ -1746,20 +1820,27 @@ export function ChatPage(): ReactElement { duration: 6000, }); } - await selectModel({ + const selection = { id: value, isLora: meta?.isLora, ggufVariant: meta?.ggufVariant, isDownloaded: meta?.isDownloaded, expectedBytes: meta?.expectedBytes, - }); + isGguf: meta?.isGguf, + }; + // "Load on selection" off: stage the model and open settings so its + // load knobs (tensor parallel, context length…) can be set, then it + // loads once via the sheet's Load button. The currently loaded model + // stays put until the user commits. + await stageOrLoad(selection); })(); }, [ + abandonStaged, activeThreadId, externalProvidersForChat, modelsFromStore, - selectModel, + stageOrLoad, view, ], ); @@ -2139,6 +2220,8 @@ export function ChatPage(): ReactElement { activeGgufVariant={activeGgufVariant} onValueChange={handleCheckpointChange} onEject={handleEject} + loadOnSelection={loadOnSelection} + onLoadOnSelectionChange={setLoadOnSelection} onFoldersChange={refreshLocalModels} onPickLocalModel={isTauri ? chooseNativeModel : undefined} onModelsChange={refreshModelLists} @@ -2180,7 +2263,7 @@ export function ChatPage(): ReactElement { selectModel(selection)} + onLoad={(selection) => stageOrLoad(selection)} /> ) : null} {loadingModel && loadToastDismissed ? ( @@ -2330,7 +2413,13 @@ export function ChatPage(): ReactElement { { + setSettingsOpen(open); + // Closing the sheet abandons a staged (not-yet-loaded) pick: cancel its + // download and revert the staged knobs so nothing lingers as a dirty + // edit (or a background download) on the loaded model. + if (!open) abandonStaged(); + }} params={inferenceParams} onParamsChange={setInferenceParams} isExternalModel={isExternalModel} @@ -2356,6 +2445,47 @@ export function ChatPage(): ReactElement { }); } }} + onLoadPendingModel={() => { + const pending = useChatRuntimeStore.getState().pendingSelection; + if (!pending) return; + const keyAtLoad = chatContextKey; + // forceReload: the staged model isn't loaded yet, so bypass the + // same-checkpoint dedupe (and selectModel clears pendingSelection). + // keepSpeculative: honor the speculative mode set on the sidebar. + void selectModel({ + ...pending, + forceReload: true, + keepSpeculative: true, + throwOnError: true, + }).catch(() => { + // Recoverable failure (expired token, gated repo, OOM…): selectModel + // cleared the pick but left the edited knobs intact. + const store = useChatRuntimeStore.getState(); + // A pick staged meanwhile owns the knobs now; leave it untouched. + if (store.pendingSelection) return; + // Restore (not re-stage, which would reset the knobs) only if the + // staged-load is still wanted: same chat context, sheet still open, + // page still mounted. + const stillWanted = + mountedRef.current && + store.settingsPanelOpen && + chatContextKeyRef.current === keyAtLoad; + if (stillWanted) { + store.setPendingSelection(pending); + } else { + // Abandoned (closed the sheet / switched chats / left chat): drop + // the orphaned staged knob edits so they don't linger as dirty + // settings over the loaded model. + store.resetModelSettingsToLoaded(); + } + }); + }} + stagedDownloadFraction={stagedDownload.progress?.fraction ?? null} + onCancelStagedDownload={() => + stagedDownload.cancelDownload( + useChatRuntimeStore.getState().pendingSelection?.ggufVariant ?? null, + ) + } />
); diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index 74c407fc3b..7e5b92ffb6 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -54,18 +54,14 @@ import { import { Slider } from "@/components/ui/slider"; import { Switch } from "@/components/ui/switch"; import { Textarea } from "@/components/ui/textarea"; -import { - Tooltip, - TooltipContent, - TooltipTrigger, -} from "@/components/ui/tooltip"; +import { InfoHint } from "@/components/ui/info-hint"; +import { Tooltip, TooltipContent } from "@/components/ui/tooltip"; import { useIsMobile } from "@/hooks/use-mobile"; import { useLlamaUpdateCheck } from "@/hooks/use-llama-update-check"; import { cn } from "@/lib/utils"; import { ArrowTurnBackwardIcon, Edit03Icon, - InformationCircleIcon, LayoutAlignRightIcon, } from "@hugeicons/core-free-icons"; import { ChevronDownStandardIcon } from "@/lib/chevron-icons"; @@ -101,7 +97,10 @@ import { providerSupportsBuiltinCodeExecution, providerSupportsFastMode, } from "./provider-capabilities"; -import { useChatRuntimeStore } from "./stores/chat-runtime-store"; +import { + isPendingGguf, + useChatRuntimeStore, +} from "./stores/chat-runtime-store"; import { RetrievalSettingsSection } from "@/features/rag/components/retrieval-settings-section"; import type { InferenceParams } from "./types/runtime"; @@ -112,33 +111,6 @@ function canUseStorage(): boolean { return typeof window !== "undefined"; } -export function InfoHint({ children }: { children: ReactNode }) { - return ( - - - - - - {children} - - - ); -} - /** * Editable numeric value display, shared by every slider value and the Context * Length input. An that looks like text (shows `displayValue ?? value`, @@ -466,6 +438,12 @@ interface ChatSettingsPanelProps { */ externalProviderType?: string | null; onReloadModel?: () => void; + /** Loads the staged `pendingSelection` (deferred "Load on selection" flow). */ + onLoadPendingModel?: () => void; + /** Download progress (0–1) for a staged GGUF being fetched, or null when idle. */ + stagedDownloadFraction?: number | null; + /** Cancels the in-flight staged download (paired with abandoning the stage). */ + onCancelStagedDownload?: () => void; } export function ChatSettingsPanel({ @@ -479,6 +457,9 @@ export function ChatSettingsPanel({ onExternalProviderChange, externalProviderType = null, onReloadModel, + onLoadPendingModel, + stagedDownloadFraction, + onCancelStagedDownload, }: ChatSettingsPanelProps) { // Local models show every knob; providerCapabilities is only consulted when // isExternalModel. Unknown providers fall back to the OpenAI-compat shape via @@ -493,9 +474,31 @@ export function ChatSettingsPanel({ const showPresencePenalty = !isExternalModel || Boolean(providerCapabilities?.presencePenalty); const isMobile = useIsMobile(); - const isGguf = useChatRuntimeStore((s) => s.activeGgufVariant) != null; + const pendingSelection = useChatRuntimeStore((s) => s.pendingSelection); + const abandonStagedModel = useChatRuntimeStore((s) => s.abandonStagedModel); + const resetModelSettingsToLoaded = useChatRuntimeStore( + (s) => s.resetModelSettingsToLoaded, + ); + // A staged GGUF pick (deferred load) shows the GGUF load knobs so they can be + // set before the single load. + const pendingIsGguf = isPendingGguf(pendingSelection); + // Short, human-readable name for the staged pick (HF ids carry an org prefix; + // native picks are already a display label). Drives the "staged, not loaded" + // callout so it's obvious the selection hasn't loaded yet. + const stagedLabel = (() => { + const id = pendingSelection?.id ?? ""; + const slash = id.lastIndexOf("/"); + const base = slash >= 0 ? id.slice(slash + 1) : id; + return base || id; + })(); + const isLoadedGguf = + useChatRuntimeStore((s) => s.activeGgufVariant) != null; + const isGguf = isLoadedGguf || pendingIsGguf; + // A staged pick is always a local GGUF, so show its Model section (and the + // Load button) even when the currently active model is external. const hasModelContent = - !isExternalModel && (isGguf || Boolean(params.checkpoint)); + pendingSelection != null || + (!isExternalModel && (isGguf || Boolean(params.checkpoint))); const speculativeType = useChatRuntimeStore((s) => s.speculativeType); const setSpeculativeType = useChatRuntimeStore((s) => s.setSpeculativeType); const loadedSpeculativeType = useChatRuntimeStore( @@ -560,8 +563,25 @@ export function ChatSettingsPanel({ const setActivePreset = useChatRuntimeStore((s) => s.setActivePreset); const settingsHydrated = useChatRuntimeStore((s) => s.settingsHydrated); - const ctxDisplayValue = customContextLength ?? ggufContextLength ?? ""; - const ctxMaxValue = ggufNativeContextLength ?? ggufContextLength ?? null; + // A staged (not-yet-loaded) GGUF carries its own header context length on + // pendingSelection, so the slider can use the staged model's real ceiling + // without reading the loaded model's `ggufContextLength`. + const stagedContextLength = pendingSelection?.contextLength ?? null; + // While staging, the sheet reflects the STAGED model, so its header context + // takes precedence over the loaded model's (which may differ or be larger). + const baseContext = pendingIsGguf ? stagedContextLength : ggufContextLength; + const baseNativeContext = pendingIsGguf + ? stagedContextLength + : ggufNativeContextLength; + // Context controls render once we actually have a ceiling: for a staged GGUF, + // once its header metadata arrives (post-download); otherwise post-load. + const showContextControl = pendingIsGguf + ? stagedContextLength != null + : isLoadedGguf; + const stagedDownloading = + stagedDownloadFraction != null && stagedDownloadFraction < 1; + const ctxDisplayValue = customContextLength ?? baseContext ?? ""; + const ctxMaxValue = baseNativeContext ?? baseContext ?? null; const kvDirty = kvCacheDtype !== loadedKvCacheDtype; const ctxDirty = customContextLength !== null; const specDirty = speculativeType !== loadedSpeculativeType; @@ -569,12 +589,6 @@ export function ChatSettingsPanel({ const tpDirty = tensorParallel !== (loadedTensorParallel ?? false); const modelSettingsDirty = kvDirty || ctxDirty || specDirty || specDraftDirty || tpDirty; - const loadedChatTemplateOverride = useChatRuntimeStore( - (s) => s.loadedChatTemplateOverride, - ); - const setChatTemplateOverride = useChatRuntimeStore( - (s) => s.setChatTemplateOverride, - ); const [presetNameInput, setPresetNameInput] = useState(activePreset); const [systemPromptEditorOpen, setSystemPromptEditorOpen] = useState(false); const [systemPromptDraft, setSystemPromptDraft] = useState(""); @@ -846,8 +860,19 @@ export function ChatSettingsPanel({ {hasModelContent && (
+ {pendingSelection && ( + + + {stagedLabel} is staged, not loaded yet + + + Set the options below, then choose Load model to load it. + + + )} {isGguf && ( <> + {showContextControl && (
@@ -857,14 +882,14 @@ export function ChatSettingsPanel({ value={ typeof ctxDisplayValue === "number" ? ctxDisplayValue - : (ggufContextLength ?? 0) + : (baseContext ?? 0) } min={128} max={ctxMaxValue ?? undefined} step={1} onChange={(v) => { setCustomContextLength( - v === (ggufContextLength ?? 0) ? null : v, + v === (baseContext ?? 0) ? null : v, ); }} ariaLabel="Context Length" @@ -879,14 +904,14 @@ export function ChatSettingsPanel({ Math.min( typeof ctxDisplayValue === "number" ? ctxDisplayValue - : (ggufContextLength ?? 4096), + : (baseContext ?? 4096), ctxMaxValue ?? 4096, ), ]} onValueChange={([v]) => { const snapped = Math.round(v); setCustomContextLength( - snapped === (ggufContextLength ?? 0) ? null : snapped, + snapped === (baseContext ?? 0) ? null : snapped, ); }} className="panel-slider" @@ -901,6 +926,7 @@ export function ChatSettingsPanel({

)}
+ )}
@@ -937,6 +963,8 @@ export function ChatSettingsPanel({
+ {isGguf && ( + <>
@@ -1043,6 +1071,8 @@ export function ChatSettingsPanel({ className="h-7 w-[76px] rounded-full border-transparent bg-black/[0.04] dark:bg-white/[0.05] hover:bg-black/[0.06] dark:hover:bg-white/[0.1] pl-3 pr-2 py-0 text-[13px] font-medium text-nav-fg outline-none focus-visible:ring-0" />
+ )} + )}
@@ -1099,8 +1129,44 @@ export function ChatSettingsPanel({ {/* Apply/Reset belongs to the model-reload settings above (context length, KV cache, speculative decoding). Render it here, before the Chat Template row, so it never reads as attached to Chat - Template (which is edited via its own dialog). */} - {modelSettingsDirty && ( + Template (which is edited via its own dialog). When a model is + staged (deferred load), Load/Cancel takes its place: there's + nothing loaded to "apply" against yet. */} + {pendingSelection ? ( +
+ {stagedDownloading && ( +

+ Downloading…{" "} + {Math.round((stagedDownloadFraction ?? 0) * 100)}% +

+ )} +
+ + +
+
+ ) : modelSettingsDirty ? (
- )} + ) : null}
@@ -1509,21 +1568,21 @@ export function ChatSettingsPanel({ : 64 } max={ - isExternalModel + // A staged GGUF caps to its own context even over an active + // external model (the staged model is what will load). + !pendingIsGguf && isExternalModel ? getExternalMaxOutputTokens( externalProviderType, externalSelection?.modelId, ) - : isGguf && ggufContextLength - ? ggufContextLength + : isGguf && baseContext + ? baseContext : 32768 } step={64} onChange={set("maxTokens")} displayValue={ - isGguf && - ggufContextLength && - params.maxTokens >= ggufContextLength + isGguf && baseContext && params.maxTokens >= baseContext ? "Max" : undefined } diff --git a/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx b/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx index 7dfa03672d..b28bf83737 100644 --- a/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx +++ b/studio/frontend/src/features/chat/components/openai-code-exec-section.tsx @@ -45,7 +45,7 @@ import { import { CHAT_HISTORY_UPDATED_EVENT } from "../api/chat-api"; import type { ExternalProviderConfig } from "../external-providers"; import { ensureThreadRecord } from "../runtime-provider"; -import { InfoHint } from "../chat-settings-sheet"; +import { InfoHint } from "@/components/ui/info-hint"; import { getStoredChatThread, listStoredChatThreads, diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index 55389e2bdf..a948d9e39d 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -43,12 +43,13 @@ import { isMultimodalResponse, } from "../types/api"; import { isExternalModelId } from "../external-providers"; +import { cancelStagedModelDownload } from "@/features/hub"; import type { ChatLoraSummary, ChatModelSummary, } from "../types/runtime"; -type SelectedModelInput = { +export type SelectedModelInput = { id: string; isLora?: boolean; ggufVariant?: string; @@ -57,7 +58,14 @@ type SelectedModelInput = { expectedBytes?: number; forceReload?: boolean; nativePathToken?: string; + /** Direct local .gguf file (no HF variant / native token) — still a GGUF + * source, so the staging flow treats it as one. */ + isGguf?: boolean; throwOnError?: boolean; + /** Keep the current speculative-decoding choice across the model switch + * instead of resetting it to the standing preference. Set by the deferred + * ("Load on selection") Load, where the user picked it for this model. */ + keepSpeculative?: boolean; }; const MODEL_LOAD_TOAST_CLASSNAMES = { @@ -370,8 +378,27 @@ export function useChatModelRuntime() { typeof selection === "string" ? false : selection.forceReload ?? false; const nativePathToken = typeof selection === "string" ? undefined : selection.nativePathToken; + const explicitIsGguf = + typeof selection === "string" ? undefined : selection.isGguf; const throwOnError = typeof selection === "string" ? false : selection.throwOnError ?? false; + const keepSpeculative = + typeof selection === "string" ? false : selection.keepSpeculative ?? false; + // Picking/loading any model abandons a staged (deferred) selection. + // Before the early-returns below so even a no-op re-select clears the + // stage, and so the Load button unmounts on first click (no double-load). + const staged = useChatRuntimeStore.getState().pendingSelection; + if (staged) { + // Loading a DIFFERENT model abandons this stage, so cancel its in-flight + // download. Loading the staged pick itself keeps it (that download feeds + // this load). + const loadingStagedPick = + staged.id === modelId && + (staged.ggufVariant ?? null) === (ggufVariant ?? null) && + (staged.nativePathToken ?? null) === (nativePathToken ?? null); + if (!loadingStagedPick) cancelStagedModelDownload(staged); + useChatRuntimeStore.getState().setPendingSelection(null); + } const currentVariant = useChatRuntimeStore.getState().activeGgufVariant; if (!forceReload && (!modelId || (params.checkpoint === modelId && (ggufVariant ?? null) === (currentVariant ?? null)))) { return; @@ -391,6 +418,7 @@ export function useChatModelRuntime() { typeof selection === "string" ? false : selection.isDownloaded ?? false; const model = models.find((entry) => entry.id === modelId); const lora = loras.find((entry) => entry.id === modelId); + const isGguf = explicitIsGguf ?? model?.isGguf ?? false; const loraIsAdapter = lora?.exportType === "lora"; const isLora = explicitIsLora ?? model?.isLora ?? loraIsAdapter ?? false; @@ -501,7 +529,10 @@ export function useChatModelRuntime() { // can't follow the user onto a model without an MTP head. // spec_draft_n_max is MTP-only and always resets. The loaded // shadow is seeded too, preventing a transient dirty Apply state. - if (currentCheckpoint && currentCheckpoint !== modelId) { + // keepSpeculative skips this for a staged Load: the user picked the + // mode for this model on the sidebar, so honor it (the backend still + // falls back at runtime if the model has no MTP head). + if (currentCheckpoint && currentCheckpoint !== modelId && !keepSpeculative) { const persistedSpeculativeType = readPersistedSpeculativeType(); useChatRuntimeStore.setState({ speculativeType: persistedSpeculativeType, @@ -525,6 +556,7 @@ export function useChatModelRuntime() { const effectiveMaxSeqLength = resolveLoadMaxSeqLength({ modelId, ggufVariant, + isGguf, customContextLength, ggufContextLength, currentCheckpoint, diff --git a/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts b/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts new file mode 100644 index 0000000000..d154b931d4 --- /dev/null +++ b/studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts @@ -0,0 +1,114 @@ +// 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 { useCallback, useEffect } from "react"; + +import { useLatestRef } from "@/features/hub/hooks/use-latest-ref"; +import { useRepoDownload } from "@/features/hub/download-manager/use-repo-download"; +import type { DownloadJob } from "@/features/hub/download-manager/use-repo-download"; + +import { fetchGgufContextLength } from "../api/chat-api"; +import { + isPendingGguf, + useChatRuntimeStore, +} from "../stores/chat-runtime-store"; + +/** + * Drives the deferred ("Load on selection" off) staging flow for a GGUF: + * download the file if needed (HF repo) or read it in place (native drag-drop / + * picked file), then read its header context length so the settings sheet can + * show the real context slider before the single GPU load. The staged context + * lands on `pendingSelection.contextLength` (scoped to the staged model, never + * the loaded model's `ggufContextLength`). Returns the live download job so the + * sheet can render progress / cancel. Mount once on the chat page. + */ +export function useStagedModelPreparation(): DownloadJob { + const pendingId = useChatRuntimeStore((s) => s.pendingSelection?.id ?? null); + const pendingVariant = useChatRuntimeStore( + (s) => s.pendingSelection?.ggufVariant ?? null, + ); + const pendingNativeToken = useChatRuntimeStore( + (s) => s.pendingSelection?.nativePathToken ?? null, + ); + // Only GGUF picks (HF variant or native file) have a header worth reading. + const pendingIsGguf = useChatRuntimeStore((s) => + isPendingGguf(s.pendingSelection), + ); + const pendingDownloaded = useChatRuntimeStore( + (s) => s.pendingSelection?.isDownloaded ?? false, + ); + const pendingHasContext = useChatRuntimeStore( + (s) => s.pendingSelection?.contextLength != null, + ); + const setPendingSelection = useChatRuntimeStore((s) => s.setPendingSelection); + + const fetchContextMetadata = useCallback(async () => { + const current = useChatRuntimeStore.getState().pendingSelection; + if (!current?.id || !isPendingGguf(current)) return; + const { id, ggufVariant, nativePathToken } = current; + try { + const contextLength = await fetchGgufContextLength({ + model_path: id, + gguf_variant: ggufVariant, + hf_token: useChatRuntimeStore.getState().hfToken || null, + nativePathToken, + }); + // Apply only if the same model is still staged (the user may have switched + // picks or loaded/cancelled while the request was in flight). Native ids + // are display labels, not paths, so two files can share an id -- compare + // the path token too, or a stale response could land on the wrong pick. + const latest = useChatRuntimeStore.getState().pendingSelection; + if ( + latest?.id === id && + (latest.ggufVariant ?? null) === (ggufVariant ?? null) && + (latest.nativePathToken ?? null) === (nativePathToken ?? null) && + contextLength != null + ) { + setPendingSelection({ ...latest, contextLength }); + } + } catch { + // Leave contextLength null: the context slider stays hidden and the user + // can still load (context fills in from the load response afterwards). + } + }, [setPendingSelection]); + + const job = useRepoDownload({ + kind: "model", + // useRepoDownload must be called unconditionally; an idle repo id keeps it + // inert until something is staged. + repoId: pendingId ?? "__staged_idle__", + activeVariant: pendingVariant, + onComplete: () => { + void fetchContextMetadata(); + }, + }); + + // job.requestStartDownload's identity changes per render; hold it in a ref so + // the staging effect re-runs only when the staged model itself changes. + const startDownloadRef = useLatestRef(job.requestStartDownload); + const fetchMetadataRef = useLatestRef(fetchContextMetadata); + + useEffect(() => { + if (!pendingId || !pendingIsGguf || pendingHasContext) return; + // Native files and already-downloaded HF files are local: read the header + // now. Otherwise download first; onComplete then reads it. + if (pendingNativeToken || pendingDownloaded) { + void fetchMetadataRef.current(); + } else { + const expectedBytes = + useChatRuntimeStore.getState().pendingSelection?.expectedBytes ?? 0; + void startDownloadRef.current(pendingVariant, expectedBytes); + } + }, [ + pendingId, + pendingVariant, + pendingNativeToken, + pendingIsGguf, + pendingDownloaded, + pendingHasContext, + startDownloadRef, + fetchMetadataRef, + ]); + + return job; +} diff --git a/studio/frontend/src/features/chat/presets/preset-policy.ts b/studio/frontend/src/features/chat/presets/preset-policy.ts index 4efbb74f11..057e4778bf 100644 --- a/studio/frontend/src/features/chat/presets/preset-policy.ts +++ b/studio/frontend/src/features/chat/presets/preset-policy.ts @@ -297,6 +297,7 @@ export function mergeBackendRecommendedInference({ export function resolveLoadMaxSeqLength({ modelId, ggufVariant, + isGguf, customContextLength, ggufContextLength, currentCheckpoint, @@ -306,6 +307,7 @@ export function resolveLoadMaxSeqLength({ }: { modelId: string; ggufVariant?: string | null; + isGguf?: boolean | null; customContextLength: number | null; ggufContextLength: number | null; currentCheckpoint: string; @@ -314,7 +316,7 @@ export function resolveLoadMaxSeqLength({ presetSource: ChatPresetSource; }): number { const isDirectGgufFile = modelId.toLowerCase().endsWith(".gguf"); - const isGgufLoad = ggufVariant != null || isDirectGgufFile; + const isGgufLoad = isGguf === true || ggufVariant != null || isDirectGgufFile; const isReloadingCurrentGguf = isGgufLoad && currentCheckpoint === modelId && diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 003a32ddb1..68b8d1d523 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -3,6 +3,7 @@ import { toast } from "@/lib/toast"; import { create } from "zustand"; +import { cancelStagedModelDownload } from "@/features/hub"; import { type ChatPresetSource, type Preset, @@ -35,6 +36,7 @@ export const CHAT_ALLOW_ARTIFACT_NETWORK_ACCESS_KEY = "unsloth_chat_allow_artifact_network_access"; export const CHAT_MCP_ENABLED_KEY = "unsloth_chat_mcp_enabled"; export const CHAT_CONFIRM_TOOL_CALLS_KEY = "unsloth_chat_confirm_tool_calls"; +export const CHAT_LOAD_ON_SELECTION_KEY = "unsloth_chat_load_on_selection"; export const CHAT_BYPASS_PERMISSIONS_KEY = "unsloth_chat_bypass_permissions"; export const CHAT_WEB_FETCH_TOOLS_ENABLED_KEY = "unsloth_chat_web_fetch_tools_enabled"; @@ -86,6 +88,7 @@ function saveRagSource(value: RagSource): void { try { window.localStorage.setItem(CHAT_RAG_SOURCE_KEY, JSON.stringify(value)); } catch { + // Ignore storage failures; the default RAG source still works for this session. } } @@ -408,6 +411,42 @@ function notifyHfTokenChanged(value: string): void { } } +/** A local model staged for a deferred load (see `pendingSelection`). Shape is + * a subset of the load hook's `SelectedModelInput`, structurally assignable. */ +export type PendingModelSelection = { + id: string; + isLora?: boolean; + ggufVariant?: string; + isDownloaded?: boolean; + expectedBytes?: number; + /** Native (drag-drop / picked-from-disk) GGUF: the path token used to read + * the header and to load. Absent for HF-repo models. */ + nativePathToken?: string; + /** Direct local .gguf file (custom folder / LM Studio): a GGUF source even + * though it carries neither an HF variant nor a native path token. */ + isGguf?: boolean; + /** Native context length read from the GGUF header once the file is local. + * Scoped here (not the shared `ggufContextLength`) so a staged model's + * metadata never pollutes the currently-loaded model's context display. */ + contextLength?: number | null; +}; + +/** A pick is a GGUF (HF variant, native file, or a direct local .gguf) and so + * has pre-load options worth staging. Works on a selection or a staged pick. */ +export function hasGgufSource(x: { + ggufVariant?: string; + nativePathToken?: string; + isGguf?: boolean; +}): boolean { + return ( + x.ggufVariant != null || x.nativePathToken != null || x.isGguf === true + ); +} + +export function isPendingGguf(pending: PendingModelSelection | null): boolean { + return pending != null && hasGgufSource(pending); +} + type ChatRuntimeStore = { settingsHydrated: boolean; params: InferenceParams; @@ -538,6 +577,13 @@ type ChatRuntimeStore = { tensorParallel: boolean; /** Backend-reported tensor-parallel state; null until first hydrated. */ loadedTensorParallel: boolean | null; + /** Persisted: when false, picking a local model stages it as + * `pendingSelection` (and opens settings) instead of loading immediately, + * so load settings can be set before the single load. */ + loadOnSelection: boolean; + /** A local model picked while `loadOnSelection` is off: staged, not loaded. + * The settings sheet shows its load knobs and a Load button. */ + pendingSelection: PendingModelSelection | null; loadedIsMultimodal: boolean; /** Active model is a block-diffusion model (DiffusionGemma): drives the * denoising-canvas artifact auto-render. */ @@ -638,7 +684,20 @@ type ChatRuntimeStore = { setKvCacheDtype: (dtype: string | null) => void; setSpeculativeType: (type: string | null) => void; setSpecDraftNMax: (value: number | null) => void; + /** Revert the editable load knobs to the loaded model's baseline (or defaults + * when nothing is loaded). Used by the settings-sheet Reset button and to + * start each deferred-staging session clean so one staged pick's settings + * don't leak onto the next. */ + resetModelSettingsToLoaded: () => void; setTensorParallel: (value: boolean) => void; + setLoadOnSelection: (value: boolean) => void; + setPendingSelection: (selection: PendingModelSelection | null) => void; + /** Stage a pick for a deferred load: revert knobs to the loaded baseline, + * record the selection, and open the settings sheet. */ + stageModel: (selection: PendingModelSelection) => void; + /** Abandon a staged pick without loading: revert the knobs to the loaded + * baseline and clear the pending selection. */ + abandonStagedModel: () => void; setCustomContextLength: (v: number | null) => void; setChatTemplateOverride: (template: string | null) => void; setPendingAudio: (base64: string, name: string) => void; @@ -839,6 +898,23 @@ function setScalarSettingVersion( saveSettingsPatch({ [key]: value }); } +/** The "revert to the loaded model" baseline for the editable load knobs. + * Shared by resetModelSettingsToLoaded (full revert) and stageModel (which + * overrides speculative to start a fresh pick from the standing default). */ +function loadedBaselineSettings(s: ChatRuntimeStore) { + const hasLoadedModel = Boolean(s.params.checkpoint); + return { + customContextLength: null, + kvCacheDtype: s.loadedKvCacheDtype, + tensorParallel: s.loadedTensorParallel ?? false, + speculativeType: hasLoadedModel + ? s.loadedSpeculativeType + : readPersistedSpeculativeType(), + specDraftNMax: hasLoadedModel ? s.loadedSpecDraftNMax : null, + chatTemplateOverride: s.loadedChatTemplateOverride, + }; +} + export const useChatRuntimeStore = create((set, get) => ({ settingsHydrated: false, // Hydrate the last external checkpoint so the external picker survives a @@ -924,6 +1000,8 @@ export const useChatRuntimeStore = create((set, get) => ({ loadedSpecDraftNMax: null, tensorParallel: false, loadedTensorParallel: null, + loadOnSelection: loadBool(CHAT_LOAD_ON_SELECTION_KEY, true), + pendingSelection: null, loadedIsMultimodal: false, loadedIsDiffusion: false, customContextLength: null, @@ -1061,6 +1139,11 @@ export const useChatRuntimeStore = create((set, get) => ({ // Clear stale per-turn usage on model change; the relaxed external-provider // render gate would otherwise show old counters until the next completion. const checkpointChanged = state.params.checkpoint !== modelId; + const pendingToClear = + checkpointChanged && state.params.checkpoint ? state.pendingSelection : null; + if (pendingToClear) { + cancelStagedModelDownload(pendingToClear); + } // Clamp maxTokens to the new model's cap when switching into an external // model so a value carried over from a local session doesn't exceed the // slider's max. @@ -1088,6 +1171,14 @@ export const useChatRuntimeStore = create((set, get) => ({ }, activeGgufVariant: ggufVariant ?? null, ...(checkpointChanged ? { contextUsage: null } : {}), + // Switching away from a loaded model (e.g. picking an external provider) + // abandons any staged pick, so its Load button and edited knobs don't + // linger over the newly active model. Same revert as abandonStagedModel. + // Guarded on a non-empty current checkpoint: an establishing set from a + // background status sync (empty -> active) must not wipe a fresh stage. + ...(pendingToClear + ? { ...loadedBaselineSettings(state), pendingSelection: null } + : {}), }; }), setActiveThreadId: (activeThreadId) => @@ -1100,6 +1191,7 @@ export const useChatRuntimeStore = create((set, get) => ({ // clear any stored external selection so the next refresh doesn't snap // back to a model the user intentionally cleared. saveLastExternalCheckpoint(null); + cancelStagedModelDownload(get().pendingSelection); return set((state) => ({ params: { ...state.params, @@ -1107,6 +1199,7 @@ export const useChatRuntimeStore = create((set, get) => ({ }, activeGgufVariant: null, activeNativePathToken: null, + pendingSelection: null, ggufContextLength: null, ggufMaxContextLength: null, ggufNativeContextLength: null, @@ -1339,6 +1432,43 @@ export const useChatRuntimeStore = create((set, get) => ({ setSpeculativeType: (speculativeType) => set({ speculativeType }), setSpecDraftNMax: (specDraftNMax) => set({ specDraftNMax }), setTensorParallel: (tensorParallel) => set({ tensorParallel }), + resetModelSettingsToLoaded: () => set((s) => loadedBaselineSettings(s)), + setLoadOnSelection: (loadOnSelection) => { + saveBool(CHAT_LOAD_ON_SELECTION_KEY, loadOnSelection); + set({ loadOnSelection }); + }, + setPendingSelection: (pendingSelection) => set({ pendingSelection }), + stageModel: (selection) => + set((s) => { + if ( + s.pendingSelection && + (s.pendingSelection.id !== selection.id || + (s.pendingSelection.ggufVariant ?? null) !== + (selection.ggufVariant ?? null)) + ) { + cancelStagedModelDownload(s.pendingSelection); + } + return { + ...loadedBaselineSettings(s), + pendingSelection: selection, + settingsPanelOpen: true, + // Speculative starts from the standing default, not the loaded model's + // mode, so a fresh pick doesn't inherit (and then carry, via the staged + // Load's keepSpeculative) a forced MTP mode onto a model that may lack it. + speculativeType: readPersistedSpeculativeType(), + specDraftNMax: null, + }; + }), + abandonStagedModel: () => { + const { pendingSelection } = get(); + if (!pendingSelection) return; + // Cancel the staged pick's in-flight download so it doesn't keep running + // after the staging UI is gone. Centralized here so every abandon path + // (sheet close, thread switch, route exit, new chat) cancels it, including + // root-level callers that have no access to the useRepoDownload hook. + cancelStagedModelDownload(pendingSelection); + set((s) => ({ ...loadedBaselineSettings(s), pendingSelection: null })); + }, setCustomContextLength: (customContextLength) => set({ customContextLength }), setChatTemplateOverride: (chatTemplateOverride) => set({ chatTemplateOverride }), diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 781172fb4f..556caa0617 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -72,6 +72,8 @@ export interface ValidateModelResponse { is_lora?: boolean; is_vision?: boolean; requires_trust_remote_code?: boolean; + /** Native context length from the local GGUF header; null until downloaded. */ + context_length?: number | null; } export interface GgufVariantDetail { diff --git a/studio/frontend/src/features/hub/download-manager/download-manager-controller.ts b/studio/frontend/src/features/hub/download-manager/download-manager-controller.ts index abfeaf3c0d..da653ddeb7 100644 --- a/studio/frontend/src/features/hub/download-manager/download-manager-controller.ts +++ b/studio/frontend/src/features/hub/download-manager/download-manager-controller.ts @@ -1,10 +1,14 @@ // 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 { DOWNLOAD_KIND } from "./constants"; import { createDownloadManagerInitialState, + jobKeyOf, removeJob, + selectActiveJob, setState, + useDownloadManagerStore, } from "./download-manager-state"; import { resetDownloadApiAdapterState } from "./download-api-adapter"; import { @@ -65,6 +69,26 @@ export const downloadManager: DownloadManagerController = { dismiss: removeJob, }; +/** Cancel the in-flight download for a staged model pick. No-op when nothing is + * downloading (e.g. a native/local file that was never fetched). Lets non-React + * callers (the chat store's abandon paths) stop a staged transfer without the + * useRepoDownload hook. */ +export function cancelStagedModelDownload( + pending: { id: string; ggufVariant?: string | null } | null, +): void { + if (!pending) return; + const variant = pending.ggufVariant ?? null; + const activeJob = selectActiveJob( + useDownloadManagerStore.getState(), + DOWNLOAD_KIND.MODEL, + pending.id, + variant, + ); + void downloadManager.cancel( + activeJob?.key ?? jobKeyOf(DOWNLOAD_KIND.MODEL, pending.id, variant), + ); +} + if (import.meta.hot) { import.meta.hot.dispose(() => { __resetDownloadManagerForTests(); diff --git a/studio/frontend/src/features/hub/download-manager/index.ts b/studio/frontend/src/features/hub/download-manager/index.ts index 60ef3851f8..dd88aaf3f0 100644 --- a/studio/frontend/src/features/hub/download-manager/index.ts +++ b/studio/frontend/src/features/hub/download-manager/index.ts @@ -20,6 +20,7 @@ export { } from "./constants"; export { __resetDownloadManagerForTests, + cancelStagedModelDownload, clearCompletedInventoryHint, downloadManager, hydrateDownloadManager, diff --git a/studio/frontend/src/features/hub/hub-page.tsx b/studio/frontend/src/features/hub/hub-page.tsx index 5ce58f4362..de78ce7438 100644 --- a/studio/frontend/src/features/hub/hub-page.tsx +++ b/studio/frontend/src/features/hub/hub-page.tsx @@ -1004,6 +1004,23 @@ export function ModelsPage() { (opts: ModelLoadOptions, isDownloaded: boolean) => { if (!selectedModel) return; const runId = selectedModel.resource.runId; + // "Load on selection" off: stage GGUF picks instead of loading, so the + // chat page's staging flow can read the header and show the load options. + // Non-GGUF models have nothing to configure pre-load, so they load now. + if ( + !useChatRuntimeStore.getState().loadOnSelection && + (opts.ggufVariant != null || selectedModel.isGguf) + ) { + useChatRuntimeStore.getState().stageModel({ + id: runId, + ggufVariant: opts.ggufVariant, + isGguf: selectedModel.isGguf, + isDownloaded, + expectedBytes: opts.expectedBytes, + }); + openNewChat(); + return; + } void selectModel({ id: runId, ggufVariant: opts.ggufVariant, @@ -1012,6 +1029,7 @@ export function ModelsPage() { throwOnError: true, }) .then(() => { + // Read fresh: the load is async, so the checkpoint may have changed. const store = useChatRuntimeStore.getState(); if (!modelIdsMatch(store.params.checkpoint, runId)) { store.setCheckpoint(runId, opts.ggufVariant ?? null); diff --git a/studio/frontend/src/features/hub/index.ts b/studio/frontend/src/features/hub/index.ts new file mode 100644 index 0000000000..ddcef146b0 --- /dev/null +++ b/studio/frontend/src/features/hub/index.ts @@ -0,0 +1,4 @@ +// 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 { cancelStagedModelDownload } from "./download-manager";