Skip to content

Commit ab29f49

Browse files
bwachmanbwachman
committed
fix(tui): keep the live selection when a catalog prefetch settles
The startup refresh after an OpenCode Go or Zen prefetch rebuilt the catalog without the live selection, so a model switched during that rebuild kept the previous model's context window. All four prefetch refresh sites now share one helper that passes it, as the connect path already did.
1 parent 0a26789 commit ab29f49

4 files changed

Lines changed: 247 additions & 115 deletions

File tree

‎src/tui/runner/catalog-refresh.test.ts‎

Lines changed: 183 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,183 @@
1+
import { afterEach, describe, expect, test } from "bun:test";
2+
import { mkdtempSync, writeFileSync, rmSync } from "node:fs";
3+
import { tmpdir } from "node:os";
4+
import { join } from "node:path";
5+
6+
import { defined } from "../../../testkit/defined.js";
7+
import { withMockedModuleDuring } from "../../../testkit/mock-module.js";
8+
import {
9+
contextWindowFor,
10+
setModelContextWindows,
11+
setProviderContextWindowOverrides,
12+
} from "../../provider/context-window.js";
13+
import { refreshCatalogAfterPrefetch } from "./catalog-refresh.js";
14+
import type { RunnerServices, RunnerState } from "./state.js";
15+
16+
const dirs: string[] = [];
17+
afterEach(() => {
18+
setModelContextWindows(undefined);
19+
setProviderContextWindowOverrides(undefined);
20+
for (const d of dirs.splice(0)) rmSync(d, { recursive: true, force: true });
21+
});
22+
23+
function settingsFile(): string {
24+
const dir = mkdtempSync(join(tmpdir(), "catalog-refresh-"));
25+
dirs.push(dir);
26+
const path = join(dir, "settings.json");
27+
writeFileSync(
28+
path,
29+
JSON.stringify({
30+
providers: {
31+
first: {
32+
baseURL: "https://first.example/v1",
33+
keyless: true,
34+
models: ["shared-model"],
35+
contextWindow: 32_000,
36+
},
37+
second: {
38+
baseURL: "https://second.example/v1",
39+
keyless: true,
40+
models: ["shared-model"],
41+
contextWindow: 64_000,
42+
},
43+
},
44+
}),
45+
);
46+
return path;
47+
}
48+
49+
function fixture(withHost: boolean) {
50+
const calls: unknown[][] = [];
51+
const host = {
52+
refreshModels: (...args: unknown[]) => {
53+
calls.push(args);
54+
},
55+
};
56+
const state = {
57+
trueGlobalSettingsPath: settingsFile(),
58+
config: {
59+
apiKey: "",
60+
baseURL: "https://first.example/v1",
61+
model: "shared-model",
62+
providerName: "first",
63+
keyless: true,
64+
providers: [],
65+
},
66+
} as unknown as RunnerState;
67+
const services = {
68+
hostHolder: withHost ? { instance: host } : {},
69+
} as unknown as RunnerServices;
70+
return { state, services, calls };
71+
}
72+
73+
function withNoOAuthProfiles<T>(run: () => Promise<T>): Promise<T> {
74+
return withMockedModuleDuring(
75+
import.meta.resolve("../../config/oauth-stores.js"),
76+
(real: typeof import("../../config/oauth-stores.js")) => ({
77+
...real,
78+
listCodexProfiles: async () => [],
79+
listXaiProfiles: async () => [],
80+
listMetaProfiles: async () => [],
81+
}),
82+
run,
83+
);
84+
}
85+
86+
const flush = () => new Promise((r) => setTimeout(r, 50));
87+
88+
describe("refreshCatalogAfterPrefetch", () => {
89+
test("rebuilds the catalog from disk and hands it to the host after the prefetch settles", async () => {
90+
const { state, services, calls } = fixture(true);
91+
const prefetch = Promise.withResolvers<undefined>();
92+
await withNoOAuthProfiles(async () => {
93+
refreshCatalogAfterPrefetch(
94+
state,
95+
services,
96+
"go",
97+
() => prefetch.promise,
98+
);
99+
await flush();
100+
expect(calls).toHaveLength(0);
101+
prefetch.resolve(undefined);
102+
await flush();
103+
});
104+
expect(calls).toHaveLength(1);
105+
expect(defined(calls[0])[2]).toBe(state.config.providers);
106+
expect(state.config.providers.map((p) => p.name).sort()).toEqual([
107+
"first",
108+
"second",
109+
]);
110+
expect(state.config.settings).toBeDefined();
111+
});
112+
113+
test("a model switch while the prefetch is in flight is the selection the refresh keeps", async () => {
114+
const { state, services } = fixture(true);
115+
const prefetch = Promise.withResolvers<undefined>();
116+
await withNoOAuthProfiles(async () => {
117+
refreshCatalogAfterPrefetch(
118+
state,
119+
services,
120+
"zen",
121+
() => prefetch.promise,
122+
);
123+
state.config = { ...state.config, providerName: "second" };
124+
prefetch.resolve(undefined);
125+
await flush();
126+
});
127+
expect(contextWindowFor("shared-model")).toBe(64_000);
128+
});
129+
130+
test("does nothing when no host is mounted", async () => {
131+
const { state, services } = fixture(false);
132+
const before = state.config;
133+
await withNoOAuthProfiles(async () => {
134+
refreshCatalogAfterPrefetch(state, services, "go", async () => undefined);
135+
await flush();
136+
});
137+
expect(state.config).toBe(before);
138+
});
139+
140+
test("a failed prefetch leaves state and the host untouched", async () => {
141+
const { state, services, calls } = fixture(true);
142+
const before = state.config;
143+
await withNoOAuthProfiles(async () => {
144+
refreshCatalogAfterPrefetch(state, services, "go", () =>
145+
Promise.reject(new Error("offline")),
146+
);
147+
await flush();
148+
});
149+
expect(calls).toHaveLength(0);
150+
expect(state.config).toBe(before);
151+
});
152+
153+
test("a model switch during the post-prefetch credential discovery is the selection the refresh keeps", async () => {
154+
const { state, services } = fixture(true);
155+
const discovery = Promise.withResolvers<[]>();
156+
const discoveryStarted = Promise.withResolvers<undefined>();
157+
await withMockedModuleDuring(
158+
import.meta.resolve("../../config/oauth-stores.js"),
159+
(real: typeof import("../../config/oauth-stores.js")) => ({
160+
...real,
161+
listCodexProfiles: () => {
162+
discoveryStarted.resolve(undefined);
163+
return discovery.promise;
164+
},
165+
listXaiProfiles: async () => [],
166+
listMetaProfiles: async () => [],
167+
}),
168+
async () => {
169+
refreshCatalogAfterPrefetch(
170+
state,
171+
services,
172+
"go",
173+
async () => undefined,
174+
);
175+
await discoveryStarted.promise;
176+
state.config = { ...state.config, providerName: "second" };
177+
discovery.resolve([]);
178+
await flush();
179+
},
180+
);
181+
expect(contextWindowFor("shared-model")).toBe(64_000);
182+
});
183+
});

