Improve organization of auth state for MCP servers (#251864)

Now they will show in the their own section in the Account Menu AND have an option to Disconnect or Sign out in the MCP Server menu depending on if the account is used by other things.

Author: TylerLeonhardt

src/vs/workbench/api/browser/mainThreadAuthentication.ts

@@ -225,7 +225,7 @@ export class MainThreadAuthentication extends Disposable implements MainThreadAu
 	}
 
 	async $registerDynamicAuthenticationProvider(id: string, label: string, authorizationServer: UriComponents, clientId: string): Promise<void> {
-		await this.$registerAuthenticationProvider(id, label, false, [authorizationServer]);
+		await this.$registerAuthenticationProvider(id, label, true, [authorizationServer]);
 		this.dynamicAuthProviderStorageService.storeClientId(id, URI.revive(authorizationServer).toString(true), clientId, label);
 	}
 

src/vs/workbench/api/browser/mainThreadMcp.ts

@@ -165,10 +165,12 @@ export class MainThreadMcp extends Disposable implements MainThreadMcpShape {
 		if (sessions.length) {
 			// If we have an existing session preference, use that. If not, we'll return any valid session at the end of this function.
 			if (matchingAccountPreferenceSession && this.authenticationMCPServerAccessService.isAccessAllowed(providerId, matchingAccountPreferenceSession.account.label, server.id)) {
+				this._mcpRegistry.setAuthenticationUsage(server.id, providerId);
 				return matchingAccountPreferenceSession.accessToken;
 			}
 			// If we only have one account for a single auth provider, lets just check if it's allowed and return it if it is.
 			if (!provider.supportsMultipleAccounts && this.authenticationMCPServerAccessService.isAccessAllowed(providerId, sessions[0].account.label, server.id)) {
+				this._mcpRegistry.setAuthenticationUsage(server.id, providerId);
 				return sessions[0].accessToken;
 			}
 		}
@@ -201,6 +203,7 @@ export class MainThreadMcp extends Disposable implements MainThreadMcpShape {
 			);
 		}
 
+		this._mcpRegistry.setAuthenticationUsage(server.id, providerId);
 		this.authenticationMCPServerAccessService.updateAllowedMcpServers(providerId, session.account.label, [{ id: server.id, name: server.label, allowed: true }]);
 		this.authenticationMcpServersService.updateAccountPreference(server.id, providerId, session.account);
 		this.authenticationMCPServerUsageService.addAccountUsage(providerId, session.account.label, scopesSupported, server.id, server.label);

src/vs/workbench/browser/parts/globalCompositeBar.ts

