diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index b5455e6437..63c97d657e 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -2268,27 +2268,72 @@ export async function autoLoadOnDeviceModel(): Promise<{ row.model_format === "gguf" || row.capabilities?.requires_variant !== true, ); - const resolutionJobs: Array<() => Promise> = [ - ...ggufRepos.map((repo) => () => resolveCachedGgufEntry(repo)), - ...cascadeLocalRows.map( - (row) => async (): Promise => { - const resolved = await resolveLocalRowCandidate( - row, - null, - isSkippedAutoLoadCandidate, - ); - if (!resolved) { - return null; + // Only GGUF directory rows hit the backend folder scan; everything else + // resolves in-process and is seeded straight into the pool so it can + // never queue behind slow scans. + const needsVariantScan = (row: LocalModelInfo): boolean => + row.model_format === "gguf" && + row.capabilities?.requires_variant === true; + await Promise.all( + cascadeLocalRows + .filter((row) => !needsVariantScan(row)) + .map(async (row) => { + try { + const resolved = await resolveLocalRowCandidate( + row, + null, + isSkippedAutoLoadCandidate, + ); + if (resolved) { + insertReady({ + type: "local", + row, + candidate: resolved.candidate, + sizeBytes: resolved.sizeBytes, + }); + } + } catch { + hadNonTrustFailure = true; } - return { - type: "local", - row, - candidate: resolved.candidate, - sizeBytes: resolved.sizeBytes, - }; - }, - ), - ]; + }), + ); + // The remaining jobs all need a backend scan and can all yield GGUF + // candidates. Cached repos and local folders interleave so a run of + // slow scans from one source cannot monopolize every worker. + const cachedScanJobs = ggufRepos.map( + (repo) => () => resolveCachedGgufEntry(repo), + ); + const localScanJobs = cascadeLocalRows.filter(needsVariantScan).map( + (row) => async (): Promise => { + const resolved = await resolveLocalRowCandidate( + row, + null, + isSkippedAutoLoadCandidate, + ); + if (!resolved) { + return null; + } + return { + type: "local", + row, + candidate: resolved.candidate, + sizeBytes: resolved.sizeBytes, + }; + }, + ); + const resolutionJobs: Array<() => Promise> = []; + for ( + let jobIndex = 0; + jobIndex < Math.max(cachedScanJobs.length, localScanJobs.length); + jobIndex += 1 + ) { + if (jobIndex < cachedScanJobs.length) { + resolutionJobs.push(cachedScanJobs[jobIndex]); + } + if (jobIndex < localScanJobs.length) { + resolutionJobs.push(localScanJobs[jobIndex]); + } + } let pendingJobs = resolutionJobs.length; let progressWaiters: Array<() => void> = []; const signalProgress = (): void => { @@ -2341,7 +2386,7 @@ export async function autoLoadOnDeviceModel(): Promise<{ }); } while (loadAttempts < MAX_AUTO_LOAD_ATTEMPTS) { - const candidate = readyPool.shift(); + const candidate = readyPool[0]; if (!candidate) { if (pendingJobs <= 0) { break; @@ -2349,6 +2394,15 @@ export async function autoLoadOnDeviceModel(): Promise<{ 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) { diff --git a/tests/studio/test_model_picker_contracts.py b/tests/studio/test_model_picker_contracts.py index 586514b352..11b558d07f 100644 --- a/tests/studio/test_model_picker_contracts.py +++ b/tests/studio/test_model_picker_contracts.py @@ -909,6 +909,27 @@ def test_send_not_blocked_by_full_inventory_resolution(): assert "clearTimeout(graceTimer)" in auto_load assert "await nextProgress();" in auto_load assert "if (pendingJobs <= 0) {" in auto_load + # Rows that need no backend scan seed the pool before the workers start, + # so they can never queue behind slow folder scans. + assert "const needsVariantScan" in auto_load + assert ".filter((row) => !needsVariantScan(row))" in auto_load + assert auto_load.index(".filter((row) => !needsVariantScan(row))") < auto_load.index( + "const cachedScanJobs" + ) + # Cached and local scans interleave so one slow source cannot + # monopolize every worker. + assert "resolutionJobs.push(cachedScanJobs[jobIndex])" in auto_load + assert "resolutionJobs.push(localScanJobs[jobIndex])" in auto_load + + +def test_pending_gguf_scans_gate_safetensors_candidates(): + """GGUF-first is the documented preference order, and incremental + consumption must not let an instantly-resolved safetensors row claim a + load slot while a pending folder scan can still yield a GGUF candidate; + resolved GGUF entries are never gated.""" + src = _read("features/chat/api/chat-adapter.ts") + auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1] + assert "if (isModelKindEntry(candidate) && pendingJobs > 0) {" in auto_load def test_autoload_keys_preserve_posix_path_case():