Merge pull request #328785 from microsoft/agents/gemini-provider-issue-investigation

Implement end-to-end BYOK image support and regression tests

src/vs/platform/agentHost/common/agentHostByokLm.ts

@@ -25,10 +25,20 @@ export interface IByokLmTextPart {
 	readonly text: string;
 }
 
+export type ByokLmImageMimeType = 'image/png' | 'image/jpeg' | 'image/gif' | 'image/webp' | 'image/bmp';
+
+export interface IByokLmImagePart {
+	readonly type: 'image';
+	readonly mimeType: ByokLmImageMimeType;
+	readonly data: string;
+}
+
+export type IByokLmContentPart = IByokLmTextPart | IByokLmImagePart;
+
 export interface IByokLmMessageItem {
 	readonly type: 'message';
 	readonly role: 'system' | 'developer' | 'user' | 'assistant';
-	readonly content: IByokLmTextPart[];
+	readonly content: IByokLmContentPart[];
 }
 
 export interface IByokLmReasoningItem {

src/vs/platform/agentHost/node/copilot/byokResponsesTranslation.ts

@@ -3,9 +3,12 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import { decodeBase64 } from '../../../../base/common/buffer.js';
 import {
+	ByokLmImageMimeType,
 	IByokLmChatRequest,
 	IByokLmChatResult,
+	IByokLmContentPart,
 	IByokLmInputItem,
 	IByokLmOutputItem,
 	IByokLmTool,
@@ -14,6 +17,7 @@ import {
 interface IResponsesContentPart {
 	readonly type?: string;
 	readonly text?: string;
+	readonly image_url?: string;
 }
 
 interface IResponsesSummaryPart {
@@ -21,6 +25,19 @@ interface IResponsesSummaryPart {
 	readonly text?: string;
 }
 
+function isSupportedImageMimeType(mimeType: string): mimeType is ByokLmImageMimeType {
+	switch (mimeType) {
+		case 'image/png':
+		case 'image/jpeg':
+		case 'image/gif':
+		case 'image/webp':
+		case 'image/bmp':
+			return true;
+		default:
+			return false;
+	}
+}
+
 interface IResponsesInputItem {
 	readonly type?: string;
 	readonly role?: string;
@@ -71,7 +88,7 @@ function toBridgeRole(role: string | undefined): 'system' | 'developer' | 'user'
 	}
 }
 
-function toTextParts(content: string | IResponsesContentPart[] | undefined, itemIndex: number): Array<{ type: 'text'; text: string }> {
+function toContentParts(content: string | IResponsesContentPart[] | undefined, itemIndex: number): IByokLmContentPart[] {
 	if (typeof content === 'string') {
 		return content ? [{ type: 'text', text: content }] : [];
 	}
@@ -82,6 +99,25 @@ function toTextParts(content: string | IResponsesContentPart[] | undefined, item
 		if ((part.type === 'input_text' || part.type === 'output_text' || part.type === 'text') && typeof part.text === 'string') {
 			return { type: 'text' as const, text: part.text };
 		}
+		if (part.type === 'input_image' && typeof part.image_url === 'string') {
+			const match = /^data:(?<mimeType>image\/[^;,]+)(?:;[^,]*)?;base64,(?<data>.*)$/.exec(part.image_url);
+			if (match?.groups) {
+				if (!isSupportedImageMimeType(match.groups.mimeType)) {
+					throw new ResponsesTranslationError(`Unsupported input[${itemIndex}].content[${contentIndex}].image_url MIME type '${match.groups.mimeType}'`);
+				}
+				try {
+					decodeBase64(match.groups.data);
+				} catch {
+					throw new ResponsesTranslationError(`Invalid input[${itemIndex}].content[${contentIndex}].image_url`);
+				}
+				return {
+					type: 'image' as const,
+					mimeType: match.groups.mimeType,
+					data: match.groups.data,
+				};
+			}
+			throw new ResponsesTranslationError(`Unsupported input[${itemIndex}].content[${contentIndex}].image_url`);
+		}
 		throw new ResponsesTranslationError(`Unsupported input[${itemIndex}].content[${contentIndex}] type '${part.type ?? ''}'`);
 	});
 }
@@ -99,7 +135,7 @@ function toBridgeInputItem(item: IResponsesInputItem, index: number): IByokLmInp
 			return {
 				type: 'message',
 				role: toBridgeRole(item.role),
-				content: toTextParts(item.content, index),
+				content: toContentParts(item.content, index),
 			};
 		case 'reasoning':
 			return {

src/vs/platform/agentHost/test/node/byokLmProxyService.test.ts

@@ -138,6 +138,127 @@ suite('ByokLmProxyService', () => {
 		assert.deepStrictEqual(captured?.input, [{ type: 'message', role: 'user', content: [{ type: 'text', text: 'hi' }] }]);
 	});
 
+	test('forwards image input on the initial and subsequent turns', async () => {
+		const captured: IByokLmChatRequest[] = [];
+		const statuses: number[] = [];
+		const imageMessage = {
+			type: 'message',
+			role: 'user',
+			content: [
+				{ type: 'input_text', text: 'What is in this image?' },
+				{ type: 'input_image', image_url: 'data:image/png;base64,iVBORw0KGgo=' },
+			],
+		};
+
+		await withProxy(
+			async request => {
+				captured.push(request);
+				return { output: [] };
+			},
+			async handle => {
+				for (const input of [
+					[imageMessage],
+					[imageMessage, { type: 'message', role: 'user', content: [{ type: 'input_text', text: 'Try again without a new image.' }] }],
+				]) {
+					const response = await fetch(responsesUrl(handle, 'gemini'), {
+						method: 'POST',
+						headers: authHeaders(handle),
+						body: JSON.stringify({ model: 'gemini-3.6-flash', input }),
+					});
+					statuses.push(response.status);
+					await response.text();
+				}
+			},
+		);
+
+		assert.deepStrictEqual({ statuses, input: captured.map(request => request.input) }, {
+			statuses: [200, 200],
+			input: [
+				[
+					{
+						type: 'message',
+						role: 'user',
+						content: [
+							{ type: 'text', text: 'What is in this image?' },
+							{ type: 'image', mimeType: 'image/png', data: 'iVBORw0KGgo=' },
+						],
+					},
+				],
+				[
+					{
+						type: 'message',
+						role: 'user',
+						content: [
+							{ type: 'text', text: 'What is in this image?' },
+							{ type: 'image', mimeType: 'image/png', data: 'iVBORw0KGgo=' },
+						],
+					},
+					{
+						type: 'message',
+						role: 'user',
+						content: [
+							{ type: 'text', text: 'Try again without a new image.' },
+						],
+					},
+				],
+			],
+		});
+	});
+
+	test('rejects image URLs that cannot be forwarded as inline data', async () => {
+		await withProxy(
+			async () => ({ output: [] }),
+			async handle => {
+				const responses: Array<{ status: number; body: unknown }> = [];
+				for (const imageUrl of ['https://example.com/image.png', 'data:image/svg+xml;base64,PHN2Zz4=', 'data:image/png;base64,not valid']) {
+					const response = await fetch(responsesUrl(handle, 'gemini'), {
+						method: 'POST',
+						headers: authHeaders(handle),
+						body: JSON.stringify({
+							model: 'gemini-3.6-flash',
+							input: [{
+								type: 'message',
+								role: 'user',
+								content: [{ type: 'input_image', image_url: imageUrl }],
+							}],
+						}),
+					});
+					responses.push({ status: response.status, body: await response.json() });
+				}
+
+				assert.deepStrictEqual(responses, [
+					{
+						status: 400,
+						body: {
+							error: {
+								message: 'Unsupported input[0].content[0].image_url',
+								type: 'invalid_request_error',
+							},
+						},
+					},
+					{
+						status: 400,
+						body: {
+							error: {
+								message: 'Unsupported input[0].content[0].image_url MIME type \'image/svg+xml\'',
+								type: 'invalid_request_error',
+							},
+						},
+					},
+					{
+						status: 400,
+						body: {
+							error: {
+								message: 'Invalid input[0].content[0].image_url',
+								type: 'invalid_request_error',
+							},
+						},
+					},
+				]);
+			},
+		);
+	});
+
 	test('forwards custom tool call history with freeform input', async () => {
 		let captured: IByokLmChatRequest | undefined;
 		await withProxy(

src/vs/workbench/contrib/chat/browser/agentSessions/agentHost/agentHostByokLmHandler.ts

@@ -6,11 +6,13 @@
 import { CancellationToken } from '../../../../../../base/common/cancellation.js';
 import { Emitter, Event } from '../../../../../../base/common/event.js';
 import { Disposable } from '../../../../../../base/common/lifecycle.js';
-import { VSBuffer } from '../../../../../../base/common/buffer.js';
+import { decodeBase64, VSBuffer } from '../../../../../../base/common/buffer.js';
 import {
+	ByokLmImageMimeType,
 	IAgentHostByokLmHandler,
 	IByokLmChatRequest,
 	IByokLmChatResult,
+	IByokLmContentPart,
 	IByokLmInputItem,
 	IByokLmModelInfo,
 	IByokLmOutputItem,
@@ -18,6 +20,7 @@ import {
 } from '../../../../../../platform/agentHost/common/agentHostByokLm.js';
 import { ILogService } from '../../../../../../platform/log/common/log.js';
 import {
+	ChatImageMimeType,
 	ChatMessageRole,
 	IChatMessage,
 	IChatMessagePart,
@@ -207,7 +210,7 @@ export class AgentHostByokLmHandler extends Disposable implements IAgentHostByok
 			case 'message':
 				return {
 					role: this._toChatRole(item.role),
-					content: [{ type: 'text', value: item.content.map(part => part.text).join('') }],
+					content: this._toChatMessageParts(item.content),
 				};
 			case 'reasoning': {
 				return {
@@ -256,6 +259,38 @@ export class AgentHostByokLmHandler extends Disposable implements IAgentHostByok
 		}
 	}
 
+	private _toChatMessageParts(parts: IByokLmContentPart[]): IChatMessagePart[] {
+		const result: IChatMessagePart[] = [];
+		for (const part of parts) {
+			if (part.type === 'text') {
+				const previous = result.at(-1);
+				if (previous?.type === 'text') {
+					previous.value += part.text;
+				} else {
+					result.push({ type: 'text', value: part.text });
+				}
+			} else {
+				result.push({ type: 'image_url', value: { mimeType: this._toChatImageMimeType(part.mimeType), data: decodeBase64(part.data) } });
+			}
+		}
+		return result.length ? result : [{ type: 'text', value: '' }];
+	}
+
+	private _toChatImageMimeType(mimeType: ByokLmImageMimeType): ChatImageMimeType {
+		switch (mimeType) {
+			case 'image/png':
+				return ChatImageMimeType.PNG;
+			case 'image/jpeg':
+				return ChatImageMimeType.JPEG;
+			case 'image/gif':
+				return ChatImageMimeType.GIF;
+			case 'image/webp':
+				return ChatImageMimeType.WEBP;
+			case 'image/bmp':
+				return ChatImageMimeType.BMP;
+		}
+	}
+
 	private _appendTextOutput(output: IByokLmOutputItem[], value: string): void {
 		const previous = output.at(-1);
 		if (previous?.type === 'message') {

src/vs/workbench/contrib/chat/test/browser/agentSessions/agentHostByokLmHandler.test.ts

@@ -251,7 +251,14 @@ suite('AgentHostByokLmHandler', () => {
 					{ type: 'custom_tool_call', callId: 't2', name: 'apply_patch', input: 'patch' },
 					{ type: 'function_call_output', callId: 't1', output: 'sunny' },
 					{ type: 'custom_tool_call_output', callId: 't2', output: 'Done!' },
-					{ type: 'message', role: 'user', content: [{ type: 'text', text: 'hi' }] },
+					{
+						type: 'message',
+						role: 'user',
+						content: [
+							{ type: 'text', text: 'hi' },
+							{ type: 'image', mimeType: 'image/png', data: 'aW1hZ2U=' },
+						],
+					},
 				],
 			},
 			CancellationToken.None,
@@ -280,7 +287,13 @@ suite('AgentHostByokLmHandler', () => {
 				},
 				{ role: ChatMessageRole.User, content: [{ type: 'tool_result', toolCallId: 't1', value: [{ type: 'text', value: 'sunny' }] }] },
 				{ role: ChatMessageRole.User, content: [{ type: 'tool_result', toolCallId: 't2', value: [{ type: 'text', value: 'Done!' }] }] },
-				{ role: ChatMessageRole.User, content: [{ type: 'text', value: 'hi' }] },
+				{
+					role: ChatMessageRole.User,
+					content: [
+						{ type: 'text', value: 'hi' },
+						{ type: 'image_url', value: { mimeType: 'image/png', data: VSBuffer.fromString('image') } },
+					],
+				},
 			],
 			options: {
 				modelOptions: { temperature: 0.5 },