feat(vscode-lm): runtime image input support

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<vscode.LanguageModelChat["sendRequest"]>().mockResolvedValue({
+					stream: asyncStreamFrom([new vscode.LanguageModelTextPart("done")]),
+					text: asyncStreamFrom(["done"]),
+				})
+				const countTokens = vi.fn<vscode.LanguageModelChat["countTokens"]>().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<vscode.LanguageModelChat["countTokens"]>().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<vscode.LanguageModelChat["sendRequest"]>().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<vscode.LanguageModelChat["countTokens"]>().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 () => {

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<string, unknown> | 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<Anthropic.Messages.ContentBlockParam>): Promise<number> {
-		// 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)}`,