Adopt mcp token flag (take 2) (microsoft#252669)

* adopt mcp token flag (microsoft/vscode-internalbacklog#5514)

* dont use constant to ChatMCP

* use common/objects equals()

Author: joshspicer

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

@@ -258,12 +258,9 @@ configurationRegistry.registerConfiguration({
 			type: 'boolean',
 			description: nls.localize('chat.mcp.enabled', "Enables integration with Model Context Protocol servers to provide additional tools and functionality."),
 			default: true,
-			tags: ['preview'],
 			policy: {
 				name: 'ChatMCP',
 				minimumVersion: '1.99',
-				previewFeature: true,
-				defaultValue: false
 			}
 		},
 		[mcpServerSamplingSection]: {

src/vs/workbench/services/accounts/common/defaultAccount.ts

@@ -36,6 +36,7 @@ export interface IDefaultAccount {
 	readonly can_signup_for_limited?: boolean;
 	readonly chat_enabled?: boolean;
 	readonly chat_preview_features_enabled?: boolean;
+	readonly mcp?: boolean;
 	readonly analytics_tracking_id?: string;
 	readonly limited_user_quotas?: {
 		readonly chat: number;
@@ -189,15 +190,15 @@ export class DefaultAccountManagementContribution extends Disposable implements
 		}
 	}
 
-	private extractFromToken(token: string, key: string): string | undefined {
+	private extractFromToken(token: string): Map<string, string> {
 		const result = new Map<string, string>();
 		const firstPart = token?.split(':')[0];
 		const fields = firstPart?.split(';');
 		for (const field of fields) {
 			const [key, value] = field.split('=');
 			result.set(key, value);
 		}
-		return result.get(key);
+		return result;
 	}
 
 	private async getDefaultAccountFromAuthenticatedSessions(authProviderId: string, enterpriseAuthProviderId: string, enterpriseAuthProviderConfig: string, scopes: string[], tokenEntitlementUrl: string, chatEntitlementUrl: string): Promise<IDefaultAccount | null> {
@@ -243,9 +244,12 @@ export class DefaultAccountManagementContribution extends Disposable implements
 
 			const chatData = await asJson<ITokenEntitlementsResponse>(chatContext);
 			if (chatData) {
+				const tokenMap = this.extractFromToken(chatData.token);
 				return {
 					// Editor preview features are disabled if the flag is present and set to 0
-					chat_preview_features_enabled: this.extractFromToken(chatData.token, 'editor_preview_features') !== '0',
+					chat_preview_features_enabled: tokenMap.get('editor_preview_features') !== '0',
+					// MCP is disabled if the flag is present and set to 0
+					mcp: tokenMap.get('mcp') !== '0',
 				};
 			}
 			this.logService.error('Failed to fetch token entitlements', 'No data returned');

src/vs/workbench/services/policies/common/accountPolicyService.ts

@@ -4,12 +4,21 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { IStringDictionary } from '../../../../base/common/collections.js';
+import { equals } from '../../../../base/common/objects.js';
 import { ILogService } from '../../../../platform/log/common/log.js';
 import { AbstractPolicyService, IPolicyService, PolicyDefinition } from '../../../../platform/policy/common/policy.js';
 import { IDefaultAccountService } from '../../accounts/common/defaultAccount.js';
 
+interface IAccountPolicy {
+	readonly chatPreviewFeaturesEnabled: boolean;
+	readonly mcpEnabled: boolean;
+}
+
 export class AccountPolicyService extends AbstractPolicyService implements IPolicyService {
-	private chatPreviewFeaturesEnabled: boolean = true;
+	private accountPolicy: IAccountPolicy = {
+		chatPreviewFeaturesEnabled: true,
+		mcpEnabled: true
+	};
 	constructor(
 		@ILogService private readonly logService: ILogService,
 		@IDefaultAccountService private readonly defaultAccountService: IDefaultAccountService
@@ -18,43 +27,61 @@ export class AccountPolicyService extends AbstractPolicyService implements IPoli
 
 		this.defaultAccountService.getDefaultAccount()
 			.then(account => {
-				this._update(account?.chat_preview_features_enabled ?? true);
-				this._register(this.defaultAccountService.onDidChangeDefaultAccount(account => this._update(account?.chat_preview_features_enabled ?? true)));
+				this._update({
+					chatPreviewFeaturesEnabled: account?.chat_preview_features_enabled ?? true,
+					mcpEnabled: account?.mcp ?? true
+				});
+				this._register(this.defaultAccountService.onDidChangeDefaultAccount(
+					account => this._update({
+						chatPreviewFeaturesEnabled: account?.chat_preview_features_enabled ?? true,
+						mcpEnabled: account?.mcp ?? true
+					})
+				));
 			});
 	}
 
-	private _update(chatPreviewFeaturesEnabled: boolean | undefined) {
-		const newValue = (chatPreviewFeaturesEnabled === undefined) || chatPreviewFeaturesEnabled;
-		if (this.chatPreviewFeaturesEnabled !== newValue) {
-			this.chatPreviewFeaturesEnabled = newValue;
+	private _update(updatedPolicy: IAccountPolicy): void {
+		if (!equals(this.accountPolicy, updatedPolicy)) {
+			this.accountPolicy = updatedPolicy;
 			this._updatePolicyDefinitions(this.policyDefinitions);
 		}
 	}
 
 	protected async _updatePolicyDefinitions(policyDefinitions: IStringDictionary<PolicyDefinition>): Promise<void> {
 		this.logService.trace(`AccountPolicyService#_updatePolicyDefinitions: Got ${Object.keys(policyDefinitions).length} policy definitions`);
+		const updated: string[] = [];
 
-		const update: string[] = [];
-		for (const key in policyDefinitions) {
-			const policy = policyDefinitions[key];
-			if (policy.previewFeature) {
-				if (this.chatPreviewFeaturesEnabled) {
+		const updateIfNeeded = (key: string, policy: PolicyDefinition, isFeatureEnabled: boolean): void => {
+			if (isFeatureEnabled) {
+				// Clear the policy if it is set
+				if (this.policies.has(key)) {
 					this.policies.delete(key);
-					update.push(key);
-					continue;
+					updated.push(key);
 				}
-				const defaultValue = policy.defaultValue;
-				const updatedValue = defaultValue === undefined ? false : defaultValue;
-				if (this.policies.get(key) === updatedValue) {
-					continue;
+			} else {
+				// Enforce the defaultValue if not already set
+				const updatedValue = policy.defaultValue === undefined ? false : policy.defaultValue;
+				if (this.policies.get(key) !== updatedValue) {
+					this.policies.set(key, updatedValue);
+					updated.push(key);
 				}
-				this.policies.set(key, updatedValue);
-				update.push(key);
+			}
+		};
+
+		for (const key in policyDefinitions) {
+			const policy = policyDefinitions[key];
+			// Preview Features
+			if (policy.previewFeature) {
+				updateIfNeeded(key, policy, this.accountPolicy?.chatPreviewFeaturesEnabled);
+			}
+			// MCP
+			else if (key === 'ChatMCP') {
+				updateIfNeeded(key, policy, this.accountPolicy?.mcpEnabled);
 			}
 		}
 
-		if (update.length) {
-			this._onDidChange.fire(update);
+		if (updated.length) {
+			this._onDidChange.fire(updated);
 		}
 	}
 }