diff --git a/packages/types/src/vscode-extension-host.ts b/packages/types/src/vscode-extension-host.ts index 8a006c0639..ed075c90ca 100644 --- a/packages/types/src/vscode-extension-host.ts +++ b/packages/types/src/vscode-extension-host.ts @@ -23,6 +23,35 @@ import type { RuleMetadata } from "./rules.js" import type { TelemetrySetting } from "./telemetry.js" import type { WorktreeIncludeStatus } from "./worktree.js" +/** Provider lookups the extension host runs for the webview, keeping provider origins out of the CSP. */ +export type ProviderInfoKind = keyof ProviderInfoResultMap + +export interface ProviderInfoParams { + apiKey?: string + baseUrl?: string + modelId?: string +} + +/** Result of each lookup. The host and the webview share it, so the kind fixes the result type. */ +export interface ProviderInfoResultMap { + openRouterKeyInfo: { + label: string + usage: number + is_free_tier: boolean + is_provisioning_key: boolean + rate_limit: { requests: number; interval: string } + limit: number | null + } | null + openRouterModelProviders: Record + requestyKeyInfo: { + name: string + monthly_limit: string + monthly_spend: string + org_balance: string + config: { aliases?: Record } + } | null +} + /** * ExtensionMessage * Extension -> Webview | CLI @@ -66,6 +95,7 @@ export interface ExtensionMessage { | "ttsStart" | "ttsStop" | "fileSearchResults" + | "providerInfo" | "toggleApiConfigPin" | "acceptInput" | "setHistoryPreviewCollapsed" @@ -549,6 +579,7 @@ export interface WebviewMessage { | "resetState" | typeof RouterModelsMessageType.flushRouterModels | typeof RouterModelsMessageType.requestRouterModels + | "requestProviderInfo" | typeof OpenAiModelsMessageType.requestOpenAiModels | typeof OllamaModelsMessageType.requestOllamaModels | typeof LmStudioModelsMessageType.requestLmStudioModels diff --git a/src/api/providers/fetchers/__tests__/openrouter.spec.ts b/src/api/providers/fetchers/__tests__/openrouter.spec.ts index 2eca76e401..ba629b7302 100644 --- a/src/api/providers/fetchers/__tests__/openrouter.spec.ts +++ b/src/api/providers/fetchers/__tests__/openrouter.spec.ts @@ -1,5 +1,8 @@ // npx vitest run api/providers/fetchers/__tests__/openrouter.spec.ts +import http from "http" +import nock from "nock" +import type { AddressInfo } from "net" import * as path from "path" import { back as nockBack } from "nock" @@ -332,11 +335,125 @@ describe("OpenRouter API", () => { "https://openrouter.ai/api/v1/models/anthropic/claude-sonnet-4/endpoints", { timeout: MODEL_CATALOG_FETCH_TIMEOUT_MS, + signal: expect.any(AbortSignal), + maxContentLength: 1024 * 1024, }, ) getSpy.mockRestore() }) + + it("returns no endpoints for an image-output model", async () => { + const axios = await import("axios") + const getSpy = vi.spyOn(axios.default, "get").mockResolvedValue({ + data: { + data: { + id: "vendor/image-model", + name: "Image Model", + architecture: { output_modalities: ["image"] }, + endpoints: [{ provider_name: "P", context_length: 1000 }], + }, + }, + }) + + await expect(getOpenRouterModelEndpoints("vendor/image-model")).resolves.toEqual({}) + + getSpy.mockRestore() + }) + + it.each([ + ["a non-numeric", "bad", false], + ["a numeric", 4096, true], + ])("with untrustedBaseUrl handles %s context_length", async (_name, contextLength, accepted) => { + const axios = await import("axios") + const endpoint = { name: "P", provider_name: "P", context_length: contextLength } + const body = { data: { id: "m", name: "M", endpoints: [endpoint] } } + const getSpy = vi.spyOn(axios.default, "get").mockResolvedValue({ data: body }) + vi.spyOn(console, "error").mockImplementation(() => {}) + + const result = await getOpenRouterModelEndpoints("m", undefined, { untrustedBaseUrl: true }) + + if (accepted) { + expect(Object.keys(result)).toEqual(["P"]) + expect(result["P"]).toMatchObject({ contextWindow: 4096 }) + } else { + expect(result).toEqual({}) + // Trusted callers keep the lenient fallback. + await expect(getOpenRouterModelEndpoints("m")).resolves.toHaveProperty("P") + } + + getSpy.mockRestore() + }) + + it("bounds the total time of a body that streams slowly after the headers arrive", async () => { + const server = http.createServer((_req, res) => { + res.writeHead(200, { "Content-Type": "application/json" }) + // One byte at a time keeps the socket active, so only a total deadline can stop it. + const timer = setInterval(() => res.write(" "), 50) + res.on("close", () => clearInterval(timer)) + }) + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)) + const { port } = server.address() as AddressInfo + vi.spyOn(console, "error").mockImplementation(() => {}) + nock.enableNetConnect("127.0.0.1") + + try { + const start = Date.now() + const result = await getOpenRouterModelEndpoints( + "m", + { openRouterBaseUrl: `http://127.0.0.1:${port}` }, + { untrustedBaseUrl: true, timeoutMs: 300 }, + ) + + expect(result).toEqual({}) + expect(Date.now() - start).toBeLessThan(2000) + } finally { + nock.disableNetConnect() + server.closeAllConnections() + await new Promise((resolve) => server.close(resolve)) + } + }) + + it("uses the caller timeout when given", async () => { + const axios = await import("axios") + const getSpy = vi.spyOn(axios.default, "get").mockResolvedValue({ + data: { data: { id: "m", name: "M", endpoints: [] } }, + }) + + await getOpenRouterModelEndpoints("m", undefined, { timeoutMs: 10_000 }) + + expect(getSpy).toHaveBeenCalledWith(expect.any(String), expect.objectContaining({ timeout: 10_000 })) + + getSpy.mockRestore() + }) + + it("with untrustedBaseUrl logs no host name when the request fails", async () => { + const axios = await import("axios") + const getSpy = vi + .spyOn(axios.default, "get") + .mockRejectedValue(new Error("getaddrinfo ENOTFOUND secret.example.com")) + const errorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + + await getOpenRouterModelEndpoints( + "m", + { openRouterBaseUrl: "https://secret.example.com" }, + { untrustedBaseUrl: true }, + ) + + expect(JSON.stringify(errorSpy.mock.calls)).not.toContain("secret.example.com") + + getSpy.mockRestore() + }) + + it("returns no endpoints for a body without an endpoints list", async () => { + const axios = await import("axios") + const getSpy = vi.spyOn(axios.default, "get").mockResolvedValue({ data: { data: { id: "m" } } }) + vi.spyOn(console, "error").mockImplementation(() => {}) + + await expect(getOpenRouterModelEndpoints("m")).resolves.toEqual({}) + + getSpy.mockRestore() + }) }) describe("parseOpenRouterModel", () => { diff --git a/src/api/providers/fetchers/openrouter.ts b/src/api/providers/fetchers/openrouter.ts index 000f83643c..dda19b54c8 100644 --- a/src/api/providers/fetchers/openrouter.ts +++ b/src/api/providers/fetchers/openrouter.ts @@ -12,10 +12,12 @@ import { import type { ApiHandlerOptions } from "../../../shared/api" import { parseApiPrice } from "../../../shared/cost" -import { throwIfAborted } from "../utils/abort-signal" +import { hasUsableTimeout, mergeAbortSignalAndTimeout, throwIfAborted, toSdkTimeoutMs } from "../utils/abort-signal" import { MODEL_CATALOG_FETCH_TIMEOUT_MS } from "./modelCache" +const MAX_ENDPOINTS_BODY_BYTES = 1024 * 1024 + /** * OpenRouterBaseModel */ @@ -149,22 +151,32 @@ export async function getOpenRouterModels( export async function getOpenRouterModelEndpoints( modelId: string, options?: ApiHandlerOptions, + // untrustedBaseUrl: the base URL comes from unsaved webview settings. Reject schema-invalid + // bodies and keep the URL out of error logs. + opts?: { untrustedBaseUrl?: boolean; timeoutMs?: number }, ): Promise> { const models: Record = {} const baseURL = options?.openRouterBaseUrl || "https://openrouter.ai/api/v1" + const untrusted = opts?.untrustedBaseUrl === true + const timeoutMs = hasUsableTimeout(opts?.timeoutMs) ? opts.timeoutMs : MODEL_CATALOG_FETCH_TIMEOUT_MS try { // 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, + timeout: toSdkTimeoutMs(timeoutMs), + // Also bound the whole body read, in case the transport stops its timer after the headers. + signal: mergeAbortSignalAndTimeout(undefined, timeoutMs), + // The base URL can come from user settings. The cap stops a huge body from filling host memory. + maxContentLength: MAX_ENDPOINTS_BODY_BYTES, }) const result = openRouterModelEndpointsResponseSchema.safeParse(response.data) const data = result.success ? result.data.data : response.data.data if (!result.success) { console.error("OpenRouter model endpoints response is invalid", result.error.format()) + if (untrusted) return models } const { id, architecture, endpoints } = data @@ -184,8 +196,11 @@ export async function getOpenRouterModelEndpoints( }) } } catch (error) { + // The error text can carry a configured host name. console.error( - `Error fetching OpenRouter model endpoints: ${JSON.stringify(error, Object.getOwnPropertyNames(error), 2)}`, + untrusted + ? `Error fetching OpenRouter model endpoints: ${error instanceof Error ? error.name : "unknown"}` + : `Error fetching OpenRouter model endpoints: ${JSON.stringify(error, Object.getOwnPropertyNames(error), 2)}`, ) } diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index 6cec3bdbe7..333c646ff9 100644 --- a/src/core/webview/ClineProvider.ts +++ b/src/core/webview/ClineProvider.ts @@ -1846,12 +1846,6 @@ export class ClineProvider const nonce = getNonce() - // Get the OpenRouter base URL from configuration - const { apiConfiguration } = await this.getState() - const openRouterBaseUrl = apiConfiguration.openRouterBaseUrl || "https://openrouter.ai" - // Extract the domain for CSP - const openRouterDomain = openRouterBaseUrl.match(/^(https?:\/\/[^/]+)/)?.[1] || "https://openrouter.ai" - const stylesUri = getUri(webview, this.contextProxy.extensionUri, [ "webview-ui", "build", @@ -1888,7 +1882,7 @@ export class ClineProvider `img-src ${webview.cspSource} https://storage.googleapis.com https://img.clerk.com https://avatars.githubusercontent.com https://lh3.googleusercontent.com data:`, `media-src ${webview.cspSource}`, `script-src 'unsafe-eval' ${webview.cspSource} https://*.posthog.com http://${localServerUrl} http://0.0.0.0:${localPort} 'nonce-${nonce}'`, - `connect-src ${webview.cspSource} ${openRouterDomain} https://* https://*.posthog.com ws://${localServerUrl} ws://0.0.0.0:${localPort} http://${localServerUrl} http://0.0.0.0:${localPort}`, + `connect-src ${webview.cspSource} https://* https://*.posthog.com ws://${localServerUrl} ws://0.0.0.0:${localPort} http://${localServerUrl} http://0.0.0.0:${localPort}`, ] return /*html*/ ` @@ -1962,12 +1956,6 @@ export class ClineProvider */ const nonce = getNonce() - // Get the OpenRouter base URL from configuration - const { apiConfiguration } = await this.getState() - const openRouterBaseUrl = apiConfiguration.openRouterBaseUrl || "https://openrouter.ai" - // Extract the domain for CSP - const openRouterDomain = openRouterBaseUrl.match(/^(https?:\/\/[^/]+)/)?.[1] || "https://openrouter.ai" - // Tip: Install the es6-string-html VS Code extension to enable code highlighting below return /*html*/ ` @@ -1976,7 +1964,7 @@ export class ClineProvider - + ", + requestyBaseUrl: "https://requesty.example", + }, + }) + vi.mocked(axios.get).mockResolvedValue({ data: "" }) + + const html = await provider[build](mockWebviewView.webview) + const csp = html.match( + /http-equiv="Content-Security-Policy" content="([\s\S]*?)">\s* { const developmentContext = { ...mockContext, extensionMode: vscode.ExtensionMode.Development } provider = new ClineProvider( @@ -3075,7 +3110,7 @@ describe("ClineProvider", () => { // Verify Content Security Policy contains the necessary PostHog domains expect(mockWebviewView.webview.html).toContain( - "connect-src vscode-webview://test-csp-source https://openrouter.ai https://api.requesty.ai https://us.i.posthog.com", + "connect-src vscode-webview://test-csp-source https://us.i.posthog.com", ) // Extract the script-src directive section and verify required security elements diff --git a/src/core/webview/__tests__/providerInfoHandler.spec.ts b/src/core/webview/__tests__/providerInfoHandler.spec.ts new file mode 100644 index 0000000000..e8c4d66f50 --- /dev/null +++ b/src/core/webview/__tests__/providerInfoHandler.spec.ts @@ -0,0 +1,233 @@ +import type { WebviewMessage } from "@roo-code/types" + +import type { ClineProvider } from "../ClineProvider" +import { getOpenRouterModelEndpoints } from "../../../api/providers/fetchers/openrouter" +import { handleRequestProviderInfo } from "../providerInfoHandler" + +vi.mock("../../../api/providers/fetchers/openrouter", () => ({ getOpenRouterModelEndpoints: vi.fn() })) + +const KEY_INFO = { + label: "sk-or-v1-...abc", + usage: 12.5, + is_free_tier: false, + is_provisioning_key: false, + rate_limit: { requests: 60, interval: "1m" }, + limit: 100, +} + +const REQUESTY_KEY_INFO = { + name: "key", + monthly_limit: "100", + monthly_spend: "5", + org_balance: "95", + config: {}, +} + +const jsonResponse = (body: unknown, status = 200) => new Response(JSON.stringify(body), { status }) + +describe("handleRequestProviderInfo", () => { + const postMessageToWebview = vi.fn() + const provider = { postMessageToWebview, log: vi.fn() } as unknown as ClineProvider // Only these two members are used. + let fetchSpy: ReturnType + + const send = (values: WebviewMessage["values"]) => + handleRequestProviderInfo(provider, { type: "requestProviderInfo", requestId: "req-1", values }) + + beforeEach(() => { + postMessageToWebview.mockReset() + vi.mocked(getOpenRouterModelEndpoints).mockReset() + fetchSpy = vi.spyOn(globalThis, "fetch") + vi.spyOn(console, "error").mockImplementation(() => {}) + }) + + afterEach(() => { + vi.restoreAllMocks() + }) + + it("fetches OpenRouter key info from the default endpoint with a bearer token", async () => { + fetchSpy.mockResolvedValue(jsonResponse({ data: KEY_INFO })) + + await send({ kind: "openRouterKeyInfo", apiKey: "sk-or-v1-test" }) + + expect(fetchSpy.mock.calls[0]![0]).toBe("https://openrouter.ai/api/v1/key") + expect(fetchSpy.mock.calls[0]![1]).toMatchObject({ headers: { Authorization: "Bearer sk-or-v1-test" } }) + expect(postMessageToWebview).toHaveBeenCalledWith({ + type: "providerInfo", + requestId: "req-1", + values: { data: KEY_INFO }, + }) + }) + + it("uses a caller-supplied OpenRouter base URL", async () => { + fetchSpy.mockResolvedValue(jsonResponse({ data: KEY_INFO })) + + await send({ kind: "openRouterKeyInfo", apiKey: "k", baseUrl: "https://proxy.example.com/api/v1" }) + + expect(fetchSpy.mock.calls[0]![0]).toBe("https://proxy.example.com/api/v1/key") + }) + + it.each(["openRouterKeyInfo", "requestyKeyInfo"])( + "returns null without a request when %s has no apiKey", + async (kind) => { + await send({ kind }) + + expect(fetchSpy).not.toHaveBeenCalled() + expect(postMessageToWebview).toHaveBeenCalledWith({ + type: "providerInfo", + requestId: "req-1", + values: { data: null }, + }) + }, + ) + + it("fetches Requesty key info from the api service URL", async () => { + fetchSpy.mockResolvedValue(jsonResponse(REQUESTY_KEY_INFO)) + + await send({ kind: "requestyKeyInfo", apiKey: "rk" }) + + expect(fetchSpy.mock.calls[0]![0]).toBe("https://api.requesty.ai/x/apikey") + expect(postMessageToWebview).toHaveBeenCalledWith({ + type: "providerInfo", + requestId: "req-1", + values: { data: REQUESTY_KEY_INFO }, + }) + }) + + it("fetches Requesty key info from a custom base URL with auth headers", async () => { + fetchSpy.mockResolvedValue(jsonResponse(REQUESTY_KEY_INFO)) + + await send({ kind: "requestyKeyInfo", apiKey: "rk", baseUrl: "https://router.custom.example/v1" }) + + const [url, init] = fetchSpy.mock.calls[0]! + expect(url).not.toBe("https://api.requesty.ai/x/apikey") + expect(String(url)).toMatch(/^https:\/\/[^/]*custom\.example\/.*x\/apikey$/) + expect(init).toMatchObject({ + headers: { Authorization: "Bearer rk", "Content-Type": "application/json" }, + }) + }) + + it("returns no providers and sends no request without a modelId", async () => { + await send({ kind: "openRouterModelProviders" }) + + expect(fetchSpy).not.toHaveBeenCalled() + expect(getOpenRouterModelEndpoints).not.toHaveBeenCalled() + expect(postMessageToWebview.mock.calls[0]![0].values).toEqual({ data: {} }) + }) + + it("passes an abort signal with the request", async () => { + fetchSpy.mockResolvedValue(jsonResponse({ data: KEY_INFO })) + + await send({ kind: "openRouterKeyInfo", apiKey: "k" }) + + expect(fetchSpy.mock.calls[0]![1]!.signal).toBeInstanceOf(AbortSignal) + }) + + it("delegates OpenRouter model providers to the shared endpoint fetcher", async () => { + const providers = { "provider-a": { contextWindow: 128000, supportsPromptCache: true } } + vi.mocked(getOpenRouterModelEndpoints).mockResolvedValue(providers) + + await send({ kind: "openRouterModelProviders", modelId: "openai/gpt-4o", baseUrl: "https://proxy.test/v1" }) + + expect(getOpenRouterModelEndpoints).toHaveBeenCalledWith( + "openai/gpt-4o", + { openRouterBaseUrl: "https://proxy.test/v1" }, + { untrustedBaseUrl: true, timeoutMs: 10_000 }, + ) + expect(postMessageToWebview).toHaveBeenCalledWith({ + type: "providerInfo", + requestId: "req-1", + values: { data: { "provider-a": { ...providers["provider-a"], label: "provider-a" } } }, + }) + }) + + it.each([ + ["openRouterKeyInfo", { apiKey: "k" }, { data: { label: 1 } }, null], + ["requestyKeyInfo", { apiKey: "k" }, { name: 1 }, null], + ])("returns the fallback and logs when %s gets a schema-invalid body", async (kind, params, body, fallback) => { + fetchSpy.mockResolvedValue(jsonResponse(body)) + + await send({ kind, ...params }) + + expect(postMessageToWebview.mock.calls[0]![0].values).toEqual({ data: fallback }) + expect(console.error).toHaveBeenCalled() + }) + + it("rejects a response body above the size cap", async () => { + const cancel = vi.fn() + const oversized = new ReadableStream({ + pull: (controller) => controller.enqueue(new Uint8Array(512 * 1024)), + cancel, + }) + fetchSpy.mockResolvedValue(new Response(oversized, { status: 200 })) + + await send({ kind: "openRouterKeyInfo", apiKey: "k" }) + + expect(cancel).toHaveBeenCalled() + expect(postMessageToWebview).toHaveBeenCalledWith({ + type: "providerInfo", + requestId: "req-1", + error: "Provider info request failed", + }) + }) + + it("accepts a body of exactly the size cap and rejects one byte more", async () => { + const MAX = 1024 * 1024 + const padded = (bytes: number) => { + const head = '{"data":null,"pad":"' + const tail = '"}' + return head + "a".repeat(bytes - head.length - tail.length) + tail + } + + fetchSpy.mockResolvedValueOnce(new Response(padded(MAX), { status: 200 })) + await send({ kind: "openRouterKeyInfo", apiKey: "k" }) + expect(postMessageToWebview.mock.calls[0]![0].error).toBeUndefined() + + fetchSpy.mockResolvedValueOnce(new Response(padded(MAX + 1), { status: 200 })) + await send({ kind: "openRouterKeyInfo", apiKey: "k" }) + expect(postMessageToWebview.mock.calls[1]![0].error).toBe("Provider info request failed") + }) + + it.each([ + ["a failed status", () => jsonResponse({}, 401)], + ["a non-JSON 2xx body", () => new Response("", { status: 200 })], + ["a rejected fetch", () => Promise.reject(new Error("https://secret.example.com/key refused"))], + ])("posts a generic error that leaks no URL on %s", async (_name, makeResponse) => { + fetchSpy.mockImplementation(async () => makeResponse()) + + await send({ kind: "openRouterKeyInfo", apiKey: "k", baseUrl: "https://secret.example.com" }) + + expect(postMessageToWebview).toHaveBeenCalledWith({ + type: "providerInfo", + requestId: "req-1", + error: "Provider info request failed", + }) + const logged = JSON.stringify([vi.mocked(provider.log).mock.calls, vi.mocked(console.error).mock.calls]) + expect(logged).not.toContain("secret.example.com") + }) + + it.each([ + ["apiKey", 5], + ["baseUrl", {}], + ["modelId", ["m"]], + ])("rejects a non-string %s without a request", async (field, value) => { + await send({ kind: "openRouterKeyInfo", [field]: value }) + + expect(fetchSpy).not.toHaveBeenCalled() + expect(postMessageToWebview).toHaveBeenCalledWith({ + type: "providerInfo", + requestId: "req-1", + error: "Invalid provider info request", + }) + }) + + it.each([undefined, "unknown", "__proto__", "constructor"])("rejects unknown kind %s", async (kind) => { + await send({ kind }) + + expect(fetchSpy).not.toHaveBeenCalled() + expect(postMessageToWebview).toHaveBeenCalledWith({ + type: "providerInfo", + requestId: "req-1", + error: "Invalid provider info request", + }) + }) +}) diff --git a/src/core/webview/__tests__/webviewMessageHandler.providerInfo.spec.ts b/src/core/webview/__tests__/webviewMessageHandler.providerInfo.spec.ts new file mode 100644 index 0000000000..5d48c8169a --- /dev/null +++ b/src/core/webview/__tests__/webviewMessageHandler.providerInfo.spec.ts @@ -0,0 +1,51 @@ +// npx vitest core/webview/__tests__/webviewMessageHandler.providerInfo.spec.ts + +import type { ClineProvider } from "../ClineProvider" +import { webviewMessageHandler } from "../webviewMessageHandler" + +vi.mock("../../../api/providers/fetchers/modelCache") + +const KEY_INFO = { name: "key", monthly_limit: "100", monthly_spend: "5", org_balance: "95", config: {} } + +describe("webviewMessageHandler - requestProviderInfo", () => { + // The handler uses only these two members for this message. + const postMessageToWebview = vi.fn() + const provider = { postMessageToWebview, log: vi.fn() } as unknown as ClineProvider + + beforeEach(() => { + postMessageToWebview.mockReset() + }) + + it("routes the message to the provider info handler and posts a providerInfo reply", async () => { + await webviewMessageHandler(provider, { + type: "requestProviderInfo", + requestId: "req-9", + values: { kind: "nope" }, + }) + + expect(postMessageToWebview).toHaveBeenCalledWith({ + type: "providerInfo", + requestId: "req-9", + error: "Invalid provider info request", + }) + }) + + it("routes a valid request to the fetcher and posts the success reply", async () => { + const fetchSpy = vi + .spyOn(globalThis, "fetch") + .mockResolvedValue(new Response(JSON.stringify(KEY_INFO), { status: 200 })) + + await webviewMessageHandler(provider, { + type: "requestProviderInfo", + requestId: "req-10", + values: { kind: "requestyKeyInfo", apiKey: "rk" }, + }) + + expect(postMessageToWebview).toHaveBeenCalledWith({ + type: "providerInfo", + requestId: "req-10", + values: { data: KEY_INFO }, + }) + fetchSpy.mockRestore() + }) +}) diff --git a/src/core/webview/providerInfoHandler.ts b/src/core/webview/providerInfoHandler.ts new file mode 100644 index 0000000000..aac43cd2cc --- /dev/null +++ b/src/core/webview/providerInfoHandler.ts @@ -0,0 +1,168 @@ +import { z } from "zod" + +import type { ProviderInfoKind, ProviderInfoParams, ProviderInfoResultMap, WebviewMessage } from "@roo-code/types" + +import type { ClineProvider } from "./ClineProvider" +import { getOpenRouterModelEndpoints } from "../../api/providers/fetchers/openrouter" +import { toRequestyServiceUrl } from "../../shared/utils/requesty" + +const FETCH_TIMEOUT_MS = 10_000 +// Far above any real provider response. The cap stops a configured endpoint from filling host memory. +const MAX_BODY_BYTES = 1024 * 1024 +const OPENROUTER_DEFAULT_BASE_URL = "https://openrouter.ai/api/v1" + +const openRouterKeyInfoSchema = z.object({ + data: z.object({ + label: z.string(), + usage: z.number(), + is_free_tier: z.boolean(), + is_provisioning_key: z.boolean(), + rate_limit: z.object({ + requests: z.number(), + interval: z.string(), + }), + limit: z.number().nullable(), + }), +}) + +const requestyKeyInfoSchema = z.object({ + name: z.string(), + monthly_limit: z.string(), + monthly_spend: z.string(), + org_balance: z.string(), + config: z.object({ + aliases: z.record(z.string(), z.any()).optional(), + }), +}) + +async function readBoundedText(response: Response): Promise { + // Count bytes as they arrive. Content-Length can be missing or wrong. + const reader = response.body?.getReader() + const chunks: Uint8Array[] = [] + let total = 0 + + while (reader) { + const { done, value } = await reader.read() + if (done) break + + total += value.byteLength + if (total > MAX_BODY_BYTES) { + await reader.cancel() + throw new Error("Response body too large") + } + chunks.push(value) + } + + return Buffer.concat(chunks).toString("utf8") +} + +async function getJson(url: string, headers?: Record): Promise { + const response = await fetch(url, { + method: "GET", + headers, + signal: AbortSignal.timeout(FETCH_TIMEOUT_MS), + }) + + if (!response.ok) { + throw new Error(`Request failed with status code ${response.status}`) + } + + return JSON.parse(await readBoundedText(response)) +} + +async function getOpenRouterKeyInfo({ apiKey, baseUrl }: ProviderInfoParams) { + if (!apiKey) return null + + const body = await getJson(`${baseUrl || OPENROUTER_DEFAULT_BASE_URL}/key`, { + Authorization: `Bearer ${apiKey}`, + }) + const result = openRouterKeyInfoSchema.safeParse(body) + + if (!result.success) { + console.error("OpenRouter API key info validation failed:", result.error) + return null + } + + return result.data.data +} + +// Reuse the shared fetcher. It returns {} on failure, which the webview treats as no providers. +// The key is the provider label, so the webview select can show it. +async function getOpenRouterModelProviders({ modelId, baseUrl }: ProviderInfoParams) { + if (!modelId) return {} + + const endpoints = await getOpenRouterModelEndpoints( + modelId, + { openRouterBaseUrl: baseUrl }, + { untrustedBaseUrl: true, timeoutMs: FETCH_TIMEOUT_MS }, + ) + + return Object.fromEntries(Object.entries(endpoints).map(([label, info]) => [label, { ...info, label }])) +} + +async function getRequestyKeyInfo({ apiKey, baseUrl }: ProviderInfoParams) { + if (!apiKey) return null + + const apiKeyUrl = new URL("x/apikey", toRequestyServiceUrl(baseUrl, "api")) + const body = await getJson(apiKeyUrl.toString(), { + Authorization: `Bearer ${apiKey}`, + "Content-Type": "application/json", + }) + const result = requestyKeyInfoSchema.safeParse(body) + + if (!result.success) { + console.error("Requesty API key info validation failed:", result.error) + return null + } + + return result.data +} + +const fetchers: { [K in ProviderInfoKind]: (params: ProviderInfoParams) => Promise } = { + openRouterKeyInfo: getOpenRouterKeyInfo, + openRouterModelProviders: getOpenRouterModelProviders, + requestyKeyInfo: getRequestyKeyInfo, +} + +const isProviderInfoKind = (kind: unknown): kind is ProviderInfoKind => + typeof kind === "string" && Object.hasOwn(fetchers, kind) + +// Non-string params fail the parse, so a bad value never reaches a fetcher. +const providerInfoParamsSchema = z.object({ + apiKey: z.string().optional(), + baseUrl: z.string().optional(), + modelId: z.string().optional(), +}) + +/** + * Fetch provider key and model info in the extension host so the webview CSP + * needs no provider origins. The webview sends unsaved settings values, so the + * handler reads credentials from the message, not from saved config. + */ +export async function handleRequestProviderInfo(provider: ClineProvider, message: WebviewMessage): Promise { + const { kind, ...rawParams } = message.values ?? {} + const parsed = providerInfoParamsSchema.safeParse(rawParams) + + if (!isProviderInfoKind(kind) || !parsed.success) { + await provider.postMessageToWebview({ + type: "providerInfo", + requestId: message.requestId, + error: "Invalid provider info request", + }) + return + } + + try { + const data = await fetchers[kind](parsed.data) + await provider.postMessageToWebview({ type: "providerInfo", requestId: message.requestId, values: { data } }) + } catch (error) { + // Log the kind only. The error text can carry a configured URL. + provider.log(`Failed to fetch provider info (${kind})`) + console.error(`Error fetching provider info (${kind}):`, error instanceof Error ? error.name : "unknown") + await provider.postMessageToWebview({ + type: "providerInfo", + requestId: message.requestId, + error: "Provider info request failed", + }) + } +} diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index c5abbc0a12..cf9d209191 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -31,6 +31,7 @@ import { VsCodeLmModelsMessageType, isTelemetryOptedIn, } from "@roo-code/types" +import { handleRequestProviderInfo } from "./providerInfoHandler" import { customToolRegistry } from "@roo-code/core" import { CloudService } from "@roo-code/cloud" import { TelemetryService } from "@roo-code/telemetry" @@ -1161,6 +1162,9 @@ export const webviewMessageHandler = async ( case "resetState": await provider.resetState() break + case "requestProviderInfo": + await handleRequestProviderInfo(provider, message) + break case RouterModelsMessageType.flushRouterModels: { const { apiConfiguration } = await provider.getState() const routerNameFlush: RouterName = toRouterName(message.text) diff --git a/webview-ui/src/components/ui/hooks/__tests__/providerInfoHooks.spec.tsx b/webview-ui/src/components/ui/hooks/__tests__/providerInfoHooks.spec.tsx new file mode 100644 index 0000000000..d50f08179e --- /dev/null +++ b/webview-ui/src/components/ui/hooks/__tests__/providerInfoHooks.spec.tsx @@ -0,0 +1,216 @@ +import React from "react" +import type { MockInstance } from "vitest" +import { act, renderHook, waitFor } from "@testing-library/react" +import { QueryClientProvider, type UseQueryResult } from "@tanstack/react-query" + +import type { WebviewMessage } from "@roo-code/types" + +import { vscode } from "@src/utils/vscode" +import { createTestQueryClient } from "@src/utils/test-utils" + +import { requestProviderInfo } from "../requestProviderInfo" +import { useOpenRouterKeyInfo } from "../useOpenRouterKeyInfo" +import { useOpenRouterModelProviders } from "../useOpenRouterModelProviders" +import { useRequestyKeyInfo } from "../useRequestyKeyInfo" + +vi.mock("@src/utils/vscode", () => ({ vscode: { postMessage: vi.fn() } })) + +const wrapper = ({ children }: React.PropsWithChildren) => ( + {children} +) + +beforeEach(() => { + vi.mocked(vscode.postMessage).mockReset() +}) + +type Reply = { data?: unknown; error?: string } + +// Answers each requestProviderInfo message like the extension host and records the messages. +const stubHost = (reply: (values: Record) => Reply | null) => { + const sent: WebviewMessage[] = [] + + vi.mocked(vscode.postMessage).mockImplementation((message: WebviewMessage) => { + sent.push(message) + const answer = reply(message.values ?? {}) + if (!answer) return + window.dispatchEvent( + new MessageEvent("message", { + data: { + type: "providerInfo", + requestId: message.requestId, + ...(answer.error ? { error: answer.error } : { values: { data: answer.data } }), + }, + }), + ) + }) + + return { sent } +} + +describe("requestProviderInfo", () => { + afterEach(() => { + vi.restoreAllMocks() + vi.useRealTimers() + }) + + // Each settlement path must remove the exact listener it added and clear its timer. + const expectCleanedUp = (add: MockInstance, remove: MockInstance) => { + const added = add.mock.calls.filter(([type]) => type === "message").map(([, listener]) => listener) + const removed = remove.mock.calls.filter(([type]) => type === "message").map(([, listener]) => listener) + expect(added).toHaveLength(1) + expect(removed).toEqual(added) + expect(vi.getTimerCount()).toBe(0) + } + + const spyOnListeners = () => { + vi.useFakeTimers() + return [vi.spyOn(window, "addEventListener"), vi.spyOn(window, "removeEventListener")] as const + } + + it("posts the kind and params with a request id and resolves with the reply data", async () => { + const [add, remove] = spyOnListeners() + const { sent } = stubHost(() => ({ data: { ok: true } })) + + await expect( + requestProviderInfo("openRouterKeyInfo", { apiKey: "k", baseUrl: "https://x.test" }), + ).resolves.toEqual({ + ok: true, + }) + + expect(sent).toEqual([ + { + type: "requestProviderInfo", + requestId: expect.any(String), + values: { kind: "openRouterKeyInfo", apiKey: "k", baseUrl: "https://x.test" }, + }, + ]) + expectCleanedUp(add, remove) + }) + + it("gives each request its own id and settles out-of-order replies by id", async () => { + const { sent } = stubHost(() => null) + + const first = requestProviderInfo("requestyKeyInfo", { apiKey: "a" }) + const second = requestProviderInfo("requestyKeyInfo", { apiKey: "b" }) + const [idA, idB] = sent.map((message) => message.requestId) + expect(idA).not.toBe(idB) + + for (const [requestId, data] of [ + [idB, { name: "second" }], + [idA, { name: "first" }], + ]) { + window.dispatchEvent( + new MessageEvent("message", { data: { type: "providerInfo", requestId, values: { data } } }), + ) + } + + await expect(first).resolves.toMatchObject({ name: "first" }) + await expect(second).resolves.toMatchObject({ name: "second" }) + }) + + it("rejects with the host error message", async () => { + const [add, remove] = spyOnListeners() + stubHost(() => ({ error: "Provider info request failed" })) + + await expect(requestProviderInfo("requestyKeyInfo", { apiKey: "k" })).rejects.toThrow( + "Provider info request failed", + ) + + expectCleanedUp(add, remove) + }) + + it("ignores replies for other request ids and rejects on timeout", async () => { + const [add, remove] = spyOnListeners() + stubHost(() => null) + + const promise = requestProviderInfo("requestyKeyInfo", { apiKey: "k" }) + const settled = expect(promise).rejects.toThrow("Provider info request timed out") + + window.dispatchEvent( + new MessageEvent("message", { + data: { type: "providerInfo", requestId: "other", values: { data: { stolen: true } } }, + }), + ) + expect(remove).not.toHaveBeenCalledWith("message", expect.anything()) + await act(() => vi.advanceTimersByTimeAsync(20_000)) + + await settled + expectCleanedUp(add, remove) + }) +}) + +describe("provider info hooks", () => { + afterEach(() => { + vi.restoreAllMocks() + }) + + it.each<[string, () => UseQueryResult]>([ + ["useOpenRouterKeyInfo", () => useOpenRouterKeyInfo(undefined)], + ["useRequestyKeyInfo", () => useRequestyKeyInfo(undefined, undefined)], + ])("%s sends nothing without an apiKey", async (_name, useHook) => { + stubHost(() => ({ data: {} })) + + const { result } = renderHook(useHook, { wrapper }) + + await waitFor(() => expect(result.current.fetchStatus).toBe("idle")) + expect(result.current.data).toBeUndefined() + expect(vscode.postMessage).not.toHaveBeenCalled() + }) + + it("useOpenRouterModelProviders returns an empty record without a modelId", async () => { + stubHost(() => ({ data: { unexpected: {} } })) + + const { result } = renderHook(() => useOpenRouterModelProviders(undefined), { wrapper }) + + await waitFor(() => expect(result.current.isSuccess).toBe(true)) + expect(result.current.data).toEqual({}) + expect(vscode.postMessage).not.toHaveBeenCalled() + }) + + it.each<[string, () => UseQueryResult, Record, unknown]>([ + [ + "useOpenRouterKeyInfo", + () => useOpenRouterKeyInfo("sk-or", "https://proxy.test/api/v1"), + { kind: "openRouterKeyInfo", apiKey: "sk-or", baseUrl: "https://proxy.test/api/v1" }, + { label: "abc" }, + ], + [ + "useRequestyKeyInfo", + () => useRequestyKeyInfo("https://custom.requesty.ai/v1", "rk"), + { kind: "requestyKeyInfo", apiKey: "rk", baseUrl: "https://custom.requesty.ai/v1" }, + { name: "my-key" }, + ], + [ + "useOpenRouterModelProviders", + () => useOpenRouterModelProviders("vendor/model", "https://proxy.test/api/v1"), + { kind: "openRouterModelProviders", modelId: "vendor/model", baseUrl: "https://proxy.test/api/v1" }, + { "provider-a": { label: "provider-a" } }, + ], + ])("%s asks the host and returns its data", async (_name, useHook, values, data) => { + const { sent } = stubHost(() => ({ data })) + + const { result } = renderHook(useHook, { wrapper }) + + await waitFor(() => expect(result.current.isSuccess).toBe(true)) + expect(result.current.data).toEqual(data) + expect(sent.map((message) => message.values)).toEqual([values]) + }) + + it.each<[string, () => UseQueryResult, unknown]>([ + ["useOpenRouterKeyInfo", () => useOpenRouterKeyInfo("sk-or"), null], + ["useRequestyKeyInfo", () => useRequestyKeyInfo(undefined, "rk"), null], + ["useOpenRouterModelProviders", () => useOpenRouterModelProviders("vendor/model"), {}], + ])("%s falls back and logs when the host reports an error", async (_name, useHook, fallback) => { + stubHost(() => ({ error: "Provider info request failed" })) + const consoleError = vi.spyOn(console, "error").mockImplementation(() => {}) + + const { result } = renderHook(useHook, { wrapper }) + + await waitFor(() => expect(result.current.isSuccess).toBe(true)) + expect(result.current.data).toEqual(fallback) + expect(consoleError).toHaveBeenCalledWith( + expect.stringContaining("Error fetching"), + expect.objectContaining({ message: "Provider info request failed" }), + ) + }) +}) diff --git a/webview-ui/src/components/ui/hooks/__tests__/useOpenRouterKeyInfo.spec.tsx b/webview-ui/src/components/ui/hooks/__tests__/useOpenRouterKeyInfo.spec.tsx deleted file mode 100644 index 2e7bf65c24..0000000000 --- a/webview-ui/src/components/ui/hooks/__tests__/useOpenRouterKeyInfo.spec.tsx +++ /dev/null @@ -1,123 +0,0 @@ -import { renderHook, waitFor } from "@testing-library/react" -import { QueryClient, QueryClientProvider } from "@tanstack/react-query" -import React from "react" - -import { stubTransport } from "@src/utils/test-utils" - -import { useOpenRouterKeyInfo } from "../useOpenRouterKeyInfo" - -const KEY_INFO = { - label: "sk-or-v1-...abc", - usage: 12.5, - is_free_tier: false, - is_provisioning_key: false, - rate_limit: { requests: 60, interval: "1m" }, - limit: 100, -} - -const wrapper = ({ children }: React.PropsWithChildren) => ( - - {children} - -) - -describe("useOpenRouterKeyInfo", () => { - afterEach(() => { - vi.restoreAllMocks() - }) - - it("returns undefined and issues no request when apiKey is missing", async () => { - const { calls } = stubTransport(() => ({ status: 200, body: { data: KEY_INFO } })) - const consoleError = vi.spyOn(console, "error").mockImplementation(() => {}) - - const { result } = renderHook(() => useOpenRouterKeyInfo(undefined), { wrapper }) - - await waitFor(() => expect(result.current.fetchStatus).toBe("idle")) - expect(result.current.data).toBeUndefined() - expect(calls).toHaveLength(0) - expect(consoleError).not.toHaveBeenCalled() - }) - - it("requests the default key endpoint with a bearer token and returns the parsed key info", async () => { - const { calls } = stubTransport(() => ({ status: 200, body: { data: KEY_INFO } })) - - const { result } = renderHook(() => useOpenRouterKeyInfo("sk-or-v1-test"), { wrapper }) - - await waitFor(() => expect(result.current.isSuccess).toBe(true)) - expect(result.current.data).toEqual(KEY_INFO) - expect(calls).toEqual([ - { - url: "https://openrouter.ai/api/v1/key", - method: "GET", - headers: { authorization: "Bearer sk-or-v1-test" }, - }, - ]) - }) - - it("uses a caller-supplied baseUrl verbatim", async () => { - const { calls } = stubTransport(() => ({ status: 200, body: { data: KEY_INFO } })) - - const { result } = renderHook(() => useOpenRouterKeyInfo("sk-or-v1-test", "https://proxy.example.com/api/v1"), { - wrapper, - }) - - await waitFor(() => expect(result.current.isSuccess).toBe(true)) - expect(calls[0].url).toBe("https://proxy.example.com/api/v1/key") - }) - - it("returns null and logs a validation failure when the payload does not match the schema", async () => { - stubTransport(() => ({ status: 200, body: { data: { label: "only-label" } } })) - const consoleError = vi.spyOn(console, "error").mockImplementation(() => {}) - - const { result } = renderHook(() => useOpenRouterKeyInfo("sk-or-v1-test"), { wrapper }) - - await waitFor(() => expect(result.current.isSuccess).toBe(true)) - expect(result.current.data).toBeNull() - expect(consoleError).toHaveBeenCalledWith( - "OpenRouter API key info validation failed:", - expect.objectContaining({ - issues: expect.arrayContaining([expect.objectContaining({ path: ["data", "usage"] })]), - }), - ) - }) - - it("returns null and logs a fetch failure when the response is not 2xx", async () => { - stubTransport(() => ({ status: 401, body: { error: { message: "No auth credentials found" } } })) - const consoleError = vi.spyOn(console, "error").mockImplementation(() => {}) - - const { result } = renderHook(() => useOpenRouterKeyInfo("sk-or-v1-test"), { wrapper }) - - await waitFor(() => expect(result.current.isSuccess).toBe(true)) - expect(result.current.data).toBeNull() - expect(consoleError).toHaveBeenCalledWith( - "Error fetching OpenRouter key info:", - expect.objectContaining({ message: "Request failed with status code 401" }), - ) - }) - - it("returns null and logs a fetch failure when the request rejects", async () => { - const failure = new Error("network down") - stubTransport(() => { - throw failure - }) - const consoleError = vi.spyOn(console, "error").mockImplementation(() => {}) - - const { result } = renderHook(() => useOpenRouterKeyInfo("sk-or-v1-test"), { wrapper }) - - await waitFor(() => expect(result.current.isSuccess).toBe(true)) - expect(result.current.data).toBeNull() - expect(consoleError).toHaveBeenCalledWith("Error fetching OpenRouter key info:", failure) - }) - - it("returns null and logs a fetch failure when a 2xx body is not valid JSON", async () => { - const parseFailure = new SyntaxError("Unexpected token '<', \"\" is not valid JSON") - stubTransport(() => ({ status: 200, body: undefined, jsonRejects: parseFailure })) - const consoleError = vi.spyOn(console, "error").mockImplementation(() => {}) - - const { result } = renderHook(() => useOpenRouterKeyInfo("sk-or-v1-test"), { wrapper }) - - await waitFor(() => expect(result.current.isSuccess).toBe(true)) - expect(result.current.data).toBeNull() - expect(consoleError).toHaveBeenCalledWith("Error fetching OpenRouter key info:", parseFailure) - }) -}) diff --git a/webview-ui/src/components/ui/hooks/__tests__/useOpenRouterModelProviders.spec.tsx b/webview-ui/src/components/ui/hooks/__tests__/useOpenRouterModelProviders.spec.tsx deleted file mode 100644 index f16edbd8a8..0000000000 --- a/webview-ui/src/components/ui/hooks/__tests__/useOpenRouterModelProviders.spec.tsx +++ /dev/null @@ -1,190 +0,0 @@ -import { renderHook, waitFor } from "@testing-library/react" -import { QueryClientProvider } from "@tanstack/react-query" -import React from "react" - -import { createTestQueryClient, stubTransport } from "@src/utils/test-utils" - -import { useOpenRouterModelProviders } from "../useOpenRouterModelProviders" - -const endpointsPayload = (overrides: Record = {}) => ({ - data: { - id: "vendor/model", - name: "Vendor Model", - description: "A test model", - architecture: { input_modalities: ["text", "image"], output_modalities: ["text"] }, - endpoints: [ - { - name: "Provider A", - tag: "provider-a", - context_length: 128000, - max_completion_tokens: 4096, - pricing: { - prompt: "0.000003", - completion: "0.000015", - input_cache_read: "0.0000003", - input_cache_write: "0.0000036", - }, - }, - { - name: "Provider B", - context_length: 32000, - pricing: { prompt: "0.000001", completion: "0.000002" }, - }, - ], - ...overrides, - }, -}) - -const wrapper = ({ children }: React.PropsWithChildren) => ( - {children} -) - -describe("useOpenRouterModelProviders", () => { - afterEach(() => { - vi.restoreAllMocks() - }) - - it("returns an empty record without a request when modelId is missing", async () => { - const { calls } = stubTransport(() => ({ status: 200, body: endpointsPayload() })) - - const { result } = renderHook(() => useOpenRouterModelProviders(undefined), { wrapper }) - - await waitFor(() => expect(result.current.isSuccess).toBe(true)) - expect(result.current.data).toEqual({}) - expect(calls).toHaveLength(0) - }) - - it("requests the endpoints URL and maps each endpoint to provider model info", async () => { - const { calls } = stubTransport(() => ({ status: 200, body: endpointsPayload() })) - - const { result } = renderHook(() => useOpenRouterModelProviders("vendor/model"), { wrapper }) - - await waitFor(() => expect(result.current.isSuccess).toBe(true)) - - expect(calls).toEqual([ - { url: "https://openrouter.ai/api/v1/models/vendor/model/endpoints", method: "GET", headers: {} }, - ]) - - expect(result.current.data).toEqual({ - "provider-a": { - maxTokens: 4096, - contextWindow: 128000, - supportsImages: true, - supportsPromptCache: true, - cacheReadsPrice: 0.3, - cacheWritesPrice: expect.closeTo(3.6, 10), - inputPrice: 3, - outputPrice: 15, - description: "A test model", - label: "provider-a", - }, - "Provider B": { - maxTokens: 32000, - contextWindow: 32000, - supportsImages: true, - supportsPromptCache: false, - cacheReadsPrice: undefined, - cacheWritesPrice: undefined, - inputPrice: 1, - outputPrice: 2, - description: "A test model", - label: "Provider B", - }, - }) - }) - - it.each([ - { name: "text-only architecture", architecture: { input_modalities: ["text"], output_modalities: ["text"] } }, - { name: "omitted architecture", architecture: undefined }, - ])("disables image support for $name", async ({ architecture }) => { - stubTransport(() => ({ status: 200, body: endpointsPayload({ architecture }) })) - - const { result } = renderHook(() => useOpenRouterModelProviders("vendor/model"), { wrapper }) - - await waitFor(() => expect(result.current.isSuccess).toBe(true)) - expect(result.current.data).toMatchObject({ - "provider-a": { supportsImages: false }, - "Provider B": { supportsImages: false }, - }) - }) - - it("uses a caller-supplied baseUrl", async () => { - const { calls } = stubTransport(() => ({ status: 200, body: endpointsPayload() })) - - const { result } = renderHook( - () => useOpenRouterModelProviders("vendor/model", "https://proxy.example.com/api/v1"), - { wrapper }, - ) - - await waitFor(() => expect(result.current.isSuccess).toBe(true)) - expect(calls[0].url).toBe("https://proxy.example.com/api/v1/models/vendor/model/endpoints") - }) - - it("returns no providers for image-generation models", async () => { - stubTransport(() => ({ - status: 200, - body: endpointsPayload({ architecture: { input_modalities: ["text"], output_modalities: ["image"] } }), - })) - - const { result } = renderHook(() => useOpenRouterModelProviders("vendor/model"), { wrapper }) - - await waitFor(() => expect(result.current.isSuccess).toBe(true)) - expect(result.current.data).toEqual({}) - }) - - it("logs a validation failure and returns no providers when the payload does not match the schema", async () => { - stubTransport(() => ({ status: 200, body: { data: { id: "vendor/model" } } })) - const consoleError = vi.spyOn(console, "error").mockImplementation(() => {}) - - const { result } = renderHook(() => useOpenRouterModelProviders("vendor/model"), { wrapper }) - - await waitFor(() => expect(result.current.isSuccess).toBe(true)) - expect(result.current.data).toEqual({}) - expect(consoleError).toHaveBeenCalledWith( - "OpenRouter API response validation failed:", - expect.objectContaining({ - issues: expect.arrayContaining([expect.objectContaining({ path: ["data", "endpoints"] })]), - }), - ) - }) - - it("logs a fetch failure and returns no providers when the response is not 2xx", async () => { - stubTransport(() => ({ status: 500, body: { error: "boom" } })) - const consoleError = vi.spyOn(console, "error").mockImplementation(() => {}) - - const { result } = renderHook(() => useOpenRouterModelProviders("vendor/model"), { wrapper }) - - await waitFor(() => expect(result.current.isSuccess).toBe(true)) - expect(result.current.data).toEqual({}) - expect(consoleError).toHaveBeenCalledWith( - "Error fetching OpenRouter providers:", - expect.objectContaining({ message: "Request failed with status code 500" }), - ) - }) - - it("logs a fetch failure and returns no providers when the request rejects", async () => { - const failure = new Error("network down") - stubTransport(() => { - throw failure - }) - const consoleError = vi.spyOn(console, "error").mockImplementation(() => {}) - - const { result } = renderHook(() => useOpenRouterModelProviders("vendor/model"), { wrapper }) - - await waitFor(() => expect(result.current.isSuccess).toBe(true)) - expect(result.current.data).toEqual({}) - expect(consoleError).toHaveBeenCalledWith("Error fetching OpenRouter providers:", failure) - }) - - it("logs a fetch failure and returns no providers when a 2xx body is not valid JSON", async () => { - const parseFailure = new SyntaxError("Unexpected token '<', \"\" is not valid JSON") - stubTransport(() => ({ status: 200, body: undefined, jsonRejects: parseFailure })) - const consoleError = vi.spyOn(console, "error").mockImplementation(() => {}) - - const { result } = renderHook(() => useOpenRouterModelProviders("vendor/model"), { wrapper }) - - await waitFor(() => expect(result.current.isSuccess).toBe(true)) - expect(result.current.data).toEqual({}) - expect(consoleError).toHaveBeenCalledWith("Error fetching OpenRouter providers:", parseFailure) - }) -}) diff --git a/webview-ui/src/components/ui/hooks/__tests__/useRequestyKeyInfo.spec.tsx b/webview-ui/src/components/ui/hooks/__tests__/useRequestyKeyInfo.spec.tsx deleted file mode 100644 index 22688db698..0000000000 --- a/webview-ui/src/components/ui/hooks/__tests__/useRequestyKeyInfo.spec.tsx +++ /dev/null @@ -1,122 +0,0 @@ -import { renderHook, waitFor } from "@testing-library/react" -import { QueryClient, QueryClientProvider } from "@tanstack/react-query" -import React from "react" - -import { stubTransport } from "@src/utils/test-utils" - -import { useRequestyKeyInfo } from "../useRequestyKeyInfo" - -const KEY_INFO = { - name: "my-key", - monthly_limit: "20", - monthly_spend: "1.25", - org_balance: "18.75", - config: { aliases: { "gpt-4o": "openai/gpt-4o" } }, -} - -const wrapper = ({ children }: React.PropsWithChildren) => ( - - {children} - -) - -describe("useRequestyKeyInfo", () => { - afterEach(() => { - vi.restoreAllMocks() - }) - - it("returns undefined and issues no request when apiKey is missing", async () => { - const { calls } = stubTransport(() => ({ status: 200, body: KEY_INFO })) - const consoleError = vi.spyOn(console, "error").mockImplementation(() => {}) - - const { result } = renderHook(() => useRequestyKeyInfo(undefined, undefined), { wrapper }) - - await waitFor(() => expect(result.current.fetchStatus).toBe("idle")) - expect(result.current.data).toBeUndefined() - expect(calls).toHaveLength(0) - expect(consoleError).not.toHaveBeenCalled() - }) - - it("requests the default api key endpoint with bearer auth and returns the parsed key info", async () => { - const { calls } = stubTransport(() => ({ status: 200, body: KEY_INFO })) - - const { result } = renderHook(() => useRequestyKeyInfo(undefined, "test-key"), { wrapper }) - - await waitFor(() => expect(result.current.isSuccess).toBe(true)) - expect(result.current.data).toEqual(KEY_INFO) - expect(calls).toEqual([ - { - url: "https://api.requesty.ai/x/apikey", - method: "GET", - headers: { authorization: "Bearer test-key", "content-type": "application/json" }, - }, - ]) - }) - - it("keeps a caller-supplied host that does not contain 'router' and drops the version suffix", async () => { - const { calls } = stubTransport(() => ({ status: 200, body: KEY_INFO })) - - const { result } = renderHook(() => useRequestyKeyInfo("https://custom.requesty.ai/v1", "test-key"), { - wrapper, - }) - - await waitFor(() => expect(result.current.isSuccess).toBe(true)) - expect(calls[0].url).toBe("https://custom.requesty.ai/x/apikey") - }) - - it("returns null and logs a validation failure when the payload does not match the schema", async () => { - stubTransport(() => ({ status: 200, body: { name: "my-key" } })) - const consoleError = vi.spyOn(console, "error").mockImplementation(() => {}) - - const { result } = renderHook(() => useRequestyKeyInfo(undefined, "test-key"), { wrapper }) - - await waitFor(() => expect(result.current.isSuccess).toBe(true)) - expect(result.current.data).toBeNull() - expect(consoleError).toHaveBeenCalledWith( - "Requesty API key info validation failed:", - expect.objectContaining({ - issues: expect.arrayContaining([expect.objectContaining({ path: ["monthly_limit"] })]), - }), - ) - }) - - it("returns null and logs a fetch failure when the response is not 2xx", async () => { - stubTransport(() => ({ status: 403, body: { detail: "forbidden" } })) - const consoleError = vi.spyOn(console, "error").mockImplementation(() => {}) - - const { result } = renderHook(() => useRequestyKeyInfo(undefined, "test-key"), { wrapper }) - - await waitFor(() => expect(result.current.isSuccess).toBe(true)) - expect(result.current.data).toBeNull() - expect(consoleError).toHaveBeenCalledWith( - "Error fetching Requesty key info:", - expect.objectContaining({ message: "Request failed with status code 403" }), - ) - }) - - it("returns null and logs a fetch failure when the request rejects", async () => { - const failure = new Error("network down") - stubTransport(() => { - throw failure - }) - const consoleError = vi.spyOn(console, "error").mockImplementation(() => {}) - - const { result } = renderHook(() => useRequestyKeyInfo(undefined, "test-key"), { wrapper }) - - await waitFor(() => expect(result.current.isSuccess).toBe(true)) - expect(result.current.data).toBeNull() - expect(consoleError).toHaveBeenCalledWith("Error fetching Requesty key info:", failure) - }) - - it("returns null and logs a fetch failure when a 2xx body is not valid JSON", async () => { - const parseFailure = new SyntaxError("Unexpected token '<', \"\" is not valid JSON") - stubTransport(() => ({ status: 200, body: undefined, jsonRejects: parseFailure })) - const consoleError = vi.spyOn(console, "error").mockImplementation(() => {}) - - const { result } = renderHook(() => useRequestyKeyInfo(undefined, "test-key"), { wrapper }) - - await waitFor(() => expect(result.current.isSuccess).toBe(true)) - expect(result.current.data).toBeNull() - expect(consoleError).toHaveBeenCalledWith("Error fetching Requesty key info:", parseFailure) - }) -}) diff --git a/webview-ui/src/components/ui/hooks/requestProviderInfo.ts b/webview-ui/src/components/ui/hooks/requestProviderInfo.ts new file mode 100644 index 0000000000..7702def0a9 --- /dev/null +++ b/webview-ui/src/components/ui/hooks/requestProviderInfo.ts @@ -0,0 +1,42 @@ +import type { ExtensionMessage, ProviderInfoKind, ProviderInfoParams, ProviderInfoResultMap } from "@roo-code/types" + +import { vscode } from "@src/utils/vscode" + +// Longer than every host-side fetch timeout (15 s at most), so the host answers first. +const TIMEOUT_MS = 20_000 + +/** Ask the extension host to run a provider lookup. The webview CSP allows no provider origins. */ +export const requestProviderInfo = ( + kind: K, + params: ProviderInfoParams, +): Promise => + new Promise((resolve, reject) => { + const requestId = crypto.randomUUID() + + const cleanup = () => { + clearTimeout(timeout) + window.removeEventListener("message", handler) + } + + const timeout = setTimeout(() => { + cleanup() + reject(new Error("Provider info request timed out")) + }, TIMEOUT_MS) + + const handler = (event: MessageEvent) => { + const message: ExtensionMessage = event.data + + if (message.type !== "providerInfo" || message.requestId !== requestId) return + + cleanup() + + if (message.error) { + reject(new Error(message.error)) + } else { + resolve(message.values?.data as ProviderInfoResultMap[K]) + } + } + + window.addEventListener("message", handler) + vscode.postMessage({ type: "requestProviderInfo", requestId, values: { kind, ...params } }) + }) diff --git a/webview-ui/src/components/ui/hooks/useOpenRouterKeyInfo.ts b/webview-ui/src/components/ui/hooks/useOpenRouterKeyInfo.ts index 406da24eb9..6dd6385491 100644 --- a/webview-ui/src/components/ui/hooks/useOpenRouterKeyInfo.ts +++ b/webview-ui/src/components/ui/hooks/useOpenRouterKeyInfo.ts @@ -1,59 +1,25 @@ -import { z } from "zod" import { useQuery, UseQueryOptions } from "@tanstack/react-query" -// Define schema for OpenRouter key response -const openRouterKeyInfoSchema = z.object({ - data: z.object({ - label: z.string(), - usage: z.number(), - is_free_tier: z.boolean(), - is_provisioning_key: z.boolean(), - rate_limit: z.object({ - requests: z.number(), - interval: z.string(), - }), - limit: z.number().nullable(), - }), -}) +import type { ProviderInfoResultMap } from "@roo-code/types" -export type OpenRouterKeyInfo = z.infer["data"] +import { requestProviderInfo } from "./requestProviderInfo" + +type OpenRouterKeyInfo = ProviderInfoResultMap["openRouterKeyInfo"] async function getOpenRouterKeyInfo(apiKey?: string, baseUrl?: string) { if (!apiKey) return null try { - // Use the provided base URL or default to OpenRouter's API URL - const apiBaseUrl = baseUrl || "https://openrouter.ai/api/v1" - - const keyEndpoint = `${apiBaseUrl}/key` - - const response = await fetch(keyEndpoint, { - method: "GET", - headers: { - Authorization: `Bearer ${apiKey}`, - }, - }) - - if (!response.ok) { - throw new Error(`Request failed with status code ${response.status}`) - } - - const result = openRouterKeyInfoSchema.safeParse(await response.json()) - if (!result.success) { - console.error("OpenRouter API key info validation failed:", result.error) - return null - } - - return result.data.data + return await requestProviderInfo("openRouterKeyInfo", { apiKey, baseUrl }) } catch (error) { console.error("Error fetching OpenRouter key info:", error) return null } } -type UseOpenRouterKeyInfoOptions = Omit, "queryKey" | "queryFn"> +type UseOpenRouterKeyInfoOptions = Omit, "queryKey" | "queryFn"> export const useOpenRouterKeyInfo = (apiKey?: string, baseUrl?: string, options?: UseOpenRouterKeyInfoOptions) => { - return useQuery({ + return useQuery({ queryKey: ["openrouter-key-info", apiKey, baseUrl], queryFn: () => getOpenRouterKeyInfo(apiKey, baseUrl), staleTime: 30 * 1000, // 30 seconds diff --git a/webview-ui/src/components/ui/hooks/useOpenRouterModelProviders.ts b/webview-ui/src/components/ui/hooks/useOpenRouterModelProviders.ts index 175ee7b8df..1150b1df89 100644 --- a/webview-ui/src/components/ui/hooks/useOpenRouterModelProviders.ts +++ b/webview-ui/src/components/ui/hooks/useOpenRouterModelProviders.ts @@ -1,99 +1,23 @@ -import { z } from "zod" import { useQuery, UseQueryOptions } from "@tanstack/react-query" -import type { ModelInfo } from "@roo-code/types" +import type { ProviderInfoResultMap } from "@roo-code/types" -import { parseApiPrice } from "@roo/cost" +import { requestProviderInfo } from "./requestProviderInfo" export const OPENROUTER_DEFAULT_PROVIDER_NAME = "[default]" -const openRouterEndpointsSchema = z.object({ - data: z.object({ - id: z.string(), - name: z.string(), - description: z.string().optional(), - architecture: z - .object({ - input_modalities: z.array(z.string()).nullish(), - output_modalities: z.array(z.string()).nullish(), - tokenizer: z.string().nullish(), - }) - .nullish(), - endpoints: z.array( - z.object({ - name: z.string(), - tag: z.string().optional(), - context_length: z.number(), - max_completion_tokens: z.number().nullish(), - pricing: z - .object({ - prompt: z.union([z.string(), z.number()]).optional(), - completion: z.union([z.string(), z.number()]).optional(), - input_cache_read: z.union([z.string(), z.number()]).optional(), - input_cache_write: z.union([z.string(), z.number()]).optional(), - }) - .optional(), - }), - ), - }), -}) - -type OpenRouterModelProvider = ModelInfo & { - label: string -} +type OpenRouterModelProvider = ProviderInfoResultMap["openRouterModelProviders"][string] async function getOpenRouterProvidersForModel(modelId: string, baseUrl?: string) { - const models: Record = {} - try { - const apiBaseUrl = baseUrl || "https://openrouter.ai/api/v1" - const response = await fetch(`${apiBaseUrl}/models/${modelId}/endpoints`, { method: "GET" }) - - if (!response.ok) { - throw new Error(`Request failed with status code ${response.status}`) - } - - const result = openRouterEndpointsSchema.safeParse(await response.json()) - - if (!result.success) { - console.error("OpenRouter API response validation failed:", result.error) - return models - } - - const { description, architecture, endpoints } = result.data.data - - // Skip image generation models (models that output images) - if (architecture?.output_modalities?.includes("image")) { - return models - } - - for (const endpoint of endpoints) { - const providerName = endpoint.tag ?? endpoint.name - const inputPrice = parseApiPrice(endpoint.pricing?.prompt) - const outputPrice = parseApiPrice(endpoint.pricing?.completion) - const cacheReadsPrice = parseApiPrice(endpoint.pricing?.input_cache_read) - const cacheWritesPrice = parseApiPrice(endpoint.pricing?.input_cache_write) - - const modelInfo: OpenRouterModelProvider = { - maxTokens: endpoint.max_completion_tokens || endpoint.context_length, - contextWindow: endpoint.context_length, - supportsImages: architecture?.input_modalities?.includes("image") ?? false, - supportsPromptCache: typeof cacheReadsPrice !== "undefined", - cacheReadsPrice, - cacheWritesPrice, - inputPrice, - outputPrice, - description, - label: providerName, - } - - models[providerName] = modelInfo - } + return await requestProviderInfo("openRouterModelProviders", { + modelId, + baseUrl, + }) } catch (error) { console.error(`Error fetching OpenRouter providers:`, error) + return {} } - - return models } type UseOpenRouterModelProvidersOptions = Omit< diff --git a/webview-ui/src/components/ui/hooks/useRequestyKeyInfo.ts b/webview-ui/src/components/ui/hooks/useRequestyKeyInfo.ts index defc295173..622e28ac8e 100644 --- a/webview-ui/src/components/ui/hooks/useRequestyKeyInfo.ts +++ b/webview-ui/src/components/ui/hooks/useRequestyKeyInfo.ts @@ -1,54 +1,25 @@ -import { z } from "zod" import { useQuery, UseQueryOptions } from "@tanstack/react-query" -import { toRequestyServiceUrl } from "@roo/utils/requesty" -const requestyKeyInfoSchema = z.object({ - name: z.string(), - monthly_limit: z.string(), - monthly_spend: z.string(), - org_balance: z.string(), - config: z.object({ - aliases: z.record(z.string(), z.any()).optional(), - }), -}) +import type { ProviderInfoResultMap } from "@roo-code/types" -export type RequestyKeyInfo = z.infer +import { requestProviderInfo } from "./requestProviderInfo" + +type RequestyKeyInfo = ProviderInfoResultMap["requestyKeyInfo"] async function getRequestyKeyInfo(baseUrl?: string, apiKey?: string) { if (!apiKey) return null - const url = toRequestyServiceUrl(baseUrl, "api") - const apiKeyUrl = new URL("x/apikey", url) - try { - const response = await fetch(apiKeyUrl.toString(), { - method: "GET", - headers: { - Authorization: `Bearer ${apiKey}`, - "Content-Type": "application/json", - }, - }) - - if (!response.ok) { - throw new Error(`Request failed with status code ${response.status}`) - } - - const result = requestyKeyInfoSchema.safeParse(await response.json()) - if (!result.success) { - console.error("Requesty API key info validation failed:", result.error) - return null - } - - return result.data + return await requestProviderInfo("requestyKeyInfo", { apiKey, baseUrl }) } catch (error) { console.error("Error fetching Requesty key info:", error) return null } } -type UseRequestyKeyInfoOptions = Omit, "queryKey" | "queryFn"> +type UseRequestyKeyInfoOptions = Omit, "queryKey" | "queryFn"> export const useRequestyKeyInfo = (baseUrl?: string, apiKey?: string, options?: UseRequestyKeyInfoOptions) => { - return useQuery({ + return useQuery({ queryKey: ["requesty-key-info", baseUrl, apiKey], queryFn: () => getRequestyKeyInfo(baseUrl, apiKey), staleTime: 30 * 1000, // 30 seconds diff --git a/webview-ui/src/utils/test-utils.tsx b/webview-ui/src/utils/test-utils.tsx index b64d88e0b6..847c401f2c 100644 --- a/webview-ui/src/utils/test-utils.tsx +++ b/webview-ui/src/utils/test-utils.tsx @@ -84,36 +84,3 @@ export * from "@testing-library/react" // override render method export { customRender as render } - -type TransportCall = { url: string; method: string; headers: Record } -type TransportResult = { status: number; body: unknown; jsonRejects?: Error } - -// Stubs global fetch and records the calls. -export const stubTransport = (handler: (call: TransportCall) => TransportResult | Promise) => { - const calls: TransportCall[] = [] - - vi.spyOn(globalThis, "fetch").mockImplementation(async (input, init) => { - const request = typeof input === "object" && "url" in input ? input : undefined - const call: TransportCall = { - url: request?.url ?? String(input), - method: init?.method ?? request?.method ?? "GET", - headers: Object.fromEntries(new Headers(init?.headers ?? request?.headers).entries()), - } - calls.push(call) - const result = await handler(call) - - // The hooks consume only ok, status, and json from this response double. - return { - ok: result.status >= 200 && result.status < 300, - status: result.status, - json: async () => { - if (result.jsonRejects) { - throw result.jsonRejects - } - return result.body - }, - } as Response - }) - - return { calls } -}