Derive chat extension enablement from chat.disableAIFeatures

src/vs/workbench/contrib/chat/browser/chatManagement/chatManagement.contribution.ts

@@ -6,6 +6,7 @@
 import { KeyCode } from '../../../../../base/common/keyCodes.js';
 import { localize, localize2 } from '../../../../../nls.js';
 import { Action2, MenuId, MenuRegistry, registerAction2 } from '../../../../../platform/actions/common/actions.js';
+import { IConfigurationService } from '../../../../../platform/configuration/common/configuration.js';
 import { ContextKeyExpr } from '../../../../../platform/contextkey/common/contextkey.js';
 import { ExtensionIdentifier } from '../../../../../platform/extensions/common/extensions.js';
 import { SyncDescriptor } from '../../../../../platform/instantiation/common/descriptors.js';
@@ -21,7 +22,7 @@ import { IEditorService } from '../../../../services/editor/common/editorService
 import { ResourceContextKey } from '../../../../common/contextkeys.js';
 import { ChatContextKeys } from '../../common/actions/chatContextKeys.js';
 import { ChatEntitlementContextKeys } from '../../../../services/chat/common/chatEntitlementService.js';
-import { CONTEXT_MODELS_EDITOR, CONTEXT_MODELS_SEARCH_FOCUS, MANAGE_CHAT_COMMAND_ID } from '../../common/constants.js';
+import { ChatAIDisabledSettingId, CONTEXT_MODELS_EDITOR, CONTEXT_MODELS_SEARCH_FOCUS, MANAGE_CHAT_COMMAND_ID } from '../../common/constants.js';
 import { CHAT_CATEGORY } from '../actions/chatActions.js';
 import { ModelsManagementEditor } from './chatManagementEditor.js';
 import { ModelsManagementEditorInput } from './chatManagementEditorInput.js';
@@ -86,18 +87,27 @@ async function ensureChatExtensionEnabled(accessor: ServicesAccessor): Promise<v
 
 	const extensionsWorkbenchService = accessor.get(IExtensionsWorkbenchService);
 	const extensionEnablementService = accessor.get(IWorkbenchExtensionEnablementService);
+	const configurationService = accessor.get(IConfigurationService);
 	const progressService = accessor.get(IProgressService);
 
 	const localExtensions = await extensionsWorkbenchService.queryLocal();
 	const chatExtension = localExtensions.find(e => ExtensionIdentifier.equals(e.identifier.id, chatExtensionId));
 	if (!chatExtension?.local || extensionEnablementService.isEnabled(chatExtension.local)) {
 		return;
 	}
+	const local = chatExtension.local;
 
 	await progressService.withProgress(
 		{ location: ProgressLocation.Window, title: localize('enableChatForByok', "Enabling AI features…") },
 		async () => {
-			await extensionsWorkbenchService.setEnablement([chatExtension], EnablementState.EnabledGlobally);
+			// Enablement is derived from `chat.disableAIFeatures` wherever it resolves to true, so
+			// the setting has to go before persisted enablement can be changed.
+			if (configurationService.getValue<boolean>(ChatAIDisabledSettingId) === true) {
+				await configurationService.updateValue(ChatAIDisabledSettingId, false);
+			}
+			if (!extensionEnablementService.isEnabled(local)) {
+				await extensionsWorkbenchService.setEnablement([chatExtension], EnablementState.EnabledGlobally);
+			}
 			await extensionsWorkbenchService.updateRunningExtensions(localize('enableChatForByokReason', "Enabling AI features"));
 		}
 	);

src/vs/workbench/contrib/chat/browser/chatSetup/chatSetupContributions.ts

@@ -791,20 +791,17 @@ export class ChatTeardownContribution extends Disposable implements IWorkbenchCo
 
 		this.registerListeners();
 		this.registerActions();
-
-		this.handleChatDisabled(false);
 	}
 
-	private handleChatDisabled(fromEvent: boolean): void {
-		const chatDisabled = this.configurationService.inspect(ChatAIDisabledSettingId);
-		if (chatDisabled.value === true) {
-			this.maybeEnableOrDisableExtension(typeof chatDisabled.workspaceValue === 'boolean' ? EnablementState.DisabledWorkspace : EnablementState.DisabledGlobally);
-			if (fromEvent) {
-				this.maybeHideAuxiliaryBar();
-			}
-		} else if (chatDisabled.value === false && fromEvent /* do not enable extensions unless its an explicit settings change */) {
-			this.maybeEnableOrDisableExtension(typeof chatDisabled.workspaceValue === 'boolean' ? EnablementState.EnabledWorkspace : EnablementState.EnabledGlobally);
+	private async handleChatDisabled(): Promise<void> {
+		const chatDisabled = this.configurationService.getValue(ChatAIDisabledSettingId) === true;
+		if (chatDisabled) {
+			this.maybeHideAuxiliaryBar();
 		}
+
+		// The chat extension is disabled by the enablement service deriving this setting,
+		// the extension host still has to be told to pick the change up.
+		await this.extensionsWorkbenchService.updateRunningExtensions(chatDisabled ? localize('restartExtensionHost.reason.disable', "Disabling AI features") : localize('restartExtensionHost.reason.enable', "Enabling AI features"));
 	}
 
 	private async registerListeners(): Promise<void> {
@@ -815,7 +812,7 @@ export class ChatTeardownContribution extends Disposable implements IWorkbenchCo
 				return;
 			}
 
-			this.handleChatDisabled(true);
+			this.handleChatDisabled();
 		}));
 
 		// Extension installation
