Fix several AI settings search bugs (microsoft#250068)

Author: rzhao271

src/vs/workbench/contrib/preferences/browser/preferencesSearch.ts

@@ -131,7 +131,8 @@ export class LocalSearchProvider implements ISearchProvider {
 		const filteredMatches = filterMatches.filter(m => (m.matchType & topKeyMatchType) || (m.matchType & alwaysAllowedMatchTypes) || m.matchType === SettingMatchType.ExactMatch);
 		return Promise.resolve({
 			filterMatches: filteredMatches,
-			exactMatch: filteredMatches.some(m => m.matchType === SettingMatchType.ExactMatch)
+			exactMatch: filteredMatches.some(m => m.matchType === SettingMatchType.ExactMatch),
+			providerName: 'local'
 		});
 	}
 

src/vs/workbench/contrib/preferences/browser/settingsEditor2.ts

@@ -633,10 +633,11 @@ export class SettingsEditor2 extends EditorPane {
 		const showSuggestions = this.configurationService.getValue<boolean>('workbench.settings.showAISearchToggle');
 		if (!setupHidden && showSuggestions) {
 			const showAiResultActionClassNames = ['action-label', ThemeIcon.asClassName(preferencesAiResultsIcon)];
+			const searchServiceEnabled = this.aiSettingsSearchService.isEnabled();
 			this.showAiResultsAction = this._register(new Action(SETTINGS_EDITOR_COMMAND_SHOW_AI_RESULTS,
-				localize('showAiResultsDescription', "Search settings with AI"), showAiResultActionClassNames.join(' '), false
+				localize('showAiResultsDescription', "Search settings with AI"), showAiResultActionClassNames.join(' '), searchServiceEnabled
 			));
-			this._register(this.aiSettingsSearchService.onDidEnable(() => {
+			this._register(this.aiSettingsSearchService.onProviderRegistered(() => {
 				this.showAiResultsAction!.enabled = true;
 			}));
 			this._register(this.showAiResultsAction.onDidChange(() => {

src/vs/workbench/contrib/preferences/browser/settingsTreeModels.ts

@@ -919,7 +919,7 @@ export const enum SearchResultIdx {
 
 export class SearchResultModel extends SettingsTreeModel {
 	private rawSearchResults: ISearchResult[] | null = null;
-	private cachedUniqueSearchResults: ISearchResult | null = null;
+	private cachedUniqueSearchResults: Map<boolean, ISearchResult | null>;
 	private newExtensionSearchResults: ISearchResult | null = null;
 	private searchResultCount: number | null = null;
 	private settingsOrderByTocIndex: Map<string, number> | null;
@@ -939,11 +939,13 @@ export class SearchResultModel extends SettingsTreeModel {
 	) {
 		super(viewState, isWorkspaceTrusted, configurationService, languageService, userDataProfileService, productService);
 		this.settingsOrderByTocIndex = settingsOrderByTocIndex;
+		this.cachedUniqueSearchResults = new Map();
 		this.update({ id: 'searchResultModel', label: '' });
 	}
 
 	set showAiResults(show: boolean) {
 		this.aiFilterEnabled = show;
+		this.updateChildren();
 	}
 
 	private sortResults(filterMatches: ISettingMatch[]): ISettingMatch[] {
@@ -986,8 +988,9 @@ export class SearchResultModel extends SettingsTreeModel {
 	}
 
 	getUniqueSearchResults(): ISearchResult | null {
-		if (this.cachedUniqueSearchResults) {
-			return this.cachedUniqueSearchResults;
+		const cachedResults = this.cachedUniqueSearchResults.get(this.aiFilterEnabled);
+		if (cachedResults) {
+			return cachedResults;
 		}
 
 		if (!this.rawSearchResults) {
@@ -1009,33 +1012,35 @@ export class SearchResultModel extends SettingsTreeModel {
 				embeddingsResult.filterMatches = embeddingsResult.filterMatches.filter(m => !aiSelectedKeys.has(m.setting.key));
 				combinedFilterMatches = combinedFilterMatches.concat(embeddingsResult.filterMatches);
 			}
-			this.cachedUniqueSearchResults = {
+			const result = {
 				filterMatches: combinedFilterMatches,
 				exactMatch: false
 			};
-		} else {
-			const localMatchKeys = new Set<string>();
-			const localResult = this.rawSearchResults[SearchResultIdx.Local];
-			if (localResult) {
-				localResult.filterMatches.forEach(m => localMatchKeys.add(m.setting.key));
-				combinedFilterMatches = localResult.filterMatches;
-			}
-
-			const remoteResult = this.rawSearchResults[SearchResultIdx.Remote];
-			if (remoteResult) {
-				remoteResult.filterMatches = remoteResult.filterMatches.filter(m => !localMatchKeys.has(m.setting.key));
-				combinedFilterMatches = combinedFilterMatches.concat(remoteResult.filterMatches);
+			this.cachedUniqueSearchResults.set(true, result);
+			return result;
+		}
 
-				this.newExtensionSearchResults = this.rawSearchResults[SearchResultIdx.NewExtensions];
-			}
-			combinedFilterMatches = this.sortResults(combinedFilterMatches);
-			this.cachedUniqueSearchResults = {
-				filterMatches: combinedFilterMatches,
-				exactMatch: localResult.exactMatch // remote results should never have an exact match
-			};
+		const localMatchKeys = new Set<string>();
+		const localResult = this.rawSearchResults[SearchResultIdx.Local];
+		if (localResult) {
+			localResult.filterMatches.forEach(m => localMatchKeys.add(m.setting.key));
+			combinedFilterMatches = localResult.filterMatches;
 		}
 
-		return this.cachedUniqueSearchResults;
+		const remoteResult = this.rawSearchResults[SearchResultIdx.Remote];
+		if (remoteResult) {
+			remoteResult.filterMatches = remoteResult.filterMatches.filter(m => !localMatchKeys.has(m.setting.key));
+			combinedFilterMatches = combinedFilterMatches.concat(remoteResult.filterMatches);
+
+			this.newExtensionSearchResults = this.rawSearchResults[SearchResultIdx.NewExtensions];
+		}
+		combinedFilterMatches = this.sortResults(combinedFilterMatches);
+		const result = {
+			filterMatches: combinedFilterMatches,
+			exactMatch: localResult.exactMatch // remote results should never have an exact match
+		};
+		this.cachedUniqueSearchResults.set(false, result);
+		return result;
 	}
 
 	getRawResults(): ISearchResult[] {
@@ -1089,7 +1094,7 @@ export class SearchResultModel extends SettingsTreeModel {
 	}
 
 	setResult(order: SearchResultIdx, result: ISearchResult | null): void {
-		this.cachedUniqueSearchResults = null;
+		this.cachedUniqueSearchResults.clear();
 		this.newExtensionSearchResults = null;
 
 		this.rawSearchResults ??= [];

src/vs/workbench/services/aiSettingsSearch/common/aiSettingsSearch.ts

@@ -29,7 +29,7 @@ export interface AiSettingsSearchProviderOptions {
 
 export interface IAiSettingsSearchService {
 	readonly _serviceBrand: undefined;
-	readonly onDidEnable: Event<void>;
+	readonly onProviderRegistered: Event<void>;
 
 	// Called from the Settings editor
 	isEnabled(): boolean;

src/vs/workbench/services/aiSettingsSearch/common/aiSettingsSearchService.ts

@@ -6,30 +6,28 @@
 import { DeferredPromise, raceCancellation } from '../../../../base/common/async.js';
 import { CancellationToken } from '../../../../base/common/cancellation.js';
 import { Emitter, Event } from '../../../../base/common/event.js';
-import { IDisposable } from '../../../../base/common/lifecycle.js';
+import { Disposable, IDisposable } from '../../../../base/common/lifecycle.js';
 import { InstantiationType, registerSingleton } from '../../../../platform/instantiation/common/extensions.js';
 import { AiSettingsSearchResult, AiSettingsSearchResultKind, IAiSettingsSearchProvider, IAiSettingsSearchService } from './aiSettingsSearch.js';
 
-export class AiSettingsSearchService implements IAiSettingsSearchService {
+export class AiSettingsSearchService extends Disposable implements IAiSettingsSearchService {
 	readonly _serviceBrand: undefined;
 	private static readonly MAX_PICKS = 5;
 
 	private _providers: IAiSettingsSearchProvider[] = [];
 	private _llmRankedResultsPromises: Map<string, DeferredPromise<string[]>> = new Map();
 	private _embeddingsResultsPromises: Map<string, DeferredPromise<string[]>> = new Map();
 
-	private _onDidEnable: Emitter<void> = new Emitter<void>();
-	readonly onDidEnable: Event<void> = this._onDidEnable.event;
+	private _onProviderRegistered: Emitter<void> = this._register(new Emitter<void>());
+	readonly onProviderRegistered: Event<void> = this._onProviderRegistered.event;
 
 	isEnabled(): boolean {
 		return this._providers.length > 0;
 	}
 
 	registerSettingsSearchProvider(provider: IAiSettingsSearchProvider): IDisposable {
 		this._providers.push(provider);
-		if (this._providers.length === 1) {
-			this._onDidEnable.fire();
-		}
+		this._onProviderRegistered.fire();
 		return {
 			dispose: () => {
 				const index = this._providers.indexOf(provider);