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

…the wallet (#24072)

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,30 @@ 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, after balance fetch
+ guard !Task.isCancelled else { return }
+ (self.userAssets, self.userNFTs) = buildAssetsAndNFTs(
+ userNetworkAssets: allUserNetworkAssets,
+ tokenBalances: tokenBalanceCache,
+ tokenPrices: tokenPricesCache,
+ nftMetadata: nftMetadataCache,
+ btcBalances: btcBalancesCache
+ )
 
  // fetch price for every user asset
  let prices: [String: String] = await assetRatioService.fetchPrices(
@@ -272,8 +281,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 +307,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 +517,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() {
+ }
+}