@eth-optimism/actions-sdk 0.2.0 → 0.4.0

package/src/lend/providers/morpho/contracts.ts

@@ -1,6 +1,20 @@
-import { baseSepolia } from 'viem/chains'
+import {
+  base,
+  baseSepolia,
+  ink,
+  mainnet,
+  mode,
+  optimism,
+  soneium,
+  unichain,
+  worldchain,
+} from 'viem/chains'
 
-import type { MorphoContractsRegistry } from '@/types/lend/contracts.js'
+import type { SupportedChainId } from '@/constants/supportedChains.js'
+import type {
+  MorphoContracts,
+  MorphoContractsRegistry,
+} from '@/types/lend/contracts.js'
 
 /**
  * Morpho Blue core contract - same address on all chains via CREATE2
@@ -8,10 +22,44 @@ import type { MorphoContractsRegistry } from '@/types/lend/contracts.js'
 const MORPHO_BLUE = '0xBBBBBbbBBb9cC5e90e3b3Af64bdAF62C37EEFFCb' as const
 
 /**
- * Contract addresses for chains NOT supported by Morpho SDK (testnets).
- * Mainnet/Base use the SDK which provides richer data including rewards.
+ * Morpho Blue contract addresses per chain.
+ * Mainnet chains use the SDK for richer data (rewards, allocations) when available.
+ * These contracts serve as the on-chain fallback and as the canonical deployment registry.
+ * @see https://github.com/morpho-org/sdks
  */
 export const MORPHO_CONTRACTS: MorphoContractsRegistry = {
+  [mainnet.id]: {
+    morphoBlue: MORPHO_BLUE,
+    irm: '0x870aC11D48B15DB9a138Cf899d20F13F79Ba00BC',
+  },
+  [optimism.id]: {
+    morphoBlue: MORPHO_BLUE,
+    irm: '0x8cD70A8F399428456b29546BC5dBe10ab6a06ef6',
+  },
+  [base.id]: {
+    morphoBlue: MORPHO_BLUE,
+    irm: '0x46415998764C29aB2a25CbeA6254146D50D22687',
+  },
+  [unichain.id]: {
+    morphoBlue: MORPHO_BLUE,
+    irm: '0x9a6061d51743B31D2c3Be75D83781Fa423f53F0E',
+  },
+  [worldchain.id]: {
+    morphoBlue: MORPHO_BLUE,
+    irm: '0x34E99D604751a72cF8d0CFDf87069292d82De472',
+  },
+  [ink.id]: {
+    morphoBlue: MORPHO_BLUE,
+    irm: '0x9515407b1512F53388ffE699524100e7270Ee57B',
+  },
+  [soneium.id]: {
+    morphoBlue: MORPHO_BLUE,
+    irm: '0x68F9b666b984527A7c145Db4103Cc6d3171C797F',
+  },
+  [mode.id]: {
+    morphoBlue: MORPHO_BLUE,
+    irm: '0xE3d46Ae190Cb39ccA3655E966DcEF96b4eAe1d1c',
+  },
   [baseSepolia.id]: {
     morphoBlue: MORPHO_BLUE,
     irm: '0x46415998764C29aB2a25CbeA6254146D50D22687',
@@ -20,18 +68,19 @@ export const MORPHO_CONTRACTS: MorphoContractsRegistry = {
 
 /**
  * Get Morpho contracts for a chain
- * @param chainId - Chain ID
- * @returns Morpho contracts if supported, undefined otherwise
  */
-export function getMorphoContracts(chainId: number) {
-  return MORPHO_CONTRACTS[chainId]
+export function getMorphoContracts(
+  chainId: number,
+): MorphoContracts | undefined {
+  return MORPHO_CONTRACTS[chainId as SupportedChainId]
 }
 
 /**
- * Check if Morpho is supported on a chain
- * @param chainId - Chain ID
- * @returns true if Morpho is deployed on this chain
+ * Get all chain IDs where Morpho contracts are deployed.
+ * Returns chains present in the local contracts registry.
+ * Filtering against ACTIONS_SUPPORTED_CHAIN_IDS and developer-configured chains
+ * is handled by the LendProvider base class.
  */
-export function isMorphoChainSupported(chainId: number): boolean {
-  return chainId in MORPHO_CONTRACTS
+export function getSupportedChainIds(): number[] {
+  return Object.keys(MORPHO_CONTRACTS).map(Number)
 }

package/src/lend/providers/morpho/sdk.ts

@@ -7,14 +7,15 @@ import {
 } from '@morpho-org/blue-sdk-viem'
 import type { Address, PublicClient } from 'viem'
 
+import { NATIVELY_SUPPORTED_ASSETS } from '@/constants/assets.js'
 import {
   fetchRewards,
   type RewardsBreakdown,
 } from '@/lend/providers/morpho/api.js'
 import { getMorphoContracts } from '@/lend/providers/morpho/contracts.js'
 import type { ChainManager } from '@/services/ChainManager.js'
-import { SUPPORTED_TOKENS } from '@/supported/tokens.js'
 import type { LendProviderConfig } from '@/types/actions.js'
+import type { Asset } from '@/types/asset.js'
 import type {
   ApyBreakdown,
   LendMarket,
@@ -27,29 +28,42 @@ import { SECONDS_PER_YEAR } from '@/utils/constants.js'
 /**
  * Fetch and calculate rewards breakdown from Morpho GraphQL API
  * @param vaultAddress - Vault address
+ * @param supportedAssets - Configured assets for reward categorization
+ * @param chainId - Chain ID for address lookup
  * @returns Promise resolving to rewards breakdown
  */
 export async function fetchAndCalculateRewards(
   vaultAddress: Address,
+  chainId: number,
+  marketAsset?: Asset,
 ): Promise<RewardsBreakdown> {
-  const vaultData = await fetchRewards(vaultAddress)
+  const vaultData = await fetchRewards(vaultAddress, chainId)
 
   if (!vaultData) {
-    // Initialize empty rewards object with all supported tokens + other
-    const emptyRewards: Record<string, number> = {
-      other: 0,
-      totalRewards: 0,
-    }
+    return buildEmptyRewards(chainId, marketAsset)
+  }
 
-    // Add all supported tokens (lowercase) to the rewards object
-    SUPPORTED_TOKENS.forEach((token) => {
-      emptyRewards[token.metadata.symbol.toLowerCase()] = 0
-    })
+  return calculateRewardsBreakdown(vaultData, chainId, marketAsset)
+}
 
-    return emptyRewards as unknown as RewardsBreakdown
+/**
+ * Build an empty rewards object with all supported asset addresses initialized to 0
+ */
+function buildEmptyRewards(
+  chainId: number,
+  marketAsset?: Asset,
+): RewardsBreakdown {
+  const assets = marketAsset
+    ? [...NATIVELY_SUPPORTED_ASSETS, marketAsset]
+    : NATIVELY_SUPPORTED_ASSETS
+  const emptyRewards: Record<string, number> = { other: 0, totalRewards: 0 }
+  for (const token of assets) {
+    const addr = token.address[chainId as keyof typeof token.address]
+    if (addr && addr !== 'native') {
+      emptyRewards[addr.toLowerCase()] = 0
+    }
   }
-
-  return calculateRewardsBreakdown(vaultData)
+  return emptyRewards as RewardsBreakdown
 }
 
 /**
@@ -238,7 +252,7 @@ async function calculateVaultApy(
 }
 
 /**
- * Fetch vault data via direct on-chain queries (for testnets)
+ * Fetch vault data via direct on-chain queries (fallback when SDK unavailable)
  */
 async function fetchVaultDataOnChain(
   marketId: LendMarketId,
@@ -332,7 +346,7 @@ export async function getVault(params: GetVaultParams): Promise<LendMarket> {
     params.marketId.chainId,
   )
 
-  // Try SDK first for supported chains (mainnets)
+  // Try SDK first for SDK-supported chains
   if (isSdkSupportedChain(params.marketId.chainId)) {
     try {
       const vault = await fetchAccrualVault(
@@ -343,9 +357,11 @@ export async function getVault(params: GetVaultParams): Promise<LendMarket> {
       // Fetch rewards data from API
       const rewardsBreakdown = await fetchAndCalculateRewards(
         params.marketId.address,
+        params.marketId.chainId,
+        marketConfig.asset,
       ).catch((error) => {
         console.error('Failed to fetch rewards data:', error)
-        return { usdc: 0, morpho: 0, other: 0, totalRewards: 0 }
+        return buildEmptyRewards(params.marketId.chainId, marketConfig.asset)
       })
 
       const apyBreakdown = calculateApyBreakdown(vault, rewardsBreakdown)
@@ -372,7 +388,7 @@ export async function getVault(params: GetVaultParams): Promise<LendMarket> {
     }
   }
 
-  // Fallback to direct on-chain queries for testnets or if SDK fails
+  // Fallback to direct on-chain queries if SDK unavailable or fails
   const contracts = getMorphoContracts(params.marketId.chainId)
   if (contracts) {
     return fetchVaultDataOnChain(
@@ -480,38 +496,59 @@ export function calculateApyBreakdown(
   }
 }
 
+/**
+ * Categorize a reward asset by its address. If the address matches a supported asset, use that.
+ * Otherwise, fall back to 'other'.
+ */
+function categorizeRewardAsset(
+  rewardAssetAddress: string | undefined,
+  knownAddresses: Set<string>,
+): string {
+  if (!rewardAssetAddress) return 'other'
+  const normalized = rewardAssetAddress.toLowerCase()
+  return knownAddresses.has(normalized) ? normalized : 'other'
+}
+
 /**
  * Calculate detailed rewards breakdown from vault and market allocations
  * @param apiVault - Vault data from GraphQL API
+ * @param supportedAssets - Configured assets for reward categorization
+ * @param chainId - Chain ID for address lookup
  * @returns Detailed rewards breakdown
  */
-export function calculateRewardsBreakdown(apiVault: any): RewardsBreakdown {
-  // Initialize rewards object with all supported tokens + other
-  const rewardsByCategory: Record<string, number> = {
-    other: 0,
+export function calculateRewardsBreakdown(
+  apiVault: any,
+  chainId: number,
+  marketAsset?: Asset,
+): RewardsBreakdown {
+  const assets = marketAsset
+    ? [...NATIVELY_SUPPORTED_ASSETS, marketAsset]
+    : NATIVELY_SUPPORTED_ASSETS
+  // Build set of known asset addresses on this chain
+  const knownAddresses = new Set<string>()
+  for (const token of assets) {
+    const addr = token.address[chainId as keyof typeof token.address]
+    if (addr && addr !== 'native') {
+      knownAddresses.add(addr.toLowerCase())
+    }
   }
 
-  // Add all supported tokens (lowercase) to the rewards object
-  SUPPORTED_TOKENS.forEach((token) => {
-    rewardsByCategory[token.metadata.symbol.toLowerCase()] = 0
-  })
+  // Initialize rewards object with all known addresses + other
+  const rewardsByCategory: Record<string, number> = { other: 0 }
+  for (const addr of knownAddresses) {
+    rewardsByCategory[addr] = 0
+  }
 
   // Calculate vault-level rewards
   if (apiVault.state?.rewards && apiVault.state.rewards.length > 0) {
-    apiVault.state.rewards.forEach((reward: any) => {
+    for (const reward of apiVault.state.rewards) {
       const rewardApr = reward.supplyApr || 0
-      const assetSymbol = reward.asset.symbol
-
-      // Use the symbol from API response for categorization
-      const category = assetSymbol ? assetSymbol.toLowerCase() : 'other'
-
-      // Add to appropriate category if supported, otherwise to 'other'
-      if (category in rewardsByCategory) {
-        rewardsByCategory[category] += rewardApr
-      } else {
-        rewardsByCategory.other += rewardApr
-      }
-    })
+      const category = categorizeRewardAsset(
+        reward.asset?.address,
+        knownAddresses,
+      )
+      rewardsByCategory[category] += rewardApr
+    }
   }
 
   // Calculate market-level rewards (weighted by allocation)
@@ -523,7 +560,7 @@ export function calculateRewardsBreakdown(apiVault: any): RewardsBreakdown {
       0,
     )
 
-    apiVault.state.allocation.forEach((allocation: any) => {
+    for (const allocation of apiVault.state.allocation) {
       if (
         allocation.market?.state?.rewards &&
         allocation.market.state.rewards.length > 0
@@ -533,23 +570,17 @@ export function calculateRewardsBreakdown(apiVault: any): RewardsBreakdown {
             ? (allocation.supplyAssetsUsd || 0) / totalSupplyUsd
             : 0
 
-        allocation.market.state.rewards.forEach((reward: any) => {
+        for (const reward of allocation.market.state.rewards) {
           const rewardApr = reward.supplyApr || 0
           const weightedRewardApr = rewardApr * weight
-          const assetSymbol = reward.asset.symbol
-
-          // Use the symbol from API response for categorization
-          const category = assetSymbol ? assetSymbol.toLowerCase() : 'other'
-
-          // Add to appropriate category if supported, otherwise to 'other'
-          if (category in rewardsByCategory) {
-            rewardsByCategory[category] += weightedRewardApr
-          } else {
-            rewardsByCategory.other += weightedRewardApr
-          }
-        })
+          const category = categorizeRewardAsset(
+            reward.asset?.address,
+            knownAddresses,
+          )
+          rewardsByCategory[category] += weightedRewardApr
+        }
       }
-    })
+    }
   }
 
   // Calculate total rewards APR

package/src/services/ChainManager.ts

@@ -12,7 +12,7 @@ import {
 import type { BundlerClient, SmartAccount } from 'viem/account-abstraction'
 import { createBundlerClient } from 'viem/account-abstraction'
 
-import type { SUPPORTED_CHAIN_IDS } from '@/constants/supportedChains.js'
+import type { SupportedChainId } from '@/constants/supportedChains.js'
 import type { ChainConfig } from '@/types/chain.js'
 
 /**
@@ -22,7 +22,7 @@ import type { ChainConfig } from '@/types/chain.js'
  */
 export class ChainManager {
   /** Map of chain IDs to their corresponding public clients */
-  private publicClients: Map<(typeof SUPPORTED_CHAIN_IDS)[number], PublicClient>
+  private publicClients: Map<SupportedChainId, PublicClient>
   /** Configuration for each supported chain */
   private chainConfigs: ChainConfig[]
 
@@ -41,7 +41,7 @@ export class ChainManager {
    * @returns PublicClient instance for the specified chain
    * @throws Error if no client is configured for the chain ID
    */
-  getPublicClient(chainId: (typeof SUPPORTED_CHAIN_IDS)[number]): PublicClient {
+  getPublicClient(chainId: SupportedChainId): PublicClient {
     const client = this.publicClients.get(chainId)
     if (!client) {
       throw new Error(`No public client configured for chain ID: ${chainId}`)
@@ -57,7 +57,7 @@ export class ChainManager {
    * @throws Error if no bundler URL is configured for the chain ID
    */
   getBundlerClient(
-    chainId: (typeof SUPPORTED_CHAIN_IDS)[number],
+    chainId: SupportedChainId,
     account: SmartAccount,
   ): BundlerClient | SmartAccountClient {
     const chainConfig = this.getChainConfig(chainId)
@@ -94,9 +94,7 @@ export class ChainManager {
    * @returns RPC URL as a string
    * @throws Error if no chain config is found for the chain ID
    */
-  getRpcUrls(
-    chainId: (typeof SUPPORTED_CHAIN_IDS)[number],
-  ): string[] | undefined {
+  getRpcUrls(chainId: SupportedChainId): string[] | undefined {
     const chainConfig = this.getChainConfig(chainId)
     return chainConfig.rpcUrls
   }
@@ -107,9 +105,7 @@ export class ChainManager {
    * @returns Bundler URL as a string or undefined if not configured
    * @throws Error if no chain config is found for the chain ID
    */
-  getBundlerUrl(
-    chainId: (typeof SUPPORTED_CHAIN_IDS)[number],
-  ): string | undefined {
+  getBundlerUrl(chainId: SupportedChainId): string | undefined {
     const chainConfig = this.getChainConfig(chainId)
     if (!chainConfig.bundler) {
       throw new Error(`No bundler configured for chain ID: ${chainId}`)
@@ -122,7 +118,7 @@ export class ChainManager {
    * @param chainId - The chain ID to retrieve information for
    * @returns Chain object containing chain details
    */
-  getChain(chainId: (typeof SUPPORTED_CHAIN_IDS)[number]): Chain {
+  getChain(chainId: SupportedChainId): Chain {
     return chainById[chainId]
   }
 
@@ -140,7 +136,7 @@ export class ChainManager {
    * @returns Transport configured with fallback RPC URLs or default http transport
    * @throws Error if no chain config is found for the chain ID
    */
-  getTransportForChain(chainId: (typeof SUPPORTED_CHAIN_IDS)[number]) {
+  getTransportForChain(chainId: SupportedChainId) {
     const rpcUrls = this.getRpcUrls(chainId)
     return rpcUrls?.length
       ? fallback(rpcUrls.map((rpcUrl) => http(rpcUrl)))
@@ -155,11 +151,8 @@ export class ChainManager {
    */
   private createPublicClients(
     chains: ChainConfig[],
-  ): Map<(typeof SUPPORTED_CHAIN_IDS)[number], PublicClient> {
-    const clients = new Map<
-      (typeof SUPPORTED_CHAIN_IDS)[number],
-      PublicClient
-    >()
+  ): Map<SupportedChainId, PublicClient> {
+    const clients = new Map<SupportedChainId, PublicClient>()
 
     for (const chainConfig of chains) {
       const chain = chainById[chainConfig.chainId]
@@ -187,9 +180,7 @@ export class ChainManager {
    * @param chainId The chain ID to retrieve the chain config for
    * @returns ChainConfig object containing chain details
    */
-  private getChainConfig(
-    chainId: (typeof SUPPORTED_CHAIN_IDS)[number],
-  ): ChainConfig {
+  private getChainConfig(chainId: SupportedChainId): ChainConfig {
     const chainConfig = this.chainConfigs.find((c) => c.chainId === chainId)
     if (!chainConfig) {
       throw new Error(`No chain config found for chain ID: ${chainId}`)
@@ -205,7 +196,7 @@ export class ChainManager {
    * @returns Pimlico bundler client for the specified chain
    */
   private getPimlicoBundlerClient(
-    chainId: (typeof SUPPORTED_CHAIN_IDS)[number],
+    chainId: SupportedChainId,
     account: SmartAccount,
     sponsorshipPolicyId?: string,
   ) {

package/src/services/__mocks__/MockChainManager.ts

@@ -4,10 +4,7 @@ import type { BundlerClient, SmartAccount } from 'viem/account-abstraction'
 import { unichain } from 'viem/chains'
 import { type MockedFunction, vi } from 'vitest'
 
-import type {
-  SUPPORTED_CHAIN_IDS,
-  SupportedChainId,
-} from '@/constants/supportedChains.js'
+import type { SupportedChainId } from '@/constants/supportedChains.js'
 
 export interface MockChainManagerConfig {
   supportedChains: SupportedChainId[]
@@ -79,15 +76,15 @@ export class MockChainManager {
     })
   }
 
-  getChain(chainId: (typeof SUPPORTED_CHAIN_IDS)[number]): Chain {
+  getChain(chainId: SupportedChainId): Chain {
     return chainById[chainId]
   }
 
-  getRpcUrls(chainId: (typeof SUPPORTED_CHAIN_IDS)[number]): string[] {
+  getRpcUrls(chainId: SupportedChainId): string[] {
     return this.getChain(chainId).rpcUrls.default.http as string[]
   }
 
-  getTransportForChain(_chainId: (typeof SUPPORTED_CHAIN_IDS)[number]) {
+  getTransportForChain(_chainId: SupportedChainId) {
     // Mock implementation returns a simple http transport
     // In tests, the actual transport behavior is typically mocked at a higher level
     return {} as unknown

package/src/services/tokenBalance.spec.ts

@@ -29,17 +29,15 @@ describe('TokenBalance', () => {
       )
 
       expect(balance).toEqual({
-        symbol: 'USDC',
-        totalBalance: 1000000n,
-        totalFormattedBalance: '1',
-        chainBalances: [
-          {
-            chainId: unichain.id,
-            balance: 1000000n,
-            formattedBalance: '1',
-            tokenAddress: MockUSDCAsset.address[unichain.id]!,
+        asset: MockUSDCAsset,
+        totalBalance: 1,
+        totalBalanceRaw: 1000000n,
+        chains: {
+          [unichain.id]: {
+            balance: 1,
+            balanceRaw: 1000000n,
           },
-        ],
+        },
       })
     })
 
@@ -63,10 +61,10 @@ describe('TokenBalance', () => {
       )
 
       expect(balance).toEqual({
-        symbol: 'UNSUPPORTED',
-        totalBalance: 0n,
-        totalFormattedBalance: '0',
-        chainBalances: [],
+        asset: unsupportedAsset,
+        totalBalance: 0,
+        totalBalanceRaw: 0n,
+        chains: {},
       })
     })
   })
@@ -76,17 +74,15 @@ describe('TokenBalance', () => {
       const balance = await fetchETHBalance(chainManager, walletAddress)
 
       expect(balance).toEqual({
-        symbol: 'ETH',
-        totalBalance: 1000000n,
-        totalFormattedBalance: '0.000000000001',
-        chainBalances: [
-          {
-            chainId: unichain.id,
-            balance: 1000000n,
-            formattedBalance: '0.000000000001',
-            tokenAddress: ETH.address[unichain.id]!,
+        asset: ETH,
+        totalBalance: 0.000000000001,
+        totalBalanceRaw: 1000000n,
+        chains: {
+          [unichain.id]: {
+            balance: 0.000000000001,
+            balanceRaw: 1000000n,
           },
-        ],
+        },
       })
     })
   })

package/src/services/tokenBalance.ts

@@ -1,6 +1,7 @@
 import type { Address } from 'viem'
 import { erc20Abi, formatEther, formatUnits } from 'viem'
 
+import { ETH } from '@/constants/assets.js'
 import type { SupportedChainId } from '@/constants/supportedChains.js'
 import type { ChainManager } from '@/services/ChainManager.js'
 import type { Asset, TokenBalance } from '@/types/asset.js'
@@ -9,7 +10,7 @@ import type { Asset, TokenBalance } from '@/types/asset.js'
  * Fetch ETH balance across all supported chains
  * @param chainManager - The chain manager
  * @param walletAddress - The wallet address
- * @returns Promise resolving to array of ETH balances
+ * @returns Promise resolving to ETH balance
  */
 export async function fetchETHBalance(
   chainManager: ChainManager,
@@ -18,26 +19,31 @@ export async function fetchETHBalance(
   const supportedChains = chainManager.getSupportedChains()
   const chainBalancePromises = supportedChains.map(async (chainId) => {
     const publicClient = chainManager.getPublicClient(chainId)
-    const balance = await publicClient.getBalance({
+    const balanceRaw = await publicClient.getBalance({
       address: walletAddress,
     })
     return {
       chainId,
-      balance,
-      tokenAddress: 'native' as const,
-      formattedBalance: formatEther(balance),
+      balanceRaw,
+      balance: parseFloat(formatEther(balanceRaw)),
     }
   })
-  const chainBalances = await Promise.all(chainBalancePromises)
-  const totalBalance = chainBalances.reduce(
-    (total, { balance }) => total + balance,
+  const chainResults = await Promise.all(chainBalancePromises)
+  const totalBalanceRaw = chainResults.reduce(
+    (total, { balanceRaw }) => total + balanceRaw,
     0n,
   )
+
+  const chains: TokenBalance['chains'] = {}
+  for (const { chainId, balance, balanceRaw } of chainResults) {
+    chains[chainId] = { balance, balanceRaw }
+  }
+
   return {
-    symbol: 'ETH',
-    totalBalance,
-    totalFormattedBalance: formatEther(totalBalance),
-    chainBalances,
+    asset: ETH,
+    totalBalance: parseFloat(formatEther(totalBalanceRaw)),
+    totalBalanceRaw,
+    chains,
   }
 }
 
@@ -55,7 +61,7 @@ export async function fetchERC20Balance(
   )
 
   const chainBalancePromises = chainsWithToken.map(async (chainId) => {
-    const { balance, tokenAddress } = await fetchERC20BalanceForChain(
+    const balanceRaw = await fetchBalanceForChain(
       asset,
       chainId,
       walletAddress,
@@ -63,35 +69,41 @@ export async function fetchERC20Balance(
     )
     return {
       chainId,
-      balance,
-      tokenAddress,
-      formattedBalance: formatUnits(balance, asset.metadata.decimals),
+      balanceRaw,
+      balance: parseFloat(formatUnits(balanceRaw, asset.metadata.decimals)),
     }
   })
 
-  const chainBalances = await Promise.all(chainBalancePromises)
-  const totalBalance = chainBalances.reduce(
-    (total, { balance }) => total + balance,
+  const chainResults = await Promise.all(chainBalancePromises)
+  const totalBalanceRaw = chainResults.reduce(
+    (total, { balanceRaw }) => total + balanceRaw,
     0n,
   )
 
+  const chains: TokenBalance['chains'] = {}
+  for (const { chainId, balance, balanceRaw } of chainResults) {
+    chains[chainId] = { balance, balanceRaw }
+  }
+
   return {
-    symbol: asset.metadata.symbol,
-    totalBalance,
-    totalFormattedBalance: formatUnits(totalBalance, asset.metadata.decimals),
-    chainBalances,
+    asset,
+    totalBalance: parseFloat(
+      formatUnits(totalBalanceRaw, asset.metadata.decimals),
+    ),
+    totalBalanceRaw,
+    chains,
   }
 }
 
 /**
  * Fetch balance for this asset on a specific chain
  */
-async function fetchERC20BalanceForChain(
+async function fetchBalanceForChain(
   asset: Asset,
   chainId: SupportedChainId,
   walletAddress: Address,
   chainManager: ChainManager,
-): Promise<{ balance: bigint; tokenAddress: Address | 'native' }> {
+): Promise<bigint> {
   const tokenAddress = asset.address[chainId]
   if (!tokenAddress) {
     throw new Error(
@@ -103,19 +115,13 @@ async function fetchERC20BalanceForChain(
 
   // Handle native ETH balance
   if (asset.type === 'native' || tokenAddress === 'native') {
-    return {
-      balance: await publicClient.getBalance({
-        address: walletAddress,
-      }),
-      tokenAddress: 'native',
-    }
+    return publicClient.getBalance({ address: walletAddress })
   }
 
-  const balance = await publicClient.readContract({
+  return publicClient.readContract({
     address: tokenAddress,
     abi: erc20Abi,
     functionName: 'balanceOf',
     args: [walletAddress],
   })
-  return { balance, tokenAddress }
 }