Fixes #3344 exposes reset on ai model picker

Renames "Reset Stored AI Key" command & adds confirmation

Author: eamodio

CHANGELOG.md

@@ -6,6 +6,14 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/) and this p
 
 ## [Unreleased]
 
+### Changed
+
+- Renames `Reset Stored AI Key` command to `Reset Stored AI Keys...` and adds confirmation prompt with options to reset only the current or all AI keys
+
+### Fixed
+
+- Fixes [#3344](https://github.com/gitkraken/vscode-gitlens/issues/3344) - Make changing the AI key easier
+
 ## [15.1.0] - 2024-06-05
 
 ### Added

package.json

@@ -5429,7 +5429,7 @@
 			},
 			{
 				"command": "gitlens.resetAIKey",
-				"title": "Reset Stored AI Key",
+				"title": "Reset Stored AI Keys...",
 				"category": "GitLens"
 			},
 			{

src/ai/aiProviderService.ts

@@ -77,7 +77,7 @@ export class AIProviderService implements Disposable {
 		this._provider?.dispose();
 	}
 
-	get providerId() {
+	get currentProviderId() {
 		return this._provider?.id;
 	}
 
@@ -108,13 +108,15 @@ export class AIProviderService implements Disposable {
 		return models.flatMap(m => getSettledValue(m, []));
 	}
 
-	private async getOrChooseModel(force?: boolean): Promise<AIModel | undefined> {
+	private async getModel(options?: { force?: boolean; silent?: boolean }): Promise<AIModel | undefined> {
 		const cfg = this.getConfiguredModel();
-		if (!force && cfg?.provider != null && cfg?.model != null) {
+		if (!options?.force && cfg?.provider != null && cfg?.model != null) {
 			const model = await this.getOrUpdateModel(cfg.provider, cfg.model);
 			if (model != null) return model;
 		}
 
+		if (options?.silent) return undefined;
+
 		const pick = await showAIModelPicker(this.container, cfg);
 		if (pick == null) return undefined;
 
@@ -221,7 +223,7 @@ export class AIProviderService implements Disposable {
 			changes = diff.contents;
 		}
 
-		const model = await this.getOrChooseModel();
+		const model = await this.getModel();
 		if (model == null) return undefined;
 
 		const provider = this._provider!;
@@ -276,7 +278,7 @@ export class AIProviderService implements Disposable {
 		const diff = await this.container.git.getDiff(commit.repoPath, commit.sha);
 		if (!diff?.contents) throw new Error('No changes found to explain.');
 
-		const model = await this.getOrChooseModel();
+		const model = await this.getModel();
 		if (model == null) return undefined;
 
 		const provider = this._provider!;
@@ -301,22 +303,64 @@ export class AIProviderService implements Disposable {
 		});
 	}
 
-	reset() {
-		const { providerId } = this;
-		if (providerId == null) return;
+	async reset() {
+		let { _provider: provider } = this;
+		if (provider == null) {
+			// If we have no provider, try to get the current model (which will load the provider)
+			await this.getModel({ silent: true });
+			provider = this._provider;
+		}
+
+		const resetCurrent: MessageItem = { title: `Reset Current` };
+		const resetAll: MessageItem = { title: 'Reset All' };
+		const cancel: MessageItem = { title: 'Cancel', isCloseAffordance: true };
+
+		let result;
+		if (provider == null) {
+			result = await window.showInformationMessage(
+				`Do you want to reset all of the stored AI keys?`,
+				{ modal: true },
+				resetAll,
+				cancel,
+			);
+		} else {
+			result = await window.showInformationMessage(
+				`Do you want to reset the stored key for the current provider (${provider.name}) or reset all of the stored AI keys?`,
+				{ modal: true },
+				resetCurrent,
+				resetAll,
+				cancel,
+			);
+		}
+
+		if (provider != null && result === resetCurrent) {
+			void env.clipboard.writeText((await this.container.storage.getSecret(`gitlens.${provider.id}.key`)) ?? '');
+			void this.container.storage.deleteSecret(`gitlens.${provider.id}.key`);
+
+			void this.container.storage.delete(`confirm:ai:tos:${provider.id}`);
+			void this.container.storage.deleteWorkspace(`confirm:ai:tos:${provider.id}`);
+		} else if (result === resetAll) {
+			const keys = [];
+			for (const [providerId] of _supportedProviderTypes) {
+				keys.push(await this.container.storage.getSecret(`gitlens.${providerId}.key`));
+			}
+			void env.clipboard.writeText(keys.join('\n'));
 
-		void this.container.storage.deleteSecret(`gitlens.${providerId}.key`);
+			for (const [providerId] of _supportedProviderTypes) {
+				void this.container.storage.deleteSecret(`gitlens.${providerId}.key`);
+			}
 
-		void this.container.storage.delete(`confirm:ai:tos:${providerId}`);
-		void this.container.storage.deleteWorkspace(`confirm:ai:tos:${providerId}`);
+			void this.container.storage.deleteWithPrefix(`confirm:ai:tos`);
+			void this.container.storage.deleteWorkspaceWithPrefix(`confirm:ai:tos`);
+		}
 	}
 
 	supports(provider: AIProviders | string) {
 		return _supportedProviderTypes.has(provider as AIProviders);
 	}
 
 	async switchModel() {
-		void (await this.getOrChooseModel(true));
+		void (await this.getModel({ force: true }));
 	}
 }
 

src/commands/resets.ts

@@ -13,7 +13,7 @@ export class ResetAIKeyCommand extends Command {
 	}
 
 	async execute() {
-		(await this.container.ai)?.reset();
+		await (await this.container.ai)?.reset();
 	}
 }
 

src/quickpicks/aiModelPicker.ts

@@ -1,8 +1,11 @@
-import type { QuickPickItem } from 'vscode';
+import type { Disposable, QuickInputButton, QuickPickItem } from 'vscode';
 import { QuickPickItemKind, ThemeIcon, window } from 'vscode';
 import type { AIModel } from '../ai/aiProviderService';
 import type { AIModels, AIProviders } from '../constants';
+import { Commands } from '../constants';
 import type { Container } from '../container';
+import { executeCommand } from '../system/command';
+import { getQuickPickIgnoreFocusOut } from '../system/utils';
 
 export interface ModelQuickPickItem extends QuickPickItem {
 	model: AIModel;
@@ -14,16 +17,15 @@ export async function showAIModelPicker(
 ): Promise<ModelQuickPickItem | undefined> {
 	const models = (await (await container.ai)?.getModels()) ?? [];
 
-	type QuickPickSeparator = { label: string; kind: QuickPickItemKind.Separator };
-	const items: (ModelQuickPickItem | QuickPickSeparator)[] = [];
+	const items: ModelQuickPickItem[] = [];
 
 	let lastProvider: AIProviders | undefined;
 	for (const m of models) {
 		if (m.hidden) continue;
 
 		if (lastProvider !== m.provider.id) {
 			lastProvider = m.provider.id;
-			items.push({ label: m.provider.name, kind: QuickPickItemKind.Separator });
+			items.push({ label: m.provider.name, kind: QuickPickItemKind.Separator } as unknown as ModelQuickPickItem);
 		}
 
 		const picked = m.provider.id === current?.provider && m.id === current?.model;
@@ -34,14 +36,48 @@ export async function showAIModelPicker(
 			// description: m.provider.name,
 			model: m,
 			picked: picked,
-		});
+		} satisfies ModelQuickPickItem);
 	}
 
-	const pick = (await window.showQuickPick(items, {
-		title: 'Choose AI Model',
-		placeHolder: 'Select an AI model to use for experimental AI features',
-		matchOnDescription: true,
-	})) as ModelQuickPickItem | undefined;
+	const quickpick = window.createQuickPick<ModelQuickPickItem>();
+	quickpick.ignoreFocusOut = getQuickPickIgnoreFocusOut();
+
+	const disposables: Disposable[] = [];
+
+	const ResetAIKeyButton: QuickInputButton = {
+		iconPath: new ThemeIcon('clear-all'),
+		tooltip: 'Reset AI Keys...',
+	};
+
+	try {
+		const pick = await new Promise<ModelQuickPickItem | undefined>(resolve => {
+			disposables.push(
+				quickpick.onDidHide(() => resolve(undefined)),
+				quickpick.onDidAccept(() => {
+					if (quickpick.activeItems.length !== 0) {
+						resolve(quickpick.activeItems[0]);
+					}
+				}),
+				quickpick.onDidTriggerButton(e => {
+					if (e === ResetAIKeyButton) {
+						void executeCommand(Commands.ResetAIKey);
+					}
+				}),
+			);
 
-	return pick;
+			quickpick.title = 'Choose AI Model';
+			quickpick.placeholder = 'Select an AI model to use for experimental AI features';
+			quickpick.matchOnDescription = true;
+			quickpick.matchOnDetail = true;
+			quickpick.buttons = [ResetAIKeyButton];
+			quickpick.items = items;
+
+			quickpick.show();
+		});
+
+		return pick;
+	} finally {
+		quickpick.dispose();
+		disposables.forEach(d => void d.dispose());
+	}
 }