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. *