@@ -358,66 +358,118 @@ export class AccountsActivityActionViewItem extends AbstractGlobalActivityAction
 	protected override async resolveMainMenuActions(accountsMenu: IMenu, disposables: DisposableStore): Promise<IAction[]> {
 		await super.resolveMainMenuActions(accountsMenu, disposables);
 
-		const providers = this.authenticationService.getProviderIds();
+		const providers = this.authenticationService.getProviderIds().filter(p => !p.startsWith(INTERNAL_AUTH_PROVIDER_PREFIX));
 		const otherCommands = accountsMenu.getActions();
 		let menus: IAction[] = [];
 
-		for (const providerId of providers) {
-			if (providerId.startsWith(INTERNAL_AUTH_PROVIDER_PREFIX)) {
-				continue;
-			}
-			if (!this.initialized) {
-				const noAccountsAvailableAction = disposables.add(new Action('noAccountsAvailable', localize('loading', "Loading..."), undefined, false));
-				menus.push(noAccountsAvailableAction);
-				break;
-			}
-			const provider = this.authenticationService.getProvider(providerId);
-			const accounts = this.groupedAccounts.get(providerId);
-			if (!accounts) {
-				if (this.problematicProviders.has(providerId)) {
-					const providerUnavailableAction = disposables.add(new Action('providerUnavailable', localize('authProviderUnavailable', '{0} is currently unavailable', provider.label), undefined, false));
-					menus.push(providerUnavailableAction);
-					// try again in the background so that if the failure was intermittent, we can resolve it on the next showing of the menu
-					try {
-						await this.addAccountsFromProvider(providerId);
-					} catch (e) {
-						this.logService.error(e);
+		const registeredProviders = providers.filter(providerId => !this.authenticationService.isDynamicAuthenticationProvider(providerId));
+		const dynamicProviders = providers.filter(providerId => this.authenticationService.isDynamicAuthenticationProvider(providerId));
+
+		if (!this.initialized) {
+			const noAccountsAvailableAction = disposables.add(new Action('noAccountsAvailable', localize('loading', "Loading..."), undefined, false));
+			menus.push(noAccountsAvailableAction);
+		} else {
+			for (const providerId of registeredProviders) {
+				const provider = this.authenticationService.getProvider(providerId);
+				const accounts = this.groupedAccounts.get(providerId);
+				if (!accounts) {
+					if (this.problematicProviders.has(providerId)) {
+						const providerUnavailableAction = disposables.add(new Action('providerUnavailable', localize('authProviderUnavailable', '{0} is currently unavailable', provider.label), undefined, false));
+						menus.push(providerUnavailableAction);
+						// try again in the background so that if the failure was intermittent, we can resolve it on the next showing of the menu
+						try {
+							await this.addAccountsFromProvider(providerId);
+						} catch (e) {
+							this.logService.error(e);
+						}
 					}
+					continue;
 				}
-				continue;
+
+				const canUseMcp = !!provider.authorizationServers?.length;
+				for (const account of accounts) {
+					const manageExtensionsAction = toAction({
+						id: `configureSessions${account.label}`,
+						label: localize('manageTrustedExtensions', "Manage Trusted Extensions"),
+						enabled: true,
+						run: () => this.commandService.executeCommand('_manageTrustedExtensionsForAccount', { providerId, accountLabel: account.label })
+					});
+
+
+					const providerSubMenuActions: IAction[] = [manageExtensionsAction];
+					if (canUseMcp) {
+						const manageMCPAction = toAction({
+							id: `configureSessions${account.label}`,
+							label: localize('manageTrustedMCPServers', "Manage Trusted MCP Servers"),
+							enabled: true,
+							run: () => this.commandService.executeCommand('_manageTrustedMCPServersForAccount', { providerId, accountLabel: account.label })
+						});
+						providerSubMenuActions.push(manageMCPAction);
+					}
+					if (account.canSignOut) {
+						providerSubMenuActions.push(toAction({
+							id: 'signOut',
+							label: localize('signOut', "Sign Out"),
+							enabled: true,
+							run: () => this.commandService.executeCommand('_signOutOfAccount', { providerId, accountLabel: account.label })
+						}));
+					}
+
+					const providerSubMenu = new SubmenuAction('activitybar.submenu', `${account.label} (${provider.label})`, providerSubMenuActions);
+					menus.push(providerSubMenu);
+				}
+			}
+
+			if (dynamicProviders.length && registeredProviders.length) {
+				menus.push(new Separator());
 			}
 
-			const canUseMcp = !!provider.authorizationServers?.length;
-			for (const account of accounts) {
-				const manageExtensionsAction = toAction({
-					id: `configureSessions${account.label}`,
-					label: localize('manageTrustedExtensions', "Manage Trusted Extensions"),
-					enabled: true,
-					run: () => this.commandService.executeCommand('_manageTrustedExtensionsForAccount', { providerId, accountLabel: account.label })
-				});
+			for (const providerId of dynamicProviders) {
+				const provider = this.authenticationService.getProvider(providerId);
+				const accounts = this.groupedAccounts.get(providerId);
+				if (!accounts) {
+					if (this.problematicProviders.has(providerId)) {
+						const providerUnavailableAction = disposables.add(new Action('providerUnavailable', localize('authProviderUnavailable', '{0} is currently unavailable', provider.label), undefined, false));
+						menus.push(providerUnavailableAction);
+						// try again in the background so that if the failure was intermittent, we can resolve it on the next showing of the menu
+						try {
+							await this.addAccountsFromProvider(providerId);
+						} catch (e) {
+							this.logService.error(e);
+						}
+					}
+					continue;
+				}
 
+				for (const account of accounts) {
+					// TODO@TylerLeonhardt: Is there a nice way to bring this back?
+					// const manageExtensionsAction = toAction({
+					// 	id: `configureSessions${account.label}`,
+					// 	label: localize('manageTrustedExtensions', "Manage Trusted Extensions"),
+					// 	enabled: true,
+					// 	run: () => this.commandService.executeCommand('_manageTrustedExtensionsForAccount', { providerId, accountLabel: account.label })
+					// });
 
-				const providerSubMenuActions: IAction[] = [manageExtensionsAction];
-				if (canUseMcp) {
+					const providerSubMenuActions: IAction[] = [];
 					const manageMCPAction = toAction({
 						id: `configureSessions${account.label}`,
 						label: localize('manageTrustedMCPServers', "Manage Trusted MCP Servers"),
 						enabled: true,
 						run: () => this.commandService.executeCommand('_manageTrustedMCPServersForAccount', { providerId, accountLabel: account.label })
 					});
 					providerSubMenuActions.push(manageMCPAction);
-				}
-				if (account.canSignOut) {
-					providerSubMenuActions.push(toAction({
-						id: 'signOut',
-						label: localize('signOut', "Sign Out"),
-						enabled: true,
-						run: () => this.commandService.executeCommand('_signOutOfAccount', { providerId, accountLabel: account.label })
-					}));
-				}
+					if (account.canSignOut) {
+						providerSubMenuActions.push(toAction({
+							id: 'signOut',
+							label: localize('signOut', "Sign Out"),
+							enabled: true,
+							run: () => this.commandService.executeCommand('_signOutOfAccount', { providerId, accountLabel: account.label })
+						}));
+					}
 
-				const providerSubMenu = new SubmenuAction('activitybar.submenu', `${account.label} (${provider.label})`, providerSubMenuActions);
-				menus.push(providerSubMenu);
+					const providerSubMenu = new SubmenuAction('activitybar.submenu', `${account.label} (${provider.label})`, providerSubMenuActions);
+					menus.push(providerSubMenu);
+				}
 			}
 		}
 

src/vs/workbench/contrib/mcp/browser/mcpCommands.ts

@@ -25,12 +25,17 @@ import { ContextKeyExpr } from '../../../../platform/contextkey/common/contextke
 import { ExtensionsLocalizedLabel } from '../../../../platform/extensionManagement/common/extensionManagement.js';
 import { IInstantiationService, ServicesAccessor } from '../../../../platform/instantiation/common/instantiation.js';
 import { IMcpGalleryService } from '../../../../platform/mcp/common/mcpManagement.js';
+import { IProductService } from '../../../../platform/product/common/productService.js';
 import { IQuickInputService, IQuickPickItem, IQuickPickSeparator } from '../../../../platform/quickinput/common/quickInput.js';
 import { StorageScope } from '../../../../platform/storage/common/storage.js';
 import { spinningLoading } from '../../../../platform/theme/common/iconRegistry.js';
 import { IWorkspaceContextService } from '../../../../platform/workspace/common/workspace.js';
 import { ActiveEditorContext, ResourceContextKey } from '../../../common/contextkeys.js';
 import { IWorkbenchContribution } from '../../../common/contributions.js';
+import { IAuthenticationAccessService } from '../../../services/authentication/browser/authenticationAccessService.js';
+import { IAuthenticationMcpAccessService } from '../../../services/authentication/browser/authenticationMcpAccessService.js';
+import { IAuthenticationMcpService } from '../../../services/authentication/browser/authenticationMcpService.js';
+import { IAuthenticationService } from '../../../services/authentication/common/authentication.js';
 import { IEditorService } from '../../../services/editor/common/editorService.js';
 import { IViewsService } from '../../../services/views/common/viewsService.js';
 import { IChatWidgetService } from '../../chat/browser/chat.js';
@@ -142,6 +147,9 @@ export class ListMcpServerCommand extends Action2 {
 	}
 }
 
+interface ActionItem extends IQuickPickItem {
+	action: 'start' | 'stop' | 'restart' | 'disconnect' | 'signout' | 'showOutput' | 'config' | 'configSampling' | 'samplingLog' | 'resources';
+}
 
 export class McpServerOptionsCommand extends Action2 {
 	constructor() {
@@ -160,6 +168,11 @@ export class McpServerOptionsCommand extends Action2 {
 		const editorService = accessor.get(IEditorService);
 		const commandService = accessor.get(ICommandService);
 		const samplingService = accessor.get(IMcpSamplingService);
+		const authenticationMcpService = accessor.get(IAuthenticationMcpService);
+		const authenticationMcpAccessService = accessor.get(IAuthenticationMcpAccessService);
+		const authenticationExtensionAccessService = accessor.get(IAuthenticationAccessService);
+		const authenticationService = accessor.get(IAuthenticationService);
+		const productService = accessor.get(IProductService);
 		const server = mcpService.servers.get().find(s => s.definition.id === id);
 		if (!server) {
 			return;
@@ -168,10 +181,6 @@ export class McpServerOptionsCommand extends Action2 {
 		const collection = mcpRegistry.collections.get().find(c => c.id === server.collection.id);
 		const serverDefinition = collection?.serverDefinitions.get().find(s => s.id === server.definition.id);
 
-		interface ActionItem extends IQuickPickItem {
-			action: 'start' | 'stop' | 'restart' | 'showOutput' | 'config' | 'configSampling' | 'samplingLog' | 'resources';
-		}
-
 		const items: (ActionItem | IQuickPickSeparator)[] = [];
 		const serverState = server.connectionState.get();
 
@@ -194,6 +203,18 @@ export class McpServerOptionsCommand extends Action2 {
 			});
 		}
 
+		const item = this._getAuthAction(
+			mcpRegistry,
+			authenticationMcpService,
+			authenticationMcpAccessService,
+			authenticationExtensionAccessService,
+			productService,
+			server.definition.id
+		);
+		if (item) {
+			items.push(item);
+		}
+
 		const configTarget = serverDefinition?.presentation?.origin || collection?.presentation?.origin;
 		if (configTarget) {
 			items.push({
@@ -254,6 +275,26 @@ export class McpServerOptionsCommand extends Action2 {
 				await server.stop();
 				await server.start({ isFromInteraction: true });
 				break;
+			case 'disconnect':
+				await this._handleAuth(
+					mcpRegistry,
+					authenticationMcpService,
+					authenticationMcpAccessService,
+					authenticationService,
+					server,
+					false
+				);
+				break;
+			case 'signout':
+				await this._handleAuth(
+					mcpRegistry,
+					authenticationMcpService,
+					authenticationMcpAccessService,
+					authenticationService,
+					server,
+					true
+				);
+				break;
 			case 'showOutput':
 				server.showOutput();
 				break;
@@ -278,6 +319,107 @@ export class McpServerOptionsCommand extends Action2 {
 				assertNever(pick.action);
 		}
 	}
+
+	private _getAuthAction(
+		mcpRegistry: IMcpRegistry,
+		authenticationMcpService: IAuthenticationMcpService,
+		authenticationMcpAccessService: IAuthenticationMcpAccessService,
+		authenticationAccessService: IAuthenticationAccessService,
+		productService: IProductService,
+		serverId: string
+	): ActionItem | undefined {
+		const providerId = mcpRegistry.getAuthenticationUsage(serverId);
+		if (!providerId) {
+			return undefined;
+		}
+		const preference = authenticationMcpService.getAccountPreference(serverId, providerId);
+		if (!preference) {
+			return undefined;
+		}
+		if (!authenticationMcpAccessService.isAccessAllowed(providerId, preference, serverId)) {
+			return undefined;
+		}
+		const allowedServers = this._getAllAllowedItems(
+			authenticationMcpAccessService,
+			authenticationAccessService,
+			productService,
+			providerId,
+			preference
+		);
+
+		// If there are multiple allowed servers/extensions, other things are using this provider
+		// so we show a disconnect action, otherwise we show a sign out action.
+		if (allowedServers.length > 1) {
+			return {
+				action: 'disconnect',
+				label: localize('mcp.disconnect', 'Disconnect Account'),
+				description: `(${preference})`,
+			};
+		}
+		return {
+			action: 'signout',
+			label: localize('mcp.signOut', 'Sign Out'),
+			description: `(${preference})`
+		};
+	}
+
+	// TODO@TylerLeonhardt: The fact that this function exists means that these classes could really use some refactoring...
+	private _getAllAllowedItems(
+		authenticationMcpAccessService: IAuthenticationMcpAccessService,
+		authenticationAccessService: IAuthenticationAccessService,
+		productService: IProductService,
+		providerId: string,
+		preference: string
+	) {
+		const trustedExtensionAuth = Array.isArray(productService.trustedExtensionAuthAccess) || !productService.trustedExtensionAuthAccess
+			? []
+			: productService.trustedExtensionAuthAccess[providerId] ?? [];
+		const trustedMcpAuth = Array.isArray(productService.trustedMcpAuthAccess) || !productService.trustedMcpAuthAccess
+			? []
+			: productService.trustedMcpAuthAccess[providerId] ?? [];
+
+		return [
+			...authenticationMcpAccessService.readAllowedMcpServers(providerId, preference).filter(s => !s.trusted),
+			...authenticationAccessService.readAllowedExtensions(providerId, preference).filter(e => !e.trusted),
+			...trustedExtensionAuth,
+			...trustedMcpAuth
+		];
+	}
+
+	private async _handleAuth(
+		mcpRegistry: IMcpRegistry,
+		authenticationMcpService: IAuthenticationMcpService,
+		authenticationMcpAccessService: IAuthenticationMcpAccessService,
+		authenticationService: IAuthenticationService,
+		server: IMcpServer,
+		signOut: boolean
+	) {
+		const providerId = mcpRegistry.getAuthenticationUsage(server.definition.id);
+		if (!providerId) {
+			return;
+		}
+		const preference = authenticationMcpService.getAccountPreference(server.definition.id, providerId);
+		if (!preference) {
+			return;
+		}
+		authenticationMcpAccessService.updateAllowedMcpServers(providerId, preference, [
+			{
+				id: server.definition.id,
+				name: server.definition.label,
+				allowed: false
+			}
+		]);
+		if (signOut) {
+			const accounts = await authenticationService.getAccounts(providerId);
+			const account = accounts.find(a => a.label === preference);
+			if (account) {
+				const sessions = await authenticationService.getSessions(providerId, undefined, { account });
+				for (const session of sessions) {
+					await authenticationService.removeSession(providerId, session.id);
+				}
+			}
+		}
+	}
 }
 
 export class MCPServerActionRendering extends Disposable implements IWorkbenchContribution {