diff --git a/packages/types/src/__tests__/provider-identifiers.test.ts b/packages/types/src/__tests__/provider-identifiers.test.ts index af4de232f9..4b1106da5a 100644 --- a/packages/types/src/__tests__/provider-identifiers.test.ts +++ b/packages/types/src/__tests__/provider-identifiers.test.ts @@ -114,6 +114,7 @@ describe("provider identifiers", () => { providerIdentifiers.nanogpt, providerIdentifiers.ioIntelligence, providerIdentifiers.kimiCode, + providerIdentifiers.friendli, ]) expect(localProviders).toEqual([providerIdentifiers.ollama, providerIdentifiers.lmstudio]) expect(internalProviders).toEqual([providerIdentifiers.vscodeLm, providerIdentifiers.githubCopilot]) diff --git a/packages/types/src/provider-settings.ts b/packages/types/src/provider-settings.ts index ebb48b041b..a61a0321af 100644 --- a/packages/types/src/provider-settings.ts +++ b/packages/types/src/provider-settings.ts @@ -102,6 +102,7 @@ export const dynamicProviders = [ providerIdentifiers.nanogpt, providerIdentifiers.ioIntelligence, providerIdentifiers.kimiCode, + providerIdentifiers.friendli, ] as const export type DynamicProvider = (typeof dynamicProviders)[number] diff --git a/src/api/providers/__tests__/friendli.spec.ts b/src/api/providers/__tests__/friendli.spec.ts index acf03687f3..8cfa35effb 100644 --- a/src/api/providers/__tests__/friendli.spec.ts +++ b/src/api/providers/__tests__/friendli.spec.ts @@ -3,17 +3,25 @@ import { Anthropic } from "@anthropic-ai/sdk" import OpenAI from "openai" -import { friendliDefaultModelId, friendliModels } from "@roo-code/types" +import { friendliDefaultModelId, friendliModels, openAiModelInfoSaneDefaults } from "@roo-code/types" import { buildApiHandler } from "../../index" import { getModelMaxOutputTokens } from "../../../shared/api" import { FriendliHandler } from "../friendli" +import { NOT_PROVIDED } from "../constants" import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" import { clearAllMocks } from "../../../test-utils/reset" import { providerIdentifiers } from "@roo-code/types/provider-identifiers" // Create mock functions const mockCreate = vi.fn() +const { mockGetModels, mockGetFriendliModels } = vi.hoisted(() => ({ + mockGetModels: vi.fn(async () => ({})), + mockGetFriendliModels: vi.fn(), +})) + +vi.mock("../fetchers/modelCache", () => ({ getModels: mockGetModels })) +vi.mock("../fetchers/friendli", () => ({ getFriendliModels: mockGetFriendliModels })) // Mock OpenAI module vi.mock("openai", () => ({ @@ -80,8 +88,280 @@ describe("FriendliHandler", () => { expect(OpenAI).toHaveBeenCalledWith(expect.objectContaining({ apiKey: friendliApiKey })) }) - it("should throw error when API key is not provided", () => { - expect(() => new FriendliHandler({})).toThrow("API key is required") + it("should construct without an API key and without fetching the catalog", () => { + expect(() => new FriendliHandler({})).not.toThrow() + expect(mockGetModels).not.toHaveBeenCalled() + }) + + it.each([ + { entry: "createMessage", run: (h: FriendliHandler) => h.createMessage("system", []).next() }, + { entry: "completePrompt", run: (h: FriendliHandler) => h.completePrompt("a") }, + ])("keyless $entry sends the placeholder key and surfaces a 401 to the caller", async ({ run }) => { + const unauthorized = Object.assign(new Error("Incorrect API key provided"), { status: 401 }) + mockCreate.mockRejectedValueOnce(unauthorized) + const keyless = new FriendliHandler({}) + + await expect(run(keyless)).rejects.toMatchObject({ status: 401 }) + expect(vi.mocked(OpenAI).mock.lastCall?.[0]).toEqual(expect.objectContaining({ apiKey: NOT_PROVIDED })) + }) + + it("should load the catalog once before building requests for fetched-only models", async () => { + const info = { maxTokens: 8192, contextWindow: 100_000, supportsPromptCache: false } + mockGetModels.mockResolvedValueOnce({ "vendor/new-model": info }) + mockCreate.mockResolvedValue({ choices: [{ message: { content: "ok" } }] }) + const handlerWithModel = new FriendliHandler({ apiModelId: "vendor/new-model", friendliApiKey: "test-key" }) + + await handlerWithModel.completePrompt("a") + await handlerWithModel.completePrompt("b") + + expect(mockGetModels).toHaveBeenCalledTimes(1) + expect(handlerWithModel.getModel()).toEqual(expect.objectContaining({ id: "vendor/new-model", info })) + }) + + it("should wait for the catalog before streaming", async () => { + const info = { maxTokens: 8192, contextWindow: 100_000, supportsPromptCache: false } + mockGetModels.mockResolvedValueOnce({ "vendor/new-model": info }) + const handlerWithModel = new FriendliHandler({ apiModelId: "vendor/new-model", friendliApiKey: "test-key" }) + + await handlerWithModel.createMessage("system", []).next() + + expect(mockCreate.mock.calls[0][0]).toEqual(expect.objectContaining({ max_tokens: 8192 })) + }) + + it("should not start a chat request when cancelled during catalog loading", async () => { + mockGetModels.mockReturnValueOnce(new Promise(() => {})) + const controller = new AbortController() + const pending = handler.createMessage("system", [], { taskId: "t", abortSignal: controller.signal }).next() + + controller.abort() + + await expect(pending).rejects.toMatchObject({ name: "AbortError" }) + expect(mockCreate).not.toHaveBeenCalled() + }) + + it("ensureModelFetched should expose the fetched context window before any request", async () => { + const info = { maxTokens: 8192, contextWindow: 100_000, supportsPromptCache: false } + mockGetModels.mockResolvedValueOnce({ "vendor/new-model": info }) + const handlerWithModel = new FriendliHandler({ apiModelId: "vendor/new-model", friendliApiKey: "k" }) + + await handlerWithModel.ensureModelFetched() + expect(handlerWithModel.getModel().info.contextWindow).toBe(100_000) + + await handlerWithModel.completePrompt("a") + expect(mockGetModels).toHaveBeenCalledTimes(1) + }) + + it("completePrompt should not start a chat request when cancelled during catalog loading", async () => { + mockGetModels.mockReturnValueOnce(new Promise(() => {})) + const controller = new AbortController() + const pending = handler.completePrompt("a", { abortSignal: controller.signal }) + + controller.abort() + + await expect(pending).rejects.toMatchObject({ name: "AbortError" }) + expect(mockCreate).not.toHaveBeenCalled() + }) + + it.each([ + { + entry: "createMessage", + run: (h: FriendliHandler, signal: AbortSignal) => + h.createMessage("system", [], { taskId: "t", abortSignal: signal }).next(), + }, + { + entry: "completePrompt", + run: (h: FriendliHandler, signal: AbortSignal) => h.completePrompt("a", { abortSignal: signal }), + }, + ])("$entry should not start a chat request when cancelled as the catalog settles", async ({ run }) => { + const controller = new AbortController() + mockGetModels.mockImplementationOnce(async () => { + controller.abort() + return { "vendor/new-model": { contextWindow: 100_000, supportsPromptCache: false } } + }) + + await expect(run(handler, controller.signal)).rejects.toMatchObject({ name: "AbortError" }) + expect(mockCreate).not.toHaveBeenCalled() + }) + + it("createMessage should forward a per-request signal and normalize a mid-stream abort", async () => { + const controller = new AbortController() + mockCreate.mockImplementationOnce(async function* () { + yield { choices: [{ delta: { content: "partial" }, index: 0 }] } + controller.abort() + throw new Error("Request was aborted.") + }) + + const stream = handler.createMessage("system", [], { taskId: "t", abortSignal: controller.signal }) + + await expect(collectStream(stream)).rejects.toMatchObject({ + name: "AbortError", + message: "The Friendli request was aborted", + }) + const forwarded = (mockCreate.mock.calls[0][1] as { signal: AbortSignal }).signal + expect(forwarded).not.toBe(controller.signal) + expect(forwarded.aborted).toBe(true) + }) + + describe("caller signal listener cleanup", () => { + // Mirrors openai@5 client.js: an anonymous abort listener that is never removed. + const leakyCreate = (body: () => AsyncGenerator) => + mockCreate.mockImplementation((_params: unknown, opts?: { signal?: AbortSignal }) => { + opts?.signal?.addEventListener("abort", () => {}) + return body() + }) + const trackListeners = (signal: AbortSignal) => { + const live = new Set() + const add = signal.addEventListener.bind(signal) + const remove = signal.removeEventListener.bind(signal) + vi.spyOn(signal, "addEventListener").mockImplementation((type, fn, opts) => { + live.add(fn) + add(type, fn, opts) + }) + vi.spyOn(signal, "removeEventListener").mockImplementation((type, fn, opts) => { + live.delete(fn) + remove(type, fn, opts) + }) + return live + } + const chunk = { choices: [{ delta: { content: "x" }, index: 0 }] } + + it("detaches from a signal reused across completed streams", async () => { + leakyCreate(() => asyncStreamFrom([chunk])) + const { signal } = new AbortController() + const live = trackListeners(signal) + + for (let i = 0; i < 3; i++) { + await collectStream(handler.createMessage("system", [], { taskId: "t", abortSignal: signal })) + } + + expect(mockCreate).toHaveBeenCalledTimes(3) + expect(live.size).toBe(0) + }) + + it.each([ + { + exit: "a failed stream", + body: async function* () { + yield chunk + throw new Error("boom") + }, + consume: (s: AsyncGenerator) => expect(collectStream(s)).rejects.toThrow("boom"), + }, + { + exit: "an early generator exit", + body: async function* () { + yield chunk + yield chunk + }, + consume: async (s: AsyncGenerator) => { + await s.next() + await s.return(undefined) + }, + }, + ])("detaches after $exit", async ({ body, consume }) => { + leakyCreate(body) + const { signal } = new AbortController() + const live = trackListeners(signal) + + await consume(handler.createMessage("system", [], { taskId: "t", abortSignal: signal })) + + expect(live.size).toBe(0) + }) + + it("completePrompt detaches from a signal reused across keyless calls that succeed and fail", async () => { + const keyless = new FriendliHandler({}) + let calls = 0 + mockCreate.mockImplementation(async (_params: unknown, opts?: { signal?: AbortSignal }) => { + opts?.signal?.addEventListener("abort", () => {}) + if (++calls === 3) throw new Error("boom") + if (calls === 4) { + // In flight until the forwarded signal aborts. + await new Promise((resolve) => opts?.signal?.addEventListener("abort", resolve)) + throw new Error("Request was aborted.") + } + return { choices: [{ message: { content: "ok" } }] } + }) + const controller = new AbortController() + const live = trackListeners(controller.signal) + + await keyless.completePrompt("a", { abortSignal: controller.signal }) + await keyless.completePrompt("b", { abortSignal: controller.signal }) + await expect(keyless.completePrompt("c", { abortSignal: controller.signal })).rejects.toThrow("boom") + + expect(mockCreate).toHaveBeenCalledTimes(3) + expect(live.size).toBe(0) + // The bridge still follows the caller after earlier calls detached. + const pending = keyless.completePrompt("d", { abortSignal: controller.signal }) + await vi.waitFor(() => expect(mockCreate).toHaveBeenCalledTimes(4)) + controller.abort() + await expect(pending).rejects.toThrow() + expect((mockCreate.mock.lastCall?.[1] as { signal: AbortSignal }).signal.aborted).toBe(true) + }) + }) + + it("should retry an empty catalog and keep the next nonempty one", async () => { + const info = { maxTokens: 8192, contextWindow: 100_000, supportsPromptCache: false } + mockGetModels.mockResolvedValueOnce({}).mockResolvedValueOnce({ "vendor/new-model": info }) + const handlerWithModel = new FriendliHandler({ apiModelId: "vendor/new-model", friendliApiKey: "k" }) + + await handlerWithModel.ensureModelFetched() + expect(handlerWithModel.getModel().info).toEqual(openAiModelInfoSaneDefaults) + + await handlerWithModel.ensureModelFetched() + expect(handlerWithModel.getModel().info).toBe(info) + + await handlerWithModel.ensureModelFetched() + expect(mockGetModels).toHaveBeenCalledTimes(2) + expect(handlerWithModel.getModel().info).toBe(info) + }) + + it("completePrompt should give the completion only the time left after the catalog wait", async () => { + let now = 1_000 + vi.spyOn(Date, "now").mockImplementation(() => now) + mockGetModels.mockImplementationOnce(async () => { + now += 1_000 + return {} + }) + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "ok" } }] }) + + await handler.completePrompt("a", { timeoutMs: 3000 }) + + expect(mockCreate).toHaveBeenLastCalledWith(expect.anything(), { timeout: 2000 }) + }) + + it("completePrompt should time out without a chat request when the catalog wait uses the deadline", async () => { + mockGetModels.mockReturnValueOnce(new Promise(() => {})) + + await expect(handler.completePrompt("a", { timeoutMs: 20 })).rejects.toMatchObject({ name: "TimeoutError" }) + expect(mockCreate).not.toHaveBeenCalled() + }) + + it("should fall back to static models when the catalog is empty", async () => { + mockCreate.mockResolvedValue({ choices: [{ message: { content: "ok" } }] }) + await handler.completePrompt("a") + + expect(handler.getModel().info).toEqual(friendliModels[friendliDefaultModelId]) + }) + + it("should keep an unknown model id with default metadata when the catalog fails", async () => { + mockGetModels.mockRejectedValueOnce(new Error("Network error")) + mockCreate.mockResolvedValue({ choices: [{ message: { content: "ok" } }] }) + const handlerWithModel = new FriendliHandler({ apiModelId: "vendor/unknown", friendliApiKey: "test-key" }) + + await handlerWithModel.completePrompt("a") + + expect(handlerWithModel.getModel()).toEqual( + expect.objectContaining({ id: "vendor/unknown", info: openAiModelInfoSaneDefaults }), + ) + }) + + it("fetchModels should load the public /v1/models catalog", async () => { + const fetched = { "vendor/new-model": { contextWindow: 100_000, supportsPromptCache: false } } + mockGetFriendliModels.mockResolvedValueOnce(fetched) + const signal = new AbortController().signal + + expect(await handler.fetchModels({ provider: providerIdentifiers.friendli }, signal)).toBe(fetched) + expect(mockGetFriendliModels).toHaveBeenCalledWith({ signal }) }) it("should return default model when no model is specified", () => { @@ -181,11 +461,13 @@ describe("FriendliHandler", () => { expect(result).toBe(expectedResponse) }) - it("completePrompt should forward the abort signal", async () => { + it("completePrompt should forward a per-request signal instead of the caller's", async () => { mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "ok" } }] }) - const { signal } = new AbortController() - await handler.completePrompt("test prompt", { abortSignal: signal }) - expect(mockCreate).toHaveBeenLastCalledWith(expect.anything(), { signal }) + const controller = new AbortController() + await handler.completePrompt("test prompt", { abortSignal: controller.signal }) + const forwarded = (mockCreate.mock.lastCall?.[1] as { signal: AbortSignal }).signal + expect(forwarded).not.toBe(controller.signal) + expect(forwarded.aborted).toBe(false) }) it.each([ @@ -193,6 +475,7 @@ describe("FriendliHandler", () => { { label: "a fractional", input: 12.5, forwarded: 13 }, { label: "an over-cap", input: 2_147_483_648, forwarded: 2_147_483_647 }, ])("completePrompt should forward $label timeoutMs as $forwarded", async ({ input, forwarded }) => { + vi.spyOn(Date, "now").mockReturnValue(1_000) // no time passes during the catalog wait mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "ok" } }] }) await handler.completePrompt("test prompt", { timeoutMs: input }) expect(mockCreate).toHaveBeenLastCalledWith(expect.anything(), { timeout: forwarded }) @@ -567,4 +850,81 @@ describe("FriendliHandler — Friendli-specific reasoning params", () => { expect(callArgs.parse_reasoning).toBe(true) expect(callArgs.include_reasoning).toBe(true) }) + + describe("toggle-only models (supportsReasoningBinary)", () => { + const toggleOnly = { contextWindow: 202_752, supportsPromptCache: false, supportsReasoningBinary: true } + + const requestFor = async (settings: { enableReasoningEffort?: boolean; reasoningEffort?: "high" }) => { + mockGetModels.mockResolvedValueOnce({ "zai-org/GLM-5.1": toggleOnly }) + mockCreate.mockImplementationOnce(() => asyncStreamFrom([])) + const handler = new FriendliHandler({ apiModelId: "zai-org/GLM-5.1", ...settings }) + await handler.createMessage("system", []).next() + return mockCreate.mock.calls[0][0] as Record + } + + it("sends enable_thinking: true with reasoning parsing when enabled, without an effort", async () => { + // A stale effort from a previous model must not leak into a toggle-only request. + const callArgs = await requestFor({ enableReasoningEffort: true, reasoningEffort: "high" }) + + expect(callArgs.chat_template_kwargs).toEqual({ enable_thinking: true }) + expect(callArgs.parse_reasoning).toBe(true) + expect(callArgs.include_reasoning).toBe(true) + expect(callArgs).not.toHaveProperty("reasoning_effort") + }) + + it.each([{ enableReasoningEffort: false }, {}])( + "sends enable_thinking: false when the toggle is off or unset (%o)", + async (settings) => { + const callArgs = await requestFor(settings) + + expect(callArgs.chat_template_kwargs).toEqual({ enable_thinking: false }) + expect(callArgs).not.toHaveProperty("parse_reasoning") + expect(callArgs).not.toHaveProperty("include_reasoning") + expect(callArgs).not.toHaveProperty("reasoning_effort") + }, + ) + + it("completePrompt sends the same toggle params", async () => { + mockGetModels.mockResolvedValueOnce({ "zai-org/GLM-5.1": toggleOnly }) + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "ok" } }] }) + const handler = new FriendliHandler({ apiModelId: "zai-org/GLM-5.1", enableReasoningEffort: true }) + + await handler.completePrompt("test") + + const callArgs = mockCreate.mock.calls[0][0] as Record + expect(callArgs.chat_template_kwargs).toEqual({ enable_thinking: true }) + expect(callArgs).not.toHaveProperty("reasoning_effort") + }) + }) +}) + +describe("FriendliHandler — reserved model IDs", () => { + beforeEach(() => { + clearAllMocks() + mockCreate.mockResolvedValue({ choices: [{ message: { content: "ok" } }] }) + }) + + it.each(["toString", "constructor", "__proto__"])( + "does not resolve inherited members for %s (catalog as a plain object)", + async (id) => { + mockGetModels.mockResolvedValueOnce({}) + const handler = new FriendliHandler({ apiModelId: id }) + + await handler.completePrompt("a") + + expect(handler.getModel()).toEqual(expect.objectContaining({ id, info: openAiModelInfoSaneDefaults })) + }, + ) + + it("resolves a fetched model whose upstream ID is __proto__", async () => { + const info = { contextWindow: 1000, supportsPromptCache: false } + const catalog: Record = Object.create(null) + catalog["__proto__"] = info + mockGetModels.mockResolvedValueOnce(catalog) + const handler = new FriendliHandler({ apiModelId: "__proto__" }) + + await handler.completePrompt("a") + + expect(handler.getModel().info).toBe(info) + }) }) diff --git a/src/api/providers/fetchers/__tests__/friendli.spec.ts b/src/api/providers/fetchers/__tests__/friendli.spec.ts new file mode 100644 index 0000000000..98bd36581e --- /dev/null +++ b/src/api/providers/fetchers/__tests__/friendli.spec.ts @@ -0,0 +1,164 @@ +// npx vitest run api/providers/fetchers/__tests__/friendli.spec.ts + +import axios from "axios" + +import { getFriendliModels } from "../friendli" + +vi.mock("axios") + +describe("getFriendliModels", () => { + it("maps chat models from /v1/models", async () => { + vi.mocked(axios.get).mockResolvedValueOnce({ + data: { + data: [ + { + id: "zai-org/GLM-5.3", + mode: "chat", + description: "GLM", + context_length: 1_048_576, + max_completion_tokens: 131_072, + pricing: { input: "0.00000126", output: "0.00000396", input_cache_read: "0.000000234" }, + input_modalities: ["text", "image"], + reasoning_options: [ + { type: "effort", values: ["low", "high", "max", "ultracode"] }, + { type: "budget_tokens", min: -1, max: 1_048_576 }, + ], + }, + { + id: "google/gemma-4-31B-it", + mode: "chat", + pricing: { input: "0.00000014", output: "0.0000004" }, + }, + { id: "zai-org/GLM-5.2", mode: "chat", input_modalities: ["text"] }, + { id: "some/embedding", mode: "embedding" }, + ], + }, + }) + + const signal = new AbortController().signal + const models = await getFriendliModels({ signal }) + + expect(axios.get).toHaveBeenCalledWith("https://api.friendli.ai/serverless/v1/models", { signal }) + expect(Object.keys(models)).toEqual(["zai-org/GLM-5.3", "google/gemma-4-31B-it", "zai-org/GLM-5.2"]) + expect(models["zai-org/GLM-5.3"]).toEqual({ + maxTokens: 131_072, + contextWindow: 1_048_576, + supportsImages: true, + supportsPromptCache: true, + supportsMaxTokens: true, + inputPrice: expect.closeTo(1.26), + outputPrice: expect.closeTo(3.96), + cacheReadsPrice: expect.closeTo(0.234), + supportsReasoningEffort: ["low", "high", "max"], + reasoningEffort: "high", + description: "GLM", + }) + expect(models["google/gemma-4-31B-it"].supportsReasoningEffort).toBeUndefined() + expect(models["google/gemma-4-31B-it"].supportsPromptCache).toBe(false) + // Missing input_modalities and a text-only list both mean no image input. + expect(models["google/gemma-4-31B-it"].supportsImages).toBe(false) + expect(models["zai-org/GLM-5.2"].supportsImages).toBe(false) + }) + + it("accepts null reasoning_options and keeps the catalog", async () => { + vi.mocked(axios.get).mockResolvedValueOnce({ + data: { + data: [ + { id: "vendor/fetched-only", reasoning_options: null }, + { id: "zai-org/GLM-5.1", reasoning_options: [{ type: "toggle" }] }, + ], + }, + }) + + const models = await getFriendliModels() + + expect(Object.keys(models)).toEqual(["vendor/fetched-only", "zai-org/GLM-5.1"]) + expect(models["vendor/fetched-only"].supportsReasoningEffort).toBeUndefined() + expect(models["vendor/fetched-only"].supportsReasoningBinary).toBeUndefined() + }) + + it("maps toggle-only reasoning to supportsReasoningBinary without inventing an effort", async () => { + vi.mocked(axios.get).mockResolvedValueOnce({ + data: { + data: [ + { + id: "zai-org/GLM-5.1", + reasoning_options: [{ type: "toggle" }, { type: "budget_tokens", min: -1, max: 202_752 }], + }, + { + id: "zai-org/GLM-5.2", + reasoning_options: [{ type: "toggle" }, { type: "effort", values: ["high", "max"] }], + }, + ], + }, + }) + + const models = await getFriendliModels() + + expect(models["zai-org/GLM-5.1"].supportsReasoningBinary).toBe(true) + expect(models["zai-org/GLM-5.1"].supportsReasoningEffort).toBeUndefined() + expect(models["zai-org/GLM-5.1"].reasoningEffort).toBeUndefined() + // A toggle next to effort levels keeps the effort selector. + expect(models["zai-org/GLM-5.2"].supportsReasoningBinary).toBeUndefined() + expect(models["zai-org/GLM-5.2"].supportsReasoningEffort).toEqual(["high", "max"]) + }) + + it("defaults to the first supported effort when high is not offered", async () => { + vi.mocked(axios.get).mockResolvedValueOnce({ + data: { + data: [ + { + id: "vendor/no-high", + reasoning_options: [{ type: "effort", values: ["ultracode", "low", "medium", "max"] }], + }, + ], + }, + }) + + const models = await getFriendliModels() + + expect(models["vendor/no-high"].supportsReasoningEffort).toEqual(["low", "medium", "max"]) + expect(models["vendor/no-high"].reasoningEffort).toBe("low") + }) + + it("keeps reserved upstream IDs as own keys", async () => { + vi.mocked(axios.get).mockResolvedValueOnce({ + data: { + data: [ + { id: "__proto__", context_length: 1000 }, + { id: "toString", context_length: 2000 }, + ], + }, + }) + + const models = await getFriendliModels() + + expect(Object.getPrototypeOf(models)).toBeNull() + expect(Object.keys(models)).toEqual(["__proto__", "toString"]) + expect(models["__proto__"].contextWindow).toBe(1000) + expect(models["toString"].contextWindow).toBe(2000) + expect(models["valueOf"]).toBeUndefined() + }) + + it("rethrows cancellation instead of returning an empty record", async () => { + const controller = new AbortController() + controller.abort() + vi.mocked(axios.get).mockRejectedValueOnce(new Error("canceled")) + + await expect(getFriendliModels({ signal: controller.signal })).rejects.toMatchObject({ name: "AbortError" }) + }) + + it("returns an empty record on an invalid response", async () => { + vi.spyOn(console, "error").mockImplementation(() => {}) + vi.mocked(axios.get).mockResolvedValueOnce({ data: { invalid: true } }) + + expect(await getFriendliModels()).toEqual({}) + }) + + it("returns an empty record on request failure", async () => { + vi.spyOn(console, "error").mockImplementation(() => {}) + vi.mocked(axios.get).mockRejectedValueOnce(new Error("Network error")) + + expect(await getFriendliModels()).toEqual({}) + }) +}) diff --git a/src/api/providers/fetchers/friendli.ts b/src/api/providers/fetchers/friendli.ts new file mode 100644 index 0000000000..6a85f6cbb5 --- /dev/null +++ b/src/api/providers/fetchers/friendli.ts @@ -0,0 +1,86 @@ +import axios from "axios" +import { z } from "zod" + +import { + type ModelInfo, + type ReasoningEffortExtended, + openAiModelInfoSaneDefaults, + reasoningEffortsExtended, +} from "@roo-code/types" + +import { parseApiPrice } from "../../../shared/cost" +import { throwIfAborted } from "../utils/abort-signal" + +const friendliModelSchema = z.object({ + id: z.string(), + mode: z.string().optional(), + description: z.string().optional(), + context_length: z.number().optional(), + max_completion_tokens: z.number().optional(), + pricing: z + .object({ + input: z.string().optional(), + output: z.string().optional(), + input_cache_read: z.string().optional(), + }) + .optional(), + input_modalities: z.array(z.string()).optional(), + reasoning_options: z.array(z.object({ type: z.string(), values: z.array(z.string()).optional() })).nullish(), +}) + +type FriendliModel = z.infer + +const friendliModelsResponseSchema = z.object({ data: z.array(friendliModelSchema) }) + +const parseFriendliModel = (model: FriendliModel): ModelInfo => { + const efforts = model.reasoning_options + ?.find((option) => option.type === "effort") + ?.values?.filter((v): v is ReasoningEffortExtended => + (reasoningEffortsExtended as readonly string[]).includes(v), + ) + // A toggle without effort levels (e.g. GLM-5.1, Gemma) is an on/off switch, not an effort scale. + const toggleOnly = !efforts?.length && !!model.reasoning_options?.some((option) => option.type === "toggle") + + return { + maxTokens: model.max_completion_tokens, + contextWindow: model.context_length ?? openAiModelInfoSaneDefaults.contextWindow, + supportsImages: model.input_modalities?.includes("image") ?? false, + supportsPromptCache: model.pricing?.input_cache_read !== undefined, + supportsMaxTokens: true, + inputPrice: parseApiPrice(model.pricing?.input), + outputPrice: parseApiPrice(model.pricing?.output), + cacheReadsPrice: parseApiPrice(model.pricing?.input_cache_read), + ...(efforts?.length && { + supportsReasoningEffort: efforts, + reasoningEffort: efforts.includes("high") ? "high" : efforts[0], + }), + ...(toggleOnly && { supportsReasoningBinary: true }), + description: model.description, + } +} + +/** Fetches the public (no auth) Friendli model catalog. */ +export async function getFriendliModels(opts?: { signal?: AbortSignal }): Promise> { + // Null prototype: upstream IDs such as "__proto__" or "toString" must stay plain keys. + const models: Record = Object.create(null) + + try { + const response = await axios.get("https://api.friendli.ai/serverless/v1/models", { signal: opts?.signal }) + const result = friendliModelsResponseSchema.safeParse(response.data) + + if (!result.success) { + console.error(`Friendli models response is invalid: ${JSON.stringify(result.error.format())}`) + return models + } + + for (const model of result.data.data) { + if (model.mode && model.mode !== "chat") continue + models[model.id] = parseFriendliModel(model) + } + } catch (error) { + throwIfAborted(opts?.signal) + console.error(`Error fetching Friendli models: ${error instanceof Error ? error.message : String(error)}`) + } + + return models +} diff --git a/src/api/providers/friendli.ts b/src/api/providers/friendli.ts index 86d57a9d2d..1c8c276ce6 100644 --- a/src/api/providers/friendli.ts +++ b/src/api/providers/friendli.ts @@ -1,9 +1,16 @@ import { Anthropic } from "@anthropic-ai/sdk" import OpenAI from "openai" -import { type FriendliModelId, friendliDefaultModelId, friendliModels } from "@roo-code/types" +import { + type FriendliModelId, + type ModelRecord, + friendliDefaultModelId, + friendliModels, + openAiModelInfoSaneDefaults, + providerIdentifiers, +} from "@roo-code/types" -import type { ApiHandlerOptions } from "../../shared/api" +import type { ApiHandlerOptions, GetModelsOptions } from "../../shared/api" import { shouldUseReasoningEffort, getModelMaxOutputTokens } from "../../shared/api" import { convertToOpenAiMessages } from "../transform/openai-format" @@ -11,7 +18,20 @@ import { getModelParams } from "../transform/model-params" import { BaseOpenAiCompatibleProvider } from "./base-openai-compatible-provider" import { handleOpenAIError } from "./utils/error-handler" -import { hasUsableTimeout, toSdkTimeoutMs } from "./utils/abort-signal" +import { + bridgeExternalAbortSignal, + createAbortError, + createTimeoutError, + hasUsableTimeout, + isRequestAborted, + isTimeoutAbort, + mergeAbortSignalAndTimeout, + toSdkTimeoutMs, +} from "./utils/abort-signal" +import { NOT_PROVIDED } from "./constants" +import { getModels } from "./fetchers/modelCache" +import { getFriendliModels } from "./fetchers/friendli" +import type { ApiStream } from "../transform/stream" import type { ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" /** @@ -55,28 +75,76 @@ type FriendliChatCompletionNonStreamingParams = Omit< * reasoning parameters that the base class doesn't know about. */ export class FriendliHandler extends BaseOpenAiCompatibleProvider { - /** - * @param options Provider settings; `friendliApiKey` is required. - */ + private models: ModelRecord = Object.create(null) + constructor(options: ApiHandlerOptions) { super({ ...options, providerName: "Friendli", baseURL: "https://api.friendli.ai/serverless/v1", - apiKey: options.friendliApiKey, + apiKey: options.friendliApiKey || NOT_PROVIDED, defaultProviderModelId: friendliDefaultModelId, providerModels: friendliModels, defaultTemperature: 0.6, }) } - override getModel() { - const id = - this.options.apiModelId && this.options.apiModelId in this.providerModels - ? (this.options.apiModelId as FriendliModelId) - : this.defaultProviderModelId + /** Loads the fetched catalog so task preparation reads its context window before any request. */ + async ensureModelFetched(signal?: AbortSignal): Promise { + await this.ensureCatalogLoaded(async () => { + this.models = await getModels({ provider: providerIdentifiers.friendli }, this) + // getModels() does not cache an empty response, so keep it retryable. + return Object.keys(this.models).length > 0 + }, signal) + } + + /** + * Waits for the catalog before a request. A failed load keeps the static `friendliModels` + * fallback; a cancelled or timed-out wait rejects so no request starts. + */ + private async waitForCatalog(abortSignal?: AbortSignal, timeoutMs?: number): Promise { + const signal = mergeAbortSignalAndTimeout(abortSignal, timeoutMs) + try { + await this.ensureModelFetched(signal) + } catch { + // Classified below: only a cut-short wait is fatal. + } + if (isTimeoutAbort(signal)) throw createTimeoutError(this.providerName) + if (signal?.aborted) throw createAbortError(this.providerName) + } + + override async *createMessage( + systemPrompt: string, + messages: Anthropic.Messages.MessageParam[], + metadata?: ApiHandlerCreateMessageMetadata, + ): ApiStream { + await this.waitForCatalog(metadata?.abortSignal) + // The SDK never removes its listener from the signal it receives, so hand it a + // per-request signal and detach that from the caller's long-lived one when the stream ends. + const { signal, dispose } = bridgeExternalAbortSignal(metadata?.abortSignal) + try { + yield* super.createMessage(systemPrompt, messages, metadata && { ...metadata, abortSignal: signal }) + } catch (error) { + // The SDK's abort error is not the AbortError contract callers check for. + if (isRequestAborted(error, signal)) throw createAbortError(this.providerName) + throw error + } finally { + dispose() + } + } - const info = this.providerModels[id] + async fetchModels(_options: GetModelsOptions, signal?: AbortSignal): Promise { + return getFriendliModels({ signal }) + } + + override getModel() { + const id = (this.options.apiModelId || this.defaultProviderModelId) as FriendliModelId + // Own-property reads: the cached catalog may come back as a plain object, so an ID + // such as "toString" must not resolve to an inherited member. + const info = + (Object.hasOwn(this.models, id) ? this.models[id] : undefined) ?? + (Object.hasOwn(this.providerModels, id) ? this.providerModels[id] : undefined) ?? + openAiModelInfoSaneDefaults const params = getModelParams({ format: "openai", modelId: id, @@ -98,11 +166,25 @@ export class FriendliHandler extends BaseOpenAiCompatibleProvider { const { info: modelInfo, reasoningEffort } = this.getModel() const extra: Partial = {} + if (modelInfo.supportsReasoningBinary && !modelInfo.supportsReasoningEffort) { + if (!this.options.enableReasoningEffort) { + // The chat template defaults thinking on, so "off" must be explicit. + extra.chat_template_kwargs = { enable_thinking: false } + return extra + } + extra.parse_reasoning = true + extra.include_reasoning = true + extra.chat_template_kwargs = { enable_thinking: true } + return extra + } + const isControllableReasoning = Array.isArray(modelInfo.supportsReasoningEffort) const useReasoningEffort = modelInfo.supportsReasoningEffort @@ -173,10 +255,11 @@ export class FriendliHandler extends BaseOpenAiCompatibleProvider { + // The catalog wait and the completion share one deadline, started before the wait. + const timeoutMs = options?.timeoutMs + const deadline = hasUsableTimeout(timeoutMs) ? Date.now() + timeoutMs : undefined + await this.waitForCatalog(options?.abortSignal, timeoutMs) + const remainingMs = deadline === undefined ? undefined : deadline - Date.now() + if (remainingMs !== undefined && remainingMs <= 0) throw createTimeoutError(this.providerName) const { id: modelId } = this.getModel() const friendliExtra = this.buildFriendliReasoningParams() @@ -193,16 +282,17 @@ export class FriendliHandler extends BaseOpenAiCompatibleProvider { "io-intelligence": mockModels, nanogpt: mockModels, "kimi-code": {}, + friendli: mockModels, }, values: undefined, }) @@ -6807,6 +6808,7 @@ describe("ClineProvider - Router Models", () => { .mockResolvedValueOnce(mockModels) // kenari (public endpoint) .mockResolvedValueOnce(mockModels) // io-intelligence (public endpoint) .mockResolvedValueOnce(mockModels) // nanogpt (public endpoint) + .mockResolvedValueOnce(mockModels) // friendli (public endpoint) await messageHandler({ type: "requestRouterModels" }) @@ -6830,6 +6832,7 @@ describe("ClineProvider - Router Models", () => { "io-intelligence": mockModels, nanogpt: mockModels, "kimi-code": {}, + friendli: mockModels, }, values: undefined, }) @@ -6936,6 +6939,7 @@ describe("ClineProvider - Router Models", () => { "io-intelligence": mockModels, nanogpt: mockModels, "kimi-code": {}, + friendli: mockModels, }, values: undefined, }) diff --git a/src/core/webview/__tests__/webviewMessageHandler.spec.ts b/src/core/webview/__tests__/webviewMessageHandler.spec.ts index bf50f08786..9cfe628214 100644 --- a/src/core/webview/__tests__/webviewMessageHandler.spec.ts +++ b/src/core/webview/__tests__/webviewMessageHandler.spec.ts @@ -1212,6 +1212,7 @@ describe("webviewMessageHandler - requestRouterModels", () => { "io-intelligence": mockModels, nanogpt: mockModels, "kimi-code": {}, + friendli: mockModels, }, values: undefined, }) @@ -1487,6 +1488,7 @@ describe("webviewMessageHandler - requestRouterModels", () => { "io-intelligence": mockModels, nanogpt: mockModels, "kimi-code": {}, + friendli: mockModels, }, values: undefined, }) @@ -1554,6 +1556,7 @@ describe("webviewMessageHandler - requestRouterModels", () => { "io-intelligence": mockModels, nanogpt: mockModels, "kimi-code": {}, + friendli: mockModels, }, values: undefined, }) diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index 4adb29264b..c5abbc0a12 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -1202,6 +1202,7 @@ export const webviewMessageHandler = async ( [providerIdentifiers.ioIntelligence]: {}, [providerIdentifiers.nanogpt]: {}, [providerIdentifiers.kimiCode]: {}, + [providerIdentifiers.friendli]: {}, } const safeGetModels = async (options: GetModelsOptions, handler: ApiHandler): Promise => { @@ -1415,6 +1416,9 @@ export const webviewMessageHandler = async ( options: { provider: providerIdentifiers.nanogpt, apiKey: nanoGptApiKey }, }) + // Friendli's /models catalog is public. + candidates.push({ key: providerIdentifiers.friendli, options: { provider: providerIdentifiers.friendli } }) + if (!providerFilter || providerFilter === providerIdentifiers.kimiCode) { const { kimiCodeOAuthManager } = await import("../../integrations/kimi-code/oauth") const kimiCodeAuthMethod = diff --git a/src/shared/api.ts b/src/shared/api.ts index d8cd547fdb..31197482ec 100644 --- a/src/shared/api.ts +++ b/src/shared/api.ts @@ -187,6 +187,7 @@ const _dynamicProviderExtras = { [providerIdentifiers.ioIntelligence]: {} as { apiKey?: string }, [providerIdentifiers.nanogpt]: {} as { apiKey?: string }, [providerIdentifiers.kimiCode]: {} as { apiKey?: string }, + [providerIdentifiers.friendli]: {} as object, } as const satisfies Record // Build the dynamic options union from the map, intersected with CommonFetchParams diff --git a/webview-ui/src/components/chat/ModelSelector.tsx b/webview-ui/src/components/chat/ModelSelector.tsx index 0e5b92b59b..063682e646 100644 --- a/webview-ui/src/components/chat/ModelSelector.tsx +++ b/webview-ui/src/components/chat/ModelSelector.tsx @@ -52,9 +52,12 @@ function useModelListForProvider(provider: ProviderName | undefined, apiConfigur const field = provider ? getModelIdField(provider) : undefined const isUnsupported = !provider || !field || UNSUPPORTED_QUICK_PICK_PROVIDERS.includes(provider) const isRouterProvider = !isStatic && !isUnsupported + // Friendli keeps a static list in MODELS_BY_PROVIDER, but it is only the fallback for the fetched catalog. + const fetchesStaticFallbackCatalog = provider === providerIdentifiers.friendli && !isUnsupported - const routerModels = useRouterModels({ provider, enabled: isRouterProvider }) - const routerModelsData = routerModels.data as Record> | undefined + const routerModels = useRouterModels({ provider, enabled: isRouterProvider || fetchesStaticFallbackCatalog }) + const allRouterModels = routerModels.data + const routerModelsData = allRouterModels as Record> | undefined const isLoading = isRouterProvider && routerModels.isLoading return React.useMemo(() => { @@ -63,11 +66,11 @@ function useModelListForProvider(provider: ProviderName | undefined, apiConfigur } const models = isStatic - ? getStaticModelsForProvider(provider, undefined, apiConfiguration) + ? getStaticModelsForProvider(provider, undefined, apiConfiguration, allRouterModels) : (routerModelsData?.[provider] ?? null) return { isUnsupported: false as const, field, models, isLoading } - }, [provider, field, isUnsupported, isStatic, apiConfiguration, routerModelsData, isLoading]) + }, [provider, field, isUnsupported, isStatic, apiConfiguration, allRouterModels, routerModelsData, isLoading]) } interface ModelSelectorProps { diff --git a/webview-ui/src/components/chat/__tests__/ModelSelector.spec.tsx b/webview-ui/src/components/chat/__tests__/ModelSelector.spec.tsx index d4dd381015..5632ccfa0f 100644 --- a/webview-ui/src/components/chat/__tests__/ModelSelector.spec.tsx +++ b/webview-ui/src/components/chat/__tests__/ModelSelector.spec.tsx @@ -1,6 +1,6 @@ import { render, screen, fireEvent } from "@/utils/test-utils" -import { providerIdentifiers, type ProviderSettings } from "@roo-code/types" +import { friendliModels, providerIdentifiers, type ProviderSettings } from "@roo-code/types" import { ModelSelector } from "../ModelSelector" @@ -373,6 +373,66 @@ describe("ModelSelector", () => { }) }) + describe("Friendli (fetched catalog with a static fallback)", () => { + const friendliConfig: ProviderSettings = { + apiProvider: providerIdentifiers.friendli, + apiModelId: "zai-org/GLM-5.2", + } + const staticIds = Object.keys(friendliModels).sort((a, b) => a.localeCompare(b)) + + const openPicker = () => { + render() + fireEvent.click(screen.getByTestId("model-selector-trigger")) + return screen.getAllByTestId("model-selector-item").map((el) => el.textContent) + } + + test("requests the Friendli catalog", () => { + openPicker() + + // useSelectedModel also calls useRouterModels, so require every Friendli call (incl. the picker's) to be enabled. + const friendliCalls = routerModelsResult.calls.filter( + (opts) => (opts as { provider?: string }).provider === providerIdentifiers.friendli, + ) + expect(friendliCalls.length).toBeGreaterThan(0) + expect(friendliCalls.every((opts) => (opts as { enabled: boolean }).enabled)).toBe(true) + }) + + test("lists the fetched catalog, including models missing from the static list", () => { + routerModelsResult.value = { + isLoading: false, + data: { friendli: { "zai-org/GLM-5.2": base, "zai-org/GLM-5.3": base } }, + } + + expect(openPicker()).toEqual(["zai-org/GLM-5.2", "zai-org/GLM-5.3"]) + }) + + test.each([ + ["while loading", { data: undefined, isLoading: true }], + ["after a failed request", { data: undefined, isLoading: false }], + ["when the catalog is empty", { data: { friendli: {} }, isLoading: false }], + ])("falls back to the static list %s", (_label, result) => { + routerModelsResult.value = result + + expect(openPicker()).toEqual(staticIds) + expect(screen.queryByText("chat:modelSelector.loading")).not.toBeInTheDocument() + }) + + test("replaces the static list with the fetched catalog in the same mounted picker", () => { + const listed = () => screen.getAllByTestId("model-selector-item").map((el) => el.textContent) + const { rerender } = render() + fireEvent.click(screen.getByTestId("model-selector-trigger")) + expect(listed()).toEqual(staticIds) + + routerModelsResult.value = { + isLoading: false, + data: { friendli: { "zai-org/GLM-5.3": base, "vendor/fetched-only": base } }, + } + rerender() + + expect(listed()).toEqual(["vendor/fetched-only", "zai-org/GLM-5.3"]) + }) + }) + test("falls back to a degraded trigger for providers without a simple model list, and jumps to Settings on click", () => { render( } -const { modelPickerMock, useOpenRouterModelProvidersMock, useSelectedModelMock } = vi.hoisted(() => ({ - modelPickerMock: vi.fn((_props: ComponentProps) => null), - useOpenRouterModelProvidersMock: vi.fn<() => OpenRouterModelProvidersQueryResult>(() => ({ data: undefined })), - useSelectedModelMock: vi.fn((configuration: ProviderSettings): UseSelectedModelReturn => ({ - provider: configuration.apiProvider, - id: configuration.apiModelId, - info: {}, - })), -})) +const { modelPickerMock, useOpenRouterModelProvidersMock, useRouterModelsMock, useSelectedModelMock } = vi.hoisted( + () => ({ + modelPickerMock: vi.fn((_props: ComponentProps) => null), + useRouterModelsMock: vi.fn((_opts: { provider?: string; enabled: boolean }) => ({ + data: {} as Record, + refetch: vi.fn(), + })), + useOpenRouterModelProvidersMock: vi.fn<() => OpenRouterModelProvidersQueryResult>(() => ({ data: undefined })), + useSelectedModelMock: vi.fn((configuration: ProviderSettings): UseSelectedModelReturn => ({ + provider: configuration.apiProvider, + id: configuration.apiModelId, + info: {}, + })), + }), +) vi.mock("@src/context/ExtensionStateContext", () => ({ useExtensionState: () => ({ @@ -57,7 +63,7 @@ vi.mock("@src/context/ExtensionStateContext", () => ({ })) vi.mock("@src/components/ui/hooks/useRouterModels", () => ({ - useRouterModels: () => ({ data: {}, refetch: vi.fn() }), + useRouterModels: useRouterModelsMock, })) vi.mock("@src/components/ui/hooks/useZooGatewayRouterModelsSync", () => ({ @@ -202,6 +208,7 @@ describe("ApiOptions interactions", () => { info: {}, })) useOpenRouterModelProvidersMock.mockImplementation(() => ({ data: undefined })) + useRouterModelsMock.mockImplementation(() => ({ data: {}, refetch: vi.fn() })) }) afterEach(() => { @@ -530,6 +537,18 @@ describe("ApiOptions interactions", () => { expect(setApiConfigurationField).toHaveBeenCalledWith("awsCustomArn", arn) }) + it("passes the fetched Friendli catalog to the settings model picker", () => { + const fetched = { "vendor/fetched-only": { contextWindow: 1000, supportsPromptCache: false } } + useRouterModelsMock.mockImplementation(() => ({ data: { friendli: fetched }, refetch: vi.fn() })) + + renderApiOptions({ + apiConfiguration: { apiProvider: providerIdentifiers.friendli, apiModelId: "vendor/fetched-only" }, + }) + + expect(useRouterModelsMock).toHaveBeenCalledWith({ enabled: true, provider: providerIdentifiers.friendli }) + expect(modelPickerMock.mock.lastCall?.[0].models).toBe(fetched) + }) + it("does not transform model IDs for non-Bedrock providers", () => { renderApiOptions({ apiConfiguration: { apiProvider: providerIdentifiers.anthropic, apiModelId: "claude-sonnet" }, diff --git a/webview-ui/src/components/settings/providers/__tests__/IOIntelligence.spec.tsx b/webview-ui/src/components/settings/providers/__tests__/IOIntelligence.spec.tsx index ae6f7bc63a..954fd92c8d 100644 --- a/webview-ui/src/components/settings/providers/__tests__/IOIntelligence.spec.tsx +++ b/webview-ui/src/components/settings/providers/__tests__/IOIntelligence.spec.tsx @@ -83,6 +83,7 @@ describe("IOIntelligence", () => { "meta-llama/Llama-3.3-70B-Instruct": { contextWindow: 128000, maxTokens: 8192, supportsPromptCache: false }, }, "kimi-code": {}, + friendli: {}, ollama: {}, lmstudio: {}, } diff --git a/webview-ui/src/components/settings/providers/__tests__/NanoGPT.spec.tsx b/webview-ui/src/components/settings/providers/__tests__/NanoGPT.spec.tsx index 698fa90e29..a78ab39a1b 100644 --- a/webview-ui/src/components/settings/providers/__tests__/NanoGPT.spec.tsx +++ b/webview-ui/src/components/settings/providers/__tests__/NanoGPT.spec.tsx @@ -105,6 +105,7 @@ describe("NanoGPT", () => { nanogpt: { "openai/test": { contextWindow: 1, maxTokens: 1, supportsPromptCache: false } }, "io-intelligence": {}, "kimi-code": {}, + friendli: {}, ollama: {}, lmstudio: {}, } diff --git a/webview-ui/src/components/settings/utils/__tests__/providerModelConfig.spec.ts b/webview-ui/src/components/settings/utils/__tests__/providerModelConfig.spec.ts index 4342ffdf93..ee0c1103d1 100644 --- a/webview-ui/src/components/settings/utils/__tests__/providerModelConfig.spec.ts +++ b/webview-ui/src/components/settings/utils/__tests__/providerModelConfig.spec.ts @@ -1,5 +1,7 @@ import { + type RouterModels, anthropicDefaultModelId, + friendliModels, ioIntelligenceDefaultModelId, mainlandZAiDefaultModelId, nanoGptDefaultModelId, @@ -212,6 +214,21 @@ describe("providerModelConfig", () => { expect(models["custom-arn"].description).toBe("Use Custom ARN") }) + it("prefers fetched Friendli models over the static fallback", () => { + const fetched = { "vendor/new-model": { contextWindow: 100_000, supportsPromptCache: false } } + const routerModels = { friendli: fetched } as unknown as RouterModels + + expect(getStaticModelsForProvider(providerIdentifiers.friendli, undefined, undefined, routerModels)).toBe( + fetched, + ) + expect(getStaticModelsForProvider(providerIdentifiers.friendli)).toBe(friendliModels) + expect( + getStaticModelsForProvider(providerIdentifiers.friendli, undefined, undefined, { + friendli: {}, + } as unknown as RouterModels), + ).toBe(friendliModels) + }) + it("returns empty object for providers without static models", () => { const models = getStaticModelsForProvider(providerIdentifiers.openrouter) expect(Object.keys(models).length).toBe(0) diff --git a/webview-ui/src/components/settings/utils/providerModelConfig.ts b/webview-ui/src/components/settings/utils/providerModelConfig.ts index 1bda52cf6f..474755e124 100644 --- a/webview-ui/src/components/settings/utils/providerModelConfig.ts +++ b/webview-ui/src/components/settings/utils/providerModelConfig.ts @@ -1,4 +1,4 @@ -import type { ProviderName, ModelInfo, ProviderSettings } from "@roo-code/types" +import type { ProviderName, ModelInfo, ProviderSettings, RouterModels } from "@roo-code/types" import { providerIdentifiers, anthropicDefaultModelId, @@ -183,7 +183,18 @@ export const getStaticModelsForProvider = ( provider: ProviderName, customArnLabel?: string, apiConfiguration?: ProviderSettings, + routerModels?: RouterModels, ): Record => { + // Friendli's static list is only a fallback for its fetched /v1/models list. + const friendliRouterModels = routerModels?.friendli + if ( + provider === providerIdentifiers.friendli && + friendliRouterModels && + Object.keys(friendliRouterModels).length > 0 + ) { + return friendliRouterModels + } + if (provider === providerIdentifiers.zai) { return getZAiModels(apiConfiguration?.zaiApiLine) } diff --git a/webview-ui/src/components/ui/hooks/__tests__/useRouterModels.spec.ts b/webview-ui/src/components/ui/hooks/__tests__/useRouterModels.spec.ts index 3744d9bc28..d8de70b7e0 100644 --- a/webview-ui/src/components/ui/hooks/__tests__/useRouterModels.spec.ts +++ b/webview-ui/src/components/ui/hooks/__tests__/useRouterModels.spec.ts @@ -1,6 +1,13 @@ +import React from "react" +import { QueryClientProvider } from "@tanstack/react-query" +import { renderHook, waitFor } from "@testing-library/react" + import { RouterModelsMessageType, providerIdentifiers } from "@roo-code/types" -import { fetchRouterModels } from "../useRouterModels" +import { createTestQueryClient } from "@/utils/test-utils" +import { vscode } from "@src/utils/vscode" + +import { fetchRouterModels, useRouterModels } from "../useRouterModels" vi.mock("@src/utils/vscode", () => ({ vscode: { @@ -84,3 +91,43 @@ describe("fetchRouterModels", () => { expect(removeSpy).toHaveBeenCalledWith("message", handler) }) }) + +describe("useRouterModels", () => { + afterEach(() => { + vi.restoreAllMocks() + }) + + test("keeps an own __proto__ model ID when a refetch replaces an empty catalog", async () => { + const provider = providerIdentifiers.friendli + // JSON.parse keeps "__proto__" as an own key, as a message from the extension host does. + const responses = [ + { [provider]: {} }, + JSON.parse( + `{"${provider}": {"__proto__": {"contextWindow": 1000}, "vendor/model": {"contextWindow": 2000}}}`, + ), + ] + vi.mocked(vscode.postMessage).mockImplementation(() => { + const routerModels = responses.shift() + queueMicrotask(() => + window.dispatchEvent( + new MessageEvent("message", { + data: { type: RouterModelsMessageType.routerModels, values: { provider }, routerModels }, + }), + ), + ) + }) + const queryClient = createTestQueryClient() + const wrapper = ({ children }: { children: React.ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children) + + const { result } = renderHook(() => useRouterModels({ provider }), { wrapper }) + await waitFor(() => expect(result.current.data).toEqual({ [provider]: {} })) + await result.current.refetch() + + await waitFor(() => expect(Object.hasOwn(result.current.data?.[provider] ?? {}, "__proto__")).toBe(true)) + const catalog = result.current.data?.[provider] + expect(Object.keys(catalog ?? {})).toEqual(["__proto__", "vendor/model"]) + expect(Object.getOwnPropertyDescriptor(catalog, "__proto__")?.value).toEqual({ contextWindow: 1000 }) + expect(catalog?.["vendor/model"]).toEqual({ contextWindow: 2000 }) + }) +}) diff --git a/webview-ui/src/components/ui/hooks/__tests__/useSelectedModel.spec.ts b/webview-ui/src/components/ui/hooks/__tests__/useSelectedModel.spec.ts index 4955c3c031..db0ac6d7b1 100644 --- a/webview-ui/src/components/ui/hooks/__tests__/useSelectedModel.spec.ts +++ b/webview-ui/src/components/ui/hooks/__tests__/useSelectedModel.spec.ts @@ -1855,7 +1855,9 @@ describe("useSelectedModel", () => { describe("friendli provider", () => { beforeEach(() => { - mockUseRouterModels.mockReturnValue(createRouterModelsResult({ openrouter: {}, requesty: {}, litellm: {} })) + mockUseRouterModels.mockReturnValue( + createRouterModelsResult({ openrouter: {}, requesty: {}, litellm: {}, friendli: {} }), + ) mockUseOpenRouterModelProviders.mockReturnValue({ data: {}, @@ -1890,6 +1892,82 @@ describe("useSelectedModel", () => { expect(result.current.id).toBe("zai-org/GLM-5.1") expect(result.current.info).toEqual(friendliModels["zai-org/GLM-5.1"]) }) + + it("should use static model info before router models load", () => { + mockUseRouterModels.mockReturnValue(createRouterModelsResult(undefined, { isLoading: true })) + const apiConfiguration: ProviderSettings = { apiProvider: providerIdentifiers.friendli } + + const { result } = renderHook(() => useSelectedModel(apiConfiguration), { wrapper: createWrapper() }) + + expect(result.current.id).toBe(friendliDefaultModelId) + expect(result.current.info).toEqual(friendliModels[friendliDefaultModelId]) + }) + + it("should prefer fetched router model info", () => { + const info = { maxTokens: 8192, contextWindow: 100_000, supportsPromptCache: false } + mockUseRouterModels.mockReturnValue(createRouterModelsResult({ friendli: { "vendor/new-model": info } })) + const apiConfiguration: ProviderSettings = { + apiProvider: providerIdentifiers.friendli, + apiModelId: "vendor/new-model", + } + + const { result } = renderHook(() => useSelectedModel(apiConfiguration), { wrapper: createWrapper() }) + + expect(result.current.id).toBe("vendor/new-model") + expect(result.current.info).toEqual(info) + }) + + it("should use the default model for an empty model ID", () => { + const { result } = renderHook( + () => useSelectedModel({ apiProvider: providerIdentifiers.friendli, apiModelId: "" }), + { wrapper: createWrapper() }, + ) + + expect(result.current.id).toBe(friendliDefaultModelId) + expect(result.current.info).toEqual(friendliModels[friendliDefaultModelId]) + }) + + it("should use fetched metadata for the default model when the model ID is empty", () => { + const info = { maxTokens: 8192, contextWindow: 123_456, supportsPromptCache: false } + mockUseRouterModels.mockReturnValue( + createRouterModelsResult({ friendli: { [friendliDefaultModelId]: info } }), + ) + + const { result } = renderHook( + () => useSelectedModel({ apiProvider: providerIdentifiers.friendli, apiModelId: "" }), + { wrapper: createWrapper() }, + ) + + expect(result.current.id).toBe(friendliDefaultModelId) + expect(result.current.info).toBe(info) + }) + + it.each(["toString", "constructor", "__proto__", "vendor/unknown"])( + "should use sane defaults, not an inherited member, for %s missing from both catalogs", + (apiModelId) => { + const { result } = renderHook( + () => useSelectedModel({ apiProvider: providerIdentifiers.friendli, apiModelId }), + { wrapper: createWrapper() }, + ) + + expect(result.current.id).toBe(apiModelId) + expect(result.current.info).toEqual(openAiModelInfoSaneDefaults) + }, + ) + + it("should resolve a fetched model whose upstream ID is __proto__", () => { + const info = { contextWindow: 1000, supportsPromptCache: false } + const fetched: Record = Object.create(null) + fetched["__proto__"] = info + mockUseRouterModels.mockReturnValue(createRouterModelsResult({ friendli: fetched })) + + const { result } = renderHook( + () => useSelectedModel({ apiProvider: providerIdentifiers.friendli, apiModelId: "__proto__" }), + { wrapper: createWrapper() }, + ) + + expect(result.current.info).toBe(info) + }) }) describe("OpenAI-compatible static catalog providers (fireworks, baseten, sambanova)", () => { diff --git a/webview-ui/src/components/ui/hooks/useRouterModels.ts b/webview-ui/src/components/ui/hooks/useRouterModels.ts index 7557562583..bcac658b04 100644 --- a/webview-ui/src/components/ui/hooks/useRouterModels.ts +++ b/webview-ui/src/components/ui/hooks/useRouterModels.ts @@ -1,4 +1,4 @@ -import { useQuery } from "@tanstack/react-query" +import { replaceEqualDeep, useQuery } from "@tanstack/react-query" import { allRouterModelsProvider, @@ -77,5 +77,16 @@ export const useRouterModels = (opts: UseRouterModelsOptions = {}) => { queryKey: [RouterModelsMessageType.routerModels, provider || allRouterModelsProvider], queryFn: ({ signal }) => fetchRouterModels(provider, signal), enabled: opts.enabled !== false, + // Structural sharing copies into `{}`, where an own "__proto__" model ID would swap the + // copy's prototype instead of becoming a key, so skip sharing for such catalogs. + structuralSharing: (oldData, newData) => + hasOwnProtoModelId(newData) ? newData : replaceEqualDeep(oldData, newData), }) } + +const hasOwnProtoModelId = (data: unknown): boolean => + typeof data === "object" && + data !== null && + Object.values(data).some( + (models) => typeof models === "object" && models !== null && Object.hasOwn(models, "__proto__"), + ) diff --git a/webview-ui/src/components/ui/hooks/useSelectedModel.ts b/webview-ui/src/components/ui/hooks/useSelectedModel.ts index 9eca0cc0f3..939dbb14ab 100644 --- a/webview-ui/src/components/ui/hooks/useSelectedModel.ts +++ b/webview-ui/src/components/ui/hooks/useSelectedModel.ts @@ -133,12 +133,15 @@ export const useSelectedModel = (apiConfiguration: ProviderSettings) => { // before the selection is resolved. const hasValidRouterData = needRouterModels && dynamicProvider - ? dynamicProvider === providerIdentifiers.litellm || dynamicProvider === providerIdentifiers.ioIntelligence - ? !routerModels.isLoading - : routerModels.data && - routerModels.data[dynamicProvider] !== undefined && - typeof routerModels.data[dynamicProvider] === "object" && - !routerModels.isLoading + ? dynamicProvider === providerIdentifiers.friendli + ? true // static friendliModels cover the window before /v1/models data arrives + : dynamicProvider === providerIdentifiers.litellm || + dynamicProvider === providerIdentifiers.ioIntelligence + ? !routerModels.isLoading + : routerModels.data && + routerModels.data[dynamicProvider] !== undefined && + typeof routerModels.data[dynamicProvider] === "object" && + !routerModels.isLoading : true const isReady = @@ -512,8 +515,14 @@ function getSelectedModel({ case providerIdentifiers.fireworks: return getStaticCatalogModel(apiConfiguration.apiModelId, fireworksModels, defaultModelId) case providerIdentifiers.friendli: { - const id = apiConfiguration.apiModelId ?? defaultModelId - const info = friendliModels[id as keyof typeof friendliModels] + const id = apiConfiguration.apiModelId || defaultModelId + // Mirror the backend getModel(): own-property reads only (an ID such as "toString" must + // not resolve to an inherited member), and sane defaults when neither catalog has the ID. + const fetched = routerModels.friendli + const info = + (fetched && Object.hasOwn(fetched, id) ? fetched[id] : undefined) ?? + (Object.hasOwn(friendliModels, id) ? friendliModels[id as keyof typeof friendliModels] : undefined) ?? + openAiModelInfoSaneDefaults return { id, info } } case providerIdentifiers.poe: { diff --git a/webview-ui/src/utils/__tests__/validate.spec.ts b/webview-ui/src/utils/__tests__/validate.spec.ts index bf25faf7f2..56930b4966 100644 --- a/webview-ui/src/utils/__tests__/validate.spec.ts +++ b/webview-ui/src/utils/__tests__/validate.spec.ts @@ -63,6 +63,7 @@ describe("Model Validation Functions", () => { nanogpt: {}, "zoo-gateway": {}, "kimi-code": {}, + friendli: {}, moonshot: {}, }