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,26 @@ 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 the setting, so it has to be cleared first.
+			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

@@ -48,7 +48,7 @@ import { IWorkbenchLayoutService, Parts } from '../../../../services/layout/brow
 import { InEditorZenModeContext } from '../../../../common/contextkeys.js';
 import { ILifecycleService } from '../../../../services/lifecycle/common/lifecycle.js';
 import { IPreferencesService } from '../../../../services/preferences/common/preferences.js';
-import { IExtension, IExtensionsWorkbenchService } from '../../../extensions/common/extensions.js';
+import { ExtensionRuntimeActionType, IExtension, IExtensionsWorkbenchService } from '../../../extensions/common/extensions.js';
 import { UpdateTitleBarEditorVisibleContext } from '../../../update/common/update.js';
 import { ChatContextKeys } from '../../common/actions/chatContextKeys.js';
 import { IChatSessionsService } from '../../common/chatSessionsService.js';
@@ -780,7 +780,8 @@ export class ChatTeardownContribution extends Disposable implements IWorkbenchCo
 		@IExtensionsWorkbenchService private readonly extensionsWorkbenchService: IExtensionsWorkbenchService,
 		@IWorkbenchExtensionEnablementService private readonly extensionEnablementService: IWorkbenchExtensionEnablementService,
 		@IViewDescriptorService private readonly viewDescriptorService: IViewDescriptorService,
-		@IWorkbenchLayoutService private readonly layoutService: IWorkbenchLayoutService
+		@IWorkbenchLayoutService private readonly layoutService: IWorkbenchLayoutService,
+		@IHostService private readonly hostService: IHostService
 	) {
 		super();
 
@@ -791,20 +792,21 @@ 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();
 		}
+
+		// Enablement is derived, but the extension host still has to be told to pick the change up.
+		const defaultChatExtension = this.extensionsWorkbenchService.local.find(value => ExtensionIdentifier.equals(value.identifier.id, defaultChat.chatExtensionId));
+		if (defaultChatExtension?.runtimeState?.action === ExtensionRuntimeActionType.ReloadWindow) {
+			return this.hostService.reload(); // a remote extension host cannot be restarted in place
+		}
+
+		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 +817,7 @@ export class ChatTeardownContribution extends Disposable implements IWorkbenchCo
 				return;
 			}
 
-			this.handleChatDisabled(true);
+			this.handleChatDisabled();
 		}));
 
 		// Extension installation
@@ -838,24 +840,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,12 @@ 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;
+		}
+		// `run` clears the setting first, so the state it owns must not gate this action.
+		if (this.extension.enablementState !== EnablementState.DisabledByAIFeaturesSetting
+			&& !this.extensionEnablementService.canChangeWorkspaceEnablement(this.extension.local)) {
 			return;
 		}
 		const inspect = this.configurationService.inspect(ChatAIDisabledSettingId);
@@ -1885,12 +1891,20 @@ 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.configurationService.getValue<boolean>(ChatAIDisabledSettingId) === true) {
+				return; // a more specific value still disables AI, there is no persisted state to fix
+			}
+		}
+
+		// A persisted disable outlives the setting, and setEnablement reports why if it cannot go.
+		if (!this.extensionEnablementService.isEnabled(this.extension.local)) {
+			await this.extensionsWorkbenchService.setEnablement(this.extension, EnablementState.EnabledWorkspace);
 		}
 	}
 }
