Skip to content

Commit 97228c7

Browse files
simurg79Bertan Ari
simurg79
and
Bertan Ari
authored
feat(vscode-lm): runtime image input support (#2159)
* feat(vscode-lm): port runtime image input support * test(webview): fix type errors in VS Code LM model hook specs * test(webview): strengthen VS Code LM hook assertions per review * test(vscode-lm): cover text-only models above the context window * test(vscode-lm): split image forwarding cases and fix context-window fixture * test(vscode-lm): assert emitted text in text-only window case --------- Co-authored-by: Bertan Ari <[email protected]>
1 parent 6a9bd2c commit 97228c7

6 files changed

Lines changed: 500 additions & 37 deletions

File tree

‎src/api/providers/__tests__/vscode-lm.spec.ts‎

Lines changed: 199 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,7 @@ import { normalizeToolSchema } from "../../../utils/json-schema"
107107
import { getMcpServerTools } from "../../../core/prompts/tools/native-tools/mcp_server"
108108
import type { McpHub } from "../../../services/mcp/McpHub"
109109
import { clearAllMocks } from "../../../test-utils/reset"
110-
import { collectStream } from "../../../test-utils/stream"
110+
import { asyncStreamFrom, collectStream } from "../../../test-utils/stream"
111111

112112
const mockLanguageModelChat = {
113113
id: "test-model",
@@ -141,6 +141,194 @@ describe("VsCodeLmHandler", () => {
141141
handler.dispose()
142142
})
143143

144+
describe("image input", () => {
145+
const image: Anthropic.ImageBlockParam = {
146+
type: "image",
147+
source: { type: "base64", media_type: "image/png", data: "iVBORw==" },
148+
}
149+
const capabilityCases = [
150+
{ capabilities: { supportsImageToText: true }, expected: true },
151+
{ capabilities: { supportsImageToText: false, imageInput: true }, expected: false },
152+
{ capabilities: { imageInput: true }, expected: true },
153+
{ capabilities: { imageInput: false }, expected: false },
154+
{ capabilities: {}, expected: undefined },
155+
{ capabilities: undefined, expected: undefined },
156+
]
157+
158+
it.each(capabilityCases)(
159+
"uses host capabilities $capabilities for model info and listing",
160+
async ({ capabilities, expected }) => {
161+
await handler.initializeClient()
162+
const client = { ...mockLanguageModelChat, family: vscodeLlmDefaultModelId, capabilities }
163+
handler["client"] = client
164+
vi.mocked(vscode.lm.selectChatModels).mockResolvedValue([client])
165+
166+
expect(handler.getModel().info.supportsImages).toBe(expected)
167+
expect((await getVsCodeLmModels())[0].modelInfo.supportsImages).toBe(expected)
168+
},
169+
)
170+
171+
it.each(
172+
(["message", "tool_result"] as const).flatMap((path) =>
173+
[true, false, undefined].map((supportsImages) => ({ path, supportsImages })),
174+
),
175+
)(
176+
"forwards image bytes in the $path path only when supported ($supportsImages)",
177+
async ({ path, supportsImages }) => {
178+
await handler.initializeClient()
179+
const sendRequest = vi.fn<vscode.LanguageModelChat["sendRequest"]>().mockResolvedValue({
180+
stream: asyncStreamFrom([new vscode.LanguageModelTextPart("done")]),
181+
text: asyncStreamFrom(["done"]),
182+
})
183+
const countTokens = vi.fn<vscode.LanguageModelChat["countTokens"]>().mockResolvedValue(37)
184+
handler["client"] = {
185+
...mockLanguageModelChat,
186+
...{ capabilities: { supportsImageToText: supportsImages } },
187+
sendRequest,
188+
countTokens,
189+
}
190+
const messages: Anthropic.Messages.MessageParam[] =
191+
path === "message"
192+
? [{ role: "user", content: [{ type: "text", text: "look" }, image] }]
193+
: [
194+
{
195+
role: "assistant",
196+
content: [{ type: "tool_use", id: "image-tool", name: "read_file", input: {} }],
197+
},
198+
{
199+
role: "user",
200+
content: [
201+
{
202+
type: "tool_result",
203+
tool_use_id: "image-tool",
204+
content: [{ type: "text", text: "look" }, image],
205+
},
206+
],
207+
},
208+
]
209+
const chunks = await collectStream(handler.createMessage("system", messages))
210+
const sent = sendRequest.mock.calls[0][0]
211+
expect(chunks).toContainEqual({ type: "usage", inputTokens: sent.length * 37, outputTokens: 37 })
212+
const counted = countTokens.mock.calls.map(([content]) => content)
213+
if (supportsImages) {
214+
expect(counted).toEqual(expect.arrayContaining(sent))
215+
} else {
216+
expect(counted.every((content) => typeof content === "string")).toBe(true)
217+
}
218+
const content = sent[sent.length - 1].content
219+
const parts =
220+
path === "tool_result" && content[0] instanceof vscode.LanguageModelToolResultPart
221+
? content[0].content
222+
: content
223+
expect(parts[0]).toMatchObject({ value: "look" })
224+
if (supportsImages) {
225+
expect(parts[1]).toMatchObject({
226+
mimeType: "image/png",
227+
data: Buffer.from("iVBORw==", "base64"),
228+
})
229+
} else {
230+
expect(parts[1]).toBeInstanceOf(vscode.LanguageModelTextPart)
231+
expect(parts[1]).toMatchObject({ value: expect.stringContaining("not supported") })
232+
}
233+
},
234+
)
235+
236+
it.each(["message", "tool_result"] as const)("counts native image bytes in %s content", async (path) => {
237+
await handler.initializeClient()
238+
const countTokens = vi.fn<vscode.LanguageModelChat["countTokens"]>().mockResolvedValue(321)
239+
handler["client"] = {
240+
...mockLanguageModelChat,
241+
...{ capabilities: { imageInput: true } },
242+
countTokens,
243+
}
244+
const content: Anthropic.Messages.ContentBlockParam[] =
245+
path === "message" ? [image] : [{ type: "tool_result", tool_use_id: "image-tool", content: [image] }]
246+
247+
await expect(handler.countTokens(content)).resolves.toBe(321)
248+
const counted = countTokens.mock.calls[0][0]
249+
expect(counted).not.toEqual(expect.any(String))
250+
if (typeof counted === "string") throw new Error("Expected an intact chat message")
251+
const part = counted.content[0]
252+
const imagePart = part instanceof vscode.LanguageModelToolResultPart ? part.content[0] : part
253+
expect(imagePart).toMatchObject({ mimeType: "image/png", data: Buffer.from("iVBORw==", "base64") })
254+
})
255+
256+
const imageCapable = { imageInput: true }
257+
it.each([
258+
{ capabilities: imageCapable, liveWindow: 100, curatedWindow: 200, inputTokens: 100, accepted: true },
259+
{ capabilities: imageCapable, liveWindow: 100, curatedWindow: 200, inputTokens: 101, accepted: false },
260+
{ capabilities: imageCapable, liveWindow: 200, curatedWindow: 100, inputTokens: 100, accepted: true },
261+
{ capabilities: imageCapable, liveWindow: 200, curatedWindow: 100, inputTokens: 101, accepted: false },
262+
{ capabilities: {}, liveWindow: 100, curatedWindow: 200, inputTokens: 101, accepted: true },
263+
{
264+
capabilities: { supportsImageToText: false },
265+
liveWindow: 200,
266+
curatedWindow: 100,
267+
inputTokens: 101,
268+
accepted: true,
269+
},
270+
])("checks measured input against both windows only for image models: %j", async (fixture) => {
271+
await handler.initializeClient()
272+
const userRole = vscode.LanguageModelChatMessage.User("").role
273+
const sendRequest = vi.fn<vscode.LanguageModelChat["sendRequest"]>().mockResolvedValue({
274+
stream: asyncStreamFrom([new vscode.LanguageModelTextPart("done")]),
275+
text: asyncStreamFrom(["done"]),
276+
})
277+
handler["client"] = {
278+
...mockLanguageModelChat,
279+
...{ capabilities: fixture.capabilities },
280+
maxInputTokens: fixture.liveWindow,
281+
sendRequest,
282+
// Image models count whole messages; text-only models count extracted strings.
283+
countTokens: vi.fn<vscode.LanguageModelChat["countTokens"]>().mockImplementation(async (content) => {
284+
if (typeof content === "string") {
285+
return fixture.capabilities === imageCapable ? 0 : fixture.inputTokens
286+
}
287+
return content.role === userRole ? fixture.inputTokens : 0
288+
}),
289+
}
290+
const windowSpy = vi.spyOn(handler, "getCondenseContextWindow").mockReturnValue(fixture.curatedWindow)
291+
try {
292+
const request = collectStream(handler.createMessage("system", [{ role: "user", content: [image] }]))
293+
if (fixture.capabilities !== imageCapable) {
294+
await expect(request).resolves.toContainEqual({ type: "text", text: "done" })
295+
expect(sendRequest).toHaveBeenCalledTimes(1)
296+
} else if (fixture.accepted) {
297+
await expect(request).resolves.toContainEqual({
298+
type: "usage",
299+
inputTokens: fixture.inputTokens,
300+
outputTokens: 0,
301+
})
302+
expect(sendRequest).toHaveBeenCalledTimes(1)
303+
} else {
304+
await expect(request).rejects.toSatisfy(checkContextWindowExceededError)
305+
expect(sendRequest).not.toHaveBeenCalled()
306+
}
307+
} finally {
308+
windowSpy.mockRestore()
309+
}
310+
})
311+
312+
it("does not advertise image input before a model has been selected", async () => {
313+
await handler.initializeClient()
314+
handler["client"] = null
315+
expect(handler.getModel().info.supportsImages).toBeUndefined()
316+
})
317+
318+
it("disables image support on hosts without image parts, even when the model reports support", async () => {
319+
await handler.initializeClient()
320+
handler["client"] = { ...mockLanguageModelChat, ...{ capabilities: { imageInput: true } } }
321+
const dataPart: { image: unknown } = Reflect.get(vscode, "LanguageModelDataPart")
322+
const original = dataPart.image
323+
try {
324+
dataPart.image = undefined
325+
expect(handler.getModel().info.supportsImages).toBe(false)
326+
} finally {
327+
dataPart.image = original
328+
}
329+
})
330+
})
331+
144332
describe("constructor", () => {
145333
it("should initialize with provided options", () => {
146334
expect(handler).toBeDefined()
@@ -4537,11 +4725,17 @@ describe("getVsCodeLmModels", () => {
45374725

45384726
beforeEach(() => vi.clearAllMocks())
45394727

4540-
it("returns the host's models as they are, with no added capability data", async () => {
4728+
it("returns serializable model metadata with host-reported image support", async () => {
45414729
const model = liveModel({ capabilities: { supportsImageToText: true } })
4542-
vi.mocked(vscode.lm.selectChatModels).mockResolvedValue([model] as never)
4543-
4544-
await expect(getVsCodeLmModels()).resolves.toEqual([model])
4730+
vi.mocked(vscode.lm.selectChatModels).mockResolvedValue([model])
4731+
4732+
const models = await getVsCodeLmModels()
4733+
expect(models[0]).toMatchObject({ id: model.id, family: model.family, modelInfo: { supportsImages: true } })
4734+
expect(models[0]).not.toHaveProperty("sendRequest")
4735+
expect(models[0]).not.toHaveProperty("countTokens")
4736+
expect(JSON.parse(JSON.stringify(models))).toMatchObject([
4737+
{ id: model.id, modelInfo: { supportsImages: true } },
4738+
])
45454739
})
45464740

45474741
it("asks the host for every model, whatever the vendor", async () => {

‎src/api/providers/vscode-lm.ts‎

Lines changed: 39 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import {
2727
import { CONTEXT_WINDOW_EXCEEDED_STATUS } from "../../core/context/context-management/context-error-handling"
2828

2929
import { BaseProvider } from "./base-provider"
30+
import { getVsCodeLmModelInfo } from "./vscode-lm-capabilities"
3031
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index"
3132

3233
/**
@@ -409,18 +410,8 @@ function declaredParamType(schema: Record<string, unknown> | undefined, paramNam
409410
* details), and only when the request would otherwise exceed the budget.
410411
*/
411412

412-
/**
413-
* Conservative characters-per-token ratio used to turn a token window into a character budget.
414-
*
415-
* `client.countTokens` is the model's real tokenizer, but it counts only a string: it cannot price
416-
* the tool schemas, image placeholders, or per-message framing the backend adds, so it cannot give
417-
* the true total for the request we are about to send. It is also an async, per-call RPC, and the
418-
* budget is needed for every message on every turn. We therefore keep a character estimate here
419-
* and stay deliberately conservative — 3 chars/token rather than the ~4 typical of English —
420-
* because the token-dense JSON, logs, and code that dominate oversized tool results tokenize to
421-
* fewer characters per token than prose. Under-counting biases toward trimming too early, which is
422-
* recoverable; over-counting sends an over-window request, which is not.
423-
*/
413+
// Text trimming uses 3 chars/token rather than ~4 for token-dense code and logs.
414+
// Native image input is counted separately by the host before sending; this budget is text-only.
424415
const VSCODE_LM_BUDGET_CHARS_PER_TOKEN = 3
425416

426417
/**
@@ -433,11 +424,7 @@ const VSCODE_LM_INPUT_BUDGET_FRACTION = 0.8
433424
/** A tool_result is never shrunk below this many characters, so a truncated result stays useful. */
434425
const MIN_TOOL_RESULT_CHARS = 2000
435426

436-
/**
437-
* Length charged for an image block. VS Code LM cannot carry image data, so
438-
* `convertToVsCodeLmMessages` replaces each image with a sentence-long textual placeholder; this
439-
* is that placeholder's approximate length.
440-
*/
427+
/** Approximate placeholder length; the text-only budget does not measure native image tokens. */
441428
const IMAGE_PLACEHOLDER_CHARS = 64
442429

443430
function readToolResultText(block: Anthropic.Messages.ContentBlockParam): string | undefined {
@@ -539,8 +526,6 @@ function estimateContentChars(content: Anthropic.Messages.MessageParam["content"
539526
} else if (type === "tool_use") {
540527
total += JSON.stringify((block as Anthropic.Messages.ToolUseBlockParam).input ?? {}).length
541528
} else if (type === "image") {
542-
// VS Code LM cannot send image data; convertToVsCodeLmMessages substitutes a textual
543-
// placeholder, so charge that placeholder's real length rather than a token-sized guess.
544529
total += IMAGE_PLACEHOLDER_CHARS
545530
}
546531
}
@@ -902,7 +887,7 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan
902887
typeof client.maxInputTokens === "number"
903888
? Math.max(0, client.maxInputTokens)
904889
: openAiModelInfoSaneDefaults.contextWindow,
905-
supportsImages: false, // VSCode Language Model API historically had no image input
890+
supportsImages: getVsCodeLmModelInfo(client).supportsImages,
906891
supportsPromptCache: true,
907892
inputPrice: 0,
908893
outputPrice: 0,
@@ -912,25 +897,38 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan
912897

913898
/** Model info reported before any client exists. */
914899
protected getFallbackModelInfo(): ModelInfo {
915-
return openAiModelInfoSaneDefaults
900+
return { ...openAiModelInfoSaneDefaults, supportsImages: undefined }
916901
}
917902

918903
/** Whether image parts are sent to the model. */
919-
protected supportsImageInput(_info: ModelInfo): boolean {
920-
return false
904+
protected supportsImageInput(info: ModelInfo): boolean {
905+
return info.supportsImages === true
921906
}
922907

923908
/** Rejects a request the selected model cannot serve before anything is sent. */
924909
protected assertRequestSupported(_messages: Anthropic.Messages.MessageParam[], _info: ModelInfo): void {}
925910

926911
/** Rejects a request whose measured input exceeds what the model accepts. */
927-
protected assertWithinContextWindow(_inputTokens: number, _info: ModelInfo): void {}
912+
protected assertWithinContextWindow(inputTokens: number, info: ModelInfo): void {
913+
// Native images are absent from the character budget; use the host's measured input too.
914+
if (
915+
this.supportsImageInput(info) &&
916+
inputTokens > Math.min(this.getCondenseContextWindow(), info.contextWindow)
917+
) {
918+
throw Object.assign(
919+
new Error(
920+
"The request exceeds the selected VS Code model's input context window. Condense the conversation or start a new task.",
921+
),
922+
{ status: CONTEXT_WINDOW_EXCEEDED_STATUS },
923+
)
924+
}
925+
}
928926

929927
/** The form in which a chat message is handed to the host's token counter. */
930928
protected prepareMessageForCounting(
931929
message: vscode.LanguageModelChatMessage,
932930
): string | vscode.LanguageModelChatMessage {
933-
return extractTextCountFromMessage(message)
931+
return this.supportsImageInput(this.getModel().info) ? message : extractTextCountFromMessage(message)
934932
}
935933

936934
/**
@@ -969,14 +967,17 @@ export class VsCodeLmHandler extends BaseProvider implements SingleCompletionHan
969967
* @returns A promise resolving to the token count
970968
*/
971969
override async countTokens(content: Array<Anthropic.Messages.ContentBlockParam>): Promise<number> {
972-
// Convert Anthropic content blocks to a string for VSCode LM token counting
970+
if (this.supportsImageInput(this.getModel().info) && containsImageBlock(content)) {
971+
return this.calculateTotalInputTokens(convertToVsCodeLmMessages([{ role: "user", content }], true))
972+
}
973+
974+
// Preserve text-only counting for models that cannot send native images.
973975
let textContent = ""
974976

975977
for (const block of content) {
976978
if (block.type === "text") {
977979
textContent += block.text || ""
978980
} else if (block.type === "image") {
979-
// VSCode LM doesn't support images directly, so we'll just use a placeholder
980981
textContent += "[IMAGE]"
981982
}
982983
}
@@ -1597,7 +1598,17 @@ export const VSCODE_LM_STATIC_BLACKLIST: string[] = ["claude-3.7-sonnet", "claud
15971598
export async function getVsCodeLmModels() {
15981599
try {
15991600
const models = (await vscode.lm.selectChatModels({})) || []
1600-
return models.filter((model) => !VSCODE_LM_STATIC_BLACKLIST.includes(model.id))
1601+
return models
1602+
.filter((model) => !VSCODE_LM_STATIC_BLACKLIST.includes(model.id))
1603+
.map((model) => ({
1604+
id: model.id,
1605+
vendor: model.vendor,
1606+
family: model.family,
1607+
version: model.version,
1608+
name: model.name,
1609+
maxInputTokens: model.maxInputTokens,
1610+
modelInfo: getVsCodeLmModelInfo(model),
1611+
}))
16011612
} catch (error) {
16021613
console.error(
16031614
`Error fetching VS Code LM models: ${JSON.stringify(error, Object.getOwnPropertyNames(error), 2)}`,

0 commit comments

Comments
 (0)