diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 41d92f2df0..2b02ffdd99 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1427,7 +1427,9 @@ function autoLoadCandidateKey( id: string, ggufVariant?: string | null, ): string { - return `${kind}:${id.toLowerCase()}:${(ggufVariant ?? "").toLowerCase()}`; + // Path-shape-aware case handling: on a case-sensitive filesystem, a skip + // key recorded for /models/Foo must not also skip /models/foo. + return `${kind}:${normalizeLoadTargetKey(id)}:${(ggufVariant ?? "").toLowerCase()}`; } function findCachedRepo( @@ -1525,6 +1527,46 @@ function sizeOrUnknownBytes(bytes?: number | null): number { return bytes && bytes > 0 ? bytes : Number.MAX_SAFE_INTEGER; } +/** + * Identifier-matching key with path-shape-aware case semantics: Windows-style + * paths (drive letter or UNC) and Hub repo ids compare case-insensitively, + * while POSIX paths keep their case, since Linux filesystems distinguish + * /models/Foo from /models/foo and folding them can match the wrong model. + */ +function normalizeLoadTargetKey(value: string): string { + const looksWindowsPath = /^(?:[A-Za-z]:[\\/]|\\\\)/.test(value); + if (!looksWindowsPath && (value.startsWith("/") || value.startsWith("~"))) { + return value; + } + return value.toLowerCase(); +} + +// Inventory scans hit disk on the backend; an unbounded fan-out over many +// indexed folders can saturate the connection pool and disk. +const AUTO_LOAD_VARIANT_SCAN_CONCURRENCY = 4; + +/** Map with at most `limit` requests in flight, preserving order. */ +async function mapWithConcurrency( + items: readonly T[], + limit: number, + fn: (item: T) => Promise, +): Promise { + const results: R[] = new Array(items.length); + let nextIndex = 0; + const workers = Array.from( + { length: Math.max(1, Math.min(limit, items.length)) }, + async () => { + while (nextIndex < items.length) { + const index = nextIndex; + nextIndex += 1; + results[index] = await fn(items[index]); + } + }, + ); + await Promise.all(workers); + return results; +} + type ResolvedLocalCandidate = { candidate: AutoLoadCandidate; /** Size of what would actually load: the resolved quant's own size for a @@ -1589,20 +1631,23 @@ function matchesRememberedLocalRow( if ((row.model_format === "gguf") !== (remembered.kind === "gguf")) { return false; } + // Inventory ids come from one backend generator on both sides, so they + // compare exactly; case folding could merge distinct case-sensitive paths + // embedded in the id. if ( remembered.inventoryId && row.inventory_id && - row.inventory_id.toLowerCase() === remembered.inventoryId.toLowerCase() + row.inventory_id === remembered.inventoryId ) { return true; } const targets = new Set( [remembered.loadId, remembered.id] .filter((value): value is string => Boolean(value)) - .map((value) => value.toLowerCase()), + .map((value) => normalizeLoadTargetKey(value)), ); return [row.load_id, row.id, row.path, row.model_id].some( - (value) => !!value && targets.has(value.toLowerCase()), + (value) => !!value && targets.has(normalizeLoadTargetKey(value)), ); } @@ -1983,7 +2028,9 @@ export async function autoLoadOnDeviceModel(): Promise<{ ...values: (string | null | undefined)[] ): void => { for (const value of values) { - if (value) seenLoadTargets.add(`${kind}:${value.toLowerCase()}`); + if (value) { + seenLoadTargets.add(`${kind}:${normalizeLoadTargetKey(value)}`); + } } }; const isSeen = ( @@ -1991,7 +2038,9 @@ export async function autoLoadOnDeviceModel(): Promise<{ ...values: (string | null | undefined)[] ): boolean => values.some( - (value) => !!value && seenLoadTargets.has(`${kind}:${value.toLowerCase()}`), + (value) => + !!value && + seenLoadTargets.has(`${kind}:${normalizeLoadTargetKey(value)}`), ); try { @@ -2127,6 +2176,8 @@ export async function autoLoadOnDeviceModel(): Promise<{ row: LocalModelInfo; candidate: AutoLoadCandidate; sizeBytes: number; + /** Re-queued quant of an already-visited row (skips the seen gate). */ + retry?: boolean; }; const bySizeAsc = (a: FallbackCandidate, b: FallbackCandidate): number => a.sizeBytes - b.sizeBytes; @@ -2146,28 +2197,28 @@ export async function autoLoadOnDeviceModel(): Promise<{ // resolved quant's own size or a folder with a small quant would lose to // a larger single-quant model. const localEntries = ( - await Promise.all( - cascadeLocalRows.map( - async (row): Promise => { - try { - const resolved = await resolveLocalRowCandidate( - row, - null, - isSkippedAutoLoadCandidate, - ); - if (!resolved) return null; - return { - type: "local" as const, - row, - candidate: resolved.candidate, - sizeBytes: resolved.sizeBytes, - }; - } catch { - hadNonTrustFailure = true; - return null; - } - }, - ), + await mapWithConcurrency( + cascadeLocalRows, + AUTO_LOAD_VARIANT_SCAN_CONCURRENCY, + async (row): Promise => { + try { + const resolved = await resolveLocalRowCandidate( + row, + null, + isSkippedAutoLoadCandidate, + ); + if (!resolved) return null; + return { + type: "local" as const, + row, + candidate: resolved.candidate, + sizeBytes: resolved.sizeBytes, + }; + } catch { + hadNonTrustFailure = true; + return null; + } + }, ) ).filter((entry): entry is FallbackCandidate => entry !== null); const ggufGroup: FallbackCandidate[] = [ @@ -2191,7 +2242,12 @@ export async function autoLoadOnDeviceModel(): Promise<{ ), ].sort(bySizeAsc); - for (const candidate of [...ggufGroup, ...modelGroup]) { + const queue: FallbackCandidate[] = [...ggufGroup, ...modelGroup]; + const isModelKindEntry = (entry: FallbackCandidate): boolean => + entry.type === "cached-model" || + (entry.type === "local" && entry.candidate.kind === "model"); + for (let queueIndex = 0; queueIndex < queue.length; queueIndex += 1) { + const candidate = queue[queueIndex]; if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) break; if (candidate.type === "cached-gguf") { const repo = candidate.repo; @@ -2266,41 +2322,61 @@ export async function autoLoadOnDeviceModel(): Promise<{ continue; } const row = candidate.row; - let localCandidate: AutoLoadCandidate | null = candidate.candidate; - if (isSeen(localCandidate.kind, row.load_id, row.id, row.path)) { + 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; } - markSeen(localCandidate.kind, row.load_id, row.id, row.path); - // Try the row's quants smallest-first: a failed LOAD (not just a - // blocked validation) marks that quant skipped and the folder's next - // complete quant is resolved and tried, so one corrupt file cannot - // abandon a folder that still holds a loadable quant. The attempt cap - // still bounds total /load calls; single-candidate rows resolve to - // null once skipped, terminating the loop. - while (localCandidate && loadAttempts < MAX_AUTO_LOAD_ATTEMPTS) { - if (!isSkippedAutoLoadCandidate(localCandidate)) { - try { - if (await loadAutoLoadCandidate(localCandidate)) { - return { loaded: true, blockedByTrustRemoteCode: false }; - } - } catch { - hadNonTrustFailure = true; - skippedAutoLoadCandidates.add( - autoLoadCandidateKey( - localCandidate.kind, - localCandidate.id, - localCandidate.ggufVariant, - ), - ); - } + try { + if (await loadAutoLoadCandidate(localCandidate)) { + return { loaded: true, blockedByTrustRemoteCode: false }; } + } catch { + hadNonTrustFailure = true; + skippedAutoLoadCandidates.add( + autoLoadCandidateKey( + localCandidate.kind, + localCandidate.id, + localCandidate.ggufVariant, + ), + ); + // A quant that passed validation can still fail /load (corrupt + // file, llama.cpp startup error). Re-enter the folder's next + // complete quant into the GLOBAL size order (still ahead of the + // safetensors group) instead of retrying inline, so one folder of + // failing quants cannot starve a smaller model elsewhere. + // Validation blocks are model-scoped, so they get no requeue. try { - localCandidate = - (await resolveLocalRowCandidate(row, null, isSkippedAutoLoadCandidate)) - ?.candidate ?? null; + const next = await resolveLocalRowCandidate( + row, + null, + isSkippedAutoLoadCandidate, + ); + if (next) { + const retryEntry: FallbackCandidate = { + type: "local", + row, + candidate: next.candidate, + sizeBytes: next.sizeBytes, + retry: true, + }; + let insertAt = queueIndex + 1; + while ( + insertAt < queue.length && + !isModelKindEntry(queue[insertAt]) && + queue[insertAt].sizeBytes <= retryEntry.sizeBytes + ) { + insertAt += 1; + } + queue.splice(insertAt, 0, retryEntry); + } } catch { hadNonTrustFailure = true; - break; } } } diff --git a/tests/studio/test_model_picker_contracts.py b/tests/studio/test_model_picker_contracts.py index 06295a161f..d7068accba 100644 --- a/tests/studio/test_model_picker_contracts.py +++ b/tests/studio/test_model_picker_contracts.py @@ -574,7 +574,9 @@ def test_autoload_deduplicates_cached_and_local_candidates(): assert "const seenLoadTargets = new Set()" in auto_load # Keys carry the model kind: a folder emitting both GGUF and safetensors # rows shares a path while holding two different models. - assert "seenLoadTargets.add(`${kind}:${value.toLowerCase()}`)" in auto_load + assert ( + "seenLoadTargets.add(`${kind}:${normalizeLoadTargetKey(value)}`)" in auto_load + ) assert 'markSeen("gguf", repo.load_id || repo.repo_id, repo.cache_path)' in auto_load assert 'markSeen("model", repo.load_id || repo.repo_id, repo.cache_path)' in auto_load assert "isSeen(localCandidate.kind, row.load_id, row.id, row.path)" in auto_load @@ -712,18 +714,66 @@ def test_local_fallback_orders_by_resolved_quant_size(): def test_cascade_retries_next_quant_after_load_failure(): """A failed /api/inference/load (not just a blocked validation) must mark - that quant skipped and try the folder's next complete quant before the - row is abandoned; single-candidate rows resolve to null once skipped so - the retry loop terminates, and the attempt cap bounds total loads.""" + that quant skipped and re-enter the folder's next complete quant into the + GLOBAL size order (still ahead of the safetensors group) instead of + retrying inline, so one folder of failing quants cannot starve a smaller + model elsewhere; single-candidate rows resolve to null once skipped, and + the attempt cap bounds total loads.""" src = _read("features/chat/api/chat-adapter.ts") auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1] - assert "while (localCandidate && loadAttempts < MAX_AUTO_LOAD_ATTEMPTS)" in auto_load - # The cascade catch records the failed quant, unlike the old generic flag. - local_loop = auto_load.split( - "while (localCandidate && loadAttempts < MAX_AUTO_LOAD_ATTEMPTS)", 1 - )[1].split("\n }", 1)[0] - assert "skippedAutoLoadCandidates.add(" in local_loop + # No inline retry loop: retries flow through the shared queue. + assert "while (localCandidate" not in auto_load + assert "const queue: FallbackCandidate[] = [...ggufGroup, ...modelGroup]" in auto_load + assert "retry: true," in auto_load + assert "queue.splice(insertAt, 0, retryEntry)" in auto_load + # Reinsertion respects the GGUF-before-safetensors group boundary and the + # ascending size order among the remaining candidates. + assert "!isModelKindEntry(queue[insertAt])" in auto_load + assert "queue[insertAt].sizeBytes <= retryEntry.sizeBytes" in auto_load + # Requeued entries bypass the seen gate; fresh rows still dedupe. + assert "if (!candidate.retry) {" in auto_load + # The cascade catch records the failed quant before requeueing. + catch_block = auto_load.split("// A quant that passed validation can still fail /load", 1)[0] + assert "skippedAutoLoadCandidates.add(" in catch_block # Termination guard: a skipped single candidate resolves to null. resolve_fn = src.split("async function resolveLocalRowCandidate", 1)[1] resolve_fn = resolve_fn.split("\nfunction ", 1)[0] assert "if (isSkippedCandidate?.(candidate)) return null;" in resolve_fn + + +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; + Windows-style paths and Hub repo ids keep case-insensitive matching.""" + src = _read("features/chat/api/chat-adapter.ts") + norm_fn = src.split("function normalizeLoadTargetKey", 1)[1] + norm_fn = norm_fn.split("\nfunction ", 1)[0].split("\nconst ", 1)[0] + assert "looksWindowsPath" in norm_fn + assert 'value.startsWith("/") || value.startsWith("~")' in norm_fn + assert "return value;" in norm_fn + assert "return value.toLowerCase();" in norm_fn + match_fn = src.split("function matchesRememberedLocalRow", 1)[1] + match_fn = match_fn.split("\nfunction ", 1)[0] + assert "normalizeLoadTargetKey" in match_fn + # Inventory ids compare exactly (same backend generator on both sides). + assert "row.inventory_id === remembered.inventoryId" in match_fn + assert "row.inventory_id.toLowerCase()" not in match_fn + # Skip keys use the same semantics: a failure recorded for /models/Foo + # must not also skip /models/foo. + key_fn = src.split("function autoLoadCandidateKey", 1)[1] + key_fn = key_fn.split("\nfunction ", 1)[0] + assert "normalizeLoadTargetKey(id)" in key_fn + assert "id.toLowerCase()" not in key_fn + + +def test_local_variant_scans_bounded_concurrency(): + """Each /gguf-variants call triggers a recursive backend directory scan, + so pre-resolution must not fan out unbounded over every indexed folder + at once.""" + src = _read("features/chat/api/chat-adapter.ts") + assert "const AUTO_LOAD_VARIANT_SCAN_CONCURRENCY" in src + assert "async function mapWithConcurrency" in src + auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1] + assert "await mapWithConcurrency(" in auto_load + assert "AUTO_LOAD_VARIANT_SCAN_CONCURRENCY," in auto_load + assert "await Promise.all(\n cascadeLocalRows.map(" not in auto_load