) => {
@@ -125,7 +124,7 @@ export function TrainingSection() {
title={t("studio.training.title")}
description={t("studio.training.description")}
accent="blue"
- className={hasMessage ? "min-h-studio-config-column" : "h-studio-config-column"}
+ className="min-h-studio-config-column"
>
{/* Loss chart */}
diff --git a/studio/frontend/src/index.css b/studio/frontend/src/index.css
index 159eba743f..333ed70f6a 100644
--- a/studio/frontend/src/index.css
+++ b/studio/frontend/src/index.css
@@ -2048,15 +2048,11 @@ html[data-chat-font] .aui-root {
}
}
- /* Fine-tuning Studio: equal default height, expandable when needed (md+) */
+ /* Fine-tuning Studio: equal minimum card height, grows with content (md+) */
.min-h-studio-config-column {
@apply md:min-h-[520px];
}
- .h-studio-config-column {
- @apply md:h-[520px];
- }
-
[data-streamdown="unordered-list"] {
list-style-type: disc;
list-style-position: outside;
From 8c83478da02b08bfcecef691eda78f15f050cbe9 Mon Sep 17 00:00:00 2001
From: Nilay <118994073+NilayYadav@users.noreply.github.com>
Date: Thu, 16 Jul 2026 20:54:33 +0530
Subject: [PATCH 020/289] Studio: use one shared Hugging Face token across
Settings and training (#7152)
---------
Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
---
.../chat/stores/chat-runtime-store.ts | 33 ++++++++-----------
studio/frontend/src/features/hub/index.ts | 5 +++
.../src/features/hub/stores/hf-token-store.ts | 25 ++++++++------
.../studio/sections/dataset-section.tsx | 14 ++------
.../training/stores/training-config-store.ts | 24 +++++++++++---
5 files changed, 56 insertions(+), 45 deletions(-)
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 abdaa871c0..af99458349 100644
--- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
+++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts
@@ -2,7 +2,11 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import type { RememberedLoadSettings } from "@/components/assistant-ui/model-selector/remembered-load-settings";
-import { cancelStagedModelDownload } from "@/features/hub";
+import {
+ cancelStagedModelDownload,
+ mirrorHfTokenInto,
+ useHfTokenStore,
+} from "@/features/hub";
import { toast } from "@/lib/toast";
import { create } from "zustand";
import { isExternalModelId, parseExternalModelId } from "../external-providers";
@@ -25,8 +29,6 @@ import {
import { useExternalProvidersStore } from "./external-providers-store";
import { PLUS_MENU_PINS_STORAGE_KEY } from "./plus-menu-prefs-store";
-const HF_TOKEN_KEY = "unsloth_hf_token";
-const HF_TOKEN_CHANGED_EVENT = "unsloth:hf-token-changed";
export const CHAT_REASONING_ENABLED_KEY = "unsloth_chat_reasoning_enabled";
export const CHAT_TOOLS_ENABLED_KEY = "unsloth_chat_tools_enabled";
export const CHAT_CODE_TOOLS_ENABLED_KEY = "unsloth_chat_code_tools_enabled";
@@ -495,17 +497,6 @@ export function saveSpeculativeType(value: string | null): void {
}
}
-function notifyHfTokenChanged(value: string): void {
- if (!canUseStorage()) return;
- try {
- window.dispatchEvent(
- new CustomEvent(HF_TOKEN_CHANGED_EVENT, { detail: value }),
- );
- } catch {
- // ignore
- }
-}
-
/** 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 = {
@@ -1144,7 +1135,7 @@ export const useChatRuntimeStore = create((set, get) => ({
runningByThreadId: {},
cancelByThreadId: {},
autoTitle: false,
- hfToken: loadString(HF_TOKEN_KEY, ""),
+ hfToken: useHfTokenStore.getState().token,
modelsError: null,
lastModelLoadError: null,
activeGgufVariant: null,
@@ -1349,11 +1340,7 @@ export const useChatRuntimeStore = create((set, get) => ({
setScalarSettingVersion("autoTitle", autoTitle, state.autoTitle);
return { autoTitle };
}),
- setHfToken: (hfToken) => {
- saveString(HF_TOKEN_KEY, hfToken);
- set({ hfToken });
- notifyHfTokenChanged(hfToken);
- },
+ setHfToken: (hfToken) => useHfTokenStore.getState().setToken(hfToken),
setModelsError: (modelsError) => set({ modelsError }),
setLastModelLoadError: (lastModelLoadError) => set({ lastModelLoadError }),
setCheckpoint: (modelId, ggufVariant) =>
@@ -1837,6 +1824,12 @@ export const useChatRuntimeStore = create((set, get) => ({
setContextUsage: (contextUsage) => set({ contextUsage }),
}));
+// Mirror token edits made through the shared store (e.g. Studio's field).
+const unsubscribeHfTokenMirror = mirrorHfTokenInto(useChatRuntimeStore);
+if (import.meta.hot) {
+ import.meta.hot.dispose(unsubscribeHfTokenMirror);
+}
+
export function resolveSpeculativeSettingsForLoad({
usePersistedPreference = false,
}: {
diff --git a/studio/frontend/src/features/hub/index.ts b/studio/frontend/src/features/hub/index.ts
index ddcef146b0..3515f6ca76 100644
--- a/studio/frontend/src/features/hub/index.ts
+++ b/studio/frontend/src/features/hub/index.ts
@@ -2,3 +2,8 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
export { cancelStagedModelDownload } from "./download-manager";
+export {
+ getHfToken,
+ mirrorHfTokenInto,
+ useHfTokenStore,
+} from "./stores/hf-token-store";
diff --git a/studio/frontend/src/features/hub/stores/hf-token-store.ts b/studio/frontend/src/features/hub/stores/hf-token-store.ts
index b1e2560f02..499ba9f644 100644
--- a/studio/frontend/src/features/hub/stores/hf-token-store.ts
+++ b/studio/frontend/src/features/hub/stores/hf-token-store.ts
@@ -5,11 +5,9 @@ import { create } from "zustand";
import { bumpInventoryVersion } from "./inventory-events";
const HF_TOKEN_KEY = "unsloth_hf_token";
-const HF_TOKEN_CHANGED_EVENT = "unsloth:hf-token-changed";
const LEGACY_TRAINING_KEY = "unsloth_training_config_v1";
let storageSyncStarted = false;
let storageSyncListener: ((event: StorageEvent) => void) | null = null;
-let tokenChangedListener: ((event: Event) => void) | null = null;
function canUseStorage(): boolean {
return typeof window !== "undefined";
@@ -63,10 +61,6 @@ function stopStorageSync(): void {
window.removeEventListener("storage", storageSyncListener);
storageSyncListener = null;
}
- if (tokenChangedListener !== null) {
- window.removeEventListener(HF_TOKEN_CHANGED_EVENT, tokenChangedListener);
- tokenChangedListener = null;
- }
storageSyncStarted = false;
}
@@ -100,10 +94,6 @@ export const useHfTokenStore = create((set) => {
applyToken(event.newValue ?? "", false);
};
window.addEventListener("storage", storageSyncListener);
- tokenChangedListener = (event) => {
- applyToken((event as CustomEvent).detail ?? "", false);
- };
- window.addEventListener(HF_TOKEN_CHANGED_EVENT, tokenChangedListener);
}
return {
@@ -117,6 +107,21 @@ export function getHfToken(): string {
return useHfTokenStore.getState().token;
}
+// Keep a plain zustand store's `hfToken` field in sync with the shared token:
+// seed the current value, then mirror later edits. Returns the unsubscribe so
+// callers can wire it to HMR disposal.
+export function mirrorHfTokenInto(store: {
+ getState: () => T;
+ setState: (partial: Partial) => void;
+}): () => void {
+ store.setState({ hfToken: getHfToken() } as Partial);
+ return useHfTokenStore.subscribe((state) => {
+ if (store.getState().hfToken !== state.token) {
+ store.setState({ hfToken: state.token } as Partial);
+ }
+ });
+}
+
// HF's JS client throws on a non-empty token that isn't `hf_...` instead of
// browsing anonymously, so treat anything malformed as no token.
export function hfApiToken(
diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx
index bb74b0d061..e7bab1f47f 100644
--- a/studio/frontend/src/features/studio/sections/dataset-section.tsx
+++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx
@@ -54,7 +54,7 @@ import {
// Imported directly from the store module rather than the "@/features/training"
// barrel to avoid an import cycle (the barrel re-exports this section's siblings).
import { hasSeparateStreamingEvalSplit } from "@/features/training/stores/training-config-store";
-import { useDebouncedValue, useHfTokenValidation } from "@/hooks";
+import { useDebouncedValue } from "@/hooks";
import { translate, useT } from "@/i18n";
import { ChevronDownStandardIcon } from "@/lib/chevron-icons";
import { toast } from "@/lib/toast";
@@ -398,9 +398,6 @@ export function DatasetSection() {
enabled: pickerTab === "huggingface",
});
- const { error: tokenValidationError, isChecking: isCheckingToken } =
- useHfTokenValidation(hfToken);
-
const hfResultIds = useMemo(() => {
const ids = hfResults.map((r) => r.id);
if (dataset && !ids.includes(dataset)) {
@@ -1005,9 +1002,9 @@ export function DatasetSection() {
- {(tokenValidationError ?? hfSearchError) && (
+ {hfSearchError && (
- {tokenValidationError ?? hfSearchError}
+ {hfSearchError}
{" — "}
)}
- {isCheckingToken && (
-
- {t("studio.dataset.checkingToken")}
-
- )}
{pickerTab !== activeSourceTab && (
{t("studio.dataset.browsingSource", {
diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts
index 306c02475c..a927f83fd8 100644
--- a/studio/frontend/src/features/training/stores/training-config-store.ts
+++ b/studio/frontend/src/features/training/stores/training-config-store.ts
@@ -3,6 +3,7 @@
import { CPT_TARGET_MODULES, DEFAULT_HYPERPARAMS, LR_DEFAULT_CPT, LR_DEFAULT_FULL, LR_DEFAULT_LORA, STEPS, TARGET_MODULES } from "@/config/training";
import { authFetch } from "@/features/auth";
+import { getHfToken, mirrorHfTokenInto, useHfTokenStore } from "@/features/hub";
import { isAdapterMethod } from "@/types/training";
import type { DatasetFormat } from "@/types/training";
import type { ModelType, StepNumber, TrainingMethod } from "@/types/training";
@@ -117,7 +118,9 @@ let _datasetFormatAutoForcedByCpt = false;
// modelType / isVisionModel / isAudioModel persist so multimodal-only UI
// paints right on reload; the model-config fetch still re-derives them.
+// hfToken mirrors the shared hf-token-store and is persisted there instead.
const NON_PERSISTED_STATE_KEYS: ReadonlySet = new Set([
+ "hfToken",
"isCheckingVision",
"isEmbeddingModel",
"isLoadingModelDefaults",
@@ -632,8 +635,7 @@ export const useTrainingConfigStore = create()(
),
);
},
- setHfToken: (hfToken) =>
- set({ hfToken: hfToken.trim().replace(/^["']+|["']+$/g, "") }),
+ setHfToken: (hfToken) => useHfTokenStore.getState().setToken(hfToken),
setDatasetSource: (datasetSource) => set({ datasetSource }),
selectHfDataset: (dataset) => {
_datasetCheckController?.abort();
@@ -923,7 +925,7 @@ export const useTrainingConfigStore = create()(
_learningRateManuallySet = false;
_yamlLearningRate = undefined;
clearCptDatasetFormatTracking();
- set(initialState);
+ set({ ...initialState, hfToken: getHfToken() });
},
resetToModelDefaults: () => {
const { selectedModel } = get();
@@ -947,7 +949,7 @@ export const useTrainingConfigStore = create()(
},
{
name: "unsloth_training_config_v1",
- version: 11,
+ version: 12,
migrate: (persisted, version) => {
const s = persisted as Record;
if (version < 2 && s.datasetSubset == null && s.datasetConfig != null) {
@@ -1000,6 +1002,15 @@ export const useTrainingConfigStore = create()(
// own version guard.
s.datasetStreaming ??= false;
}
+ if (version < 12) {
+ // hfToken moved to the shared hf-token-store; seed it once so an
+ // existing Studio-only token isn't lost.
+ const legacyToken = typeof s.hfToken === "string" ? s.hfToken.trim() : "";
+ if (legacyToken && !getHfToken()) {
+ useHfTokenStore.getState().setToken(legacyToken);
+ }
+ delete s.hfToken;
+ }
return s as unknown as TrainingConfigStore;
},
partialize: partializePersistedState,
@@ -1022,3 +1033,8 @@ export const useTrainingConfigStore = create()(
},
),
);
+
+const unsubscribeHfTokenMirror = mirrorHfTokenInto(useTrainingConfigStore);
+if (import.meta.hot) {
+ import.meta.hot.dispose(unsubscribeHfTokenMirror);
+}
From 3555dbdda7cf1fe5eb7f7036045b5076d56ca6bd Mon Sep 17 00:00:00 2001
From: oobabooga
Date: Thu, 16 Jul 2026 19:47:22 -0300
Subject: [PATCH 021/289] Studio: don't drop parallel tool calls after an
internal no-op (#7157)
---
studio/backend/core/inference/llama_cpp.py | 10 ++-
.../core/inference/safetensors_agentic.py | 10 ++-
.../core/inference/tool_loop_controller.py | 20 ++++-
.../backend/tests/test_llama_cpp_tool_loop.py | 74 +++++++++++++++++++
.../tests/test_safetensors_tool_loop.py | 55 ++++++++++++++
.../tests/test_tool_loop_controller.py | 28 ++++++-
6 files changed, 188 insertions(+), 9 deletions(-)
diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index 06b68ea831..dadfdfd38d 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -82,6 +82,7 @@ from core.inference.tool_call_parser import (
)
from core.inference.tool_loop_controller import (
ToolLoopController,
+ append_deferred_nudges,
tool_event_provenance,
)
from state.tool_approvals import (
@@ -10111,6 +10112,9 @@ class LlamaCppBackend:
assistant_msg: dict = {"role": "assistant", "content": content_text}
assistant_appended = False
+ # Collect no-op nudges and flush them after the batch, so a no-op
+ # doesn't abort it and drop the parallel calls that follow.
+ deferred_noop_msgs: list = []
# The text-path provisional card uses the parser's default id ("call_0");
# a Mistral-style call carries its own id and would open a duplicate. Reuse
@@ -10153,14 +10157,14 @@ class LlamaCppBackend:
"provenance": decision.provenance,
}
completion = tool_controller.record_noop(decision)
- conversation.append(completion.model_message())
+ deferred_noop_msgs.append(completion.model_message())
if _forced_tool_call_pending:
_forced_tool_call_pending = False
logger.info(
"Suppressed local GGUF tool call as internal no-op: "
f"action={decision.action} tool={decision.tool_name}"
)
- break
+ continue
if not assistant_appended:
assistant_msg["tool_calls"] = [decision.as_assistant_tool_call()]
@@ -10279,6 +10283,8 @@ class LlamaCppBackend:
if _forced_tool_call_pending:
_forced_tool_call_pending = False
+ append_deferred_nudges(conversation, deferred_noop_msgs)
+
# Close provisional cards not resolved by execution/no-op handling.
for _pid, _pname in provisional_started_tool_calls.items():
if _pid not in resolved_provisional_tool_call_ids:
diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py
index f4c243d1bf..43b72110ff 100644
--- a/studio/backend/core/inference/safetensors_agentic.py
+++ b/studio/backend/core/inference/safetensors_agentic.py
@@ -57,6 +57,7 @@ from core.tool_healing import (
)
from core.inference.tool_loop_controller import (
ToolLoopController,
+ append_deferred_nudges,
coerce_tool_arguments,
status_for_tool,
tool_event_provenance,
@@ -1099,6 +1100,9 @@ def run_safetensors_tool_loop(
assistant_msg: dict = {"role": "assistant", "content": content_text}
assistant_appended = False
+ # Collect no-op nudges and flush them after the batch, so a no-op doesn't
+ # abort it and drop the parallel calls that follow.
+ deferred_noop_msgs: list = []
for tc in tool_calls or []:
func = tc.get("function", {}) or {}
@@ -1127,12 +1131,12 @@ def run_safetensors_tool_loop(
"provenance": decision.provenance,
}
completion = tool_controller.record_noop(decision)
- conversation.append(completion.model_message())
+ deferred_noop_msgs.append(completion.model_message())
logger.info(
"Suppressed local safetensors tool call as internal no-op: "
f"action={decision.action} tool={decision.tool_name}"
)
- break
+ continue
if not assistant_appended:
assistant_msg["tool_calls"] = [decision.as_assistant_tool_call()]
@@ -1243,6 +1247,8 @@ def run_safetensors_tool_loop(
yield completion.tool_end_event()
conversation.append(completion.tool_message())
+ append_deferred_nudges(conversation, deferred_noop_msgs)
+
# Clear the status badge before the next turn.
yield {"type": "status", "text": ""}
diff --git a/studio/backend/core/inference/tool_loop_controller.py b/studio/backend/core/inference/tool_loop_controller.py
index f595531b90..f7ed450d11 100644
--- a/studio/backend/core/inference/tool_loop_controller.py
+++ b/studio/backend/core/inference/tool_loop_controller.py
@@ -266,6 +266,17 @@ def strip_result_for_model(result: str) -> str:
return result
+def append_deferred_nudges(conversation: list, msgs: Sequence[dict]) -> None:
+ """Append a batch's no-op nudges as one deduped ``role=user`` message.
+
+ Deferred to after the batch's tool results so a no-op never splits an
+ assistant's ``tool_calls`` from their ``role=tool`` results.
+ """
+ contents = list(dict.fromkeys(msg["content"] for msg in msgs))
+ if contents:
+ conversation.append({"role": "user", "content": "\n\n".join(contents)})
+
+
def _tool_name_from_schema(tool: Mapping[str, Any]) -> str:
function = tool.get("function")
if not isinstance(function, Mapping):
@@ -277,8 +288,9 @@ def _tool_name_from_schema(tool: Mapping[str, Any]) -> str:
def _noop_result(reason: NoopReason, tool_name: str) -> str:
if reason == "duplicate":
return (
- "The previous tool request was not executed because this exact "
- "tool call already completed successfully. Do not repeat the same "
+ f"One earlier request to call tool '{tool_name}' in this batch was "
+ "not executed because an identical call had already completed "
+ "successfully. Do not repeat the same "
"tool call. Continue with a different enabled tool if that would "
"materially help, or provide the final answer if you have enough "
"information."
@@ -291,8 +303,8 @@ def _noop_result(reason: NoopReason, tool_name: str) -> str:
"the requested final note or answer."
)
return (
- f"The previous tool request was not executed because tool "
- f"'{tool_name}' is not enabled for this request. Provide the "
+ f"One earlier request to call tool '{tool_name}' in this batch was "
+ "not executed because that tool is not enabled for this request. Provide the "
"final answer now without calling more tools."
)
diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py
index c3161c5714..bd2c008589 100644
--- a/studio/backend/tests/test_llama_cpp_tool_loop.py
+++ b/studio/backend/tests/test_llama_cpp_tool_loop.py
@@ -1061,6 +1061,80 @@ def test_same_turn_duplicate_web_search_is_internal_noop(monkeypatch):
]
+def test_same_turn_duplicate_does_not_drop_later_parallel_call(monkeypatch):
+ # One batch: search(a), search(a) [duplicate], search(b). The duplicate is an
+ # internal no-op, but the distinct search(b) after it must still run, and the
+ # no-op nudge must land after the tool results rather than splitting them.
+ batch = [
+ _sse(
+ {
+ "tool_calls": [
+ {
+ "index": 0,
+ "id": "call_a1",
+ "type": "function",
+ "function": {"name": "web_search", "arguments": json.dumps({"query": "a"})},
+ },
+ {
+ "index": 1,
+ "id": "call_a2",
+ "type": "function",
+ "function": {"name": "web_search", "arguments": json.dumps({"query": "a"})},
+ },
+ {
+ "index": 2,
+ "id": "call_b",
+ "type": "function",
+ "function": {"name": "web_search", "arguments": json.dumps({"query": "b"})},
+ },
+ ]
+ }
+ ),
+ _done(),
+ ]
+ final_stream = [_sse({"content": "Final answer."}), _done()]
+ payloads: list[dict] = []
+ backend = _make_backend(monkeypatch, [batch, final_stream], payloads)
+
+ calls: list[dict] = []
+
+ def fake_execute_tool(name, arguments, **_kwargs):
+ calls.append(arguments)
+ return "search-result"
+
+ monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
+
+ events = list(
+ backend.generate_chat_completion_with_tools(
+ messages = [{"role": "user", "content": "search"}],
+ tools = [{"type": "function", "function": {"name": "web_search"}}],
+ max_tool_iterations = 3,
+ )
+ )
+
+ # Both distinct calls ran; the duplicate did not (old `break` dropped search(b)).
+ assert calls == [{"query": "a"}, {"query": "b"}]
+ assert [e.get("tool_call_id") for e in events if e.get("type") == "tool_end"] == [
+ "call_a1",
+ "call_b",
+ ]
+
+ # The next generation's conversation must be well-formed: the assistant lists
+ # only the executed calls (no orphan for the duplicate), the two tool results
+ # follow contiguously, and the no-op nudge lands after them, never between.
+ conv = payloads[1]["messages"]
+ asst = next(m for m in conv if m["role"] == "assistant" and m.get("tool_calls"))
+ assert [tc.get("id") for tc in asst["tool_calls"]] == ["call_a1", "call_b"]
+ after = conv[conv.index(asst) + 1 :]
+ assert [m["role"] for m in after[:2]] == ["tool", "tool"]
+ assert [m.get("tool_call_id") for m in after[:2]] == ["call_a1", "call_b"]
+ assert after[2]["role"] == "user" # deferred duplicate nudge, after the results
+ assert after[2]["content"].startswith(
+ "One earlier request to call tool 'web_search' in this batch was not executed"
+ )
+ assert "previous tool request" not in after[2]["content"].lower()
+
+
def test_same_turn_repeated_render_html_does_not_emit_second_provisional_start(monkeypatch):
same_turn_render_calls = [
_sse(
diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py
index eae1a75161..e3633de289 100644
--- a/studio/backend/tests/test_safetensors_tool_loop.py
+++ b/studio/backend/tests/test_safetensors_tool_loop.py
@@ -2843,6 +2843,61 @@ class TestLoopBehaviour:
]
assert len(duplicate_nudges) == 1
+ def test_same_turn_duplicate_does_not_drop_later_parallel_call(self):
+ # Turn 1 runs search(x). Turn 2's batch is [search(x) duplicate, python]:
+ # the duplicate is a no-op, but python after it must still run, and the
+ # no-op nudge must land after python's result rather than splitting it.
+ captured_messages: list[list[dict]] = []
+ turns = iter(
+ [
+ ['{"name":"web_search","arguments":{"query":"x"}}'],
+ [
+ '{"name":"web_search","arguments":{"query":"x"}}'
+ '{"name":"python","arguments":{"code":"print(1)"}}'
+ ],
+ ["final"],
+ ]
+ )
+
+ def fake_single_turn(messages, active_tools = None):
+ captured_messages.append([dict(m) for m in messages])
+ chunks = next(turns)
+ acc = ""
+ for chunk in chunks:
+ acc += chunk
+ yield acc
+
+ exec_fn = FakeExecuteTool(["search-x", "py-result"])
+ _collect_events(
+ run_safetensors_tool_loop(
+ single_turn = fake_single_turn,
+ messages = [{"role": "user", "content": "hi"}],
+ tools = [
+ {"type": "function", "function": {"name": "web_search"}},
+ {"type": "function", "function": {"name": "python"}},
+ ],
+ execute_tool = exec_fn,
+ max_tool_iterations = 4,
+ )
+ )
+
+ # Turn-1 search and turn-2 python both ran; the turn-2 duplicate search did not.
+ assert exec_fn.calls == [
+ ("web_search", {"query": "x"}),
+ ("python", {"code": "print(1)"}),
+ ]
+
+ conv = captured_messages[-1]
+ turn2 = [m for m in conv if m.get("role") == "assistant" and m.get("tool_calls")][-1]
+ assert [tc["function"]["name"] for tc in turn2["tool_calls"]] == ["python"]
+ after = conv[conv.index(turn2) + 1 :]
+ assert after[0]["role"] == "tool" and after[0]["content"] == "py-result"
+ assert after[1]["role"] == "user" # deferred duplicate nudge, after the result
+ assert after[1]["content"].startswith(
+ "One earlier request to call tool 'web_search' in this batch was not executed"
+ )
+ assert "previous tool request" not in after[1]["content"].lower()
+
def test_duplicate_tool_call_internal_noop_allows_distinct_followup_tool(self):
captured_messages: list[list[dict]] = []
captured_tool_names: list[list[str]] = []
diff --git a/studio/backend/tests/test_tool_loop_controller.py b/studio/backend/tests/test_tool_loop_controller.py
index 0e8ae798af..496c30ac13 100644
--- a/studio/backend/tests/test_tool_loop_controller.py
+++ b/studio/backend/tests/test_tool_loop_controller.py
@@ -13,6 +13,7 @@ if _BACKEND_DIR not in sys.path:
from core.inference.tool_loop_controller import (
ToolLoopController,
+ append_deferred_nudges,
canonical_tool_call_key,
coerce_tool_arguments,
status_for_tool,
@@ -21,6 +22,22 @@ from core.inference.tool_loop_controller import (
)
+def test_append_deferred_nudges_merges_deduped_into_one_message():
+ conversation = [{"role": "assistant", "tool_calls": [1]}, {"role": "tool", "content": "r"}]
+ nudges = [
+ {"role": "user", "content": "duplicate"},
+ {"role": "user", "content": "duplicate"}, # dropped: same content
+ {"role": "user", "content": "disabled foo"},
+ ]
+ append_deferred_nudges(conversation, nudges)
+ # One user message, after the results, with distinct contents joined.
+ assert conversation[2:] == [{"role": "user", "content": "duplicate\n\ndisabled foo"}]
+ # Empty is a no-op.
+ before = list(conversation)
+ append_deferred_nudges(conversation, [])
+ assert conversation == before
+
+
def _tool(name: str) -> dict:
return {"type": "function", "function": {"name": name}}
@@ -111,6 +128,10 @@ def test_successful_duplicate_is_internal_noop_and_keeps_remaining_tools():
assert not duplicate.should_execute
assert not duplicate.emit_visible_events
duplicate_nudge = completion.model_message()["content"]
+ assert duplicate_nudge.startswith(
+ "One earlier request to call tool 'web_search' in this batch was not executed"
+ )
+ assert "previous tool request" not in duplicate_nudge.lower()
assert "already completed successfully" in duplicate_nudge
assert "different enabled tool" in duplicate_nudge
assert completion.model_message()["role"] == "user"
@@ -165,7 +186,12 @@ def test_empty_enabled_tool_list_blocks_all_tool_calls():
assert decision.action == "disabled"
assert not decision.emit_visible_events
assert completion.model_message()["role"] == "user"
- assert "not enabled" in completion.model_message()["content"]
+ disabled_nudge = completion.model_message()["content"]
+ assert disabled_nudge.startswith(
+ "One earlier request to call tool 'web_search' in this batch was not executed"
+ )
+ assert "previous tool request" not in disabled_nudge.lower()
+ assert "not enabled" in disabled_nudge
assert controller.force_final_answer
assert controller.active_tools() == []
From c4e6dd4f6c60c95ac6d2951693a2ffc40abcde6a Mon Sep 17 00:00:00 2001
From: oobabooga
Date: Thu, 16 Jul 2026 19:48:12 -0300
Subject: [PATCH 022/289] Studio: extract text from PDF web results (#7154)
---
studio/backend/core/inference/tools.py | 117 +++++++++++++--
studio/backend/core/rag/parsers.py | 55 +++++--
studio/backend/tests/test_rag_parsing.py | 50 +++++++
.../tests/test_web_fetch_binary_guard.py | 138 ++++++++++++++++--
4 files changed, 324 insertions(+), 36 deletions(-)
diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py
index dd268a6bb7..5fd57e1b2c 100644
--- a/studio/backend/core/inference/tools.py
+++ b/studio/backend/core/inference/tools.py
@@ -3600,14 +3600,18 @@ _MAX_PAGE_CHARS = 16000 # cap fetched page text (after HTML-to-MD conversion)
# Raw download cap > _MAX_PAGE_CHARS since SSR pages embed large sections
# stripped during conversion; 512 KB still reaches article content.
_MAX_FETCH_BYTES = 512 * 1024
+# PDF cross-reference data lives at EOF, so extraction needs the whole body.
+_MAX_PDF_FETCH_BYTES = 10 * 1024 * 1024
+_MAX_WEB_PDF_PAGES = 50
# Control/undecodable chars, excluding text whitespace and ESC (for ANSI logs).
# Binary when they exceed 12.5%, after allowing 16 minor encoding glitches.
_BINARY_CHAR_RE = re.compile("[\\x00-\\x08\\x0b\\x0c\\x0e-\\x1a\\x1c-\\x1f\\x7f-\\x9f\\ufffd]")
_MIN_BINARY_CHARS = 16
_BINARY_CHAR_DIVISOR = 8
# Common binary signatures that can otherwise look text-heavy when mislabeled.
+_PDF_MAGIC = b"%PDF-"
_BINARY_MAGIC = (
- b"%PDF-", # PDF
+ _PDF_MAGIC,
b"PK\x03\x04", # zip / docx / xlsx / pptx / epub / jar
b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1", # OLE / legacy Office
b"\x89PNG\r\n\x1a\n", # PNG
@@ -3641,14 +3645,22 @@ def _looks_binary(text: str) -> bool:
)
-def _has_binary_magic(data: bytes) -> bool:
- """Whether a common binary signature follows optional BOM or whitespace."""
+def _magic_head(data: bytes) -> bytes:
head = data[:1024].lstrip()
for bom, _codec in _UNICODE_BOM_CODECS:
if head.startswith(bom):
head = head.removeprefix(bom).lstrip()
break
- return head.startswith(_BINARY_MAGIC)
+ return head
+
+
+def _has_pdf_magic(data: bytes) -> bool:
+ return _magic_head(data).startswith(_PDF_MAGIC)
+
+
+def _has_binary_magic(data: bytes) -> bool:
+ """Whether a common binary signature follows optional BOM or whitespace."""
+ return _magic_head(data).startswith(_BINARY_MAGIC)
def _has_single_byte_text_evidence(data: bytes) -> bool:
@@ -3659,6 +3671,45 @@ def _has_single_byte_text_evidence(data: bytes) -> bool:
return ascii_text_bytes / len(data) >= _MIN_SINGLE_BYTE_ASCII_RATIO
+def _extract_pdf_text(data: bytes) -> str:
+ """Extract page-delimited text with the same parser used by RAG ingestion."""
+ from ..rag.parsers import parse_pdf_bytes
+
+ pages, total_pages = parse_pdf_bytes(data, max_pages = _MAX_WEB_PDF_PAGES)
+ page_limit_reached = total_pages > _MAX_WEB_PDF_PAGES
+ parts: list[str] = []
+ length = 0
+ text_limited = False
+ for page in pages:
+ page_text = page.text.strip()
+ if not page_text:
+ continue
+ section = f"## Page {page.page_number}\n\n{page_text}"
+ piece = ("\n\n" if parts else "") + section
+ remaining = _MAX_PAGE_CHARS - length
+ if len(piece) > remaining:
+ parts.append(piece[:remaining])
+ text_limited = True
+ break
+ parts.append(piece)
+ length += len(piece)
+
+ text = "".join(parts).rstrip()
+ if not text:
+ if page_limit_reached:
+ return f"(PDF contains no extractable text in the first {_MAX_WEB_PDF_PAGES} pages)"
+ return ""
+ limits = []
+ if text_limited:
+ limits.append(f"text limited to {_MAX_PAGE_CHARS:,} characters")
+ if page_limit_reached:
+ limits.append(f"page processing capped at {_MAX_WEB_PDF_PAGES} pages")
+ if limits:
+ marker = f"\n\n... (PDF extraction {'; '.join(limits)})"
+ text = text[: _MAX_PAGE_CHARS - len(marker)].rstrip() + marker
+ return text
+
+
_USER_AGENTS = (
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
@@ -4054,29 +4105,71 @@ def _fetch_url_raw(
return reason2, "", ""
current_host = rp.hostname
continue
+
+ # get_content_type() defaults to "text/plain" when the header is
+ # absent (RFC 2045); report "" instead so callers can tell a missing
+ # header apart from a server that really declared text/plain.
+ if resp.headers.get("Content-Type") is None:
+ content_type = ""
+ else:
+ content_type = (resp.headers.get_content_type() or "").lower()
+
# Success: read the capped body enforcing the budget between chunks
# (see _read_capped_body), so a slow-drip server can't stretch a
# single resp.read past the deadline.
+ declared_pdf = content_type == "application/pdf"
+ read_limit = _MAX_PDF_FETCH_BYTES + 1 if declared_pdf else max_bytes
body_error, raw_bytes = _read_capped_body(
resp,
- max_bytes,
+ read_limit,
timeout,
deadline,
cancel_event,
)
if body_error is not None:
return body_error, "", ""
+
+ # A missing or wrong PDF MIME type is common: once the initial text-sized
+ # read identifies PDF magic, finish the bounded download to reach the EOF xref.
+ if not declared_pdf and len(raw_bytes) == max_bytes and _has_pdf_magic(raw_bytes):
+ tail_error, tail = _read_capped_body(
+ resp,
+ _MAX_PDF_FETCH_BYTES - max_bytes + 1,
+ timeout,
+ deadline,
+ cancel_event,
+ )
+ if tail_error is not None:
+ return tail_error, "", ""
+ raw_bytes += tail
break
else:
return "Failed to fetch URL: too many redirects.", "", ""
- # get_content_type() defaults to "text/plain" when the header is
- # absent (RFC 2045); report "" instead so callers can tell a missing
- # header apart from a server that really declared text/plain.
- if resp.headers.get("Content-Type") is None:
- content_type = ""
- else:
- content_type = (resp.headers.get_content_type() or "").lower()
+ is_pdf = declared_pdf or _has_pdf_magic(raw_bytes)
+ if is_pdf:
+ if len(raw_bytes) > _MAX_PDF_FETCH_BYTES:
+ return (
+ "(PDF content exceeds the download limit; not readable as text)",
+ "",
+ content_type,
+ )
+ budget_error = _fetch_budget_exceeded(deadline, cancel_event)
+ if budget_error is not None:
+ return budget_error, "", content_type
+ try:
+ pdf_text = _extract_pdf_text(raw_bytes)
+ except Exception as exc:
+ logger.debug("web PDF text extraction failed (%s)", type(exc).__name__)
+ return "(PDF content could not be read as text)", "", content_type
+ budget_error = _fetch_budget_exceeded(deadline, cancel_event)
+ if budget_error is not None:
+ return budget_error, "", content_type
+ if not pdf_text:
+ pdf_text = "(PDF contains no extractable text)"
+ # Report the true type even for a mislabeled body so the caller's "html"
+ # check routes the extracted text to the plain-text path, not html_to_markdown.
+ return None, pdf_text, "application/pdf"
# Reject known-binary MIME types before decoding. Binary is returned as the
# error string so the caller surfaces the placeholder, not replacement chars.
diff --git a/studio/backend/core/rag/parsers.py b/studio/backend/core/rag/parsers.py
index 9afddf1d9e..0b42906b85 100644
--- a/studio/backend/core/rag/parsers.py
+++ b/studio/backend/core/rag/parsers.py
@@ -103,7 +103,7 @@ def _markdown_incomplete(markdown: str, plain: str) -> bool:
return markdown_letters < _PDF_INCOMPLETE_RATIO * plain_letters
-def _pdf_markdown(doc) -> list[str] | None:
+def _pdf_markdown(doc, pages: range | None = None) -> list[str] | None:
"""Per-page layout-aware Markdown (tables, headings, lists) via pymupdf4llm; index
i maps to page i+1. Returns None when the lib is missing, extraction fails, or the
page count does not line up, so the caller falls back to plain PyMuPDF text."""
@@ -112,28 +112,44 @@ def _pdf_markdown(doc) -> list[str] | None:
except Exception:
return None
try:
- chunks = pymupdf4llm.to_markdown(
- doc,
- page_chunks = True,
- show_progress = False,
- )
+ kwargs = {"page_chunks": True, "show_progress": False}
+ if pages is not None:
+ kwargs["pages"] = list(pages)
+ chunks = pymupdf4llm.to_markdown(doc, **kwargs)
except Exception: # noqa: BLE001 - never let Markdown extraction break ingestion
logger.warning("pymupdf4llm extraction failed; using plain text", exc_info = True)
return None
- if not isinstance(chunks, list) or len(chunks) != doc.page_count:
+ expected_pages = doc.page_count if pages is None else len(pages)
+ if not isinstance(chunks, list) or len(chunks) != expected_pages:
return None
return [str(c.get("text") or "") for c in chunks]
-def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]:
+def _pdf(
+ source: str | bytes,
+ want_images: bool,
+ max_pages: int | None = None,
+) -> tuple[list[Page], list[ParsedImage], int]:
import fitz # PyMuPDF
pages: list[Page] = []
images: list[ParsedImage] = []
- doc = fitz.open(path)
+ doc = (
+ fitz.open(stream = source, filetype = "pdf") if isinstance(source, bytes) else fitz.open(source)
+ )
try:
- md = _pdf_markdown(doc) if config.PDF_MARKDOWN else None
- for i, page in enumerate(doc):
+ if doc.needs_pass:
+ raise ValueError("encrypted PDF requires a password")
+ total_pages = doc.page_count
+ page_numbers = range(total_pages if max_pages is None else min(total_pages, max_pages))
+ if not config.PDF_MARKDOWN:
+ md = None
+ elif max_pages is None:
+ md = _pdf_markdown(doc)
+ else:
+ md = _pdf_markdown(doc, page_numbers)
+ for i, page_number in enumerate(page_numbers):
+ page = doc[page_number]
plain = page.get_text("text") or ""
candidate = md[i] if md else ""
# Prefer layout-aware Markdown (keeps tables/headings legible for retrieval),
@@ -147,7 +163,7 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]:
text = candidate
else:
text = plain
- pages.append(_page(text, i + 1))
+ pages.append(_page(text, page_number + 1))
if want_images:
for img in page.get_images(full = True):
xref = img[0]
@@ -161,13 +177,22 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]:
images.append(
ParsedImage(
image_bytes = image_bytes,
- page_number = i + 1,
+ page_number = page_number + 1,
xref = xref,
)
)
finally:
doc.close()
- return pages, images
+ return pages, images, total_pages
+
+
+def parse_pdf_bytes(data: bytes, *, max_pages: int | None = None) -> tuple[list[Page], int]:
+ """Extract PDF pages from an in-memory download using the ingestion parser.
+
+ Returns the (capped) pages plus the document's full page count, so a caller
+ that set ``max_pages`` can tell a fully-read short PDF from a truncated one."""
+ pages, _images, total_pages = _pdf(data, want_images = False, max_pages = max_pages)
+ return pages, total_pages
def _merge_rects(boxes: list) -> list:
@@ -416,7 +441,7 @@ def parse(path: str, *, want_images: bool = False):
ext = os.path.splitext(path)[1].lower()
if ext == ".pdf":
- pages, images = _pdf(path, want_images)
+ pages, images, _total = _pdf(path, want_images)
return (pages, images) if want_images else pages
if ext == ".docx":
diff --git a/studio/backend/tests/test_rag_parsing.py b/studio/backend/tests/test_rag_parsing.py
index 14ab0efe2e..3e259f6bd0 100644
--- a/studio/backend/tests/test_rag_parsing.py
+++ b/studio/backend/tests/test_rag_parsing.py
@@ -54,6 +54,56 @@ def test_pdf_markdown_off_uses_plain_text(tmp_path, monkeypatch):
assert "#" not in text and "|" not in text # plain text path emits no Markdown markup
+def test_pdf_bytes_use_same_extraction_path(tmp_path, monkeypatch):
+ from core.rag import config, parsers
+
+ monkeypatch.setattr(config, "PDF_MARKDOWN", False)
+ pdf = tmp_path / "table.pdf"
+ _table_pdf(pdf)
+ from_file = parsers.parse(str(pdf))
+ from_bytes, total_pages = parsers.parse_pdf_bytes(pdf.read_bytes())
+ assert [page.text for page in from_bytes] == [page.text for page in from_file]
+ assert total_pages == len(from_file)
+
+
+def test_pdf_bytes_limit_pages_before_extraction(monkeypatch):
+ import pymupdf
+
+ from core.rag import config, parsers
+
+ monkeypatch.setattr(config, "PDF_MARKDOWN", False)
+ doc = pymupdf.open()
+ for marker in ("page one", "page two", "page three"):
+ page = doc.new_page()
+ page.insert_text((40, 40), marker)
+ data = doc.tobytes()
+ doc.close()
+
+ pages, total_pages = parsers.parse_pdf_bytes(data, max_pages = 2)
+ assert len(pages) == 2
+ assert "page two" in pages[-1].text
+ assert total_pages == 3 # full count, not the 2 extracted
+
+
+def test_pdf_markdown_receives_page_limit(monkeypatch):
+ from core.rag import parsers
+
+ captured = {}
+
+ class _FakePymupdf4llm:
+ @staticmethod
+ def to_markdown(doc, **kwargs):
+ captured.update(kwargs)
+ return [{"text": "page"} for _ in kwargs["pages"]]
+
+ class _Doc:
+ page_count = 100
+
+ monkeypatch.setitem(__import__("sys").modules, "pymupdf4llm", _FakePymupdf4llm)
+ assert parsers._pdf_markdown(_Doc(), range(2)) == ["page", "page"]
+ assert captured == {"page_chunks": True, "show_progress": False, "pages": [0, 1]}
+
+
def test_pdf_markdown_passes_only_supported_legacy_kwargs(monkeypatch):
# The pinned PyMuPDF4LLM legacy path ignores unknown kwargs; do not pass the
# newer layout-only OCR knobs or Markdown extraction silently loses policy control.
diff --git a/studio/backend/tests/test_web_fetch_binary_guard.py b/studio/backend/tests/test_web_fetch_binary_guard.py
index 3041ed5c34..10db953913 100644
--- a/studio/backend/tests/test_web_fetch_binary_guard.py
+++ b/studio/backend/tests/test_web_fetch_binary_guard.py
@@ -59,6 +59,18 @@ def _fetch_with(monkeypatch, body: bytes, content_type: str | None) -> str:
return tools._fetch_page_text("https://example.com/thing", timeout = 5)
+def _pdf_bytes(*page_texts: str) -> bytes:
+ pymupdf = pytest.importorskip("pymupdf")
+ doc = pymupdf.open()
+ for text in page_texts:
+ page = doc.new_page()
+ if text:
+ page.insert_textbox(pymupdf.Rect(40, 40, 550, 750), text, fontsize = 11)
+ data = doc.tobytes()
+ doc.close()
+ return data
+
+
@pytest.mark.parametrize(
"content_type,expected",
[
@@ -90,10 +102,119 @@ def test_is_text_candidate_content_type(content_type, expected):
assert tools._is_text_candidate_content_type(content_type) is expected
-def test_pdf_rejected_by_content_type(monkeypatch):
- out = _fetch_with(monkeypatch, b"%PDF-1.7\n\xff\xd8\xff\x00\x89PNG" * 200, "application/pdf")
- assert "�" not in out
- assert "non-text content" in out and "application/pdf" in out
+@pytest.mark.parametrize(
+ "content_type",
+ ["application/pdf", "application/octet-stream", "text/html", "text/plain", None],
+)
+def test_pdf_text_extracted(monkeypatch, content_type):
+ out = _fetch_with(
+ monkeypatch,
+ _pdf_bytes("First page marker", "Second page marker"),
+ content_type,
+ )
+ assert "## Page 1\n\nFirst page marker" in out
+ assert "## Page 2" in out and "Second page marker" in out
+ assert "binary content" not in out and "non-text content" not in out
+
+
+@pytest.mark.parametrize("content_type", ["application/pdf", "text/plain"])
+def test_malformed_pdf_returns_safe_placeholder(monkeypatch, content_type):
+ out = _fetch_with(monkeypatch, b"%PDF-1.7\nnot a complete PDF", content_type)
+ assert out == "(PDF content could not be read as text)"
+
+
+def test_pdf_without_text_layer_reported(monkeypatch):
+ out = _fetch_with(monkeypatch, _pdf_bytes(""), "application/pdf")
+ assert out == "(PDF contains no extractable text)"
+
+
+def test_encrypted_pdf_returns_safe_placeholder(monkeypatch):
+ pymupdf = pytest.importorskip("pymupdf")
+ doc = pymupdf.open()
+ doc.new_page().insert_text((40, 40), "private text")
+ data = doc.tobytes(
+ encryption = pymupdf.PDF_ENCRYPT_AES_256,
+ owner_pw = "owner",
+ user_pw = "secret",
+ )
+ doc.close()
+ out = _fetch_with(monkeypatch, data, "application/pdf")
+ assert out == "(PDF content could not be read as text)"
+
+
+def test_pdf_download_limit_enforced(monkeypatch):
+ monkeypatch.setattr(tools, "_MAX_PDF_FETCH_BYTES", 256)
+ out = _fetch_with(monkeypatch, _pdf_bytes("Readable but oversized"), "application/pdf")
+ assert out == "(PDF content exceeds the download limit; not readable as text)"
+
+
+def test_mislabeled_pdf_is_read_past_text_download_cap(monkeypatch):
+ body = _pdf_bytes("Cross-reference data was fetched")
+ monkeypatch.setattr(tools, "_MAX_FETCH_BYTES", 128)
+ monkeypatch.setattr(tools, "_MAX_PDF_FETCH_BYTES", len(body) + 100)
+ out = _fetch_with(monkeypatch, body, "text/plain")
+ assert "Cross-reference data was fetched" in out
+
+
+def test_pdf_extraction_caps_pages_and_intermediate_text(monkeypatch):
+ from core.rag.parsers import Page
+
+ seen = {}
+
+ def fake_parse(data, *, max_pages = None):
+ seen["max_pages"] = max_pages
+ pages = [Page(text = "x" * 1000, page_number = i, char_count = 1000) for i in range(1, 51)]
+ return pages, 60 # document actually has more pages than the cap
+
+ monkeypatch.setattr("core.rag.parsers.parse_pdf_bytes", fake_parse)
+ text = tools._extract_pdf_text(b"unused")
+ assert seen["max_pages"] == tools._MAX_WEB_PDF_PAGES
+ assert len(text) <= tools._MAX_PAGE_CHARS
+ assert "text limited to 16,000 characters" in text
+ assert "page processing capped at 50 pages" in text
+
+
+def test_pdf_exactly_at_page_cap_not_marked_capped(monkeypatch):
+ from core.rag.parsers import Page
+
+ # Exactly _MAX_WEB_PDF_PAGES pages are fully read, so no "capped" marker.
+ monkeypatch.setattr(
+ "core.rag.parsers.parse_pdf_bytes",
+ lambda data, *, max_pages = None: (
+ [Page(text = "short", page_number = i, char_count = 5) for i in range(1, 51)],
+ 50,
+ ),
+ )
+ text = tools._extract_pdf_text(b"unused")
+ assert "page processing capped" not in text
+ assert "## Page 50\n\nshort" in text
+
+
+def test_pdf_page_cap_does_not_claim_later_pages_are_textless(monkeypatch):
+ from core.rag.parsers import Page
+ monkeypatch.setattr(
+ "core.rag.parsers.parse_pdf_bytes",
+ lambda data, *, max_pages = None: (
+ [Page(text = "", page_number = i, char_count = 0) for i in range(1, 51)],
+ 60,
+ ),
+ )
+ assert tools._extract_pdf_text(b"unused") == (
+ "(PDF contains no extractable text in the first 50 pages)"
+ )
+
+
+def test_pdf_result_discarded_after_fetch_deadline(monkeypatch):
+ clock = {"time": 1000.0}
+ monkeypatch.setattr(tools.time, "monotonic", lambda: clock["time"])
+
+ def slow_extract(data):
+ clock["time"] += 10.0
+ return "late PDF text"
+
+ monkeypatch.setattr(tools, "_extract_pdf_text", slow_extract)
+ out = _fetch_with(monkeypatch, _pdf_bytes("Readable text"), "application/pdf")
+ assert out == "Failed to fetch URL: timed out."
def test_text_octet_stream_kept_after_sniffing(monkeypatch):
@@ -154,7 +275,6 @@ def test_valid_utf8_binary_caught_by_control_chars(monkeypatch):
@pytest.mark.parametrize(
"magic",
[
- b"%PDF-",
b"PK\x03\x04",
b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1",
b"\x1f\x8b",
@@ -180,8 +300,8 @@ def test_text_labeled_binary_caught_by_magic(monkeypatch, magic):
b"\t\xef\xbb\xbf ",
],
)
-def test_pdf_magic_after_harmless_prefix(monkeypatch, prefix):
- body = prefix + b"%PDF-1.7\n" + b"1 0 obj<>endobj\n" * 100
+def test_binary_magic_after_harmless_prefix(monkeypatch, prefix):
+ body = prefix + b"\x1f\x8b" + b" printable text-heavy body" * 100
out = _fetch_with(monkeypatch, body, "text/plain")
assert "binary content" in out
@@ -243,10 +363,10 @@ def test_html_page_unaffected(monkeypatch):
def test_content_type_sanitized_in_message(monkeypatch):
# Do not echo obs-folded header content into the model response.
- out = _fetch_with(monkeypatch, b"\x00\x01\x02" * 500, "application/pdf\r\n data: injected")
+ out = _fetch_with(monkeypatch, b"PK\x03\x04" * 500, "application/zip\r\n data: injected")
assert "\n" not in out and "\r" not in out
assert "injected" not in out
- assert "application/pdf" in out
+ assert "application/zip" in out
@pytest.mark.parametrize(
From 1777aae37ee8ecd6e5ba5b303f4505d975382751 Mon Sep 17 00:00:00 2001
From: Nilay <118994073+NilayYadav@users.noreply.github.com>
Date: Fri, 17 Jul 2026 04:54:05 +0530
Subject: [PATCH 023/289] don't kill live llama-servers when a new Studio
instance starts (#7182)
---
studio/backend/core/inference/llama_cpp.py | 8 ++++
.../test_llama_cpp_wait_for_vram_settle.py | 43 +++++++++++++++++++
2 files changed, 51 insertions(+)
diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py
index dadfdfd38d..3e93a26787 100644
--- a/studio/backend/core/inference/llama_cpp.py
+++ b/studio/backend/core/inference/llama_cpp.py
@@ -8232,6 +8232,11 @@ class LlamaCppBackend:
if not is_ours:
continue
+ # A live parent means a running Studio (or the user's
+ # shell) still owns it -- not an orphan.
+ if LlamaCppBackend._pid_parent_is_alive(proc.info["pid"]):
+ continue
+
proc.kill()
killed += 1
logger.info(
@@ -8284,6 +8289,9 @@ class LlamaCppBackend:
if not owned:
continue
+ if LlamaCppBackend._pid_parent_is_alive(pid):
+ continue
+
try:
os.kill(pid, signal.SIGKILL)
killed += 1
diff --git a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
index 3c21f41701..d0213f6079 100644
--- a/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
+++ b/studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
@@ -373,6 +373,7 @@ def test_kill_orphaned_servers_returns_count():
with (
patch.dict(sys.modules, {"psutil": fake_psutil}),
patch.dict(os.environ, {"LLAMA_SERVER_PATH": fake_path}),
+ patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)),
):
n = LlamaCppBackend._kill_orphaned_servers()
assert n == 1, "only the Studio-owned orphan should be counted"
@@ -384,11 +385,53 @@ def test_kill_orphaned_servers_returns_count():
with (
patch.dict(sys.modules, {"psutil": fake_psutil}),
patch.dict(os.environ, {"LLAMA_SERVER_PATH": fake_path}),
+ patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)),
):
assert LlamaCppBackend._kill_orphaned_servers() == 0
assert killed == []
+def test_kill_orphaned_servers_spares_live_parent():
+ """A Studio-owned llama-server whose parent is still running is not an
+ orphan (a live Studio or the user's shell owns it) and must never be
+ killed; only the true orphan (parent gone) is reaped."""
+ import os
+
+ mypid = os.getpid()
+ fake_path = "/tmp/unsloth-test-llama/llama-server"
+ killed: list[int] = []
+
+ class _FakeProc:
+ def __init__(self, pid, name, exe):
+ self.info = {"pid": pid, "name": name, "exe": exe}
+
+ def kill(self):
+ killed.append(self.info["pid"])
+
+ live_parent = _FakeProc(mypid + 1, "llama-server", fake_path)
+ true_orphan = _FakeProc(mypid + 2, "llama-server", fake_path)
+
+ fake_psutil = _types.ModuleType("psutil")
+ fake_psutil.NoSuchProcess = type("NoSuchProcess", (Exception,), {})
+ fake_psutil.AccessDenied = type("AccessDenied", (Exception,), {})
+ fake_psutil.ZombieProcess = type("ZombieProcess", (Exception,), {})
+ fake_psutil.process_iter = lambda attrs = None: [live_parent, true_orphan]
+
+ with (
+ patch.dict(sys.modules, {"psutil": fake_psutil}),
+ patch.dict(os.environ, {"LLAMA_SERVER_PATH": fake_path}),
+ patch.object(LlamaCppBackend, "_reap_recorded_pid", staticmethod(lambda: 0)),
+ patch.object(
+ LlamaCppBackend,
+ "_pid_parent_is_alive",
+ staticmethod(lambda pid: pid == mypid + 1),
+ ),
+ ):
+ n = LlamaCppBackend._kill_orphaned_servers()
+ assert n == 1, "only the true orphan should be reaped"
+ assert killed == [mypid + 2], "the live-parent server must be spared"
+
+
def test_startup_reaper_arms_settle_timestamp():
"""__init__ arms ``_last_kill_monotonic`` when the startup reaper kills an
orphan (so the first load_model waits for VRAM to settle), and leaves the
From b508c8fe89a60df5995aeafe0c17589a93cecd50 Mon Sep 17 00:00:00 2001
From: Andrew Chen <48723787+chuenchen309@users.noreply.github.com>
Date: Fri, 17 Jul 2026 16:55:34 +0800
Subject: [PATCH 024/289] fix(save):
unsloth_push_to_hub_gguf(save_method="lora") raises NameError (#7193)
* fix(save): unsloth_push_to_hub_gguf(save_method="lora") raises NameError
unsloth_push_to_hub_gguf reads is_main_process at save.py:3181 but never
declares it. Its twin unsloth_save_pretrained_gguf declares it (2783) and
uses it the same way (2839) -- the LoRA branch was copied between the twins,
the parameter it depends on was not. There is no module-level global, so the
name resolves as a global load and the branch raises NameError 100% of the
time.
save_pretrained_gguf(save_method="lora", push_to_hub=True) raises a
ValueError that tells users to "use .push_to_hub_gguf(save_method='lora')
instead" -- the documented escape hatch is the broken call.
Add is_main_process to the signature, positioned as in the twin, and forward
it to unsloth_save_pretrained_gguf on the merged path so the parameter is not
silently ignored there. Default stays True, so nothing changes for existing
callers.
Co-Authored-By: Claude Opus 4.8
* fix(save): preserve GGUF push compatibility
---------
Co-authored-by: Claude Opus 4.8
Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com>
---
tests/saving/test_export_dispatch.py | 76 ++++++++++++++++++++++++++++
unsloth/save.py | 6 ++-
2 files changed, 80 insertions(+), 2 deletions(-)
diff --git a/tests/saving/test_export_dispatch.py b/tests/saving/test_export_dispatch.py
index 3870d6269d..ad6d51e0b7 100644
--- a/tests/saving/test_export_dispatch.py
+++ b/tests/saving/test_export_dispatch.py
@@ -8,6 +8,8 @@ regressions that pure AST checks cannot (e.g. wrong scheme/suffix/outtype passed
from __future__ import annotations
+import inspect
+
import pytest
import unsloth.save as save_mod
@@ -126,6 +128,80 @@ def test_gguf_lora_push_to_hub_is_rejected(tmp_path):
)
+# The above rejection points users at push_to_hub_gguf(save_method='lora'), so that path
+# has to work; it is only ever exercised here.
+
+
+def test_push_to_hub_gguf_lora_dispatches(monkeypatch):
+ seen = {}
+ monkeypatch.setattr(
+ save_mod,
+ "_unsloth_save_lora_gguf",
+ lambda model, tok, sd, **kw: seen.update(kw),
+ )
+ save_mod.unsloth_push_to_hub_gguf(
+ _FakeModel(),
+ "repo/id",
+ tokenizer = object(),
+ save_method = "lora",
+ quantization_method = "q8_0",
+ )
+ assert seen.get("outtype") == "q8_0"
+ assert seen.get("push_to_hub") is True
+
+
+def test_push_to_hub_gguf_lora_skips_non_main_process(monkeypatch):
+ calls = []
+ monkeypatch.setattr(
+ save_mod,
+ "_unsloth_save_lora_gguf",
+ lambda *a, **kw: calls.append(kw),
+ )
+ result = save_mod.unsloth_push_to_hub_gguf(
+ _FakeModel(),
+ "repo/id",
+ tokenizer = object(),
+ save_method = "lora",
+ is_main_process = False,
+ )
+ assert result is None
+ assert calls == []
+
+
+def test_push_to_hub_gguf_skips_non_main_process_before_merged_conversion(monkeypatch):
+ calls = []
+ monkeypatch.setattr(
+ save_mod,
+ "unsloth_save_pretrained_gguf",
+ lambda **kw: calls.append(kw),
+ )
+ result = save_mod.unsloth_push_to_hub_gguf(
+ _FakeModel(),
+ "repo/id",
+ tokenizer = object(),
+ is_main_process = False,
+ )
+ assert result is None
+ assert calls == []
+
+
+def test_push_to_hub_gguf_preserves_positional_max_shard_size():
+ bound = inspect.signature(save_mod.unsloth_push_to_hub_gguf).bind(
+ _FakeModel(),
+ "repo/id",
+ object(),
+ "q4_k_m",
+ None,
+ None,
+ None,
+ None,
+ "token",
+ "50GB",
+ )
+ assert bound.arguments["max_shard_size"] == "50GB"
+ assert "is_main_process" not in bound.arguments
+
+
# -- torchao PTQ / QAT dispatch ------------------------------------------------------------
diff --git a/unsloth/save.py b/unsloth/save.py
index 50a180f823..7d5774aa97 100644
--- a/unsloth/save.py
+++ b/unsloth/save.py
@@ -3143,6 +3143,7 @@ def unsloth_push_to_hub_gguf(
datasets: Optional[List[str]] = None,
save_method: str = None,
imatrix_file = None,
+ is_main_process: bool = True,
):
"""
Same as .push_to_hub(...) except 4bit weights are auto
@@ -3175,11 +3176,11 @@ def unsloth_push_to_hub_gguf(
"""
if tokenizer is None:
raise ValueError("Unsloth: Saving to GGUF must have a tokenizer.")
+ if not is_main_process:
+ return None
# save_method="lora" exports the adapter itself as a GGUF LoRA (not a merged model).
if save_method is not None and str(save_method).lower() == "lora":
- if not is_main_process:
- return None # only the main rank converts and uploads, like the local lora branch
_qm = quantization_method
if isinstance(_qm, (list, tuple)) and len(_qm) == 1:
_qm = _qm[0] # the gguf API allows a list; unwrap a single outtype
@@ -3233,6 +3234,7 @@ def unsloth_push_to_hub_gguf(
first_conversion = first_conversion,
push_to_hub = False, # Never push from here
token = token, # forwarded so imatrix_file=True can read a gated/private upstream
+ is_main_process = is_main_process,
max_shard_size = max_shard_size,
safe_serialization = safe_serialization,
temporary_location = temporary_location,
From 8cbdfbe355a83b6cc0706e2ed8ec1c737b71c3f3 Mon Sep 17 00:00:00 2001
From: Eyera
Date: Fri, 17 Jul 2026 15:08:01 +0200
Subject: [PATCH 025/289] Feat/model picker per model config (#6647)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
* refactor(studio): move chat model picker into features/model-picker
Relocate model-selector + its support files from components/assistant-ui
into a self-contained features/model-picker feature (own barrel), mirroring
the modular Hub layout. Pure move + import repoint; no behaviour change.
* feat(model-picker): add per-model config persistence layer
Superset PerModelConfig (customContextLength, kvCacheDtype, speculativeType,
specDraftNMax, tensorParallel, chatTemplateOverride, trustRemoteCode) persisted
to localStorage (unsloth_model_configs) with schema versioning + LRU budget.
KV-dtype and speculative value sets match main's sidebar (no q4_0/ngram-simple).
Reuses features/hub/lib/model-identity for normalization; adds storage-key layer
and applyPerModelConfigToRuntime (sets tensorParallel, which the old PR omitted).
* feat(picker): modular backend for chat-template validate + default fetch
New studio/backend/picker package (schemas/service/routes) mounted at /api/picker:
- POST /api/picker/validate-chat-template (Jinja syntax validation, no false positives)
- GET /api/picker/chat-template/{model_name} (default template from tokenizer_config.json,
reusing get_cache_path/resolve_cached_repo_id_case; graceful null, no model-code exec)
Frontend api/templates.ts client + hooks/use-model-defaults lazy cache. No backend
changes to the existing inference load route (per-model load fields already supported).
* feat(model-picker): bind picker on-device list to shared hub inventory
Picker now sources cached + local models from useHubInventory (the Hub's shared
store) via a thin adapter, replacing its own /api/models/* fetchers + module
caches. Hub, download manager, and picker now share one source of truth, so
completed downloads reflect in the picker automatically. Partial/live-download
rows are filtered from the cached lists (unchanged rendering). Local naming/search
preserved via additive LocalInventoryRow modelId/displayName. Variant expander,
scan-folder management, recommended-fit, search, external providers untouched.
Known minor: cached 'Downloaded date' sort tiebreak degrades to alphabetical
(hub cached rows carry no mtime); default 'recent' (load-time) sort preserved.
* feat(model-picker): per-model config step inside the picker
Picking a (non-external) model now opens an in-picker config view built from
main's current load controls (context length, KV cache dtype, speculative
decoding, draft tokens, tensor parallel) plus a chat-template editor backed by
the picker validate/default endpoints. 'Remember for this model' persists the
config per model+variant; Run forwards the config to the existing load flow via
meta.config. External models bypass the step. Two-view orchestration lives in
model-selector (single interception point); pickers.tsx call sites untouched.
trustRemoteCode dropped from PerModelConfig to preserve main's per-load consent.
* feat(chat): apply/persist per-model config through the load flow
handleCheckpointChange threads meta.config into the selection; stageOrLoad and
the autoload/Hub-run paths now apply the picker config (explicit pick or saved
remembered config) via applyPerModelConfigToRuntime before staging/loading, with
keepSpeculative set so a remembered speculative mode survives the model switch.
Replaces the old remembered-load-settings seeding (resolveInitialConfig now the
single source). SelectedModelInput carries config.
* refactor(chat): remove per-model load config from the right sidebar
The load knobs (context, KV cache, speculative, draft tokens, tensor parallel)
and the chat-template editor now live only in the picker config step. The sheet's
Model section keeps the staged Load/Cancel flow (config is applied at pick time);
sampling params, system prompt, and RAG are unchanged. Deletes the superseded
remembered-load-settings module + the store's applyRememberedLoadSettings action,
removes the now-dead sheet state/imports, and points the settings reset at
unsloth_model_configs. Delete-cleanup deferred (stale config is LRU-capped).
* fix(model-picker): remove leftover sidebar-staging cogwheel + empty Model section
The downloaded-variant gear (ModelLoadSettingsAction) staged a model straight
into the right-sidebar Run-settings flow -- the old 'configure before load' path
now fully replaced by the in-picker config step. Removed the gear + its component.
Also gate the sheet's 'Model' section to staged picks only (pendingSelection):
after the load-knob strip its content is staged-only, so it was rendering an
empty section header whenever a model was merely loaded.
* chore(chat): remove dead per-model-config setters + modelControlsDisabled
After the load-config UI moved into the picker, the store's per-model setters
(setKvCacheDtype/setSpeculativeType/setSpecDraftNMax/setTensorParallel/
setCustomContextLength/setChatTemplateOverride) had zero callers
(applyPerModelConfigToRuntime writes via setState), and the sheet's
modelControlsDisabled was unreferenced. Verified dead across the whole tree.
* fix(chat): config-step Load actually loads (ignore Load-on-selection)
Root cause: with Settings > Chat > 'Load on selection' turned OFF, the config
step's load went down the deferred-staging path -- opening the right sidebar with
' is staged, not loaded yet / Choose Load model'. The in-picker config step
IS the deliberate load action, so its Load now loads immediately (or downloads +
auto-loads when not cached) regardless of the toggle. Renamed the button
'Run model' -> 'Load model' to match. Native/dropped picks still honor the toggle.
* refactor(chat,hub): retire 'Load on selection' — config step is the only load flow
The in-picker config step (and the Hub Run button) now fully supersede the old
stage-to-sidebar flow, so the Load-on-selection toggle is removed everywhere:
- chat stageOrLoad: every pick loads immediately, or downloads + auto-loads when
not cached (the previous default behaviour, now universal).
- hub Run: drops the stage branch; downloaded GGUFs load directly with their saved
per-model config (no collision with the chat config step — both end at selectModel).
- store: removed loadOnSelection field/setter/key/default; Settings>Chat toggle and
its settings-reset entry removed.
- staged sidebar section is now a download-progress view (auto-loads on completion).
No manual staging remains; stageModel is used only for background auto-load downloads.
* feat(model-picker): default chat template from GGUF + thread variant through config flow
Read the embedded tokenizer.chat_template from GGUF files (read_gguf_chat_template
in gguf_metadata) and use it as the per-model default. Plumb gguf_variant through
the picker service, /api/picker/chat-template route, frontend templates API, and
use-model-defaults so the right variant's template is fetched.
Also refine the picker config-page/model-selector wiring, drop the dead
ggufNativeContextLength runtime path, and add the per-model-config storage keys to
the settings prefs export.
* feat(model-picker): read safetensors chat template + hide editor where it has no effect
Resolve the default chat template for safetensors models: prefer the modern
chat_template.jinja, fall back to the tokenizer_config.json chat_template field,
then chat_template.json (multimodal processor), then the GGUF embedded template.
Applied to local dirs, the HF cache snapshot scan, and the HF remote fetch.
Hide the chat-template editor in the picker for safetensors models — the override
is only applied at load by the GGUF/llama.cpp backend, so editing it on safetensors
currently has no effect. GGUF keeps the editor. Nothing removed; the dialog stays
for when the safetensors apply path is wired up in a later branch.
* fix(model-picker): set legacy-migration flag only after the write succeeds
Set unsloth_model_configs_migrated only once writeMap confirms the migrated
map persisted, so a quota/storage failure no longer marks migration done and
silently drops the user's pre-existing remembered settings — the next load retries.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* MVP model picker fixes
* MVP picker config fix
* MVP safetensors config
* MVP max seq config
* MVP max seq fix
* Fix static max tokens cap ignoring model context
* Fix picker GGUF scan parity
* fix(studio): harden model picker config loading
Apply remembered per-model configs consistently from picker and Hub loads, keep default configs from overriding standing speculative settings, add config access for direct local GGUF files, and support saving or forgetting active model settings without a reload.
* Fix model picker config flow
* Fix model picker config loads
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Avoid recursive per-model config migration reads
* Apply the displayed context length when loading a GGUF
* Fix template validation, cached template lookup, and failed load rollback
- Validate chat templates with the loopcontrols extension so templates
that use break or continue tags pass the picker validator, matching the
inference renderer that already accepts them.
- Read the default chat template from the newest cache snapshot rather than
an arbitrary iterdir order, so an older cached revision no longer prefills
a stale template.
- Capture the runtime per-model config before a load and reapply it when the
load fails, so a failed switch leaves the active model context, KV cache,
template, and speculative settings as they were.
* Make chat template view only for safetensors models
Custom chat template overrides are applied at inference only for GGUF
models, which pass the template to llama-server. The safetensors backend
renders with the model built-in template and ignores the override, so
editing it would save a value that never loads. For safetensors the
config page now opens the template as a read-only preview with a note
that editing is not available yet. This can become editable once
inference support for custom safetensors templates lands in main.
* Fix model picker config edge cases
- Restore prior runtime config when a load no-ops for the active model
- Cap the picker validator request body via the protected prefixes
- Keep the GGUF context slider max above the loaded context
- Fetch subfolder chat templates for uncached Hub repos
- Show the compare side config when reopening the picker
* Keep saved GGUF context above the fallback ceiling
* Show the model config in the run settings sidebar
* Fix model config sidebar reset and context slider
- Stack the remember toggle and action buttons in the sidebar
- Reset the config to defaults instead of the loaded values
- Fetch the native context so the slider max is not the loaded value
* Fix model picker config and download regressions
- Run picker chat template routes off the event loop
- Depth and root guard local template directory scans
- Restore download manager flow for uncached hub picks
- Apply per model context length on reload
- Import model picker symbols from the feature barrel
* Fix model picker config and cached download sorting
- Restore load settings when a Hub run is rejected mid load
- Reuse one NumericValueInput instead of a duplicate copy
- Fix double decode of the model name in the template route
- Remove the unused reset-to-loaded settings action
- Fix cached model download sorting
* Fix model picker per-model config edge cases
Honor a saved or typed max seq length above the model's native context so
RoPE extended values are no longer clamped and silently overwritten. Allow
typing past native while the slider keeps native as a soft ceiling.
Guard the fetch success paths in use-model-defaults against an aborted
signal, and refetch when the HF token changes.
Hash the chat template content in the sidebar remount key instead of its
length. Enable reset for a GGUF whose native context is unknown, and floor
the context slider max so it can never fall below the min.
* Fix GGUF context auto-fit and gated model config token
Stop forcing a 32768 context when a GGUF native context is unknown so the backend auto-fits to VRAM again, while still honoring an explicit context edit.
Send the HF token as a query param so gated safetensors models resolve their max position embeddings.
Derive model default state during render to drop the set-state-in-effect calls.
* Fix native GGUF context ceiling and guard picker template reads
Restore the native context store field so the sidebar slider keeps the
full ceiling for drag and drop GGUFs. Limit local chat template reads to
the browse allowlist, skip malformed repo ids, and drop unused model
picker exports.
* Fix model picker lint boundaries
* Fix model picker review findings
Chat template editor never seeded its draft. Radix only calls onOpenChange
from internal events, so the seed in the nextOpen branch was dead and a model
with a saved override opened empty. Saving then cleared the override. Drop the
dead branch, treat draft as an untouched sentinel, and reset it on every close.
Uncached Hub picks could auto load a model after the user left the chat. Main
detached the staged pick on route exit and on chat context change. Carry the
context key on the pending pick and skip the load when it no longer matches.
Also clear configTarget when the picker closes, restore the onUpdated ref so
variant rows stop resubscribing on every parent render, skip the LRU write when
the entry is already most recent, import NumericValueInput relatively, and drop
the unused ModelUpdateAction barrel export.
* Preserve GGUF context on active reload
* Fix model picker per-model config regressions
- Stop reloading the already loaded model on re-pick
- Hide infra models from the chat picker
- Detect vision support on cached GGUF repos
- Honor saved maxSeqLength on auto load
- Restore default chat template for local GGUFs
- Warn on save failure and revert config on cancel
- Refetch picker inventory on open
- Persist read only per model config safely
* Fix stale model auto load
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix model picker numeric input sizing and constraints
Size value inputs to their content so long context lengths are not clipped,
restrict them to numeric characters, and stop the speculative decoding label
from truncating in the sidebar.
* Fix picker CI tests and harden chat template resolution for PR #6647
- tests: point the descender guard at the moved model-selector.tsx path
- tests: exclude the disabled Reload model button from the regenerate locator so .first targets the real Regenerate
- picker/service.py: reject symlinked template/gguf leaves that resolve outside the browse allowlist (HF cache reads unchanged)
- compare mode: resolve each pane's own remembered chat template instead of inheriting the other pane's from the store
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Protect future-schema per-model configs from deletion for PR #6647
savePerModelConfig already refuses to overwrite a stored config whose schema version is newer than this client understands, but deletePerModelConfig did not. Unchecking Remember on an older client therefore silently destroyed a newer client's saved config. Apply the same guard on delete and surface the blocked case through the existing saveFailed toast.
* Protect future-schema per-model configs from quota eviction for PR #6647
The save and delete guards already refuse to touch a stored config whose schema version is newer than this client understands, but the quota-eviction path did not, so a full store on an older client could still evict a newer client's config. Skip future-schema entries when evicting and fail the save if the budget cannot be met without them.
* Fix GGUF context persistence, compare context, and rollback settings for PR #6647
Persist a GGUF context override from the user's intent instead of collapsing it against the loaded context, which reintroduced the context-reset (f4838782cb reverted the native-baseline fix). model-config-page now collapses the saved value against native, and use-chat-model-runtime and chat-adapter retain the requested context on load so re-saving another setting keeps the override; a null request stays null so a VRAM auto-fit never becomes a stored override.
shared-composer: a compare pane with no explicit GGUF context now loads at native (0) like single-view, not the session maxSeqLength that silently shrank the shown context.
use-chat-model-runtime: restore the previous model's KV cache dtype and chat template on a failed-load rollback so it runs as it was, not with backend defaults.
* Preserve native path token when reloading the active model for PR #6647
handleReloadActiveModel rebuilt the selection without the store's activeNativePathToken, so reloading a file-picked GGUF after a settings change validated the display label as a repo/path and failed. Thread the active native token through the reload selection so native-loaded models reopen correctly.
* Make picker template validation resilient and accept HF generation tags for PR #6647
Import Jinja lazily inside validate_chat_template so a backend without the optional jinja2 package (GGUF-only installs) still starts instead of raising ModuleNotFoundError at import time. Register a no-op extension for the Transformers {% generation %} assistant-mask tag so pasting a valid HF chat template validates, matching the renderer, rather than being rejected as an unknown tag.
* Honor remembered compare config and parse processor chat_template.json for PR #6647
* Fix failed-load rollback context and processor template map fallback for PR #6647
* Restore speculative decoding config on failed-switch rollback
When a model switch fails after the previous model was unloaded, the
rollback reload restored tensor_parallel, KV cache dtype and the chat
template override, but omitted speculative_type and spec_draft_n_max and
cleared their loaded shadows to null. The previous model therefore came
back running at backend defaults (speculation off) while the UI still
showed it enabled, and the status resync confirmed the off state. Resend
the previous model's speculative settings in the rollback load and keep
the store's active and loaded speculative fields in sync with them.
* Reset max sequence length when a model has no saved config
applyPerModelConfigToRuntime reset every per-model field except
maxSeqLength, which it only wrote when the incoming config had one.
maxSeqLength is the sole field carried on store.params, so selecting a
model with no remembered config left the previous model's value in place
and later loaded the new model at that leaked length. Fall back to the
standing default so an unremembered model loads at its own default.
* Surface a message when a variant update cannot start
startManagedUpdate handled the conflict and error start outcomes but let
busy fall through as if the update began, so the confirm dialog closed
with no job created and the cached variant stayed stale. Show an info
message when the repo is busy with a sibling transfer so the click is
not silently dropped.
* Keep per-model speculative choices out of the global default
A staged load with a per-model or one-off config sets keepSpeculative,
which already skips reading the global speculative preference. The
matching save still ran unconditionally, so the model-specific choice was
written to the global unsloth_chat_speculative_type and a later model with
no saved config started from it instead of Auto. Skip saveSpeculativeType
when keepSpeculative so the per-model choice stays isolated.
* Seed non-active model settings from the app default max length
The Run settings page captured initialMaxSeqLength from the loaded
model's runtime params and fell back to it for a model with no saved
config. Opening settings for a different, unloaded model and clicking
Load then sent the active model's context (for example 64k) instead of
the 4096 default, risking validation failures or OOMs. Seed the default
for non-active models and keep the runtime value only for the active one.
* Prefer sidecar tokenizer chat template over the GGUF copy for variants
_chat_template_from_dir returned the embedded GGUF template first when a
variant was selected, reversing the tokenizer-first precedence of the
no-variant path. A model whose chat_template.jinja or tokenizer_config.json
supersedes a stale embedded template then got the wrong template on
variant selection. Keep tokenizer files first regardless of variant; the
variant only picks which GGUF is the fallback. Adds regression tests for
both the tokenizer-wins and gguf-fallback cases.
* Keep per-model speculative choices load-local in autoload and compare
The interactive load path treats a per-model speculative choice as
load-local and skips writing it to the global default. Autoload and
generalized compare still called saveSpeculativeType unconditionally, so a
remembered off or ngram setting leaked into unsloth_chat_speculative_type
and later models with no saved config inherited it. Persist the global
preference only when the value came from the global settings.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: record the compare pane's loaded context in runtime state so the active model's settings and any reload or save use it, not the previous context
* Studio: notify the user when a Hub autoload can't start because another download for the model is already running, instead of silently dropping it
* Studio: drop the merge's orphaned staged-model store helpers and unused alert imports
The main merge left isPendingGguf and pendingSelectionMatches referencing the
removed PendingModelSelection type, and the alert-dialog/alert imports unused
after the permission-mode dropdown replaced the bypass dialog, so tsc -b failed.
* Studio: cache a null default chat template so the viewer stops re-fetching it
A model with no sidecar or embedded template resolves to a terminal null, but
that result was never cached, so reopening the template viewer re-ran the
backend and Hugging Face lookup every time.
* Studio: detect direct-file GGUFs in run settings so Max Tokens uses their context
A GGUF loaded from a local file or custom folder has no variant label, so the
run-settings panel treated it as non-GGUF and clamped Max Tokens to the session
max_seq_length instead of the loaded GGUF context. Detect it via the reported
GGUF context and the .gguf checkpoint suffix, matching the chat page.
* Studio: prompt to re-select a local model file when its lease expired before reload
A file-picked GGUF is reachable only through a native path token that the
desktop host prunes after a TTL. Reloading reused that token blindly, so a
reload long after the initial load failed with an opaque error. Track the
token's expiry and, when it has passed, ask the user to re-select the file
instead of attempting a doomed reload.
* Fix descender-clipping test to tolerate sidebar layout utilities
The sidebar account-block div carries layout utilities (min-w-0, flex-1)
between 'flex' and 'flex-col', so the descender-clipping guard's regex,
which required 'flex' immediately followed by 'flex-col', no longer matched
and the test failed to locate the account-block div. Generalize the prefix
to allow intervening flex utilities while still capturing the leading-*
class before the collapsible visibility utility and asserting leading-tight,
so the guard against clipped glyph descenders is fully preserved.
* Harden picker chat-template resolution
Enforce the 64 KiB chat-template contract at the validate endpoint's request
model so a direct caller cannot submit a template far larger than the frontend
allows (MaxBodyMiddleware only bounds the whole request body, not this field);
oversized templates now return a clean 422.
Apply sidecar-over-GGUF template precedence globally across cached snapshots
instead of per snapshot. A repo with multiple cached revisions previously
returned the first snapshot's template, so a newer GGUF-only revision could
win over an older revision's maintained chat_template.jinja sidecar, which
contradicted the documented intent that sidecars supersede the embedded copy.
* Guard per-model config against future-schema and lossy migration
Two forward-compatibility gaps in the versioned per-model config store:
- The load/apply path returned and normalized a stored record without checking
its schema version, so a record written by a newer client was reinterpreted
under the current schema and applied to a live model load, even though save,
delete and eviction all refuse to touch future-schema records. Reject
future-schema records on load too.
- The one-time legacy migration enforced the storage budget without protecting
the entries it had just migrated and set the completion flag unconditionally.
When storage was already full of future-schema records (which are unevictable
by an older client), the migrated entries were the only evictable ones and
could be dropped while migration was still marked complete. Protect the
migrated keys during eviction and only mark migration complete when they
survive, so it retries once space frees up.
* Discard chat-template validation results after the dialog closes
Server-side template validation is async, but closing or cancelling the editor
did not abort it, so a late-arriving valid response still called onSave and
applied a template the user had already dismissed. Track a validation token
that is bumped on close and ignore any validation result whose token is stale.
* Record native lease expiry when loading a picked GGUF from the chip
The pending-native-model chip loaded via stageOrLoad directly, bypassing
loadNativeModelIntent, so activeNativePathExpiresAtMs was never recorded for a
chip-loaded file. A later reload then either skipped the lease-expiry guard
entirely (expiry left null) or compared against a previously loaded file's
stale expiry, so reload could reuse an already-pruned token or wrongly block a
still-valid one. Route the chip through loadNativeModelIntent, which builds the
same selection and records the expiry.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Prefer sidecar template for a directly selected local GGUF file
A direct .gguf file path read its embedded chat template without checking the
parent directory for a maintained sidecar (chat_template.jinja /
tokenizer_config.json), while directory and variant selections already prefer
the sidecar. That let the config editor preview or save a stale embedded
template for the same model depending on how it was selected. Check the parent
directory sidecars first, then fall back to the embedded copy, and cover both
paths with tests.
* Resolve cached chat template per revision, newest first
The earlier change searched every cached snapshot for a sidecar before
considering any snapshot's embedded GGUF template, which let an obsolete sidecar
from an older revision override the newest revision's template. Restore
per-snapshot resolution (newest first): a revision's sidecar still supersedes
its own embedded GGUF copy, but a newer revision is no longer overridden by an
older revision's sidecar.
* Preserve autoload transport conflicts and surface background busy downloads
- When a Hub autoload hits a transport conflict, keep pendingHubAutoLoad bound
instead of clearing it. Clearing it re-keyed the download surface and its
cleanup cancelled the conflict the toast tells the user to resolve, so the
Hub resume affordance was gone the moment it appeared. Return early on
conflict, mirroring the started branch, so resolving it from the Hub still
auto-loads on completion.
- The background-download branch handled started and conflict but silently
dropped a busy outcome, leaving the user with no feedback when a peer variant
of the same repo was already downloading. Surface the same busy toast the
autoload path uses.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: danielhanchen
---
studio/backend/hub/schemas/inventory.py | 1 +
.../hub/services/models/cache_inventory.py | 74 +-
studio/backend/main.py | 3 +
studio/backend/picker/__init__.py | 2 +
studio/backend/picker/routes/__init__.py | 6 +
studio/backend/picker/routes/templates.py | 42 +
studio/backend/picker/schemas.py | 32 +
studio/backend/picker/service.py | 361 +++++++
.../tests/test_model_update_robustness.py | 46 +
studio/backend/tests/test_picker_service.py | 162 +++
studio/backend/utils/models/gguf_metadata.py | 79 ++
studio/frontend/src/app/routes/__root.tsx | 7 -
.../remembered-load-settings.ts | 69 --
.../src/features/chat/api/chat-adapter.ts | 50 +-
.../frontend/src/features/chat/chat-page.tsx | 517 ++++++----
.../src/features/chat/chat-settings-sheet.tsx | 947 +++---------------
.../chat/hooks/use-chat-model-runtime.ts | 105 +-
.../hooks/use-staged-model-preparation.ts | 155 ---
studio/frontend/src/features/chat/index.ts | 18 +
.../lib/apply-inference-status-to-store.ts | 5 -
.../src/features/chat/shared-composer.tsx | 111 +-
.../chat/stores/chat-runtime-store.ts | 186 +---
.../export/components/export-run-panel.tsx | 71 +-
.../hub/catalog/models-catalog-rows.tsx | 41 +-
.../hub/catalog/on-device-folders-dialog.tsx | 33 +-
.../download-manager-controller.ts | 23 -
.../features/hub/download-manager/index.ts | 1 -
studio/frontend/src/features/hub/hub-page.tsx | 129 +--
studio/frontend/src/features/hub/index.ts | 58 +-
.../src/features/hub/inventory/api.ts | 2 +
.../src/features/hub/inventory/types.ts | 3 +
.../src/features/hub/inventory/view-models.ts | 9 +
.../model-picker/api/model-metadata.ts | 20 +
.../features/model-picker/api/templates.ts | 51 +
.../chat-template-editor-dialog.tsx | 191 ++++
.../components/model-config-page.tsx | 742 ++++++++++++++
.../components}/model-selector.tsx | 130 ++-
.../model-selector/folder-browser.tsx | 79 +-
.../model-selector/model-capabilities.ts | 0
.../model-selector/model-delete-action.tsx | 9 +-
.../model-load-settings-action.tsx | 19 +-
.../model-selector/model-update-action.tsx | 25 +-
.../components}/model-selector/model-usage.ts | 3 +-
.../components}/model-selector/pickers.tsx | 645 +++++++-----
.../components}/model-selector/pill-tabs.tsx | 3 +-
.../model-selector/recommended-fit.ts | 0
.../components}/model-selector/row-meta.ts | 0
.../components}/model-selector/source-tabs.ts | 0
.../components}/model-selector/types.ts | 13 +
.../components/numeric-value-input.tsx | 113 +++
.../components/sidebar-model-config.tsx | 89 ++
.../model-picker/hooks/use-model-defaults.ts | 176 ++++
.../src/features/model-picker/index.ts | 30 +
.../inventory/use-chat-picker-inventory.ts | 118 +++
.../model-config/apply-per-model-config.ts | 85 ++
.../model-config/model-identity.ts | 69 ++
.../model-config/per-model-config.ts | 571 +++++++++++
.../src/features/settings/tabs/chat-tab.tsx | 39 -
.../features/settings/tabs/general-tab.tsx | 3 +-
.../frontend/src/features/training/index.ts | 4 +-
tests/studio/playwright_chat_ui.py | 14 +-
.../test_studio_text_descender_clipping.py | 17 +-
62 files changed, 4478 insertions(+), 2128 deletions(-)
create mode 100644 studio/backend/picker/__init__.py
create mode 100644 studio/backend/picker/routes/__init__.py
create mode 100644 studio/backend/picker/routes/templates.py
create mode 100644 studio/backend/picker/schemas.py
create mode 100644 studio/backend/picker/service.py
create mode 100644 studio/backend/tests/test_picker_service.py
delete mode 100644 studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts
delete mode 100644 studio/frontend/src/features/chat/hooks/use-staged-model-preparation.ts
create mode 100644 studio/frontend/src/features/model-picker/api/model-metadata.ts
create mode 100644 studio/frontend/src/features/model-picker/api/templates.ts
create mode 100644 studio/frontend/src/features/model-picker/components/chat-template-editor-dialog.tsx
create mode 100644 studio/frontend/src/features/model-picker/components/model-config-page.tsx
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector.tsx (86%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/folder-browser.tsx (88%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/model-capabilities.ts (100%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/model-delete-action.tsx (90%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/model-load-settings-action.tsx (66%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/model-update-action.tsx (82%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/model-usage.ts (93%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/pickers.tsx (89%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/pill-tabs.tsx (98%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/recommended-fit.ts (100%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/row-meta.ts (100%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/source-tabs.ts (100%)
rename studio/frontend/src/{components/assistant-ui => features/model-picker/components}/model-selector/types.ts (76%)
create mode 100644 studio/frontend/src/features/model-picker/components/numeric-value-input.tsx
create mode 100644 studio/frontend/src/features/model-picker/components/sidebar-model-config.tsx
create mode 100644 studio/frontend/src/features/model-picker/hooks/use-model-defaults.ts
create mode 100644 studio/frontend/src/features/model-picker/index.ts
create mode 100644 studio/frontend/src/features/model-picker/inventory/use-chat-picker-inventory.ts
create mode 100644 studio/frontend/src/features/model-picker/model-config/apply-per-model-config.ts
create mode 100644 studio/frontend/src/features/model-picker/model-config/model-identity.ts
create mode 100644 studio/frontend/src/features/model-picker/model-config/per-model-config.ts
diff --git a/studio/backend/hub/schemas/inventory.py b/studio/backend/hub/schemas/inventory.py
index ef95efe2f2..f81c9a3498 100644
--- a/studio/backend/hub/schemas/inventory.py
+++ b/studio/backend/hub/schemas/inventory.py
@@ -160,6 +160,7 @@ class CachedRepoBase(BaseModel):
repo_id: str
size_bytes: int = 0
cache_path: Optional[str] = None
+ last_modified: Optional[float] = None
partial: bool = False
partial_transport: Optional[str] = None
inventory_id: Optional[str] = None
diff --git a/studio/backend/hub/services/models/cache_inventory.py b/studio/backend/hub/services/models/cache_inventory.py
index 1f38af9381..ba3a266bfb 100644
--- a/studio/backend/hub/services/models/cache_inventory.py
+++ b/studio/backend/hub/services/models/cache_inventory.py
@@ -31,6 +31,7 @@ from hub.services.models.common import (
_is_checkpoint_weight_name,
_is_gguf_filename,
_is_main_gguf_filename,
+ _is_mmproj_filename,
_is_transformers_safetensors_weight_name,
_local_inventory_id,
_prefer_complete_larger,
@@ -125,6 +126,34 @@ def _repo_has_gguf_files(repo_info) -> bool:
return _repo_gguf_size_bytes(repo_info) > 0
+def _blob_mtime(file_obj) -> float:
+ ts = getattr(file_obj, "blob_last_modified", None)
+ if isinstance(ts, (int, float)) and ts > 0:
+ return float(ts)
+ blob_path = getattr(file_obj, "blob_path", None)
+ if blob_path:
+ try:
+ return float(Path(blob_path).stat().st_mtime)
+ except OSError:
+ pass
+ return 0.0
+
+
+def _repo_gguf_last_modified(repo_info) -> float:
+ latest = 0.0
+ for revision in repo_info.revisions:
+ for f in revision.files:
+ if _is_main_gguf_filename(f.file_name):
+ latest = max(latest, _blob_mtime(f))
+ return latest
+
+
+def _repo_has_mmproj(repo_info) -> bool:
+ return any(
+ _is_mmproj_filename(f.file_name) for revision in repo_info.revisions for f in revision.files
+ )
+
+
def _cached_repo_file_name(file_obj) -> str:
file_path = getattr(file_obj, "file_path", None)
if file_path:
@@ -266,6 +295,7 @@ def _scan_cached_gguf() -> list[dict]:
continue
key = repo_id.lower()
existing = seen_lower.get(key)
+ last_modified = _repo_gguf_last_modified(repo_info)
row = {
"repo_id": repo_id,
"size_bytes": max(total_size, variant_state_size),
@@ -275,6 +305,9 @@ def _scan_cached_gguf() -> list[dict]:
# per-variant detail lives on GgufVariantDetail.
"partial_transport": None,
}
+ last_modified = max(last_modified, (existing or {}).get("last_modified", 0.0))
+ if last_modified > 0:
+ row["last_modified"] = last_modified
row.update(
_cache_inventory_fields(
repo_id,
@@ -283,8 +316,12 @@ def _scan_cached_gguf() -> list[dict]:
requires_variant = True,
)
)
+ if _repo_has_mmproj(repo_info):
+ row["capabilities"]["supports_vision"] = True
if _prefer_cache_row(row, existing):
seen_lower[key] = row
+ elif last_modified > existing.get("last_modified", 0.0):
+ existing["last_modified"] = last_modified
except Exception as e:
repo_label = getattr(repo_info, "repo_id", "")
logger.warning(f"Skipping cached GGUF repo {repo_label}: {e}")
@@ -312,13 +349,14 @@ class _CachedNonGgufPayload(NamedTuple):
size_bytes: int
has_runnable_weights: bool
model_format: ModelFormat
+ last_modified: float
def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload:
- all_weight_blobs: dict[str, int] = {}
- adapter_blobs: dict[str, int] = {}
- safetensors_blobs: dict[str, int] = {}
- checkpoint_blobs: dict[str, int] = {}
+ all_weight_blobs: dict[str, tuple[int, float]] = {}
+ adapter_blobs: dict[str, tuple[int, float]] = {}
+ safetensors_blobs: dict[str, tuple[int, float]] = {}
+ checkpoint_blobs: dict[str, tuple[int, float]] = {}
has_config = False
has_adapter_config = False
has_adapter_weights = False
@@ -326,12 +364,15 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload:
has_transformers_safetensors = False
has_checkpoint = False
- def _record_blob(target: dict[str, int], file_obj, rev_id: str, file_name: str) -> None:
+ def _record_blob(
+ target: dict[str, tuple[int, float]], file_obj, rev_id: str, file_name: str
+ ) -> None:
blob_path = getattr(file_obj, "blob_path", None)
size = int(file_obj.size_on_disk or 0)
key = str(blob_path) if blob_path else f"{rev_id}:{file_name}"
- target[key] = size
- all_weight_blobs[key] = size
+ value = (size, _blob_mtime(file_obj))
+ target[key] = value
+ all_weight_blobs[key] = value
for revision in repo_info.revisions:
rev_id = getattr(revision, "commit_hash", None) or str(id(revision))
@@ -375,18 +416,19 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload:
or "unknown"
)
if model_format == "adapter":
- size_bytes = sum(adapter_blobs.values())
+ selected_blobs = adapter_blobs
elif model_format == "safetensors":
- size_bytes = sum(safetensors_blobs.values())
+ selected_blobs = safetensors_blobs
elif model_format == "checkpoint":
- size_bytes = sum(checkpoint_blobs.values())
+ selected_blobs = checkpoint_blobs
else:
- size_bytes = sum(all_weight_blobs.values())
+ selected_blobs = all_weight_blobs
return _CachedNonGgufPayload(
- size_bytes = size_bytes,
+ size_bytes = sum(size for size, _mtime in selected_blobs.values()),
has_runnable_weights = model_format != "unknown",
model_format = model_format,
+ last_modified = max((mtime for _size, mtime in selected_blobs.values()), default = 0.0),
)
@@ -508,6 +550,12 @@ def _scan_cached_models() -> list[dict]:
),
**_cached_model_local_metadata(repo_path),
}
+ last_modified = max(
+ payload.last_modified,
+ (existing or {}).get("last_modified", 0.0),
+ )
+ if last_modified > 0:
+ row["last_modified"] = last_modified
row.update(
_cache_inventory_fields(
repo_id,
@@ -517,6 +565,8 @@ def _scan_cached_models() -> list[dict]:
)
if _prefer_cache_row(row, existing):
seen_lower[key] = row
+ elif last_modified > existing.get("last_modified", 0.0):
+ existing["last_modified"] = last_modified
except Exception as e:
repo_label = getattr(repo_info, "repo_id", "")
logger.warning(f"Skipping cached model repo {repo_label}: {e}")
diff --git a/studio/backend/main.py b/studio/backend/main.py
index e64048dc00..bd0d26cf8f 100644
--- a/studio/backend/main.py
+++ b/studio/backend/main.py
@@ -313,6 +313,7 @@ from hub.routes import (
inventory_router as hub_inventory_router,
datasets_router as hub_datasets_router,
)
+from picker.routes import templates_router as picker_templates_router
from hub.schemas.downloads import TransportCapabilities
from hub.utils.download_registry import (
get_download_transport_capabilities,
@@ -745,6 +746,7 @@ _BODY_PROTECTED_PREFIXES = (
"/v1/completions",
"/p/",
"/api/inference",
+ "/api/picker",
"/api/data-recipe",
"/api/datasets",
"/api/hub",
@@ -975,6 +977,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(picker_templates_router, prefix = "/api/picker", tags = ["picker"])
# 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/picker/__init__.py b/studio/backend/picker/__init__.py
new file mode 100644
index 0000000000..32014236c6
--- /dev/null
+++ b/studio/backend/picker/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
diff --git a/studio/backend/picker/routes/__init__.py b/studio/backend/picker/routes/__init__.py
new file mode 100644
index 0000000000..c0e988c8bb
--- /dev/null
+++ b/studio/backend/picker/routes/__init__.py
@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+from .templates import router as templates_router
+
+__all__ = ["templates_router"]
diff --git a/studio/backend/picker/routes/templates.py b/studio/backend/picker/routes/templates.py
new file mode 100644
index 0000000000..03707669fa
--- /dev/null
+++ b/studio/backend/picker/routes/templates.py
@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+from __future__ import annotations
+
+import asyncio
+from typing import Optional
+
+from fastapi import APIRouter, Body, Depends, Query
+
+from auth.authentication import get_current_subject
+from hub.dependencies import get_hf_token
+
+from ..schemas import (
+ ModelTemplateResponse,
+ ValidateChatTemplateRequest,
+ ValidateChatTemplateResponse,
+)
+from ..service import read_default_chat_template, validate_chat_template
+
+router = APIRouter()
+
+
+@router.post("/validate-chat-template", response_model = ValidateChatTemplateResponse)
+async def validate_chat_template_route(
+ body: ValidateChatTemplateRequest = Body(...),
+ current_subject: str = Depends(get_current_subject),
+) -> ValidateChatTemplateResponse:
+ return await asyncio.to_thread(validate_chat_template, body.template)
+
+
+@router.get("/chat-template/{model_name:path}", response_model = ModelTemplateResponse)
+async def get_default_chat_template_route(
+ model_name: str,
+ gguf_variant: Optional[str] = Query(None),
+ hf_token: Optional[str] = Depends(get_hf_token),
+ current_subject: str = Depends(get_current_subject),
+) -> ModelTemplateResponse:
+ template = await asyncio.to_thread(
+ read_default_chat_template, model_name, hf_token, gguf_variant
+ )
+ return ModelTemplateResponse(model_name = model_name, chat_template = template)
diff --git a/studio/backend/picker/schemas.py b/studio/backend/picker/schemas.py
new file mode 100644
index 0000000000..b4f956188f
--- /dev/null
+++ b/studio/backend/picker/schemas.py
@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+from typing import Optional
+
+from pydantic import BaseModel, Field, field_validator
+
+# Mirror the frontend's 64 KiB chat-template contract (per-model-config.ts) at
+# the API boundary so a direct caller cannot make Jinja parse an oversized
+# template. MaxBodyMiddleware only caps the whole request body, not this field.
+MAX_CHAT_TEMPLATE_BYTES = 65_536
+
+
+class ValidateChatTemplateRequest(BaseModel):
+ template: str = Field(default = "")
+
+ @field_validator("template")
+ @classmethod
+ def _enforce_template_size(cls, value: str) -> str:
+ 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
+
+
+class ValidateChatTemplateResponse(BaseModel):
+ valid: bool
+ error: Optional[str] = None
+
+
+class ModelTemplateResponse(BaseModel):
+ model_name: str
+ chat_template: Optional[str] = None
diff --git a/studio/backend/picker/service.py b/studio/backend/picker/service.py
new file mode 100644
index 0000000000..f5994dc550
--- /dev/null
+++ b/studio/backend/picker/service.py
@@ -0,0 +1,361 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import re
+from pathlib import Path
+from typing import Optional
+
+from hub.services.models.folder_browser import (
+ _build_browse_allowlist,
+ _is_path_inside_allowlist,
+)
+from hub.utils.gguf import iter_hf_cache_snapshots
+from utils.models.gguf_metadata import read_gguf_chat_template
+from utils.models.model_config import (
+ _extract_quant_label,
+ _is_big_endian_gguf_path,
+ _is_mmproj,
+ _is_mtp_drafter,
+)
+from utils.paths.path_utils import (
+ is_local_path,
+ normalize_path,
+ resolve_cached_repo_id_case,
+)
+
+from .schemas import ValidateChatTemplateResponse
+
+logger = logging.getLogger(__name__)
+
+_VALID_REPO_ID = re.compile(r"^[A-Za-z0-9._-]+/[A-Za-z0-9._-]+$")
+
+
+def _is_valid_repo_id(repo_id: str) -> bool:
+ return bool(_VALID_REPO_ID.fullmatch(repo_id))
+
+
+_TOKENIZER_CONFIG_PATHS = ("tokenizer_config.json", "LLM/tokenizer_config.json")
+_JINJA_TEMPLATE_PATHS = ("chat_template.jinja", "LLM/chat_template.jinja")
+_PROCESSOR_TEMPLATE_PATHS = ("chat_template.json", "LLM/chat_template.json")
+
+
+def _leaf_inside_allowlist(path: Path, allow_roots: Optional[list[Path]]) -> bool:
+ # Block symlinked children from escaping the validated directory
+ # (realpath-checked). None = trusted caller (HF cache / remote download).
+ return allow_roots is None or _is_path_inside_allowlist(path, allow_roots)
+
+
+def validate_chat_template(template: str) -> ValidateChatTemplateResponse:
+ text = (template or "").strip()
+ if not text:
+ return ValidateChatTemplateResponse(valid = True, error = None)
+ # Import Jinja lazily: it is optional at runtime (e.g. GGUF-only installs),
+ # so a missing dependency must not crash API startup through this module.
+ try:
+ from jinja2 import TemplateError
+ from jinja2.ext import Extension
+ from jinja2.sandbox import ImmutableSandboxedEnvironment
+ except ImportError:
+ return ValidateChatTemplateResponse(valid = True, error = None)
+
+ class _GenerationTag(Extension):
+ # Accept Transformers' {% generation %}...{% endgeneration %} assistant
+ # mask tag so a pasted HF chat template validates (we only parse it).
+ tags = {"generation"}
+
+ def parse(self, parser):
+ next(parser.stream)
+ return parser.parse_statements(["name:endgeneration"], drop_needle = True)
+
+ try:
+ env = ImmutableSandboxedEnvironment(
+ trim_blocks = True,
+ lstrip_blocks = True,
+ extensions = ["jinja2.ext.loopcontrols", _GenerationTag],
+ )
+ env.parse(text)
+ return ValidateChatTemplateResponse(valid = True, error = None)
+ except TemplateError as exc:
+ message = getattr(exc, "message", None) or str(exc)
+ lineno = getattr(exc, "lineno", None)
+ if lineno:
+ message = f"Line {lineno}: {message}"
+ return ValidateChatTemplateResponse(valid = False, error = message)
+ except Exception as exc:
+ return ValidateChatTemplateResponse(valid = False, error = str(exc))
+
+
+def _chat_template_from_tokenizer_config(config: dict) -> Optional[str]:
+ if not isinstance(config, dict):
+ return None
+ raw = config.get("chat_template")
+ if isinstance(raw, str) and raw.strip():
+ return raw
+ if isinstance(raw, list):
+ fallback: Optional[str] = None
+ for entry in raw:
+ if not isinstance(entry, dict):
+ continue
+ template = entry.get("template")
+ if not isinstance(template, str):
+ continue
+ if entry.get("name") == "default":
+ return template
+ if fallback is None:
+ fallback = template
+ return fallback
+ return None
+
+
+def _chat_template_from_jinja_file(
+ dir_path: Path, allow_roots: Optional[list[Path]] = None
+) -> Optional[str]:
+ for rel in _JINJA_TEMPLATE_PATHS:
+ template_file = dir_path / rel
+ if not template_file.exists() or not _leaf_inside_allowlist(template_file, allow_roots):
+ continue
+ try:
+ template = template_file.read_text(encoding = "utf-8")
+ except Exception:
+ continue
+ if template.strip():
+ return template
+ return None
+
+
+def _chat_template_from_processor_payload(payload: object) -> Optional[str]:
+ # processor chat_template.json may be the template string itself or a
+ # {name: template} map, not only a tokenizer_config-shaped object.
+ if isinstance(payload, str):
+ return payload if payload.strip() else None
+ template = _chat_template_from_tokenizer_config(payload) # type: ignore[arg-type]
+ if template:
+ return template
+ if isinstance(payload, dict):
+ # Named-template map: prefer "default", else the first non-empty entry
+ # (mirrors the tokenizer-config list fallback).
+ default = payload.get("default")
+ if isinstance(default, str) and default.strip():
+ return default
+ for value in payload.values():
+ if isinstance(value, str) and value.strip():
+ return value
+ return None
+
+
+def _chat_template_from_processor_json(
+ dir_path: Path, allow_roots: Optional[list[Path]] = None
+) -> Optional[str]:
+ for rel in _PROCESSOR_TEMPLATE_PATHS:
+ config_file = dir_path / rel
+ if not config_file.exists() or not _leaf_inside_allowlist(config_file, allow_roots):
+ continue
+ try:
+ payload = json.loads(config_file.read_text(encoding = "utf-8"))
+ except Exception:
+ continue
+ template = _chat_template_from_processor_payload(payload)
+ if template:
+ return template
+ return None
+
+
+def _chat_template_from_tokenizer_dir(
+ dir_path: Path, allow_roots: Optional[list[Path]] = None
+) -> Optional[str]:
+ jinja = _chat_template_from_jinja_file(dir_path, allow_roots)
+ if jinja:
+ return jinja
+ for rel in _TOKENIZER_CONFIG_PATHS:
+ config_file = dir_path / rel
+ if not config_file.exists() or not _leaf_inside_allowlist(config_file, allow_roots):
+ continue
+ try:
+ config = json.loads(config_file.read_text(encoding = "utf-8"))
+ except Exception:
+ continue
+ template = _chat_template_from_tokenizer_config(config)
+ if template:
+ return template
+ return _chat_template_from_processor_json(dir_path, allow_roots)
+
+
+_GGUF_SCAN_MAX_DEPTH = 2
+
+
+def _iter_ggufs(dir_path: Path) -> list[Path]:
+ if dir_path == dir_path.parent:
+ return []
+ root = str(dir_path)
+ found: list[Path] = []
+ for current, dirs, files in os.walk(root, followlinks = False):
+ rel = os.path.relpath(current, root)
+ depth = 0 if rel == os.curdir else rel.count(os.sep) + 1
+ if depth >= _GGUF_SCAN_MAX_DEPTH:
+ dirs[:] = []
+ for name in files:
+ if not name.lower().endswith(".gguf") or _is_mmproj(name):
+ continue
+ path = Path(current) / name
+ try:
+ rel = path.relative_to(dir_path).as_posix()
+ except ValueError:
+ rel = name
+ quant = _extract_quant_label(rel)
+ if _is_mtp_drafter(rel) or _is_big_endian_gguf_path(rel, quant):
+ continue
+ found.append(path)
+ return found
+
+
+def _variant_matches(relative_path: str, needle: str) -> bool:
+ quant = _extract_quant_label(relative_path).lower()
+ if quant == needle:
+ return True
+ prefix = f"{needle}-"
+ if not quant.startswith(prefix):
+ return False
+ suffix = quant[len(prefix) :]
+ if not suffix.endswith("bpw"):
+ return False
+ value = suffix[:-3]
+ return bool(value) and value.replace(".", "", 1).isdigit()
+
+
+def _find_gguf_in_dir(dir_path: Path, gguf_variant: Optional[str]) -> Optional[Path]:
+ try:
+ ggufs = sorted(_iter_ggufs(dir_path))
+ except OSError:
+ return None
+ if not ggufs:
+ return None
+ needle = (gguf_variant or "").strip().lower()
+ if needle:
+ for path in ggufs:
+ try:
+ relative = path.relative_to(dir_path).as_posix()
+ except ValueError:
+ relative = path.name
+ if _variant_matches(relative, needle):
+ return path
+ return None
+ try:
+ return max(ggufs, key = lambda path: path.stat().st_size)
+ except OSError:
+ return ggufs[0]
+
+
+def _chat_template_from_dir(
+ dir_path: Path,
+ gguf_variant: Optional[str] = None,
+ allow_roots: Optional[list[Path]] = None,
+) -> Optional[str]:
+ def from_gguf() -> Optional[str]:
+ gguf = _find_gguf_in_dir(dir_path, gguf_variant)
+ if gguf is None or not _leaf_inside_allowlist(gguf, allow_roots):
+ return None
+ return read_gguf_chat_template(str(gguf))
+
+ # Sidecar tokenizer files (chat_template.jinja / tokenizer_config.json) are
+ # the model author's maintained template and supersede the GGUF's embedded
+ # copy, which can be stale. The variant only selects which GGUF to fall back
+ # to, so keep tokenizer-first precedence whether or not a variant is given.
+ return _chat_template_from_tokenizer_dir(dir_path, allow_roots) or from_gguf()
+
+
+def read_default_chat_template(
+ model_name: str,
+ hf_token: Optional[str] = None,
+ gguf_variant: Optional[str] = None,
+) -> Optional[str]:
+ if not isinstance(model_name, str) or not model_name.strip():
+ return None
+ name = model_name.strip()
+
+ if is_local_path(name):
+ try:
+ target = Path(normalize_path(name)).expanduser()
+ allow_roots = _build_browse_allowlist()
+ if not _is_path_inside_allowlist(target, allow_roots):
+ logger.debug("Refused chat template read outside allowed folders: %s", name)
+ return None
+ if name.lower().endswith(".gguf"):
+ # Prefer a maintained sidecar template (chat_template.jinja /
+ # tokenizer_config.json) next to the file over the GGUF's embedded
+ # copy, matching the tokenizer-first precedence used for directory
+ # and variant selections.
+ sidecar = _chat_template_from_tokenizer_dir(target.parent, allow_roots)
+ if sidecar:
+ return sidecar
+ return read_gguf_chat_template(str(target))
+ return _chat_template_from_dir(target, gguf_variant, allow_roots)
+ except Exception as exc:
+ logger.debug("Could not read local chat template for %s: %s", name, exc)
+ return None
+
+ if not _is_valid_repo_id(name):
+ return None
+
+ resolved = resolve_cached_repo_id_case(name)
+
+ try:
+ # Resolve within each cached revision, newest first. A revision's
+ # maintained sidecar (chat_template.jinja / tokenizer_config.json)
+ # supersedes its own embedded GGUF copy, but a newer revision must not be
+ # overridden by an older revision's sidecar, so precedence stays
+ # per-snapshot rather than searching all sidecars globally first.
+ for snapshot in iter_hf_cache_snapshots(resolved):
+ template = _chat_template_from_dir(snapshot, gguf_variant)
+ if template:
+ return template
+ except Exception as exc:
+ logger.debug("Could not read cached chat template for %s: %s", resolved, exc)
+
+ try:
+ from huggingface_hub import hf_hub_download
+
+ def _download_text(rel: str) -> Optional[str]:
+ try:
+ path = hf_hub_download(resolved, rel, token = hf_token)
+ return Path(path).read_text(encoding = "utf-8")
+ except Exception:
+ return None
+
+ for rel in _JINJA_TEMPLATE_PATHS:
+ template = _download_text(rel)
+ if template and template.strip():
+ return template
+
+ for rel in _TOKENIZER_CONFIG_PATHS:
+ raw = _download_text(rel)
+ if not raw:
+ continue
+ try:
+ config = json.loads(raw)
+ except Exception:
+ continue
+ template = _chat_template_from_tokenizer_config(config)
+ if template:
+ return template
+
+ for rel in _PROCESSOR_TEMPLATE_PATHS:
+ raw = _download_text(rel)
+ if not raw:
+ continue
+ try:
+ payload = json.loads(raw)
+ except Exception:
+ continue
+ template = _chat_template_from_processor_payload(payload)
+ if template:
+ return template
+
+ return None
+ except Exception as exc:
+ logger.debug("Could not fetch chat template for %s: %s", resolved, exc)
+ return None
diff --git a/studio/backend/tests/test_model_update_robustness.py b/studio/backend/tests/test_model_update_robustness.py
index edf55812e2..4c5822e662 100644
--- a/studio/backend/tests/test_model_update_robustness.py
+++ b/studio/backend/tests/test_model_update_robustness.py
@@ -314,6 +314,7 @@ def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path):
file_name = "model.safetensors",
size_on_disk = 100,
blob_path = str(repo_path / "blobs" / "modelsha"),
+ blob_last_modified = 3_000.0,
),
]
)
@@ -336,6 +337,51 @@ def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path):
assert rows[0]["repo_id"] == "Org/SafeTensorRepo"
assert rows[0]["model_format"] == "safetensors"
assert rows[0]["size_bytes"] == 100
+ assert rows[0]["last_modified"] == 3_000.0
+
+
+def test_cached_gguf_scan_keeps_download_timestamp(monkeypatch, tmp_path):
+ repo_path = tmp_path / "models--Org--GgufRepo"
+ repo = SimpleNamespace(
+ repo_id = "Org/GgufRepo",
+ repo_type = "model",
+ repo_path = repo_path,
+ revisions = [
+ SimpleNamespace(
+ files = [
+ SimpleNamespace(
+ file_name = "model-Q4_K_M.gguf",
+ size_on_disk = 100,
+ blob_path = None,
+ blob_last_modified = 5_000.0,
+ ),
+ ]
+ )
+ ],
+ )
+ monkeypatch.setattr(
+ CI,
+ "all_hf_cache_scans",
+ lambda: [SimpleNamespace(repos = [repo])],
+ )
+ monkeypatch.setattr(
+ CI.hf_cache_scan,
+ "is_gguf_repo_partial",
+ lambda *args, **kwargs: False,
+ )
+ monkeypatch.setattr(
+ CI,
+ "_gguf_variant_state_summary",
+ lambda _repo_id: (False, 0),
+ )
+
+ rows = CI._scan_cached_gguf()
+
+ assert len(rows) == 1
+ assert rows[0]["repo_id"] == "Org/GgufRepo"
+ assert rows[0]["model_format"] == "gguf"
+ assert rows[0]["size_bytes"] == 100
+ assert rows[0]["last_modified"] == 5_000.0
# ── hf_hub_download_with_xet_fallback force_download bypass (X2/F2) ───
diff --git a/studio/backend/tests/test_picker_service.py b/studio/backend/tests/test_picker_service.py
new file mode 100644
index 0000000000..fc835bf019
--- /dev/null
+++ b/studio/backend/tests/test_picker_service.py
@@ -0,0 +1,162 @@
+# 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 json
+
+from picker.service import (
+ _chat_template_from_dir,
+ _chat_template_from_tokenizer_config,
+ _chat_template_from_tokenizer_dir,
+ _find_gguf_in_dir,
+ _iter_ggufs,
+ read_default_chat_template,
+ validate_chat_template,
+)
+
+
+def test_iter_ggufs_skips_gguf_companions(tmp_path):
+ mtp_dir = tmp_path / "MTP"
+ mtp_dir.mkdir()
+ main = tmp_path / "model-Q8_0.gguf"
+ main.write_bytes(b"")
+ (tmp_path / "mmproj-F16.gguf").write_bytes(b"")
+ (tmp_path / "mtp-model-Q8_0.gguf").write_bytes(b"")
+ (mtp_dir / "model-Q8_0-MTP.gguf").write_bytes(b"")
+ (tmp_path / "model-Q8_0-be.gguf").write_bytes(b"")
+
+ assert _iter_ggufs(tmp_path) == [main]
+
+
+def test_find_gguf_in_dir_matches_quant_label(tmp_path):
+ mtp_dir = tmp_path / "MTP"
+ mtp_dir.mkdir()
+ main = tmp_path / "model-Q8_0.gguf"
+ main.write_bytes(b"")
+ (mtp_dir / "model-Q8_0-MTP.gguf").write_bytes(b"")
+ (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"")
+
+ assert _find_gguf_in_dir(tmp_path, "Q8_0") == main
+ assert _find_gguf_in_dir(tmp_path, "Q4_K") is None
+
+
+def test_find_gguf_in_dir_without_variant_prefers_largest_model(tmp_path):
+ smaller = tmp_path / "a-model-Q4_K_M.gguf"
+ larger = tmp_path / "z-model-Q8_0.gguf"
+ smaller.write_bytes(b"0")
+ larger.write_bytes(b"00")
+
+ assert _find_gguf_in_dir(tmp_path, None) == larger
+
+
+def test_find_gguf_in_dir_matches_bpw_variant_base_label(tmp_path):
+ target = tmp_path / "model-IQ4_XS-3.53bpw.gguf"
+ target.write_bytes(b"")
+ (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"")
+
+ assert _find_gguf_in_dir(tmp_path, "IQ4_XS") == target
+ assert _find_gguf_in_dir(tmp_path, "IQ4_XS-3.53bpw") == target
+ assert _find_gguf_in_dir(tmp_path, "Q4_K") is None
+
+
+def test_validate_chat_template_accepts_valid_and_empty():
+ assert validate_chat_template("{{ messages[0].content }}").valid is True
+ assert validate_chat_template("").valid is True
+ assert validate_chat_template(" ").valid is True
+
+
+def test_validate_chat_template_reports_syntax_error_with_line():
+ result = validate_chat_template("{% if %}{% endif %}")
+ assert result.valid is False
+ assert result.error is not None
+ assert result.error.startswith("Line ")
+
+
+def test_chat_template_from_tokenizer_config_reads_string():
+ assert _chat_template_from_tokenizer_config({"chat_template": "HELLO"}) == "HELLO"
+ assert _chat_template_from_tokenizer_config({"chat_template": " "}) is None
+ assert _chat_template_from_tokenizer_config({}) is None
+
+
+def test_chat_template_from_tokenizer_config_prefers_named_default():
+ config = {
+ "chat_template": [
+ {"name": "tool_use", "template": "TOOL"},
+ {"name": "default", "template": "DEFAULT"},
+ ]
+ }
+ assert _chat_template_from_tokenizer_config(config) == "DEFAULT"
+
+
+def test_chat_template_from_tokenizer_config_falls_back_to_first_entry():
+ config = {
+ "chat_template": [
+ {"name": "tool_use", "template": "TOOL"},
+ {"name": "other", "template": "OTHER"},
+ ]
+ }
+ assert _chat_template_from_tokenizer_config(config) == "TOOL"
+
+
+def test_chat_template_from_tokenizer_dir_prefers_jinja_file(tmp_path):
+ (tmp_path / "chat_template.jinja").write_text("FROM_JINJA", encoding = "utf-8")
+ (tmp_path / "tokenizer_config.json").write_text(
+ json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
+ )
+ assert _chat_template_from_tokenizer_dir(tmp_path) == "FROM_JINJA"
+
+
+def test_chat_template_from_tokenizer_dir_reads_tokenizer_config(tmp_path):
+ (tmp_path / "tokenizer_config.json").write_text(
+ json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
+ )
+ assert _chat_template_from_tokenizer_dir(tmp_path) == "FROM_CONFIG"
+
+
+def test_chat_template_from_dir_without_variant_prefers_tokenizer(tmp_path):
+ (tmp_path / "tokenizer_config.json").write_text(
+ json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
+ )
+ assert _chat_template_from_dir(tmp_path) == "FROM_CONFIG"
+
+
+def test_chat_template_from_dir_with_variant_still_prefers_tokenizer(tmp_path, monkeypatch):
+ (tmp_path / "tokenizer_config.json").write_text(
+ json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
+ )
+ (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"")
+ monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF")
+ # Selecting a variant must not flip precedence to the embedded GGUF template.
+ assert _chat_template_from_dir(tmp_path, "Q4_K_M") == "FROM_CONFIG"
+
+
+def test_chat_template_from_dir_with_variant_falls_back_to_gguf(tmp_path, monkeypatch):
+ (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"")
+ monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF")
+ # With no tokenizer sidecar, the embedded GGUF template is still the fallback.
+ assert _chat_template_from_dir(tmp_path, "Q4_K_M") == "FROM_GGUF"
+
+
+def test_chat_template_from_dir_returns_none_when_absent(tmp_path):
+ assert _chat_template_from_dir(tmp_path) is None
+
+
+def test_read_default_chat_template_direct_gguf_prefers_sidecar(tmp_path, monkeypatch):
+ gguf = tmp_path / "model-Q4_K_M.gguf"
+ gguf.write_bytes(b"")
+ (tmp_path / "tokenizer_config.json").write_text(
+ json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
+ )
+ monkeypatch.setattr("picker.service._build_browse_allowlist", lambda: [tmp_path])
+ monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF")
+ # A directly selected .gguf file must prefer a maintained sidecar template
+ # over its embedded copy, matching directory/variant precedence.
+ assert read_default_chat_template(str(gguf)) == "FROM_CONFIG"
+
+
+def test_read_default_chat_template_direct_gguf_falls_back_to_embedded(tmp_path, monkeypatch):
+ gguf = tmp_path / "model-Q4_K_M.gguf"
+ gguf.write_bytes(b"")
+ monkeypatch.setattr("picker.service._build_browse_allowlist", lambda: [tmp_path])
+ monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF")
+ # With no sidecar next to the file, the embedded GGUF template is the fallback.
+ assert read_default_chat_template(str(gguf)) == "FROM_GGUF"
diff --git a/studio/backend/utils/models/gguf_metadata.py b/studio/backend/utils/models/gguf_metadata.py
index c24ec28e1d..5e25ce1927 100644
--- a/studio/backend/utils/models/gguf_metadata.py
+++ b/studio/backend/utils/models/gguf_metadata.py
@@ -50,6 +50,8 @@ _CACHE_MAX_ENTRIES = 4096
# keyed by (file cache key, wanted key). None = key absent / file unreadable.
_BOOL_CACHE: Dict[Tuple[_CacheKey, str], Optional[bool]] = {}
+_STRING_CACHE: Dict[Tuple[_CacheKey, str], Optional[str]] = {}
+
# 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]] = {}
@@ -353,6 +355,83 @@ def _read_gguf_bool(path: str, wanted_key: str) -> Optional[bool]:
return result
+def _parse_gguf_string(path: str, wanted_key: str) -> Optional[str]:
+ 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:
+ 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:
+ break
+ sbytes = f.read(slen)
+ if len(sbytes) < slen:
+ break
+ return sbytes.decode("utf-8", "replace")
+ if not _skip_gguf_value(f, vtype):
+ break
+ except (struct.error, UnicodeDecodeError):
+ break
+ except OSError as e:
+ logger.debug(f"_parse_gguf_string: cannot open {path}: {e}")
+ return None
+ except Exception as e:
+ logger.debug(f"_parse_gguf_string: parse failure on {path}: {e}")
+ return None
+ return None
+
+
+def _read_gguf_string(path: str, wanted_key: str) -> Optional[str]:
+ fkey = _cache_key(path)
+ if fkey is None:
+ return None
+ ckey = (fkey, wanted_key)
+ with _CACHE_LOCK:
+ if ckey in _STRING_CACHE:
+ return _STRING_CACHE[ckey]
+ result = _parse_gguf_string(path, wanted_key)
+ with _CACHE_LOCK:
+ while len(_STRING_CACHE) >= _CACHE_MAX_ENTRIES:
+ try:
+ _STRING_CACHE.pop(next(iter(_STRING_CACHE)))
+ except StopIteration:
+ break
+ _STRING_CACHE[ckey] = result
+ return result
+
+
+def read_gguf_chat_template(path: str) -> Optional[str]:
+ template = _read_gguf_string(path, "tokenizer.chat_template")
+ if isinstance(template, str) and template.strip():
+ return template
+ return None
+
+
def read_mmproj_audio_capability(path: str) -> Optional[bool]:
"""``clip.has_audio_encoder`` from an mmproj GGUF (e.g. Gemma 4's
gemma4ua): ``True``/``False`` if present, ``None`` if absent/unreadable.
diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx
index ba56ce7525..6c2505f1ca 100644
--- a/studio/frontend/src/app/routes/__root.tsx
+++ b/studio/frontend/src/app/routes/__root.tsx
@@ -195,9 +195,6 @@ function RootLayout() {
chatRuntime.setActiveThreadId(null);
chatRuntime.setActiveProjectId(null);
chatRuntime.setIncognito(false);
- // Detach the staging UI but keep any in-flight download running, like Hub.
- if (chatRuntime.pendingSelection)
- chatRuntime.abandonStagedModel({ keepDownload: true });
void navigate({
to: "/chat",
search: { new: crypto.randomUUID() },
@@ -220,10 +217,6 @@ function RootLayout() {
chatRuntime.setActiveProjectId(null);
chatRuntime.setActiveThreadId(null);
chatRuntime.setIncognito(false);
- // Leaving chat must not kill an in-flight download: detach the staging UI
- // but keep the transfer running in the manager, like a Hub download.
- if (chatRuntime.pendingSelection)
- chatRuntime.abandonStagedModel({ keepDownload: true });
}, [isChatRoute]);
return (
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts b/studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts
deleted file mode 100644
index ec75b17f20..0000000000
--- a/studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts
+++ /dev/null
@@ -1,69 +0,0 @@
-// SPDX-License-Identifier: AGPL-3.0-only
-// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
-
-// Per-model pre-load inference settings, persisted in localStorage so the load
-// dialog can offer "Remember settings for ".
-
-const KEY = "unsloth_load_settings";
-
-export interface RememberedLoadSettings {
- contextLength: number | null;
- kvCacheDtype: string | null;
- speculativeType: string | null;
- specDraftNMax: number | null;
- tensorParallel: boolean;
-}
-
-// Storage key for a pick's remembered settings. The remembered knobs are
-// VRAM-budget driven (context override, KV-cache dtype, tensor-parallel), so the
-// right values differ per quant. An HF repo collapses all its GGUF variants into
-// one `id`, so fold the variant in to scope settings per quant. Local .gguf
-// paths key by their file path (already file-specific); native drag-drop files
-// key by display label, so same-named files in different folders share an entry.
-export function rememberedLoadSettingsKey(selection: {
- id: string;
- ggufVariant?: string | null;
-}): string {
- return selection.ggufVariant
- ? `${selection.id}::${selection.ggufVariant}`
- : selection.id;
-}
-
-function readAll(): Record {
- try {
- return JSON.parse(localStorage.getItem(KEY) ?? "{}");
- } catch {
- return {};
- }
-}
-
-function writeAll(all: Record) {
- try {
- localStorage.setItem(KEY, JSON.stringify(all));
- } catch {
- // Ignore quota / unavailable storage.
- }
-}
-
-export function loadRememberedLoadSettings(
- key: string,
-): RememberedLoadSettings | null {
- return readAll()[key] ?? null;
-}
-
-export function saveRememberedLoadSettings(
- key: string,
- settings: RememberedLoadSettings,
-) {
- const all = readAll();
- all[key] = settings;
- writeAll(all);
-}
-
-export function clearRememberedLoadSettings(key: string) {
- const all = readAll();
- if (key in all) {
- delete all[key];
- writeAll(all);
- }
-}
diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts
index 0bf46e7343..d3bdcb6898 100644
--- a/studio/frontend/src/features/chat/api/chat-adapter.ts
+++ b/studio/frontend/src/features/chat/api/chat-adapter.ts
@@ -2,10 +2,7 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { getAuthToken } from "@/features/auth";
-import {
- loadRememberedLoadSettings,
- rememberedLoadSettingsKey,
-} from "@/components/assistant-ui/model-selector/remembered-load-settings";
+import { resolveInitialConfig } from "@/features/model-picker";
import { projectHasSources } from "@/features/rag/api/rag-api";
import { apiUrl } from "@/lib/api-base";
import { parseParamCountB } from "@/lib/model-size";
@@ -1520,27 +1517,25 @@ async function autoLoadSmallestModel(): Promise<{
return false;
}
const currentStore = useChatRuntimeStore.getState();
- const remembered = loadRememberedLoadSettings(
- rememberedLoadSettingsKey({
- id: candidate.id,
- ggufVariant: candidate.ggufVariant,
- }),
- );
+ const { config } = resolveInitialConfig(candidate.id, candidate.ggufVariant);
const effectiveMaxSeqLength = resolveLoadMaxSeqLength({
modelId: candidate.id,
ggufVariant: candidate.ggufVariant,
isGguf: candidate.kind === "gguf",
- customContextLength: remembered?.contextLength ?? null,
+ customContextLength: config.customContextLength,
ggufContextLength: null,
currentCheckpoint: currentStore.params.checkpoint,
activeGgufVariant: currentStore.activeGgufVariant,
- maxSeqLength: candidate.maxSeqLength,
+ maxSeqLength: config.maxSeqLength ?? candidate.maxSeqLength,
presetSource: currentStore.activePresetSource,
});
const effectiveSpeculativeType =
- remembered?.speculativeType ?? specSettings.speculativeType;
+ config.speculativeType ?? specSettings.speculativeType;
const effectiveSpecDraftNMax =
- remembered?.specDraftNMax ?? specSettings.specDraftNMax;
+ config.specDraftNMax ?? specSettings.specDraftNMax;
+ const effectiveChatTemplateOverride = config.chatTemplateOverride?.trim()
+ ? config.chatTemplateOverride
+ : null;
if (
!(await canAutoLoad({
model_path: candidate.id,
@@ -1563,12 +1558,18 @@ async function autoLoadSmallestModel(): Promise<{
is_lora: false,
gguf_variant: candidate.ggufVariant,
trust_remote_code: trustRemoteCode,
- cache_type_kv: remembered?.kvCacheDtype ?? null,
+ chat_template_override: effectiveChatTemplateOverride,
+ cache_type_kv: config.kvCacheDtype,
speculative_type: effectiveSpeculativeType,
spec_draft_n_max: effectiveSpecDraftNMax,
- tensor_parallel: remembered?.tensorParallel ?? false,
+ tensor_parallel: config.tensorParallel,
});
- saveSpeculativeType(effectiveSpeculativeType);
+ // Only persist the global preference when the value came from the global
+ // settings. A per-model config's choice must stay load-local, or autoloading
+ // a remembered model on startup would rewrite the global default.
+ if (config.speculativeType == null) {
+ saveSpeculativeType(effectiveSpeculativeType);
+ }
useChatRuntimeStore
.getState()
.setCheckpoint(candidate.id, candidate.ggufVariant ?? undefined);
@@ -1578,6 +1579,9 @@ async function autoLoadSmallestModel(): Promise<{
);
store.setParams({
...store.params,
+ ...(candidate.kind === "gguf"
+ ? {}
+ : { maxSeqLength: effectiveMaxSeqLength }),
maxTokens:
candidate.kind === "gguf"
? loadResp.context_length ?? 131072
@@ -1614,8 +1618,11 @@ async function autoLoadSmallestModel(): Promise<{
tensorParallel: loadResp.tensor_parallel ?? false,
loadedTensorParallel: loadResp.tensor_parallel ?? false,
defaultChatTemplate: loadResp.chat_template ?? null,
- chatTemplateOverride: null,
- loadedChatTemplateOverride: null,
+ chatTemplateOverride: effectiveChatTemplateOverride,
+ loadedChatTemplateOverride: effectiveChatTemplateOverride,
+ // Retain the saved requested context so re-saving the config keeps the
+ // override; null stays null (auto/VRAM-fit).
+ customContextLength: config.customContextLength,
loadedIsMultimodal: isMultimodalResponse(loadResp),
loadedIsDiffusion: loadResp.is_diffusion ?? false,
...resolveLoadedSpeculativeSettings(loadResp),
@@ -1634,8 +1641,9 @@ async function autoLoadSmallestModel(): Promise<{
tensorParallel: loadResp.tensor_parallel ?? false,
loadedTensorParallel: loadResp.tensor_parallel ?? false,
defaultChatTemplate: loadResp.chat_template ?? null,
- chatTemplateOverride: null,
- loadedChatTemplateOverride: null,
+ chatTemplateOverride: effectiveChatTemplateOverride,
+ loadedChatTemplateOverride: effectiveChatTemplateOverride,
+ customContextLength: null,
...resolveLoadedSpeculativeSettings(loadResp),
loadedIsMultimodal: isMultimodalResponse(loadResp),
loadedIsDiffusion: loadResp.is_diffusion ?? false,
diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx
index 380ce0e0ab..731a551c53 100644
--- a/studio/frontend/src/features/chat/chat-page.tsx
+++ b/studio/frontend/src/features/chat/chat-page.tsx
@@ -2,16 +2,18 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import {
+ applyModelLoadConfigToRuntime,
+ currentRuntimePerModelConfig,
type DeletedModelRef,
type ExternalModelOption,
type LoraModelOption,
type ModelOption,
ModelSelector,
-} from "@/components/assistant-ui/model-selector";
-import {
- loadRememberedLoadSettings,
- rememberedLoadSettingsKey,
-} from "@/components/assistant-ui/model-selector/remembered-load-settings";
+ type ModelSelectorChangeMeta,
+ type PerModelConfig,
+ resolveInitialConfig,
+ SidebarModelConfig,
+} from "@/features/model-picker";
import { ProjectComposer, Thread } from "@/components/assistant-ui/thread";
import { CopyableErrorChip } from "@/components/ui/copyable-error-chip";
import {
@@ -27,10 +29,10 @@ import {
} from "@/components/ui/resizable";
import { useSidebar } from "@/components/ui/sidebar";
import { Tooltip, TooltipContent } from "@/components/ui/tooltip";
-import { useLatestRef } from "@/features/hub/hooks/use-latest-ref";
import {
DOWNLOAD_KIND,
downloadManager,
+ useRepoDownload,
} from "@/features/hub/download-manager";
import {
type NativeIntent,
@@ -93,7 +95,6 @@ import {
renameChatItem,
useChatSidebarItems,
} from "./hooks/use-chat-sidebar-items";
-import { useStagedModelPreparation } from "./hooks/use-staged-model-preparation";
import {
clearTrainingCompareHandoff,
getTrainingCompareHandoff,
@@ -128,10 +129,8 @@ import {
hasGgufSource,
isDownloadableHubRepo,
loadOptionalBool,
- pendingSelectionMatches,
useChatRuntimeStore,
} from "./stores/chat-runtime-store";
-import type { PendingModelSelection } from "./stores/chat-runtime-store";
import { useChatPreferencesStore } from "./stores/chat-preferences-store";
import { useExternalProvidersStore } from "./stores/external-providers-store";
import { buildChatTourSteps } from "./tour";
@@ -385,6 +384,7 @@ type CompareModelSelection = {
id: string;
isLora: boolean;
ggufVariant?: string;
+ config?: PerModelConfig;
};
function modelMatchesDeleted(
@@ -645,6 +645,8 @@ function GeneralCompareHeader({
loraModels,
externalModels,
value,
+ selectedConfig,
+ selectedGgufVariant,
onValueChange,
onFoldersChange,
onModelsChange,
@@ -655,9 +657,11 @@ function GeneralCompareHeader({
loraModels: LoraModelOption[];
externalModels: ExternalModelOption[];
value: string;
+ selectedConfig?: PerModelConfig | null;
+ selectedGgufVariant?: string | null;
onValueChange: (
id: string,
- meta: { isLora: boolean; ggufVariant?: string },
+ meta: ModelSelectorChangeMeta,
) => void;
onFoldersChange?: () => void;
onModelsChange?: (deletedModel?: DeletedModelRef) => void;
@@ -684,6 +688,8 @@ function GeneralCompareHeader({
loraModels={loraModels}
externalModels={externalModels}
value={value}
+ selectedConfig={selectedConfig}
+ selectedGgufVariant={selectedGgufVariant}
onValueChange={onValueChange}
onFoldersChange={onFoldersChange}
onModelsChange={onModelsChange}
@@ -811,11 +817,14 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
loraModels={loraModels}
externalModels={externalModels}
value={model1.id}
+ selectedConfig={model1.config}
+ selectedGgufVariant={model1.ggufVariant}
onValueChange={(id, meta) =>
setModel1({
id,
isLora: meta.isLora,
ggufVariant: meta.ggufVariant,
+ config: meta.config,
})
}
onFoldersChange={onFoldersChange}
@@ -838,11 +847,14 @@ const GeneralCompareContent = memo(function GeneralCompareContent({
loraModels={loraModels}
externalModels={externalModels}
value={model2.id}
+ selectedConfig={model2.config}
+ selectedGgufVariant={model2.ggufVariant}
onValueChange={(id, meta) =>
setModel2({
id,
isLora: meta.isLora,
ggufVariant: meta.ggufVariant,
+ config: meta.config,
})
}
onFoldersChange={onFoldersChange}
@@ -1236,6 +1248,13 @@ export function validateChatSearch(search: Record): ChatSearch
};
}
+type PendingHubAutoLoad = {
+ selection: SelectedModelInput;
+ contextKey: string;
+ originCheckpoint: string;
+ originGgufVariant: string | null;
+};
+
// `search` comes from RootLayout (not useSearch) so ChatPage stays mounted off-route
// (keeping an in-flight generation alive), frozen to the last /chat search. `active`
// is false off-route: close body-portaled surfaces and stop route-specific listeners
@@ -1248,30 +1267,6 @@ export function ChatPage({
const settingsOpen = useChatRuntimeStore((s) => s.settingsPanelOpen);
const setSettingsOpen = useChatRuntimeStore((s) => s.setSettingsPanelOpen);
- // 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.
- // autoLoad picks instead load the cached file as soon as the download ends;
- // selectModel is defined below, so the load runs through a ref.
- const autoLoadStagedRef = useRef<
- ((pending: PendingModelSelection) => void) | null
- >(null);
- const stagedDownload = useStagedModelPreparation({
- onAutoLoad: (pending) => autoLoadStagedRef.current?.(pending),
- });
- // 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();
- }, []);
- // Detach a staged pick on navigation without cancelling its download: the
- // transfer keeps running in the manager and lands in cache, like Hub.
- const detachStaged = useCallback(() => {
- useChatRuntimeStore.getState().abandonStagedModel({ keepDownload: true });
- }, []);
- // 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
@@ -1363,6 +1358,9 @@ export function ChatPage({
const ggufContextLength = useChatRuntimeStore(
(state) => state.ggufContextLength,
);
+ const ggufNativeContextLength = useChatRuntimeStore(
+ (state) => state.ggufNativeContextLength,
+ );
const contextUsage = useChatRuntimeStore((state) => state.contextUsage);
const modelsFromStore = useChatRuntimeStore((state) => state.models);
const lorasFromStore = useChatRuntimeStore((state) => state.loras);
@@ -1440,37 +1438,82 @@ export function ChatPage({
refreshRef.current = refresh;
selectModelRef.current = selectModel;
}, [refresh, selectModel]);
- // Load a cached autoLoad pick once its download finishes. The sheet was never
- // opened, so on a load failure just drop the orphaned staged knobs. The knobs
- // were already seeded on stage, so keepSpeculative only when a config was
- // saved -- otherwise the standing speculative preference should win.
- autoLoadStagedRef.current = (pending) => {
- const remembered = loadRememberedLoadSettings(
- rememberedLoadSettingsKey(pending),
- );
- void selectModel({
- ...pending,
- isDownloaded: true,
- forceReload: true,
- keepSpeculative: remembered != null,
- throwOnError: true,
- }).catch(() => {
- const store = useChatRuntimeStore.getState();
- // selectModel only clears pendingSelection on success, so a failed
- // auto-load leaves our staged pick (and its edited load knobs) behind.
- // Abandon it when it is still the active stage; otherwise just revert the
- // settings if the stage was already cleared by something else.
- if (pendingSelectionMatches(store.pendingSelection, pending)) {
- store.abandonStagedModel();
- } else if (!store.pendingSelection) {
- store.resetModelSettingsToLoaded();
- }
- });
- };
+ const rememberedConfigFor = useCallback(
+ (selection: {
+ id: string;
+ ggufVariant?: string | null;
+ source?: string;
+ }) => {
+ if (selection.source === "external") return null;
+ const resolved = resolveInitialConfig(selection.id, selection.ggufVariant);
+ return resolved.remembered ? resolved.config : null;
+ },
+ [],
+ );
const isExternalModel = useMemo(
() => isExternalModelId(inferenceParams.checkpoint),
[inferenceParams.checkpoint],
);
+ const runtimeCustomContextLength = useChatRuntimeStore(
+ (s) => s.customContextLength,
+ );
+ const runtimeKvCacheDtype = useChatRuntimeStore((s) => s.kvCacheDtype);
+ const runtimeSpeculativeType = useChatRuntimeStore((s) => s.speculativeType);
+ const runtimeSpecDraftNMax = useChatRuntimeStore((s) => s.specDraftNMax);
+ const runtimeTensorParallel = useChatRuntimeStore((s) => s.tensorParallel);
+ const runtimeChatTemplateOverride = useChatRuntimeStore(
+ (s) => s.chatTemplateOverride,
+ );
+ const activeModelConfig = useMemo(() => {
+ if (!inferenceParams.checkpoint || isExternalModel) return null;
+ const activeModelIsGguf =
+ activeGgufVariant != null ||
+ ggufContextLength != null ||
+ inferenceParams.checkpoint.toLowerCase().endsWith(".gguf");
+ return {
+ customContextLength: runtimeCustomContextLength ?? null,
+ maxSeqLength: activeModelIsGguf ? null : inferenceParams.maxSeqLength,
+ kvCacheDtype: runtimeKvCacheDtype ?? null,
+ speculativeType: runtimeSpeculativeType ?? "auto",
+ specDraftNMax: runtimeSpecDraftNMax ?? null,
+ tensorParallel: runtimeTensorParallel ?? false,
+ chatTemplateOverride: runtimeChatTemplateOverride ?? null,
+ };
+ }, [
+ inferenceParams.checkpoint,
+ inferenceParams.maxSeqLength,
+ isExternalModel,
+ activeGgufVariant,
+ ggufContextLength,
+ runtimeCustomContextLength,
+ runtimeKvCacheDtype,
+ runtimeSpeculativeType,
+ runtimeSpecDraftNMax,
+ runtimeTensorParallel,
+ runtimeChatTemplateOverride,
+ ]);
+ const activeModelIsGguf = useMemo(() => {
+ const checkpoint = inferenceParams.checkpoint;
+ if (!checkpoint || isExternalModel) return false;
+ return (
+ activeGgufVariant != null ||
+ ggufContextLength != null ||
+ checkpoint.toLowerCase().endsWith(".gguf")
+ );
+ }, [
+ inferenceParams.checkpoint,
+ isExternalModel,
+ activeGgufVariant,
+ ggufContextLength,
+ ]);
+ const activeModelIsLora = useMemo(() => {
+ const checkpoint = inferenceParams.checkpoint;
+ if (!checkpoint || isExternalModel) return false;
+ const model = modelsFromStore.find((entry) => entry.id === checkpoint);
+ if (model) return model.isLora;
+ const lora = lorasFromStore.find((entry) => entry.id === checkpoint);
+ return lora?.exportType === "lora";
+ }, [inferenceParams.checkpoint, isExternalModel, modelsFromStore, lorasFromStore]);
const reasoningEnabled = useChatRuntimeStore((s) => s.reasoningEnabled);
const reasoningStyle = useChatRuntimeStore((s) => s.reasoningStyle);
const reasoningEffort = useChatRuntimeStore((s) => s.reasoningEffort);
@@ -1783,75 +1826,21 @@ export function ChatPage({
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;
- detachStaged();
- }, [chatContextKey, detachStaged]);
-
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 chatContextKey = `${view.mode}|${activeThreadId ?? ""}|${search.new ?? ""}|${search.project ?? ""}`;
+ const [pendingHubAutoLoad, setPendingHubAutoLoad] =
+ useState(null);
const stageOrLoad = useCallback(
async (selection: SelectedModelInput) => {
const store = useChatRuntimeStore.getState();
- // An un-cached HF repo (GGUF variant or a full non-GGUF snapshot) downloads
- // through the manager first (global indicator), then auto-loads. Everything
- // else -- cached picks, local/native files, LoRA, external -- loads now.
const wantManagerDownload =
isDownloadableHubRepo(selection) && !selection.isDownloaded;
- if (
- (!hasGgufSource(selection) && !wantManagerDownload) ||
- (store.loadOnSelection && selection.isDownloaded)
- ) {
- // Detach any staged pick 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. Detach
- // (not abandon) keeps its download running.
- detachStaged();
- // Load-on-selection skips the sheet, so seed the saved knobs here the
- // way the sheet's restore effect would; the switch would otherwise reset
- // the remembered speculative choice (keepSpeculative below prevents it).
- const remembered = hasGgufSource(selection)
- ? loadRememberedLoadSettings(rememberedLoadSettingsKey(selection))
- : null;
- if (remembered) store.applyRememberedLoadSettings(remembered);
- await selectModel(
- remembered ? { ...selection, keepSpeculative: true } : selection,
- );
- return;
- }
- // Loads can't queue behind each other, but a download is independent: if
- // the pick needs downloading, start it in the manager so it runs alongside
- // the load. Nothing to download (already on device) just waits.
if (store.modelLoading) {
- // Both an uncached non-GGUF snapshot (wantManagerDownload) and an
- // uncached remote GGUF quant download through the manager, so either can
- // run in the background while another model loads. wantManagerDownload
- // excludes GGUF by design, so the GGUF case is checked separately.
const wantBackgroundDownload =
wantManagerDownload ||
(selection.source === "hub" &&
hasGgufSource(selection) &&
!selection.isDownloaded);
- // The model currently loading already downloads as part of its own load
- // (the /load flow fetches before setting the checkpoint), so re-picking
- // it must not kick off a second transfer against the same cache.
const isLoadingThisPick =
!!loadingModel &&
normalizeModelRef(loadingModel.id) ===
@@ -1862,11 +1851,6 @@ export function ChatPage({
description: "It's downloading as part of the load in progress.",
});
} else if (wantBackgroundDownload) {
- // Only claim the download started once a job is actually created. A
- // transport conflict records state that is only resolvable from the
- // Hub download card, so point the user there instead of showing a
- // success toast for a transfer that never began; "busy" and "error"
- // already surface their own toasts.
const outcome = await downloadManager.requestStart({
kind: DOWNLOAD_KIND.MODEL,
repoId: selection.id,
@@ -1883,6 +1867,11 @@ export function ChatPage({
description:
"An earlier partial download used a different transport. Open the Hub tab to resume or restart it.",
});
+ } else if (outcome === "busy") {
+ toast.info("Download already in progress", {
+ description:
+ "Another download for this model is still running. Reselect it once that finishes to load it.",
+ });
}
} else {
toast.info("Another model is already loading", {
@@ -1891,23 +1880,118 @@ export function ChatPage({
}
return;
}
- // Detach the prior staged pick (keeping its download) before rebinding, so
- // a second pick downloads alongside the first instead of cancelling it.
- detachStaged();
- store.stageModel({
- id: selection.id,
- isLora: selection.isLora,
- ggufVariant: selection.ggufVariant,
- isDownloaded: selection.isDownloaded,
- expectedBytes: selection.expectedBytes,
- nativePathToken: selection.nativePathToken,
- isGguf: selection.isGguf,
- isHubRepo: wantManagerDownload || undefined,
- autoLoad: store.loadOnSelection,
+ const wantManagerStage =
+ wantManagerDownload ||
+ (selection.source === "hub" &&
+ hasGgufSource(selection) &&
+ !selection.isDownloaded);
+ if (wantManagerStage) {
+ setPendingHubAutoLoad({
+ selection,
+ contextKey: chatContextKey,
+ originCheckpoint: store.params.checkpoint,
+ originGgufVariant: store.activeGgufVariant,
+ });
+ return;
+ }
+ setPendingHubAutoLoad(null);
+ const previousConfig = currentRuntimePerModelConfig({
+ includeMaxSeqLength: true,
+ });
+ const hasAppliedConfig = applyModelLoadConfigToRuntime(
+ selection.config ?? rememberedConfigFor(selection),
+ );
+ await selectModel({
+ ...selection,
+ ...(hasAppliedConfig ? { keepSpeculative: true } : {}),
+ previousConfig,
});
},
- [detachStaged, selectModel, loadingModel],
+ [selectModel, loadingModel, rememberedConfigFor, chatContextKey],
);
+ useRepoDownload({
+ kind: DOWNLOAD_KIND.MODEL,
+ repoId: pendingHubAutoLoad?.selection.id ?? "__hub_autoload_idle__",
+ activeVariant: pendingHubAutoLoad?.selection.ggufVariant ?? null,
+ onComplete: (variant) => {
+ const pending = pendingHubAutoLoad;
+ if (
+ !pending ||
+ (pending.selection.ggufVariant ?? null) !== (variant ?? null)
+ ) {
+ return;
+ }
+ setPendingHubAutoLoad(null);
+ const store = useChatRuntimeStore.getState();
+ if (
+ !active ||
+ pending.contextKey !== chatContextKey ||
+ normalizeModelRef(pending.originCheckpoint) !==
+ normalizeModelRef(store.params.checkpoint) ||
+ pending.originGgufVariant !== store.activeGgufVariant
+ ) {
+ return;
+ }
+ void stageOrLoad({ ...pending.selection, isDownloaded: true });
+ },
+ onError: (variant) => {
+ if (
+ pendingHubAutoLoad &&
+ (pendingHubAutoLoad.selection.ggufVariant ?? null) === (variant ?? null)
+ ) {
+ setPendingHubAutoLoad(null);
+ }
+ },
+ onCancelled: (variant) => {
+ if (
+ pendingHubAutoLoad &&
+ (pendingHubAutoLoad.selection.ggufVariant ?? null) === (variant ?? null)
+ ) {
+ setPendingHubAutoLoad(null);
+ }
+ },
+ });
+ useEffect(() => {
+ const pending = pendingHubAutoLoad;
+ if (!pending) return;
+ let active = true;
+ void (async () => {
+ const outcome = await downloadManager.requestStart({
+ kind: DOWNLOAD_KIND.MODEL,
+ repoId: pending.selection.id,
+ variant: pending.selection.ggufVariant ?? null,
+ expectedBytes: pending.selection.expectedBytes ?? 0,
+ });
+ if (!active) return;
+ if (outcome === "started") {
+ toast.info("Downloading model", {
+ description: "It'll load automatically once the download finishes.",
+ });
+ return;
+ }
+ if (outcome === "conflict") {
+ // Keep pendingHubAutoLoad bound so this surface's cleanup does not wipe
+ // the conflict just recorded by requestStart (which the toast points the
+ // user to); resolving it from the Hub completes the download and this
+ // surface's onComplete auto-loads, mirroring the "started" branch.
+ toast.info("Resume this download from the Hub", {
+ description:
+ "An earlier partial download used a different transport. Open the Hub tab to resume or restart it.",
+ });
+ return;
+ }
+ if (outcome === "busy") {
+ toast.info("Download already in progress", {
+ description:
+ "Another download for this model is still running. Reselect it once that finishes to load it.",
+ });
+ }
+ setPendingHubAutoLoad((current) => (current === pending ? null : current));
+ })();
+ return () => {
+ active = false;
+ };
+ }, [pendingHubAutoLoad]);
const loadNativeModelIntent = useCallback(
async (intent: NativeIntent, loadingDescription: string) => {
const label =
@@ -1920,6 +2004,11 @@ export function ChatPage({
forceReload: true,
throwOnError: true,
});
+ // Record when this file lease expires so a later reload can prompt
+ // re-selection instead of reusing a token the host has already pruned.
+ useChatRuntimeStore.setState({
+ activeNativePathExpiresAtMs: intent.path.expiresAtMs ?? null,
+ });
useNativeIntentStore.getState().clearModelIntent(intent.id);
},
[stageOrLoad],
@@ -1965,28 +2054,20 @@ export function ChatPage({
const handleCheckpointChange = useCallback(
(
value: string,
- meta?: {
- source?: string;
- isLora: boolean;
- ggufVariant?: string;
- isDownloaded?: boolean;
- expectedBytes?: number;
- isGguf?: boolean;
- },
+ meta?: ModelSelectorChangeMeta,
) => {
const store = useChatRuntimeStore.getState();
const currentCheckpoint = store.params.checkpoint;
const currentVariant = store.activeGgufVariant;
- if (
- !value ||
- (value === currentCheckpoint &&
- (meta?.ggufVariant ?? null) === (currentVariant ?? null))
- )
+ if (!value) return;
+ setPendingHubAutoLoad(null);
+ const isSameLoadedModel =
+ value === currentCheckpoint &&
+ (meta?.ggufVariant ?? null) === (currentVariant ?? null);
+ if (isSameLoadedModel && !meta?.forceReload) {
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(
@@ -2158,19 +2239,17 @@ export function ChatPage({
source: meta?.source,
isLora: meta?.isLora,
ggufVariant: meta?.ggufVariant,
- isDownloaded: meta?.isDownloaded,
+ isDownloaded: meta?.isDownloaded || isSameLoadedModel,
expectedBytes: meta?.expectedBytes,
isGguf: meta?.isGguf,
+ config: meta?.config,
+ nativePathToken: meta?.nativePathToken,
+ forceReload: isSameLoadedModel || undefined,
};
- // "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,
@@ -2178,6 +2257,44 @@ export function ChatPage({
view,
],
);
+ const handleReloadActiveModel = useCallback(
+ (config: PerModelConfig) => {
+ const checkpoint = inferenceParams.checkpoint;
+ if (!checkpoint) return;
+ const runtime = useChatRuntimeStore.getState();
+ const nativeToken = runtime.activeNativePathToken;
+ const nativeExpiry = runtime.activeNativePathExpiresAtMs;
+ // A file-picked GGUF is reachable only via its native path token, which
+ // the desktop host prunes after a TTL. Reusing an expired token makes the
+ // reload fail with an opaque error, so prompt the user to re-select the
+ // file instead.
+ if (nativeToken && nativeExpiry != null && Date.now() >= nativeExpiry) {
+ toast.error("This local model file's access has expired.", {
+ description: "Re-select the model file to reload it.",
+ });
+ return;
+ }
+ handleCheckpointChange(checkpoint, {
+ source: "local",
+ isLora: activeModelIsLora,
+ ggufVariant: activeGgufVariant ?? undefined,
+ // Without the native token the reload validates the display label as a
+ // repo and fails.
+ nativePathToken: nativeToken ?? undefined,
+ isGguf: activeModelIsGguf,
+ isDownloaded: true,
+ config,
+ forceReload: true,
+ });
+ },
+ [
+ inferenceParams.checkpoint,
+ activeGgufVariant,
+ activeModelIsLora,
+ activeModelIsGguf,
+ handleCheckpointChange,
+ ],
+ );
const handleEject = useCallback(() => {
void (async () => {
if (await ejectModel()) {
@@ -2580,6 +2697,8 @@ export function ChatPage({
externalModels={externalModels}
value={inferenceParams.checkpoint}
activeGgufVariant={activeGgufVariant}
+ activeModelConfig={activeModelConfig}
+ activeGgufContextLength={ggufContextLength}
onValueChange={handleCheckpointChange}
onEject={handleEject}
onFoldersChange={refreshLocalModels}
@@ -2633,7 +2752,12 @@ export function ChatPage({
stageOrLoad(selection)}
+ onLoad={() =>
+ loadNativeModelIntent(
+ pendingNativeModelIntent,
+ "Loading selected local GGUF model.",
+ )
+ }
/>
) : null}
{loadingModel && loadToastDismissed ? (
@@ -2790,13 +2914,22 @@ export function ChatPage({
open={active && settingsOpen}
onOpenChange={(open) => {
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}
+ modelConfig={
+ view.mode !== "compare" && activeModelConfig && !modelLoading ? (
+
+ ) : null
+ }
isExternalModel={isExternalModel}
providerCapabilities={activeProviderCapabilities}
activeExternalProvider={activeExternalProvider}
@@ -2808,62 +2941,6 @@ export function ChatPage({
);
}}
externalProviderType={activeExternalProviderType}
- loadingModel={loadingModel}
- onReloadModel={() => {
- const state = useChatRuntimeStore.getState();
- if (state.params.checkpoint) {
- selectModel({
- id: state.params.checkpoint,
- ggufVariant: state.activeGgufVariant ?? undefined,
- forceReload: true,
- isDownloaded: true,
- loadingDescription: "Reloading with updated chat template.",
- });
- }
- }}
- 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. 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…): the pick is
- // cleared only on success, so it normally stays staged with edited
- // knobs intact — nothing to restore.
- const store = useChatRuntimeStore.getState();
- // Still staged (this pick, or a newer one queued meanwhile): leave it.
- if (store.pendingSelection) return;
- // Cleared mid-load (sheet closed / switched chats). Re-stage 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,
- )
- }
/>