Skip to content

Commit 46d5648

Browse files
committed
feat(api): add ensureModelFetched to the six remaining catalog providers
Add BaseProvider.ensureCatalogLoaded, a single-flight one-time catalog load mirroring RouterProvider.ensureModelFetched: concurrent callers share one fetch, a successful load is memoized, a failed load is retried, and a caller's signal detaches only that caller's wait. Wire the hook into openrouter, requesty, unbound, poe, lm-studio and native-ollama, and drop OpenRouter's eager constructor fetch so the catalog loads on demand instead of at handler construction. Extend the shared adapter harness with a pre-warm matrix (memoization, retry, single flight, pre-aborted caller, mid-flight waiter detach) plus per-provider regression tests and direct ensureCatalogLoaded unit tests.
1 parent ca05b73 commit 46d5648

15 files changed

Lines changed: 465 additions & 32 deletions

‎src/api/providers/__tests__/base-provider.spec.ts‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,11 @@ class TestProvider extends BaseProvider {
3232
public testConvertToolsForOpenAI(tools: any[] | undefined, strict?: boolean): any[] | undefined {
3333
return this.convertToolsForOpenAI(tools, strict)
3434
}
35+
36+
// Expose protected method for testing
37+
public testEnsureCatalogLoaded(load: () => Promise<void>, signal?: AbortSignal): Promise<void> {
38+
return this.ensureCatalogLoaded(load, signal)
39+
}
3540
}
3641

3742
describe("BaseProvider", () => {
@@ -527,4 +532,83 @@ describe("BaseProvider", () => {
527532
expect(result?.[0]).toEqual(tools[0])
528533
})
529534
})
535+
536+
describe("ensureCatalogLoaded", () => {
537+
it("loads once and memoizes a successful load", async () => {
538+
const load = vi.fn().mockResolvedValue(undefined)
539+
540+
await provider.testEnsureCatalogLoaded(load)
541+
await provider.testEnsureCatalogLoaded(load)
542+
543+
expect(load).toHaveBeenCalledTimes(1)
544+
})
545+
546+
it("single-flights concurrent callers", async () => {
547+
let resolveLoad: () => void = () => {}
548+
const load = vi.fn(
549+
() =>
550+
new Promise<void>((resolve) => {
551+
resolveLoad = resolve
552+
}),
553+
)
554+
555+
const first = provider.testEnsureCatalogLoaded(load)
556+
const second = provider.testEnsureCatalogLoaded(load)
557+
558+
expect(load).toHaveBeenCalledTimes(1)
559+
560+
resolveLoad()
561+
await Promise.all([first, second])
562+
})
563+
564+
it("retries after a failed load", async () => {
565+
const load = vi.fn().mockRejectedValueOnce(new Error("catalog unavailable")).mockResolvedValue(undefined)
566+
567+
await expect(provider.testEnsureCatalogLoaded(load)).rejects.toThrow("catalog unavailable")
568+
await expect(provider.testEnsureCatalogLoaded(load)).resolves.toBeUndefined()
569+
570+
expect(load).toHaveBeenCalledTimes(2)
571+
})
572+
573+
it("rejects a pre-aborted caller without starting a load", async () => {
574+
const controller = new AbortController()
575+
controller.abort(new Error("caller gave up"))
576+
const load = vi.fn().mockResolvedValue(undefined)
577+
578+
await expect(provider.testEnsureCatalogLoaded(load, controller.signal)).rejects.toThrow("caller gave up")
579+
expect(load).not.toHaveBeenCalled()
580+
})
581+
582+
it("detaches only the aborted waiter while the shared load still settles", async () => {
583+
let resolveLoad: () => void = () => {}
584+
const load = vi.fn(
585+
() =>
586+
new Promise<void>((resolve) => {
587+
resolveLoad = resolve
588+
}),
589+
)
590+
591+
const controller = new AbortController()
592+
const wait = provider.testEnsureCatalogLoaded(load, controller.signal)
593+
controller.abort(new Error("caller gave up"))
594+
595+
await expect(wait).rejects.toThrow("caller gave up")
596+
597+
resolveLoad()
598+
await new Promise((resolve) => setTimeout(resolve, 0))
599+
600+
// The shared load completed, so a later caller is a memoized no-op.
601+
await expect(provider.testEnsureCatalogLoaded(load)).resolves.toBeUndefined()
602+
expect(load).toHaveBeenCalledTimes(1)
603+
})
604+
605+
it("does not forward the caller signal to the shared load", async () => {
606+
const controller = new AbortController()
607+
const load = vi.fn().mockResolvedValue(undefined)
608+
609+
await provider.testEnsureCatalogLoaded(load, controller.signal)
610+
611+
expect(load).toHaveBeenCalledWith()
612+
})
613+
})
530614
})

‎src/api/providers/__tests__/lmstudio.spec.ts‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,9 +62,21 @@ vi.mock("openai", () => {
6262

6363
import type { Anthropic } from "@anthropic-ai/sdk"
6464

65+
import { providerIdentifiers } from "@roo-code/types"
66+
6567
import { LmStudioHandler } from "../lm-studio"
6668
import type { ApiHandlerOptions } from "../../../shared/api"
6769

70+
const { mockGetModels, mockGetModelsFromCache } = vi.hoisted(() => ({
71+
mockGetModels: vi.fn(),
72+
mockGetModelsFromCache: vi.fn(),
73+
}))
74+
75+
vi.mock("../fetchers/modelCache", () => ({
76+
getModels: mockGetModels,
77+
getModelsFromCache: mockGetModelsFromCache,
78+
}))
79+
6880
describe("LmStudioHandler", () => {
6981
let handler: LmStudioHandler
7082
let mockOptions: ApiHandlerOptions
@@ -77,6 +89,9 @@ describe("LmStudioHandler", () => {
7789
}
7890
handler = new LmStudioHandler(mockOptions)
7991
mockCreate.mockClear()
92+
mockGetModels.mockReset()
93+
mockGetModelsFromCache.mockReset()
94+
mockGetModelsFromCache.mockReturnValue(undefined)
8095
})
8196

8297
describe("constructor", () => {
@@ -258,4 +273,31 @@ describe("LmStudioHandler", () => {
258273
expect(modelInfo.info.contextWindow).toBe(128_000)
259274
})
260275
})
276+
277+
describe("ensureModelFetched", () => {
278+
it("returns default info before the pre-warm and cached info after", async () => {
279+
// The pre-warm populates the shared cache that getModel() reads, so the
280+
// cache starts empty and is filled by the mocked getModels() call.
281+
let cacheContents: Record<string, { maxTokens: number; contextWindow: number }> | undefined
282+
mockGetModelsFromCache.mockImplementation(() => cacheContents)
283+
mockGetModels.mockImplementation(async () => {
284+
cacheContents = { "local-model": { maxTokens: 4096, contextWindow: 32_768 } }
285+
return cacheContents
286+
})
287+
288+
expect(handler.getModel().info.contextWindow).toBe(128_000)
289+
290+
await handler.ensureModelFetched()
291+
292+
expect(mockGetModels).toHaveBeenCalledWith(
293+
{
294+
provider: providerIdentifiers.lmstudio,
295+
baseUrl: "http://localhost:1234",
296+
},
297+
handler,
298+
)
299+
expect(handler.getModel().info.contextWindow).toBe(32_768)
300+
expect(handler.getModel().info.maxTokens).toBe(4096)
301+
})
302+
})
261303
})

0 commit comments

Comments
 (0)