Prevent enterprise MCP tools from disappearing after reconnect

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

@@ -3,6 +3,8 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import { addDisposableListener } from '../../../../base/browser/dom.js';
+import { mainWindow } from '../../../../base/browser/window.js';
 import { distinct } from '../../../../base/common/arrays.js';
 import { Barrier, RunOnceScheduler, ThrottledDelayer, timeout } from '../../../../base/common/async.js';
 import { CancellationToken } from '../../../../base/common/cancellation.js';
@@ -446,6 +448,13 @@ export class DefaultAccountProvider extends Disposable implements IDefaultAccoun
 				this.refetchDefaultAccount();
 			}
 		}));
+
+		this._register(addDisposableListener(mainWindow, 'online', () => {
+			this.logService.debug('[DefaultAccount] Network is online, refreshing default account');
+			this.updateDefaultAccount({ forceRefresh: true }).catch(error => {
+				this.logService.error('[DefaultAccount] Failed to refresh default account after network came online', getErrorMessage(error));
+			});
+		}));
 	}
 
 	private async whenDefaultAccountAuthenticationProviderAvailable(): Promise<void> {
@@ -670,7 +679,7 @@ export class DefaultAccountProvider extends Disposable implements IDefaultAccoun
 			const managedSettingsCompatibilityError = managedSettingsResult
 				? managedSettingsResult.compatibilityError
 				: this._managedSettingsCompatibilityError;
-			let mcpRegistryDataFetchedAt: number | undefined;
+			let mcpRegistryDataFetchedAt: number | undefined = accountPolicyData?.mcpRegistryDataFetchedAt;
 			let policyData: Mutable<IPolicyData> | undefined = accountPolicyData?.policyData ? { ...accountPolicyData.policyData } : undefined;
 			if (entitlementsData) {
 				policyData = policyData ?? {};
@@ -684,9 +693,14 @@ export class DefaultAccountProvider extends Disposable implements IDefaultAccoun
 				policyData.mcp = tokenEntitlementsData.policyData.mcp;
 				if (policyData.mcp) {
 					const mcpRegistryResult = await this.getMcpRegistryProvider(sessions, accountPolicyData, options);
-					mcpRegistryDataFetchedAt = mcpRegistryResult?.fetchedAt;
-					policyData.mcpRegistryUrl = mcpRegistryResult?.data?.url;
-					policyData.mcpAccess = mcpRegistryResult?.data?.registry_access;
+					mcpRegistryDataFetchedAt = mcpRegistryResult?.fetchedAt ?? accountPolicyData?.mcpRegistryDataFetchedAt;
+					if (mcpRegistryResult?.data) {
+						policyData.mcpRegistryUrl = mcpRegistryResult.data.url;
+						policyData.mcpAccess = mcpRegistryResult.data.registry_access;
+					} else if (mcpRegistryResult) {
+						policyData.mcpRegistryUrl = undefined;
+						policyData.mcpAccess = undefined;
+					}
 				} else {
 					policyData.mcpRegistryUrl = undefined;
 					policyData.mcpAccess = undefined;

src/vs/workbench/services/accounts/test/browser/defaultAccount.test.ts

@@ -229,7 +229,95 @@ suite('DefaultAccountProvider managed settings', () => {
 		});
 	});
 
-	async function createProvider(requestService: TestRequestService): Promise<DefaultAccountProvider> {
+	test('transient MCP registry failure preserves cached enterprise registry data', async () => {
+		const requestService = new TestRequestService(async options => {
+			if (options.url?.endsWith('/copilot_internal/user')) {
+				return jsonResponse({ chat_enabled: true });
+			}
+			if (options.url?.endsWith('/copilot_internal/v2/token')) {
+				return jsonResponse({ token: 'mcp=1' });
+			}
+			if (options.url?.includes('/copilot/mcp_registry')) {
+				// 5xx is a transient failure: enterprise registry data must be preserved.
+				return jsonResponse({}, 500);
+			}
+			throw new Error(`Unexpected request: ${options.url}`);
+		});
+		const provider = await createProvider(requestService, {
+			tokenEntitlementUrl: 'https://api.github.com/copilot_internal/v2/token',
+			mcpRegistryDataUrl: 'https://api.github.com/copilot/mcp_registry',
+			managedSettingsUrl: '',
+		});
+		provider['_policyData'] = {
+			accountId,
+			policyData: {
+				mcp: true,
+				mcpRegistryUrl: 'https://example.com/enterprise-registry',
+				mcpAccess: 'registry_only',
+			},
+			mcpRegistryDataFetchedAt: Date.now(),
+		};
+
+		const result = await provider['getDefaultAccountFromAuthenticatedSessions'](
+			{ id: 'github', name: 'GitHub', enterprise: false },
+			sessions,
+			{ forceRefresh: true }
+		);
+
+		assert.deepStrictEqual({
+			mcpRegistryUrl: result?.policyData?.policyData.mcpRegistryUrl,
+			mcpAccess: result?.policyData?.policyData.mcpAccess,
+		}, {
+			mcpRegistryUrl: 'https://example.com/enterprise-registry',
+			mcpAccess: 'registry_only',
+		});
+	});
+
+	test('definitive MCP registry removal clears cached registry data', async () => {
+		const requestService = new TestRequestService(async options => {
+			if (options.url?.endsWith('/copilot_internal/user')) {
+				return jsonResponse({ chat_enabled: true });
+			}
+			if (options.url?.endsWith('/copilot_internal/v2/token')) {
+				return jsonResponse({ token: 'mcp=1' });
+			}
+			if (options.url?.includes('/copilot/mcp_registry')) {
+				// 4xx is a definitive "no registry available" signal: cached data must be cleared.
+				return jsonResponse({}, 404);
+			}
+			throw new Error(`Unexpected request: ${options.url}`);
+		});
+		const provider = await createProvider(requestService, {
+			tokenEntitlementUrl: 'https://api.github.com/copilot_internal/v2/token',
+			mcpRegistryDataUrl: 'https://api.github.com/copilot/mcp_registry',
+			managedSettingsUrl: '',
+		});
+		provider['_policyData'] = {
+			accountId,
+			policyData: {
+				mcp: true,
+				mcpRegistryUrl: 'https://example.com/enterprise-registry',
+				mcpAccess: 'registry_only',
+			},
+			mcpRegistryDataFetchedAt: Date.now(),
+		};
+
+		const result = await provider['getDefaultAccountFromAuthenticatedSessions'](
+			{ id: 'github', name: 'GitHub', enterprise: false },
+			sessions,
+			{ forceRefresh: true }
+		);
+
+		assert.deepStrictEqual({
+			mcpRegistryUrl: result?.policyData?.policyData.mcpRegistryUrl,
+			mcpAccess: result?.policyData?.policyData.mcpAccess,
+		}, {
+			mcpRegistryUrl: undefined,
+			mcpAccess: undefined,
+		});
+	});
+
+	async function createProvider(requestService: TestRequestService, configOverrides?: Partial<{ tokenEntitlementUrl: string; mcpRegistryDataUrl: string; managedSettingsUrl: string }>): Promise<DefaultAccountProvider> {
 		const instantiationService = disposables.add(new TestInstantiationService());
 		instantiationService.stub(IConfigurationService, new TestConfigurationService());
 		instantiationService.stub(IAuthenticationService, {
@@ -280,6 +368,7 @@ suite('DefaultAccountProvider managed settings', () => {
 			entitlementUrl: 'https://api.github.com/copilot_internal/user',
 			mcpRegistryDataUrl: '',
 			managedSettingsUrl: 'https://api.github.com/copilot_internal/managed_settings',
+			...configOverrides,
 		}));
 		await provider.refresh();
 		return provider;