From db682a870ffe60051e522f1a0cbd376631ec10e6 Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Mon, 14 Sep 2026 17:12:53 +0100 Subject: [PATCH 1/8] feat(assets-controller): add transient per-account loading state for user-visible fetches Add a non-persisted 'assetsLoadingStatus' state map (Record) that is populated while getAssets runs for user-visible moments and cleared when the fetch settles. Only the two user-action fetch sites (account group change, unlock/startup refresh) pass the new optional 'trigger' option, so background fetches (polling, price/tx/currency refreshes) never flap the indicator. The field is marked persist: false, so no migration is needed. Also export selectors (getAccountLoadingStatus, isAccountLoading, getAccountsLoadingStatus, isAnyAccountLoading) for the UX layer. getAssets is refactored into a thin lifecycle wrapper around the extracted #getAssetsInternal pipeline (body unchanged), and tempHealAssetsInfoMetadata is made generic over the state slice so its return type stays assignable to the widened controller state. --- packages/assets-controller/CHANGELOG.md | 5 + .../AssetsController-method-action-types.ts | 28 ++ .../src/AssetsController.test.ts | 286 +++++++++++++++++- .../assets-controller/src/AssetsController.ts | 140 +++++++++ packages/assets-controller/src/index.ts | 8 + .../src/migrations/healAssetsInfoMetadata.ts | 12 +- .../src/selectors/loading.test.ts | 125 ++++++++ .../src/selectors/loading.ts | 99 ++++++ packages/assets-controller/src/types.ts | 12 + 9 files changed, 710 insertions(+), 5 deletions(-) create mode 100644 packages/assets-controller/src/selectors/loading.test.ts create mode 100644 packages/assets-controller/src/selectors/loading.ts diff --git a/packages/assets-controller/CHANGELOG.md b/packages/assets-controller/CHANGELOG.md index 8e1bbf0f5bd..1a45aa1b744 100644 --- a/packages/assets-controller/CHANGELOG.md +++ b/packages/assets-controller/CHANGELOG.md @@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Add a transient, non-persisted per-account assets loading state. `AssetsControllerState` now includes `assetsLoadingStatus: Record` (`'accountSwitch' | 'unlock'`), populated while `getAssets` runs for user-visible moments and cleared when the fetch settles; the field is never persisted. `getAssets` accepts a new optional `trigger` option that opts a fetch into this behavior — only account switches and unlock-time startup refreshes pass it ([#PR](https://github.com/MetaMask/core/pull/PR)) +- Add loading-state selectors `getAccountLoadingStatus`, `isAccountLoading`, `getAccountsLoadingStatus`, and `isAnyAccountLoading` for the UX layer to read `assetsLoadingStatus` ([#PR](https://github.com/MetaMask/core/pull/PR)) + ### Changed - Bump `@metamask/account-tree-controller` from `^10.0.0` to `^10.0.1` ([#10166](https://github.com/MetaMask/core/pull/10166)) diff --git a/packages/assets-controller/src/AssetsController-method-action-types.ts b/packages/assets-controller/src/AssetsController-method-action-types.ts index 0e0479ff245..dbde23af24b 100644 --- a/packages/assets-controller/src/AssetsController-method-action-types.ts +++ b/packages/assets-controller/src/AssetsController-method-action-types.ts @@ -5,6 +5,34 @@ import type { AssetsController } from './AssetsController.js'; +/** + * Fetch assets for the given accounts, publishing a transient per-account + * loading state while user-visible fetches are in flight. + * + * When `options.trigger` is set, each requested account is marked in + * `state.assetsLoadingStatus` with the trigger for the duration of the + * fetch (set before the fetch starts, removed when it settles — success or + * failure). This exists so the UI can show a loading indicator during + * user-visible moments (account switch, unlock) without inferring it from + * balance state. Pass a trigger **only** for those moments; background + * fetches (polling, price refreshes, post-transaction refreshes) must omit + * it so the indicator does not flap. + * + * @param accounts - Accounts to fetch assets for. + * @param options - Fetch options. + * @param options.chainIds - Chains to fetch for; defaults to enabled chains. + * @param options.assetTypes - Asset types to fetch (fungible, native). + * @param options.forceUpdate - Skip cache and fetch fresh data. + * @param options.bypassServerCache - Also bypass server-side HTTP caches. + * @param options.dataTypes - Data types to fetch (balances, info, prices). + * @param options.assetsForPriceUpdate - Asset IDs to fetch prices for. + * @param options.updateMode - `'merge'` to combine with existing state + * instead of replacing it. + * @param options.trigger - User-visible moment this fetch belongs to; when + * set, per-account loading status is published while the fetch is in + * flight. Omit for background fetches. + * @returns The combined assets per account, read from state after the fetch. + */ export type AssetsControllerGetAssetsAction = { type: `AssetsController:getAssets`; handler: AssetsController['getAssets']; diff --git a/packages/assets-controller/src/AssetsController.test.ts b/packages/assets-controller/src/AssetsController.test.ts index dc31450f492..b8f86c11846 100644 --- a/packages/assets-controller/src/AssetsController.test.ts +++ b/packages/assets-controller/src/AssetsController.test.ts @@ -1,5 +1,6 @@ -import { clientControllerSelectors } from '@metamask/client-controller'; /* eslint-disable jest/unbound-method */ +import { deriveStateFromMetadata } from '@metamask/base-controller'; +import { clientControllerSelectors } from '@metamask/client-controller'; import type { TraceCallback, TraceRequest } from '@metamask/controller-utils'; import type { ApiPlatformClient } from '@metamask/core-backend'; import type { InternalAccount } from '@metamask/keyring-internal-api'; @@ -89,6 +90,63 @@ function createMockQueryApiClient(): ApiPlatformClient { } as unknown as ApiPlatformClient; } +/** + * A query API client whose Accounts API calls can be frozen mid-flight so a + * forced `getAssets()` run is observable while in flight. + * + * While "armed", every Accounts API network/balance call returns a promise + * that stays pending until `release()` is called, then resolves with a valid + * (empty) response shape. While "disarmed", calls resolve immediately. + * + * @returns The gated client plus `arm()`/`release()` controls. + */ +function createGatedQueryApiClient(): { + client: ApiPlatformClient; + arm: () => void; + release: () => void; +} { + let armed = false; + let releaseGate: () => void = () => undefined; + const gate = new Promise((resolve) => { + releaseGate = resolve; + }); + + const gatedCall = + (value: () => Result): (() => Promise) => + () => + armed ? gate.then(value) : Promise.resolve(value()); + + const client = { + ...createMockQueryApiClient(), + accounts: { + fetchV2SupportedNetworks: gatedCall(() => ({ + fullSupport: ['eip155:1'], + partialSupport: [], + })), + fetchV5MultiAccountBalances: gatedCall(() => ({ + balances: [], + unprocessedNetworks: [], + })), + fetchV6MultiAccountBalances: gatedCall(() => ({ + accounts: [], + unprocessedNetworks: [], + unprocessedIncludeAssetIds: [], + })), + }, + } as unknown as ApiPlatformClient; + + return { + client, + arm: (): void => { + armed = true; + }, + release: (): void => { + armed = false; + releaseGate(); + }, + }; +} + type AllActions = MessengerActions; type AllEvents = MessengerEvents; @@ -334,6 +392,7 @@ describe('AssetsController', () => { customAssets: {}, assetPreferences: {}, selectedCurrency: 'usd', + assetsLoadingStatus: {}, }); }); @@ -363,6 +422,7 @@ describe('AssetsController', () => { customAssets: {}, assetPreferences: {}, selectedCurrency: 'usd', + assetsLoadingStatus: {}, }); }); }); @@ -520,6 +580,7 @@ describe('AssetsController', () => { assetsPrice: {}, customAssets: {}, selectedCurrency: 'usd', + assetsLoadingStatus: {}, }); // Action handlers should be registered @@ -3765,6 +3826,229 @@ describe('AssetsController', () => { }); }); + describe('assets loading status', () => { + it('marks each requested account as loading while a triggered getAssets is in flight, then clears it', async () => { + const { client, arm, release } = createGatedQueryApiClient(); + + await withController( + { queryApiClient: client }, + async ({ controller }) => { + const accountA = createMockInternalAccount(); + const accountB = createMockInternalAccount({ + id: 'mock-account-id-2', + }); + + // Let the Accounts API data source finish its construction-time + // supported-networks lookup so the armed fetch is gated on the + // balances call (deterministic in-flight window). + await flushPromises(); + + arm(); + const fetchPromise = controller.getAssets([accountA, accountB], { + forceUpdate: true, + trigger: 'accountSwitch', + }); + await flushPromises(); + + // In flight: both accounts are marked loading with the trigger. + expect(controller.state.assetsLoadingStatus).toStrictEqual({ + [accountA.id]: 'accountSwitch', + [accountB.id]: 'accountSwitch', + }); + + release(); + await fetchPromise; + + // Settled: the transient entries are removed. + expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + }, + ); + }); + + it('clears the loading status when a triggered getAssets rejects', async () => { + await withController(async ({ controller }) => { + const account = createMockInternalAccount(); + // Force the fetch to fail before the pipeline runs. + const customAssetsSpy = jest + .spyOn(controller, 'getCustomAssets') + .mockImplementation(() => { + throw new Error('fetch failed'); + }); + + await expect( + controller.getAssets([account], { + forceUpdate: true, + trigger: 'accountSwitch', + }), + ).rejects.toThrow('fetch failed'); + + expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + customAssetsSpy.mockRestore(); + }); + }); + + it('does not set the loading status for fetches without a trigger (polling/refresh paths)', async () => { + const { client, arm, release } = createGatedQueryApiClient(); + + await withController( + { queryApiClient: client }, + async ({ controller }) => { + const account = createMockInternalAccount(); + + await flushPromises(); + + arm(); + const fetchPromise = controller.getAssets([account], { + forceUpdate: true, + }); + await flushPromises(); + + // In flight, but no trigger: no loading status is published. + expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + + release(); + await fetchPromise; + + expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + }, + ); + }); + + it('marks accounts as loading with the unlock trigger when tracking restarts after unlock', async () => { + const { client, arm, release } = createGatedQueryApiClient(); + + await withController( + { + queryApiClient: client, + clientControllerState: { isUiOpen: true }, + }, + async ({ controller, messenger }) => { + // Start tracking with the gate disarmed so the first startup + // refresh settles and the Accounts API active chains are known. + await activateTracking(messenger); + expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + + // Lock, then unlock with the gate armed: #start() runs a fresh + // #runStartupRefresh() -> getAssets(trigger: 'unlock'), frozen + // mid-flight by the gated balances call. + messenger.publish('KeyringController:lock'); + arm(); + messenger.publish('KeyringController:unlock'); + await flushPromises(); + + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'unlock', + ); + + release(); + await flushPromises(); + + expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + }, + ); + }); + + it('marks accounts as loading with the accountSwitch trigger on selected group change', async () => { + const { client, arm, release } = createGatedQueryApiClient(); + + await withController( + { queryApiClient: client }, + async ({ controller, messenger }) => { + // Start tracking with the gate disarmed so the startup refresh settles. + await activateTracking(messenger); + expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + + arm(); + (messenger.publish as CallableFunction)( + 'AccountTreeController:selectedAccountGroupChange', + 'entropy:mock-keyring-id-1/1', + 'entropy:mock-keyring-id-1/0', + ); + await flushPromises(); + + // The group-change refresh is frozen mid-flight by the gated client. + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'accountSwitch', + ); + + release(); + await flushPromises(); + + expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + }, + ); + }); + + it('emits state change events when the loading status is set and cleared', async () => { + const { client, arm, release } = createGatedQueryApiClient(); + + await withController( + { queryApiClient: client }, + async ({ controller, messenger }) => { + const stateChanges: AssetsControllerState[] = []; + messenger.subscribe('AssetsController:stateChanged', (state) => { + stateChanges.push(state); + }); + + const account = createMockInternalAccount(); + + await flushPromises(); + + arm(); + const fetchPromise = controller.getAssets([account], { + forceUpdate: true, + trigger: 'accountSwitch', + }); + await flushPromises(); + + const whileLoading = stateChanges.find( + (state) => + state.assetsLoadingStatus[account.id] === 'accountSwitch', + ); + expect(whileLoading).toBeDefined(); + + release(); + await fetchPromise; + + const afterCleared = stateChanges.at(-1); + expect(afterCleared?.assetsLoadingStatus[account.id]).toBeUndefined(); + }, + ); + }); + + it('does not persist the loading status in the persisted state snapshot', async () => { + const { client, arm, release } = createGatedQueryApiClient(); + + await withController( + { queryApiClient: client }, + async ({ controller }) => { + const account = createMockInternalAccount(); + + await flushPromises(); + + arm(); + const fetchPromise = controller.getAssets([account], { + forceUpdate: true, + trigger: 'accountSwitch', + }); + await flushPromises(); + + // Transient state: excluded from the persisted snapshot even while + // an entry is present. + const persisted = deriveStateFromMetadata( + controller.state, + controller.metadata, + 'persist', + ); + expect(persisted).not.toHaveProperty('assetsLoadingStatus'); + + release(); + await fetchPromise; + }, + ); + }); + }); + describe('account tree initialized', () => { it('triggers start when the tree initializes after unlock with empty accounts', async () => { const getAccountsMock = jest.fn().mockReturnValue([]); diff --git a/packages/assets-controller/src/AssetsController.ts b/packages/assets-controller/src/AssetsController.ts index e3dffd5630b..98f1bb09bad 100644 --- a/packages/assets-controller/src/AssetsController.ts +++ b/packages/assets-controller/src/AssetsController.ts @@ -139,6 +139,7 @@ import type { AssetBalance, AccountWithSupportedChains, AssetType, + AssetsLoadingTrigger, DataType, DataRequest, DataResponse, @@ -261,6 +262,13 @@ export type AssetsControllerState = { assetPreferences: { [assetId: string]: AssetPreferences }; /** Currently-active ISO 4217 currency code */ selectedCurrency: SupportedCurrency; + /** + * Transient per-account loading state for user-visible asset fetches + * (account switch, unlock). An entry is present while the fetch for that + * account is in flight and its value is the trigger; entries are removed + * when the fetch settles. Never persisted. + */ + assetsLoadingStatus: Record; }; /** @@ -283,6 +291,8 @@ export function getDefaultAssetsControllerState(): AssetsControllerState { customAssets: {}, assetPreferences: {}, selectedCurrency: 'usd', + // Transient loading state — never restored from persisted state. + assetsLoadingStatus: {}, }; } @@ -504,6 +514,14 @@ const stateMetadata: StateMetadata = { includeInDebugSnapshot: false, usedInUi: true, }, + assetsLoadingStatus: { + // Transient in-flight fetch markers (account switch, unlock). Never + // persisted — a restarted client always starts with no fetch in flight. + persist: false, + includeInStateLogs: true, + includeInDebugSnapshot: true, + usedInUi: true, + }, }; // ============================================================================ @@ -1362,6 +1380,7 @@ export class AssetsController extends BaseController< await this.getAssets(accounts, { chainIds: [...this.#enabledChains], forceUpdate: true, + trigger: 'unlock', }); // Seed before subscribe so the price poll / update fetch sees natives // and default tracked assets that were never returned by balance APIs. @@ -1604,6 +1623,34 @@ export class AssetsController extends BaseController< // PUBLIC API: QUERY METHODS // ============================================================================ + /** + * Fetch assets for the given accounts, publishing a transient per-account + * loading state while user-visible fetches are in flight. + * + * When `options.trigger` is set, each requested account is marked in + * `state.assetsLoadingStatus` with the trigger for the duration of the + * fetch (set before the fetch starts, removed when it settles — success or + * failure). This exists so the UI can show a loading indicator during + * user-visible moments (account switch, unlock) without inferring it from + * balance state. Pass a trigger **only** for those moments; background + * fetches (polling, price refreshes, post-transaction refreshes) must omit + * it so the indicator does not flap. + * + * @param accounts - Accounts to fetch assets for. + * @param options - Fetch options. + * @param options.chainIds - Chains to fetch for; defaults to enabled chains. + * @param options.assetTypes - Asset types to fetch (fungible, native). + * @param options.forceUpdate - Skip cache and fetch fresh data. + * @param options.bypassServerCache - Also bypass server-side HTTP caches. + * @param options.dataTypes - Data types to fetch (balances, info, prices). + * @param options.assetsForPriceUpdate - Asset IDs to fetch prices for. + * @param options.updateMode - `'merge'` to combine with existing state + * instead of replacing it. + * @param options.trigger - User-visible moment this fetch belongs to; when + * set, per-account loading status is published while the fetch is in + * flight. Omit for background fetches. + * @returns The combined assets per account, read from state after the fetch. + */ async getAssets( accounts: InternalAccount[], options?: { @@ -1621,6 +1668,58 @@ export class AssetsController extends BaseController< assetsForPriceUpdate?: Caip19AssetId[]; /** When set to `'merge'`, fetch result is merged with existing state instead of replacing. Use for partial fetches (e.g. newly added chains). */ updateMode?: AssetsUpdateMode; + /** + * User-visible moment this fetch belongs to. When set, per-account + * loading status is published in `state.assetsLoadingStatus` while the + * fetch is in flight. Omit for background fetches. + */ + trigger?: AssetsLoadingTrigger; + }, + ): Promise>> { + const { trigger } = options ?? {}; + + if (!trigger || accounts.length === 0) { + return this.#getAssetsInternal(accounts, options); + } + + this.#setAssetsLoadingStatus(accounts, trigger); + try { + return await this.#getAssetsInternal(accounts, options); + } finally { + this.#clearAssetsLoadingStatus(accounts, trigger); + } + } + + /** + * Fetch pipeline and state read for {@link AssetsController.getAssets}. + * Kept separate so the public method stays a thin loading-state lifecycle + * wrapper around this fetch logic. + * + * @param accounts - Accounts to fetch assets for. + * @param options - Fetch options (forwarded from `getAssets`). + * @param options.chainIds - Chains to fetch for; defaults to enabled chains. + * @param options.assetTypes - Asset types to fetch (fungible, native). + * @param options.forceUpdate - Skip cache and fetch fresh data. + * @param options.bypassServerCache - Also bypass server-side HTTP caches. + * @param options.dataTypes - Data types to fetch (balances, info, prices). + * @param options.assetsForPriceUpdate - Asset IDs to fetch prices for. + * @param options.updateMode - `'merge'` to combine with existing state + * instead of replacing it. + * @param options.trigger - User-visible moment this fetch belongs to + * (loading status is handled by the `getAssets` wrapper). + * @returns The combined assets per account, read from state after the fetch. + */ + async #getAssetsInternal( + accounts: InternalAccount[], + options?: { + chainIds?: ChainId[]; + assetTypes?: AssetType[]; + forceUpdate?: boolean; + bypassServerCache?: boolean; + dataTypes?: DataType[]; + assetsForPriceUpdate?: Caip19AssetId[]; + updateMode?: AssetsUpdateMode; + trigger?: AssetsLoadingTrigger; }, ): Promise>> { const chainIds = options?.chainIds ?? [...this.#enabledChains]; @@ -1798,6 +1897,46 @@ export class AssetsController extends BaseController< return result; } + /** + * Mark the given accounts as loading for a user-visible fetch trigger. + * + * @param accounts - Accounts whose fetch is starting. + * @param trigger - The user-visible moment the fetch belongs to. + */ + #setAssetsLoadingStatus( + accounts: InternalAccount[], + trigger: AssetsLoadingTrigger, + ): void { + this.update((state) => { + for (const account of accounts) { + state.assetsLoadingStatus[account.id] = trigger; + } + }); + } + + /** + * Clear the loading marker this fetch set for the given accounts. + * + * Only removes entries still pointing at `trigger`, so a newer fetch for + * the same account (a different trigger) is never clobbered by an older + * fetch settling. + * + * @param accounts - Accounts whose fetch settled. + * @param trigger - The trigger the entries were set with. + */ + #clearAssetsLoadingStatus( + accounts: InternalAccount[], + trigger: AssetsLoadingTrigger, + ): void { + this.update((state) => { + for (const account of accounts) { + if (state.assetsLoadingStatus[account.id] === trigger) { + delete state.assetsLoadingStatus[account.id]; + } + } + }); + } + async getAssetsBalance( accounts: InternalAccount[], options?: { @@ -3637,6 +3776,7 @@ export class AssetsController extends BaseController< await this.getAssets(accounts, { chainIds: [...this.#enabledChains], forceUpdate: true, + trigger: 'accountSwitch', }); } diff --git a/packages/assets-controller/src/index.ts b/packages/assets-controller/src/index.ts index fca7ca96476..64e761eb8ef 100644 --- a/packages/assets-controller/src/index.ts +++ b/packages/assets-controller/src/index.ts @@ -93,6 +93,7 @@ export type { DataRequest, DataResponse, AssetsUpdateMode, + AssetsLoadingTrigger, // Middleware types Context, NextFunction, @@ -199,6 +200,13 @@ export { getInternalAccountsForGroup, } from './selectors/balance.js'; +export { + getAccountLoadingStatus, + getAccountsLoadingStatus, + isAccountLoading, + isAnyAccountLoading, +} from './selectors/loading.js'; + export type { AccountGroupBalance, AccountsById, diff --git a/packages/assets-controller/src/migrations/healAssetsInfoMetadata.ts b/packages/assets-controller/src/migrations/healAssetsInfoMetadata.ts index 75c0c6ae868..0af51f28559 100644 --- a/packages/assets-controller/src/migrations/healAssetsInfoMetadata.ts +++ b/packages/assets-controller/src/migrations/healAssetsInfoMetadata.ts @@ -99,9 +99,11 @@ export type AssetsInfoHealingPatch = { const log = createModuleLogger(projectLogger, 'tempHealAssetsInfoMetadata'); -export type TempHealAssetsInfoMetadataOptions = { +export type TempHealAssetsInfoMetadataOptions< + State extends AssetsControllerStateInternal = AssetsControllerStateInternal, +> = { /** Current `AssetsController` state the healing patch is computed against. */ - state: AssetsControllerStateInternal; + state: State; /** * Host-provided getter for the untrusted legacy state root (see * `AssetsControllerOptions.tempMigrateAssetsInfoMetadataAssets3346`). @@ -121,11 +123,13 @@ export type TempHealAssetsInfoMetadataOptions = { * @returns Updated controller state with the healing patch applied, or the * original state when there is nothing to heal or healing fails. */ -export function tempHealAssetsInfoMetadata({ +export function tempHealAssetsInfoMetadata< + State extends AssetsControllerStateInternal = AssetsControllerStateInternal, +>({ state, getMigrationState, captureException, -}: TempHealAssetsInfoMetadataOptions): AssetsControllerStateInternal { +}: TempHealAssetsInfoMetadataOptions): State { const reportError = (error: unknown): void => { log('Failed to heal assetsInfo metadata', error); captureException?.( diff --git a/packages/assets-controller/src/selectors/loading.test.ts b/packages/assets-controller/src/selectors/loading.test.ts new file mode 100644 index 00000000000..55861b2ec60 --- /dev/null +++ b/packages/assets-controller/src/selectors/loading.test.ts @@ -0,0 +1,125 @@ +import type { AssetsControllerState } from '../AssetsController.js'; +import type { AccountId } from '../types.js'; +import { + getAccountLoadingStatus, + getAccountsLoadingStatus, + isAccountLoading, + isAnyAccountLoading, +} from './loading.js'; + +const ACCOUNT_ID_A: AccountId = 'mock-account-id-1'; +const ACCOUNT_ID_B: AccountId = 'mock-account-id-2'; +const ACCOUNT_ID_C: AccountId = 'mock-account-id-3'; + +const createState = ( + assetsLoadingStatus: AssetsControllerState['assetsLoadingStatus'], +): Pick => ({ + assetsLoadingStatus, +}); + +describe('loading selectors', () => { + describe('getAccountLoadingStatus', () => { + it('returns the trigger for an account with an in-flight fetch', () => { + const state = createState({ + [ACCOUNT_ID_A]: 'accountSwitch', + [ACCOUNT_ID_B]: 'unlock', + }); + + expect(getAccountLoadingStatus(state, ACCOUNT_ID_A)).toBe( + 'accountSwitch', + ); + expect(getAccountLoadingStatus(state, ACCOUNT_ID_B)).toBe('unlock'); + }); + + it('returns undefined when the account is not loading', () => { + const state = createState({ [ACCOUNT_ID_A]: 'accountSwitch' }); + + expect(getAccountLoadingStatus(state, ACCOUNT_ID_B)).toBeUndefined(); + }); + + it('returns undefined when no account is loading', () => { + const state = createState({}); + + expect(getAccountLoadingStatus(state, ACCOUNT_ID_A)).toBeUndefined(); + }); + }); + + describe('isAccountLoading', () => { + it('returns true while the account has an in-flight fetch', () => { + const state = createState({ + [ACCOUNT_ID_A]: 'accountSwitch', + [ACCOUNT_ID_B]: 'unlock', + }); + + expect(isAccountLoading(state, ACCOUNT_ID_A)).toBe(true); + expect(isAccountLoading(state, ACCOUNT_ID_B)).toBe(true); + }); + + it('returns false when the account is not loading', () => { + const state = createState({ [ACCOUNT_ID_A]: 'accountSwitch' }); + + expect(isAccountLoading(state, ACCOUNT_ID_B)).toBe(false); + }); + + it('returns false when no account is loading', () => { + const state = createState({}); + + expect(isAccountLoading(state, ACCOUNT_ID_A)).toBe(false); + }); + }); + + describe('getAccountsLoadingStatus', () => { + it('returns loading entries only for the requested account ids', () => { + const state = createState({ + [ACCOUNT_ID_A]: 'accountSwitch', + [ACCOUNT_ID_B]: 'unlock', + }); + + expect( + getAccountsLoadingStatus(state, [ACCOUNT_ID_A, ACCOUNT_ID_C]), + ).toStrictEqual({ [ACCOUNT_ID_A]: 'accountSwitch' }); + }); + + it('returns an empty record when none of the requested accounts are loading', () => { + const state = createState({ [ACCOUNT_ID_B]: 'unlock' }); + + expect( + getAccountsLoadingStatus(state, [ACCOUNT_ID_A, ACCOUNT_ID_C]), + ).toStrictEqual({}); + }); + + it('returns an empty record for an empty account list', () => { + const state = createState({ [ACCOUNT_ID_A]: 'accountSwitch' }); + + expect(getAccountsLoadingStatus(state, [])).toStrictEqual({}); + }); + }); + + describe('isAnyAccountLoading', () => { + it('returns true when any requested account is loading', () => { + const state = createState({ + [ACCOUNT_ID_A]: 'accountSwitch', + [ACCOUNT_ID_B]: 'unlock', + }); + + expect(isAnyAccountLoading(state, [ACCOUNT_ID_B, ACCOUNT_ID_C])).toBe( + true, + ); + }); + + it('returns false when none of the requested accounts are loading', () => { + const state = createState({ [ACCOUNT_ID_A]: 'accountSwitch' }); + + expect(isAnyAccountLoading(state, [ACCOUNT_ID_B, ACCOUNT_ID_C])).toBe( + false, + ); + }); + + it('defaults to all accounts when no account ids are given', () => { + expect(isAnyAccountLoading(createState({}))).toBe(false); + expect( + isAnyAccountLoading(createState({ [ACCOUNT_ID_C]: 'accountSwitch' })), + ).toBe(true); + }); + }); +}); diff --git a/packages/assets-controller/src/selectors/loading.ts b/packages/assets-controller/src/selectors/loading.ts new file mode 100644 index 00000000000..a51bdee9734 --- /dev/null +++ b/packages/assets-controller/src/selectors/loading.ts @@ -0,0 +1,99 @@ +import type { AssetsControllerState } from '../AssetsController.js'; +import type { AccountId, AssetsLoadingTrigger } from '../types.js'; + +/** + * Loading-state selectors over + * {@link AssetsControllerState.assetsLoadingStatus}. + * + * `assetsLoadingStatus` is transient (never persisted): an entry is present + * while a user-visible asset fetch (account switch, unlock) is in flight for + * that account, and its value is the trigger. Absence means "not loading". + * + * All selectors are synchronous and read from the controller state slice, + * so they can be used directly in `useSelector`-style subscriptions to + * `AssetsController:stateChanged`. + */ + +/** + * Get the loading trigger for a single account, if a user-visible fetch is + * currently in flight for it. + * + * @param state - AssetsController state slice. + * @param accountId - The account id (`InternalAccount.id`). + * @returns The trigger for the in-flight fetch, or `undefined` when the + * account is not loading. + */ +export function getAccountLoadingStatus( + state: Pick, + accountId: AccountId, +): AssetsLoadingTrigger | undefined { + return state.assetsLoadingStatus?.[accountId]; +} + +/** + * Check whether a user-visible asset fetch is in flight for an account. + * + * @param state - AssetsController state slice. + * @param accountId - The account id (`InternalAccount.id`). + * @returns True while the account's assets are loading. + */ +export function isAccountLoading( + state: Pick, + accountId: AccountId, +): boolean { + return getAccountLoadingStatus(state, accountId) !== undefined; +} + +/** + * Get the in-flight loading triggers for a set of accounts (e.g. every + * account in the selected account group). + * + * @param state - AssetsController state slice. + * @param accountIds - The account ids to report on. + * @returns A record containing an entry for each requested account that is + * currently loading, keyed by account id. + */ +export function getAccountsLoadingStatus( + state: Pick, + accountIds: AccountId[], +): Record { + const result: Record = {}; + const loadingStatus = state.assetsLoadingStatus ?? {}; + for (const accountId of accountIds) { + const trigger = loadingStatus[accountId]; + if (trigger !== undefined) { + result[accountId] = trigger; + } + } + return result; +} + +/** + * Check whether any of the given accounts has a user-visible fetch in + * flight. When `accountIds` is omitted, checks every account with a loading + * entry in state. + * + * @param state - AssetsController state slice. + * @param accountIds - Optional account ids to restrict the check to. + * @returns True if at least one of the accounts is loading. + */ +export function isAnyAccountLoading( + state: Pick, + accountIds?: AccountId[], +): boolean { + const loadingStatus = state.assetsLoadingStatus ?? {}; + if (accountIds) { + for (const accountId of accountIds) { + if (loadingStatus[accountId] !== undefined) { + return true; + } + } + return false; + } + for (const accountId in loadingStatus) { + if (loadingStatus[accountId] !== undefined) { + return true; + } + } + return false; +} diff --git a/packages/assets-controller/src/types.ts b/packages/assets-controller/src/types.ts index 501220b6e5f..a99c5035774 100644 --- a/packages/assets-controller/src/types.ts +++ b/packages/assets-controller/src/types.ts @@ -408,6 +408,18 @@ export type DataResponse = { */ export type AssetsUpdateMode = 'full' | 'merge' | 'update'; +/** + * User-visible moments for which an assets fetch is surfaced as a loading + * state in {@link AssetsControllerState.assetsLoadingStatus}. + * + * - **accountSwitch**: The selected account group changed (e.g. the user + * switched accounts), so balances for the newly selected accounts are + * being fetched. + * - **unlock**: The wallet was unlocked (or tracking started after unlock), + * so balances are being fetched before anything is renderable. + */ +export type AssetsLoadingTrigger = 'accountSwitch' | 'unlock'; + // ============================================================================ // DATA SOURCE <-> CONTROLLER (DIRECT CALLS, NO MESSENGER PER SOURCE) // ============================================================================ From 47b7803dfde74b597b06ea47b07374650ecee9d4 Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Mon, 14 Sep 2026 17:15:03 +0100 Subject: [PATCH 2/8] docs(assets-controller): link changelog entries to PR --- packages/assets-controller/CHANGELOG.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/assets-controller/CHANGELOG.md b/packages/assets-controller/CHANGELOG.md index 1a45aa1b744..742aae65282 100644 --- a/packages/assets-controller/CHANGELOG.md +++ b/packages/assets-controller/CHANGELOG.md @@ -9,8 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- Add a transient, non-persisted per-account assets loading state. `AssetsControllerState` now includes `assetsLoadingStatus: Record` (`'accountSwitch' | 'unlock'`), populated while `getAssets` runs for user-visible moments and cleared when the fetch settles; the field is never persisted. `getAssets` accepts a new optional `trigger` option that opts a fetch into this behavior — only account switches and unlock-time startup refreshes pass it ([#PR](https://github.com/MetaMask/core/pull/PR)) -- Add loading-state selectors `getAccountLoadingStatus`, `isAccountLoading`, `getAccountsLoadingStatus`, and `isAnyAccountLoading` for the UX layer to read `assetsLoadingStatus` ([#PR](https://github.com/MetaMask/core/pull/PR)) +- Add a transient, non-persisted per-account assets loading state. `AssetsControllerState` now includes `assetsLoadingStatus: Record` (`'accountSwitch' | 'unlock'`), populated while `getAssets` runs for user-visible moments and cleared when the fetch settles; the field is never persisted. `getAssets` accepts a new optional `trigger` option that opts a fetch into this behavior — only account switches and unlock-time startup refreshes pass it ([#10229](https://github.com/MetaMask/core/pull/10229)) +- Add loading-state selectors `getAccountLoadingStatus`, `isAccountLoading`, `getAccountsLoadingStatus`, and `isAnyAccountLoading` for the UX layer to read `assetsLoadingStatus` ([#10229](https://github.com/MetaMask/core/pull/10229)) ### Changed From b081c9f3ca41b40b375135466bfe1f792f1cd051 Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Mon, 14 Sep 2026 18:18:15 +0100 Subject: [PATCH 3/8] docs(assets-controller): point changelog entries at same-repo PR --- packages/assets-controller/CHANGELOG.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/assets-controller/CHANGELOG.md b/packages/assets-controller/CHANGELOG.md index 742aae65282..b4990411e94 100644 --- a/packages/assets-controller/CHANGELOG.md +++ b/packages/assets-controller/CHANGELOG.md @@ -9,8 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- Add a transient, non-persisted per-account assets loading state. `AssetsControllerState` now includes `assetsLoadingStatus: Record` (`'accountSwitch' | 'unlock'`), populated while `getAssets` runs for user-visible moments and cleared when the fetch settles; the field is never persisted. `getAssets` accepts a new optional `trigger` option that opts a fetch into this behavior — only account switches and unlock-time startup refreshes pass it ([#10229](https://github.com/MetaMask/core/pull/10229)) -- Add loading-state selectors `getAccountLoadingStatus`, `isAccountLoading`, `getAccountsLoadingStatus`, and `isAnyAccountLoading` for the UX layer to read `assetsLoadingStatus` ([#10229](https://github.com/MetaMask/core/pull/10229)) +- Add a transient, non-persisted per-account assets loading state. `AssetsControllerState` now includes `assetsLoadingStatus: Record` (`'accountSwitch' | 'unlock'`), populated while `getAssets` runs for user-visible moments and cleared when the fetch settles; the field is never persisted. `getAssets` accepts a new optional `trigger` option that opts a fetch into this behavior — only account switches and unlock-time startup refreshes pass it ([#10230](https://github.com/MetaMask/core/pull/10230)) +- Add loading-state selectors `getAccountLoadingStatus`, `isAccountLoading`, `getAccountsLoadingStatus`, and `isAnyAccountLoading` for the UX layer to read `assetsLoadingStatus` ([#10230](https://github.com/MetaMask/core/pull/10230)) ### Changed From 1b282cc5ce747e18c13f44ebb1382722e7fb8dbb Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Mon, 14 Sep 2026 20:50:34 +0100 Subject: [PATCH 4/8] refactor(assets-controller): simplify loading state to loading/loaded, fix queued-switch gap - Replace the trigger union with a simple 'loading' | 'loaded' status, set before acquiring the refresh mutex so queued switches are marked immediately, and settled to 'loaded' when the fetch finishes - Move the lifecycle into the account-switch and unlock handlers; revert getAssets to its original form (no trigger option), keeping the generated action-types file unchanged - Trim comments and shorten the changelog entry --- packages/assets-controller/CHANGELOG.md | 3 +- .../AssetsController-method-action-types.ts | 28 --- .../src/AssetsController.test.ts | 226 ++++++++---------- .../assets-controller/src/AssetsController.ts | 205 +++++----------- packages/assets-controller/src/index.ts | 2 +- .../src/selectors/loading.test.ts | 59 ++--- .../src/selectors/loading.ts | 66 ++--- packages/assets-controller/src/types.ts | 12 +- 8 files changed, 212 insertions(+), 389 deletions(-) diff --git a/packages/assets-controller/CHANGELOG.md b/packages/assets-controller/CHANGELOG.md index b4990411e94..b86348cde3c 100644 --- a/packages/assets-controller/CHANGELOG.md +++ b/packages/assets-controller/CHANGELOG.md @@ -9,8 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- Add a transient, non-persisted per-account assets loading state. `AssetsControllerState` now includes `assetsLoadingStatus: Record` (`'accountSwitch' | 'unlock'`), populated while `getAssets` runs for user-visible moments and cleared when the fetch settles; the field is never persisted. `getAssets` accepts a new optional `trigger` option that opts a fetch into this behavior — only account switches and unlock-time startup refreshes pass it ([#10230](https://github.com/MetaMask/core/pull/10230)) -- Add loading-state selectors `getAccountLoadingStatus`, `isAccountLoading`, `getAccountsLoadingStatus`, and `isAnyAccountLoading` for the UX layer to read `assetsLoadingStatus` ([#10230](https://github.com/MetaMask/core/pull/10230)) +- Add a transient, non-persisted per-account assets loading state (`assetsLoadingStatus`) that marks accounts as loading during account switches and unlock, with selectors to read it ([#10230](https://github.com/MetaMask/core/pull/10230)) ### Changed diff --git a/packages/assets-controller/src/AssetsController-method-action-types.ts b/packages/assets-controller/src/AssetsController-method-action-types.ts index dbde23af24b..0e0479ff245 100644 --- a/packages/assets-controller/src/AssetsController-method-action-types.ts +++ b/packages/assets-controller/src/AssetsController-method-action-types.ts @@ -5,34 +5,6 @@ import type { AssetsController } from './AssetsController.js'; -/** - * Fetch assets for the given accounts, publishing a transient per-account - * loading state while user-visible fetches are in flight. - * - * When `options.trigger` is set, each requested account is marked in - * `state.assetsLoadingStatus` with the trigger for the duration of the - * fetch (set before the fetch starts, removed when it settles — success or - * failure). This exists so the UI can show a loading indicator during - * user-visible moments (account switch, unlock) without inferring it from - * balance state. Pass a trigger **only** for those moments; background - * fetches (polling, price refreshes, post-transaction refreshes) must omit - * it so the indicator does not flap. - * - * @param accounts - Accounts to fetch assets for. - * @param options - Fetch options. - * @param options.chainIds - Chains to fetch for; defaults to enabled chains. - * @param options.assetTypes - Asset types to fetch (fungible, native). - * @param options.forceUpdate - Skip cache and fetch fresh data. - * @param options.bypassServerCache - Also bypass server-side HTTP caches. - * @param options.dataTypes - Data types to fetch (balances, info, prices). - * @param options.assetsForPriceUpdate - Asset IDs to fetch prices for. - * @param options.updateMode - `'merge'` to combine with existing state - * instead of replacing it. - * @param options.trigger - User-visible moment this fetch belongs to; when - * set, per-account loading status is published while the fetch is in - * flight. Omit for background fetches. - * @returns The combined assets per account, read from state after the fetch. - */ export type AssetsControllerGetAssetsAction = { type: `AssetsController:getAssets`; handler: AssetsController['getAssets']; diff --git a/packages/assets-controller/src/AssetsController.test.ts b/packages/assets-controller/src/AssetsController.test.ts index b8f86c11846..74d1de0d0e1 100644 --- a/packages/assets-controller/src/AssetsController.test.ts +++ b/packages/assets-controller/src/AssetsController.test.ts @@ -3827,94 +3827,89 @@ describe('AssetsController', () => { }); describe('assets loading status', () => { - it('marks each requested account as loading while a triggered getAssets is in flight, then clears it', async () => { + it('marks the selected accounts as loading on account switch, then loaded once the fetch settles', async () => { const { client, arm, release } = createGatedQueryApiClient(); await withController( { queryApiClient: client }, - async ({ controller }) => { - const accountA = createMockInternalAccount(); - const accountB = createMockInternalAccount({ - id: 'mock-account-id-2', - }); - - // Let the Accounts API data source finish its construction-time - // supported-networks lookup so the armed fetch is gated on the - // balances call (deterministic in-flight window). - await flushPromises(); + async ({ controller, messenger }) => { + await activateTracking(messenger); arm(); - const fetchPromise = controller.getAssets([accountA, accountB], { - forceUpdate: true, - trigger: 'accountSwitch', - }); + (messenger.publish as CallableFunction)( + 'AccountTreeController:selectedAccountGroupChange', + 'entropy:mock-keyring-id-1/1', + 'entropy:mock-keyring-id-1/0', + ); await flushPromises(); - // In flight: both accounts are marked loading with the trigger. - expect(controller.state.assetsLoadingStatus).toStrictEqual({ - [accountA.id]: 'accountSwitch', - [accountB.id]: 'accountSwitch', - }); + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loading', + ); release(); - await fetchPromise; + await flushPromises(); - // Settled: the transient entries are removed. - expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loaded', + ); }, ); }); - it('clears the loading status when a triggered getAssets rejects', async () => { - await withController(async ({ controller }) => { - const account = createMockInternalAccount(); - // Force the fetch to fail before the pipeline runs. - const customAssetsSpy = jest - .spyOn(controller, 'getCustomAssets') - .mockImplementation(() => { - throw new Error('fetch failed'); - }); - - await expect( - controller.getAssets([account], { - forceUpdate: true, - trigger: 'accountSwitch', - }), - ).rejects.toThrow('fetch failed'); - - expect(controller.state.assetsLoadingStatus).toStrictEqual({}); - customAssetsSpy.mockRestore(); - }); - }); - - it('does not set the loading status for fetches without a trigger (polling/refresh paths)', async () => { + it('marks a queued account switch as loading immediately, while a previous refresh still holds the refresh mutex', async () => { const { client, arm, release } = createGatedQueryApiClient(); await withController( { queryApiClient: client }, - async ({ controller }) => { - const account = createMockInternalAccount(); + async ({ controller, messenger, getSelectedAccountsMock }) => { + await activateTracking(messenger); + arm(); + (messenger.publish as CallableFunction)( + 'AccountTreeController:selectedAccountGroupChange', + 'entropy:mock-keyring-id-1/1', + 'entropy:mock-keyring-id-1/0', + ); await flushPromises(); - arm(); - const fetchPromise = controller.getAssets([account], { - forceUpdate: true, + // First switch is frozen mid-flight and holds the refresh mutex. + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loading', + ); + + const accountB = createMockInternalAccount({ + id: 'mock-account-id-2', }); + getSelectedAccountsMock.mockReturnValue([accountB]); + (messenger.publish as CallableFunction)( + 'AccountTreeController:selectedAccountGroupChange', + 'entropy:mock-keyring-id-1/2', + 'entropy:mock-keyring-id-1/1', + ); await flushPromises(); - // In flight, but no trigger: no loading status is published. - expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + // The queued switch is already marked loading for the new group's + // accounts, without waiting for the previous refresh to finish. + expect(controller.state.assetsLoadingStatus).toStrictEqual({ + [MOCK_ACCOUNT_ID]: 'loading', + [accountB.id]: 'loading', + }); release(); - await fetchPromise; + await flushPromises(); - expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + expect(controller.state.assetsLoadingStatus).toStrictEqual({ + [MOCK_ACCOUNT_ID]: 'loaded', + [accountB.id]: 'loaded', + }); + + getSelectedAccountsMock.mockClear(); }, ); }); - it('marks accounts as loading with the unlock trigger when tracking restarts after unlock', async () => { + it('marks accounts as loading on unlock, then loaded once the fetch settles', async () => { const { client, arm, release } = createGatedQueryApiClient(); await withController( @@ -3923,129 +3918,120 @@ describe('AssetsController', () => { clientControllerState: { isUiOpen: true }, }, async ({ controller, messenger }) => { - // Start tracking with the gate disarmed so the first startup - // refresh settles and the Accounts API active chains are known. await activateTracking(messenger); - expect(controller.state.assetsLoadingStatus).toStrictEqual({}); - // Lock, then unlock with the gate armed: #start() runs a fresh - // #runStartupRefresh() -> getAssets(trigger: 'unlock'), frozen - // mid-flight by the gated balances call. messenger.publish('KeyringController:lock'); arm(); messenger.publish('KeyringController:unlock'); await flushPromises(); expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( - 'unlock', + 'loading', ); release(); await flushPromises(); - expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loaded', + ); }, ); }); - it('marks accounts as loading with the accountSwitch trigger on selected group change', async () => { + it('marks accounts as loaded even when the startup fetch fails', async () => { + await withController(async ({ controller, messenger }) => { + const getAssetsSpy = jest + .spyOn(controller, 'getAssets') + .mockRejectedValue(new Error('fetch failed')); + + await activateTracking(messenger); + + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loaded', + ); + + getAssetsSpy.mockRestore(); + }); + }); + + it('does not set the loading status for direct background fetches', async () => { const { client, arm, release } = createGatedQueryApiClient(); await withController( { queryApiClient: client }, - async ({ controller, messenger }) => { - // Start tracking with the gate disarmed so the startup refresh settles. - await activateTracking(messenger); - expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + async ({ controller }) => { + const account = createMockInternalAccount(); + + await flushPromises(); arm(); - (messenger.publish as CallableFunction)( - 'AccountTreeController:selectedAccountGroupChange', - 'entropy:mock-keyring-id-1/1', - 'entropy:mock-keyring-id-1/0', - ); + const fetchPromise = controller.getAssets([account], { + forceUpdate: true, + }); await flushPromises(); - // The group-change refresh is frozen mid-flight by the gated client. - expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( - 'accountSwitch', - ); + expect(controller.state.assetsLoadingStatus).toStrictEqual({}); release(); - await flushPromises(); + await fetchPromise; expect(controller.state.assetsLoadingStatus).toStrictEqual({}); }, ); }); - it('emits state change events when the loading status is set and cleared', async () => { + it('emits state change events when the loading status is set and settled', async () => { const { client, arm, release } = createGatedQueryApiClient(); await withController( { queryApiClient: client }, - async ({ controller, messenger }) => { + async ({ messenger }) => { const stateChanges: AssetsControllerState[] = []; messenger.subscribe('AssetsController:stateChanged', (state) => { stateChanges.push(state); }); - const account = createMockInternalAccount(); - - await flushPromises(); + await activateTracking(messenger); arm(); - const fetchPromise = controller.getAssets([account], { - forceUpdate: true, - trigger: 'accountSwitch', - }); + (messenger.publish as CallableFunction)( + 'AccountTreeController:selectedAccountGroupChange', + 'entropy:mock-keyring-id-1/1', + 'entropy:mock-keyring-id-1/0', + ); await flushPromises(); const whileLoading = stateChanges.find( - (state) => - state.assetsLoadingStatus[account.id] === 'accountSwitch', + (state) => state.assetsLoadingStatus[MOCK_ACCOUNT_ID] === 'loading', ); expect(whileLoading).toBeDefined(); release(); - await fetchPromise; + await flushPromises(); - const afterCleared = stateChanges.at(-1); - expect(afterCleared?.assetsLoadingStatus[account.id]).toBeUndefined(); + expect( + stateChanges.at(-1)?.assetsLoadingStatus[MOCK_ACCOUNT_ID], + ).toBe('loaded'); }, ); }); it('does not persist the loading status in the persisted state snapshot', async () => { - const { client, arm, release } = createGatedQueryApiClient(); - - await withController( - { queryApiClient: client }, - async ({ controller }) => { - const account = createMockInternalAccount(); - - await flushPromises(); - - arm(); - const fetchPromise = controller.getAssets([account], { - forceUpdate: true, - trigger: 'accountSwitch', - }); - await flushPromises(); + await withController(async ({ controller, messenger }) => { + await activateTracking(messenger); - // Transient state: excluded from the persisted snapshot even while - // an entry is present. - const persisted = deriveStateFromMetadata( - controller.state, - controller.metadata, - 'persist', - ); - expect(persisted).not.toHaveProperty('assetsLoadingStatus'); + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loaded', + ); - release(); - await fetchPromise; - }, - ); + const persisted = deriveStateFromMetadata( + controller.state, + controller.metadata, + 'persist', + ); + expect(persisted).not.toHaveProperty('assetsLoadingStatus'); + }); }); }); diff --git a/packages/assets-controller/src/AssetsController.ts b/packages/assets-controller/src/AssetsController.ts index 98f1bb09bad..7bcb16a6f53 100644 --- a/packages/assets-controller/src/AssetsController.ts +++ b/packages/assets-controller/src/AssetsController.ts @@ -139,7 +139,7 @@ import type { AssetBalance, AccountWithSupportedChains, AssetType, - AssetsLoadingTrigger, + AssetsLoadingStatus, DataType, DataRequest, DataResponse, @@ -262,13 +262,7 @@ export type AssetsControllerState = { assetPreferences: { [assetId: string]: AssetPreferences }; /** Currently-active ISO 4217 currency code */ selectedCurrency: SupportedCurrency; - /** - * Transient per-account loading state for user-visible asset fetches - * (account switch, unlock). An entry is present while the fetch for that - * account is in flight and its value is the trigger; entries are removed - * when the fetch settles. Never persisted. - */ - assetsLoadingStatus: Record; + assetsLoadingStatus: Record; }; /** @@ -291,7 +285,6 @@ export function getDefaultAssetsControllerState(): AssetsControllerState { customAssets: {}, assetPreferences: {}, selectedCurrency: 'usd', - // Transient loading state — never restored from persisted state. assetsLoadingStatus: {}, }; } @@ -515,8 +508,6 @@ const stateMetadata: StateMetadata = { usedInUi: true, }, assetsLoadingStatus: { - // Transient in-flight fetch markers (account switch, unlock). Never - // persisted — a restarted client always starts with no fetch in flight. persist: false, includeInStateLogs: true, includeInDebugSnapshot: true, @@ -1375,28 +1366,36 @@ export class AssetsController extends BaseController< * @param accounts - Selected accounts to refresh. */ async #runStartupRefresh(accounts: InternalAccount[]): Promise { - const releaseLock = await this.#accountRefreshMutex.acquire(); + this.#setAssetsLoadingStatus(accounts); try { - await this.getAssets(accounts, { - chainIds: [...this.#enabledChains], - forceUpdate: true, - trigger: 'unlock', - }); - // Seed before subscribe so the price poll / update fetch sees natives - // and default tracked assets that were never returned by balance APIs. - this.#ensureNativeBalancesDefaultZero(); - this.#ensureDefaultTrackedAssetsSeeded(); - // Balances were just force-fetched — skip AccountsApi's subscribe-time poll. - this.#subscribeAssets({ skipInitialFetch: true }); - this.#fetchMissingPricesWithoutCache(accounts, [...this.#enabledChains]); - } catch (error) { - log('Failed to fetch assets on startup', error); - this.#ensureNativeBalancesDefaultZero(); - this.#ensureDefaultTrackedAssetsSeeded(); - this.#subscribeAssets({ skipInitialFetch: true }); - this.#fetchMissingPricesWithoutCache(accounts, [...this.#enabledChains]); + const releaseLock = await this.#accountRefreshMutex.acquire(); + try { + await this.getAssets(accounts, { + chainIds: [...this.#enabledChains], + forceUpdate: true, + }); + // Seed before subscribe so the price poll / update fetch sees natives + // and default tracked assets that were never returned by balance APIs. + this.#ensureNativeBalancesDefaultZero(); + this.#ensureDefaultTrackedAssetsSeeded(); + // Balances were just force-fetched — skip AccountsApi's subscribe-time poll. + this.#subscribeAssets({ skipInitialFetch: true }); + this.#fetchMissingPricesWithoutCache(accounts, [ + ...this.#enabledChains, + ]); + } catch (error) { + log('Failed to fetch assets on startup', error); + this.#ensureNativeBalancesDefaultZero(); + this.#ensureDefaultTrackedAssetsSeeded(); + this.#subscribeAssets({ skipInitialFetch: true }); + this.#fetchMissingPricesWithoutCache(accounts, [ + ...this.#enabledChains, + ]); + } finally { + releaseLock(); + } } finally { - releaseLock(); + this.#markAssetsLoaded(accounts); } } @@ -1623,34 +1622,6 @@ export class AssetsController extends BaseController< // PUBLIC API: QUERY METHODS // ============================================================================ - /** - * Fetch assets for the given accounts, publishing a transient per-account - * loading state while user-visible fetches are in flight. - * - * When `options.trigger` is set, each requested account is marked in - * `state.assetsLoadingStatus` with the trigger for the duration of the - * fetch (set before the fetch starts, removed when it settles — success or - * failure). This exists so the UI can show a loading indicator during - * user-visible moments (account switch, unlock) without inferring it from - * balance state. Pass a trigger **only** for those moments; background - * fetches (polling, price refreshes, post-transaction refreshes) must omit - * it so the indicator does not flap. - * - * @param accounts - Accounts to fetch assets for. - * @param options - Fetch options. - * @param options.chainIds - Chains to fetch for; defaults to enabled chains. - * @param options.assetTypes - Asset types to fetch (fungible, native). - * @param options.forceUpdate - Skip cache and fetch fresh data. - * @param options.bypassServerCache - Also bypass server-side HTTP caches. - * @param options.dataTypes - Data types to fetch (balances, info, prices). - * @param options.assetsForPriceUpdate - Asset IDs to fetch prices for. - * @param options.updateMode - `'merge'` to combine with existing state - * instead of replacing it. - * @param options.trigger - User-visible moment this fetch belongs to; when - * set, per-account loading status is published while the fetch is in - * flight. Omit for background fetches. - * @returns The combined assets per account, read from state after the fetch. - */ async getAssets( accounts: InternalAccount[], options?: { @@ -1668,58 +1639,6 @@ export class AssetsController extends BaseController< assetsForPriceUpdate?: Caip19AssetId[]; /** When set to `'merge'`, fetch result is merged with existing state instead of replacing. Use for partial fetches (e.g. newly added chains). */ updateMode?: AssetsUpdateMode; - /** - * User-visible moment this fetch belongs to. When set, per-account - * loading status is published in `state.assetsLoadingStatus` while the - * fetch is in flight. Omit for background fetches. - */ - trigger?: AssetsLoadingTrigger; - }, - ): Promise>> { - const { trigger } = options ?? {}; - - if (!trigger || accounts.length === 0) { - return this.#getAssetsInternal(accounts, options); - } - - this.#setAssetsLoadingStatus(accounts, trigger); - try { - return await this.#getAssetsInternal(accounts, options); - } finally { - this.#clearAssetsLoadingStatus(accounts, trigger); - } - } - - /** - * Fetch pipeline and state read for {@link AssetsController.getAssets}. - * Kept separate so the public method stays a thin loading-state lifecycle - * wrapper around this fetch logic. - * - * @param accounts - Accounts to fetch assets for. - * @param options - Fetch options (forwarded from `getAssets`). - * @param options.chainIds - Chains to fetch for; defaults to enabled chains. - * @param options.assetTypes - Asset types to fetch (fungible, native). - * @param options.forceUpdate - Skip cache and fetch fresh data. - * @param options.bypassServerCache - Also bypass server-side HTTP caches. - * @param options.dataTypes - Data types to fetch (balances, info, prices). - * @param options.assetsForPriceUpdate - Asset IDs to fetch prices for. - * @param options.updateMode - `'merge'` to combine with existing state - * instead of replacing it. - * @param options.trigger - User-visible moment this fetch belongs to - * (loading status is handled by the `getAssets` wrapper). - * @returns The combined assets per account, read from state after the fetch. - */ - async #getAssetsInternal( - accounts: InternalAccount[], - options?: { - chainIds?: ChainId[]; - assetTypes?: AssetType[]; - forceUpdate?: boolean; - bypassServerCache?: boolean; - dataTypes?: DataType[]; - assetsForPriceUpdate?: Caip19AssetId[]; - updateMode?: AssetsUpdateMode; - trigger?: AssetsLoadingTrigger; }, ): Promise>> { const chainIds = options?.chainIds ?? [...this.#enabledChains]; @@ -1897,41 +1816,19 @@ export class AssetsController extends BaseController< return result; } - /** - * Mark the given accounts as loading for a user-visible fetch trigger. - * - * @param accounts - Accounts whose fetch is starting. - * @param trigger - The user-visible moment the fetch belongs to. - */ - #setAssetsLoadingStatus( - accounts: InternalAccount[], - trigger: AssetsLoadingTrigger, - ): void { + #setAssetsLoadingStatus(accounts: InternalAccount[]): void { this.update((state) => { for (const account of accounts) { - state.assetsLoadingStatus[account.id] = trigger; + state.assetsLoadingStatus[account.id] = 'loading'; } }); } - /** - * Clear the loading marker this fetch set for the given accounts. - * - * Only removes entries still pointing at `trigger`, so a newer fetch for - * the same account (a different trigger) is never clobbered by an older - * fetch settling. - * - * @param accounts - Accounts whose fetch settled. - * @param trigger - The trigger the entries were set with. - */ - #clearAssetsLoadingStatus( - accounts: InternalAccount[], - trigger: AssetsLoadingTrigger, - ): void { + #markAssetsLoaded(accounts: InternalAccount[]): void { this.update((state) => { for (const account of accounts) { - if (state.assetsLoadingStatus[account.id] === trigger) { - delete state.assetsLoadingStatus[account.id]; + if (state.assetsLoadingStatus[account.id] === 'loading') { + state.assetsLoadingStatus[account.id] = 'loaded'; } } }); @@ -3770,22 +3667,28 @@ export class AssetsController extends BaseController< previousGroupId, }); - const releaseLock = await this.#accountRefreshMutex.acquire(); + this.#setAssetsLoadingStatus(accounts); try { - if (accounts.length > 0) { - await this.getAssets(accounts, { - chainIds: [...this.#enabledChains], - forceUpdate: true, - trigger: 'accountSwitch', - }); - } + const releaseLock = await this.#accountRefreshMutex.acquire(); + try { + if (accounts.length > 0) { + await this.getAssets(accounts, { + chainIds: [...this.#enabledChains], + forceUpdate: true, + }); + } - this.#ensureNativeBalancesDefaultZero(); - this.#ensureDefaultTrackedAssetsSeeded(); - this.#subscribeAssets({ skipInitialFetch: true }); - this.#fetchMissingPricesWithoutCache(accounts, [...this.#enabledChains]); + this.#ensureNativeBalancesDefaultZero(); + this.#ensureDefaultTrackedAssetsSeeded(); + this.#subscribeAssets({ skipInitialFetch: true }); + this.#fetchMissingPricesWithoutCache(accounts, [ + ...this.#enabledChains, + ]); + } finally { + releaseLock(); + } } finally { - releaseLock(); + this.#markAssetsLoaded(accounts); } } diff --git a/packages/assets-controller/src/index.ts b/packages/assets-controller/src/index.ts index 64e761eb8ef..4d939605e28 100644 --- a/packages/assets-controller/src/index.ts +++ b/packages/assets-controller/src/index.ts @@ -93,7 +93,7 @@ export type { DataRequest, DataResponse, AssetsUpdateMode, - AssetsLoadingTrigger, + AssetsLoadingStatus, // Middleware types Context, NextFunction, diff --git a/packages/assets-controller/src/selectors/loading.test.ts b/packages/assets-controller/src/selectors/loading.test.ts index 55861b2ec60..eb42966aa59 100644 --- a/packages/assets-controller/src/selectors/loading.test.ts +++ b/packages/assets-controller/src/selectors/loading.test.ts @@ -19,25 +19,23 @@ const createState = ( describe('loading selectors', () => { describe('getAccountLoadingStatus', () => { - it('returns the trigger for an account with an in-flight fetch', () => { + it('returns the loading status for a loading account', () => { const state = createState({ - [ACCOUNT_ID_A]: 'accountSwitch', - [ACCOUNT_ID_B]: 'unlock', + [ACCOUNT_ID_A]: 'loading', + [ACCOUNT_ID_B]: 'loaded', }); - expect(getAccountLoadingStatus(state, ACCOUNT_ID_A)).toBe( - 'accountSwitch', - ); - expect(getAccountLoadingStatus(state, ACCOUNT_ID_B)).toBe('unlock'); + expect(getAccountLoadingStatus(state, ACCOUNT_ID_A)).toBe('loading'); + expect(getAccountLoadingStatus(state, ACCOUNT_ID_B)).toBe('loaded'); }); - it('returns undefined when the account is not loading', () => { - const state = createState({ [ACCOUNT_ID_A]: 'accountSwitch' }); + it('returns undefined when the account has no loading status', () => { + const state = createState({ [ACCOUNT_ID_A]: 'loading' }); expect(getAccountLoadingStatus(state, ACCOUNT_ID_B)).toBeUndefined(); }); - it('returns undefined when no account is loading', () => { + it('returns undefined when no account has a loading status', () => { const state = createState({}); expect(getAccountLoadingStatus(state, ACCOUNT_ID_A)).toBeUndefined(); @@ -45,23 +43,23 @@ describe('loading selectors', () => { }); describe('isAccountLoading', () => { - it('returns true while the account has an in-flight fetch', () => { + it('returns true only while the account is loading', () => { const state = createState({ - [ACCOUNT_ID_A]: 'accountSwitch', - [ACCOUNT_ID_B]: 'unlock', + [ACCOUNT_ID_A]: 'loading', + [ACCOUNT_ID_B]: 'loaded', }); expect(isAccountLoading(state, ACCOUNT_ID_A)).toBe(true); - expect(isAccountLoading(state, ACCOUNT_ID_B)).toBe(true); + expect(isAccountLoading(state, ACCOUNT_ID_B)).toBe(false); }); - it('returns false when the account is not loading', () => { - const state = createState({ [ACCOUNT_ID_A]: 'accountSwitch' }); + it('returns false when the account has no loading status', () => { + const state = createState({ [ACCOUNT_ID_A]: 'loading' }); expect(isAccountLoading(state, ACCOUNT_ID_B)).toBe(false); }); - it('returns false when no account is loading', () => { + it('returns false when no account has a loading status', () => { const state = createState({}); expect(isAccountLoading(state, ACCOUNT_ID_A)).toBe(false); @@ -69,19 +67,19 @@ describe('loading selectors', () => { }); describe('getAccountsLoadingStatus', () => { - it('returns loading entries only for the requested account ids', () => { + it('returns statuses only for the requested account ids', () => { const state = createState({ - [ACCOUNT_ID_A]: 'accountSwitch', - [ACCOUNT_ID_B]: 'unlock', + [ACCOUNT_ID_A]: 'loading', + [ACCOUNT_ID_B]: 'loaded', }); expect( getAccountsLoadingStatus(state, [ACCOUNT_ID_A, ACCOUNT_ID_C]), - ).toStrictEqual({ [ACCOUNT_ID_A]: 'accountSwitch' }); + ).toStrictEqual({ [ACCOUNT_ID_A]: 'loading' }); }); - it('returns an empty record when none of the requested accounts are loading', () => { - const state = createState({ [ACCOUNT_ID_B]: 'unlock' }); + it('returns an empty record when none of the requested accounts have a status', () => { + const state = createState({ [ACCOUNT_ID_B]: 'loaded' }); expect( getAccountsLoadingStatus(state, [ACCOUNT_ID_A, ACCOUNT_ID_C]), @@ -89,7 +87,7 @@ describe('loading selectors', () => { }); it('returns an empty record for an empty account list', () => { - const state = createState({ [ACCOUNT_ID_A]: 'accountSwitch' }); + const state = createState({ [ACCOUNT_ID_A]: 'loading' }); expect(getAccountsLoadingStatus(state, [])).toStrictEqual({}); }); @@ -98,17 +96,17 @@ describe('loading selectors', () => { describe('isAnyAccountLoading', () => { it('returns true when any requested account is loading', () => { const state = createState({ - [ACCOUNT_ID_A]: 'accountSwitch', - [ACCOUNT_ID_B]: 'unlock', + [ACCOUNT_ID_A]: 'loading', + [ACCOUNT_ID_B]: 'loaded', }); - expect(isAnyAccountLoading(state, [ACCOUNT_ID_B, ACCOUNT_ID_C])).toBe( + expect(isAnyAccountLoading(state, [ACCOUNT_ID_A, ACCOUNT_ID_C])).toBe( true, ); }); it('returns false when none of the requested accounts are loading', () => { - const state = createState({ [ACCOUNT_ID_A]: 'accountSwitch' }); + const state = createState({ [ACCOUNT_ID_A]: 'loaded' }); expect(isAnyAccountLoading(state, [ACCOUNT_ID_B, ACCOUNT_ID_C])).toBe( false, @@ -118,7 +116,10 @@ describe('loading selectors', () => { it('defaults to all accounts when no account ids are given', () => { expect(isAnyAccountLoading(createState({}))).toBe(false); expect( - isAnyAccountLoading(createState({ [ACCOUNT_ID_C]: 'accountSwitch' })), + isAnyAccountLoading(createState({ [ACCOUNT_ID_C]: 'loaded' })), + ).toBe(false); + expect( + isAnyAccountLoading(createState({ [ACCOUNT_ID_C]: 'loading' })), ).toBe(true); }); }); diff --git a/packages/assets-controller/src/selectors/loading.ts b/packages/assets-controller/src/selectors/loading.ts index a51bdee9734..a7dfbcc1241 100644 --- a/packages/assets-controller/src/selectors/loading.ts +++ b/packages/assets-controller/src/selectors/loading.ts @@ -1,37 +1,23 @@ import type { AssetsControllerState } from '../AssetsController.js'; -import type { AccountId, AssetsLoadingTrigger } from '../types.js'; +import type { AccountId, AssetsLoadingStatus } from '../types.js'; /** - * Loading-state selectors over - * {@link AssetsControllerState.assetsLoadingStatus}. - * - * `assetsLoadingStatus` is transient (never persisted): an entry is present - * while a user-visible asset fetch (account switch, unlock) is in flight for - * that account, and its value is the trigger. Absence means "not loading". - * - * All selectors are synchronous and read from the controller state slice, - * so they can be used directly in `useSelector`-style subscriptions to - * `AssetsController:stateChanged`. - */ - -/** - * Get the loading trigger for a single account, if a user-visible fetch is - * currently in flight for it. + * Get the loading status for a single account. * * @param state - AssetsController state slice. * @param accountId - The account id (`InternalAccount.id`). - * @returns The trigger for the in-flight fetch, or `undefined` when the - * account is not loading. + * @returns `'loading'` while the account's assets are loading, `'loaded'` + * after its fetch has settled, or `undefined` if no fetch was triggered. */ export function getAccountLoadingStatus( state: Pick, accountId: AccountId, -): AssetsLoadingTrigger | undefined { +): AssetsLoadingStatus | undefined { return state.assetsLoadingStatus?.[accountId]; } /** - * Check whether a user-visible asset fetch is in flight for an account. + * Check whether an account's assets are currently loading. * * @param state - AssetsController state slice. * @param accountId - The account id (`InternalAccount.id`). @@ -41,37 +27,35 @@ export function isAccountLoading( state: Pick, accountId: AccountId, ): boolean { - return getAccountLoadingStatus(state, accountId) !== undefined; + return getAccountLoadingStatus(state, accountId) === 'loading'; } /** - * Get the in-flight loading triggers for a set of accounts (e.g. every - * account in the selected account group). + * Get the loading statuses for a set of accounts. * * @param state - AssetsController state slice. * @param accountIds - The account ids to report on. - * @returns A record containing an entry for each requested account that is - * currently loading, keyed by account id. + * @returns A record containing an entry for each requested account that has + * a loading status, keyed by account id. */ export function getAccountsLoadingStatus( state: Pick, accountIds: AccountId[], -): Record { - const result: Record = {}; +): Record { + const result: Record = {}; const loadingStatus = state.assetsLoadingStatus ?? {}; for (const accountId of accountIds) { - const trigger = loadingStatus[accountId]; - if (trigger !== undefined) { - result[accountId] = trigger; + const status = loadingStatus[accountId]; + if (status !== undefined) { + result[accountId] = status; } } return result; } /** - * Check whether any of the given accounts has a user-visible fetch in - * flight. When `accountIds` is omitted, checks every account with a loading - * entry in state. + * Check whether any of the given accounts is loading. When `accountIds` is + * omitted, checks every account with a loading status in state. * * @param state - AssetsController state slice. * @param accountIds - Optional account ids to restrict the check to. @@ -82,18 +66,6 @@ export function isAnyAccountLoading( accountIds?: AccountId[], ): boolean { const loadingStatus = state.assetsLoadingStatus ?? {}; - if (accountIds) { - for (const accountId of accountIds) { - if (loadingStatus[accountId] !== undefined) { - return true; - } - } - return false; - } - for (const accountId in loadingStatus) { - if (loadingStatus[accountId] !== undefined) { - return true; - } - } - return false; + const ids = accountIds ?? Object.keys(loadingStatus); + return ids.some((accountId) => loadingStatus[accountId] === 'loading'); } diff --git a/packages/assets-controller/src/types.ts b/packages/assets-controller/src/types.ts index a99c5035774..a4c875e4aae 100644 --- a/packages/assets-controller/src/types.ts +++ b/packages/assets-controller/src/types.ts @@ -408,17 +408,7 @@ export type DataResponse = { */ export type AssetsUpdateMode = 'full' | 'merge' | 'update'; -/** - * User-visible moments for which an assets fetch is surfaced as a loading - * state in {@link AssetsControllerState.assetsLoadingStatus}. - * - * - **accountSwitch**: The selected account group changed (e.g. the user - * switched accounts), so balances for the newly selected accounts are - * being fetched. - * - **unlock**: The wallet was unlocked (or tracking started after unlock), - * so balances are being fetched before anything is renderable. - */ -export type AssetsLoadingTrigger = 'accountSwitch' | 'unlock'; +export type AssetsLoadingStatus = 'loading' | 'loaded'; // ============================================================================ // DATA SOURCE <-> CONTROLLER (DIRECT CALLS, NO MESSENGER PER SOURCE) From 8f92a386d91e629f6b45001ffc9f1a82fe1919d8 Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Mon, 14 Sep 2026 21:09:17 +0100 Subject: [PATCH 5/8] fix(assets-controller): keep newer loading marker through older refresh settle Track a per-account token when the loading marker is set; only the invocation that owns the token settles it to 'loaded'. An older refresh finishing while a newer queued refresh owns the marker now leaves it 'loading' instead of clobbering it to 'loaded'. --- .../src/AssetsController.test.ts | 70 +++++++++++++++++++ .../assets-controller/src/AssetsController.ts | 37 +++++++--- 2 files changed, 97 insertions(+), 10 deletions(-) diff --git a/packages/assets-controller/src/AssetsController.test.ts b/packages/assets-controller/src/AssetsController.test.ts index 74d1de0d0e1..5faa028a362 100644 --- a/packages/assets-controller/src/AssetsController.test.ts +++ b/packages/assets-controller/src/AssetsController.test.ts @@ -3909,6 +3909,76 @@ describe('AssetsController', () => { ); }); + it('does not let an older refresh mark an account loaded while a newer queued refresh owns the marker', async () => { + const { client, arm, release } = createGatedQueryApiClient(); + + await withController( + { queryApiClient: client }, + async ({ controller, messenger }) => { + await activateTracking(messenger); + + let settleSecondFetch!: () => void; + const secondFetch = new Promise< + Awaited> + >((resolve): void => { + settleSecondFetch = (): void => resolve({}); + }); + const realGetAssets = controller.getAssets.bind(controller); + let fetchCallCount = 0; + const getAssetsSpy = jest + .spyOn(controller, 'getAssets') + .mockImplementation( + (...args: Parameters) => { + fetchCallCount += 1; + return fetchCallCount === 1 + ? realGetAssets(...args) + : secondFetch; + }, + ); + + // First switch is frozen mid-flight, holding the refresh mutex. + arm(); + (messenger.publish as CallableFunction)( + 'AccountTreeController:selectedAccountGroupChange', + 'entropy:mock-keyring-id-1/1', + 'entropy:mock-keyring-id-1/0', + ); + await flushPromises(); + + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loading', + ); + + // A second switch back to the same account queues behind the mutex + // and takes ownership of the marker. + (messenger.publish as CallableFunction)( + 'AccountTreeController:selectedAccountGroupChange', + 'entropy:mock-keyring-id-1/2', + 'entropy:mock-keyring-id-1/1', + ); + await flushPromises(); + + release(); + await flushPromises(); + + // The older refresh has settled, but the newer refresh (queued, in + // flight) still owns the marker, so the account stays loading. + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loading', + ); + + settleSecondFetch(); + await flushPromises(); + + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loaded', + ); + + getAssetsSpy.mockRestore(); + }, + ); + }); + it('marks accounts as loading on unlock, then loaded once the fetch settles', async () => { const { client, arm, release } = createGatedQueryApiClient(); diff --git a/packages/assets-controller/src/AssetsController.ts b/packages/assets-controller/src/AssetsController.ts index 7bcb16a6f53..d7b04827231 100644 --- a/packages/assets-controller/src/AssetsController.ts +++ b/packages/assets-controller/src/AssetsController.ts @@ -767,6 +767,10 @@ export class AssetsController extends BaseController< /** Serializes account-switch fetch + subscribe to prevent overlapping races. */ readonly #accountRefreshMutex = new Mutex(); + #loadingTokenCounter = 0; + + readonly #loadingTokens = new Map(); + /** * Active balance subscriptions keyed by account ID. * Each account has one logical subscription that may span multiple data sources. @@ -1366,7 +1370,7 @@ export class AssetsController extends BaseController< * @param accounts - Selected accounts to refresh. */ async #runStartupRefresh(accounts: InternalAccount[]): Promise { - this.#setAssetsLoadingStatus(accounts); + const loadingToken = this.#setAssetsLoadingStatus(accounts); try { const releaseLock = await this.#accountRefreshMutex.acquire(); try { @@ -1395,7 +1399,7 @@ export class AssetsController extends BaseController< releaseLock(); } } finally { - this.#markAssetsLoaded(accounts); + this.#markAssetsLoaded(accounts, loadingToken); } } @@ -1816,22 +1820,35 @@ export class AssetsController extends BaseController< return result; } - #setAssetsLoadingStatus(accounts: InternalAccount[]): void { + #setAssetsLoadingStatus(accounts: InternalAccount[]): number { + this.#loadingTokenCounter += 1; + const token = this.#loadingTokenCounter; + for (const account of accounts) { + this.#loadingTokens.set(account.id, token); + } this.update((state) => { for (const account of accounts) { state.assetsLoadingStatus[account.id] = 'loading'; } }); + return token; } - #markAssetsLoaded(accounts: InternalAccount[]): void { + #markAssetsLoaded(accounts: InternalAccount[], token: number): void { + const ownedAccounts = accounts.filter( + (account) => this.#loadingTokens.get(account.id) === token, + ); + if (ownedAccounts.length === 0) { + return; + } this.update((state) => { - for (const account of accounts) { - if (state.assetsLoadingStatus[account.id] === 'loading') { - state.assetsLoadingStatus[account.id] = 'loaded'; - } + for (const account of ownedAccounts) { + state.assetsLoadingStatus[account.id] = 'loaded'; } }); + for (const account of ownedAccounts) { + this.#loadingTokens.delete(account.id); + } } async getAssetsBalance( @@ -3667,7 +3684,7 @@ export class AssetsController extends BaseController< previousGroupId, }); - this.#setAssetsLoadingStatus(accounts); + const loadingToken = this.#setAssetsLoadingStatus(accounts); try { const releaseLock = await this.#accountRefreshMutex.acquire(); try { @@ -3688,7 +3705,7 @@ export class AssetsController extends BaseController< releaseLock(); } } finally { - this.#markAssetsLoaded(accounts); + this.#markAssetsLoaded(accounts, loadingToken); } } From ec06807b7ecc7fbb1a1fc1a1dc1d31787fc944de Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Mon, 14 Sep 2026 21:24:24 +0100 Subject: [PATCH 6/8] fix(assets-controller): group loading selectors and deterministic loading tests --- .../src/AssetsController.test.ts | 239 +++++++++--------- .../assets-controller/src/AssetsController.ts | 7 +- packages/assets-controller/src/index.ts | 5 +- .../src/migrations/healAssetsInfoMetadata.ts | 12 +- .../src/selectors/loading.test.ts | 234 ++++++++++------- .../src/selectors/loading.ts | 61 +++-- 6 files changed, 321 insertions(+), 237 deletions(-) diff --git a/packages/assets-controller/src/AssetsController.test.ts b/packages/assets-controller/src/AssetsController.test.ts index 5faa028a362..37114be00d8 100644 --- a/packages/assets-controller/src/AssetsController.test.ts +++ b/packages/assets-controller/src/AssetsController.test.ts @@ -13,6 +13,7 @@ import type { import type { NetworkState } from '@metamask/network-controller'; import { registerKeyringUnlockMock } from './__fixtures__/MockAssetControllerMessenger.js'; +import { waitFor } from './__fixtures__/test-utils.js'; import { AssetsController, getDefaultAssetsControllerState, @@ -91,30 +92,34 @@ function createMockQueryApiClient(): ApiPlatformClient { } /** - * A query API client whose Accounts API calls can be frozen mid-flight so a - * forced `getAssets()` run is observable while in flight. + * Fake accounts API client whose calls can be frozen and released one at a + * time, so tests can hold specific fetches in flight while the rest of the + * pipeline keeps running. `arm()` installs a fresh gate that the next gated + * call waits on; `release()` resolves the oldest pending gate. * - * While "armed", every Accounts API network/balance call returns a promise - * that stays pending until `release()` is called, then resolves with a valid - * (empty) response shape. While "disarmed", calls resolve immediately. - * - * @returns The gated client plus `arm()`/`release()` controls. + * @returns The gated client plus `arm()`/`release()` controls and a counter + * of calls currently frozen on a gate. */ function createGatedQueryApiClient(): { client: ApiPlatformClient; arm: () => void; release: () => void; + getGatedCallCount: () => number; } { let armed = false; - let releaseGate: () => void = () => undefined; - const gate = new Promise((resolve) => { - releaseGate = resolve; - }); + let gatedCallCount = 0; + const pendingGates: (() => void)[] = []; + let currentGate = new Promise(() => undefined); const gatedCall = (value: () => Result): (() => Promise) => - () => - armed ? gate.then(value) : Promise.resolve(value()); + () => { + if (!armed) { + return Promise.resolve(value()); + } + gatedCallCount += 1; + return currentGate.then(value); + }; const client = { ...createMockQueryApiClient(), @@ -139,14 +144,46 @@ function createGatedQueryApiClient(): { client, arm: (): void => { armed = true; + currentGate = new Promise((resolve) => { + pendingGates.push(resolve); + }); }, release: (): void => { - armed = false; - releaseGate(); + const resolveGate = pendingGates.shift(); + if (resolveGate) { + resolveGate(); + } else { + armed = false; + } }, + getGatedCallCount: (): number => gatedCallCount, }; } +/** + * Fake accounts API client whose balance endpoints always fail, to test how + * the loading status settles when a fetch errors out. + * + * @returns The failing client. + */ +function createFailingQueryApiClient(): ApiPlatformClient { + return { + ...createMockQueryApiClient(), + accounts: { + fetchV2SupportedNetworks: jest.fn().mockResolvedValue({ + fullSupport: ['eip155:1'], + partialSupport: [], + }), + fetchV5MultiAccountBalances: jest + .fn() + .mockRejectedValue(new Error('fetch failed')), + fetchV6MultiAccountBalances: jest + .fn() + .mockRejectedValue(new Error('fetch failed')), + }, + } as unknown as ApiPlatformClient; +} + type AllActions = MessengerActions; type AllEvents = MessengerEvents; @@ -3841,17 +3878,19 @@ describe('AssetsController', () => { 'entropy:mock-keyring-id-1/1', 'entropy:mock-keyring-id-1/0', ); - await flushPromises(); - expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( - 'loading', + await waitFor(() => + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loading', + ), ); release(); - await flushPromises(); - expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( - 'loaded', + await waitFor(() => + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loaded', + ), ); }, ); @@ -3871,11 +3910,12 @@ describe('AssetsController', () => { 'entropy:mock-keyring-id-1/1', 'entropy:mock-keyring-id-1/0', ); - await flushPromises(); // First switch is frozen mid-flight and holds the refresh mutex. - expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( - 'loading', + await waitFor(() => + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loading', + ), ); const accountB = createMockInternalAccount({ @@ -3887,22 +3927,24 @@ describe('AssetsController', () => { 'entropy:mock-keyring-id-1/2', 'entropy:mock-keyring-id-1/1', ); - await flushPromises(); // The queued switch is already marked loading for the new group's // accounts, without waiting for the previous refresh to finish. - expect(controller.state.assetsLoadingStatus).toStrictEqual({ - [MOCK_ACCOUNT_ID]: 'loading', - [accountB.id]: 'loading', - }); + await waitFor(() => + expect(controller.state.assetsLoadingStatus).toStrictEqual({ + [MOCK_ACCOUNT_ID]: 'loading', + [accountB.id]: 'loading', + }), + ); release(); - await flushPromises(); - expect(controller.state.assetsLoadingStatus).toStrictEqual({ - [MOCK_ACCOUNT_ID]: 'loaded', - [accountB.id]: 'loaded', - }); + await waitFor(() => + expect(controller.state.assetsLoadingStatus).toStrictEqual({ + [MOCK_ACCOUNT_ID]: 'loaded', + [accountB.id]: 'loaded', + }), + ); getSelectedAccountsMock.mockClear(); }, @@ -3910,32 +3952,14 @@ describe('AssetsController', () => { }); it('does not let an older refresh mark an account loaded while a newer queued refresh owns the marker', async () => { - const { client, arm, release } = createGatedQueryApiClient(); + const { client, arm, release, getGatedCallCount } = + createGatedQueryApiClient(); await withController( { queryApiClient: client }, async ({ controller, messenger }) => { await activateTracking(messenger); - let settleSecondFetch!: () => void; - const secondFetch = new Promise< - Awaited> - >((resolve): void => { - settleSecondFetch = (): void => resolve({}); - }); - const realGetAssets = controller.getAssets.bind(controller); - let fetchCallCount = 0; - const getAssetsSpy = jest - .spyOn(controller, 'getAssets') - .mockImplementation( - (...args: Parameters) => { - fetchCallCount += 1; - return fetchCallCount === 1 - ? realGetAssets(...args) - : secondFetch; - }, - ); - // First switch is frozen mid-flight, holding the refresh mutex. arm(); (messenger.publish as CallableFunction)( @@ -3943,38 +3967,36 @@ describe('AssetsController', () => { 'entropy:mock-keyring-id-1/1', 'entropy:mock-keyring-id-1/0', ); - await flushPromises(); - - expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( - 'loading', - ); + await waitFor(() => expect(getGatedCallCount()).toBeGreaterThan(0)); + const firstFetchCallCount = getGatedCallCount(); // A second switch back to the same account queues behind the mutex // and takes ownership of the marker. + arm(); (messenger.publish as CallableFunction)( 'AccountTreeController:selectedAccountGroupChange', 'entropy:mock-keyring-id-1/2', 'entropy:mock-keyring-id-1/1', ); - await flushPromises(); release(); - await flushPromises(); - // The older refresh has settled, but the newer refresh (queued, in + // The older refresh settles, but the newer refresh (queued, in // flight) still owns the marker, so the account stays loading. + await waitFor(() => + expect(getGatedCallCount()).toBeGreaterThan(firstFetchCallCount), + ); expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( 'loading', ); - settleSecondFetch(); - await flushPromises(); + release(); - expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( - 'loaded', + await waitFor(() => + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loaded', + ), ); - - getAssetsSpy.mockRestore(); }, ); }); @@ -3993,62 +4015,45 @@ describe('AssetsController', () => { messenger.publish('KeyringController:lock'); arm(); messenger.publish('KeyringController:unlock'); - await flushPromises(); - expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( - 'loading', + await waitFor(() => + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loading', + ), ); release(); - await flushPromises(); - expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( - 'loaded', + await waitFor(() => + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loaded', + ), ); }, ); }); it('marks accounts as loaded even when the startup fetch fails', async () => { - await withController(async ({ controller, messenger }) => { - const getAssetsSpy = jest - .spyOn(controller, 'getAssets') - .mockRejectedValue(new Error('fetch failed')); - - await activateTracking(messenger); - - expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( - 'loaded', - ); + await withController( + { queryApiClient: createFailingQueryApiClient() }, + async ({ controller, messenger }) => { + await activateTracking(messenger); - getAssetsSpy.mockRestore(); - }); + expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + 'loaded', + ); + }, + ); }); it('does not set the loading status for direct background fetches', async () => { - const { client, arm, release } = createGatedQueryApiClient(); - - await withController( - { queryApiClient: client }, - async ({ controller }) => { - const account = createMockInternalAccount(); - - await flushPromises(); - - arm(); - const fetchPromise = controller.getAssets([account], { - forceUpdate: true, - }); - await flushPromises(); + await withController(async ({ controller }) => { + const account = createMockInternalAccount(); - expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + await controller.getAssets([account], { forceUpdate: true }); - release(); - await fetchPromise; - - expect(controller.state.assetsLoadingStatus).toStrictEqual({}); - }, - ); + expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + }); }); it('emits state change events when the loading status is set and settled', async () => { @@ -4070,19 +4075,23 @@ describe('AssetsController', () => { 'entropy:mock-keyring-id-1/1', 'entropy:mock-keyring-id-1/0', ); - await flushPromises(); - const whileLoading = stateChanges.find( - (state) => state.assetsLoadingStatus[MOCK_ACCOUNT_ID] === 'loading', + await waitFor(() => + expect( + stateChanges.find( + (state) => + state.assetsLoadingStatus[MOCK_ACCOUNT_ID] === 'loading', + ), + ).toBeDefined(), ); - expect(whileLoading).toBeDefined(); release(); - await flushPromises(); - expect( - stateChanges.at(-1)?.assetsLoadingStatus[MOCK_ACCOUNT_ID], - ).toBe('loaded'); + await waitFor(() => + expect( + stateChanges.at(-1)?.assetsLoadingStatus[MOCK_ACCOUNT_ID], + ).toBe('loaded'), + ); }, ); }); diff --git a/packages/assets-controller/src/AssetsController.ts b/packages/assets-controller/src/AssetsController.ts index d7b04827231..927d4b0d1ee 100644 --- a/packages/assets-controller/src/AssetsController.ts +++ b/packages/assets-controller/src/AssetsController.ts @@ -936,13 +936,14 @@ export class AssetsController extends BaseController< // TEMPORARY: heal assetsInfo metadata wiped by a prior defect // (see extension migration #215 / ASSETS-3346). Remove in a future release. if (tempMigrateAssetsInfoMetadataAssets3346) { - this.update(() => - tempHealAssetsInfoMetadata({ + this.update(() => ({ + ...this.state, + ...tempHealAssetsInfoMetadata({ state: this.state, getMigrationState: tempMigrateAssetsInfoMetadataAssets3346, captureException, }), - ); + })); } this.#initializeNativeAssetsMap(queryApiClient); diff --git a/packages/assets-controller/src/index.ts b/packages/assets-controller/src/index.ts index 4d939605e28..178336facf4 100644 --- a/packages/assets-controller/src/index.ts +++ b/packages/assets-controller/src/index.ts @@ -201,10 +201,11 @@ export { } from './selectors/balance.js'; export { + getAccountGroupLoadingStatus, getAccountLoadingStatus, - getAccountsLoadingStatus, + getIsAssetsLoadingForSelectedAccountGroup, + isAccountGroupLoading, isAccountLoading, - isAnyAccountLoading, } from './selectors/loading.js'; export type { diff --git a/packages/assets-controller/src/migrations/healAssetsInfoMetadata.ts b/packages/assets-controller/src/migrations/healAssetsInfoMetadata.ts index 0af51f28559..75c0c6ae868 100644 --- a/packages/assets-controller/src/migrations/healAssetsInfoMetadata.ts +++ b/packages/assets-controller/src/migrations/healAssetsInfoMetadata.ts @@ -99,11 +99,9 @@ export type AssetsInfoHealingPatch = { const log = createModuleLogger(projectLogger, 'tempHealAssetsInfoMetadata'); -export type TempHealAssetsInfoMetadataOptions< - State extends AssetsControllerStateInternal = AssetsControllerStateInternal, -> = { +export type TempHealAssetsInfoMetadataOptions = { /** Current `AssetsController` state the healing patch is computed against. */ - state: State; + state: AssetsControllerStateInternal; /** * Host-provided getter for the untrusted legacy state root (see * `AssetsControllerOptions.tempMigrateAssetsInfoMetadataAssets3346`). @@ -123,13 +121,11 @@ export type TempHealAssetsInfoMetadataOptions< * @returns Updated controller state with the healing patch applied, or the * original state when there is nothing to heal or healing fails. */ -export function tempHealAssetsInfoMetadata< - State extends AssetsControllerStateInternal = AssetsControllerStateInternal, ->({ +export function tempHealAssetsInfoMetadata({ state, getMigrationState, captureException, -}: TempHealAssetsInfoMetadataOptions): State { +}: TempHealAssetsInfoMetadataOptions): AssetsControllerStateInternal { const reportError = (error: unknown): void => { log('Failed to heal assetsInfo metadata', error); captureException?.( diff --git a/packages/assets-controller/src/selectors/loading.test.ts b/packages/assets-controller/src/selectors/loading.test.ts index eb42966aa59..2d78dfbe0f4 100644 --- a/packages/assets-controller/src/selectors/loading.test.ts +++ b/packages/assets-controller/src/selectors/loading.test.ts @@ -1,126 +1,178 @@ +import type { AccountTreeControllerState } from '@metamask/account-tree-controller'; + import type { AssetsControllerState } from '../AssetsController.js'; import type { AccountId } from '../types.js'; import { + getAccountGroupLoadingStatus, getAccountLoadingStatus, - getAccountsLoadingStatus, + getIsAssetsLoadingForSelectedAccountGroup, + isAccountGroupLoading, isAccountLoading, - isAnyAccountLoading, } from './loading.js'; -const ACCOUNT_ID_A: AccountId = 'mock-account-id-1'; -const ACCOUNT_ID_B: AccountId = 'mock-account-id-2'; -const ACCOUNT_ID_C: AccountId = 'mock-account-id-3'; +const ACCOUNT_ID_A = 'mock-account-id-a' as AccountId; +const ACCOUNT_ID_B = 'mock-account-id-b' as AccountId; +const ACCOUNT_ID_C = 'mock-account-id-c' as AccountId; +const GROUP_ID_1 = 'mock-group-id-1'; +const GROUP_ID_2 = 'mock-group-id-2'; const createState = ( - assetsLoadingStatus: AssetsControllerState['assetsLoadingStatus'], + assetsLoadingStatus: Record, ): Pick => ({ assetsLoadingStatus, }); -describe('loading selectors', () => { - describe('getAccountLoadingStatus', () => { - it('returns the loading status for a loading account', () => { - const state = createState({ - [ACCOUNT_ID_A]: 'loading', - [ACCOUNT_ID_B]: 'loaded', - }); - - expect(getAccountLoadingStatus(state, ACCOUNT_ID_A)).toBe('loading'); - expect(getAccountLoadingStatus(state, ACCOUNT_ID_B)).toBe('loaded'); - }); - - it('returns undefined when the account has no loading status', () => { - const state = createState({ [ACCOUNT_ID_A]: 'loading' }); - - expect(getAccountLoadingStatus(state, ACCOUNT_ID_B)).toBeUndefined(); - }); - - it('returns undefined when no account has a loading status', () => { - const state = createState({}); - - expect(getAccountLoadingStatus(state, ACCOUNT_ID_A)).toBeUndefined(); +const createAccountTreeState = ( + selectedAccountGroup: string, +): AccountTreeControllerState => + ({ + accountTree: { + wallets: { + 'mock-wallet-id': { + groups: { + [GROUP_ID_1]: { accounts: [ACCOUNT_ID_A, ACCOUNT_ID_B] }, + [GROUP_ID_2]: { accounts: [ACCOUNT_ID_C] }, + }, + }, + }, + }, + selectedAccountGroup, + }) as unknown as AccountTreeControllerState; + +describe('getAccountLoadingStatus', () => { + it('returns the loading status for an account', () => { + const state = createState({ + [ACCOUNT_ID_A]: 'loading', + [ACCOUNT_ID_B]: 'loaded', }); + expect(getAccountLoadingStatus(state, ACCOUNT_ID_A)).toBe('loading'); + expect(getAccountLoadingStatus(state, ACCOUNT_ID_B)).toBe('loaded'); }); - describe('isAccountLoading', () => { - it('returns true only while the account is loading', () => { - const state = createState({ - [ACCOUNT_ID_A]: 'loading', - [ACCOUNT_ID_B]: 'loaded', - }); + it('returns undefined for an account with no loading status', () => { + expect( + getAccountLoadingStatus(createState({}), ACCOUNT_ID_C), + ).toBeUndefined(); + }); +}); - expect(isAccountLoading(state, ACCOUNT_ID_A)).toBe(true); - expect(isAccountLoading(state, ACCOUNT_ID_B)).toBe(false); - }); +describe('isAccountLoading', () => { + it('returns true while an account is loading', () => { + expect( + isAccountLoading( + createState({ [ACCOUNT_ID_A]: 'loading' }), + ACCOUNT_ID_A, + ), + ).toBe(true); + }); - it('returns false when the account has no loading status', () => { - const state = createState({ [ACCOUNT_ID_A]: 'loading' }); + it('returns false once the account is loaded or has no status', () => { + expect( + isAccountLoading(createState({ [ACCOUNT_ID_A]: 'loaded' }), ACCOUNT_ID_A), + ).toBe(false); + expect(isAccountLoading(createState({}), ACCOUNT_ID_A)).toBe(false); + }); +}); - expect(isAccountLoading(state, ACCOUNT_ID_B)).toBe(false); +describe('getAccountGroupLoadingStatus', () => { + it('returns statuses for the accounts in the group', () => { + const state = createState({ + [ACCOUNT_ID_A]: 'loading', + [ACCOUNT_ID_B]: 'loaded', }); - - it('returns false when no account has a loading status', () => { - const state = createState({}); - - expect(isAccountLoading(state, ACCOUNT_ID_A)).toBe(false); + expect( + getAccountGroupLoadingStatus( + state, + createAccountTreeState(GROUP_ID_1), + GROUP_ID_1, + ), + ).toStrictEqual({ + [ACCOUNT_ID_A]: 'loading', + [ACCOUNT_ID_B]: 'loaded', }); }); - describe('getAccountsLoadingStatus', () => { - it('returns statuses only for the requested account ids', () => { - const state = createState({ - [ACCOUNT_ID_A]: 'loading', - [ACCOUNT_ID_B]: 'loaded', - }); - - expect( - getAccountsLoadingStatus(state, [ACCOUNT_ID_A, ACCOUNT_ID_C]), - ).toStrictEqual({ [ACCOUNT_ID_A]: 'loading' }); - }); - - it('returns an empty record when none of the requested accounts have a status', () => { - const state = createState({ [ACCOUNT_ID_B]: 'loaded' }); - - expect( - getAccountsLoadingStatus(state, [ACCOUNT_ID_A, ACCOUNT_ID_C]), - ).toStrictEqual({}); - }); + it('omits accounts in the group that have no loading status', () => { + const state = createState({ [ACCOUNT_ID_A]: 'loading' }); + expect( + getAccountGroupLoadingStatus( + state, + createAccountTreeState(GROUP_ID_1), + GROUP_ID_1, + ), + ).toStrictEqual({ [ACCOUNT_ID_A]: 'loading' }); + }); - it('returns an empty record for an empty account list', () => { - const state = createState({ [ACCOUNT_ID_A]: 'loading' }); + it('returns an empty object for an unknown group', () => { + const state = createState({ [ACCOUNT_ID_A]: 'loading' }); + expect( + getAccountGroupLoadingStatus(state, createAccountTreeState(''), 'nope'), + ).toStrictEqual({}); + }); +}); - expect(getAccountsLoadingStatus(state, [])).toStrictEqual({}); +describe('isAccountGroupLoading', () => { + it('returns true while any account in the group is loading', () => { + const state = createState({ + [ACCOUNT_ID_A]: 'loaded', + [ACCOUNT_ID_B]: 'loading', }); + expect( + isAccountGroupLoading( + state, + createAccountTreeState(GROUP_ID_1), + GROUP_ID_1, + ), + ).toBe(true); }); - describe('isAnyAccountLoading', () => { - it('returns true when any requested account is loading', () => { - const state = createState({ - [ACCOUNT_ID_A]: 'loading', - [ACCOUNT_ID_B]: 'loaded', - }); - - expect(isAnyAccountLoading(state, [ACCOUNT_ID_A, ACCOUNT_ID_C])).toBe( - true, - ); + it('returns false when no account in the group is loading', () => { + const state = createState({ + [ACCOUNT_ID_A]: 'loaded', + [ACCOUNT_ID_B]: 'loaded', + [ACCOUNT_ID_C]: 'loading', }); + expect( + isAccountGroupLoading( + state, + createAccountTreeState(GROUP_ID_1), + GROUP_ID_1, + ), + ).toBe(false); + }); +}); - it('returns false when none of the requested accounts are loading', () => { - const state = createState({ [ACCOUNT_ID_A]: 'loaded' }); +describe('getIsAssetsLoadingForSelectedAccountGroup', () => { + it('returns true while any account in the selected group is loading', () => { + const state = createState({ [ACCOUNT_ID_C]: 'loading' }); + expect( + getIsAssetsLoadingForSelectedAccountGroup( + state, + createAccountTreeState(GROUP_ID_2), + ), + ).toBe(true); + }); - expect(isAnyAccountLoading(state, [ACCOUNT_ID_B, ACCOUNT_ID_C])).toBe( - false, - ); + it('returns false when the selected group is not loading', () => { + const state = createState({ + [ACCOUNT_ID_A]: 'loaded', + [ACCOUNT_ID_B]: 'loading', }); + expect( + getIsAssetsLoadingForSelectedAccountGroup( + state, + createAccountTreeState(GROUP_ID_2), + ), + ).toBe(false); + }); - it('defaults to all accounts when no account ids are given', () => { - expect(isAnyAccountLoading(createState({}))).toBe(false); - expect( - isAnyAccountLoading(createState({ [ACCOUNT_ID_C]: 'loaded' })), - ).toBe(false); - expect( - isAnyAccountLoading(createState({ [ACCOUNT_ID_C]: 'loading' })), - ).toBe(true); - }); + it('returns false when no group is selected', () => { + const state = createState({ [ACCOUNT_ID_A]: 'loading' }); + expect( + getIsAssetsLoadingForSelectedAccountGroup( + state, + createAccountTreeState(''), + ), + ).toBe(false); }); }); diff --git a/packages/assets-controller/src/selectors/loading.ts b/packages/assets-controller/src/selectors/loading.ts index a7dfbcc1241..3c1cfd19872 100644 --- a/packages/assets-controller/src/selectors/loading.ts +++ b/packages/assets-controller/src/selectors/loading.ts @@ -1,5 +1,8 @@ +import type { AccountTreeControllerState } from '@metamask/account-tree-controller'; + import type { AssetsControllerState } from '../AssetsController.js'; import type { AccountId, AssetsLoadingStatus } from '../types.js'; +import { getAccountIdsForGroup } from './balance.js'; /** * Get the loading status for a single account. @@ -13,7 +16,7 @@ export function getAccountLoadingStatus( state: Pick, accountId: AccountId, ): AssetsLoadingStatus | undefined { - return state.assetsLoadingStatus?.[accountId]; + return state.assetsLoadingStatus[accountId]; } /** @@ -31,20 +34,22 @@ export function isAccountLoading( } /** - * Get the loading statuses for a set of accounts. + * Get the loading statuses for every account in an account group. * * @param state - AssetsController state slice. - * @param accountIds - The account ids to report on. - * @returns A record containing an entry for each requested account that has - * a loading status, keyed by account id. + * @param accountTreeState - AccountTreeController state slice. + * @param groupId - The account group id. + * @returns A record containing an entry for each account in the group that + * has a loading status, keyed by account id. */ -export function getAccountsLoadingStatus( +export function getAccountGroupLoadingStatus( state: Pick, - accountIds: AccountId[], + accountTreeState: AccountTreeControllerState, + groupId: string, ): Record { + const loadingStatus = state.assetsLoadingStatus; const result: Record = {}; - const loadingStatus = state.assetsLoadingStatus ?? {}; - for (const accountId of accountIds) { + for (const accountId of getAccountIdsForGroup(accountTreeState, groupId)) { const status = loadingStatus[accountId]; if (status !== undefined) { result[accountId] = status; @@ -54,18 +59,38 @@ export function getAccountsLoadingStatus( } /** - * Check whether any of the given accounts is loading. When `accountIds` is - * omitted, checks every account with a loading status in state. + * Check whether any account in an account group is currently loading. + * + * @param state - AssetsController state slice. + * @param accountTreeState - AccountTreeController state slice. + * @param groupId - The account group id. + * @returns True while at least one account in the group is loading. + */ +export function isAccountGroupLoading( + state: Pick, + accountTreeState: AccountTreeControllerState, + groupId: string, +): boolean { + const loadingStatus = state.assetsLoadingStatus; + return getAccountIdsForGroup(accountTreeState, groupId).some( + (accountId) => loadingStatus[accountId] === 'loading', + ); +} + +/** + * Check whether the selected account group is currently loading its assets. * * @param state - AssetsController state slice. - * @param accountIds - Optional account ids to restrict the check to. - * @returns True if at least one of the accounts is loading. + * @param accountTreeState - AccountTreeController state slice. + * @returns True while any account in the selected account group is loading. */ -export function isAnyAccountLoading( +export function getIsAssetsLoadingForSelectedAccountGroup( state: Pick, - accountIds?: AccountId[], + accountTreeState: AccountTreeControllerState, ): boolean { - const loadingStatus = state.assetsLoadingStatus ?? {}; - const ids = accountIds ?? Object.keys(loadingStatus); - return ids.some((accountId) => loadingStatus[accountId] === 'loading'); + const groupId = accountTreeState.selectedAccountGroup; + if (!groupId) { + return false; + } + return isAccountGroupLoading(state, accountTreeState, groupId); } From 5355bd15586f81c61c9d55e291f79b9700cb4e9d Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Mon, 14 Sep 2026 21:58:01 +0100 Subject: [PATCH 7/8] Move assets loading marker into getAssets via decorator Per review feedback, mark the assets loading status inside getAssets itself using a trackAssetsLoading method decorator instead of in the startup-refresh and account-group-change handlers, keeping the token ownership logic so an older overlapping fetch cannot clobber a newer fetch's loading marker. --- packages/assets-controller/CHANGELOG.md | 2 +- .../src/AssetsController.test.ts | 112 +++++++++------ .../assets-controller/src/AssetsController.ts | 115 +++++---------- .../src/trackAssetsLoading.ts | 131 ++++++++++++++++++ 4 files changed, 237 insertions(+), 123 deletions(-) create mode 100644 packages/assets-controller/src/trackAssetsLoading.ts diff --git a/packages/assets-controller/CHANGELOG.md b/packages/assets-controller/CHANGELOG.md index b86348cde3c..7d96936d678 100644 --- a/packages/assets-controller/CHANGELOG.md +++ b/packages/assets-controller/CHANGELOG.md @@ -9,7 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- Add a transient, non-persisted per-account assets loading state (`assetsLoadingStatus`) that marks accounts as loading during account switches and unlock, with selectors to read it ([#10230](https://github.com/MetaMask/core/pull/10230)) +- Add a transient, non-persisted per-account assets loading state (`assetsLoadingStatus`) that marks accounts as loading while `getAssets` fetches their assets, with selectors to read it ([#10230](https://github.com/MetaMask/core/pull/10230)) ### Changed diff --git a/packages/assets-controller/src/AssetsController.test.ts b/packages/assets-controller/src/AssetsController.test.ts index 37114be00d8..86c1bc73745 100644 --- a/packages/assets-controller/src/AssetsController.test.ts +++ b/packages/assets-controller/src/AssetsController.test.ts @@ -95,10 +95,11 @@ function createMockQueryApiClient(): ApiPlatformClient { * Fake accounts API client whose calls can be frozen and released one at a * time, so tests can hold specific fetches in flight while the rest of the * pipeline keeps running. `arm()` installs a fresh gate that the next gated - * call waits on; `release()` resolves the oldest pending gate. + * call waits on; `release()` resolves the oldest pending gate and lets calls + * made after the release pass through. * - * @returns The gated client plus `arm()`/`release()` controls and a counter - * of calls currently frozen on a gate. + * @returns The gated client plus `arm()`/`release()` controls and a count of + * calls made while a gate was armed. */ function createGatedQueryApiClient(): { client: ApiPlatformClient; @@ -152,6 +153,9 @@ function createGatedQueryApiClient(): { const resolveGate = pendingGates.shift(); if (resolveGate) { resolveGate(); + // Calls made after this release pass through, while calls that + // already captured a pending gate stay frozen until it is released. + currentGate = Promise.resolve(); } else { armed = false; } @@ -3896,7 +3900,7 @@ describe('AssetsController', () => { ); }); - it('marks a queued account switch as loading immediately, while a previous refresh still holds the refresh mutex', async () => { + it('runs a queued account switch after the previous refresh finishes, marking its accounts loading then loaded', async () => { const { client, arm, release } = createGatedQueryApiClient(); await withController( @@ -3928,17 +3932,16 @@ describe('AssetsController', () => { 'entropy:mock-keyring-id-1/1', ); - // The queued switch is already marked loading for the new group's - // accounts, without waiting for the previous refresh to finish. - await waitFor(() => - expect(controller.state.assetsLoadingStatus).toStrictEqual({ - [MOCK_ACCOUNT_ID]: 'loading', - [accountB.id]: 'loading', - }), - ); + // The queued switch has not fetched yet, so its accounts are not + // marked while the previous refresh still holds the mutex. + expect( + controller.state.assetsLoadingStatus[accountB.id], + ).toBeUndefined(); release(); + // Once the previous refresh finishes, the queued switch runs, marks + // its own accounts, and settles them. await waitFor(() => expect(controller.state.assetsLoadingStatus).toStrictEqual({ [MOCK_ACCOUNT_ID]: 'loaded', @@ -3951,7 +3954,7 @@ describe('AssetsController', () => { ); }); - it('does not let an older refresh mark an account loaded while a newer queued refresh owns the marker', async () => { + it('does not let an older fetch mark an account loaded while a newer overlapping fetch owns the marker', async () => { const { client, arm, release, getGatedCallCount } = createGatedQueryApiClient(); @@ -3960,42 +3963,42 @@ describe('AssetsController', () => { async ({ controller, messenger }) => { await activateTracking(messenger); - // First switch is frozen mid-flight, holding the refresh mutex. + const account = createMockInternalAccount(); + + // Two overlapping getAssets calls for the same account: the older + // one is frozen mid-flight, then the newer one takes the marker. arm(); - (messenger.publish as CallableFunction)( - 'AccountTreeController:selectedAccountGroupChange', - 'entropy:mock-keyring-id-1/1', - 'entropy:mock-keyring-id-1/0', - ); + const olderFetch = controller.getAssets([account], { + forceUpdate: true, + }); await waitFor(() => expect(getGatedCallCount()).toBeGreaterThan(0)); - const firstFetchCallCount = getGatedCallCount(); - // A second switch back to the same account queues behind the mutex - // and takes ownership of the marker. arm(); - (messenger.publish as CallableFunction)( - 'AccountTreeController:selectedAccountGroupChange', - 'entropy:mock-keyring-id-1/2', - 'entropy:mock-keyring-id-1/1', + const newerFetch = controller.getAssets([account], { + forceUpdate: true, + }); + await waitFor(() => expect(getGatedCallCount()).toBeGreaterThan(1)); + + await waitFor(() => + expect(controller.state.assetsLoadingStatus[account.id]).toBe( + 'loading', + ), ); release(); + await olderFetch; - // The older refresh settles, but the newer refresh (queued, in - // flight) still owns the marker, so the account stays loading. - await waitFor(() => - expect(getGatedCallCount()).toBeGreaterThan(firstFetchCallCount), - ); - expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( + // The older fetch has settled, but the newer one still owns the + // marker, so the account stays loading. + expect(controller.state.assetsLoadingStatus[account.id]).toBe( 'loading', ); release(); + await newerFetch; - await waitFor(() => - expect(controller.state.assetsLoadingStatus[MOCK_ACCOUNT_ID]).toBe( - 'loaded', - ), + expect(controller.state.assetsLoadingStatus[account.id]).toBe( + 'loaded', ); }, ); @@ -4046,11 +4049,40 @@ describe('AssetsController', () => { ); }); - it('does not set the loading status for direct background fetches', async () => { - await withController(async ({ controller }) => { - const account = createMockInternalAccount(); + it('marks the loading status for direct getAssets calls as well', async () => { + const { client, arm, release } = createGatedQueryApiClient(); + + await withController( + { queryApiClient: client }, + async ({ controller, messenger }) => { + await activateTracking(messenger); + + const account = createMockInternalAccount(); + + arm(); + const fetchPromise = controller.getAssets([account], { + forceUpdate: true, + }); - await controller.getAssets([account], { forceUpdate: true }); + await waitFor(() => + expect(controller.state.assetsLoadingStatus[account.id]).toBe( + 'loading', + ), + ); + + release(); + await fetchPromise; + + expect(controller.state.assetsLoadingStatus[account.id]).toBe( + 'loaded', + ); + }, + ); + }); + + it('does not mark anything when getAssets is called with no accounts', async () => { + await withController(async ({ controller }) => { + await controller.getAssets([], { forceUpdate: true }); expect(controller.state.assetsLoadingStatus).toStrictEqual({}); }); diff --git a/packages/assets-controller/src/AssetsController.ts b/packages/assets-controller/src/AssetsController.ts index 927d4b0d1ee..c1a9036b228 100644 --- a/packages/assets-controller/src/AssetsController.ts +++ b/packages/assets-controller/src/AssetsController.ts @@ -125,6 +125,7 @@ import { isUnlockCleanupEnabled, tempHealAssetsInfoMetadata, } from './migrations/healAssetsInfoMetadata.js'; +import { trackAssetsLoading } from './trackAssetsLoading.js'; import type { AccountId, AssetPreferences, @@ -767,10 +768,6 @@ export class AssetsController extends BaseController< /** Serializes account-switch fetch + subscribe to prevent overlapping races. */ readonly #accountRefreshMutex = new Mutex(); - #loadingTokenCounter = 0; - - readonly #loadingTokens = new Map(); - /** * Active balance subscriptions keyed by account ID. * Each account has one logical subscription that may span multiple data sources. @@ -1371,36 +1368,27 @@ export class AssetsController extends BaseController< * @param accounts - Selected accounts to refresh. */ async #runStartupRefresh(accounts: InternalAccount[]): Promise { - const loadingToken = this.#setAssetsLoadingStatus(accounts); + const releaseLock = await this.#accountRefreshMutex.acquire(); try { - const releaseLock = await this.#accountRefreshMutex.acquire(); - try { - await this.getAssets(accounts, { - chainIds: [...this.#enabledChains], - forceUpdate: true, - }); - // Seed before subscribe so the price poll / update fetch sees natives - // and default tracked assets that were never returned by balance APIs. - this.#ensureNativeBalancesDefaultZero(); - this.#ensureDefaultTrackedAssetsSeeded(); - // Balances were just force-fetched — skip AccountsApi's subscribe-time poll. - this.#subscribeAssets({ skipInitialFetch: true }); - this.#fetchMissingPricesWithoutCache(accounts, [ - ...this.#enabledChains, - ]); - } catch (error) { - log('Failed to fetch assets on startup', error); - this.#ensureNativeBalancesDefaultZero(); - this.#ensureDefaultTrackedAssetsSeeded(); - this.#subscribeAssets({ skipInitialFetch: true }); - this.#fetchMissingPricesWithoutCache(accounts, [ - ...this.#enabledChains, - ]); - } finally { - releaseLock(); - } + await this.getAssets(accounts, { + chainIds: [...this.#enabledChains], + forceUpdate: true, + }); + // Seed before subscribe so the price poll / update fetch sees natives + // and default tracked assets that were never returned by balance APIs. + this.#ensureNativeBalancesDefaultZero(); + this.#ensureDefaultTrackedAssetsSeeded(); + // Balances were just force-fetched — skip AccountsApi's subscribe-time poll. + this.#subscribeAssets({ skipInitialFetch: true }); + this.#fetchMissingPricesWithoutCache(accounts, [...this.#enabledChains]); + } catch (error) { + log('Failed to fetch assets on startup', error); + this.#ensureNativeBalancesDefaultZero(); + this.#ensureDefaultTrackedAssetsSeeded(); + this.#subscribeAssets({ skipInitialFetch: true }); + this.#fetchMissingPricesWithoutCache(accounts, [...this.#enabledChains]); } finally { - this.#markAssetsLoaded(accounts, loadingToken); + releaseLock(); } } @@ -1627,6 +1615,7 @@ export class AssetsController extends BaseController< // PUBLIC API: QUERY METHODS // ============================================================================ + @trackAssetsLoading async getAssets( accounts: InternalAccount[], options?: { @@ -1821,37 +1810,6 @@ export class AssetsController extends BaseController< return result; } - #setAssetsLoadingStatus(accounts: InternalAccount[]): number { - this.#loadingTokenCounter += 1; - const token = this.#loadingTokenCounter; - for (const account of accounts) { - this.#loadingTokens.set(account.id, token); - } - this.update((state) => { - for (const account of accounts) { - state.assetsLoadingStatus[account.id] = 'loading'; - } - }); - return token; - } - - #markAssetsLoaded(accounts: InternalAccount[], token: number): void { - const ownedAccounts = accounts.filter( - (account) => this.#loadingTokens.get(account.id) === token, - ); - if (ownedAccounts.length === 0) { - return; - } - this.update((state) => { - for (const account of ownedAccounts) { - state.assetsLoadingStatus[account.id] = 'loaded'; - } - }); - for (const account of ownedAccounts) { - this.#loadingTokens.delete(account.id); - } - } - async getAssetsBalance( accounts: InternalAccount[], options?: { @@ -3685,28 +3643,21 @@ export class AssetsController extends BaseController< previousGroupId, }); - const loadingToken = this.#setAssetsLoadingStatus(accounts); + const releaseLock = await this.#accountRefreshMutex.acquire(); try { - const releaseLock = await this.#accountRefreshMutex.acquire(); - try { - if (accounts.length > 0) { - await this.getAssets(accounts, { - chainIds: [...this.#enabledChains], - forceUpdate: true, - }); - } - - this.#ensureNativeBalancesDefaultZero(); - this.#ensureDefaultTrackedAssetsSeeded(); - this.#subscribeAssets({ skipInitialFetch: true }); - this.#fetchMissingPricesWithoutCache(accounts, [ - ...this.#enabledChains, - ]); - } finally { - releaseLock(); + if (accounts.length > 0) { + await this.getAssets(accounts, { + chainIds: [...this.#enabledChains], + forceUpdate: true, + }); } + + this.#ensureNativeBalancesDefaultZero(); + this.#ensureDefaultTrackedAssetsSeeded(); + this.#subscribeAssets({ skipInitialFetch: true }); + this.#fetchMissingPricesWithoutCache(accounts, [...this.#enabledChains]); } finally { - this.#markAssetsLoaded(accounts, loadingToken); + releaseLock(); } } diff --git a/packages/assets-controller/src/trackAssetsLoading.ts b/packages/assets-controller/src/trackAssetsLoading.ts new file mode 100644 index 00000000000..2c4eeae2029 --- /dev/null +++ b/packages/assets-controller/src/trackAssetsLoading.ts @@ -0,0 +1,131 @@ +import type { InternalAccount } from '@metamask/keyring-internal-api'; + +import type { AccountId, AssetsLoadingStatus } from './types.js'; + +type AssetsLoadingStateUpdater = { + update: ( + callback: (state: { + assetsLoadingStatus: Record; + }) => void, + ) => void; +}; + +const loadingTokens = new WeakMap< + AssetsLoadingStateUpdater, + Map +>(); +let loadingTokenCounter = 0; + +/** + * Get the per-controller ownership token map, creating it on first use. + * + * @param controller - The controller whose token map to get. + * @returns The map of account ID to the token of the invocation that + * currently owns that account's loading marker. + */ +function getTokens( + controller: AssetsLoadingStateUpdater, +): Map { + let tokens = loadingTokens.get(controller); + if (tokens === undefined) { + tokens = new Map(); + loadingTokens.set(controller, tokens); + } + return tokens; +} + +/** + * Mark the given accounts as `'loading'` and return an ownership token that + * identifies this invocation as the owner of those markers. + * + * @param controller - The controller whose state should be updated. + * @param accounts - The accounts being fetched. + * @returns The ownership token, or `0` when there is nothing to mark. + */ +function markAssetsLoading( + controller: AssetsLoadingStateUpdater, + accounts: InternalAccount[], +): number { + if (accounts.length === 0) { + return 0; + } + loadingTokenCounter += 1; + const token = loadingTokenCounter; + const tokens = getTokens(controller); + for (const account of accounts) { + tokens.set(account.id, token); + } + controller.update((state) => { + for (const account of accounts) { + state.assetsLoadingStatus[account.id] = 'loading'; + } + }); + return token; +} + +/** + * Settle the markers owned by the given token to `'loaded'`. Accounts whose + * marker was re-claimed by a newer invocation are left untouched, so an older + * overlapping fetch cannot clobber a newer one's loading status. + * + * @param controller - The controller whose state should be updated. + * @param accounts - The accounts the invocation fetched. + * @param token - The ownership token returned by {@link markAssetsLoading}. + */ +function markAssetsSettled( + controller: AssetsLoadingStateUpdater, + accounts: InternalAccount[], + token: number, +): void { + if (token === 0) { + return; + } + const tokens = getTokens(controller); + const ownedAccounts = accounts.filter( + (account) => tokens.get(account.id) === token, + ); + if (ownedAccounts.length === 0) { + return; + } + controller.update((state) => { + for (const account of ownedAccounts) { + state.assetsLoadingStatus[account.id] = 'loaded'; + } + }); + for (const account of ownedAccounts) { + tokens.delete(account.id); + } +} + +/** + * Method decorator that tracks a per-account assets loading status around the + * decorated fetch method: the requested accounts are marked `'loading'` when + * the method is invoked and settle to `'loaded'` once it completes, whether it + * succeeded or failed. + * + * @param target - The decorated fetch method. + * @param _context - The decorator context. + * @returns The wrapped fetch method. + */ +export function trackAssetsLoading< + This, + Args extends [InternalAccount[], ...unknown[]], + Return, +>( + target: (this: This, ...args: Args) => Promise, + _context: ClassMethodDecoratorContext< + This, + (this: This, ...args: Args) => Promise + >, +): (this: This, ...args: Args) => Promise { + return async function (this: This, ...args: Args): Promise { + const [accounts] = args; + const controller = this as unknown as AssetsLoadingStateUpdater; + const token = markAssetsLoading(controller, accounts); + try { + return await target.call(this, ...args); + } finally { + markAssetsSettled(controller, accounts, token); + } + }; +} From d02e8ba6b6ae983764f3a5217b9571335bb36bed Mon Sep 17 00:00:00 2001 From: Prithpal Sooriya Date: Mon, 14 Sep 2026 22:19:38 +0100 Subject: [PATCH 8/8] Track assets loading tokens in controller state, skip cache reads Per review feedback: base the decorator's controller type on the real AssetsController (picking its state), store the loading ownership tokens in the controller's non-persisted state instead of a module WeakMap, and add isolated tests for the decorator. Calls that do not force an update are cache reads and no longer mark the loading status. --- packages/assets-controller/CHANGELOG.md | 2 +- .../src/AssetsController.test.ts | 3 + .../assets-controller/src/AssetsController.ts | 8 + .../src/trackAssetsLoading.test.ts | 153 ++++++++++++++++++ .../src/trackAssetsLoading.ts | 56 +++---- 5 files changed, 184 insertions(+), 38 deletions(-) create mode 100644 packages/assets-controller/src/trackAssetsLoading.test.ts diff --git a/packages/assets-controller/CHANGELOG.md b/packages/assets-controller/CHANGELOG.md index 7d96936d678..8322264144b 100644 --- a/packages/assets-controller/CHANGELOG.md +++ b/packages/assets-controller/CHANGELOG.md @@ -9,7 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- Add a transient, non-persisted per-account assets loading state (`assetsLoadingStatus`) that marks accounts as loading while `getAssets` fetches their assets, with selectors to read it ([#10230](https://github.com/MetaMask/core/pull/10230)) +- Add a transient, non-persisted per-account assets loading state (`assetsLoadingStatus`, with `assetsLoadingTokens` tracking in-flight `getAssets` fetches) that marks accounts as loading while their assets are fetched, with selectors to read it ([#10230](https://github.com/MetaMask/core/pull/10230)) ### Changed diff --git a/packages/assets-controller/src/AssetsController.test.ts b/packages/assets-controller/src/AssetsController.test.ts index 86c1bc73745..0d7cfe0e6a0 100644 --- a/packages/assets-controller/src/AssetsController.test.ts +++ b/packages/assets-controller/src/AssetsController.test.ts @@ -434,6 +434,7 @@ describe('AssetsController', () => { assetPreferences: {}, selectedCurrency: 'usd', assetsLoadingStatus: {}, + assetsLoadingTokens: {}, }); }); @@ -464,6 +465,7 @@ describe('AssetsController', () => { assetPreferences: {}, selectedCurrency: 'usd', assetsLoadingStatus: {}, + assetsLoadingTokens: {}, }); }); }); @@ -622,6 +624,7 @@ describe('AssetsController', () => { customAssets: {}, selectedCurrency: 'usd', assetsLoadingStatus: {}, + assetsLoadingTokens: {}, }); // Action handlers should be registered diff --git a/packages/assets-controller/src/AssetsController.ts b/packages/assets-controller/src/AssetsController.ts index c1a9036b228..2389c938b70 100644 --- a/packages/assets-controller/src/AssetsController.ts +++ b/packages/assets-controller/src/AssetsController.ts @@ -264,6 +264,7 @@ export type AssetsControllerState = { /** Currently-active ISO 4217 currency code */ selectedCurrency: SupportedCurrency; assetsLoadingStatus: Record; + assetsLoadingTokens: Record; }; /** @@ -287,6 +288,7 @@ export function getDefaultAssetsControllerState(): AssetsControllerState { assetPreferences: {}, selectedCurrency: 'usd', assetsLoadingStatus: {}, + assetsLoadingTokens: {}, }; } @@ -514,6 +516,12 @@ const stateMetadata: StateMetadata = { includeInDebugSnapshot: true, usedInUi: true, }, + assetsLoadingTokens: { + persist: false, + includeInStateLogs: false, + includeInDebugSnapshot: false, + usedInUi: false, + }, }; // ============================================================================ diff --git a/packages/assets-controller/src/trackAssetsLoading.test.ts b/packages/assets-controller/src/trackAssetsLoading.test.ts new file mode 100644 index 00000000000..9fd5d945f97 --- /dev/null +++ b/packages/assets-controller/src/trackAssetsLoading.test.ts @@ -0,0 +1,153 @@ +import type { InternalAccount } from '@metamask/keyring-internal-api'; + +import type { AssetsControllerState } from './AssetsController.js'; +import { trackAssetsLoading } from './trackAssetsLoading.js'; + +type FakeControllerState = Pick< + AssetsControllerState, + 'assetsLoadingStatus' | 'assetsLoadingTokens' +>; + +const ACCOUNT_1: InternalAccount = { + id: 'account-1', + address: '0x1234567890123456789012345678901234567890', + options: {}, + methods: [], + type: 'eip155:eoa', + scopes: ['eip155:1'], + metadata: { + name: 'Test Account 1', + keyring: { type: 'HD Key Tree' }, + importTime: 0, + lastSelected: 0, + }, +}; + +const ACCOUNT_2: InternalAccount = { + ...ACCOUNT_1, + id: 'account-2', + metadata: { ...ACCOUNT_1.metadata, name: 'Test Account 2' }, +}; + +/** + * Minimal stand-in for the controller: just the state and `update` the + * decorator needs, with a fetch method that stays in flight until released. + */ +class FakeAssetsController { + readonly holds: (() => void)[] = []; + + state: FakeControllerState = { + assetsLoadingStatus: {}, + assetsLoadingTokens: {}, + }; + + update(callback: (state: FakeControllerState) => void): void { + callback(this.state); + } + + /** + * Let the oldest fetch that is still in flight finish. + */ + releaseNextFetch(): void { + const resolve = this.holds.shift(); + resolve?.(); + } + + @trackAssetsLoading + async getAssets( + _accounts: InternalAccount[], + options?: { forceUpdate?: boolean; fail?: boolean }, + ): Promise { + if (options?.fail) { + throw new Error('fetch failed'); + } + await new Promise((resolve) => { + this.holds.push(resolve); + }); + } +} + +describe('trackAssetsLoading', () => { + it('marks accounts as loading while the fetch is in flight and loaded once it settles', async () => { + const controller = new FakeAssetsController(); + const fetchPromise = controller.getAssets([ACCOUNT_1, ACCOUNT_2], { + forceUpdate: true, + }); + + expect(controller.state.assetsLoadingStatus).toStrictEqual({ + [ACCOUNT_1.id]: 'loading', + [ACCOUNT_2.id]: 'loading', + }); + expect(controller.state.assetsLoadingTokens[ACCOUNT_1.id]).toBeDefined(); + + controller.releaseNextFetch(); + await fetchPromise; + + expect(controller.state.assetsLoadingStatus).toStrictEqual({ + [ACCOUNT_1.id]: 'loaded', + [ACCOUNT_2.id]: 'loaded', + }); + expect(controller.state.assetsLoadingTokens).toStrictEqual({}); + }); + + it('settles the loading status when the fetch fails', async () => { + const controller = new FakeAssetsController(); + + await expect( + controller.getAssets([ACCOUNT_1], { forceUpdate: true, fail: true }), + ).rejects.toThrow('fetch failed'); + + expect(controller.state.assetsLoadingStatus).toStrictEqual({ + [ACCOUNT_1.id]: 'loaded', + }); + expect(controller.state.assetsLoadingTokens).toStrictEqual({}); + }); + + it('does not let an older fetch settle an account a newer overlapping fetch owns', async () => { + const controller = new FakeAssetsController(); + const olderFetch = controller.getAssets([ACCOUNT_1], { + forceUpdate: true, + }); + const newerFetch = controller.getAssets([ACCOUNT_1], { + forceUpdate: true, + }); + + expect(controller.state.assetsLoadingStatus[ACCOUNT_1.id]).toBe('loading'); + + controller.releaseNextFetch(); + await olderFetch; + + expect(controller.state.assetsLoadingStatus[ACCOUNT_1.id]).toBe('loading'); + + controller.releaseNextFetch(); + await newerFetch; + + expect(controller.state.assetsLoadingStatus[ACCOUNT_1.id]).toBe('loaded'); + }); + + it('does not mark anything when there are no accounts', async () => { + const controller = new FakeAssetsController(); + const fetchPromise = controller.getAssets([], { forceUpdate: true }); + + expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + expect(controller.state.assetsLoadingTokens).toStrictEqual({}); + + controller.releaseNextFetch(); + await fetchPromise; + + expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + }); + + it('does not track calls that do not force an update', async () => { + const controller = new FakeAssetsController(); + const fetchPromise = controller.getAssets([ACCOUNT_1]); + + expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + expect(controller.state.assetsLoadingTokens).toStrictEqual({}); + + controller.releaseNextFetch(); + await fetchPromise; + + expect(controller.state.assetsLoadingStatus).toStrictEqual({}); + }); +}); diff --git a/packages/assets-controller/src/trackAssetsLoading.ts b/packages/assets-controller/src/trackAssetsLoading.ts index 2c4eeae2029..33d721c3722 100644 --- a/packages/assets-controller/src/trackAssetsLoading.ts +++ b/packages/assets-controller/src/trackAssetsLoading.ts @@ -1,41 +1,25 @@ import type { InternalAccount } from '@metamask/keyring-internal-api'; +import type { AssetsController } from './AssetsController.js'; import type { AccountId, AssetsLoadingStatus } from './types.js'; -type AssetsLoadingStateUpdater = { +/** + * The controller surface the loading tracker needs. `update` is spelled out + * here because it is protected on the controller and so cannot be picked. + */ +type AssetsLoadingStateUpdater = Pick & { update: ( callback: (state: { assetsLoadingStatus: Record; + assetsLoadingTokens: Record; }) => void, ) => void; }; -const loadingTokens = new WeakMap< - AssetsLoadingStateUpdater, - Map ->(); let loadingTokenCounter = 0; /** - * Get the per-controller ownership token map, creating it on first use. - * - * @param controller - The controller whose token map to get. - * @returns The map of account ID to the token of the invocation that - * currently owns that account's loading marker. - */ -function getTokens( - controller: AssetsLoadingStateUpdater, -): Map { - let tokens = loadingTokens.get(controller); - if (tokens === undefined) { - tokens = new Map(); - loadingTokens.set(controller, tokens); - } - return tokens; -} - -/** - * Mark the given accounts as `'loading'` and return an ownership token that + * Mark the given accounts as `'loading'` and record an ownership token that * identifies this invocation as the owner of those markers. * * @param controller - The controller whose state should be updated. @@ -51,12 +35,9 @@ function markAssetsLoading( } loadingTokenCounter += 1; const token = loadingTokenCounter; - const tokens = getTokens(controller); - for (const account of accounts) { - tokens.set(account.id, token); - } controller.update((state) => { for (const account of accounts) { + state.assetsLoadingTokens[account.id] = token; state.assetsLoadingStatus[account.id] = 'loading'; } }); @@ -70,7 +51,7 @@ function markAssetsLoading( * * @param controller - The controller whose state should be updated. * @param accounts - The accounts the invocation fetched. - * @param token - The ownership token returned by {@link markAssetsLoading}. + * @param token - The ownership token recorded by {@link markAssetsLoading}. */ function markAssetsSettled( controller: AssetsLoadingStateUpdater, @@ -80,9 +61,8 @@ function markAssetsSettled( if (token === 0) { return; } - const tokens = getTokens(controller); const ownedAccounts = accounts.filter( - (account) => tokens.get(account.id) === token, + (account) => controller.state.assetsLoadingTokens[account.id] === token, ); if (ownedAccounts.length === 0) { return; @@ -90,18 +70,17 @@ function markAssetsSettled( controller.update((state) => { for (const account of ownedAccounts) { state.assetsLoadingStatus[account.id] = 'loaded'; + delete state.assetsLoadingTokens[account.id]; } }); - for (const account of ownedAccounts) { - tokens.delete(account.id); - } } /** * Method decorator that tracks a per-account assets loading status around the * decorated fetch method: the requested accounts are marked `'loading'` when * the method is invoked and settle to `'loaded'` once it completes, whether it - * succeeded or failed. + * succeeded or failed. Calls that do not force an update are cache reads and + * are not tracked. * * @param target - The decorated fetch method. * @param _context - The decorator context. @@ -109,7 +88,7 @@ function markAssetsSettled( */ export function trackAssetsLoading< This, - Args extends [InternalAccount[], ...unknown[]], + Args extends [InternalAccount[], { forceUpdate?: boolean }?, ...unknown[]], Return, >( target: (this: This, ...args: Args) => Promise, @@ -119,7 +98,10 @@ export function trackAssetsLoading< >, ): (this: This, ...args: Args) => Promise { return async function (this: This, ...args: Args): Promise { - const [accounts] = args; + const [accounts, options] = args; + if (options?.forceUpdate !== true) { + return target.call(this, ...args); + } const controller = this as unknown as AssetsLoadingStateUpdater; const token = markAssetsLoading(controller, accounts); try {