‎src/tui/runner/catalog-refresh.ts‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
import { getLogger } from "@intx/log";
2+
import { LOG_NAMESPACE_ROOT } from "../../branding.js";
3+
import { refreshLiveProviderCatalog } from "../../config/index.js";
4+
import {
5+
listFavoriteModels,
6+
listRecentModels,
7+
loadSettings,
8+
type ResolvedProvider,
9+
} from "../../config/settings.js";
10+
import { caughtDebug } from "../../logging/caught.js";
11+
import type { RunnerServices, RunnerState } from "./state.js";
12+
13+
const tuiLogger = getLogger([LOG_NAMESPACE_ROOT, "tui"]);
14+
15+
/**
16+
* Once a live model catalog prefetch settles, rebuild the provider catalog
17+
* from disk and hand it to the model picker. The resolved provider is read
18+
* when the prefetch settles, not when it starts, so a model switched in the
19+
* meantime is the one the catalog keeps.
20+
*/
21+
export function refreshCatalogAfterPrefetch(
22+
state: RunnerState,
23+
services: RunnerServices,
24+
catalogName: string,
25+
prefetch: () => Promise<unknown>,
26+
): void {
27+
void prefetch()
28+
.then(async () => {
29+
const host = services.hostHolder.instance;
30+
if (host === undefined) return;
31+
const onDisk = await loadSettings(state.trueGlobalSettingsPath);
32+
const resolvedForCatalog: ResolvedProvider = {
33+
apiKey: state.config.apiKey,
34+
baseURL: state.config.baseURL,
35+
model: state.config.model,
36+
providerName: state.config.providerName,
37+
...(state.config.keyless !== undefined
38+
? { keyless: state.config.keyless }
39+
: {}),
40+
};
41+
const providers = await refreshLiveProviderCatalog(
42+
onDisk,
43+
resolvedForCatalog,
44+
() => state.config,
45+
);
46+
state.config = {
47+
...state.config,
48+
providers,
49+
...(onDisk !== null ? { settings: onDisk } : {}),
50+
};
51+
host.refreshModels(
52+
listRecentModels(state.config.settings ?? { providers: {} }),
53+
listFavoriteModels(state.config.settings ?? { providers: {} }),
54+
providers,
55+
);
56+
})
57+
.catch(caughtDebug(tuiLogger, `${catalogName} model prefetch failed`));
58+
}

