Handle smart image routing in auto (#3504)

Author: lramos15

src/extension/conversation/vscode-node/languageModelAccess.ts

@@ -278,7 +278,7 @@ export class LanguageModelAccess extends Disposable implements IExtensionContrib
 				},
 				isUserSelectable: endpoint.showInModelPicker,
 				capabilities: {
-					imageInput: endpoint.supportsVision,
+					imageInput: endpoint instanceof AutoChatEndpoint ? true : endpoint.supportsVision,
 					toolCalling: endpoint.supportsToolCalls,
 				}
 			};

src/extension/prompts/node/panel/image.tsx

@@ -44,13 +44,14 @@ export class HistoricalImage extends PromptElement<HistoricalImageProps, unknown
 	constructor(
 		props: HistoricalImageProps,
 		@IPromptEndpoint private readonly promptEndpoint: IPromptEndpoint,
+		@IAuthenticationService private readonly authService: IAuthenticationService,
 	) {
 		super(props);
 	}
 
 	override async render(_state: unknown, sizing: PromptSizing) {
 		// If the model doesn't support vision, omit historical images
-		if (!this.promptEndpoint.supportsVision) {
+		if (!this.promptEndpoint.supportsVision || !this.authService.copilotToken?.isEditorPreviewFeaturesEnabled()) {
 			return undefined;
 		}
 
@@ -77,7 +78,7 @@ export class Image extends PromptElement<ImageProps, unknown> {
 		const fillerUri: Uri = this.props.reference ?? Uri.parse('Attached Image');
 
 		try {
-			if (!this.promptEndpoint.supportsVision) {
+			if (!this.promptEndpoint.supportsVision || !this.authService.copilotToken?.isEditorPreviewFeaturesEnabled()) {
 				if (this.props.omitReferences) {
 					return;
 				}

src/platform/endpoint/node/automodeService.ts

@@ -94,10 +94,11 @@ class AutoModeTokenBank extends Disposable {
 			: 'copilotchat.autoModelHint';
 
 		const autoModeHint = this._expService.getTreatmentVariable<string>(expName) || 'auto';
+		console.log(`AutoModeService: Using auto mode hint '${autoModeHint}' for location '${this._location}'.`);
 
 		const response = await this._capiClientService.makeRequest<Response>({
 			json: {
-				'auto_mode': { 'model_hints': [autoModeHint] }
+				'auto_mode': { 'model_hints': ['grok-code-fast-1'] }
 			},
 			headers,
 			method: 'POST'
@@ -229,6 +230,8 @@ export class AutomodeService extends Disposable implements IAutomodeService {
 				throw new Error(errorMsg);
 			}
 		}
+		selectedModel = this._applyVisionFallback(chatRequest, selectedModel, reserveToken.available_models, knownEndpoints);
+
 		const existingEndpoints = entry?.endpoints || [];
 		let autoEndpoint = existingEndpoints.find(e => e.model === selectedModel.model);
 		if (!autoEndpoint) {
@@ -259,7 +262,15 @@ export class AutomodeService extends Disposable implements IAutomodeService {
 				}
 				entry.endpoints = [this._instantiationService.createInstance(AutoChatEndpoint, newModel, entryToken.session_token, entryToken.discounted_costs?.[newModel.model] || 0, this._calculateDiscountRange(entryToken.discounted_costs))];
 			}
-			return entry.endpoints[0];
+			// Apply vision fallback even on cached entries, since the cached model may not support images
+			const cachedEndpoint = entry.endpoints[0];
+			const fallbackEndpoint = this._applyVisionFallback(chatRequest, cachedEndpoint, entryToken.available_models, knownEndpoints);
+			if (fallbackEndpoint !== cachedEndpoint) {
+				const autoEndpoint = this._instantiationService.createInstance(AutoChatEndpoint, fallbackEndpoint, entryToken.session_token, entryToken.discounted_costs?.[fallbackEndpoint.model] || 0, this._calculateDiscountRange(entryToken.discounted_costs));
+				entry.endpoints[0] = autoEndpoint;
+				return autoEndpoint;
+			}
+			return cachedEndpoint;
 		}
 
 		// No cached entry, use the reserve token
@@ -269,18 +280,38 @@ export class AutomodeService extends Disposable implements IAutomodeService {
 		reserveTokenBank.debugName = conversationId;
 
 		const reserveToken = await reserveTokenBank.getToken();
-		const selectedModel = knownEndpoints.find(e => e.model === reserveToken.selected_model);
+		let selectedModel = knownEndpoints.find(e => e.model === reserveToken.selected_model);
 		if (!selectedModel) {
 			const errorMsg = `Auto mode failed: selected model '${reserveToken.selected_model}' not found in known endpoints.`;
 			this._logService.error(errorMsg);
 			throw new Error(errorMsg);
 		}
+		selectedModel = this._applyVisionFallback(chatRequest, selectedModel, reserveToken.available_models, knownEndpoints);
 		const autoEndpoint = this._instantiationService.createInstance(AutoChatEndpoint, selectedModel, reserveToken.session_token, reserveToken.discounted_costs?.[selectedModel.model] || 0, this._calculateDiscountRange(reserveToken.discounted_costs));
 
 		this._autoModelCache.set(conversationId, { endpoints: [autoEndpoint], tokenBank: reserveTokenBank });
 		return autoEndpoint;
 	}
 
+	/**
+	 * If the request contains an image and the selected model doesn't support vision,
+	 * fall back to the first vision-capable model from the available models.
+	 */
+	private _applyVisionFallback(chatRequest: ChatRequest | undefined, selectedModel: IChatEndpoint, availableModels: string[], knownEndpoints: IChatEndpoint[]): IChatEndpoint {
+		if (!hasImage(chatRequest) || selectedModel.supportsVision) {
+			return selectedModel;
+		}
+		const visionModel = availableModels
+			.map(model => knownEndpoints.find(e => e.model === model))
+			.find(endpoint => endpoint?.supportsVision);
+		if (visionModel) {
+			this._logService.trace(`Selected model '${selectedModel.model}' does not support vision, falling back to '${visionModel.model}'.`);
+			return visionModel;
+		}
+		this._logService.warn(`Request contains an image but no vision-capable model is available.`);
+		return selectedModel;
+	}
+
 	private _calculateDiscountRange(discounts: Record<string, number> | undefined): { low: number; high: number } {
 		if (!discounts) {
 			return { low: 0, high: 0 };
@@ -311,5 +342,19 @@ function getConversationId(chatRequest: ChatRequest | undefined): string {
 	if (!chatRequest) {
 		return 'unknown';
 	}
-	return (chatRequest?.toolInvocationToken as { sessionId: string })?.sessionId || 'unknown';
+	return chatRequest?.sessionId || 'unknown';
+}
+
+function hasImage(chatRequest: ChatRequest | undefined): boolean {
+	if (!chatRequest || !chatRequest.references) {
+		return false;
+	}
+	return chatRequest.references.some(ref => {
+		const value = ref.value;
+		return typeof value === 'object' &&
+			value !== null &&
+			'mimeType' in value &&
+			typeof value.mimeType === 'string'
+			&& value.mimeType.startsWith('image/');
+	});
 }