feat: wallet balance rollup (#101)

src/background/connections/extensionConnection/models.ts

@@ -25,6 +25,7 @@ export enum ExtensionRequest {
   BALANCES_GET = 'balances_get',
   BALANCES_START_POLLING = 'balances_start_polling',
   BALANCES_STOP_POLLING = 'balances_stop_polling',
+  BALANCES_GET_TOTAL_FOR_WALLET = 'balance_get_total_for_wallet',
   NETWORK_BALANCES_UPDATE = 'network_balances_update',
   NFT_BALANCES_GET = 'nft_balances_get',
   NFT_REFRESH_METADATA = 'nft_refresh_metadata',

src/background/connections/extensionConnection/registry.ts

@@ -129,6 +129,7 @@ import { StopBalancesPollingHandler } from '@src/background/services/balances/ha
 import { BalancesUpdatedEvents } from '@src/background/services/balances/events/balancesUpdatedEvent';
 import { UnifiedBridgeTrackTransfer } from '@src/background/services/unifiedBridge/handlers/unifiedBridgeTrackTransfer';
 import { UpdateActionTxDataHandler } from '@src/background/services/actions/handlers/updateTxData';
+import { GetTotalBalanceForWalletHandler } from '@src/background/services/balances/handlers/getTotalBalanceForWallet/getTotalBalanceForWallet';
 
 /**
  * TODO: GENERATE THIS FILE AS PART OF THE BUILD PROCESS
@@ -373,6 +374,10 @@ import { UpdateActionTxDataHandler } from '@src/background/services/actions/hand
     token: 'ExtensionRequestHandler',
     useToken: StopBalancesPollingHandler,
   },
+  {
+    token: 'ExtensionRequestHandler',
+    useToken: GetTotalBalanceForWalletHandler,
+  },
 ])
 export class ExtensionRequestHandlerRegistry {}
 

src/background/services/accounts/AccountsService.ts

@@ -333,6 +333,10 @@ export class AccountsService implements OnLock, OnUnlock {
     );
   }
 
+  getPrimaryAccountsByWalletId(walletId: string) {
+    return this.accounts.primary[walletId] ?? [];
+  }
+
   #buildAccount(
     accountData,
     importType: ImportType,

src/background/services/accounts/handlers/avalanche_getAddressesInRange.ts

@@ -2,7 +2,6 @@ import { ethErrors } from 'eth-rpc-errors';
 import { injectable } from 'tsyringe';
 import { DAppRequestHandler } from '@src/background/connections/dAppConnection/DAppRequestHandler';
 import { DAppProviderRequest } from '@src/background/connections/dAppConnection/models';
-import { Avalanche } from '@avalabs/core-wallets-sdk';
 import { SecretsService } from '../../secrets/SecretsService';
 import { NetworkService } from '../../network/NetworkService';
 import { canSkipApproval } from '@src/utils/canSkipApproval';
@@ -21,6 +20,7 @@ type Params = [
   internalLimit: number
 ];
 import { AccountsService } from '../AccountsService';
+import { getAddressesInRange } from '../utils/getAddressesInRange';
 
 const EXPOSED_DOMAINS = [
   'develop.avacloud-app.pages.dev',
@@ -71,38 +71,23 @@ export class AvalancheGetAddressesInRangeHandler extends DAppRequestHandler<
 
     if (secrets?.xpubXP) {
       if (externalLimit > 0) {
-        for (
-          let index = externalStart;
-          index < externalStart + externalLimit;
-          index++
-        ) {
-          addresses.external.push(
-            Avalanche.getAddressFromXpub(
-              secrets.xpubXP,
-              index,
-              provXP,
-              'X'
-            ).split('-')[1] as string // since addresses are the same for X/P we return them without the chain alias prefix (e.g.: fuji1jsduya7thx2ayrawf9dnw7v9jz7vc6xjycra2m)
-          );
-        }
+        addresses.external = getAddressesInRange(
+          secrets.xpubXP,
+          provXP,
+          false,
+          externalStart,
+          externalLimit
+        );
       }
 
       if (internalLimit > 0) {
-        for (
-          let index = internalStart;
-          index < internalStart + internalLimit;
-          index++
-        ) {
-          addresses.internal.push(
-            Avalanche.getAddressFromXpub(
-              secrets.xpubXP,
-              index,
-              provXP,
-              'X',
-              true
-            ).split('-')[1] as string // only X has "internal" (change) addresses, but we remove the chain alias here as well to make it consistent with the external address list
-          );
-        }
+        addresses.internal = getAddressesInRange(
+          secrets.xpubXP,
+          provXP,
+          true,
+          internalStart,
+          internalLimit
+        );
       }
     }
 

src/background/services/accounts/utils/getAddressesInRange.ts

@@ -0,0 +1,21 @@
+import { Avalanche } from '@avalabs/core-wallets-sdk';
+
+export function getAddressesInRange(
+  xpubXP: string,
+  providerXP: Avalanche.JsonRpcProvider,
+  internal = false,
+  start = 0,
+  limit = 64
+) {
+  const addresses: string[] = [];
+
+  for (let i = start; i < start + limit; i++) {
+    addresses.push(
+      Avalanche.getAddressFromXpub(xpubXP, i, providerXP, 'P', internal).split(
+        '-'
+      )[1] as string
+    );
+  }
+
+  return addresses;
+}

src/background/services/balances/BalanceAggregatorService.test.ts

@@ -218,6 +218,63 @@ describe('src/background/services/balances/BalanceAggregatorService.ts', () => {
       );
     });
 
+    it('only updates the balances of the requested accounts', async () => {
+      // Mock the existing balances for other accounts
+      (balancesServiceMock.getBalancesForNetwork as jest.Mock).mockReset();
+
+      balancesServiceMock.getBalancesForNetwork
+        .mockResolvedValueOnce({
+          [account2.addressC]: {
+            [networkToken1.symbol]: network1TokenBalance,
+          },
+        })
+        .mockResolvedValueOnce({
+          [account1.addressC]: {
+            [networkToken1.symbol]: network1TokenBalance,
+          },
+        });
+
+      // Get balances for the `account2` so they get cached
+      await service.getBalancesForNetworks(
+        [network1.chainId],
+        [account2],
+        [TokenType.NATIVE, TokenType.ERC20, TokenType.ERC721]
+      );
+
+      expect(balancesServiceMock.getBalancesForNetwork).toHaveBeenCalledTimes(
+        1
+      );
+
+      expect(service.balances).toEqual({
+        [network1.chainId]: {
+          [account2.addressC]: {
+            [networkToken1.symbol]: network1TokenBalance,
+          },
+        },
+      });
+
+      // Now get the balances for the first account and verify the `account2` balances are kept in cache
+      await service.getBalancesForNetworks(
+        [network1.chainId],
+        [account1],
+        [TokenType.NATIVE, TokenType.ERC20, TokenType.ERC721]
+      );
+
+      expect(balancesServiceMock.getBalancesForNetwork).toHaveBeenCalledTimes(
+        2
+      );
+      expect(service.balances).toEqual({
+        [network1.chainId]: {
+          [account1.addressC]: {
+            [networkToken1.symbol]: network1TokenBalance,
+          },
+          [account2.addressC]: {
+            [networkToken1.symbol]: network1TokenBalance,
+          },
+        },
+      });
+    });
+
     it('can fetch the balance for multiple networks and one account', async () => {
       const balances = await service.getBalancesForNetworks(
         [network1.chainId, network2.chainId],

src/background/services/balances/BalanceAggregatorService.ts

@@ -6,7 +6,7 @@ import { BalancesService } from './BalancesService';
 import { NetworkService } from '../network/NetworkService';
 import { EventEmitter } from 'events';
 import * as Sentry from '@sentry/browser';
-import { isEqual, pick } from 'lodash';
+import { isEqual, omit, pick } from 'lodash';
 
 import { LockService } from '../lock/LockService';
 import { StorageService } from '../storage/StorageService';
@@ -55,7 +55,8 @@ export class BalanceAggregatorService implements OnLock, OnUnlock {
   async getBalancesForNetworks(
     chainIds: number[],
     accounts: Account[],
-    tokenTypes: TokenType[]
+    tokenTypes: TokenType[],
+    cacheResponse = true
   ): Promise<{ tokens: Balances; nfts: Balances<NftTokenWithBalance> }> {
     const sentryTracker = Sentry.startTransaction({
       name: 'BalanceAggregatorService: getBatchedUpdatedBalancesForNetworks',
@@ -143,14 +144,15 @@ export class BalanceAggregatorService implements OnLock, OnUnlock {
       for (const [chainId, chainBalances] of freshData) {
         for (const [address, addressBalance] of Object.entries(chainBalances)) {
           aggregatedBalances[chainId] = {
+            ...omit(aggregatedBalances[chainId], address), // Keep cached balances for other accounts
             ...chainBalances,
             [address]: addressBalance,
           };
         }
       }
     }
 
-    if (hasChanges && !this.lockService.locked) {
+    if (cacheResponse && hasChanges && !this.lockService.locked) {
       this.#balances = aggregatedBalances;
       this.#nfts = aggregatedNfts;