@@ -838,24 +835,6 @@ export class ChatTeardownContribution extends Disposable implements IWorkbenchCo
 		}));
 	}
 
-	private async maybeEnableOrDisableExtension(state: EnablementState.EnabledGlobally | EnablementState.EnabledWorkspace | EnablementState.DisabledGlobally | EnablementState.DisabledWorkspace): Promise<void> {
-		const defaultChatExtension = this.extensionsWorkbenchService.local.find(value => ExtensionIdentifier.equals(value.identifier.id, defaultChat.chatExtensionId));
-		if (!defaultChatExtension?.local) {
-			return;
-		}
-
-		const workspace = state === EnablementState.EnabledWorkspace || state === EnablementState.DisabledWorkspace;
-		const canChange = workspace
-			? this.extensionEnablementService.canChangeWorkspaceEnablement(defaultChatExtension.local)
-			: this.extensionEnablementService.canChangeEnablement(defaultChatExtension.local);
-		if (!canChange) {
-			return;
-		}
-
-		await this.extensionsWorkbenchService.setEnablement([defaultChatExtension], state);
-		await this.extensionsWorkbenchService.updateRunningExtensions(state === EnablementState.EnabledGlobally || state === EnablementState.EnabledWorkspace ? localize('restartExtensionHost.reason.enable', "Enabling AI features") : localize('restartExtensionHost.reason.disable', "Disabling AI features"));
-	}
-
 	private maybeHideAuxiliaryBar(): void {
 		const activeContainers = this.viewDescriptorService.getViewContainersByLocation(ViewContainerLocation.AuxiliaryBar).filter(
 			container => this.viewDescriptorService.getViewContainerModel(container).activeViewDescriptors.length > 0

src/vs/workbench/contrib/extensions/browser/extensionsActions.ts

@@ -1847,6 +1847,7 @@ export class EnableAIFeaturesInWorkspaceAction extends ExtensionAction {
 		@IExtensionsWorkbenchService private readonly extensionsWorkbenchService: IExtensionsWorkbenchService,
 		@IConfigurationService private readonly configurationService: IConfigurationService,
 		@IWorkbenchExtensionEnablementService private readonly extensionEnablementService: IWorkbenchExtensionEnablementService,
+		@IWorkspaceContextService private readonly contextService: IWorkspaceContextService,
 	) {
 		super(EnableAIFeaturesInWorkspaceAction.ID, EnableAIFeaturesInWorkspaceAction.LABEL, ExtensionAction.LABEL_ACTION_CLASS);
 		this.tooltip = localize('enableAIInWorkspaceActionToolTip', "Enable AI features in this workspace");
@@ -1866,7 +1867,14 @@ export class EnableAIFeaturesInWorkspaceAction extends ExtensionAction {
 		if (!ExtensionIdentifier.equals(this.extension.identifier.id, this.productService.defaultChatAgent?.chatExtensionId)) {
 			return;
 		}
-		if (!this.extensionEnablementService.canChangeWorkspaceEnablement(this.extension.local)) {
+		if (this.contextService.getWorkbenchState() === WorkbenchState.EMPTY) {
+			return;
+		}
+		// While the setting resolves to true it owns enablement and canChangeWorkspaceEnablement
+		// reports false. `run` clears the setting before touching persisted state, so this action
+		// stays the way out of it.
+		if (this.extension.enablementState !== EnablementState.DisabledByAIFeaturesSetting
+			&& !this.extensionEnablementService.canChangeWorkspaceEnablement(this.extension.local)) {
 			return;
 		}
 		const inspect = this.configurationService.inspect(ChatAIDisabledSettingId);
@@ -1885,13 +1893,17 @@ export class EnableAIFeaturesInWorkspaceAction extends ExtensionAction {
 	}
 
 	override async run(): Promise<void> {
-		if (!this.extension) {
+		if (!this.extension?.local) {
 			return;
 		}
-		await this.extensionsWorkbenchService.setEnablement(this.extension, EnablementState.EnabledWorkspace);
+
 		if (this.configurationService.getValue<boolean>(ChatAIDisabledSettingId) === true) {
 			await this.configurationService.updateValue(ChatAIDisabledSettingId, false, ConfigurationTarget.WORKSPACE);
 		}
+
+		if (!this.extensionEnablementService.isEnabled(this.extension.local)) {
+			await this.extensionsWorkbenchService.setEnablement(this.extension, EnablementState.EnabledWorkspace);
+		}
 	}
 }
 

src/vs/workbench/services/extensionManagement/browser/extensionEnablementService.ts

@@ -136,6 +136,13 @@ export class ExtensionEnablementService extends Disposable implements IWorkbench
 			}
 		}));
 
+		this._register(this.configurationService.onDidChangeConfiguration(e => {
+			if (e.affectsConfiguration(ChatAIDisabledSettingId)) {
+				this.reconcileChatExtensionGlobalDisablement();
+				this._onEnablementChanged.fire(this.extensionsManager.extensions.filter(ext => ext.identifier.id.toLowerCase() === this._chatExtensionId));
+			}
+		}));
+
 		// delay notification for extensions disabled until workbench restored
 		if (this.allUserExtensionsDisabled) {
 			this.lifecycleService.when(LifecyclePhase.Eventually).then(() => {
@@ -150,6 +157,7 @@ export class ExtensionEnablementService extends Disposable implements IWorkbench
 		}
 
 		this.ensureChatExtensionInitialDisabledState();
+		this.reconcileChatExtensionGlobalDisablement();
 	}
 
 	private ensureChatExtensionInitialDisabledState(): void {
@@ -190,6 +198,25 @@ export class ExtensionEnablementService extends Disposable implements IWorkbench
 		}
 	}
 
+	private reconcileChatExtensionGlobalDisablement(): void {
+		if (!this._chatExtensionId) {
+			return;
+		}
+
+		// A globally disabled chat extension applies to every window on this machine, whereas
+		// `chat.disableAIFeatures` applies only where its configuration resolves. Releases that
+		// implemented the setting as a global disable leak into windows the setting does not cover,
+		// so drop the entry that the derived enablement state already accounts for. Runs on every
+		// change too, because the migration above settles the setting asynchronously.
+		if (this.configurationService.getValue(ChatAIDisabledSettingId) !== true || !this._isDisabledGlobally({ id: this._chatExtensionId })) {
+			return;
+		}
+
+		this.logService.debug('Removing global disablement of builtin chat extension in favor of chat.disableAIFeatures');
+		this.globalExtensionEnablementService.enableExtension({ id: this._chatExtensionId }, SOURCE)
+			.catch(err => this.logService.error('Failed to remove global disablement of builtin chat extension', err));
+	}
+
 	private get hasWorkspace(): boolean {
 		return this.contextService.getWorkbenchState() !== WorkbenchState.EMPTY;
 	}
@@ -278,6 +305,8 @@ export class ExtensionEnablementService extends Disposable implements IWorkbench
 				throw new Error(localize('cannot change enablement virtual workspace', "Cannot change enablement of {0} extension because it does not support virtual workspaces", extension.manifest.displayName || extension.identifier.id));
 			case EnablementState.DisabledByExtensionKind:
 				throw new Error(localize('cannot change enablement extension kind', "Cannot change enablement of {0} extension because of its extension kind", extension.manifest.displayName || extension.identifier.id));
+			case EnablementState.DisabledByAIFeaturesSetting:
+				throw new Error(localize('cannot change enablement ai features', "Cannot change enablement of {0} extension because AI features are disabled in settings", extension.manifest.displayName || extension.identifier.id));
 			case EnablementState.DisabledByAllowlist:
 				throw new Error(localize('cannot change disallowed extension enablement', "Cannot change enablement of {0} extension because it is disallowed", extension.manifest.displayName || extension.identifier.id));
 			case EnablementState.DisabledByInvalidExtension:
@@ -497,6 +526,10 @@ export class ExtensionEnablementService extends Disposable implements IWorkbench
 			enablementState = EnablementState.DisabledByEnvironment;
 		}
 
