diff --git a/packages/coding-agent/src/extensions/llama/ui.ts b/packages/coding-agent/src/extensions/llama/ui.ts index 25f64c07f..a2ba8e4d4 100644 --- a/packages/coding-agent/src/extensions/llama/ui.ts +++ b/packages/coding-agent/src/extensions/llama/ui.ts @@ -110,6 +110,8 @@ class HuggingFaceSearch extends Container implements Focusable { private results: HuggingFaceModel[] = []; private filteredResults: HuggingFaceModel[] = []; private selectedIndex = 0; + // Whether the user moved the selection since the query last changed. + private selectionTouched = false; private query = ""; private status = "Type at least 2 characters"; private debounce: ReturnType | undefined; @@ -186,14 +188,16 @@ class HuggingFaceSearch extends Container implements Focusable { this.tui.requestRender(); } - private filterResults(): void { + private filterResults(preferredId?: string): void { if (this.query) { const matches = new Set(fuzzyFilter(this.results, this.query, (model) => model.id).map((model) => model.id)); this.filteredResults = this.results.filter((model) => matches.has(model.id)); } else { this.filteredResults = this.results; } - this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredResults.length - 1)); + const preferredIndex = preferredId ? this.filteredResults.findIndex((model) => model.id === preferredId) : -1; + if (preferredIndex >= 0) this.selectedIndex = preferredIndex; + else this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredResults.length - 1)); this.updateResults(); } @@ -225,10 +229,12 @@ class HuggingFaceSearch extends Container implements Focusable { const results = await this.search(query, request.signal); this.cache.set(query.toLowerCase(), results); if (this.closed || request.signal.aborted || this.query !== query) return; + // Keep the row the user browsed to while the search was in flight; otherwise show the best match. + const selectedId = this.selectionTouched ? this.filteredResults[this.selectedIndex]?.id : undefined; this.results = results; this.selectedIndex = 0; this.status = results.length === 0 ? "No GGUF models found" : ""; - this.filterResults(); + this.filterResults(selectedId); } catch (error) { if (this.closed || request.signal.aborted || this.query !== query) return; this.results = []; @@ -251,6 +257,7 @@ class HuggingFaceSearch extends Container implements Focusable { if (this.keybindings.matches(data, "tui.select.up")) { if (this.filteredResults.length > 0) { this.selectedIndex = this.selectedIndex === 0 ? this.filteredResults.length - 1 : this.selectedIndex - 1; + this.selectionTouched = true; this.updateResults(); } return; @@ -258,6 +265,7 @@ class HuggingFaceSearch extends Container implements Focusable { if (this.keybindings.matches(data, "tui.select.down")) { if (this.filteredResults.length > 0) { this.selectedIndex = this.selectedIndex === this.filteredResults.length - 1 ? 0 : this.selectedIndex + 1; + this.selectionTouched = true; this.updateResults(); } return; @@ -272,6 +280,7 @@ class HuggingFaceSearch extends Container implements Focusable { if (pageIndex !== undefined) { if (this.filteredResults.length > 0) { this.selectedIndex = pageIndex; + this.selectionTouched = true; this.updateResults(); } return; @@ -290,6 +299,7 @@ class HuggingFaceSearch extends Container implements Focusable { const query = this.input.getValue().trim(); if (query === this.query) return; this.query = query; + this.selectionTouched = false; this.scheduleSearch(); } } diff --git a/packages/coding-agent/src/modes/interactive/components/model-selector.ts b/packages/coding-agent/src/modes/interactive/components/model-selector.ts index f33eecf27..181d9b42c 100644 --- a/packages/coding-agent/src/modes/interactive/components/model-selector.ts +++ b/packages/coding-agent/src/modes/interactive/components/model-selector.ts @@ -209,8 +209,10 @@ export class ModelSelectorComponent extends Container implements Focusable { this.refreshStatusSuccess = true; } } + const selectedModel = this.filteredModels[this.selectedIndex]?.model; this.loadModelsFromSnapshot(); - this.filterModels(this.searchInput.getValue()); + // The user may have moved while the refresh was in flight, and the refresh can shift rows. + this.filterModels(this.searchInput.getValue(), selectedModel); this.tui.requestRender(); } catch (error) { if (this.closed) return; @@ -280,7 +282,7 @@ export class ModelSelectorComponent extends Container implements Focusable { } } - private filterModels(query: string): void { + private filterModels(query: string, preferredModel?: Model): void { if (query) { const filtered = fuzzyFilter(this.activeModels, query, (item) => { const defaultText = this.isDefaultModel(item.model) ? " default" : ""; @@ -299,10 +301,14 @@ export class ModelSelectorComponent extends Container implements Focusable { } else { this.filteredModels = this.activeModels; } - // When filtering by a query, move the selector to the top row so the best - // match is highlighted. When the query is cleared, keep the current position - // clamped to the (restored) list length. - this.selectedIndex = query ? 0 : Math.min(this.selectedIndex, Math.max(0, this.filteredModels.length - 1)); + // Keep the preferred model selected if it is still listed. Otherwise, when filtering + // by a query, move the selector to the top row so the best match is highlighted. + // When the query is cleared, keep the current position clamped to the (restored) list length. + const preferredIndex = preferredModel + ? this.filteredModels.findIndex((item) => modelsAreEqual(preferredModel, item.model)) + : -1; + if (preferredIndex >= 0) this.selectedIndex = preferredIndex; + else this.selectedIndex = query ? 0 : Math.min(this.selectedIndex, Math.max(0, this.filteredModels.length - 1)); this.updateList(); } diff --git a/packages/coding-agent/test/llama-ui.test.ts b/packages/coding-agent/test/llama-ui.test.ts new file mode 100644 index 000000000..f8c4998a5 --- /dev/null +++ b/packages/coding-agent/test/llama-ui.test.ts @@ -0,0 +1,122 @@ +import { type Component, setKeybindings, type TUI } from "@earendil-works/pi-tui"; +import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; +import type { ExtensionCommandContext } from "../src/core/extensions/types.ts"; +import { KeybindingsManager } from "../src/core/keybindings.ts"; +import type { HuggingFaceModel } from "../src/extensions/llama/huggingface.ts"; +import { type LlamaUi, showLlamaUi } from "../src/extensions/llama/ui.ts"; +import { initTheme, theme } from "../src/modes/interactive/theme/theme.ts"; +import { stripAnsi } from "../src/utils/ansi.ts"; + +const DOWN = "\x1b[B"; +const ENTER = "\r"; + +function model(id: string): HuggingFaceModel { + return { id, downloads: 1000 }; +} + +function selectedId(view: Component): string | undefined { + const row = stripAnsi(view.render(120).join("\n")) + .split("\n") + .find((line) => line.includes("→ ")); + return row?.split("→ ")[1]?.split(" ")[0]; +} + +// Opens the Hugging Face search with each query's results held until the test resolves them. +async function openSearch(): Promise<{ + view: Component; + type: (text: string) => Promise; + finishSearch: (query: string, results: HuggingFaceModel[]) => Promise; + selected: Promise; +}> { + const pending = new Map void>(); + const search = (query: string) => + new Promise((resolve) => { + pending.set(query, resolve); + }); + let view: Component | undefined; + let resolveSelected: (model: string | undefined) => void = () => {}; + const selected = new Promise((resolve) => { + resolveSelected = resolve; + }); + const ctx = { + ui: { + custom: (factory: (tui: TUI, th: typeof theme, kb: KeybindingsManager, done: () => void) => Component) => { + view = factory({ requestRender: () => {} } as unknown as TUI, theme, new KeybindingsManager(), () => {}); + return Promise.resolve(); + }, + notify: () => {}, + }, + } as unknown as ExtensionCommandContext; + void showLlamaUi(ctx, async (ui: LlamaUi) => { + resolveSelected(await ui.searchModels(search)); + }); + if (!view) throw new Error("llama UI was not shown"); + const shown = view; + const type = async (text: string) => { + for (const char of text) shown.handleInput?.(char); + await vi.advanceTimersByTimeAsync(500); + }; + const finishSearch = async (query: string, results: HuggingFaceModel[]) => { + const resolve = pending.get(query); + if (!resolve) throw new Error(`no search pending for ${query}`); + resolve(results); + await vi.advanceTimersByTimeAsync(0); + }; + return { view: shown, type, finishSearch, selected }; +} + +describe("llama.cpp Hugging Face search", () => { + beforeAll(() => { + initTheme("dark"); + }); + + beforeEach(() => { + setKeybindings(new KeybindingsManager()); + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + + const firstResults = [model("qwen/a-GGUF"), model("qwen/b-GGUF"), model("qwen/c-GGUF")]; + + it("keeps the browsed model selected when the next search lands", async () => { + const { view, type, finishSearch, selected } = await openSearch(); + await type("qw"); + await finishSearch("qw", firstResults); + await type("e"); + view.handleInput?.(DOWN); + view.handleInput?.(DOWN); + expect(selectedId(view)).toBe("qwen/c-GGUF"); + + await finishSearch("qwe", [model("qwen/new-GGUF"), ...firstResults]); + + expect(selectedId(view)).toBe("qwen/c-GGUF"); + view.handleInput?.(ENTER); + await expect(selected).resolves.toBe("qwen/c-GGUF"); + }); + + it("selects the best match when the user did not browse", async () => { + const { view, type, finishSearch } = await openSearch(); + await type("qw"); + await finishSearch("qw", firstResults); + await type("e"); + + await finishSearch("qwe", [model("qwen/new-GGUF"), ...firstResults]); + + expect(selectedId(view)).toBe("qwen/new-GGUF"); + }); + + it("selects the best match when the browsed model is not in the new results", async () => { + const { view, type, finishSearch } = await openSearch(); + await type("qw"); + await finishSearch("qw", firstResults); + await type("e"); + view.handleInput?.(DOWN); + + await finishSearch("qwe", [model("qwen/new-GGUF"), model("qwen/a-GGUF")]); + + expect(selectedId(view)).toBe("qwen/new-GGUF"); + }); +}); diff --git a/packages/coding-agent/test/model-selector.test.ts b/packages/coding-agent/test/model-selector.test.ts index df1a570e6..78fd63d56 100644 --- a/packages/coding-agent/test/model-selector.test.ts +++ b/packages/coding-agent/test/model-selector.test.ts @@ -1,3 +1,4 @@ +import type { ModelsRefreshResult } from "@earendil-works/pi-ai"; import { setKeybindings, type TUI } from "@earendil-works/pi-tui"; import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; import { KeybindingsManager } from "../src/core/keybindings.ts"; @@ -10,6 +11,13 @@ function createFakeTui(): TUI { return { requestRender: () => {} } as unknown as TUI; } +function selectedModelId(selector: ModelSelectorComponent): string | undefined { + const row = stripAnsi(selector.render(120).join("\n")) + .split("\n") + .find((line) => line.startsWith("→")); + return row?.replace(/^→\s*(✓\s*)?/, "").split(" ")[0]; +} + describe("model selector", () => { let harness: Harness | undefined; @@ -153,4 +161,92 @@ describe("model selector", () => { expect(rendered).toContain("Could not refresh 2 model catalogs (openai, anthropic); showing cached models."); }); }); + + // #9109 + describe("when the catalog refresh lands", () => { + // Opens the selector with the refresh in flight, so a test can browse first and then let it land. + async function openWithPendingRefresh(): Promise<{ + activeHarness: Harness; + selector: ModelSelectorComponent; + finishRefresh: () => Promise; + }> { + const activeHarness = await createHarness({ + models: [{ id: "alpha-1" }, { id: "alpha-2" }, { id: "alpha-3" }, { id: "beta-1" }], + }); + harness = activeHarness; + const { promise, resolve } = Promise.withResolvers(); + vi.spyOn(activeHarness.session.modelRuntime, "refresh").mockReturnValue(promise); + const selector = new ModelSelectorComponent( + createFakeTui(), + activeHarness.getModel("alpha-1"), + activeHarness.session.modelRuntime, + [], + () => {}, + () => {}, + ); + const finishRefresh = async () => { + resolve({ aborted: false, errors: new Map() }); + await vi.waitFor(() => { + expect(stripAnsi(selector.render(120).join("\n"))).toContain("Model catalogs refreshed."); + }); + }; + return { activeHarness, selector, finishRefresh }; + } + + it("keeps the selection on the browsed model", async () => { + const { selector, finishRefresh } = await openWithPendingRefresh(); + selector.handleInput("\x1b[B"); + selector.handleInput("\x1b[B"); + + await finishRefresh(); + + expect(selectedModelId(selector)).toBe("alpha-3"); + selector.dispose(); + }); + + it("keeps the selection while a search query is active", async () => { + const { selector, finishRefresh } = await openWithPendingRefresh(); + for (const char of "alpha") selector.handleInput(char); + selector.handleInput("\x1b[B"); + const browsed = selectedModelId(selector); + + await finishRefresh(); + + expect(selectedModelId(selector)).toBe(browsed); + expect(browsed).not.toBe("alpha-1"); + selector.dispose(); + }); + + it("follows the selected model when new rows shift its position", async () => { + const { activeHarness, selector, finishRefresh } = await openWithPendingRefresh(); + selector.handleInput("\x1b[B"); + const modelRuntime = activeHarness.session.modelRuntime; + const snapshot = modelRuntime.getAvailableSnapshot(); + vi.spyOn(modelRuntime, "getAvailableSnapshot").mockReturnValue([ + { ...snapshot[0], id: "added-model" }, + ...snapshot, + ]); + + await finishRefresh(); + + expect(selectedModelId(selector)).toBe("alpha-2"); + selector.handleInput("\x1b[A"); + selector.handleInput("\x1b[A"); + expect(selectedModelId(selector)).toBe("alpha-1"); + selector.dispose(); + }); + + it("falls back to the current model when the selected model disappears", async () => { + const { activeHarness, selector, finishRefresh } = await openWithPendingRefresh(); + selector.handleInput("\x1b[B"); + const modelRuntime = activeHarness.session.modelRuntime; + const remaining = modelRuntime.getAvailableSnapshot().filter((model) => model.id !== "alpha-2"); + vi.spyOn(modelRuntime, "getAvailableSnapshot").mockReturnValue(remaining); + + await finishRefresh(); + + expect(selectedModelId(selector)).toBe("alpha-1"); + selector.dispose(); + }); + }); });