From b3b85696d5eebc17ab3f87c55e03be1fa058f887 Mon Sep 17 00:00:00 2001 From: Povilas Kanapickas Date: Sun, 25 Oct 2026 14:32:50 +0000 Subject: [PATCH] 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 load that produced models is memoized, a failed load is retried, and a caller's signal detaches only that caller's wait. A loader reports whether the catalog it assigned is non-empty, so an empty response stays retryable. getModels() deliberately does not cache an empty result, and memoizing it here would pin the handler to the fallback model info for the rest of its lifetime. 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. The native Ollama path bypasses modelCache's single-flight, so it bounds its own fetch with the shared catalog timeout signal. Extend the shared adapter harness with a pre-warm matrix (memoization, empty-catalog retry, retry, single flight, pre-aborted caller, mid-flight waiter detach) plus per-provider regression tests and direct ensureCatalogLoaded unit tests covering abort-listener cleanup on both settlement paths. --- .../providers/__tests__/base-provider.spec.ts | 158 ++++++++++++++++++ src/api/providers/__tests__/lmstudio.spec.ts | 42 +++++ .../__tests__/model-catalog-adapters.spec.ts | 130 +++++++++++++- .../providers/__tests__/native-ollama.spec.ts | 58 +++++++ .../providers/__tests__/openrouter.spec.ts | 105 +++++++++++- src/api/providers/__tests__/poe.spec.ts | 34 +++- src/api/providers/__tests__/requesty.spec.ts | 45 +++++ src/api/providers/__tests__/unbound.spec.ts | 13 ++ src/api/providers/base-provider.ts | 44 +++++ .../fetchers/__tests__/modelCache.spec.ts | 31 ++++ .../fetchers/__tests__/openrouter.spec.ts | 23 +++ src/api/providers/fetchers/modelCache.ts | 5 +- src/api/providers/fetchers/openrouter.ts | 9 +- src/api/providers/lm-studio.ts | 17 +- src/api/providers/native-ollama.ts | 13 ++ src/api/providers/openrouter.ts | 49 +++--- src/api/providers/poe.ts | 18 +- src/api/providers/requesty.ts | 26 ++- src/api/providers/router-provider.ts | 16 +- src/api/providers/unbound.ts | 9 + .../utils/__tests__/abort-signal.spec.ts | 130 ++++++++++++++ src/api/providers/utils/abort-signal.ts | 52 ++++-- 22 files changed, 965 insertions(+), 62 deletions(-) diff --git a/src/api/providers/__tests__/base-provider.spec.ts b/src/api/providers/__tests__/base-provider.spec.ts index 998b08e3f4..7a68262bdd 100644 --- a/src/api/providers/__tests__/base-provider.spec.ts +++ b/src/api/providers/__tests__/base-provider.spec.ts @@ -32,6 +32,11 @@ class TestProvider extends BaseProvider { public testConvertToolsForOpenAI(tools: any[] | undefined, strict?: boolean): any[] | undefined { return this.convertToolsForOpenAI(tools, strict) } + + // Expose protected method for testing + public testEnsureCatalogLoaded(load: () => Promise, signal?: AbortSignal): Promise { + return this.ensureCatalogLoaded(load, signal) + } } describe("BaseProvider", () => { @@ -527,4 +532,157 @@ describe("BaseProvider", () => { expect(result?.[0]).toEqual(tools[0]) }) }) + + describe("ensureCatalogLoaded", () => { + it("loads once and memoizes a successful load", async () => { + const load = vi.fn().mockResolvedValue(undefined) + + await provider.testEnsureCatalogLoaded(load) + await provider.testEnsureCatalogLoaded(load) + + expect(load).toHaveBeenCalledTimes(1) + }) + + it("single-flights concurrent callers", async () => { + let resolveLoad: () => void = () => {} + const load = vi.fn( + () => + new Promise((resolve) => { + resolveLoad = resolve + }), + ) + + const first = provider.testEnsureCatalogLoaded(load) + const second = provider.testEnsureCatalogLoaded(load) + + expect(load).toHaveBeenCalledTimes(1) + + resolveLoad() + await Promise.all([first, second]) + }) + + it("retries after a failed load", async () => { + const load = vi.fn().mockRejectedValueOnce(new Error("catalog unavailable")).mockResolvedValue(undefined) + + await expect(provider.testEnsureCatalogLoaded(load)).rejects.toThrow("catalog unavailable") + await expect(provider.testEnsureCatalogLoaded(load)).resolves.toBeUndefined() + + expect(load).toHaveBeenCalledTimes(2) + }) + + it("rejects a pre-aborted caller without starting a load", async () => { + const controller = new AbortController() + controller.abort(new Error("caller gave up")) + const load = vi.fn().mockResolvedValue(undefined) + + await expect(provider.testEnsureCatalogLoaded(load, controller.signal)).rejects.toThrow("caller gave up") + expect(load).not.toHaveBeenCalled() + }) + + it("detaches only the aborted waiter while the shared load still settles", async () => { + let resolveLoad: () => void = () => {} + const load = vi.fn( + () => + new Promise((resolve) => { + resolveLoad = resolve + }), + ) + + const controller = new AbortController() + const wait = provider.testEnsureCatalogLoaded(load, controller.signal) + controller.abort(new Error("caller gave up")) + + await expect(wait).rejects.toThrow("caller gave up") + + resolveLoad() + await new Promise((resolve) => setTimeout(resolve, 0)) + + // The shared load completed, so a later caller is a memoized no-op. + await expect(provider.testEnsureCatalogLoaded(load)).resolves.toBeUndefined() + expect(load).toHaveBeenCalledTimes(1) + }) + + it("does not forward the caller signal to the shared load", async () => { + const controller = new AbortController() + const load = vi.fn().mockResolvedValue(undefined) + + await provider.testEnsureCatalogLoaded(load, controller.signal) + + expect(load).toHaveBeenCalledWith() + }) + + it("propagates a shared-load failure to a signal-carrying waiter", async () => { + const controller = new AbortController() + const load = vi.fn().mockRejectedValueOnce(new Error("catalog unavailable")).mockResolvedValue(undefined) + + // The waiter is watching a signal that never aborts, so the rejection can + // only come from the shared fetch failing -- the branch that must reject + // the waiter rather than hang until the caller's signal fires. + await expect(provider.testEnsureCatalogLoaded(load, controller.signal)).rejects.toThrow( + "catalog unavailable", + ) + + // The failed load cleared the shared promise, so the next caller retries. + await expect(provider.testEnsureCatalogLoaded(load)).resolves.toBeUndefined() + expect(load).toHaveBeenCalledTimes(2) + }) + + it("removes the exact abort listener it registered after a successful load", async () => { + const controller = new AbortController() + const addListener = vi.spyOn(controller.signal, "addEventListener") + const removeListener = vi.spyOn(controller.signal, "removeEventListener") + const load = vi.fn().mockResolvedValue(undefined) + + await provider.testEnsureCatalogLoaded(load, controller.signal) + + // A long-lived signal must not accumulate listeners: the removal has to + // target the same function reference that was registered. + expect(addListener).toHaveBeenCalledWith("abort", expect.any(Function), { once: true }) + const registered = addListener.mock.calls[0]?.[1] + expect(removeListener).toHaveBeenCalledWith("abort", registered) + expect(removeListener).toHaveBeenCalledTimes(1) + }) + + it("removes the exact abort listener it registered after a rejected load", async () => { + const controller = new AbortController() + const addListener = vi.spyOn(controller.signal, "addEventListener") + const removeListener = vi.spyOn(controller.signal, "removeEventListener") + const load = vi.fn().mockRejectedValueOnce(new Error("catalog unavailable")) + + await expect(provider.testEnsureCatalogLoaded(load, controller.signal)).rejects.toThrow( + "catalog unavailable", + ) + + expect(addListener).toHaveBeenCalledWith("abort", expect.any(Function), { once: true }) + const registered = addListener.mock.calls[0]?.[1] + expect(removeListener).toHaveBeenCalledWith("abort", registered) + expect(removeListener).toHaveBeenCalledTimes(1) + }) + + it("does not memoize a load that reports an empty catalog", async () => { + const load = vi.fn().mockResolvedValueOnce(false).mockResolvedValueOnce(true) + + await provider.testEnsureCatalogLoaded(load) + // The empty result must stay retryable, so the next caller loads again. + await provider.testEnsureCatalogLoaded(load) + expect(load).toHaveBeenCalledTimes(2) + + // The second load produced a catalog, so it is memoized. + await provider.testEnsureCatalogLoaded(load) + expect(load).toHaveBeenCalledTimes(2) + }) + + it("memoizes a load that returns true or void", async () => { + const returningTrue = vi.fn().mockResolvedValue(true) + await provider.testEnsureCatalogLoaded(returningTrue) + await provider.testEnsureCatalogLoaded(returningTrue) + expect(returningTrue).toHaveBeenCalledTimes(1) + + const other = new TestProvider() + const returningVoid = vi.fn().mockResolvedValue(undefined) + await other.testEnsureCatalogLoaded(returningVoid) + await other.testEnsureCatalogLoaded(returningVoid) + expect(returningVoid).toHaveBeenCalledTimes(1) + }) + }) }) diff --git a/src/api/providers/__tests__/lmstudio.spec.ts b/src/api/providers/__tests__/lmstudio.spec.ts index f3b87df7a9..5e298bad87 100644 --- a/src/api/providers/__tests__/lmstudio.spec.ts +++ b/src/api/providers/__tests__/lmstudio.spec.ts @@ -62,9 +62,21 @@ vi.mock("openai", () => { import type { Anthropic } from "@anthropic-ai/sdk" +import { providerIdentifiers } from "@roo-code/types" + import { LmStudioHandler } from "../lm-studio" import type { ApiHandlerOptions } from "../../../shared/api" +const { mockGetModels, mockGetModelsFromCache } = vi.hoisted(() => ({ + mockGetModels: vi.fn(), + mockGetModelsFromCache: vi.fn(), +})) + +vi.mock("../fetchers/modelCache", () => ({ + getModels: mockGetModels, + getModelsFromCache: mockGetModelsFromCache, +})) + describe("LmStudioHandler", () => { let handler: LmStudioHandler let mockOptions: ApiHandlerOptions @@ -77,6 +89,9 @@ describe("LmStudioHandler", () => { } handler = new LmStudioHandler(mockOptions) mockCreate.mockClear() + mockGetModels.mockReset() + mockGetModelsFromCache.mockReset() + mockGetModelsFromCache.mockReturnValue(undefined) }) describe("constructor", () => { @@ -258,4 +273,31 @@ describe("LmStudioHandler", () => { expect(modelInfo.info.contextWindow).toBe(128_000) }) }) + + describe("ensureModelFetched", () => { + it("returns default info before the pre-warm and cached info after", async () => { + // The pre-warm populates the shared cache that getModel() reads, so the + // cache starts empty and is filled by the mocked getModels() call. + let cacheContents: Record | undefined + mockGetModelsFromCache.mockImplementation(() => cacheContents) + mockGetModels.mockImplementation(async () => { + cacheContents = { "local-model": { maxTokens: 4096, contextWindow: 32_768 } } + return cacheContents + }) + + expect(handler.getModel().info.contextWindow).toBe(128_000) + + await handler.ensureModelFetched() + + expect(mockGetModels).toHaveBeenCalledWith( + { + provider: providerIdentifiers.lmstudio, + baseUrl: "http://localhost:1234", + }, + handler, + ) + expect(handler.getModel().info.contextWindow).toBe(32_768) + expect(handler.getModel().info.maxTokens).toBe(4096) + }) + }) }) diff --git a/src/api/providers/__tests__/model-catalog-adapters.spec.ts b/src/api/providers/__tests__/model-catalog-adapters.spec.ts index fdcbde7b7f..691cce3b84 100644 --- a/src/api/providers/__tests__/model-catalog-adapters.spec.ts +++ b/src/api/providers/__tests__/model-catalog-adapters.spec.ts @@ -28,6 +28,20 @@ const fetchers = vi.hoisted(() => ({ getZooGatewayModels: vi.fn(), })) +// The pre-warm hook reaches the fetcher either directly (fetchModel) or through +// getModels(). Delegating to the handler's own fetchModels keeps the fetcher +// call counts in this file meaningful and avoids touching the real disk cache. +const modelCache = vi.hoisted(() => ({ + getModels: vi.fn(), + getModelsFromCache: vi.fn(), +})) + +vi.mock("../fetchers/modelCache", async (importOriginal) => ({ + ...(await importOriginal()), + getModels: modelCache.getModels, + getModelsFromCache: modelCache.getModelsFromCache, +})) + vi.mock("../fetchers/deepseek", async (importOriginal) => ({ ...(await importOriginal()), getDeepSeekModels: fetchers.getDeepSeekModels, @@ -126,6 +140,8 @@ interface AdapterCase { expectedArgs: unknown[] /** Whether the adapter forwards an AbortSignal to its fetcher. */ forwardsSignal: boolean + /** Whether the handler implements the ensureModelFetched pre-warm hook. */ + prewarms: boolean } const cases: AdapterCase[] = [ @@ -136,6 +152,7 @@ const cases: AdapterCase[] = [ scope: { urlScoped: true, keyScoped: false, authScoped: false }, expectedArgs: [BASE_URL, API_KEY], forwardsSignal: true, + prewarms: false, }, { name: "IO Intelligence", @@ -144,6 +161,7 @@ const cases: AdapterCase[] = [ scope: { urlScoped: false, keyScoped: true, authScoped: false }, expectedArgs: [API_KEY], forwardsSignal: true, + prewarms: false, }, { name: "Kenari", @@ -152,6 +170,7 @@ const cases: AdapterCase[] = [ scope: { urlScoped: false, keyScoped: false, authScoped: false }, expectedArgs: [API_KEY], forwardsSignal: true, + prewarms: false, }, { name: "Kimi Code", @@ -160,6 +179,7 @@ const cases: AdapterCase[] = [ scope: { urlScoped: false, keyScoped: true, authScoped: true }, expectedArgs: [API_KEY], forwardsSignal: false, + prewarms: false, }, { name: "LiteLLM", @@ -168,6 +188,7 @@ const cases: AdapterCase[] = [ scope: { urlScoped: true, keyScoped: true, authScoped: false }, expectedArgs: [API_KEY, BASE_URL], forwardsSignal: true, + prewarms: false, }, { name: "LM Studio", @@ -176,6 +197,7 @@ const cases: AdapterCase[] = [ scope: { urlScoped: true, keyScoped: false, authScoped: false }, expectedArgs: [BASE_URL], forwardsSignal: true, + prewarms: true, }, { name: "Moonshot", @@ -184,6 +206,7 @@ const cases: AdapterCase[] = [ scope: { urlScoped: true, keyScoped: true, authScoped: false }, expectedArgs: [BASE_URL, API_KEY], forwardsSignal: true, + prewarms: false, }, { name: "NanoGPT", @@ -192,6 +215,7 @@ const cases: AdapterCase[] = [ scope: { urlScoped: false, keyScoped: true, authScoped: false }, expectedArgs: [API_KEY], forwardsSignal: true, + prewarms: false, }, { name: "Ollama", @@ -200,6 +224,7 @@ const cases: AdapterCase[] = [ scope: { urlScoped: true, keyScoped: false, authScoped: false }, expectedArgs: [BASE_URL, API_KEY], forwardsSignal: true, + prewarms: true, }, { name: "OpenCode Go", @@ -208,6 +233,7 @@ const cases: AdapterCase[] = [ scope: { urlScoped: false, keyScoped: false, authScoped: false }, expectedArgs: [API_KEY], forwardsSignal: true, + prewarms: false, }, { name: "OpenRouter", @@ -216,6 +242,7 @@ const cases: AdapterCase[] = [ scope: { urlScoped: false, keyScoped: false, authScoped: false }, expectedArgs: [undefined], forwardsSignal: true, + prewarms: true, }, { name: "Poe", @@ -224,6 +251,7 @@ const cases: AdapterCase[] = [ scope: { urlScoped: true, keyScoped: true, authScoped: false }, expectedArgs: [API_KEY, BASE_URL], forwardsSignal: true, + prewarms: true, }, { name: "Requesty", @@ -232,6 +260,7 @@ const cases: AdapterCase[] = [ scope: { urlScoped: true, keyScoped: true, authScoped: false }, expectedArgs: [BASE_URL, API_KEY], forwardsSignal: true, + prewarms: true, }, { name: "Unbound", @@ -240,6 +269,7 @@ const cases: AdapterCase[] = [ scope: { urlScoped: false, keyScoped: false, authScoped: false }, expectedArgs: [API_KEY], forwardsSignal: true, + prewarms: true, }, { name: "Vercel AI Gateway", @@ -248,6 +278,7 @@ const cases: AdapterCase[] = [ scope: { urlScoped: false, keyScoped: false, authScoped: false }, expectedArgs: [undefined], forwardsSignal: true, + prewarms: false, }, { name: "Zoo Gateway", @@ -256,6 +287,7 @@ const cases: AdapterCase[] = [ scope: { urlScoped: true, keyScoped: true, authScoped: true }, expectedArgs: [{ zooSessionToken: API_KEY, zooGatewayBaseUrl: BASE_URL }], forwardsSignal: false, + prewarms: false, }, ] @@ -267,9 +299,16 @@ describe("model catalog adapters", () => { fetcher.mockReset() fetcher.mockResolvedValue(CATALOG) } + + modelCache.getModels.mockReset() + modelCache.getModels.mockImplementation((opts: GetModelsOptions, handler: ApiHandler) => + handler.fetchModels!(opts), + ) + modelCache.getModelsFromCache.mockReset() + modelCache.getModelsFromCache.mockReturnValue(undefined) }) - describe.each(cases)("$name", ({ create, fetcher, scope, expectedArgs, forwardsSignal }) => { + describe.each(cases)("$name", ({ create, fetcher, scope, expectedArgs, forwardsSignal, prewarms }) => { it("declares its model cache scope", () => { expect(create().getModelCacheScope()).toEqual(scope) }) @@ -301,6 +340,94 @@ describe("model catalog adapters", () => { await expect(create().fetchModels?.(options)).rejects.toThrow("catalog unavailable") }) + + if (prewarms) { + it("implements the ensureModelFetched pre-warm hook", () => { + expect(create().ensureModelFetched).toBeTypeOf("function") + }) + } + }) + + describe.each(cases.filter(({ prewarms }) => prewarms))("$name pre-warm", ({ create, fetcher }) => { + it("memoizes a successful load", async () => { + const handler = create() + + await handler.ensureModelFetched!() + await handler.ensureModelFetched!() + + expect(fetchers[fetcher]).toHaveBeenCalledTimes(1) + }) + + it("retries a failed load on the next call", async () => { + const handler = create() + fetchers[fetcher].mockRejectedValueOnce(new Error("catalog unavailable")) + + await expect(handler.ensureModelFetched!()).rejects.toThrow("catalog unavailable") + await expect(handler.ensureModelFetched!()).resolves.toBeUndefined() + + expect(fetchers[fetcher]).toHaveBeenCalledTimes(2) + }) + + it("retries when the provider returns an empty catalog", async () => { + const handler = create() + // getModels() deliberately does not cache an empty response, so the + // pre-warm must not memoize it either -- otherwise the handler would be + // pinned to the fallback model info for the rest of its lifetime. + fetchers[fetcher].mockResolvedValueOnce({}) + + await handler.ensureModelFetched!() + expect(fetchers[fetcher]).toHaveBeenCalledTimes(1) + + await handler.ensureModelFetched!() + expect(fetchers[fetcher]).toHaveBeenCalledTimes(2) + + // The second response was non-empty, so it is memoized. + await handler.ensureModelFetched!() + expect(fetchers[fetcher]).toHaveBeenCalledTimes(2) + }) + + it("single-flights concurrent calls", async () => { + const handler = create() + + await Promise.all([handler.ensureModelFetched!(), handler.ensureModelFetched!()]) + + expect(fetchers[fetcher]).toHaveBeenCalledTimes(1) + }) + + it("rejects a pre-aborted caller without fetching", async () => { + const handler = create() + const controller = new AbortController() + controller.abort(new Error("caller gave up")) + + await expect(handler.ensureModelFetched!(controller.signal)).rejects.toThrow("caller gave up") + expect(fetchers[fetcher]).not.toHaveBeenCalled() + }) + + it("detaches only the aborted waiter while the shared fetch still settles", async () => { + const handler = create() + let settle: (models: ModelRecord) => void = () => {} + fetchers[fetcher].mockImplementation( + () => + new Promise((resolve) => { + settle = resolve + }), + ) + + const controller = new AbortController() + const wait = handler.ensureModelFetched!(controller.signal) + controller.abort(new Error("caller gave up")) + + await expect(wait).rejects.toThrow("caller gave up") + + settle(CATALOG) + // Let the shared load's promise chain settle (load -> .then -> .finally). + await new Promise((resolve) => setTimeout(resolve, 0)) + + // The load completed despite the detached waiter, so the next call + // is a memoized no-op rather than a second fetch. + await expect(handler.ensureModelFetched!()).resolves.toBeUndefined() + expect(fetchers[fetcher]).toHaveBeenCalledTimes(1) + }) }) it("defaults to an unscoped cache and no catalog fetcher for static providers", () => { @@ -308,5 +435,6 @@ describe("model catalog adapters", () => { expect(handler.getModelCacheScope()).toEqual({ urlScoped: false, keyScoped: false, authScoped: false }) expect(handler.fetchModels).toBeUndefined() + expect(handler.ensureModelFetched).toBeUndefined() }) }) diff --git a/src/api/providers/__tests__/native-ollama.spec.ts b/src/api/providers/__tests__/native-ollama.spec.ts index 764d509090..e3bb383b84 100644 --- a/src/api/providers/__tests__/native-ollama.spec.ts +++ b/src/api/providers/__tests__/native-ollama.spec.ts @@ -5,6 +5,7 @@ import { Anthropic } from "@anthropic-ai/sdk" import { NativeOllamaHandler } from "../native-ollama" import { ApiHandlerOptions } from "../../../shared/api" import { getOllamaModels } from "../fetchers/ollama" +import { MODEL_CATALOG_FETCH_TIMEOUT_MS } from "../fetchers/modelCache" import { makeCreateMessageMetadata } from "../../../test-utils/api" import { clearAllMocks } from "../../../test-utils/reset" @@ -1352,6 +1353,63 @@ describe("NativeOllamaHandler", () => { }) }) + describe("ensureModelFetched", () => { + it("bounds the catalog fetch with a fetch-local timeout signal", async () => { + mockGetOllamaModels.mockResolvedValue({ + llama2: { + contextWindow: 8192, + maxTokens: 4096, + supportsImages: false, + supportsPromptCache: false, + }, + }) + + const timeoutSpy = vi.spyOn(AbortSignal, "timeout") + try { + const controller = new AbortController() + await handler.ensureModelFetched(controller.signal) + + // This path bypasses getModels()'s single-flight, so it must carry its + // own bound. The signal is fetch-local: the caller's signal is never + // forwarded, so one caller aborting cannot cancel the shared fetch. + const forwarded = mockGetOllamaModels.mock.calls[0]?.[2]?.signal + // Pinning the delay and the exact forwarded signal proves the bound is + // real: a plain new AbortController().signal would satisfy a weaker + // "is an AbortSignal" check while removing the timeout entirely. + expect(timeoutSpy).toHaveBeenCalledWith(MODEL_CATALOG_FETCH_TIMEOUT_MS) + expect(forwarded).toBe(timeoutSpy.mock.results[0]?.value) + expect(forwarded).not.toBe(controller.signal) + expect(forwarded?.aborted).toBe(false) + } finally { + timeoutSpy.mockRestore() + } + }) + + it("fetches the catalog once and populates this.models", async () => { + mockGetOllamaModels.mockResolvedValue({ + llama2: { + contextWindow: 8192, + maxTokens: 4096, + supportsImages: false, + supportsPromptCache: false, + }, + }) + + // Before the pre-warm the instance catalog is empty, so getModel() + // returns the sane defaults rather than the discovered metadata. + expect(handler.getModel().info.contextWindow).not.toBe(8192) + + await handler.ensureModelFetched() + + expect(mockGetOllamaModels).toHaveBeenCalledTimes(1) + expect(handler.getModel().info.contextWindow).toBe(8192) + + // A second call is a memoized no-op. + await handler.ensureModelFetched() + expect(mockGetOllamaModels).toHaveBeenCalledTimes(1) + }) + }) + describe("tool calling", () => { it("should include tools when tools are provided", async () => { // Model metadata should not gate tool inclusion; metadata.tools controls it. diff --git a/src/api/providers/__tests__/openrouter.spec.ts b/src/api/providers/__tests__/openrouter.spec.ts index 9e07bcd515..1bf6669324 100644 --- a/src/api/providers/__tests__/openrouter.spec.ts +++ b/src/api/providers/__tests__/openrouter.spec.ts @@ -17,7 +17,7 @@ const MOCK_TIMEOUT_MS = 300_000 import { Anthropic } from "@anthropic-ai/sdk" import OpenAI from "openai" -import { providerIdentifiers } from "@roo-code/types" +import { OPENROUTER_DEFAULT_PROVIDER_NAME, openRouterDefaultModelInfo, providerIdentifiers } from "@roo-code/types" import type { ModelInfo } from "@roo-code/types" import { applyOpenRouterMoonshotK3Profile } from "../fetchers/openrouter" @@ -1041,5 +1041,108 @@ describe("OpenRouterHandler", () => { const handler = new OpenRouterHandler(mockOptions) expect(handler.getModelCacheScope()).toEqual({ urlScoped: false, keyScoped: false, authScoped: false }) }) + + it("does not fetch the catalog in the constructor", async () => { + const { getModels } = await import("../fetchers/modelCache") + + new OpenRouterHandler(mockOptions) + + expect(getModels).not.toHaveBeenCalled() + expect(mockGetModelEndpoints).not.toHaveBeenCalled() + }) + + it("populates models and endpoints when ensureModelFetched resolves", async () => { + mockGetModelEndpoints.mockResolvedValue({ + moonshotai: { + maxTokens: 1000, + contextWindow: 2000, + supportsPromptCache: false, + description: "Moonshot endpoint", + }, + }) + + const handler = new OpenRouterHandler( + makeApiHandlerOptions({ + ...mockOptions, + openRouterSpecificProvider: "moonshotai", + }), + ) + + // Before the pre-warm both catalogs are empty, so getModel() falls back + // to the hardcoded default info. + expect(handler.getModel().info).toBe(openRouterDefaultModelInfo) + + await handler.ensureModelFetched() + + // The endpoint catalog is populated and selected for the specific provider. + expect(handler.getModel().info.description).toBe("Moonshot endpoint") + expect(handler.getModel().info.contextWindow).toBe(2000) + + // The model catalog is populated too: without a specific provider the + // fetched record (not the hardcoded default) is used. + const plainHandler = new OpenRouterHandler(mockOptions) + await plainHandler.ensureModelFetched() + expect(plainHandler.getModel().info.description).toBe("Claude 3.7 Sonnet") + }) + + it("retries while a selected provider's endpoint catalog is empty", async () => { + const handler = new OpenRouterHandler( + makeApiHandlerOptions({ + ...mockOptions, + openRouterSpecificProvider: "moonshotai", + }), + ) + + // The model-level catalog resolves, but the endpoint request produced + // nothing (e.g. it timed out). getModel() would fall back to the + // model-level record, so the load must stay retryable. + await handler.ensureModelFetched() + expect(mockGetModelEndpoints).toHaveBeenCalledTimes(1) + + await handler.ensureModelFetched() + expect(mockGetModelEndpoints).toHaveBeenCalledTimes(2) + + mockGetModelEndpoints.mockResolvedValue({ + moonshotai: { + maxTokens: 1000, + contextWindow: 2000, + supportsPromptCache: false, + description: "Moonshot endpoint", + }, + }) + + await handler.ensureModelFetched() + expect(mockGetModelEndpoints).toHaveBeenCalledTimes(3) + expect(handler.getModel().info.description).toBe("Moonshot endpoint") + + // Both catalogs are populated, so the load is now memoized. + await handler.ensureModelFetched() + expect(mockGetModelEndpoints).toHaveBeenCalledTimes(3) + }) + + it("memoizes the default-provider load even when the endpoint catalog is empty", async () => { + const handler = new OpenRouterHandler(mockOptions) + + await handler.ensureModelFetched() + await handler.ensureModelFetched() + + // Without a selected provider the model-level catalog is what getModel() + // uses, so an empty endpoint response must not force a refetch. + expect(mockGetModelEndpoints).toHaveBeenCalledTimes(1) + }) + + it("treats the [default] provider sentinel as the default provider", async () => { + const handler = new OpenRouterHandler( + makeApiHandlerOptions({ + ...mockOptions, + openRouterSpecificProvider: OPENROUTER_DEFAULT_PROVIDER_NAME, + }), + ) + + await handler.ensureModelFetched() + await handler.ensureModelFetched() + + expect(mockGetModelEndpoints).toHaveBeenCalledTimes(1) + }) }) }) diff --git a/src/api/providers/__tests__/poe.spec.ts b/src/api/providers/__tests__/poe.spec.ts index 1384be0f88..9d9698aa0f 100644 --- a/src/api/providers/__tests__/poe.spec.ts +++ b/src/api/providers/__tests__/poe.spec.ts @@ -5,12 +5,13 @@ import { getModelsFromCache } from "../fetchers/modelCache" import { clearAllMocks } from "../../../test-utils/reset" -const { mockStreamText, mockGenerateText, mockCreatePoe, mockGetModelsFromCache, mockCaptureException } = +const { mockStreamText, mockGenerateText, mockCreatePoe, mockGetModels, mockGetModelsFromCache, mockCaptureException } = vitest.hoisted(() => ({ mockStreamText: vitest.fn(), mockGenerateText: vitest.fn(), mockCreatePoe: vitest.fn(), mockCaptureException: vitest.fn(), + mockGetModels: vitest.fn(), mockGetModelsFromCache: vitest.fn(), })) @@ -90,6 +91,7 @@ vitest.mock("ai", async (importOriginal) => { }) vitest.mock("../fetchers/modelCache", () => ({ + getModels: mockGetModels, getModelsFromCache: mockGetModelsFromCache, })) @@ -172,6 +174,36 @@ describe("PoeHandler", () => { }) }) + describe("ensureModelFetched", () => { + it("returns default info before the pre-warm and cached info after", async () => { + // The pre-warm populates the shared cache that getModel() reads, so the + // cache starts empty and is filled by the mocked getModels() call. + let cacheContents: typeof cachedModels | undefined + mockGetModelsFromCache.mockImplementation(() => cacheContents) + mockGetModels.mockImplementation(async () => { + cacheContents = cachedModels + return cachedModels + }) + + const handler = new PoeHandler({ poeApiKey: "key", apiModelId: "anthropic/claude-sonnet-4" }) + + expect(handler.getModel().info.maxTokens).toBe(8192) + + await handler.ensureModelFetched() + + expect(mockGetModels).toHaveBeenCalledWith( + { + provider: providerIdentifiers.poe, + apiKey: "key", + baseUrl: undefined, + }, + handler, + ) + expect(handler.getModel().info.maxTokens).toBe(10_000) + expect(handler.getModel().info.contextWindow).toBe(200_000) + }) + }) + describe("createMessage", () => { it("streams text chunks", async () => { const handler = new PoeHandler({ poeApiKey: "key", apiModelId: "anthropic/claude-sonnet-4" }) diff --git a/src/api/providers/__tests__/requesty.spec.ts b/src/api/providers/__tests__/requesty.spec.ts index c71d976e9b..346ea4f14c 100644 --- a/src/api/providers/__tests__/requesty.spec.ts +++ b/src/api/providers/__tests__/requesty.spec.ts @@ -9,6 +9,9 @@ const MOCK_TIMEOUT_MS = 300_000 import { Anthropic } from "@anthropic-ai/sdk" import OpenAI from "openai" +import { providerIdentifiers, requestyDefaultModelInfo } from "@roo-code/types" + +import { getModels } from "../fetchers/modelCache" import { RequestyHandler } from "../requesty" import { Package } from "../../../shared/package" import { ApiHandlerCreateMessageMetadata } from "../../index" @@ -709,5 +712,47 @@ describe("RequestyHandler", () => { const handler = new RequestyHandler(mockOptions) expect(handler.getModelCacheScope()).toEqual({ urlScoped: true, keyScoped: true, authScoped: false }) }) + + it("populates this.models so getModel() stops returning the default info", async () => { + const handler = new RequestyHandler(mockOptions) + + expect(handler.getModel().info.description).toBe(requestyDefaultModelInfo.description) + + await handler.ensureModelFetched() + + expect(handler.getModel().info.description).toBe("Claude 4 Sonnet") + expect(handler.getModel().info.contextWindow).toBe(200000) + }) + + it("fetches the catalog with the configured API key", async () => { + // Requesty enforces per-key model allowlists, so an unauthenticated fetch would + // both miss the key's models and write a URL-only cache entry that another key + // could then be served from. + const handler = new RequestyHandler(mockOptions) + + await handler.ensureModelFetched() + + expect(getModels).toHaveBeenCalledWith( + { + provider: providerIdentifiers.requesty, + apiKey: "test-key", + baseUrl: "https://router.requesty.ai/v1", + }, + handler, + ) + }) + + it("fetches the catalog with each handler's own API key", async () => { + // Two handlers on the same server must resolve different key-scoped cache + // entries, which is only possible if each passes its own key to getModels(). + const handlerA = new RequestyHandler(makeApiHandlerOptions({ requestyApiKey: "key-a" })) + const handlerB = new RequestyHandler(makeApiHandlerOptions({ requestyApiKey: "key-b" })) + + await handlerA.ensureModelFetched() + await handlerB.ensureModelFetched() + + const apiKeys = vitest.mocked(getModels).mock.calls.map(([options]) => options.apiKey) + expect(apiKeys).toEqual(["key-a", "key-b"]) + }) }) }) diff --git a/src/api/providers/__tests__/unbound.spec.ts b/src/api/providers/__tests__/unbound.spec.ts index 15f7612bfa..c0f142d8ad 100644 --- a/src/api/providers/__tests__/unbound.spec.ts +++ b/src/api/providers/__tests__/unbound.spec.ts @@ -1,6 +1,8 @@ import { Anthropic } from "@anthropic-ai/sdk" import OpenAI from "openai" +import { unboundDefaultModelInfo } from "@roo-code/types" + import { UnboundHandler } from "../unbound" import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" import { clearAllMocks } from "../../../test-utils/reset" @@ -207,5 +209,16 @@ describe("UnboundHandler", () => { const handler = new UnboundHandler({ unboundApiKey: "test-key", unboundModelId: "openai/gpt-4o" }) expect(handler.getModelCacheScope()).toEqual({ urlScoped: false, keyScoped: false, authScoped: false }) }) + + it("populates this.models so getModel() stops returning the default info", async () => { + const handler = new UnboundHandler({ unboundApiKey: "test-key", unboundModelId: "openai/gpt-4o" }) + + expect(handler.getModel().info.description).toBe(unboundDefaultModelInfo.description) + + await handler.ensureModelFetched() + + expect(handler.getModel().info.description).toBe("GPT-4o") + expect(handler.getModel().info.maxTokens).toBe(4096) + }) }) }) diff --git a/src/api/providers/base-provider.ts b/src/api/providers/base-provider.ts index 713636595d..2829c28b86 100644 --- a/src/api/providers/base-provider.ts +++ b/src/api/providers/base-provider.ts @@ -7,6 +7,7 @@ import { ApiStream } from "../transform/stream" import { sanitizeSurrogates, sanitizeSurrogatesDeep } from "../transform/sanitize-surrogates" import { countTokens } from "../../utils/countTokens" import { isMcpTool } from "../../utils/mcp-name" +import { detachOnAbort } from "./utils/abort-signal" import { getApiRequestTimeout } from "./utils/timeout-config" /** @@ -15,6 +16,9 @@ import { getApiRequestTimeout } from "./utils/timeout-config" export abstract class BaseProvider implements ApiHandler { protected readonly timeoutMs: number = getApiRequestTimeout() + private catalogFetch?: Promise + private catalogLoaded = false + abstract createMessage( systemPrompt: string, messages: Anthropic.Messages.MessageParam[], @@ -157,4 +161,44 @@ export abstract class BaseProvider implements ApiHandler { getModelCacheScope(): ModelCacheScope { return { urlScoped: false, keyScoped: false, authScoped: false } } + + /** + * Single-flight the one-time catalog load behind ensureModelFetched(). + * + * Semantics mirror RouterProvider.ensureModelFetched(): + * - concurrent callers share one fetch; + * - a load that produced models is never repeated (catalogLoaded); + * - a failed load is retried by the next caller (the promise is cleared); + * - a load that resolved with an empty catalog is retried too: getModels() + * deliberately does not cache an empty response, so memoizing it here would + * permanently pin the handler to the fallback model info; + * - a caller's signal detaches only that caller's wait, never the shared fetch, + * because the fetch is already bounded by modelCache's 15 s single-flight. + * + * A loader returns `false` when the catalog it assigned is empty. Returning + * `void` (or `true`) means the load counts as successful. + */ + protected async ensureCatalogLoaded(load: () => Promise, signal?: AbortSignal): Promise { + if (this.catalogLoaded) return + if (signal?.aborted) throw signal.reason + + this.catalogFetch ??= load() + .then((hasCatalog) => { + this.catalogLoaded = hasCatalog !== false + }) + .finally(() => { + this.catalogFetch = undefined + }) + + const fetch = this.catalogFetch + + if (!signal) { + await fetch + return + } + + // Detach this waiter as soon as the signal aborts; the shared in-flight + // fetch continues for any other waiter and still populates the cache. + await detachOnAbort(fetch, signal) + } } diff --git a/src/api/providers/fetchers/__tests__/modelCache.spec.ts b/src/api/providers/fetchers/__tests__/modelCache.spec.ts index 9df780a170..3081b5f3eb 100644 --- a/src/api/providers/fetchers/__tests__/modelCache.spec.ts +++ b/src/api/providers/fetchers/__tests__/modelCache.spec.ts @@ -1198,6 +1198,37 @@ describe("key-scoped cache key derivation", () => { const discriminator = cacheKey.split(":").pop() as string expect(discriminator).toMatch(/^[0-9a-f]{8}$/) }) + + it("never serves one key's catalog to another key on the same server", async () => { + // Requesty enforces per-key model allowlists, so a shared entry would expose one + // key's catalog to another. Drive the real getModels() cache path with a stateful + // in-memory store so a cache hit is observable. + const store = new Map() + // NodeCache keys are `string | number`; the cache only ever writes string keys here. + mockCache.get.mockImplementation((key: string | number) => store.get(String(key)) as never) + mockCache.set.mockImplementation(((key: string | number, value: unknown) => { + store.set(String(key), value) + return true + }) as never) + + const keyOneModels = { "key-one/model": mockModels["key-scoped/model"] } + const keyTwoModels = { "key-two/model": mockModels["key-scoped/model"] } + const fetchKeyOne = vi.fn().mockResolvedValue(keyOneModels) + const fetchKeyTwo = vi.fn().mockResolvedValue(keyTwoModels) + const handlerOne = makeHandler(REQUESTY_SCOPE, fetchKeyOne) + const handlerTwo = makeHandler(REQUESTY_SCOPE, fetchKeyTwo) + + const keyOneOptions = { provider: keyScopedProvider, apiKey: "key-one" } as const + const keyTwoOptions = { provider: keyScopedProvider, apiKey: "key-two" } as const + + await expect(getModels(keyOneOptions, handlerOne)).resolves.toEqual(keyOneModels) + await expect(getModels(keyTwoOptions, handlerTwo)).resolves.toEqual(keyTwoModels) + + // Re-reading key one must hit its own entry, not key two's. + await expect(getModels(keyOneOptions, handlerOne)).resolves.toEqual(keyOneModels) + expect(fetchKeyOne).toHaveBeenCalledTimes(1) + expect(fetchKeyTwo).toHaveBeenCalledTimes(1) + }) }) describe("NanoGPT key-scoped cache isolation", () => { diff --git a/src/api/providers/fetchers/__tests__/openrouter.spec.ts b/src/api/providers/fetchers/__tests__/openrouter.spec.ts index 8e33706836..2eca76e401 100644 --- a/src/api/providers/fetchers/__tests__/openrouter.spec.ts +++ b/src/api/providers/fetchers/__tests__/openrouter.spec.ts @@ -6,6 +6,8 @@ import { back as nockBack } from "nock" import type { ModelInfo } from "@roo-code/types" +import { MODEL_CATALOG_FETCH_TIMEOUT_MS } from "../modelCache" + import { applyOpenRouterMoonshotK3Profile, getOpenRouterModelEndpoints, @@ -314,6 +316,27 @@ describe("OpenRouter API", () => { getSpy.mockRestore() }) + + it("bounds the endpoints request with the model-catalog fetch timeout", async () => { + const axios = await import("axios") + const getSpy = vi.spyOn(axios.default, "get").mockResolvedValue({ + data: { data: { id: "anthropic/claude-sonnet-4", name: "Claude Sonnet 4", endpoints: [] } }, + }) + + await getOpenRouterModelEndpoints("anthropic/claude-sonnet-4") + + // This request bypasses getModels()' single-flight, so it must carry its own + // bound: without a timeout a stalled response keeps the caller's shared + // catalog load pending and later callers join it instead of retrying. + expect(getSpy).toHaveBeenCalledWith( + "https://openrouter.ai/api/v1/models/anthropic/claude-sonnet-4/endpoints", + { + timeout: MODEL_CATALOG_FETCH_TIMEOUT_MS, + }, + ) + + getSpy.mockRestore() + }) }) describe("parseOpenRouterModel", () => { diff --git a/src/api/providers/fetchers/modelCache.ts b/src/api/providers/fetchers/modelCache.ts index 0134ec47cb..4380535373 100644 --- a/src/api/providers/fetchers/modelCache.ts +++ b/src/api/providers/fetchers/modelCache.ts @@ -35,7 +35,10 @@ const inFlightRefresh = new Map() // an in-flight entry pending indefinitely. The value is the maximum of the 5–15 s bounds the // individual fetchers it subsumes used to apply, relaxing rather than tightening endpoints that // already had a bound. -const MODEL_CATALOG_FETCH_TIMEOUT_MS = 15_000 +// +// Exported so providers that call a fetcher outside getModels() (e.g. the native Ollama +// pre-warm, which bypasses the single-flight) can apply the same bound. +export const MODEL_CATALOG_FETCH_TIMEOUT_MS = 15_000 /** * State of one shared (single-flight) provider fetch. diff --git a/src/api/providers/fetchers/openrouter.ts b/src/api/providers/fetchers/openrouter.ts index 758ee4778b..000f83643c 100644 --- a/src/api/providers/fetchers/openrouter.ts +++ b/src/api/providers/fetchers/openrouter.ts @@ -14,6 +14,8 @@ import { parseApiPrice } from "../../../shared/cost" import { throwIfAborted } from "../utils/abort-signal" +import { MODEL_CATALOG_FETCH_TIMEOUT_MS } from "./modelCache" + /** * OpenRouterBaseModel */ @@ -152,7 +154,12 @@ export async function getOpenRouterModelEndpoints( const baseURL = options?.openRouterBaseUrl || "https://openrouter.ai/api/v1" try { - const response = await axios.get(`${baseURL}/models/${modelId}/endpoints`) + // This request runs outside getModels()' single-flight, so it needs its own bound: + // a stalled endpoint response would otherwise keep the caller's catalog load pending + // forever, and later callers would join that pending promise instead of retrying. + const response = await axios.get(`${baseURL}/models/${modelId}/endpoints`, { + timeout: MODEL_CATALOG_FETCH_TIMEOUT_MS, + }) const result = openRouterModelEndpointsResponseSchema.safeParse(response.data) const data = result.success ? result.data.data : response.data.data diff --git a/src/api/providers/lm-studio.ts b/src/api/providers/lm-studio.ts index a8f93738e5..6f5979a469 100644 --- a/src/api/providers/lm-studio.ts +++ b/src/api/providers/lm-studio.ts @@ -23,7 +23,7 @@ import type { CompletePromptOptions, ModelCacheScope, } from "../index" -import { getModelsFromCache } from "./fetchers/modelCache" +import { getModels, getModelsFromCache } from "./fetchers/modelCache" import { getLMStudioModels } from "./fetchers/lmstudio" import { handleOpenAIError } from "./utils/error-handler" import { extractReasoningFromDelta } from "./utils/extract-reasoning" @@ -56,6 +56,21 @@ export class LmStudioHandler extends BaseProvider implements SingleCompletionHan return getLMStudioModels(options.baseUrl, ...fetchOpts) } + async ensureModelFetched(signal?: AbortSignal): Promise { + await this.ensureCatalogLoaded(async () => { + const models = await getModels( + { + provider: providerIdentifiers.lmstudio, + baseUrl: this.options.lmStudioBaseUrl, + }, + this, + ) + // getModels() does not cache an empty response, so an empty catalog must + // stay retryable instead of being memoized as a successful load. + return Object.keys(models).length > 0 + }, signal) + } + override async *createMessage( systemPrompt: string, messages: Anthropic.Messages.MessageParam[], diff --git a/src/api/providers/native-ollama.ts b/src/api/providers/native-ollama.ts index 86229b4918..dbae7106a7 100644 --- a/src/api/providers/native-ollama.ts +++ b/src/api/providers/native-ollama.ts @@ -6,6 +6,7 @@ import { ApiStream } from "../transform/stream" import { BaseProvider } from "./base-provider" import type { ApiHandlerOptions, GetModelsOptions } from "../../shared/api" import { getOllamaModels, isSecureOllamaEndpoint } from "./fetchers/ollama" +import { MODEL_CATALOG_FETCH_TIMEOUT_MS } from "./fetchers/modelCache" import { createAbortError, hasUsableTimeout, @@ -666,6 +667,18 @@ export class NativeOllamaHandler extends BaseProvider implements SingleCompletio return this.getModel() } + async ensureModelFetched(signal?: AbortSignal): Promise { + await this.ensureCatalogLoaded(async () => { + // This path bypasses getModels()'s single-flight, so it carries its own + // bound: a hung Ollama endpoint would otherwise leave the shared + // catalogFetch pending forever and block every later retry. The timeout + // is fetch-local so the caller's signal stays independent of it. + await this.fetchModel(AbortSignal.timeout(MODEL_CATALOG_FETCH_TIMEOUT_MS)) + // An empty catalog must stay retryable instead of being memoized. + return Object.keys(this.models).length > 0 + }, signal) + } + override getModel(): { id: string; info: ModelInfo } { const modelId = this.options.ollamaModelId || "" return { diff --git a/src/api/providers/openrouter.ts b/src/api/providers/openrouter.ts index 5251b046d8..36e09aa883 100644 --- a/src/api/providers/openrouter.ts +++ b/src/api/providers/openrouter.ts @@ -149,11 +149,6 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH const apiKey = this.options.openRouterApiKey ?? NOT_PROVIDED this.client = new OpenAI({ baseURL, apiKey, defaultHeaders: DEFAULT_HEADERS, timeout: this.timeoutMs }) - - // Load models asynchronously to populate cache before getModel() is called - this.loadDynamicModels().catch((error) => { - console.error("[OpenRouterHandler] Failed to load dynamic models:", error) - }) } override getModelCacheScope(): ModelCacheScope { @@ -165,28 +160,30 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH return getOpenRouterModels(undefined, ...fetchOpts) } - private async loadDynamicModels(): Promise { - try { - const [models, endpoints] = await Promise.all([ - getModels({ provider: providerIdentifiers.openrouter }, this), - getModelEndpoints( - { - router: providerIdentifiers.openrouter, - modelId: this.options.openRouterModelId, - endpoint: this.options.openRouterSpecificProvider, - }, - this, - ), - ]) + /** + * Whether the user pinned a concrete OpenRouter upstream provider. The + * "[default]" sentinel means "let OpenRouter choose", so it behaves like no + * selection at all. + */ + private get hasSpecificProvider(): boolean { + return ( + Boolean(this.options.openRouterSpecificProvider) && + this.options.openRouterSpecificProvider !== OPENROUTER_DEFAULT_PROVIDER_NAME + ) + } - this.models = models - this.endpoints = endpoints - } catch (error) { - console.error("[OpenRouterHandler] Error loading dynamic models:", { - error: error instanceof Error ? error.message : String(error), - stack: error instanceof Error ? error.stack : undefined, - }) - } + async ensureModelFetched(signal?: AbortSignal): Promise { + await this.ensureCatalogLoaded(async () => { + await this.fetchModel() + // fetchModel() returns model info, not the catalog, so report emptiness + // from the catalogs it assigned: an empty response must stay retryable. + // With a specific provider selected, the endpoint catalog is what + // getModel() uses, so an empty endpoint response is a failed load even + // when the model-level catalog arrived. + const hasModels = Object.keys(this.models).length > 0 + const hasEndpoints = Object.keys(this.endpoints).length > 0 + return hasModels && (!this.hasSpecificProvider || hasEndpoints) + }, signal) } getReasoningDetails(): any[] | undefined { diff --git a/src/api/providers/poe.ts b/src/api/providers/poe.ts index e4df8b98ed..29333aa056 100644 --- a/src/api/providers/poe.ts +++ b/src/api/providers/poe.ts @@ -32,7 +32,7 @@ import type { CompletePromptOptions, ModelCacheScope, } from "../index" -import { getModelsFromCache } from "./fetchers/modelCache" +import { getModels, getModelsFromCache } from "./fetchers/modelCache" import { getPoeModels } from "./fetchers/poe" const DEFAULT_THINKING_BUDGET = 8192 @@ -59,6 +59,22 @@ export class PoeHandler extends BaseProvider implements SingleCompletionHandler return getPoeModels(options.apiKey, options.baseUrl, ...fetchOpts) } + async ensureModelFetched(signal?: AbortSignal): Promise { + await this.ensureCatalogLoaded(async () => { + const models = await getModels( + { + provider: providerIdentifiers.poe, + apiKey: this.options.poeApiKey, + baseUrl: this.options.poeBaseUrl, + }, + this, + ) + // getModels() does not cache an empty response, so an empty catalog must + // stay retryable instead of being memoized as a successful load. + return Object.keys(models).length > 0 + }, signal) + } + override getModel() { const id = this.options.apiModelId ?? poeDefaultModelId const cached = getModelsFromCache( diff --git a/src/api/providers/requesty.ts b/src/api/providers/requesty.ts index 3a9d5e064a..f8d23c0ce6 100644 --- a/src/api/providers/requesty.ts +++ b/src/api/providers/requesty.ts @@ -95,10 +95,34 @@ export class RequestyHandler extends BaseProvider implements SingleCompletionHan } public async fetchModel() { - this.models = await getModels({ provider: providerIdentifiers.requesty, baseUrl: this.baseURL }, this) + // Pass the API key so the catalog request is authenticated and the cache entry is + // key-scoped. Requesty enforces per-key model allowlists, so a keyless (URL-only) + // entry would let one key's catalog be served to another key on the same server. + this.models = await getModels( + { + provider: providerIdentifiers.requesty, + apiKey: this.options.requestyApiKey, + baseUrl: this.baseURL, + }, + this, + ) return this.getModel() } + async ensureModelFetched(signal?: AbortSignal): Promise { + await this.ensureCatalogLoaded(async () => { + // The caller's signal is deliberately not threaded into fetchModel(): this load + // is shared by every waiter on the handler, so binding it to one caller's signal + // would cancel the fetch for the others. The signal still detaches this caller's + // own wait (see ensureCatalogLoaded), and the shared fetch stays bounded by + // modelCache's single-flight timeout. + await this.fetchModel() + // fetchModel() returns model info, not the catalog, so report emptiness + // from the catalog it assigned: an empty response must stay retryable. + return Object.keys(this.models).length > 0 + }, signal) + } + override getModel() { const id = this.options.requestyModelId ?? requestyDefaultModelId const cachedInfo = this.models[id] ?? requestyDefaultModelInfo diff --git a/src/api/providers/router-provider.ts b/src/api/providers/router-provider.ts index 746310d62b..6c206d9ddb 100644 --- a/src/api/providers/router-provider.ts +++ b/src/api/providers/router-provider.ts @@ -6,6 +6,7 @@ import { ApiHandlerOptions, RouterName } from "../../shared/api" import { BaseProvider } from "./base-provider" import { getModels, getModelsFromCache, refreshModels } from "./fetchers/modelCache" +import { detachOnAbort } from "./utils/abort-signal" import { DEFAULT_HEADERS, NOT_PROVIDED } from "./constants" @@ -123,20 +124,7 @@ export abstract class RouterProvider extends BaseProvider { // Detach this waiter as soon as the signal aborts; the shared in-flight // fetch continues for any other waiter and still populates the cache. - await new Promise((resolve, reject) => { - const onAbort = () => reject(signal.reason) - signal.addEventListener("abort", onAbort, { once: true }) - fetch.then( - () => { - signal.removeEventListener("abort", onAbort) - resolve() - }, - (error: unknown) => { - signal.removeEventListener("abort", onAbort) - reject(error) - }, - ) - }) + await detachOnAbort(fetch, signal) } override getModel(): { id: string; info: ModelInfo } { diff --git a/src/api/providers/unbound.ts b/src/api/providers/unbound.ts index 4b8ba242c5..59c7bb3b66 100644 --- a/src/api/providers/unbound.ts +++ b/src/api/providers/unbound.ts @@ -99,6 +99,15 @@ export class UnboundHandler extends BaseProvider implements SingleCompletionHand return this.getModel() } + async ensureModelFetched(signal?: AbortSignal): Promise { + await this.ensureCatalogLoaded(async () => { + await this.fetchModel() + // fetchModel() returns model info, not the catalog, so report emptiness + // from the catalog it assigned: an empty response must stay retryable. + return Object.keys(this.models).length > 0 + }, signal) + } + override getModel() { const id = this.options.unboundModelId ?? unboundDefaultModelId const cachedInfo = this.models[id] ?? unboundDefaultModelInfo diff --git a/src/api/providers/utils/__tests__/abort-signal.spec.ts b/src/api/providers/utils/__tests__/abort-signal.spec.ts index fc65239db2..44ac0b8f27 100644 --- a/src/api/providers/utils/__tests__/abort-signal.spec.ts +++ b/src/api/providers/utils/__tests__/abort-signal.spec.ts @@ -1,6 +1,7 @@ import { bridgeExternalAbortSignal, createAbortError, + detachOnAbort, isRequestAborted, mergeAbortSignalAndTimeout, hasUsableTimeout, @@ -130,6 +131,135 @@ describe("rejectOnAbort", () => { }) }) +describe("detachOnAbort", () => { + it("resolves with the pending value when it settles before the signal aborts", async () => { + const controller = new AbortController() + + await expect(withSettleGuard(detachOnAbort(Promise.resolve("done"), controller.signal))).resolves.toBe("done") + expect(controller.signal.aborted).toBe(false) + }) + + it("rejects with the signal's own reason when the signal aborts first", async () => { + const controller = new AbortController() + const reason = new Error("caller gave up") + // Never settles: the wait must end purely via the abort. + const pending = new Promise(() => {}) + const wait = detachOnAbort(pending, controller.signal) + controller.abort(reason) + + // Not createAbortError(): the caller's reason must surface unchanged. + await expect(withSettleGuard(wait)).rejects.toBe(reason) + }) + + it("rejects with the default AbortError reason when abort() carries no reason", async () => { + const controller = new AbortController() + const wait = detachOnAbort(new Promise(() => {}), controller.signal) + controller.abort() + + await expect(withSettleGuard(wait)).rejects.toMatchObject({ name: "AbortError" }) + }) + + it("rejects immediately when the signal is already aborted", async () => { + const controller = new AbortController() + const reason = new Error("caller gave up") + controller.abort(reason) + + // An already-aborted signal never fires its listener, so waiting would hang. + await expect(withSettleGuard(detachOnAbort(new Promise(() => {}), controller.signal))).rejects.toBe( + reason, + ) + }) + + it("does not attach a listener when the signal is already aborted", async () => { + const controller = new AbortController() + controller.abort() + const addSpy = vi.spyOn(controller.signal, "addEventListener") + + await expect(withSettleGuard(detachOnAbort(Promise.resolve("done"), controller.signal))).rejects.toBeDefined() + + expect(addSpy).not.toHaveBeenCalled() + addSpy.mockRestore() + }) + + it("does not leave a rejecting pending unhandled when the signal is already aborted", async () => { + // The pre-aborted branch never consumes `pending`, so it must still + // handle a rejection or it surfaces as unhandled. + const controller = new AbortController() + controller.abort() + const pending = Promise.reject(new Error("shared fetch failed")) + + const unhandled: unknown[] = [] + const onUnhandledRejection = (reason: unknown) => unhandled.push(reason) + process.on("unhandledRejection", onUnhandledRejection) + try { + await expect(withSettleGuard(detachOnAbort(pending, controller.signal))).rejects.toBeDefined() + + // Node reports an unhandled rejection only after the microtask queue drains. + await new Promise((resolve) => setTimeout(resolve, 0)) + expect(unhandled).toHaveLength(0) + } finally { + process.off("unhandledRejection", onUnhandledRejection) + } + }) + + it("propagates the pending rejection when the signal stays active", async () => { + const controller = new AbortController() + const boom = new Error("shared fetch failed") + + await expect(withSettleGuard(detachOnAbort(Promise.reject(boom), controller.signal))).rejects.toBe(boom) + }) + + it("detaches the abort listener once the pending settles", async () => { + const controller = new AbortController() + const addSpy = vi.spyOn(controller.signal, "addEventListener") + const removeSpy = vi.spyOn(controller.signal, "removeEventListener") + + await expect(withSettleGuard(detachOnAbort(Promise.resolve("done"), controller.signal))).resolves.toBe("done") + + // Must remove the exact registered listener, not just any function. + expect(addSpy).toHaveBeenCalledTimes(1) + const registeredListener = addSpy.mock.calls[0]?.[1] as EventListener | undefined + expect(typeof registeredListener).toBe("function") + expect(removeSpy).toHaveBeenCalledWith("abort", registeredListener) + addSpy.mockRestore() + removeSpy.mockRestore() + }) + + it("detaches the abort listener when the pending rejects", async () => { + const controller = new AbortController() + const addSpy = vi.spyOn(controller.signal, "addEventListener") + const removeSpy = vi.spyOn(controller.signal, "removeEventListener") + const lookupError = new Error("shared fetch failed") + + await expect(withSettleGuard(detachOnAbort(Promise.reject(lookupError), controller.signal))).rejects.toBe( + lookupError, + ) + + expect(addSpy).toHaveBeenCalledTimes(1) + const registeredListener = addSpy.mock.calls[0]?.[1] as EventListener | undefined + expect(typeof registeredListener).toBe("function") + expect(removeSpy).toHaveBeenCalledWith("abort", registeredListener) + addSpy.mockRestore() + removeSpy.mockRestore() + }) + + it("leaves the shared pending running after this waiter detaches", async () => { + // Aborting one caller must not cancel work other callers still await. + const controller = new AbortController() + let resolveShared: (value: string) => void = () => {} + const shared = new Promise((resolve) => { + resolveShared = resolve + }) + const wait = detachOnAbort(shared, controller.signal) + controller.abort(new Error("caller gave up")) + + await expect(withSettleGuard(wait)).rejects.toThrow("caller gave up") + + resolveShared("late value") + await expect(shared).resolves.toBe("late value") + }) +}) + describe("resolveModelWithAbort", () => { it("fast-fails with the provider abort error when the signal is already aborted", async () => { const controller = new AbortController() diff --git a/src/api/providers/utils/abort-signal.ts b/src/api/providers/utils/abort-signal.ts index c26bce4f6f..7f389ac6fb 100644 --- a/src/api/providers/utils/abort-signal.ts +++ b/src/api/providers/utils/abort-signal.ts @@ -155,32 +155,33 @@ export function createAbortError(providerName: string): Error { } /** - * Await `pending` but reject with the provider's abort error when `signal` - * aborts first. For async phases that have no native signal support (model - * discovery) yet must still settle promptly on cancellation. The underlying - * promise keeps running (its settlement is ignored) — cancellation is - * cooperative at this boundary. + * Await `pending` but settle early with `abortError` when `signal` aborts. + * + * The underlying promise keeps running (its settlement is ignored) — + * cancellation is cooperative at this boundary. The abort listener is detached + * once `pending` settles (success or failure), so repeated waits on one + * long-lived signal do not accumulate listeners. * - * The abort listener is detached once `pending` settles (success or - * failure), so repeated calls on one signal do not accumulate listeners. + * An already-aborted signal rejects without attaching a listener: it would + * never fire, so the wait would hang. */ -export function rejectOnAbort(pending: Promise, signal: AbortSignal, providerName: string): Promise { +function raceAbort(pending: Promise, signal: AbortSignal, abortError: () => unknown): Promise { if (signal.aborted) { - // The caller has already cancelled, so this branch never consumes `pending`; - // attach a handler anyway so a rejecting `pending` cannot surface as an - // unhandled rejection. + // This branch never consumes `pending`; attach a handler anyway so a + // rejecting `pending` cannot surface as an unhandled rejection. void pending.catch(() => {}) - return Promise.reject(createAbortError(providerName)) + return Promise.reject(abortError()) } + return new Promise((resolve, reject) => { - const onAbort = () => reject(createAbortError(providerName)) + const onAbort = () => reject(abortError()) signal.addEventListener("abort", onAbort, { once: true }) void pending.then( (value) => { signal.removeEventListener("abort", onAbort) resolve(value) }, - (error) => { + (error: unknown) => { signal.removeEventListener("abort", onAbort) reject(error) }, @@ -188,6 +189,29 @@ export function rejectOnAbort(pending: Promise, signal: AbortSignal, provi }) } +/** + * Await `pending` but reject with the provider's abort error when `signal` + * aborts first. For async phases that have no native signal support (model + * discovery) yet must still settle promptly on cancellation. + * + * Use {@link detachOnAbort} instead when the awaited work is shared with other + * callers and the caller's own abort reason must propagate unchanged. + */ +export function rejectOnAbort(pending: Promise, signal: AbortSignal, providerName: string): Promise { + return raceAbort(pending, signal, () => createAbortError(providerName)) +} + +/** + * Await `pending` but detach this waiter when `signal` aborts, rejecting with + * `signal.reason`. The shared work keeps running for other callers. + * + * Unlike {@link rejectOnAbort}, the caller's own reason is preserved instead of + * being normalized to `createAbortError(providerName)`. + */ +export function detachOnAbort(pending: Promise, signal: AbortSignal): Promise { + return raceAbort(pending, signal, () => signal.reason) +} + /** * Resolve a provider's model metadata inside a cancellation scope. *