Gemini BYOK: Handle invalid key (microsoft#2733)

* Handle invalid API Key

* update test

---------

Co-authored-by: vijay upadya <vj@example.com>

Author: vijayupadya

src/extension/byok/vscode-node/geminiNativeProvider.ts

@@ -33,6 +33,20 @@ export class GeminiNativeBYOKLMProvider implements BYOKModelProvider<LanguageMod
 		@IRequestLogger private readonly _requestLogger: IRequestLogger
 	) { }
 
+	private _isInvalidApiKeyError(error: unknown): boolean {
+		if (!error) {
+			return false;
+		}
+
+		const message = typeof error === 'string' ? error : (error as Error).message;
+		if (typeof message !== 'string') {
+			return false;
+		}
+
+		const lower = message.toLowerCase();
+		return lower.includes('api key not valid') || lower.includes('api_key_invalid') || lower.includes('api key invalid');
+	}
+
 	private async _getOrReadApiKey(): Promise<string | undefined> {
 		if (!this._apiKey) {
 			this._apiKey = await this._byokStorageService.getAPIKey(GeminiNativeBYOKLMProvider.providerName);
@@ -107,7 +121,20 @@ export class GeminiNativeBYOKLMProvider implements BYOKModelProvider<LanguageMod
 					return [];
 				}
 			}
-		} catch {
+		} catch (error) {
+			if (this._isInvalidApiKeyError(error)) {
+				if (options.silent) {
+					return [];
+				}
+				await this.updateAPIKey();
+				if (this._apiKey) {
+					try {
+						return byokKnownModelsToAPIInfo(GeminiNativeBYOKLMProvider.providerName, await this.getAllModels(this._apiKey));
+					} catch (retryError) {
+						this._logService.error(`Error after re-prompting for API key: ${toErrorMessage(retryError, true)}`);
+					}
+				}
+			}
 			return [];
 		}
 	}

src/extension/byok/vscode-node/test/geminiNativeProvider.spec.ts

@@ -217,4 +217,76 @@ describe('GeminiNativeBYOKLMProvider', () => {
 			tokenSource.token
 		)).rejects.toThrow(/No API key configured/i);
 	});
+
+	it('prompts for a new API key when listing models fails with an invalid key', async () => {
+		const { GeminiNativeBYOKLMProvider } = await import('../geminiNativeProvider');
+		const genai = await import('@google/genai');
+		const MockGoogleGenAI = genai.GoogleGenAI as unknown as { listModelsResult: AsyncIterable<any> };
+		// Simulate the models.list() call throwing an invalid API key error when iterated
+		MockGoogleGenAI.listModelsResult = (async function* () {
+			throw new Error('ApiError: {"error":{"message":"API key not valid. Please pass a valid API key.","details":[{"reason":"API_KEY_INVALID"}]}}');
+		})();
+
+		const storage = createStorageService({
+			getAPIKey: vi.fn().mockResolvedValue('bad_key'),
+		});
+
+		mockHandleAPIKeyUpdate.mockResolvedValue({ apiKey: undefined, deleted: false, cancelled: true });
+
+		const provider = new GeminiNativeBYOKLMProvider(undefined, storage, new TestLogService(), createRequestLogger());
+		const tokenSource = new vscode.CancellationTokenSource();
+		const models = await provider.provideLanguageModelChatInformation({ silent: false }, tokenSource.token);
+
+		// When the key is invalid, we should re-prompt for a new one
+		// and handle the failure gracefully by returning an empty list.
+		expect(models).toEqual([]);
+		expect(mockHandleAPIKeyUpdate).toHaveBeenCalled();
+	});
+
+	it('retries listing models after re-prompting with a valid API key', async () => {
+		const { GeminiNativeBYOKLMProvider } = await import('../geminiNativeProvider');
+		const genai = await import('@google/genai');
+		const MockGoogleGenAI = genai.GoogleGenAI as unknown as { listModelsResult: AsyncIterable<any> };
+
+		let iterationCount = 0;
+		let hasThrown = false;
+		const modelId = 'test-model';
+
+		MockGoogleGenAI.listModelsResult = {
+			async *[Symbol.asyncIterator]() {
+				iterationCount++;
+				if (!hasThrown) {
+					hasThrown = true;
+					throw new Error('ApiError: {"error":{"message":"API key not valid. Please pass a valid API key.","details":[{"reason":"API_KEY_INVALID"}]}}');
+				}
+				yield { name: modelId };
+			}
+		};
+
+		const storage = createStorageService({
+			getAPIKey: vi.fn().mockResolvedValue('bad_key'),
+		});
+
+		mockHandleAPIKeyUpdate.mockResolvedValue({ apiKey: 'k_new', deleted: false, cancelled: false });
+
+		const knownModels = {
+			[modelId]: {
+				name: 'Test Model',
+				maxInputTokens: 1000,
+				maxOutputTokens: 1000,
+				toolCalling: false,
+				vision: false
+			}
+		};
+
+		const provider = new GeminiNativeBYOKLMProvider(knownModels, storage, new TestLogService(), createRequestLogger());
+		const tokenSource = new vscode.CancellationTokenSource();
+		const models = await provider.provideLanguageModelChatInformation({ silent: false }, tokenSource.token);
+
+		// First attempt should fail with invalid key, then after re-prompting
+		// we should retry listing models and succeed with the new key.
+		expect(models.map(m => m.id)).toEqual([modelId]);
+		expect(iterationCount).toBe(2);
+		expect(mockHandleAPIKeyUpdate).toHaveBeenCalled();
+	});
 });