Fix #6921: Persist Asset Balance in Portfolio

Sources/BraveWallet/Crypto/Portfolio/AddCustomAssetView.swift

@@ -15,7 +15,6 @@ struct AddCustomAssetView: View {
   @ObservedObject var userAssetStore: UserAssetsStore
   var tokenNeedsTokenId: BraveWallet.BlockchainToken?
   var supportedTokenTypes: [TokenType] = [.token, .nft]
-  var onNewAssetAdded: (() -> Void)?
 
   @Environment(\.presentationMode) @Binding private var presentationMode
 
@@ -377,7 +376,6 @@ struct AddCustomAssetView: View {
     }
     userAssetStore.addUserAsset(token) { [self] success in
       if success {
-        onNewAssetAdded?()
         presentationMode.dismiss()
       } else {
         showError = true

Sources/BraveWallet/Crypto/Portfolio/EditUserAssetsView.swift

@@ -75,7 +75,6 @@ struct EditUserAssetsView: View {
   var networkStore: NetworkStore
   var keyringStore: KeyringStore
   @ObservedObject var userAssetsStore: UserAssetsStore
-  var assetsUpdated: () -> Void
 
   @Environment(\.presentationMode) @Binding private var presentationMode
   @State private var query = ""
@@ -173,7 +172,6 @@ struct EditUserAssetsView: View {
         }
         ToolbarItemGroup(placement: .confirmationAction) {
           Button(action: {
-            assetsUpdated()
             presentationMode.dismiss()
           }) {
             Text(Strings.done)

Sources/BraveWallet/Crypto/Portfolio/PortfolioView.swift

@@ -59,9 +59,7 @@ struct PortfolioView: View {
           networkStore: networkStore,
           keyringStore: keyringStore,
           userAssetsStore: portfolioStore.userAssetsStore
-        ) {
-          cryptoStore.updateAssets()
-        }
+        )
       })
     .background(Color.clear
       .sheet(isPresented: $isPresentingAssetsFilters) {
@@ -93,9 +91,7 @@ struct PortfolioView: View {
           keyringStore: keyringStore,
           userAssetStore: cryptoStore.nftStore.userAssetsStore,
           supportedTokenTypes: [.nft]
-        ) {
-          cryptoStore.updateAssets()
-        }
+        ) 
       })
     .background(Color.clear
       .sheet(isPresented: $isPresentingNFTsFilters) {

Sources/BraveWallet/Crypto/Stores/CryptoStore.swift

@@ -154,7 +154,12 @@ public class CryptoStore: ObservableObject, WalletObserverStore {
     self.ethTxManagerProxy = ethTxManagerProxy
     self.solTxManagerProxy = solTxManagerProxy
     self.ipfsApi = ipfsApi
-    self.userAssetManager = WalletUserAssetManager(rpcService: rpcService, walletService: walletService)
+    self.userAssetManager = WalletUserAssetManager(
+      keyringService: keyringService,
+      rpcService: rpcService,
+      walletService: walletService,
+      txService: txService
+    )
     self.origin = origin
 
     self.networkStore = .init(
@@ -246,7 +251,7 @@ public class CryptoStore: ObservableObject, WalletObserverStore {
         guard let isUpdatingUserAssets = self?.isUpdatingUserAssets, !isUpdatingUserAssets else { return }
         self?.isUpdatingUserAssets = true
         Preferences.Wallet.migrateCoreToWalletUserAssetCompleted.reset()
-        WalletUserAssetGroup.removeAllGroup() {
+        self?.userAssetManager.removeUserAssetsAndBalance(for: nil) {
           self?.userAssetManager.migrateUserAssets(completion: {
             self?.updateAssets()
             self?.isUpdatingUserAssets = false
@@ -334,6 +339,9 @@ public class CryptoStore: ObservableObject, WalletObserverStore {
     marketStore.setupObservers()
     settingsStore.setupObservers()
 
+    // user asset manager's observers
+    userAssetManager.setupObservers()
+
     accountActivityStore?.setupObservers()
     assetDetailStore?.setupObservers()
     nftDetailStore?.setupObservers()
@@ -358,6 +366,9 @@ public class CryptoStore: ObservableObject, WalletObserverStore {
     marketStore.tearDown()
     settingsStore.tearDown()
 
+    // user asset manager
+    userAssetManager.tearDown()
+
     accountActivityStore?.tearDown()
     assetDetailStore?.tearDown()
     nftDetailStore?.tearDown()
@@ -591,11 +602,13 @@ public class CryptoStore: ObservableObject, WalletObserverStore {
   }
 
   func updateAutoDiscoveredAssets() {
-    // at this point, all auto-discovered assets have been added to CD
-    // update `Portfolio/Assets`
-    portfolioStore.update()
-    // fetch junk NFTs from SimpleHash which will also update `Portfolio/NFTs`
-    nftStore.fetchJunkNFTs()
+    // at this point, all auto-discovered assets have been added to CD. We now need to fetch and cache their balance
+    userAssetManager.refreshBalances { [weak self] in
+      // update `Portfolio/Assets`
+      self?.portfolioStore.update()
+      // fetch junk NFTs from SimpleHash which will also update `Portfolio/NFTs`
+      self?.nftStore.fetchJunkNFTs()
+    }
   }
 
   func prepare(isInitialOpen: Bool = false) {

Sources/BraveWallet/Crypto/Stores/NFTStore.swift

@@ -194,6 +194,9 @@ public class NFTStore: ObservableObject, WalletObserverStore {
     self.assetManager = userAssetManager
     self.txService = txService
 
+    // user asset data update observer
+    self.assetManager.addUserAssetDataObserver(self)
+
     self.setupObservers()
 
     keyringService.isLocked { [self] isLocked in
@@ -267,6 +270,9 @@ public class NFTStore: ObservableObject, WalletObserverStore {
   func update(forceUpdateNFTBalances: Bool = false) {
     self.updateTask?.cancel()
     self.updateTask = Task { @MainActor in
+      let isLocked = await keyringService.isLocked()
+      guard !isLocked else { return } // `update() will be called after unlock`
+
       self.allAccounts = await keyringService.allAccounts().accounts
         .filter { account in
           WalletConstants.supportedCoinTypes().contains(account.coin)
@@ -662,6 +668,16 @@ extension NFTStore: PreferencesObserver {
   }
 }
 
+extension NFTStore: WalletUserAssetDataObserver {
+  public func cachedBalanceRefreshed() {
+    update()
+  }
+
+  public func userAssetUpdated() {
+    update()
+  }
+}
+
 private extension Array where Element == NFTAssetViewModel {
   /// Optionally filters out NFTs not belonging to the given `selectedAccounts`.
   func optionallyFilterUnownedNFTs(

Sources/BraveWallet/Crypto/Stores/NetworkStore.swift

@@ -311,7 +311,8 @@ public class NetworkStore: ObservableObject, WalletObserverStore {
           rpcService.setNetwork(BraveWallet.MainnetChainId, coin: .eth, origin: nil, completion: { _ in })
         }
         // delete local stored user assets that in this custom network
-        assetManager.removeGroup(for: network.walletUserAssetGroupId, completion: nil)
+        // delete local stored user assets' balances that in this custom network
+        assetManager.removeUserAssetsAndBalance(for: network, completion: nil)
         Task {
           await updateChainList()
         }

Sources/BraveWallet/Crypto/Stores/PortfolioStore.swift

@@ -337,6 +337,9 @@ public class PortfolioStore: ObservableObject, WalletObserverStore {
     self.blockchainRegistry = blockchainRegistry
     self.ipfsApi = ipfsApi
     self.assetManager = userAssetManager
+
+    // cache balance update observer
+    self.assetManager.addUserAssetDataObserver(self)
 
     self.setupObservers()
 
@@ -412,6 +415,9 @@ public class PortfolioStore: ObservableObject, WalletObserverStore {
   func update() {
     self.updateTask?.cancel()
     self.updateTask = Task { @MainActor in
+      let isLocked = await keyringService.isLocked()
+      guard !isLocked else { return } // `update() will be called after unlock`
+
       self.isLoadingBalances = true
       self.allAccounts = await keyringService.allAccounts().accounts
         .filter { account in
@@ -455,31 +461,47 @@ public class PortfolioStore: ObservableObject, WalletObserverStore {
           )
         }
       }
-
-      /// Fetch balance for each token, for all accounts applicable to that token
-      tokenBalancesCache = await withTaskGroup(
-        of: [String: [String: Double]].self,
-        body: { @MainActor [tokenBalancesCache, rpcService] group in
-          for tokenNetworkAccounts in allTokenNetworkAccounts { // for each token
-            group.addTask { @MainActor in
-              let token = tokenNetworkAccounts.token
-              var tokenBalances = tokenBalancesCache[token.id] ?? [:]
-              for account in tokenNetworkAccounts.accounts { // fetch balance for this token for each account
-                let balanceForToken = await rpcService.balance(
-                  for: token,
-                  in: account,
-                  network: tokenNetworkAccounts.network
-                )
-                tokenBalances.merge(with: [account.address: balanceForToken ?? 0])
+
+      // Looping through `allTokenNetworkAccounts` to get token's cached balance
+      for tokenNetworkAccounts in allTokenNetworkAccounts {
+        if let tokenBalances = assetManager.getBalances(for: tokenNetworkAccounts.token, account: nil) {
+          var result: [String: Double] = [:]
+          for balancePerAccount in tokenBalances {
+            result.merge(with: [balancePerAccount.accountAddress: Double(balancePerAccount.balance) ?? 0])
+          }
+          tokenBalancesCache.merge(with: [tokenNetworkAccounts.token.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`
+          let fetchedTokenBalances = await withTaskGroup(
+            of: [String: [String: Double]].self,
+            body: { @MainActor [tokenBalancesCache, rpcService, assetManager] group in
+              group.addTask { @MainActor in
+                let token = tokenNetworkAccounts.token
+                var tokenBalances = tokenBalancesCache[token.id] ?? [:]
+                for account in tokenNetworkAccounts.accounts { // fetch balance for this token for each account
+                  let balanceForToken = await rpcService.balance(
+                    for: token,
+                    in: account,
+                    network: tokenNetworkAccounts.network
+                  )
+                  tokenBalances.merge(with: [account.address: balanceForToken ?? 0])
+                  assetManager.updateBalance(for: token, account: account.address, balance: "\(balanceForToken ?? 0)", completion: nil)
+                }
+                return [token.id: tokenBalances]
               }
-              return [token.id: tokenBalances]
+              return await group.reduce(into: [String: [String: Double]](), { partialResult, new in
+                partialResult.merge(with: new)
+              })
             }
-          }
-          return await group.reduce(into: [String: [String: Double]](), { partialResult, new in
-            partialResult.merge(with: new)
-          })
-        })
-
+          )
+          tokenBalancesCache.merge(with: fetchedTokenBalances)
+        }
+      }
+
       // fetch price for every token
       let allTokens = allVisibleUserAssets.flatMap(\.tokens)
       let allAssetRatioIds = allTokens.map(\.assetRatioId)
@@ -754,6 +776,16 @@ extension PortfolioStore: PreferencesObserver {
   }
 }
 
+extension PortfolioStore: WalletUserAssetDataObserver {
+  public func cachedBalanceRefreshed() {
+    update()
+  }
+
+  public func userAssetUpdated() {
+    update()
+  }
+}
+
 extension Array {
   /// `filter` helper that skips iterating through the entire array when not applying any filtering.
   @inlinable public func optionallyFilter(shouldFilter: Bool, isIncluded: (Element) throws -> Bool) rethrows -> [Element] {

Sources/BraveWallet/Crypto/Stores/UserAssetsStore.swift

@@ -13,6 +13,8 @@ public class AssetStore: ObservableObject, Equatable, WalletObserverStore {
   @Published var token: BraveWallet.BlockchainToken
   @Published var isVisible: Bool {
     didSet {
+      // update user asset's visibility. `assetManager` will broadcast to all its user assets data
+      // observers to update views. `assetManager` will also fetch user asset's balance
       assetManager.updateUserAsset(for: token, visible: isVisible, isSpam: false, isDeletedByUser: false, completion: nil)
     }
   }