feat(wallet): Apply persistent token/account balance for the rest of the wallet

ios/brave-ios/Sources/BraveWallet/Crypto/Stores/AccountActivityStore.swift

@@ -202,13 +202,6 @@ class AccountActivityStore: ObservableObject, WalletObserverStore {
       let allTokens = await blockchainRegistry.allTokens(in: networksForAccountCoin).flatMap(
         \.tokens
       )
-      (self.userAssets, self.userNFTs) = buildAssetsAndNFTs(
-        userNetworkAssets: allUserNetworkAssets,
-        tokenBalances: tokenBalanceCache,
-        tokenPrices: tokenPricesCache,
-        nftMetadata: nftMetadataCache,
-        btcBalances: btcBalancesCache
-      )
       let allAccountsForCoin = await keyringService.allAccounts().accounts.filter {
         $0.coin == account.coin
       }
@@ -228,8 +221,6 @@ class AccountActivityStore: ObservableObject, WalletObserverStore {
       )
 
       self.isLoadingAccountFiat = true
-      // TODO: cleanup with balance caching with issue
-      // https://github.com/brave/brave-browser/issues/36764
       var tokenBalances: [String: Double] = [:]
       if account.coin == .btc {
         let networkAsset = allUserNetworkAssets.first {
@@ -244,12 +235,40 @@ class AccountActivityStore: ObservableObject, WalletObserverStore {
           tokenBalances = [btcToken.id: btcTotalBalance]
         }
       } else {
-        tokenBalances = await self.rpcService.fetchBalancesForTokens(
-          account: account,
-          networkAssets: allUserNetworkAssets
-        )
+        if let accountBalances = self.assetManager.getBalances(for: nil, account: account.id) {
+          tokenBalances = accountBalances.reduce(into: [String: Double]()) {
+            let tokenId =
+              $1.contractAddress + $1.chainId
+              + $1.symbol + $1.tokenId
+            $0[tokenId] = Double($1.balance) ?? 0
+          }
+        } else {
+          tokenBalances = await self.rpcService.fetchBalancesForTokens(
+            account: account,
+            networkAssets: allUserNetworkAssets
+          )
+        }
       }
       tokenBalanceCache.merge(with: tokenBalances)
+      // update assets, NFTs, transactions after balance fetch
+      guard !Task.isCancelled else { return }
+      (self.userAssets, self.userNFTs) = buildAssetsAndNFTs(
+        userNetworkAssets: allUserNetworkAssets,
+        tokenBalances: tokenBalanceCache,
+        tokenPrices: tokenPricesCache,
+        nftMetadata: nftMetadataCache,
+        btcBalances: btcBalancesCache
+      )
+      self.transactionSections = buildTransactionSections(
+        transactions: transactions,
+        networksForCoin: [account.coin: networksForAccountCoin],
+        accountInfos: allAccountsForCoin,
+        userAssets: allUserAssets,
+        allTokens: allTokens,
+        tokenPrices: tokenPricesCache,
+        nftMetadata: nftMetadataCache,
+        solEstimatedTxFees: solEstimatedTxFeesCache
+      )
 
       // fetch price for every user asset
       let prices: [String: String] = await assetRatioService.fetchPrices(
@@ -272,8 +291,8 @@ class AccountActivityStore: ObservableObject, WalletObserverStore {
       self.accountTotalFiat = currencyFormatter.formatAsFiat(totalFiat) ?? "$0.00"
       self.isLoadingAccountFiat = false
 
-      guard !Task.isCancelled else { return }
       // update assets, NFTs, transactions after balance & price fetch
+      guard !Task.isCancelled else { return }
       (self.userAssets, self.userNFTs) = buildAssetsAndNFTs(
         userNetworkAssets: allUserNetworkAssets,
         tokenBalances: tokenBalanceCache,
@@ -298,7 +317,7 @@ class AccountActivityStore: ObservableObject, WalletObserverStore {
         ipfsApi: ipfsApi
       )
       nftMetadataCache.merge(with: allNFTMetadata)
-
+      // update assets, NFTs, transactions after balance & price & metadata fetch
       guard !Task.isCancelled else { return }
       (self.userAssets, self.userNFTs) = buildAssetsAndNFTs(
         userNetworkAssets: allUserNetworkAssets,
@@ -508,6 +527,7 @@ class AccountActivityStore: ObservableObject, WalletObserverStore {
 
 extension AccountActivityStore: WalletUserAssetDataObserver {
   public func cachedBalanceRefreshed() {
+    update()
   }
 
   public func userAssetUpdated() {

ios/brave-ios/Sources/BraveWallet/Crypto/Stores/AccountsStore.swift

@@ -34,7 +34,7 @@ class AccountsStore: ObservableObject, WalletObserverStore {
   let currencyFormatter: NumberFormatter = .usdCurrencyFormatter
 
   private typealias TokenBalanceCache = [String: [String: Double]]
-  /// Cache of token balances for each account. [account.cacheBalanceKey: [token.id: balance]]
+  /// Cache of token balances for each account. `[account.id: [token.id: balance]]`
   private var tokenBalancesCache: TokenBalanceCache = [:]
   /// Cache of prices for each token. The key is the token's `assetRatioId`.
   private var pricesCache: [String: String] = [:]
@@ -72,6 +72,7 @@ class AccountsStore: ObservableObject, WalletObserverStore {
 
   func setupObservers() {
     guard !isObserving else { return }
+    self.userAssetManager.addUserAssetDataObserver(self)
     self.keyringServiceObserver = KeyringServiceObserver(
       keyringService: keyringService,
       _accountsChanged: { [weak self] in
@@ -141,35 +142,69 @@ class AccountsStore: ObservableObject, WalletObserverStore {
     for accounts: [BraveWallet.AccountInfo],
     networkAssets allNetworkAssets: [NetworkAssets]
   ) async {
+    // Update BTC account balance
+    if accounts.contains(where: { $0.coin == .btc }) {
+      await withTaskGroup(
+        of: [String: [String: Double]].self
+      ) { [bitcoinWalletService] group in
+        for account in accounts where account.coin == .btc {
+          group.addTask {
+            let btcBalance =
+              await bitcoinWalletService.fetchBTCBalance(
+                accountId: account.accountId,
+                type: .total
+              ) ?? 0
+            if let btcToken = allNetworkAssets.first(where: {
+              $0.network.supportedKeyrings.contains(
+                account.keyringId.rawValue as NSNumber
+              )
+            })?.tokens.first {
+              return [account.id: [btcToken.id: btcBalance]]
+            }
+            return [:]
+          }
+        }
+        for await accountBTCBalances in group {
+          tokenBalancesCache.merge(with: accountBTCBalances)
+        }
+      }
+    }
+    // Update non-BTC account balance
     let balancesForAccounts = await withTaskGroup(
       of: TokenBalanceCache.self,
       body: { group in
-        for account in accounts {
-          group.addTask {
-            // TODO: cleanup with balance caching with issue
-            // https://github.com/brave/brave-browser/issues/36764
-            var balancesForTokens: [String: Double] = [:]
-            if account.coin == .btc {
-              let networkAssets = allNetworkAssets.first {
-                $0.network.supportedKeyrings.contains(account.keyringId.rawValue as NSNumber)
-              }
-              if let btc = networkAssets?.tokens.first,
-                let btcBalance = await self.bitcoinWalletService.fetchBTCBalance(
-                  accountId: account.accountId,
-                  type: .total
-                )
-              {
-                balancesForTokens = [btc.id: btcBalance]
-              }
-            } else {
+        for account in accounts where account.coin != .btc {
+          if let allTokenBalance = self.userAssetManager.getBalances(
+            for: nil,
+            account: account.id
+          ) {
+            var result: [String: Double] = [:]
+            for balancePerToken in allTokenBalance {
+              let tokenId =
+                balancePerToken.contractAddress + balancePerToken.chainId
+                + balancePerToken.symbol + balancePerToken.tokenId
+              result.merge(with: [
+                tokenId: Double(balancePerToken.balance) ?? 0
+              ])
+            }
+            self.tokenBalancesCache.merge(with: [account.id: result])
+          } else {
+            // 1. We have a user asset from CD but wallet has never
+            // fetched it's balance. Should never happen. But we will fetch its
+            // balance and cache it in CD.
+            // 2. Test Cases will come here, we will fetch balance using
+            // a mock `rpcService` and `bitcoinWalletService`
+            group.addTask {
+              var balancesForTokens: [String: Double] = [:]
               balancesForTokens = await self.rpcService.fetchBalancesForTokens(
                 account: account,
                 networkAssets: allNetworkAssets
               )
+              return [account.id: balancesForTokens]
             }
-            return [account.id: balancesForTokens]
           }
         }
+
         return await group.reduce(
           into: TokenBalanceCache(),
           { partialResult, new in
@@ -292,3 +327,12 @@ class AccountsStore: ObservableObject, WalletObserverStore {
     tokenBalancesCache[account.id]?[tokenId]
   }
 }
+
+extension AccountsStore: WalletUserAssetDataObserver {
+  func cachedBalanceRefreshed() {
+    update()
+  }
+
+  func userAssetUpdated() {
+  }
+}

ios/brave-ios/Sources/BraveWallet/Crypto/Stores/AssetDetailStore.swift

@@ -170,6 +170,7 @@ class AssetDetailStore: ObservableObject, WalletObserverStore {
 
   func setupObservers() {
     guard !isObserving else { return }
+    self.assetManager.addUserAssetDataObserver(self)
     self.keyringServiceObserver = KeyringServiceObserver(
       keyringService: keyringService,
       _accountsChanged: { [weak self] in
@@ -457,20 +458,25 @@ class AssetDetailStore: ObservableObject, WalletObserverStore {
       @MainActor group -> [AccountBalance] in
       for accountAssetViewModel in accountAssetViewModels {
         group.addTask { @MainActor in
-          // TODO: cleanup with balance caching with issue
-          // https://github.com/brave/brave-browser/issues/36764
           var tokenBalance: Double?
           if accountAssetViewModel.account.coin == .btc {
             tokenBalance = await self.bitcoinWalletService.fetchBTCBalance(
               accountId: accountAssetViewModel.account.accountId,
               type: .total
             )
           } else {
-            tokenBalance = await self.rpcService.balance(
+            if let assetBalancePerAccount = self.assetManager.getBalances(
               for: token,
-              in: accountAssetViewModel.account,
-              network: network
-            )
+              account: accountAssetViewModel.account.id
+            )?.first {
+              tokenBalance = Double(assetBalancePerAccount.balance)
+            } else {
+              tokenBalance = await self.rpcService.balance(
+                for: token,
+                in: accountAssetViewModel.account,
+                network: network
+              )
+            }
           }
           return [AccountBalance(accountAssetViewModel.account, tokenBalance)]
         }
@@ -694,3 +700,12 @@ extension AssetDetailStore: BraveWalletBraveWalletServiceObserver {
   func onResetWallet() {
   }
 }
+
+extension AssetDetailStore: WalletUserAssetDataObserver {
+  func cachedBalanceRefreshed() {
+    update()
+  }
+
+  func userAssetUpdated() {
+  }
+}

ios/brave-ios/Sources/BraveWallet/Crypto/Stores/NFTDetailStore.swift

@@ -118,6 +118,8 @@ class NFTDetailStore: ObservableObject, WalletObserverStore {
   }
 
   func setupObservers() {
+    guard !isObserving else { return }
+    self.assetManager.addUserAssetDataObserver(self)
     self.txServiceObserver = TxServiceObserver(
       txService: txService,
       _onTransactionStatusChanged: { [weak self] txInfo in
@@ -142,9 +144,7 @@ class NFTDetailStore: ObservableObject, WalletObserverStore {
         networkInfo = network
       }
 
-      if owner == nil {
-        updateOwner()
-      }
+      updateOwner()
 
       if nftMetadata == nil {
         isLoading = true
@@ -182,15 +182,22 @@ class NFTDetailStore: ObservableObject, WalletObserverStore {
       let accounts = await keyringService.allAccounts().accounts
       let nftBalances: [String: Int] = await withTaskGroup(
         of: [String: Int].self,
-        body: { @MainActor [rpcService, nft] group in
+        body: { @MainActor [assetManager, rpcService, nft] group in
           for account in accounts where account.coin == nft.coin {
             group.addTask { @MainActor in
-              let balanceForToken = await rpcService.balance(
+              if let assetBalance = assetManager.getBalances(
                 for: nft,
-                in: account,
-                network: network
-              )
-              return [account.id: Int(balanceForToken ?? 0)]
+                account: account.id
+              )?.first {
+                return [account.id: (assetBalance.balance as NSString).integerValue]
+              } else {
+                let balanceForToken = await rpcService.balance(
+                  for: nft,
+                  in: account,
+                  network: network
+                )
+                return [account.id: Int(balanceForToken ?? 0)]
+              }
             }
           }
           return await group.reduce(
@@ -217,3 +224,12 @@ class NFTDetailStore: ObservableObject, WalletObserverStore {
     }
   }
 }
+
+extension NFTDetailStore: WalletUserAssetDataObserver {
+  func cachedBalanceRefreshed() {
+    update()
+  }
+
+  func userAssetUpdated() {
+  }
+}