‎src/tui/runner/settings.ts‎

Lines changed: 3 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ import { resolveWaitForApproval } from "../../agent/tool-execution-watchdog.js";
6565
import { hostOf, type RunnerServices, type RunnerState } from "./state.js";
6666
import type { ProductHostConnectRequest } from "../product-host.js";
6767
import { LOG_NAMESPACE_ROOT } from "../../branding.js";
68+
import { refreshCatalogAfterPrefetch } from "./catalog-refresh.js";
6869

6970
const tuiLogger = getLogger([LOG_NAMESPACE_ROOT, "tui"]);
7071

@@ -378,50 +379,10 @@ export async function wireSettings(
378379
`Connected ${connectedName}. Open /model to pick a model.`,
379380
);
380381
if (isOpenCodeGoProvider({ name: providerName })) {
381-
void prefetchGoModels()
382-
.then(async () => {
383-
if (services.hostHolder.instance === undefined) return;
384-
const nextDisk = await loadSettings(trueGlobalSettingsPath);
385-
const nextProviders = await refreshLiveProviderCatalog(
386-
nextDisk,
387-
resolvedForCatalog,
388-
() => state.config,
389-
);
390-
state.config = {
391-
...state.config,
392-
providers: nextProviders,
393-
...(nextDisk !== null ? { settings: nextDisk } : {}),
394-
};
395-
services.hostHolder.instance.refreshModels(
396-
listRecentModels(state.config.settings ?? { providers: {} }),
397-
listFavoriteModels(state.config.settings ?? { providers: {} }),
398-
nextProviders,
399-
);
400-
})
401-
.catch(caughtDebug(tuiLogger, "go model prefetch failed"));
382+
refreshCatalogAfterPrefetch(state, services, "go", prefetchGoModels);
402383
}
403384
if (isZenProvider({ name: providerName })) {
404-
void prefetchZenModels()
405-
.then(async () => {
406-
if (services.hostHolder.instance === undefined) return;
407-
const nextDisk = await loadSettings(trueGlobalSettingsPath);
408-
const nextProviders = await refreshLiveProviderCatalog(
409-
nextDisk,
410-
resolvedForCatalog,
411-
() => state.config,
412-
);
413-
state.config = {
414-
...state.config,
415-
providers: nextProviders,
416-
...(nextDisk !== null ? { settings: nextDisk } : {}),
417-
};
418-
services.hostHolder.instance.refreshModels(
419-
listRecentModels(state.config.settings ?? { providers: {} }),
420-
listFavoriteModels(state.config.settings ?? { providers: {} }),
421-
nextProviders,
422-
);
423-
})
424-
.catch(caughtDebug(tuiLogger, "zen model prefetch failed"));
385+
refreshCatalogAfterPrefetch(state, services, "zen", prefetchZenModels);
425386
}
426387
})().catch(caughtDebug(tuiLogger, "provider connect failed"));
427388
};

0 commit comments

Comments
 (0)