+		else if (this._isDisabledByAIFeaturesSetting(extension)) {
+			enablementState = EnablementState.DisabledByAIFeaturesSetting;
+		}
+
 		else if (isEnabled && this._isDisabledByExtensionDependency(extension, extensions, workspaceType, computedEnablementStates)) {
 			enablementState = EnablementState.DisabledByExtensionDependency;
 		}
@@ -691,6 +724,12 @@ export class ExtensionEnablementService extends Disposable implements IWorkbench
 		return !this.extensionManifestPropertiesService.canExecuteOnSessionsWindow(extension.manifest);
 	}
 
+	private _isDisabledByAIFeaturesSetting(extension: IExtension): boolean {
+		return !!this._chatExtensionId
+			&& extension.identifier.id.toLowerCase() === this._chatExtensionId
+			&& this.configurationService.getValue(ChatAIDisabledSettingId) === true;
+	}
+
 	private _enableExtension(identifier: IExtensionIdentifier): Promise<boolean> {
 		this._removeFromWorkspaceDisabledExtensions(identifier);
 		this._removeFromWorkspaceEnabledExtensions(identifier);

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

@@ -109,6 +109,7 @@ export const enum EnablementState {
 	DisabledByAllowlist,
 	DisabledByExtensionDependency,
 	DisabledByUnification, // Temporary TODO@benibenj remove when unification transition is complete
+	DisabledByAIFeaturesSetting,
 	DisabledGlobally,
 	DisabledWorkspace,
 	EnabledGlobally,