diff --git a/src/vs/workbench/services/accounts/browser/defaultAccount.ts b/src/vs/workbench/services/accounts/browser/defaultAccount.ts index f0a645aea45..66b782d7c91 100644 --- a/src/vs/workbench/services/accounts/browser/defaultAccount.ts +++ b/src/vs/workbench/services/accounts/browser/defaultAccount.ts @@ -18,7 +18,7 @@ import { IHostService } from '../../host/browser/host.js'; import { IConfigurationService } from '../../../../platform/configuration/common/configuration.js'; import { getErrorMessage } from '../../../../base/common/errors.js'; import { IDefaultAccount, IDefaultAccountAuthenticationProvider, IEntitlementsData, IPolicyData } from '../../../../base/common/defaultAccount.js'; -import { isString, Mutable } from '../../../../base/common/types.js'; +import { isString, isUndefined, Mutable } from '../../../../base/common/types.js'; import { IStorageService, StorageScope, StorageTarget } from '../../../../platform/storage/common/storage.js'; import { IWorkbenchEnvironmentService } from '../../environment/common/environmentService.js'; import { isWeb } from '../../../../base/common/platform.js'; @@ -60,7 +60,7 @@ const enum DefaultAccountStatus { const CONTEXT_DEFAULT_ACCOUNT_STATE = new RawContextKey('defaultAccountStatus', DefaultAccountStatus.Uninitialized); const CACHED_POLICY_DATA_KEY = 'defaultAccount.cachedPolicyData'; -const ACCOUNT_DATA_POLL_INTERVAL_MS = 15 * 60 * 1000; // 15 minutes +const ACCOUNT_DATA_POLL_INTERVAL_MS = 60 * 60 * 1000; // 1 hour interface ITokenEntitlementsResponse { token: string; @@ -69,7 +69,7 @@ interface ITokenEntitlementsResponse { interface IMcpRegistryProvider { readonly url: string; readonly registry_access: 'allow_all' | 'registry_only'; - readonly owner: { + readonly owner?: { readonly login: string; readonly id: number; readonly type: string; @@ -189,6 +189,8 @@ export class DefaultAccountService extends Disposable implements IDefaultAccount interface IAccountPolicyData { readonly accountId: string; readonly policyData: IPolicyData; + readonly isTokenEntitlementsDataFetched: boolean; + readonly isMcpRegistryDataFetched: boolean; } interface IDefaultAccountData { @@ -226,7 +228,7 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid private initialized = false; private readonly initPromise: Promise; private readonly updateThrottler = this._register(new ThrottledDelayer(100)); - private readonly accountDataPollScheduler = this._register(new RunOnceScheduler(() => this.updateDefaultAccount(), ACCOUNT_DATA_POLL_INTERVAL_MS)); + private readonly accountDataPollScheduler = this._register(new RunOnceScheduler(() => this.refetchDefaultAccount(), ACCOUNT_DATA_POLL_INTERVAL_MS)); constructor( private readonly defaultAccountConfig: IDefaultAccountConfig, @@ -259,7 +261,7 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid const { accountId, policyData } = JSON.parse(cached); if (accountId && policyData) { this.logService.debug('[DefaultAccount] Initializing with cached policy data'); - return { accountId, policyData }; + return { accountId, policyData, isTokenEntitlementsDataFetched: false, isMcpRegistryDataFetched: false }; } } catch (error) { this.logService.error('[DefaultAccount] Failed to parse cached policy data', getErrorMessage(error)); @@ -282,7 +284,7 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid } this.logService.debug('[DefaultAccount] Starting initialization'); - await this.doUpdateDefaultAccount(); + await this.doUpdateDefaultAccount(false); this.logService.debug('[DefaultAccount] Initialization complete'); this._register(this.onDidChangeDefaultAccount(account => { @@ -330,11 +332,11 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid })); this._register(this.hostService.onDidChangeFocus(focused => { - if (focused && this._defaultAccount) { - // Update default account when window gets focused + // Refresh default account when window gets focused and policy data is not fully fetched, to ensure we have the latest policy data. + if (focused && this._policyData && (!this._policyData.isMcpRegistryDataFetched || !this._policyData.isTokenEntitlementsDataFetched)) { this.accountDataPollScheduler.cancel(); this.logService.debug('[DefaultAccount] Window focused, updating default account'); - this.updateDefaultAccount(); + this.refresh(); } })); } @@ -350,13 +352,23 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid return this.defaultAccount; } - private async updateDefaultAccount(): Promise { - await this.updateThrottler.trigger(() => this.doUpdateDefaultAccount()); + private async refetchDefaultAccount(): Promise { + if (!this.hostService.hasFocus) { + this.scheduleAccountDataPoll(); + this.logService.debug('[DefaultAccount] Skipping refetching default account because window is not focused'); + return; + } + this.logService.debug('[DefaultAccount] Refetching default account'); + await this.updateDefaultAccount(true); } - private async doUpdateDefaultAccount(): Promise { + private async updateDefaultAccount(donotUseLastFetchedData: boolean = false): Promise { + await this.updateThrottler.trigger(() => this.doUpdateDefaultAccount(donotUseLastFetchedData)); + } + + private async doUpdateDefaultAccount(donotUseLastFetchedData: boolean): Promise { try { - const defaultAccount = await this.fetchDefaultAccount(); + const defaultAccount = await this.fetchDefaultAccount(donotUseLastFetchedData); this.setDefaultAccount(defaultAccount); this.scheduleAccountDataPoll(); } catch (error) { @@ -364,7 +376,7 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid } } - private async fetchDefaultAccount(): Promise { + private async fetchDefaultAccount(donotUseLastFetchedData: boolean): Promise { const defaultAccountProvider = this.getDefaultAccountAuthenticationProvider(); this.logService.debug('[DefaultAccount] Default account provider ID:', defaultAccountProvider.id); @@ -374,7 +386,7 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid return null; } - return await this.getDefaultAccountForAuthenticationProvider(defaultAccountProvider); + return await this.getDefaultAccountForAuthenticationProvider(defaultAccountProvider, donotUseLastFetchedData); } private setDefaultAccount(account: IDefaultAccountData | null): void { @@ -437,7 +449,7 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid return result; } - private async getDefaultAccountForAuthenticationProvider(authenticationProvider: IDefaultAccountAuthenticationProvider): Promise { + private async getDefaultAccountForAuthenticationProvider(authenticationProvider: IDefaultAccountAuthenticationProvider, donotUseLastFetchedData: boolean): Promise { try { this.logService.debug('[DefaultAccount] Getting Default Account from authenticated sessions for provider:', authenticationProvider.id); const sessions = await this.findMatchingProviderSession(authenticationProvider.id, this.defaultAccountConfig.authenticationProvider.scopes); @@ -447,33 +459,42 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid return null; } - return this.getDefaultAccountFromAuthenticatedSessions(authenticationProvider, sessions); + return this.getDefaultAccountFromAuthenticatedSessions(authenticationProvider, sessions, donotUseLastFetchedData); } catch (error) { this.logService.error('[DefaultAccount] Failed to get default account for provider:', authenticationProvider.id, getErrorMessage(error)); return null; } } - private async getDefaultAccountFromAuthenticatedSessions(authenticationProvider: IDefaultAccountAuthenticationProvider, sessions: AuthenticationSession[]): Promise { + private async getDefaultAccountFromAuthenticatedSessions(authenticationProvider: IDefaultAccountAuthenticationProvider, sessions: AuthenticationSession[], donotUseLastFetchedData: boolean): Promise { try { const accountId = sessions[0].account.id; + const accountPolicyData = !donotUseLastFetchedData && this._policyData?.accountId === accountId ? this._policyData : undefined; + const [entitlementsData, tokenEntitlementsData] = await Promise.all([ this.getEntitlements(sessions), - this.getTokenEntitlements(sessions), + this.getTokenEntitlements(sessions, accountPolicyData), ]); - let policyData: Mutable | undefined = this._policyData?.accountId === accountId ? { ...this._policyData.policyData } : undefined; + let isTokenEntitlementsDataFetched = false; + let isMcpRegistryDataFetched = false; + let policyData: Mutable | undefined = accountPolicyData?.policyData ? { ...accountPolicyData.policyData } : undefined; if (tokenEntitlementsData) { + isTokenEntitlementsDataFetched = true; policyData = policyData ?? {}; policyData.chat_agent_enabled = tokenEntitlementsData.chat_agent_enabled; policyData.chat_preview_features_enabled = tokenEntitlementsData.chat_preview_features_enabled; policyData.mcp = tokenEntitlementsData.mcp; if (policyData.mcp) { - const mcpRegistryProvider = await this.getMcpRegistryProvider(sessions); - if (mcpRegistryProvider) { - policyData.mcpRegistryUrl = mcpRegistryProvider.url; - policyData.mcpAccess = mcpRegistryProvider.registry_access; + const mcpRegistryProvider = await this.getMcpRegistryProvider(sessions, accountPolicyData); + if (!isUndefined(mcpRegistryProvider)) { + isMcpRegistryDataFetched = true; + policyData.mcpRegistryUrl = mcpRegistryProvider?.url; + policyData.mcpAccess = mcpRegistryProvider?.registry_access; } + } else { + policyData.mcpRegistryUrl = undefined; + policyData.mcpAccess = undefined; } } @@ -484,7 +505,7 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid entitlementsData, }; this.logService.debug('[DefaultAccount] Successfully created default account for provider:', authenticationProvider.id); - return { defaultAccount, policyData: policyData ? { accountId, policyData } : null }; + return { defaultAccount, policyData: policyData ? { accountId, policyData, isTokenEntitlementsDataFetched, isMcpRegistryDataFetched } : null }; } catch (error) { this.logService.error('[DefaultAccount] Failed to create default account for provider:', authenticationProvider.id, getErrorMessage(error)); return null; @@ -539,7 +560,15 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid return expectedScopes.every(scope => scopes.includes(scope)); } - private async getTokenEntitlements(sessions: AuthenticationSession[]): Promise | undefined> { + private async getTokenEntitlements(sessions: AuthenticationSession[], accountPolicyData: IAccountPolicyData | undefined): Promise | undefined> { + if (accountPolicyData?.isTokenEntitlementsDataFetched) { + this.logService.debug('[DefaultAccount] Using last fetched token entitlements data'); + return accountPolicyData.policyData; + } + return await this.requestTokenEntitlements(sessions); + } + + private async requestTokenEntitlements(sessions: AuthenticationSession[]): Promise | undefined> { const tokenEntitlementsUrl = this.getTokenEntitlementUrl(); if (!tokenEntitlementsUrl) { this.logService.debug('[DefaultAccount] No token entitlements URL found'); @@ -610,11 +639,19 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid return undefined; } - private async getMcpRegistryProvider(sessions: AuthenticationSession[]): Promise { + private async getMcpRegistryProvider(sessions: AuthenticationSession[], accountPolicyData: IAccountPolicyData | undefined): Promise { + if (accountPolicyData?.isMcpRegistryDataFetched) { + this.logService.debug('[DefaultAccount] Using last fetched MCP registry data'); + return accountPolicyData.policyData.mcpRegistryUrl && accountPolicyData.policyData.mcpAccess ? { url: accountPolicyData.policyData.mcpRegistryUrl, registry_access: accountPolicyData.policyData.mcpAccess } : null; + } + return await this.requestMcpRegistryProvider(sessions); + } + + private async requestMcpRegistryProvider(sessions: AuthenticationSession[]): Promise { const mcpRegistryDataUrl = this.getMcpRegistryDataUrl(); if (!mcpRegistryDataUrl) { this.logService.debug('[DefaultAccount] No MCP registry data URL found'); - return undefined; + return null; } this.logService.debug('[DefaultAccount] Fetching MCP registry data from:', mcpRegistryDataUrl); @@ -625,20 +662,23 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid if (response.res.statusCode && response.res.statusCode !== 200) { this.logService.trace(`[DefaultAccount] unexpected status code ${response.res.statusCode} while fetching MCP registry data`); - return undefined; + return response.res.statusCode === 404 /* mcp not configured */ + ? null + : undefined; } try { const data = await asJson(response); if (data) { this.logService.debug('Fetched MCP registry providers', data.mcp_registries); - return data.mcp_registries[0]; + return data.mcp_registries[0] ?? null; } - this.logService.debug('Failed to fetch MCP registry providers', 'No data returned'); + this.logService.debug('No MCP registry providers content found in response'); + return null; } catch (error) { this.logService.error('Failed to fetch MCP registry providers', getErrorMessage(error)); + return undefined; } - return undefined; } private async request(url: string, type: 'GET', body: undefined, sessions: AuthenticationSession[], token: CancellationToken): Promise; @@ -681,11 +721,6 @@ class DefaultAccountProvider extends Disposable implements IDefaultAccountProvid return undefined; } - if (lastResponse.res.statusCode && lastResponse.res.statusCode !== 200) { - this.logService.trace(`[DefaultAccount]: unexpected status code ${lastResponse.res.statusCode} for request`, url); - return undefined; - } - return lastResponse; }