Skip to content

Commit b84b028

Browse files
Authentication polish around account access (#252452)
This adopts the AuthenticationQueryService in: 1. Manage Trusted Extensions quick pick 2. An MCP server's Quick Pick (shows Sign Out or Disconnect account) 3. An MCP server's view (shows Sign Out or Disconnect account) 4. Depend on this service instead of some sily map in the MCP Registry 5. Have the `AuthenticationAccessService` own looking at the product.json so that implementers don't need to know about it This also contains a number of squashed bugs... 1. We were not cleaning up storage state when an extension or mcp server was uninstalled 2. The Extension Access service wasn't using `ExtensionIdentifier.toKey`
1 parent f462d2d commit b84b028

14 files changed

+1200
-317
lines changed

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

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -165,13 +165,11 @@ export class MainThreadMcp extends Disposable implements MainThreadMcpShape {
165165
if (sessions.length) {
166166
// If we have an existing session preference, use that. If not, we'll return any valid session at the end of this function.
167167
if (matchingAccountPreferenceSession && this.authenticationMCPServerAccessService.isAccessAllowed(providerId, matchingAccountPreferenceSession.account.label, server.id)) {
168-
this._mcpRegistry.setAuthenticationUsage(server.id, providerId);
169168
this.authenticationMCPServerUsageService.addAccountUsage(providerId, matchingAccountPreferenceSession.account.label, scopesSupported, server.id, server.label);
170169
return matchingAccountPreferenceSession.accessToken;
171170
}
172171
// If we only have one account for a single auth provider, lets just check if it's allowed and return it if it is.
173172
if (!provider.supportsMultipleAccounts && this.authenticationMCPServerAccessService.isAccessAllowed(providerId, sessions[0].account.label, server.id)) {
174-
this._mcpRegistry.setAuthenticationUsage(server.id, providerId);
175173
this.authenticationMCPServerUsageService.addAccountUsage(providerId, sessions[0].account.label, scopesSupported, server.id, server.label);
176174
return sessions[0].accessToken;
177175
}
@@ -205,7 +203,6 @@ export class MainThreadMcp extends Disposable implements MainThreadMcpShape {
205203
);
206204
}
207205

208-
this._mcpRegistry.setAuthenticationUsage(server.id, providerId);
209206
this.authenticationMCPServerAccessService.updateAllowedMcpServers(providerId, session.account.label, [{ id: server.id, name: server.label, allowed: true }]);
210207
this.authenticationMcpServersService.updateAccountPreference(server.id, providerId, session.account);
211208
this.authenticationMCPServerUsageService.addAccountUsage(providerId, session.account.label, scopesSupported, server.id, server.label);

src/vs/workbench/contrib/authentication/browser/actions/manageTrustedExtensionsForAccountAction.ts

Lines changed: 82 additions & 106 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,9 @@ import { Action2 } from '../../../../../platform/actions/common/actions.js';
1212
import { ICommandService } from '../../../../../platform/commands/common/commands.js';
1313
import { IDialogService } from '../../../../../platform/dialogs/common/dialogs.js';
1414
import { IInstantiationService, ServicesAccessor } from '../../../../../platform/instantiation/common/instantiation.js';
15-
import { IProductService } from '../../../../../platform/product/common/productService.js';
1615
import { IQuickInputService, IQuickPickItem, IQuickPickSeparator } from '../../../../../platform/quickinput/common/quickInput.js';
17-
import { IAuthenticationAccessService } from '../../../../services/authentication/browser/authenticationAccessService.js';
18-
import { IAuthenticationUsageService } from '../../../../services/authentication/browser/authenticationUsageService.js';
1916
import { AllowedExtension, IAuthenticationService } from '../../../../services/authentication/common/authentication.js';
17+
import { IAuthenticationQueryService, IAccountQuery } from '../../../../services/authentication/common/authenticationQuery.js';
2018
import { IExtensionService } from '../../../../services/extensions/common/extensions.js';
2119

2220
export class ManageTrustedExtensionsForAccountAction extends Action2 {
@@ -42,133 +40,110 @@ interface TrustedExtensionsQuickPickItem extends IQuickPickItem {
4240

4341
class ManageTrustedExtensionsForAccountActionImpl {
4442
constructor(
45-
@IProductService private readonly _productService: IProductService,
4643
@IExtensionService private readonly _extensionService: IExtensionService,
4744
@IDialogService private readonly _dialogService: IDialogService,
4845
@IQuickInputService private readonly _quickInputService: IQuickInputService,
4946
@IAuthenticationService private readonly _authenticationService: IAuthenticationService,
50-
@IAuthenticationUsageService private readonly _authenticationUsageService: IAuthenticationUsageService,
51-
@IAuthenticationAccessService private readonly _authenticationAccessService: IAuthenticationAccessService,
47+
@IAuthenticationQueryService private readonly _authenticationQueryService: IAuthenticationQueryService,
5248
@ICommandService private readonly _commandService: ICommandService
5349
) { }
5450

5551
async run(options?: { providerId: string; accountLabel: string }) {
56-
const { providerId, accountLabel } = await this._resolveProviderAndAccountLabel(options?.providerId, options?.accountLabel);
57-
if (!providerId || !accountLabel) {
52+
const accountQuery = await this._resolveAccountQuery(options?.providerId, options?.accountLabel);
53+
if (!accountQuery) {
5854
return;
5955
}
6056

61-
const items = await this._getItems(providerId, accountLabel);
57+
const items = await this._getItems(accountQuery);
6258
if (!items.length) {
6359
return;
6460
}
65-
const disposables = new DisposableStore();
66-
const picker = this._createQuickPick(disposables, providerId, accountLabel);
61+
const picker = this._createQuickPick(accountQuery);
6762
picker.items = items;
6863
picker.selectedItems = items.filter((i): i is TrustedExtensionsQuickPickItem => i.type !== 'separator' && !!i.picked);
6964
picker.show();
7065
}
7166

72-
private async _resolveProviderAndAccountLabel(providerId: string | undefined, accountLabel: string | undefined) {
73-
if (!providerId || !accountLabel) {
74-
const accounts = new Array<{ providerId: string; providerLabel: string; accountLabel: string }>();
75-
for (const id of this._authenticationService.getProviderIds()) {
76-
const providerLabel = this._authenticationService.getProvider(id).label;
77-
const sessions = await this._authenticationService.getSessions(id);
78-
const uniqueAccountLabels = new Set<string>();
79-
for (const session of sessions) {
80-
if (!uniqueAccountLabels.has(session.account.label)) {
81-
uniqueAccountLabels.add(session.account.label);
82-
accounts.push({ providerId: id, providerLabel, accountLabel: session.account.label });
83-
}
84-
}
85-
}
67+
//#region Account Query Resolution
8668

87-
const pick = await this._quickInputService.pick(
88-
accounts.map(account => ({
89-
providerId: account.providerId,
90-
label: account.accountLabel,
91-
description: account.providerLabel
92-
})),
93-
{
94-
placeHolder: localize('pickAccount', "Pick an account to manage trusted extensions for"),
95-
matchOnDescription: true,
96-
}
97-
);
98-
99-
if (pick) {
100-
providerId = pick.providerId;
101-
accountLabel = pick.label;
102-
} else {
103-
return { providerId: undefined, accountLabel: undefined };
104-
}
69+
private async _resolveAccountQuery(providerId: string | undefined, accountLabel: string | undefined): Promise<IAccountQuery | undefined> {
70+
if (providerId && accountLabel) {
71+
return this._authenticationQueryService.provider(providerId).account(accountLabel);
10572
}
106-
return { providerId, accountLabel };
73+
74+
const accounts = await this._getAllAvailableAccounts();
75+
const pick = await this._quickInputService.pick(accounts, {
76+
placeHolder: localize('pickAccount', "Pick an account to manage trusted extensions for"),
77+
matchOnDescription: true,
78+
});
79+
80+
return pick ? this._authenticationQueryService.provider(pick.providerId).account(pick.label) : undefined;
10781
}
10882

109-
private async _getItems(providerId: string, accountLabel: string) {
110-
let allowedExtensions = this._authenticationAccessService.readAllowedExtensions(providerId, accountLabel);
111-
// only include extensions that are installed
112-
const resolvedExtensions = await Promise.all(allowedExtensions.map(ext => this._extensionService.getExtension(ext.id)));
113-
allowedExtensions = resolvedExtensions
114-
.map((ext, i) => ext ? allowedExtensions[i] : undefined)
115-
.filter(ext => !!ext);
116-
const trustedExtensionAuthAccess = this._productService.trustedExtensionAuthAccess;
117-
const trustedExtensionIds =
118-
// Case 1: trustedExtensionAuthAccess is an array
119-
Array.isArray(trustedExtensionAuthAccess)
120-
? trustedExtensionAuthAccess
121-
// Case 2: trustedExtensionAuthAccess is an object
122-
: typeof trustedExtensionAuthAccess === 'object'
123-
? trustedExtensionAuthAccess[providerId] ?? []
124-
: [];
125-
for (const extensionId of trustedExtensionIds) {
126-
const allowedExtension = allowedExtensions.find(ext => ext.id === extensionId);
127-
if (!allowedExtension) {
128-
// Add the extension to the allowedExtensions list
129-
const extension = await this._extensionService.getExtension(extensionId);
130-
if (extension) {
131-
allowedExtensions.push({
132-
id: extensionId,
133-
name: extension.displayName || extension.name,
134-
allowed: true,
135-
trusted: true
83+
private async _getAllAvailableAccounts() {
84+
const accounts = [];
85+
for (const providerId of this._authenticationService.getProviderIds()) {
86+
const provider = this._authenticationService.getProvider(providerId);
87+
const sessions = await this._authenticationService.getSessions(providerId);
88+
const uniqueLabels = new Set<string>();
89+
90+
for (const session of sessions) {
91+
if (!uniqueLabels.has(session.account.label)) {
92+
uniqueLabels.add(session.account.label);
93+
accounts.push({
94+
providerId,
95+
label: session.account.label,
96+
description: provider.label
13697
});
13798
}
138-
} else {
139-
// Update the extension to be allowed
140-
allowedExtension.allowed = true;
141-
allowedExtension.trusted = true;
14299
}
143100
}
101+
return accounts;
102+
}
144103

145-
if (!allowedExtensions.length) {
146-
this._dialogService.info(localize('noTrustedExtensions', "This account has not been used by any extensions."));
147-
return [];
148-
}
104+
//#endregion
149105

150-
const usages = this._authenticationUsageService.readAccountUsages(providerId, accountLabel);
151-
const trustedExtensions = [];
152-
const otherExtensions = [];
153-
for (const extension of allowedExtensions) {
154-
const usage = usages.find(usage => extension.id === usage.extensionId);
155-
extension.lastUsed = usage?.lastUsed;
156-
if (extension.trusted) {
157-
trustedExtensions.push(extension);
158-
} else {
159-
otherExtensions.push(extension);
106+
//#region Item Retrieval and Quick Pick Creation
107+
108+
private async _getItems(accountQuery: IAccountQuery) {
109+
const allowedExtensions = accountQuery.extensions().getAllowedExtensions();
110+
const extensionIdToDisplayName = new Map<string, string>();
111+
112+
// Get display names for all allowed extensions
113+
const resolvedExtensions = await Promise.all(allowedExtensions.map(ext => this._extensionService.getExtension(ext.id)));
114+
resolvedExtensions.forEach((resolved, i) => {
115+
if (resolved) {
116+
extensionIdToDisplayName.set(allowedExtensions[i].id, resolved.displayName || resolved.name);
160117
}
118+
});
119+
120+
// Filter out extensions that are not currently installed and enrich with display names
121+
const filteredExtensions = allowedExtensions
122+
.filter(ext => extensionIdToDisplayName.has(ext.id))
123+
.map(ext => {
124+
const usage = accountQuery.extension(ext.id).getUsage();
125+
return {
126+
...ext,
127+
// Use the extension display name from the extension service
128+
name: extensionIdToDisplayName.get(ext.id)!,
129+
lastUsed: usage.length > 0 ? Math.max(...usage.map(u => u.lastUsed)) : ext.lastUsed
130+
};
131+
});
132+
133+
if (!filteredExtensions.length) {
134+
this._dialogService.info(localize('noTrustedExtensions', "This account has not been used by any extensions."));
135+
return [];
161136
}
162137

138+
const trustedExtensions = filteredExtensions.filter(e => e.trusted);
139+
const otherExtensions = filteredExtensions.filter(e => !e.trusted);
163140
const sortByLastUsed = (a: AllowedExtension, b: AllowedExtension) => (b.lastUsed || 0) - (a.lastUsed || 0);
164141

165-
const items = [
142+
return [
166143
...otherExtensions.sort(sortByLastUsed).map(this._toQuickPickItem),
167144
{ type: 'separator', label: localize('trustedExtensions', "Trusted by Microsoft") } satisfies IQuickPickSeparator,
168145
...trustedExtensions.sort(sortByLastUsed).map(this._toQuickPickItem)
169146
];
170-
171-
return items;
172147
}
173148

174149
private _toQuickPickItem(extension: AllowedExtension): TrustedExtensionsQuickPickItem {
@@ -196,38 +171,39 @@ class ManageTrustedExtensionsForAccountActionImpl {
196171
};
197172
}
198173

199-
private _createQuickPick(disposableStore: DisposableStore, providerId: string, accountLabel: string) {
174+
private _createQuickPick(accountQuery: IAccountQuery) {
175+
const disposableStore = new DisposableStore();
200176
const quickPick = disposableStore.add(this._quickInputService.createQuickPick<TrustedExtensionsQuickPickItem>({ useSeparators: true }));
177+
178+
// Configure quick pick
201179
quickPick.canSelectMany = true;
202180
quickPick.customButton = true;
203181
quickPick.customLabel = localize('manageTrustedExtensions.cancel', 'Cancel');
204-
205182
quickPick.title = localize('manageTrustedExtensions', "Manage Trusted Extensions");
206183
quickPick.placeholder = localize('manageExtensions', "Choose which extensions can access this account");
207184

185+
// Set up event handlers
208186
disposableStore.add(quickPick.onDidAccept(() => {
209187
const updatedAllowedList = quickPick.items
210188
.filter((item): item is TrustedExtensionsQuickPickItem => item.type !== 'separator')
211189
.map(i => i.extension);
212190

213191
const allowedExtensionsSet = new Set(quickPick.selectedItems.map(i => i.extension));
214-
updatedAllowedList.forEach(extension => {
215-
extension.allowed = allowedExtensionsSet.has(extension);
216-
});
217-
this._authenticationAccessService.updateAllowedExtensions(providerId, accountLabel, updatedAllowedList);
192+
for (const extension of updatedAllowedList) {
193+
const allowed = allowedExtensionsSet.has(extension);
194+
accountQuery.extension(extension.id).setAccessAllowed(allowed, extension.name);
195+
}
218196
quickPick.hide();
219197
}));
220198

221-
disposableStore.add(quickPick.onDidHide(() => {
222-
disposableStore.dispose();
223-
}));
224-
225-
disposableStore.add(quickPick.onDidCustom(() => {
226-
quickPick.hide();
227-
}));
199+
disposableStore.add(quickPick.onDidHide(() => disposableStore.dispose()));
200+
disposableStore.add(quickPick.onDidCustom(() => quickPick.hide()));
228201
disposableStore.add(quickPick.onDidTriggerItemButton(e =>
229-
this._commandService.executeCommand('_manageAccountPreferencesForExtension', e.item.extension.id, providerId)
202+
this._commandService.executeCommand('_manageAccountPreferencesForExtension', e.item.extension.id, accountQuery.providerId)
230203
));
204+
231205
return quickPick;
232206
}
207+
208+
//#endregion
233209
}

0 commit comments

Comments
 (0)