diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 63c97d657e..6d8a109ef0 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -15,6 +15,8 @@ import { isHiddenModelId, } from "@/features/hub/lib/hidden-models"; import { resolveInitialConfig } from "@/features/model-picker"; +import { isMlxId } from "@/features/model-picker/components/model-selector/recommended-fit"; +import { fetchDeviceType, usePlatformStore } from "@/config/env"; import { projectHasSources } from "@/features/rag/api/rag-api"; import { apiUrl } from "@/lib/api-base"; import { parseParamCountB } from "@/lib/model-size"; @@ -1479,15 +1481,51 @@ const AUTO_LOAD_LOCAL_SOURCES: ReadonlySet = new Set([ "custom", ]); +/** The picker's chat-only platform snapshot, read once per auto-load run. */ +type AutoLoadPlatform = { + chatOnly: boolean; + isMac: boolean; +}; + +// Mirrors the picker's localModelIsGguf / localModelIsMlx checks: the backend +// format hint is authoritative for indexed rows, with the same name/path +// fallbacks the picker applies. +function localRowIsGgufLike(row: LocalModelInfo): boolean { + return ( + row.model_format === "gguf" || row.path.toLowerCase().endsWith(".gguf") + ); +} + +function localRowIsMlxNamed(row: LocalModelInfo): boolean { + return ( + isMlxId(row.id) || + isMlxId(row.display_name ?? "") || + isMlxId(row.model_id ?? "") + ); +} + /** * Backend-indexed local rows eligible for background auto-load: same policy * as the on-device picker (complete, chat-capable, not hidden infra), plus * no variant requirement, since a background load cannot ask for a quant. */ -function isAutoLoadableLocalRow(row: LocalModelInfo): boolean { +function isAutoLoadableLocalRow( + row: LocalModelInfo, + platform: AutoLoadPlatform, +): boolean { if (!AUTO_LOAD_LOCAL_SOURCES.has(row.source)) return false; if (row.capabilities?.can_chat !== true) return false; if (row.partial) return false; + // Chat-only installs run GGUF (any host) and MLX (Mac only); the picker + // hides other local formats there, so the background load must not pick a + // row the user could not have selected (mirrors sortedLocalDir's gate). + if ( + platform.chatOnly && + !localRowIsGgufLike(row) && + !(platform.isMac && localRowIsMlxNamed(row)) + ) { + return false; + } // Adapters are chat-capable but load by resolving their base model, which // for a Hub-id base can trigger the implicit remote fetch a background // auto-load must never start. Adapters stay interactive-only. @@ -2007,9 +2045,20 @@ export async function autoLoadOnDeviceModel(): Promise<{ }; } + // Resolve the platform snapshot the picker's format gates key on. Cached + // after the app's boot fetch and never throws (falls back to client-side + // detection when the backend is not ready). + await fetchDeviceType(); + const platformState = usePlatformStore.getState(); + const platform: AutoLoadPlatform = { + chatOnly: platformState.isChatOnly(), + isMac: platformState.deviceType === "mac", + }; const ggufRepos = allGgufRepos.filter(isAutoLoadableCachedRepo); const modelRepos = allModelRepos.filter(isAutoLoadableCachedRepo); - const localRows = allLocalRows.filter(isAutoLoadableLocalRow); + const localRows = allLocalRows.filter((row) => + isAutoLoadableLocalRow(row, platform), + ); // Dedupe candidates that resolve to the SAME load target (e.g. a custom // scan folder pointing into an HF cache). Keyed on kind + load target / // on-disk path: a shared model_id does not mean the same files (a distinct @@ -2218,8 +2267,12 @@ export async function autoLoadOnDeviceModel(): Promise<{ } readyPool.splice(at, 0, entry); }; - // Non-GGUF cached repos need no scan: their snapshot loads whole. - for (const repo of modelRepos) { + // Non-GGUF cached repos need no scan: their snapshot loads whole. The + // picker hides cached non-GGUF rows entirely on chat-only installs, so + // the automatic cascade must not pick one there either; the remembered + // path above stays ungated since a recorded load is user precedent that + // the model runs on this install (e.g. an MLX repo loaded on a Mac). + for (const repo of platform.chatOnly ? [] : modelRepos) { insertReady({ type: "cached-model", repo, @@ -2335,6 +2388,11 @@ export async function autoLoadOnDeviceModel(): Promise<{ } } let pendingJobs = resolutionJobs.length; + // Once auto-load reaches a terminal result (a model loaded, the attempt + // cap was hit, or the pool drained), the workers stop claiming jobs so a + // successful early load does not leave background scans hammering the + // backend while inference is already running. In-flight scans finish. + let resolutionStopped = false; let progressWaiters: Array<() => void> = []; const signalProgress = (): void => { const waiters = progressWaiters; @@ -2356,7 +2414,7 @@ export async function autoLoadOnDeviceModel(): Promise<{ ), }, async () => { - while (nextJob < resolutionJobs.length) { + while (!resolutionStopped && nextJob < resolutionJobs.length) { const job = resolutionJobs[nextJob]; nextJob += 1; let entry: FallbackCandidate | null = null; @@ -2385,151 +2443,165 @@ export async function autoLoadOnDeviceModel(): Promise<{ }); }); } - while (loadAttempts < MAX_AUTO_LOAD_ATTEMPTS) { - const candidate = readyPool[0]; - if (!candidate) { - if (pendingJobs <= 0) { - break; + try { + while (loadAttempts < MAX_AUTO_LOAD_ATTEMPTS) { + const candidate = readyPool[0]; + if (!candidate) { + if (pendingJobs <= 0) { + break; + } + await nextProgress(); + continue; } - await nextProgress(); - continue; - } - // Every pending scan job can still yield a GGUF candidate (no-scan - // rows were seeded upfront), and GGUF outranks safetensors in the - // documented order, so the model-kind group stays gated until the - // scans settle; resolved GGUF entries keep flowing immediately. - if (isModelKindEntry(candidate) && pendingJobs > 0) { - await nextProgress(); - continue; - } - readyPool.shift(); - if (candidate.type === "cached-gguf") { - const repo = candidate.repo; - if (!candidate.retry) { - // A shared load target may already have been visited through an - // indexed local row (e.g. a scan folder aliasing this cache). - if (isSeen("gguf", repo.load_id || repo.repo_id, repo.cache_path)) { + // The final attempt is precious: while scans are still pending, a + // smaller candidate can still enter the pool, so the last slot is + // not spent until resolution settles and the global order is + // complete. Earlier attempts keep flowing incrementally. + if (pendingJobs > 0 && loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS - 1) { + await nextProgress(); + continue; + } + // Every pending scan job can still yield a GGUF candidate (no-scan + // rows were seeded upfront), and GGUF outranks safetensors in the + // documented order, so the model-kind group stays gated until the + // scans settle; resolved GGUF entries keep flowing immediately. + if (isModelKindEntry(candidate) && pendingJobs > 0) { + await nextProgress(); + continue; + } + readyPool.shift(); + if (candidate.type === "cached-gguf") { + const repo = candidate.repo; + if (!candidate.retry) { + // A shared load target may already have been visited through an + // indexed local row (e.g. a scan folder aliasing this cache). + if (isSeen("gguf", repo.load_id || repo.repo_id, repo.cache_path)) { + continue; + } + markSeen("gguf", repo.load_id || repo.repo_id, repo.cache_path); + } + const skipKey = autoLoadCandidateKey( + "gguf", + repo.load_id || repo.repo_id, + candidate.variant.quant, + ); + if (skippedAutoLoadCandidates.has(skipKey)) { continue; } - markSeen("gguf", repo.load_id || repo.repo_id, repo.cache_path); + try { + if ( + await loadAutoLoadCandidate({ + id: repo.repo_id, + loadId: repo.load_id, + kind: "gguf", + ggufVariant: candidate.variant.quant, + maxSeqLength: 0, + successLabel: `Loaded ${repo.repo_id} (${candidate.variant.quant})`, + inventoryId: repo.inventory_id ?? null, + source: "hf_cache", + }) + ) { + return { loaded: true, blockedByTrustRemoteCode: false }; + } + } catch { + hadNonTrustFailure = true; + skippedAutoLoadCandidates.add(skipKey); + // A quant that passed validation can still fail /load (corrupt + // file, llama.cpp startup error). Re-enter the repo's next + // complete quant into the global size order, so one repo of + // failing quants cannot starve a smaller model elsewhere. + // Validation blocks are model-scoped, so they get no requeue. + try { + const next = await resolveCachedGgufEntry(repo); + if (next) { + insertReady({ ...next, retry: true }); + } + } catch { + hadNonTrustFailure = true; + } + } + continue; } - const skipKey = autoLoadCandidateKey( - "gguf", - repo.load_id || repo.repo_id, - candidate.variant.quant, - ); - if (skippedAutoLoadCandidates.has(skipKey)) { + if (candidate.type === "cached-model") { + const repo = candidate.repo; + if (isSeen("model", repo.load_id || repo.repo_id, repo.cache_path)) { + continue; + } + markSeen("model", repo.load_id || repo.repo_id, repo.cache_path); + if ( + skippedAutoLoadCandidates.has( + autoLoadCandidateKey("model", repo.load_id || repo.repo_id), + ) + ) { + continue; + } + try { + if ( + await loadAutoLoadCandidate({ + id: repo.repo_id, + loadId: repo.load_id, + kind: "model", + ggufVariant: null, + maxSeqLength: 4096, + successLabel: `Loaded ${repo.repo_id}`, + inventoryId: repo.inventory_id ?? null, + source: "hf_cache", + }) + ) { + return { loaded: true, blockedByTrustRemoteCode: false }; + } + } catch { + hadNonTrustFailure = true; + skippedAutoLoadCandidates.add( + autoLoadCandidateKey("model", repo.load_id || repo.repo_id), + ); + } + continue; + } + const row = candidate.row; + const localCandidate = candidate.candidate; + if (!candidate.retry) { + if (isSeen(localCandidate.kind, row.load_id, row.id, row.path)) { + continue; + } + markSeen(localCandidate.kind, row.load_id, row.id, row.path); + } + if (isSkippedAutoLoadCandidate(localCandidate)) { continue; } try { - if ( - await loadAutoLoadCandidate({ - id: repo.repo_id, - loadId: repo.load_id, - kind: "gguf", - ggufVariant: candidate.variant.quant, - maxSeqLength: 0, - successLabel: `Loaded ${repo.repo_id} (${candidate.variant.quant})`, - inventoryId: repo.inventory_id ?? null, - source: "hf_cache", - }) - ) { + if (await loadAutoLoadCandidate(localCandidate)) { return { loaded: true, blockedByTrustRemoteCode: false }; } } catch { hadNonTrustFailure = true; - skippedAutoLoadCandidates.add(skipKey); - // A quant that passed validation can still fail /load (corrupt - // file, llama.cpp startup error). Re-enter the repo's next - // complete quant into the global size order, so one repo of - // failing quants cannot starve a smaller model elsewhere. - // Validation blocks are model-scoped, so they get no requeue. + skippedAutoLoadCandidates.add(autoLoadSkipKey(localCandidate)); + // Same requeue as the cached-gguf branch: the folder's next complete + // quant re-enters the global size order instead of retrying inline. try { - const next = await resolveCachedGgufEntry(repo); + const next = await resolveLocalRowCandidate( + row, + null, + isSkippedAutoLoadCandidate, + ); if (next) { - insertReady({ ...next, retry: true }); + insertReady({ + type: "local", + row, + candidate: next.candidate, + sizeBytes: next.sizeBytes, + retry: true, + }); } } catch { hadNonTrustFailure = true; } } - continue; - } - if (candidate.type === "cached-model") { - const repo = candidate.repo; - if (isSeen("model", repo.load_id || repo.repo_id, repo.cache_path)) { - continue; - } - markSeen("model", repo.load_id || repo.repo_id, repo.cache_path); - if ( - skippedAutoLoadCandidates.has( - autoLoadCandidateKey("model", repo.load_id || repo.repo_id), - ) - ) { - continue; - } - try { - if ( - await loadAutoLoadCandidate({ - id: repo.repo_id, - loadId: repo.load_id, - kind: "model", - ggufVariant: null, - maxSeqLength: 4096, - successLabel: `Loaded ${repo.repo_id}`, - inventoryId: repo.inventory_id ?? null, - source: "hf_cache", - }) - ) { - return { loaded: true, blockedByTrustRemoteCode: false }; - } - } catch { - hadNonTrustFailure = true; - skippedAutoLoadCandidates.add( - autoLoadCandidateKey("model", repo.load_id || repo.repo_id), - ); - } - continue; - } - const row = candidate.row; - const localCandidate = candidate.candidate; - if (!candidate.retry) { - if (isSeen(localCandidate.kind, row.load_id, row.id, row.path)) { - continue; - } - markSeen(localCandidate.kind, row.load_id, row.id, row.path); - } - if (isSkippedAutoLoadCandidate(localCandidate)) { - continue; - } - try { - if (await loadAutoLoadCandidate(localCandidate)) { - return { loaded: true, blockedByTrustRemoteCode: false }; - } - } catch { - hadNonTrustFailure = true; - skippedAutoLoadCandidates.add(autoLoadSkipKey(localCandidate)); - // Same requeue as the cached-gguf branch: the folder's next complete - // quant re-enters the global size order instead of retrying inline. - try { - const next = await resolveLocalRowCandidate( - row, - null, - isSkippedAutoLoadCandidate, - ); - if (next) { - insertReady({ - type: "local", - row, - candidate: next.candidate, - sizeBytes: next.sizeBytes, - retry: true, - }); - } - } catch { - hadNonTrustFailure = true; - } } + } finally { + // Runs on every exit (successful return, cap, drained pool, or a + // thrown error) so no worker keeps scanning after the outcome is set. + resolutionStopped = true; } // No auto-loadable on-device model (or the attempt cap was hit). Never diff --git a/tests/studio/test_model_picker_contracts.py b/tests/studio/test_model_picker_contracts.py index 11b558d07f..2b6d67dcc7 100644 --- a/tests/studio/test_model_picker_contracts.py +++ b/tests/studio/test_model_picker_contracts.py @@ -846,7 +846,9 @@ def test_fallback_orders_by_resolved_quant_size(): assert "sizeBytes: sizeOrUnknownBytes(variant.size_bytes)" in auto_load # The all-variant row sum only orders non-GGUF cached repos, whose # snapshot loads whole. - seed_block = auto_load.split("for (const repo of modelRepos)", 1)[1] + seed_block = auto_load.split( + "for (const repo of platform.chatOnly ? [] : modelRepos)", 1 + )[1] seed_block = seed_block.split("const resolveCachedGgufEntry", 1)[0] assert "sizeOrUnknownBytes(repo.size_bytes)" in seed_block assert auto_load.count("sizeOrUnknownBytes(repo.size_bytes)") == 1 @@ -932,6 +934,57 @@ def test_pending_gguf_scans_gate_safetensors_candidates(): assert "if (isModelKindEntry(candidate) && pendingJobs > 0) {" in auto_load +def test_final_attempt_waits_for_pending_scans(): + """Fast-resolving candidates whose loads fail must not exhaust the attempt + cap while pending scans can still yield a smaller loadable quant: the + final attempt is only spent once resolution has settled and the global + order is complete.""" + src = _read("features/chat/api/chat-adapter.ts") + auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1] + assert ( + "if (pendingJobs > 0 && loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS - 1) {" + in auto_load + ) + + +def test_resolution_workers_stop_on_terminal_result(): + """A successful early load (or any other terminal outcome) must stop the + workers from claiming further folder scans, so autoload cannot leave + background scans contending with inference for the backend and disk.""" + src = _read("features/chat/api/chat-adapter.ts") + auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1] + assert "let resolutionStopped = false;" in auto_load + assert ( + "while (!resolutionStopped && nextJob < resolutionJobs.length)" + in auto_load + ) + # The flag is set in a finally so every exit path (return, break, throw) + # stops the workers. + assert "resolutionStopped = true;" in auto_load + assert auto_load.index("} finally {") < auto_load.index( + "resolutionStopped = true;" + ) + + +def test_local_rows_apply_picker_platform_gate(): + """Chat-only installs run GGUF (any host) and MLX (Mac only); the picker + hides other local formats and all cached non-GGUF rows there, so the + background cascade must not load a row the user could not have picked. + The remembered path stays ungated: a recorded load is user precedent.""" + src = _read("features/chat/api/chat-adapter.ts") + local_fn = src.split("function isAutoLoadableLocalRow", 1)[1] + local_fn = local_fn.split("\nfunction ", 1)[0] + assert "platform.chatOnly" in local_fn + assert "localRowIsGgufLike(row)" in local_fn + assert "platform.isMac && localRowIsMlxNamed(row)" in local_fn + auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1] + # The platform snapshot hydrates through the cached, non-throwing fetch. + assert "await fetchDeviceType();" in auto_load + # Cascade seeding of cached non-GGUF repos mirrors the picker's + # chat-only exclusion; the remembered lookup above it stays unfiltered. + assert "platform.chatOnly ? [] : modelRepos" in auto_load + + def test_autoload_keys_preserve_posix_path_case(): """Linux filesystems distinguish /models/Foo from /models/foo, so seen keys and remembered-model matching must not fold case on POSIX paths;