diff --git a/src/api/providers/__tests__/vscode-lm.spec.ts b/src/api/providers/__tests__/vscode-lm.spec.ts index 06681fe374..101306b4fa 100644 --- a/src/api/providers/__tests__/vscode-lm.spec.ts +++ b/src/api/providers/__tests__/vscode-lm.spec.ts @@ -107,7 +107,7 @@ import { normalizeToolSchema } from "../../../utils/json-schema" import { getMcpServerTools } from "../../../core/prompts/tools/native-tools/mcp_server" import type { McpHub } from "../../../services/mcp/McpHub" import { clearAllMocks } from "../../../test-utils/reset" -import { collectStream } from "../../../test-utils/stream" +import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" const mockLanguageModelChat = { id: "test-model", @@ -141,6 +141,194 @@ describe("VsCodeLmHandler", () => { handler.dispose() }) + describe("image input", () => { + const image: Anthropic.ImageBlockParam = { + type: "image", + source: { type: "base64", media_type: "image/png", data: "iVBORw==" }, + } + const capabilityCases = [ + { capabilities: { supportsImageToText: true }, expected: true }, + { capabilities: { supportsImageToText: false, imageInput: true }, expected: false }, + { capabilities: { imageInput: true }, expected: true }, + { capabilities: { imageInput: false }, expected: false }, + { capabilities: {}, expected: undefined }, + { capabilities: undefined, expected: undefined }, + ] + + it.each(capabilityCases)( + "uses host capabilities $capabilities for model info and listing", + async ({ capabilities, expected }) => { + await handler.initializeClient() + const client = { ...mockLanguageModelChat, family: vscodeLlmDefaultModelId, capabilities } + handler["client"] = client + vi.mocked(vscode.lm.selectChatModels).mockResolvedValue([client]) + + expect(handler.getModel().info.supportsImages).toBe(expected) + expect((await getVsCodeLmModels())[0].modelInfo.supportsImages).toBe(expected) + }, + ) + + it.each( + (["message", "tool_result"] as const).flatMap((path) => + [true, false, undefined].map((supportsImages) => ({ path, supportsImages })), + ), + )( + "forwards image bytes in the $path path only when supported ($supportsImages)", + async ({ path, supportsImages }) => { + await handler.initializeClient() + const sendRequest = vi.fn().mockResolvedValue({ + stream: asyncStreamFrom([new vscode.LanguageModelTextPart("done")]), + text: asyncStreamFrom(["done"]), + }) + const countTokens = vi.fn().mockResolvedValue(37) + handler["client"] = { + ...mockLanguageModelChat, + ...{ capabilities: { supportsImageToText: supportsImages } }, + sendRequest, + countTokens, + } + const messages: Anthropic.Messages.MessageParam[] = + path === "message" + ? [{ role: "user", content: [{ type: "text", text: "look" }, image] }] + : [ + { + role: "assistant", + content: [{ type: "tool_use", id: "image-tool", name: "read_file", input: {} }], + }, + { + role: "user", + content: [ + { + type: "tool_result", + tool_use_id: "image-tool", + content: [{ type: "text", text: "look" }, image], + }, + ], + }, + ] + const chunks = await collectStream(handler.createMessage("system", messages)) + const sent = sendRequest.mock.calls[0][0] + expect(chunks).toContainEqual({ type: "usage", inputTokens: sent.length * 37, outputTokens: 37 }) + const counted = countTokens.mock.calls.map(([content]) => content) + if (supportsImages) { + expect(counted).toEqual(expect.arrayContaining(sent)) + } else { + expect(counted.every((content) => typeof content === "string")).toBe(true) + } + const content = sent[sent.length - 1].content + const parts = + path === "tool_result" && content[0] instanceof vscode.LanguageModelToolResultPart + ? content[0].content + : content + expect(parts[0]).toMatchObject({ value: "look" }) + if (supportsImages) { + expect(parts[1]).toMatchObject({ + mimeType: "image/png", + data: Buffer.from("iVBORw==", "base64"), + }) + } else { + expect(parts[1]).toBeInstanceOf(vscode.LanguageModelTextPart) + expect(parts[1]).toMatchObject({ value: expect.stringContaining("not supported") }) + } + }, + ) + + it.each(["message", "tool_result"] as const)("counts native image bytes in %s content", async (path) => { + await handler.initializeClient() + const countTokens = vi.fn().mockResolvedValue(321) + handler["client"] = { + ...mockLanguageModelChat, + ...{ capabilities: { imageInput: true } }, + countTokens, + } + const content: Anthropic.Messages.ContentBlockParam[] = + path === "message" ? [image] : [{ type: "tool_result", tool_use_id: "image-tool", content: [image] }] + + await expect(handler.countTokens(content)).resolves.toBe(321) + const counted = countTokens.mock.calls[0][0] + expect(counted).not.toEqual(expect.any(String)) + if (typeof counted === "string") throw new Error("Expected an intact chat message") + const part = counted.content[0] + const imagePart = part instanceof vscode.LanguageModelToolResultPart ? part.content[0] : part + expect(imagePart).toMatchObject({ mimeType: "image/png", data: Buffer.from("iVBORw==", "base64") }) + }) + + const imageCapable = { imageInput: true } + it.each([ + { capabilities: imageCapable, liveWindow: 100, curatedWindow: 200, inputTokens: 100, accepted: true }, + { capabilities: imageCapable, liveWindow: 100, curatedWindow: 200, inputTokens: 101, accepted: false }, + { capabilities: imageCapable, liveWindow: 200, curatedWindow: 100, inputTokens: 100, accepted: true }, + { capabilities: imageCapable, liveWindow: 200, curatedWindow: 100, inputTokens: 101, accepted: false }, + { capabilities: {}, liveWindow: 100, curatedWindow: 200, inputTokens: 101, accepted: true }, + { + capabilities: { supportsImageToText: false }, + liveWindow: 200, + curatedWindow: 100, + inputTokens: 101, + accepted: true, + }, + ])("checks measured input against both windows only for image models: %j", async (fixture) => { + await handler.initializeClient() + const userRole = vscode.LanguageModelChatMessage.User("").role + const sendRequest = vi.fn().mockResolvedValue({ + stream: asyncStreamFrom([new vscode.LanguageModelTextPart("done")]), + text: asyncStreamFrom(["done"]), + }) + handler["client"] = { + ...mockLanguageModelChat, + ...{ capabilities: fixture.capabilities }, + maxInputTokens: fixture.liveWindow, + sendRequest, + // Image models count whole messages; text-only models count extracted strings. + countTokens: vi.fn().mockImplementation(async (content) => { + if (typeof content === "string") { + return fixture.capabilities === imageCapable ? 0 : fixture.inputTokens + } + return content.role === userRole ? fixture.inputTokens : 0 + }), + } + const windowSpy = vi.spyOn(handler, "getCondenseContextWindow").mockReturnValue(fixture.curatedWindow) + try { + const request = collectStream(handler.createMessage("system", [{ role: "user", content: [image] }])) + if (fixture.capabilities !== imageCapable) { + await expect(request).resolves.toBeDefined() + expect(sendRequest).toHaveBeenCalledTimes(1) + } else if (fixture.accepted) { + await expect(request).resolves.toContainEqual({ + type: "usage", + inputTokens: fixture.inputTokens, + outputTokens: 0, + }) + expect(sendRequest).toHaveBeenCalledTimes(1) + } else { + await expect(request).rejects.toSatisfy(checkContextWindowExceededError) + expect(sendRequest).not.toHaveBeenCalled() + } + } finally { + windowSpy.mockRestore() + } + }) + + it("does not advertise image input before a model has been selected", async () => { + await handler.initializeClient() + handler["client"] = null + expect(handler.getModel().info.supportsImages).toBeUndefined() + }) + + it("disables image support on hosts without image parts, even when the model reports support", async () => { + await handler.initializeClient() + handler["client"] = { ...mockLanguageModelChat, ...{ capabilities: { imageInput: true } } } + const dataPart: { image: unknown } = Reflect.get(vscode, "LanguageModelDataPart") + const original = dataPart.image + try { + dataPart.image = undefined + expect(handler.getModel().info.supportsImages).toBe(false) + } finally { + dataPart.image = original + } + }) + }) + describe("constructor", () => { it("should initialize with provided options", () => { expect(handler).toBeDefined() @@ -4540,11 +4728,17 @@ describe("getVsCodeLmModels", () => { beforeEach(() => vi.clearAllMocks()) - it("returns the host's models as they are, with no added capability data", async () => { + it("returns serializable model metadata with host-reported image support", async () => { const model = liveModel({ capabilities: { supportsImageToText: true } }) - vi.mocked(vscode.lm.selectChatModels).mockResolvedValue([model] as never) - - await expect(getVsCodeLmModels()).resolves.toEqual([model]) + vi.mocked(vscode.lm.selectChatModels).mockResolvedValue([model]) + + const models = await getVsCodeLmModels() + expect(models[0]).toMatchObject({ id: model.id, family: model.family, modelInfo: { supportsImages: true } }) + expect(models[0]).not.toHaveProperty("sendRequest") + expect(models[0]).not.toHaveProperty("countTokens") + expect(JSON.parse(JSON.stringify(models))).toMatchObject([ + { id: model.id, modelInfo: { supportsImages: true } }, + ]) }) it("asks the host for every model, whatever the vendor", async () => { diff --git a/src/api/providers/vscode-lm.ts b/src/api/providers/vscode-lm.ts index 7842326b2f..38f04ee8da 100644 --- a/src/api/providers/vscode-lm.ts +++ b/src/api/providers/vscode-lm.ts @@ -27,6 +27,7 @@ import { import { CONTEXT_WINDOW_EXCEEDED_STATUS } from "../../core/context/context-management/context-error-handling" import { BaseProvider } from "./base-provider" +import { getVsCodeLmModelInfo } from "./vscode-lm-capabilities" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" /** @@ -409,18 +410,8 @@ function declaredParamType(schema: Record | undefined, paramNam * details), and only when the request would otherwise exceed the budget. */ -/** - * Conservative characters-per-token ratio used to turn a token window into a character budget. - * - * `client.countTokens` is the model's real tokenizer, but it counts only a string: it cannot price - * the tool schemas, image placeholders, or per-message framing the backend adds, so it cannot give - * the true total for the request we are about to send. It is also an async, per-call RPC, and the - * budget is needed for every message on every turn. We therefore keep a character estimate here - * and stay deliberately conservative — 3 chars/token rather than the ~4 typical of English — - * because the token-dense JSON, logs, and code that dominate oversized tool results tokenize to - * fewer characters per token than prose. Under-counting biases toward trimming too early, which is - * recoverable; over-counting sends an over-window request, which is not. - */ +// Text trimming uses 3 chars/token rather than ~4 for token-dense code and logs. +// Native image input is counted separately by the host before sending; this budget is text-only. const VSCODE_LM_BUDGET_CHARS_PER_TOKEN = 3 /** @@ -433,11 +424,7 @@ const VSCODE_LM_INPUT_BUDGET_FRACTION = 0.8 /** A tool_result is never shrunk below this many characters, so a truncated result stays useful. */ const MIN_TOOL_RESULT_CHARS = 2000 -/** - * Length charged for an image block. VS Code LM cannot carry image data, so - * `convertToVsCodeLmMessages` replaces each image with a sentence-long textual placeholder; this - * is that placeholder's approximate length. - */ +/** Approximate placeholder length; the text-only budget does not measure native image tokens. */ const IMAGE_PLACEHOLDER_CHARS = 64 function readToolResultText(block: Anthropic.Messages.ContentBlockParam): string | undefined { @@ -539,8 +526,6 @@ function estimateContentChars(content: Anthropic.Messages.MessageParam["content" } else if (type === "tool_use") { total += JSON.stringify((block as Anthropic.Messages.ToolUseBlockParam).input ?? {}).length } else if (type === "image") { - // VS Code LM cannot send image data; convertToVsCodeLmMessages substitutes a textual - // placeholder, so charge that placeholder's real length rather than a token-sized guess. total += IMAGE_PLACEHOLDER_CHARS } } @@ -902,7 +887,7 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan typeof client.maxInputTokens === "number" ? Math.max(0, client.maxInputTokens) : openAiModelInfoSaneDefaults.contextWindow, - supportsImages: false, // VSCode Language Model API historically had no image input + supportsImages: getVsCodeLmModelInfo(client).supportsImages, supportsPromptCache: true, inputPrice: 0, outputPrice: 0, @@ -912,25 +897,38 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan /** Model info reported before any client exists. */ protected getFallbackModelInfo(): ModelInfo { - return openAiModelInfoSaneDefaults + return { ...openAiModelInfoSaneDefaults, supportsImages: undefined } } /** Whether image parts are sent to the model. */ - protected supportsImageInput(_info: ModelInfo): boolean { - return false + protected supportsImageInput(info: ModelInfo): boolean { + return info.supportsImages === true } /** Rejects a request the selected model cannot serve before anything is sent. */ protected assertRequestSupported(_messages: Anthropic.Messages.MessageParam[], _info: ModelInfo): void {} /** Rejects a request whose measured input exceeds what the model accepts. */ - protected assertWithinContextWindow(_inputTokens: number, _info: ModelInfo): void {} + protected assertWithinContextWindow(inputTokens: number, info: ModelInfo): void { + // Native images are absent from the character budget; use the host's measured input too. + if ( + this.supportsImageInput(info) && + inputTokens > Math.min(this.getCondenseContextWindow(), info.contextWindow) + ) { + throw Object.assign( + new Error( + "The request exceeds the selected VS Code model's input context window. Condense the conversation or start a new task.", + ), + { status: CONTEXT_WINDOW_EXCEEDED_STATUS }, + ) + } + } /** The form in which a chat message is handed to the host's token counter. */ protected prepareMessageForCounting( message: vscode.LanguageModelChatMessage, ): string | vscode.LanguageModelChatMessage { - return extractTextCountFromMessage(message) + return this.supportsImageInput(this.getModel().info) ? message : extractTextCountFromMessage(message) } /** @@ -969,14 +967,17 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan * @returns A promise resolving to the token count */ override async countTokens(content: Array): Promise { - // Convert Anthropic content blocks to a string for VSCode LM token counting + if (this.supportsImageInput(this.getModel().info) && containsImageBlock(content)) { + return this.calculateTotalInputTokens(convertToVsCodeLmMessages([{ role: "user", content }], true)) + } + + // Preserve text-only counting for models that cannot send native images. let textContent = "" for (const block of content) { if (block.type === "text") { textContent += block.text || "" } else if (block.type === "image") { - // VSCode LM doesn't support images directly, so we'll just use a placeholder textContent += "[IMAGE]" } } @@ -1597,7 +1598,17 @@ export const VSCODE_LM_STATIC_BLACKLIST: string[] = ["claude-3.7-sonnet", "claud export async function getVsCodeLmModels() { try { const models = (await vscode.lm.selectChatModels({})) || [] - return models.filter((model) => !VSCODE_LM_STATIC_BLACKLIST.includes(model.id)) + return models + .filter((model) => !VSCODE_LM_STATIC_BLACKLIST.includes(model.id)) + .map((model) => ({ + id: model.id, + vendor: model.vendor, + family: model.family, + version: model.version, + name: model.name, + maxInputTokens: model.maxInputTokens, + modelInfo: getVsCodeLmModelInfo(model), + })) } catch (error) { console.error( `Error fetching VS Code LM models: ${JSON.stringify(error, Object.getOwnPropertyNames(error), 2)}`, 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 3004bee9b8..61f6f8d0e2 100644 --- a/webview-ui/src/components/ui/hooks/__tests__/useSelectedModel.spec.ts +++ b/webview-ui/src/components/ui/hooks/__tests__/useSelectedModel.spec.ts @@ -59,6 +59,9 @@ import { useLmStudioModels } from "../useLmStudioModels" import { useOllamaModels } from "../useOllamaModels" import { useGitHubCopilotModels } from "../useGitHubCopilotModels" +import { useVsCodeLmModels } from "../useVsCodeLmModels" + +vi.mock("../useVsCodeLmModels") vi.mock("../useRouterModels") vi.mock("../useOpenRouterModelProviders") vi.mock("../useLmStudioModels") @@ -116,6 +119,7 @@ describe("useSelectedModel", () => { mockUseLmStudioModels.mockReturnValue(createLocalModelsResult({})) mockUseOllamaModels.mockReturnValue(createLocalModelsResult({})) mockUseGitHubCopilotModels.mockReturnValue({ data: undefined } as ReturnType) + vi.mocked(useVsCodeLmModels).mockReturnValue({ data: undefined } as ReturnType) }) const dynamicProviderCases = [ @@ -1671,6 +1675,101 @@ describe("useSelectedModel", () => { } as any) }) + it.each([true, false, undefined])( + "uses the generic host's image capability (%s), not the catalog", + (supportsImages) => { + vi.mocked(useVsCodeLmModels).mockReturnValue({ + data: [ + { + vendor: "other", + family: vscodeLlmDefaultModelId, + modelInfo: { contextWindow: 100, supportsImages }, + }, + ], + } as ReturnType) + const { result } = renderHook(() => + useSelectedModel({ + apiProvider: providerIdentifiers.vscodeLm, + vsCodeLmModelSelector: { vendor: "other", family: vscodeLlmDefaultModelId }, + }), + ) + expect(result.current.info?.supportsImages).toBe(supportsImages) + expect(result.current.info?.contextWindow).toBe(vscodeLlmModels[vscodeLlmDefaultModelId].maxInputTokens) + expect(useVsCodeLmModels).toHaveBeenLastCalledWith(true) + }, + ) + + it("does not request generic VS Code LM models for other providers", () => { + renderHook(() => useSelectedModel({ apiProvider: providerIdentifiers.openaiNative })) + expect(useVsCodeLmModels).toHaveBeenLastCalledWith(false) + }) + + it.each([ + { vendor: "other", family: "unknown" }, + { vendor: "wanted", family: "unknown", id: "missing" }, + { vendor: "wanted", family: "unknown", version: "missing" }, + ])("does not borrow image support from a different selector: %j", (selector) => { + vi.mocked(useVsCodeLmModels).mockReturnValue({ + data: [ + { + vendor: "wanted", + family: "unknown", + id: "present", + version: "1", + modelInfo: { contextWindow: 100, supportsImages: true }, + }, + ], + } as ReturnType) + const { result } = renderHook(() => + useSelectedModel({ + apiProvider: providerIdentifiers.vscodeLm, + vsCodeLmModelSelector: selector, + }), + ) + expect(result.current.info?.supportsImages).toBeUndefined() + }) + + it.each([false, true])("keeps ambiguous generic selectors unknown (reverse order: %s)", (reverse) => { + const models = [true, false].map((supportsImages) => ({ + vendor: "other", + family: "shared-family", + id: supportsImages ? "vision" : "text", + version: "1", + modelInfo: { contextWindow: 100, supportsImages }, + })) + vi.mocked(useVsCodeLmModels).mockReturnValue({ + data: reverse ? models.reverse() : models, + } as ReturnType) + + const { result, rerender } = renderHook( + ({ selector }: { selector: ProviderSettings["vsCodeLmModelSelector"] }) => + useSelectedModel({ apiProvider: providerIdentifiers.vscodeLm, vsCodeLmModelSelector: selector }), + { initialProps: { selector: undefined as ProviderSettings["vsCodeLmModelSelector"] } }, + ) + expect(result.current.info?.supportsImages).toBeUndefined() + rerender({ selector: { vendor: "other", family: "shared-family" } }) + expect(result.current.info?.supportsImages).toBeUndefined() + + for (const model of models) { + const { id, vendor, family, version } = model + rerender({ selector: { id, vendor, family, version } }) + expect(result.current.info?.supportsImages).toBe(model.modelInfo.supportsImages) + } + }) + + it("enables an unlisted family when the host reports image support", () => { + vi.mocked(useVsCodeLmModels).mockReturnValue({ + data: [{ vendor: "other", family: "unknown", modelInfo: { contextWindow: 100, supportsImages: true } }], + } as ReturnType) + const { result } = renderHook(() => + useSelectedModel({ + apiProvider: providerIdentifiers.vscodeLm, + vsCodeLmModelSelector: { vendor: "other", family: "unknown" }, + }), + ) + expect(result.current.info?.supportsImages).toBe(true) + }) + it("resolves a listed family's contextWindow to its maxInputTokens", () => { const family = vscodeLlmDefaultModelId const apiConfiguration: ProviderSettings = { @@ -1685,7 +1784,7 @@ describe("useSelectedModel", () => { expect(result.current.id).toBe(`copilot/${family}`) // The bar and the condense gate share one source of truth: contextWindow === maxInputTokens. expect(result.current.info?.contextWindow).toBe(vscodeLlmModels[family].maxInputTokens) - expect(result.current.info?.supportsImages).toBe(false) + expect(result.current.info?.supportsImages).toBeUndefined() }) it.each(["claude-opus-4.8", "claude-opus-5", "gemini-3.6-flash"] as const)( @@ -1704,7 +1803,7 @@ describe("useSelectedModel", () => { expect(result.current.info?.contextWindow).toBe(197897) expect(result.current.info?.contextWindow).toBe(vscodeLlmModels[family].maxInputTokens) expect(result.current.info?.contextWindow).not.toBe(vscodeLlmModels[family].contextWindow) - expect(result.current.info?.supportsImages).toBe(false) + expect(result.current.info?.supportsImages).toBeUndefined() }, ) @@ -1720,7 +1819,7 @@ describe("useSelectedModel", () => { // A family miss must not use the 128000 sane-defaults window; use the default model's instead. expect(result.current.info?.contextWindow).not.toBe(128000) expect(result.current.info?.contextWindow).toBe(vscodeLlmModels[vscodeLlmDefaultModelId].maxInputTokens) - expect(result.current.info?.supportsImages).toBe(false) + expect(result.current.info?.supportsImages).toBeUndefined() }) }) diff --git a/webview-ui/src/components/ui/hooks/__tests__/useVsCodeLmModels.spec.tsx b/webview-ui/src/components/ui/hooks/__tests__/useVsCodeLmModels.spec.tsx new file mode 100644 index 0000000000..7765e8d0b5 --- /dev/null +++ b/webview-ui/src/components/ui/hooks/__tests__/useVsCodeLmModels.spec.tsx @@ -0,0 +1,109 @@ +import React from "react" +import { QueryClientProvider } from "@tanstack/react-query" +import { act, renderHook, waitFor } from "@testing-library/react" + +import { type ExtensionMessage, providerIdentifiers, VsCodeLmModelsMessageType } from "@roo-code/types" + +import { createTestQueryClient } from "@src/utils/test-utils" +import { vscode } from "@src/utils/vscode" + +import { useVsCodeLmModels } from "../useVsCodeLmModels" + +vi.mock("@src/utils/vscode", () => ({ vscode: { postMessage: vi.fn() } })) + +const send = (data: ExtensionMessage) => act(() => void window.dispatchEvent(new MessageEvent("message", { data }))) +const model = { + id: "vision", + vendor: "other", + family: "vision", + modelInfo: { contextWindow: 100, supportsPromptCache: false, supportsImages: true }, +} + +const setup = (enabled = true) => { + const client = createTestQueryClient() + const wrapper = ({ children }: { children: React.ReactNode }) => ( + {children} + ) + return { + client, + ...renderHook(({ enabled }) => useVsCodeLmModels(enabled), { wrapper, initialProps: { enabled } }), + } +} +const cachedModels = (client: ReturnType) => client.getQueryData(["vsCodeLmModels"]) + +describe("useVsCodeLmModels", () => { + beforeEach(() => vi.clearAllMocks()) + + it("requests generic models when enabled, leaving capabilities unknown until the reply", () => { + const { result } = setup() + expect(result.current.data).toBeUndefined() + expect(vscode.postMessage).toHaveBeenCalledTimes(1) + expect(vscode.postMessage).toHaveBeenCalledWith({ + type: VsCodeLmModelsMessageType.requestVsCodeLmModels, + apiConfiguration: { apiProvider: providerIdentifiers.vscodeLm }, + }) + }) + + it("updates capabilities and clears models on an empty reply", async () => { + const { result } = setup() + send({ type: VsCodeLmModelsMessageType.vsCodeLmModels, vsCodeLmModels: [model] }) + await waitFor(() => expect(result.current.data?.[0].modelInfo?.supportsImages).toBe(true)) + send({ + type: VsCodeLmModelsMessageType.vsCodeLmModels, + vsCodeLmModels: [{ ...model, modelInfo: { ...model.modelInfo, supportsImages: false } }], + }) + await waitFor(() => expect(result.current.data?.[0].modelInfo?.supportsImages).toBe(false)) + send({ type: VsCodeLmModelsMessageType.vsCodeLmModels, vsCodeLmModels: [] }) + await waitFor(() => expect(result.current.data).toEqual([])) + }) + + it.each([ + { type: VsCodeLmModelsMessageType.githubCopilotModels, vsCodeLmModels: [model] }, + { type: VsCodeLmModelsMessageType.githubCopilotSignInResult, vsCodeLmModels: [model] }, + { type: VsCodeLmModelsMessageType.vsCodeLmModels }, + ])("ignores non-generic or unsuccessful replies: %j", (message) => { + const { client } = setup() + send(message) + expect(cachedModels(client)).toBeUndefined() + }) + + it.each([{ models: undefined }, { models: [model] }])( + "clears stale capabilities on an explicit error: %j", + async ({ models }) => { + const { result } = setup() + send({ type: VsCodeLmModelsMessageType.vsCodeLmModels, vsCodeLmModels: [model] }) + await waitFor(() => expect(result.current.data?.[0].modelInfo?.supportsImages).toBe(true)) + + send({ type: VsCodeLmModelsMessageType.vsCodeLmModels, error: "unavailable", vsCodeLmModels: models }) + await waitFor(() => expect(result.current.data).toEqual([])) + + send({ type: VsCodeLmModelsMessageType.vsCodeLmModels, vsCodeLmModels: [model] }) + await waitFor(() => expect(result.current.data?.[0].modelInfo?.supportsImages).toBe(true)) + }, + ) + + it("does not request or listen when disabled, and removes its listener on unmount", () => { + const add = vi.spyOn(window, "addEventListener") + const remove = vi.spyOn(window, "removeEventListener") + const lastMessageHandler = () => add.mock.calls.filter(([type]) => type === "message").at(-1)?.[1] + const { client, rerender, unmount } = setup(false) + expect(vscode.postMessage).not.toHaveBeenCalled() + expect(lastMessageHandler()).toBeUndefined() + send({ type: VsCodeLmModelsMessageType.vsCodeLmModels, vsCodeLmModels: [model] }) + expect(cachedModels(client)).toBeUndefined() + rerender({ enabled: true }) + expect(vscode.postMessage).toHaveBeenCalledTimes(1) + const firstHandler = lastMessageHandler() + expect(firstHandler).toBeDefined() + rerender({ enabled: false }) + expect(remove).toHaveBeenCalledWith("message", firstHandler) + rerender({ enabled: true }) + const secondHandler = lastMessageHandler() + expect(secondHandler).toBeDefined() + remove.mockClear() + unmount() + expect(remove).toHaveBeenCalledWith("message", secondHandler) + add.mockRestore() + remove.mockRestore() + }) +}) diff --git a/webview-ui/src/components/ui/hooks/useSelectedModel.ts b/webview-ui/src/components/ui/hooks/useSelectedModel.ts index 1e9dddeda3..32d2392a63 100644 --- a/webview-ui/src/components/ui/hooks/useSelectedModel.ts +++ b/webview-ui/src/components/ui/hooks/useSelectedModel.ts @@ -51,6 +51,7 @@ import { useOpenRouterModelProviders } from "./useOpenRouterModelProviders" import { useLmStudioModels } from "./useLmStudioModels" import { useOllamaModels } from "./useOllamaModels" import { useGitHubCopilotModels, type GitHubCopilotModels } from "./useGitHubCopilotModels" +import { useVsCodeLmModels, type VsCodeLmModels } from "./useVsCodeLmModels" /** * Helper to get a validated model ID for dynamic providers. @@ -118,6 +119,7 @@ export const useSelectedModel = (apiConfiguration?: ProviderSettings) => { const ollamaModels = useOllamaModels(ollamaModelId) // Not part of readiness: until the host reports, the curated catalog describes the model. const copilotModels = useGitHubCopilotModels(activeProvider === providerIdentifiers.githubCopilot) + const vsCodeLmModels = useVsCodeLmModels(activeProvider === providerIdentifiers.vscodeLm) // Compute readiness only for the data actually needed for the selected provider const needRouterModels = shouldFetchRouterModels @@ -156,6 +158,7 @@ export const useSelectedModel = (apiConfiguration?: ProviderSettings) => { lmStudioModels: (lmStudioModels.data || undefined) as ModelRecord | undefined, ollamaModels: (ollamaModels.data || undefined) as ModelRecord | undefined, copilotModels: copilotModels.data, + vsCodeLmModels: vsCodeLmModels.data, }) : activeProvider === providerIdentifiers.kimiCode && apiConfiguration ? { @@ -189,6 +192,7 @@ function getSelectedModel({ lmStudioModels, ollamaModels, copilotModels, + vsCodeLmModels, }: { provider: ProviderName apiConfiguration: ProviderSettings @@ -197,6 +201,7 @@ function getSelectedModel({ lmStudioModels: ModelRecord | undefined ollamaModels: ModelRecord | undefined copilotModels?: GitHubCopilotModels + vsCodeLmModels?: VsCodeLmModels }): { id: string; info: ModelInfo | undefined } { // the `undefined` case are used to show the invalid selection to prevent // users from seeing the default model if their selection is invalid @@ -469,6 +474,14 @@ function getSelectedModel({ } } case providerIdentifiers.vscodeLm: { + const selector = apiConfiguration.vsCodeLmModelSelector + const matches = vsCodeLmModels?.filter((model) => + (["id", "vendor", "family", "version"] as const).every( + (key) => selector?.[key] === undefined || model[key] === selector[key], + ), + ) + // The host selects its first match; a separate listing cannot establish that order. + const reported = matches?.length === 1 ? matches[0] : undefined const id = apiConfiguration?.vsCodeLmModelSelector ? `${apiConfiguration.vsCodeLmModelSelector.vendor}/${apiConfiguration.vsCodeLmModelSelector.family}` : vscodeLlmDefaultModelId @@ -484,7 +497,7 @@ function getSelectedModel({ ...openAiModelInfoSaneDefaults, ...listedModel, contextWindow: listedModel.maxInputTokens, - supportsImages: false, // VSCode LM API currently doesn't support images. + supportsImages: reported?.modelInfo?.supportsImages, } return { id, info } } diff --git a/webview-ui/src/components/ui/hooks/useVsCodeLmModels.ts b/webview-ui/src/components/ui/hooks/useVsCodeLmModels.ts new file mode 100644 index 0000000000..0bf67902b1 --- /dev/null +++ b/webview-ui/src/components/ui/hooks/useVsCodeLmModels.ts @@ -0,0 +1,37 @@ +import { useEffect } from "react" +import { skipToken, useQuery, useQueryClient } from "@tanstack/react-query" + +import { type ExtensionMessage, providerIdentifiers, VsCodeLmModelsMessageType } from "@roo-code/types" + +import { vscode } from "@src/utils/vscode" + +export type VsCodeLmModels = NonNullable + +const VSCODE_LM_MODELS_QUERY_KEY = ["vsCodeLmModels"] as const + +export const useVsCodeLmModels = (enabled: boolean) => { + const queryClient = useQueryClient() + + useEffect(() => { + if (!enabled) return + + const handler = (event: MessageEvent) => { + const message = event.data + if (message.type !== VsCodeLmModelsMessageType.vsCodeLmModels) return + if (message.error) { + queryClient.setQueryData(VSCODE_LM_MODELS_QUERY_KEY, []) + } else if (message.vsCodeLmModels) { + queryClient.setQueryData(VSCODE_LM_MODELS_QUERY_KEY, message.vsCodeLmModels) + } + } + + window.addEventListener("message", handler) + vscode.postMessage({ + type: VsCodeLmModelsMessageType.requestVsCodeLmModels, + apiConfiguration: { apiProvider: providerIdentifiers.vscodeLm }, + }) + return () => window.removeEventListener("message", handler) + }, [enabled, queryClient]) + + return useQuery({ queryKey: VSCODE_LM_MODELS_QUERY_KEY, queryFn: skipToken }) +}