chore: balance sdk method refactoring (#56)

* chore: fixed the type that is returned for the getEarningBalance method of SDK

* chore: changeset

* chore: fix

Author: 0xdeval

.changeset/spotty-crabs-listen.md

@@ -0,0 +1,20 @@
+---
+'@mycelium-sdk/core': major
+---
+
+## What was changed?
+
+- Introduced AddressBalance: balance is now returned as:{ overall: OverallAddressBalance; perVault: VaultBalance[] } instead of VaultBalance[]
+- getBalances() (SparkProtocol, ProxyProtocol) now returns AddressBalance (overall totals plus per-vault list)
+- getEarnBalances() on the wallet now returns AddressBalance instead of VaultBalance[]
+- AddressBalance is exported from the SDK public types
+
+## Why was changed/added?
+
+- Mycelium Cloud started to return a different format of data for the balance
+- Callers need both an aggregated view (overall balance) and per-vault breakdown; a single type with overall and perVault supports both without extra requests
+- Aligns protocol and wallet APIs and keeps balance shape consistent across the SDK
+
+## How to use the change?
+
+Use the new return type: result.overall for aggregated numbers, result.perVault for the list of vault balances (same structure as before, but under perVault). Example: const { overall, perVault } = await wallet.getEarnBalances(); and then use perVault[0].balance, perVault[0].vaultInfo, etc. Replace any code that treated the result as a plain array (e.g. result[0]) with result.perVault[0] and use result.overall when you need totals

packages/cli/src/cli.ts

@@ -4,8 +4,9 @@ import {
   type MyceliumSDKConfig,
   type ProxyBalance,
   type SmartWallet,
-  type VaultBalance,
+  type AddressBalance,
   type Vaults,
+  type VaultBalance,
 } from '@mycelium-sdk/core';
 import { formatBalancesToDisplay, formatVaultInfoToDisplay, getEnv } from './utils/formatters';
 import { WalletDatabase } from './libs/database';
@@ -241,16 +242,16 @@ export class CLI {
       return;
     }
 
-    const earningBalances = (await this.wallet.getEarnBalances()) as VaultBalance[];
+    const earningBalances = (await this.wallet.getEarnBalances()) as AddressBalance;
 
     if (!earningBalances) {
       logError('No earning balances found. You have not deposited any funds yet');
       return;
     }
 
-    const formattedBalances = earningBalances
-      .filter((balance) => balance.balance !== null)
-      .map((balance) => {
+    const formattedBalances = earningBalances.perVault
+      .filter((balance: VaultBalance) => balance.balance !== null)
+      .map((balance: VaultBalance) => {
         const currentBalance = balance.balance as ProxyBalance;
         return {
           vaultInfo: balance.vaultInfo,
@@ -371,7 +372,7 @@ export class CLI {
       return;
     }
 
-    const balances = [...earningBalances].map((vault, index) => {
+    const balances = earningBalances.perVault.map((vault: VaultBalance, index: number) => {
       const currentBalance = vault.balance as ProxyBalance;
       return {
         optionId: index + 1,

packages/sdk/src/protocols/base/BaseProtocol.ts

@@ -15,10 +15,10 @@ import type { SupportedChainId } from '@/constants/chains';
 import type { SmartWallet } from '@/wallet/base/wallets/SmartWallet';
 import type {
   VaultInfo,
-  VaultBalance,
   VaultTxnResult,
   Vaults,
   ProtocolsSecurityConfig,
+  AddressBalance,
 } from '@/types/protocols/general';
 import type { ApiClient } from '@/tools/ApiClient';
 
@@ -108,7 +108,7 @@ export abstract class BaseProtocol {
    * @param protocolId Protocol ID to get balances for
    * @returns Balance of deposited funds
    */
-  abstract getBalances(walletAddress: Address, protocolId?: string): Promise<VaultBalance[]>;
+  abstract getBalances(walletAddress: Address, protocolId?: string): Promise<AddressBalance>;
 
   /**
    * Approve a token for protocol use

packages/sdk/src/protocols/implementations/ProxyProtocol.spec.ts

@@ -6,6 +6,7 @@ import type { ChainManager } from '@mycelium-sdk/core/tools/ChainManager';
 import type { ApiClient } from '@mycelium-sdk/core/tools/ApiClient';
 import type { SmartWallet } from '@mycelium-sdk/core/wallet/base/wallets/SmartWallet';
 import type {
+  AddressBalance,
   ProtocolsSecurityConfig,
   VaultInfo,
   VaultBalance,
@@ -63,6 +64,14 @@ describe('ProxyProtocol integration tests', () => {
     earned90dUpdatedAt: '2024-01-01T00:00:00Z',
   };
 
+  const createMockAddressBalance = (perVault: VaultBalance[]): AddressBalance => ({
+    overall: {
+      currentBalance: 1000,
+      actualCurrentBalance: 1000,
+    },
+    perVault,
+  });
+
   beforeEach(() => {
     vi.clearAllMocks();
     vi.resetAllMocks();
@@ -380,14 +389,9 @@ describe('ProxyProtocol integration tests', () => {
     });
 
     it('should withdraw specified amount from vault', async () => {
-      const mockEarningBalances: VaultBalance[] = [
-        {
-          vaultInfo: mockVaultInfo,
-          balance: mockProxyBalance,
-        },
-      ];
-
-      vi.mocked(smartWallet.getEarnBalances).mockResolvedValue(mockEarningBalances);
+      vi.mocked(smartWallet.getEarnBalances).mockResolvedValue(
+        createMockAddressBalance([{ vaultInfo: mockVaultInfo, balance: mockProxyBalance }]),
+      );
 
       const mockOperationData: TransactionData = {
         to: mockVaultInfo.vaultAddress,
@@ -419,14 +423,9 @@ describe('ProxyProtocol integration tests', () => {
     });
 
     it('should withdraw all balance when amount is not specified', async () => {
-      const mockEarningBalances: VaultBalance[] = [
-        {
-          vaultInfo: mockVaultInfo,
-          balance: mockProxyBalance,
-        },
-      ];
-
-      vi.mocked(smartWallet.getEarnBalances).mockResolvedValue(mockEarningBalances);
+      vi.mocked(smartWallet.getEarnBalances).mockResolvedValue(
+        createMockAddressBalance([{ vaultInfo: mockVaultInfo, balance: mockProxyBalance }]),
+      );
 
       const mockOperationData: TransactionData = {
         to: mockVaultInfo.vaultAddress,
@@ -464,29 +463,21 @@ describe('ProxyProtocol integration tests', () => {
     });
 
     it('should throw error when vault balance not found in earning balances', async () => {
-      const mockEarningBalances: VaultBalance[] = [
-        {
-          vaultInfo: { ...mockVaultInfo, id: 'different-vault' },
-          balance: mockProxyBalance,
-        },
-      ];
-
-      vi.mocked(smartWallet.getEarnBalances).mockResolvedValue(mockEarningBalances);
+      vi.mocked(smartWallet.getEarnBalances).mockResolvedValue(
+        createMockAddressBalance([
+          { vaultInfo: { ...mockVaultInfo, id: 'different-vault' }, balance: mockProxyBalance },
+        ]),
+      );
 
       await expect(proxyProtocol.withdraw(mockVaultInfo, smartWallet, '500')).rejects.toThrow(
         'No earning balance found',
       );
     });
 
     it('should throw error when API fails to return withdraw operations', async () => {
-      const mockEarningBalances: VaultBalance[] = [
-        {
-          vaultInfo: mockVaultInfo,
-          balance: mockProxyBalance,
-        },
-      ];
-
-      vi.mocked(smartWallet.getEarnBalances).mockResolvedValue(mockEarningBalances);
+      vi.mocked(smartWallet.getEarnBalances).mockResolvedValue(
+        createMockAddressBalance([{ vaultInfo: mockVaultInfo, balance: mockProxyBalance }]),
+      );
 
       (apiClient.sendRequest as ReturnType<typeof vi.fn>).mockResolvedValue({
         success: false,
@@ -499,14 +490,9 @@ describe('ProxyProtocol integration tests', () => {
     });
 
     it('should log operation after successful withdrawal', async () => {
-      const mockEarningBalances: VaultBalance[] = [
-        {
-          vaultInfo: mockVaultInfo,
-          balance: mockProxyBalance,
-        },
-      ];
-
-      vi.mocked(smartWallet.getEarnBalances).mockResolvedValue(mockEarningBalances);
+      vi.mocked(smartWallet.getEarnBalances).mockResolvedValue(
+        createMockAddressBalance([{ vaultInfo: mockVaultInfo, balance: mockProxyBalance }]),
+      );
 
       const mockOperationData: TransactionData = {
         to: mockVaultInfo.vaultAddress,
@@ -541,14 +527,10 @@ describe('ProxyProtocol integration tests', () => {
     it('should withdraw with paymaster token when paymaster token equals withdraw token', async () => {
       const paymasterToken = mockVaultInfo.tokenAddress;
       const mockPublicClient = chainManager.getPublicClient(8453);
-      const mockEarningBalances: VaultBalance[] = [
-        {
-          vaultInfo: mockVaultInfo,
-          balance: mockProxyBalance,
-        },
-      ];
 
-      vi.mocked(smartWallet.getEarnBalances).mockResolvedValue(mockEarningBalances);
+      vi.mocked(smartWallet.getEarnBalances).mockResolvedValue(
+        createMockAddressBalance([{ vaultInfo: mockVaultInfo, balance: mockProxyBalance }]),
+      );
 
       // Mock balance check for gas reserve
       vi.mocked(mockPublicClient.readContract as ReturnType<typeof vi.fn>).mockResolvedValue(
@@ -591,14 +573,10 @@ describe('ProxyProtocol integration tests', () => {
     it('should throw error when wallet balance is insufficient for gas payment', async () => {
       const paymasterToken = mockVaultInfo.tokenAddress;
       const mockPublicClient = chainManager.getPublicClient(8453);
-      const mockEarningBalances: VaultBalance[] = [
-        {
-          vaultInfo: mockVaultInfo,
-          balance: mockProxyBalance,
-        },
-      ];
 
-      vi.mocked(smartWallet.getEarnBalances).mockResolvedValue(mockEarningBalances);
+      vi.mocked(smartWallet.getEarnBalances).mockResolvedValue(
+        createMockAddressBalance([{ vaultInfo: mockVaultInfo, balance: mockProxyBalance }]),
+      );
 
       // Mock balance that's too low for gas reserve
       vi.mocked(mockPublicClient.readContract as ReturnType<typeof vi.fn>).mockResolvedValue(
@@ -618,16 +596,13 @@ describe('ProxyProtocol integration tests', () => {
 
     it('should fetch and return earning balances of a user by a provided address', async () => {
       const walletAddress = '0x1234567890123456789012345678901234567890' as Address;
-      const mockBalances: VaultBalance[] = [
-        {
-          vaultInfo: mockVaultInfo,
-          balance: mockProxyBalance,
-        },
-      ];
+      const mockAddressBalance = createMockAddressBalance([
+        { vaultInfo: mockVaultInfo, balance: mockProxyBalance },
+      ]);
 
       (apiClient.sendRequest as ReturnType<typeof vi.fn>).mockResolvedValue({
         success: true,
-        data: mockBalances,
+        data: mockAddressBalance,
       });
 
       const result = await proxyProtocol.getBalances(walletAddress);
@@ -638,21 +613,18 @@ describe('ProxyProtocol integration tests', () => {
         userAddress: walletAddress,
       });
 
-      expect(result).toEqual(mockBalances);
+      expect(result).toEqual(mockAddressBalance);
     });
 
     it('should fetch balances for specific protocol ID', async () => {
       const walletAddress = '0x1234567890123456789012345678901234567890' as Address;
-      const mockBalances: VaultBalance[] = [
-        {
-          vaultInfo: mockVaultInfo,
-          balance: mockProxyBalance,
-        },
-      ];
+      const mockAddressBalance = createMockAddressBalance([
+        { vaultInfo: mockVaultInfo, balance: mockProxyBalance },
+      ]);
 
       (apiClient.sendRequest as ReturnType<typeof vi.fn>).mockResolvedValue({
         success: true,
-        data: mockBalances,
+        data: mockAddressBalance,
       });
 
       await proxyProtocol.getBalances(walletAddress, 'spark');

packages/sdk/src/protocols/implementations/ProxyProtocol.ts

@@ -1,8 +1,8 @@
 import { BaseProtocol } from '@/protocols/base/BaseProtocol';
 import type { ChainManager } from '@/tools/ChainManager';
 import type {
+  AddressBalance,
   ProtocolsSecurityConfig,
-  VaultBalance,
   VaultInfo,
   Vaults,
   VaultTxnResult,
@@ -259,7 +259,9 @@ export class ProxyProtocol extends BaseProtocol {
       throw new Error('No earning balances found');
     }
 
-    const earningBalance = earningBalances.find((balance) => balance.vaultInfo.id === vaultInfo.id);
+    const earningBalance = earningBalances.perVault.find(
+      (balance) => balance.vaultInfo.id === vaultInfo.id,
+    );
 
     if (!earningBalance) {
       throw new Error('No earning balance found');
@@ -320,7 +322,7 @@ export class ProxyProtocol extends BaseProtocol {
    * @param protocolId Protocol ID to get the balances for. Optional, default is undefined
    * @returns Balances of the user in the protocol vaults
    */
-  async getBalances(walletAddress: Address, protocolId?: string): Promise<VaultBalance[]> {
+  async getBalances(walletAddress: Address, protocolId?: string): Promise<AddressBalance> {
     const pathParams = {
       chain_id: this.getSelectedChainId().toString(),
       protocol_id: protocolId || '',
@@ -333,7 +335,7 @@ export class ProxyProtocol extends BaseProtocol {
       throw new Error(apiResponse.error || 'Failed to get balances');
     }
 
-    const balances: VaultBalance[] = apiResponse.data as unknown as VaultBalance[];
+    const balances: AddressBalance = apiResponse.data as unknown as AddressBalance;
 
     return balances;
   }

packages/sdk/src/protocols/implementations/SparkProtocol.spec.ts

@@ -488,14 +488,15 @@ describe('SparkProtocol integration tests', () => {
 
       expect(formatUnits).toHaveBeenCalledWith(mockAssets, mockVaultInfo.tokenDecimals);
 
-      expect(result).toHaveLength(1);
-      expect(result[0]?.vaultInfo.id).toBe('sUSDC');
-      expect(result[0]?.balance).toBe('1050.0');
-      expect(result[0]?.vaultInfo.metadata?.apy).toBeDefined();
-      expect(typeof result[0]?.vaultInfo.metadata?.apy).toBe('number');
+      expect(result.overall).toEqual({ currentBalance: 1050, actualCurrentBalance: 1050 });
+      expect(result.perVault).toHaveLength(1);
+      expect(result.perVault[0]?.vaultInfo.id).toBe('sUSDC');
+      expect(result.perVault[0]?.balance).toBe('1050.0');
+      expect(result.perVault[0]?.vaultInfo.metadata?.apy).toBeDefined();
+      expect(typeof result.perVault[0]?.vaultInfo.metadata?.apy).toBe('number');
     });
 
-    it('should return null balance when wallet has no shares', async () => {
+    it('should return zero overall and perVault balance when wallet has no shares', async () => {
       const walletAddress = '0x1234567890123456789012345678901234567890' as Address;
       const mockPublicClient = chainManager.getPublicClient(8453);
       const mockSSR = BigInt('1050000000000000000000000000');
@@ -512,11 +513,12 @@ describe('SparkProtocol integration tests', () => {
         functionName: 'getSSR',
       });
 
-      expect(result).toHaveLength(1);
-      expect(result[0]?.balance).toBeNull();
-      expect(result[0]?.vaultInfo.id).toBe('sUSDC');
-      expect(result[0]?.vaultInfo.metadata?.apy).toBeDefined();
-      expect(typeof result[0]?.vaultInfo.metadata?.apy).toBe('number');
+      expect(result.overall).toEqual({ currentBalance: 0, actualCurrentBalance: 0 });
+      expect(result.perVault).toHaveLength(1);
+      expect(result.perVault[0]?.balance).toBe('0');
+      expect(result.perVault[0]?.vaultInfo.id).toBe('sUSDC');
+      expect(result.perVault[0]?.vaultInfo.metadata?.apy).toBeDefined();
+      expect(typeof result.perVault[0]?.vaultInfo.metadata?.apy).toBe('number');
     });
 
     it('should throw error when public client is not initialized', async () => {