@@ -3,7 +3,6 @@ import type {
33 AccountsControllerSelectedEvmAccountChangeEvent ,
44 AccountsControllerGetSelectedAccountAction ,
55 AccountsControllerListAccountsAction ,
6- AccountsControllerSelectedAccountChangeEvent ,
76} from '@metamask/accounts-controller' ;
87import type {
98 ControllerStateChangeEvent ,
@@ -16,16 +15,17 @@ import {
1615 toChecksumHexAddress ,
1716} from '@metamask/controller-utils' ;
1817import EthQuery from '@metamask/eth-query' ;
18+ import type { KeyringControllerUnlockEvent } from '@metamask/keyring-controller' ;
1919import type { InternalAccount } from '@metamask/keyring-internal-api' ;
2020import type { Messenger } from '@metamask/messenger' ;
2121import type {
2222 NetworkClient ,
2323 NetworkClientId ,
2424 NetworkControllerGetNetworkClientByIdAction ,
2525 NetworkControllerGetStateAction ,
26+ NetworkControllerNetworkAddedEvent ,
2627} from '@metamask/network-controller' ;
2728import { StaticIntervalPollingController } from '@metamask/polling-controller' ;
28- import type { PreferencesControllerGetStateAction } from '@metamask/preferences-controller' ;
2929import type {
3030 TransactionControllerTransactionConfirmedEvent ,
3131 TransactionControllerUnapprovedTransactionAddedEvent ,
@@ -174,7 +174,10 @@ export type AccountTrackerControllerActions =
174174 */
175175export type AllowedActions =
176176 | AccountsControllerListAccountsAction
177- | PreferencesControllerGetStateAction
177+ | {
178+ type : 'PreferencesController:getState' ;
179+ handler : ( ) => { isMultiAccountBalancesEnabled : boolean } ;
180+ }
178181 | AccountsControllerGetSelectedAccountAction
179182 | NetworkControllerGetStateAction
180183 | NetworkControllerGetNetworkClientByIdAction ;
@@ -199,9 +202,10 @@ export type AccountTrackerControllerEvents =
199202 */
200203export type AllowedEvents =
201204 | AccountsControllerSelectedEvmAccountChangeEvent
202- | AccountsControllerSelectedAccountChangeEvent
203205 | TransactionControllerUnapprovedTransactionAddedEvent
204- | TransactionControllerTransactionConfirmedEvent ;
206+ | TransactionControllerTransactionConfirmedEvent
207+ | NetworkControllerNetworkAddedEvent
208+ | KeyringControllerUnlockEvent ;
205209
206210/**
207211 * The messenger of the {@link AccountTrackerController}.
@@ -236,6 +240,8 @@ export class AccountTrackerController extends StaticIntervalPollingController<Ac
236240
237241 readonly #balanceFetchers: BalanceFetcher [ ] ;
238242
243+ readonly #fetchingEnabled: ( ) => boolean ;
244+
239245 /**
240246 * Creates an AccountTracker instance.
241247 *
@@ -247,6 +253,7 @@ export class AccountTrackerController extends StaticIntervalPollingController<Ac
247253 * @param options.includeStakedAssets - Whether to include staked assets in the account balances.
248254 * @param options.accountsApiChainIds - Function that returns array of chainIds that should use Accounts-API strategy (if supported by API).
249255 * @param options.allowExternalServices - Disable external HTTP calls (privacy / offline mode).
256+ * @param options.fetchingEnabled - Function that returns whether the controller is fetching enabled.
250257 */
251258 constructor ( {
252259 interval = 10000 ,
@@ -256,6 +263,7 @@ export class AccountTrackerController extends StaticIntervalPollingController<Ac
256263 includeStakedAssets = false ,
257264 accountsApiChainIds = ( ) => [ ] ,
258265 allowExternalServices = ( ) => true ,
266+ fetchingEnabled = ( ) => true ,
259267 } : {
260268 interval ?: number ;
261269 state ?: Partial < AccountTrackerControllerState > ;
@@ -264,6 +272,7 @@ export class AccountTrackerController extends StaticIntervalPollingController<Ac
264272 includeStakedAssets ?: boolean ;
265273 accountsApiChainIds ?: ( ) => ChainIdHex [ ] ;
266274 allowExternalServices ?: ( ) => boolean ;
275+ fetchingEnabled ?: ( ) => boolean ;
267276 } ) {
268277 const { selectedNetworkClientId } = messenger . call (
269278 'NetworkController:getState' ,
@@ -302,6 +311,8 @@ export class AccountTrackerController extends StaticIntervalPollingController<Ac
302311 ) ,
303312 ] ;
304313
314+ this . #fetchingEnabled = fetchingEnabled ;
315+
305316 this . setIntervalLength ( interval ) ;
306317
307318 this . messenger . subscribe (
@@ -316,23 +327,39 @@ export class AccountTrackerController extends StaticIntervalPollingController<Ac
316327 ( event ) : string => event . address ,
317328 ) ;
318329
330+ this . messenger . subscribe ( 'NetworkController:networkAdded' , async ( ) => {
331+ await this . refresh ( this . #getNetworkClientIds( ) ) ;
332+ } ) ;
333+
334+ this . messenger . subscribe ( 'KeyringController:unlock' , async ( ) => {
335+ await this . refresh ( this . #getNetworkClientIds( ) ) ;
336+ } ) ;
337+
319338 this . messenger . subscribe (
320339 'TransactionController:unapprovedTransactionAdded' ,
321340 async ( transactionMeta : TransactionMeta ) => {
322- await this . #refreshAddress(
323- [ transactionMeta . networkClientId ] ,
324- transactionMeta . txParams . from ,
325- ) ;
341+ const addresses = [ transactionMeta . txParams . from ] ;
342+ if ( transactionMeta . txParams . to ) {
343+ addresses . push ( transactionMeta . txParams . to ) ;
344+ }
345+ await this . refreshAddresses ( {
346+ networkClientIds : [ transactionMeta . networkClientId ] ,
347+ addresses,
348+ } ) ;
326349 } ,
327350 ) ;
328351
329352 this . messenger . subscribe (
330353 'TransactionController:transactionConfirmed' ,
331354 async ( transactionMeta : TransactionMeta ) => {
332- await this . #refreshAddress(
333- [ transactionMeta . networkClientId ] ,
334- transactionMeta . txParams . from ,
335- ) ;
355+ const addresses = [ transactionMeta . txParams . from ] ;
356+ if ( transactionMeta . txParams . to ) {
357+ addresses . push ( transactionMeta . txParams . to ) ;
358+ }
359+ await this . refreshAddresses ( {
360+ networkClientIds : [ transactionMeta . networkClientId ] ,
361+ addresses,
362+ } ) ;
336363 } ,
337364 ) ;
338365
@@ -540,13 +567,28 @@ export class AccountTrackerController extends StaticIntervalPollingController<Ac
540567 } ) ;
541568 }
542569
543- async #refreshAddress( networkClientIds : NetworkClientId [ ] , address : string ) {
544- const checksumAddress = toChecksumHexAddress ( address ) as ChecksumAddress ;
570+ async refreshAddresses ( {
571+ networkClientIds,
572+ addresses,
573+ } : {
574+ networkClientIds : NetworkClientId [ ] ;
575+ addresses : string [ ] ;
576+ } ) {
577+ const checksummedAddresses = addresses . map ( ( address ) =>
578+ toChecksumHexAddress ( address ) ,
579+ ) ;
580+
581+ const accounts = this . messenger
582+ . call ( 'AccountsController:listAccounts' )
583+ . filter ( ( account ) =>
584+ checksummedAddresses . includes ( toChecksumHexAddress ( account . address ) ) ,
585+ ) ;
586+
545587 await this . #refreshAccounts( {
546588 networkClientIds,
547- queryAllAccounts : false ,
548- selectedAccount : checksumAddress ,
549- allAccounts : [ ] ,
589+ queryAllAccounts : true ,
590+ selectedAccount : '0x0' ,
591+ allAccounts : accounts ,
550592 } ) ;
551593 }
552594
@@ -570,6 +612,10 @@ export class AccountTrackerController extends StaticIntervalPollingController<Ac
570612
571613 this . syncAccounts ( chainIds ) ;
572614
615+ if ( ! this . #fetchingEnabled( ) ) {
616+ return ;
617+ }
618+
573619 // Use balance fetchers with fallback strategy
574620 const aggregated : ProcessedBalance [ ] = [ ] ;
575621 let remainingChains = [ ...chainIds ] as ChainIdHex [ ] ;
0 commit comments