@@ -1935,31 +1949,40 @@ class DisableAIFeaturesInWorkspaceAction extends ExtensionAction {
 
 	constructor(
 		@IProductService private readonly productService: IProductService,
-		@IExtensionsWorkbenchService private readonly extensionsWorkbenchService: IExtensionsWorkbenchService,
+		@IConfigurationService private readonly configurationService: IConfigurationService,
 		@IWorkbenchExtensionEnablementService private readonly extensionEnablementService: IWorkbenchExtensionEnablementService,
+		@IWorkspaceContextService private readonly contextService: IWorkspaceContextService,
 		@IExtensionService private readonly extensionService: IExtensionService,
 	) {
 		super(DisableAIFeaturesInWorkspaceAction.ID, DisableAIFeaturesInWorkspaceAction.LABEL, ExtensionAction.LABEL_ACTION_CLASS);
 		this.tooltip = localize('disableAIInWorkspaceActionToolTip', "Disable AI features in this workspace");
 		this.update();
 		this._register(this.extensionService.onDidChangeExtensions(() => this.update()));
+		this._register(this.configurationService.onDidChangeConfiguration(e => {
+			if (e.affectsConfiguration(ChatAIDisabledSettingId)) {
+				this.update();
+			}
+		}));
 	}
 
 	update(): void {
 		this.enabled = false;
-		if (this.extension && this.extension.local && ExtensionIdentifier.equals(this.extension.identifier.id, this.productService.defaultChatAgent?.chatExtensionId)) {
-			this.enabled = this.extension.state === ExtensionState.Installed
-				&& (this.extension.enablementState === EnablementState.EnabledGlobally || this.extension.enablementState === EnablementState.EnabledWorkspace)
-				&& this.extensionEnablementService.canChangeWorkspaceEnablement(this.extension.local);
+		if (!this.extension?.local) {
+			return;
 		}
+		if (!ExtensionIdentifier.equals(this.extension.identifier.id, this.productService.defaultChatAgent?.chatExtensionId)) {
+			return;
+		}
+		if (this.contextService.getWorkbenchState() === WorkbenchState.EMPTY) {
+			return;
+		}
+		this.enabled = this.extension.state === ExtensionState.Installed
+			&& this.configurationService.getValue<boolean>(ChatAIDisabledSettingId) !== true
+			&& this.extensionEnablementService.isEnabled(this.extension.local);
 	}
 
 	override async run(): Promise<void> {
-		if (!this.extension) {
-			return;
-		}
-		await this.extensionsWorkbenchService.setEnablement(this.extension, EnablementState.DisabledWorkspace);
-		await this.extensionsWorkbenchService.updateRunningExtensions(localize('restartExtensionHost.reason.disable', "Disabling AI features"));
+		await this.configurationService.updateValue(ChatAIDisabledSettingId, true, ConfigurationTarget.WORKSPACE);
 	}
 }
 

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.reconcileChatExtensionDisablement();
+				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.reconcileChatExtensionDisablement();
 	}
 
 	private ensureChatExtensionInitialDisabledState(): void {
@@ -190,6 +198,45 @@ export class ExtensionEnablementService extends Disposable implements IWorkbench
 		}
 	}
 
+	// Releases that applied `chat.disableAIFeatures` by persisting enablement left an entry that
+	// outlives the setting: a global one reaches every window on this machine, a workspace one
+	// survives the setting being cleared. Drop whichever the derived state now accounts for. Runs
+	// at startup and on profile switch only, so it cannot discard a disable made after either.
+	private reconcileChatExtensionDisablement(): void {
+		if (!this._chatExtensionId) {
+			return;
+		}
+
+		const identifier = { id: this._chatExtensionId };
+		const inspect = this.configurationService.inspect<boolean>(ChatAIDisabledSettingId);
+
+		if (inspect.workspaceValue === true && this._getWorkspaceDisabledExtensions().some(e => areSameExtensions(e, identifier))) {
+			this.logService.debug('Removing workspace disablement of builtin chat extension in favor of chat.disableAIFeatures');
+			this._removeFromWorkspaceDisabledExtensions(identifier)
+				.catch(err => this.logService.error('Failed to remove workspace disablement of builtin chat extension', err));
+		}
+
+		// A workspace value can mask the one the global entry was written for, so that entry is
+		// reconciled against the user and application values rather than the resolved one.
+		if (inspect.userValue !== true && inspect.applicationValue !== true) {
+			return;
+		}
+		if (!this._isDisabledGlobally(identifier)) {
+			return;
+		}
+
+		// The disable for a profile where chat setup never completed belongs to
+		// ensureChatExtensionInitialDisabledState and has to survive.
+		const context = (this.chatEntitlementService as ChatEntitlementService).context;
+		if (context && !context.value.state.completed) {
+			return;
+		}
+
+		this.logService.debug('Removing global disablement of builtin chat extension in favor of chat.disableAIFeatures');
+		this.globalExtensionEnablementService.enableExtension(identifier, 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 +325,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 +546,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 +744,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);
@@ -807,6 +866,11 @@ export class ExtensionEnablementService extends Disposable implements IWorkbench
 	}
 
 	private _onDidChangeExtensions(added: ReadonlyArray<IExtension>, removed: ReadonlyArray<IExtension>, isProfileSwitch: boolean): void {
+		if (isProfileSwitch) {
+			// The disabled extension lists are profile scoped, so the new profile carries its own entries
+			this.reconcileChatExtensionDisablement();
+		}
+
 		const changedExtensions: IExtension[] = added.filter(e => !this.isEnabledEnablementState(this.getEnablementState(e)));
 		const existingDisabledExtensions = this.extensionsDisabledExtensions;
 		this.extensionsDisabledExtensions = this.extensionsManager.extensions.filter(extension => {

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,