Skip to content

Commit 591199d

Browse files
sandy081sandy081
authored
do not fetch policy data if already fetched (#293996)
* do not fetch policy data if already fetched * reset mcp data when not configured and refetch on polling * refetch always * reschedule when window has no focus * another change * - handle 429 for mcp registry - do not fetch entitlements on refreshing policies
1 parent 7cc9148 commit 591199d

File tree

1 file changed

+90
-41
lines changed

1 file changed

+90
-41
lines changed

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

Lines changed: 90 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ import { IHostService } from '../../host/browser/host.js';
1818
import { IConfigurationService } from '../../../../platform/configuration/common/configuration.js';
1919
import { getErrorMessage } from '../../../../base/common/errors.js';
2020
import { IDefaultAccount, IDefaultAccountAuthenticationProvider, IEntitlementsData, IPolicyData } from '../../../../base/common/defaultAccount.js';
21-
import { isString, Mutable } from '../../../../base/common/types.js';
21+
import { isString, isUndefined, Mutable } from '../../../../base/common/types.js';
2222
import { IStorageService, StorageScope, StorageTarget } from '../../../../platform/storage/common/storage.js';
2323
import { IWorkbenchEnvironmentService } from '../../environment/common/environmentService.js';
2424
import { isWeb } from '../../../../base/common/platform.js';
@@ -60,7 +60,7 @@ const enum DefaultAccountStatus {
6060

6161
const CONTEXT_DEFAULT_ACCOUNT_STATE = new RawContextKey<string>('defaultAccountStatus', DefaultAccountStatus.Uninitialized);
6262
const CACHED_POLICY_DATA_KEY = 'defaultAccount.cachedPolicyData';
63-
const ACCOUNT_DATA_POLL_INTERVAL_MS = 15 * 60 * 1000; // 15 minutes
63+
const ACCOUNT_DATA_POLL_INTERVAL_MS = 60 * 60 * 1000; // 1 hour
6464

6565
interface ITokenEntitlementsResponse {
6666
token: string;
@@ -69,7 +69,7 @@ interface ITokenEntitlementsResponse {
6969
interface IMcpRegistryProvider {
7070
readonly url: string;
7171
readonly registry_access: 'allow_all' | 'registry_only';
72-
readonly owner: {
72+
readonly owner?: {
7373
readonly login: string;
7474
readonly id: number;
7575
readonly type: string;
@@ -189,6 +189,8 @@ export class DefaultAccountService extends Disposable implements IDefaultAccount
189189
interface IAccountPolicyData {
190190
readonly accountId: string;
191191
readonly policyData: IPolicyData;
192+
readonly isTokenEntitlementsDataFetched: boolean;
193+
readonly isMcpRegistryDataFetched: boolean;
192194
}
193195

194196
interface IDefaultAccountData {
@@ -226,7 +228,7 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid
226228
private initialized = false;
227229
private readonly initPromise: Promise<void>;
228230
private readonly updateThrottler = this._register(new ThrottledDelayer(100));
229-
private readonly accountDataPollScheduler = this._register(new RunOnceScheduler(() => this.updateDefaultAccount(), ACCOUNT_DATA_POLL_INTERVAL_MS));
231+
private readonly accountDataPollScheduler = this._register(new RunOnceScheduler(() => this.refetchPolicyData(), ACCOUNT_DATA_POLL_INTERVAL_MS));
230232

231233
constructor(
232234
private readonly defaultAccountConfig: IDefaultAccountConfig,
@@ -259,7 +261,7 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid
259261
const { accountId, policyData } = JSON.parse(cached);
260262
if (accountId && policyData) {
261263
this.logService.debug('[DefaultAccount] Initializing with cached policy data');
262-
return { accountId, policyData };
264+
return { accountId, policyData, isTokenEntitlementsDataFetched: false, isMcpRegistryDataFetched: false };
263265
}
264266
} catch (error) {
265267
this.logService.error('[DefaultAccount] Failed to parse cached policy data', getErrorMessage(error));
@@ -282,7 +284,7 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid
282284
}
283285

284286
this.logService.debug('[DefaultAccount] Starting initialization');
285-
await this.doUpdateDefaultAccount();
287+
await this.doUpdateDefaultAccount(false, false);
286288
this.logService.debug('[DefaultAccount] Initialization complete');
287289

288290
this._register(this.onDidChangeDefaultAccount(account => {
@@ -330,41 +332,51 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid
330332
}));
331333

332334
this._register(this.hostService.onDidChangeFocus(focused => {
333-
if (focused && this._defaultAccount) {
334-
// Update default account when window gets focused
335+
// Refresh default account when window gets focused and we have cached policy data, which likely means we haven't successfully fetched data since the last time the window was focused (e.g. due to network issues), so we should try again to fetch the data.
336+
if (focused && this._defaultAccount && this._policyData && (!this._policyData.isMcpRegistryDataFetched || !this._policyData.isTokenEntitlementsDataFetched)) {
335337
this.accountDataPollScheduler.cancel();
336338
this.logService.debug('[DefaultAccount] Window focused, updating default account');
337-
this.updateDefaultAccount();
339+
this.refresh(true);
338340
}
339341
}));
340342
}
341343

342-
async refresh(): Promise<IDefaultAccount | null> {
344+
async refresh(donotFetchEntitlements: boolean = false): Promise<IDefaultAccount | null> {
343345
if (!this.initialized) {
344346
await this.initPromise;
345347
return this.defaultAccount;
346348
}
347349

348350
this.logService.debug('[DefaultAccount] Refreshing default account');
349-
await this.updateDefaultAccount();
351+
await this.updateDefaultAccount(false, donotFetchEntitlements);
350352
return this.defaultAccount;
351353
}
352354

353-
private async updateDefaultAccount(): Promise<void> {
354-
await this.updateThrottler.trigger(() => this.doUpdateDefaultAccount());
355+
private async refetchPolicyData(): Promise<void> {
356+
if (!this.hostService.hasFocus) {
357+
this.scheduleAccountDataPoll();
358+
this.logService.debug('[DefaultAccount] Skipping refetching policy data because window is not focused');
359+
return;
360+
}
361+
this.logService.debug('[DefaultAccount] Refetching policy data for current default account');
362+
await this.updateDefaultAccount(true);
363+
}
364+
365+
private async updateDefaultAccount(donotUseCache: boolean = false, donotFetchEntitlements: boolean = false): Promise<void> {
366+
await this.updateThrottler.trigger(() => this.doUpdateDefaultAccount(donotUseCache, donotFetchEntitlements));
355367
}
356368

357-
private async doUpdateDefaultAccount(): Promise<void> {
369+
private async doUpdateDefaultAccount(donotUseCache: boolean, donotFetchEntitlements: boolean): Promise<void> {
358370
try {
359-
const defaultAccount = await this.fetchDefaultAccount();
371+
const defaultAccount = await this.fetchDefaultAccount(donotUseCache, donotFetchEntitlements);
360372
this.setDefaultAccount(defaultAccount);
361373
this.scheduleAccountDataPoll();
362374
} catch (error) {
363375
this.logService.error('[DefaultAccount] Error while updating default account', getErrorMessage(error));
364376
}
365377
}
366378

367-
private async fetchDefaultAccount(): Promise<IDefaultAccountData | null> {
379+
private async fetchDefaultAccount(donotUseCache: boolean, donotFetchEntitlements: boolean): Promise<IDefaultAccountData | null> {
368380
const defaultAccountProvider = this.getDefaultAccountAuthenticationProvider();
369381
this.logService.debug('[DefaultAccount] Default account provider ID:', defaultAccountProvider.id);
370382

@@ -374,7 +386,7 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid
374386
return null;
375387
}
376388

377-
return await this.getDefaultAccountForAuthenticationProvider(defaultAccountProvider);
389+
return await this.getDefaultAccountForAuthenticationProvider(defaultAccountProvider, donotUseCache, donotFetchEntitlements);
378390
}
379391

380392
private setDefaultAccount(account: IDefaultAccountData | null): void {
@@ -437,7 +449,7 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid
437449
return result;
438450
}
439451

440-
private async getDefaultAccountForAuthenticationProvider(authenticationProvider: IDefaultAccountAuthenticationProvider): Promise<IDefaultAccountData | null> {
452+
private async getDefaultAccountForAuthenticationProvider(authenticationProvider: IDefaultAccountAuthenticationProvider, donotUseCache: boolean, donotFetchEntitlements: boolean): Promise<IDefaultAccountData | null> {
441453
try {
442454
this.logService.debug('[DefaultAccount] Getting Default Account from authenticated sessions for provider:', authenticationProvider.id);
443455
const sessions = await this.findMatchingProviderSession(authenticationProvider.id, this.defaultAccountConfig.authenticationProvider.scopes);
@@ -447,33 +459,42 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid
447459
return null;
448460
}
449461

450-
return this.getDefaultAccountFromAuthenticatedSessions(authenticationProvider, sessions);
462+
return this.getDefaultAccountFromAuthenticatedSessions(authenticationProvider, sessions, donotUseCache, donotFetchEntitlements);
451463
} catch (error) {
452464
this.logService.error('[DefaultAccount] Failed to get default account for provider:', authenticationProvider.id, getErrorMessage(error));
453465
return null;
454466
}
455467
}
456468

457-
private async getDefaultAccountFromAuthenticatedSessions(authenticationProvider: IDefaultAccountAuthenticationProvider, sessions: AuthenticationSession[]): Promise<IDefaultAccountData | null> {
469+
private async getDefaultAccountFromAuthenticatedSessions(authenticationProvider: IDefaultAccountAuthenticationProvider, sessions: AuthenticationSession[], donotUseCache: boolean, donotFetchEntitlements: boolean): Promise<IDefaultAccountData | null> {
458470
try {
459471
const accountId = sessions[0].account.id;
472+
const accountPolicyData = this._policyData?.accountId === accountId ? this._policyData : undefined;
473+
460474
const [entitlementsData, tokenEntitlementsData] = await Promise.all([
461-
this.getEntitlements(sessions),
462-
this.getTokenEntitlements(sessions),
475+
donotFetchEntitlements && accountPolicyData && this._defaultAccount ? this._defaultAccount.defaultAccount.entitlementsData : this.getEntitlements(sessions),
476+
this.getTokenEntitlements(sessions, donotUseCache ? undefined : accountPolicyData),
463477
]);
464478

465-
let policyData: Mutable<IPolicyData> | undefined = this._policyData?.accountId === accountId ? { ...this._policyData.policyData } : undefined;
479+
let isTokenEntitlementsDataFetched = false;
480+
let isMcpRegistryDataFetched = false;
481+
let policyData: Mutable<IPolicyData> | undefined = accountPolicyData?.policyData ? { ...accountPolicyData.policyData } : undefined;
466482
if (tokenEntitlementsData) {
483+
isTokenEntitlementsDataFetched = true;
467484
policyData = policyData ?? {};
468485
policyData.chat_agent_enabled = tokenEntitlementsData.chat_agent_enabled;
469486
policyData.chat_preview_features_enabled = tokenEntitlementsData.chat_preview_features_enabled;
470487
policyData.mcp = tokenEntitlementsData.mcp;
471488
if (policyData.mcp) {
472-
const mcpRegistryProvider = await this.getMcpRegistryProvider(sessions);
473-
if (mcpRegistryProvider) {
474-
policyData.mcpRegistryUrl = mcpRegistryProvider.url;
475-
policyData.mcpAccess = mcpRegistryProvider.registry_access;
489+
const mcpRegistryProvider = await this.getMcpRegistryProvider(sessions, donotUseCache ? undefined : accountPolicyData);
490+
if (!isUndefined(mcpRegistryProvider)) {
491+
isMcpRegistryDataFetched = true;
492+
policyData.mcpRegistryUrl = mcpRegistryProvider?.url;
493+
policyData.mcpAccess = mcpRegistryProvider?.registry_access;
476494
}
495+
} else {
496+
policyData.mcpRegistryUrl = undefined;
497+
policyData.mcpAccess = undefined;
477498
}
478499
}
479500

@@ -484,7 +505,7 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid
484505
entitlementsData,
485506
};
486507
this.logService.debug('[DefaultAccount] Successfully created default account for provider:', authenticationProvider.id);
487-
return { defaultAccount, policyData: policyData ? { accountId, policyData } : null };
508+
return { defaultAccount, policyData: policyData ? { accountId, policyData, isTokenEntitlementsDataFetched, isMcpRegistryDataFetched } : null };
488509
} catch (error) {
489510
this.logService.error('[DefaultAccount] Failed to create default account for provider:', authenticationProvider.id, getErrorMessage(error));
490511
return null;
@@ -539,7 +560,15 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid
539560
return expectedScopes.every(scope => scopes.includes(scope));
540561
}
541562

542-
private async getTokenEntitlements(sessions: AuthenticationSession[]): Promise<Partial<IPolicyData> | undefined> {
563+
private async getTokenEntitlements(sessions: AuthenticationSession[], accountPolicyData: IAccountPolicyData | undefined): Promise<Partial<IPolicyData> | undefined> {
564+
if (accountPolicyData?.isTokenEntitlementsDataFetched) {
565+
this.logService.debug('[DefaultAccount] Using last fetched token entitlements data');
566+
return accountPolicyData.policyData;
567+
}
568+
return await this.requestTokenEntitlements(sessions);
569+
}
570+
571+
private async requestTokenEntitlements(sessions: AuthenticationSession[]): Promise<Partial<IPolicyData> | undefined> {
543572
const tokenEntitlementsUrl = this.getTokenEntitlementUrl();
544573
if (!tokenEntitlementsUrl) {
545574
this.logService.debug('[DefaultAccount] No token entitlements URL found');
@@ -610,11 +639,19 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid
610639
return undefined;
611640
}
612641

613-
private async getMcpRegistryProvider(sessions: AuthenticationSession[]): Promise<IMcpRegistryProvider | undefined> {
642+
private async getMcpRegistryProvider(sessions: AuthenticationSession[], accountPolicyData: IAccountPolicyData | undefined): Promise<IMcpRegistryProvider | null | undefined> {
643+
if (accountPolicyData?.isMcpRegistryDataFetched) {
644+
this.logService.debug('[DefaultAccount] Using last fetched MCP registry data');
645+
return accountPolicyData.policyData.mcpRegistryUrl && accountPolicyData.policyData.mcpAccess ? { url: accountPolicyData.policyData.mcpRegistryUrl, registry_access: accountPolicyData.policyData.mcpAccess } : null;
646+
}
647+
return await this.requestMcpRegistryProvider(sessions);
648+
}
649+
650+
private async requestMcpRegistryProvider(sessions: AuthenticationSession[]): Promise<IMcpRegistryProvider | null | undefined> {
614651
const mcpRegistryDataUrl = this.getMcpRegistryDataUrl();
615652
if (!mcpRegistryDataUrl) {
616653
this.logService.debug('[DefaultAccount] No MCP registry data URL found');
617-
return undefined;
654+
return null;
618655
}
619656

620657
this.logService.debug('[DefaultAccount] Fetching MCP registry data from:', mcpRegistryDataUrl);
@@ -624,21 +661,34 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid
624661
}
625662

626663
if (response.res.statusCode && response.res.statusCode !== 200) {
627-
this.logService.trace(`[DefaultAccount] unexpected status code ${response.res.statusCode} while fetching MCP registry data`);
664+
if (response.res.statusCode === 401) {
665+
this.logService.debug('[DefaultAccount] Unauthorized (401) when fetching MCP registry data, treating as no registry available');
666+
return null;
667+
}
668+
if (response.res.statusCode === 404) {
669+
this.logService.debug('[DefaultAccount] MCP registry endpoint not found (404), treating as no registry available');
670+
return null;
671+
}
672+
if (response.res.statusCode === 429) {
673+
this.logService.debug(`[DefaultAccount] Received 429 Too Many Requests for MCP registry data, treating as no registry available and not retrying immediately.`);
674+
return null;
675+
}
676+
this.logService.debug(`[DefaultAccount] unexpected status code ${response.res.statusCode} while fetching MCP registry data`);
628677
return undefined;
629678
}
630679

631680
try {
632681
const data = await asJson<IMcpRegistryResponse>(response);
633682
if (data) {
634683
this.logService.debug('Fetched MCP registry providers', data.mcp_registries);
635-
return data.mcp_registries[0];
684+
return data.mcp_registries[0] ?? null;
636685
}
637-
this.logService.debug('Failed to fetch MCP registry providers', 'No data returned');
686+
this.logService.debug('No MCP registry providers content found in response');
687+
return null;
638688
} catch (error) {
639689
this.logService.error('Failed to fetch MCP registry providers', getErrorMessage(error));
690+
return undefined;
640691
}
641-
return undefined;
642692
}
643693

644694
private async request(url: string, type: 'GET', body: undefined, sessions: AuthenticationSession[], token: CancellationToken): Promise<IRequestContext | undefined>;
@@ -664,14 +714,18 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid
664714

665715
const status = response.res.statusCode;
666716
if (status && status !== 200) {
717+
if (status === 429) {
718+
this.logService.warn(`[DefaultAccount] Received 429 Too Many Requests for ${url}.`);
719+
return response;
720+
}
667721
lastResponse = response;
668722
continue; // try next session
669723
}
670724

671725
return response;
672726
} catch (error) {
673727
if (!token.isCancellationRequested) {
674-
this.logService.error(`[chat entitlement] request: error ${error}`);
728+
this.logService.error(`[DefaultAccount] request: error ${error}`, url);
675729
}
676730
}
677731
}
@@ -681,11 +735,6 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid
681735
return undefined;
682736
}
683737

684-
if (lastResponse.res.statusCode && lastResponse.res.statusCode !== 200) {
685-
this.logService.trace(`[DefaultAccount]: unexpected status code ${lastResponse.res.statusCode} for request`, url);
686-
return undefined;
687-
}
688-
689738
return lastResponse;
690739
}
691740

0 commit comments

Comments
 (0)