chore: refactoring ProxyProtocol and SparkProtocol; updated tests

Author: 0xdeval

.changeset/late-clowns-add.md

@@ -4,17 +4,17 @@
 
 ## What was changed?
 
-1. Added paymaster to ProxyProtocol to use with DeFi opportunities with Mycelium Cloud
-2. Updated tests for ProxyProtocol
-3. Added Paymaster support for DefaultSmartWallet
-4. Updated tests for DefaultSmartWallet.ts
-5. Added paymaster to SparkProtocol to use with Spark vaults
-6. Updated SparkProtocol tests
+- Added paymaster to ProxyProtocol to use with DeFi opportunities with Mycelium Cloud
+- Updated tests for ProxyProtocol
+- Added Paymaster support for DefaultSmartWallet
+- Updated tests for DefaultSmartWallet.ts
+- Added paymaster to SparkProtocol to use with Spark vaults
+- Updated SparkProtocol tests
 
 ## Why was changed/added?
 
-1. Support usage of paymaster with DeFi opportunities from Mycelium Cloud
-2. Supported usage of paymaster with DeFi opportunities with core SDK without Mycelium Cloud
+- Support usage of paymaster with DeFi opportunities from Mycelium Cloud
+- Supported usage of paymaster with DeFi opportunities with core SDK without Mycelium Cloud
 
 ## How to use the change?
 

packages/sdk/src/constants/paymaster.ts

@@ -3,7 +3,7 @@
 export const ERC20_PAYMASTER_ADDRESS = '0x6666666666667849c56f2850848ce1c4da65c68b';
 
 // Percentage of balance to reserve for gas payment for tokens with high decimals (>8)
-export const GAS_RESERVE_PERCENTAGE = 1;
+export const GAS_RESERVE_PERCENTAGE = '1';
 
 // Minimum gas reserve amount in token units for tokens with low decimals (<8)
-export const GAS_RESERVE_MINIMUM = 0.01;
+export const GAS_RESERVE_MINIMUM = '0.01';

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

@@ -6,8 +6,11 @@ import {
   createWalletClient,
   http,
   parseGwei,
+  parseUnits,
+  formatUnits,
   type PublicClient,
 } from 'viem';
+import { GAS_RESERVE_MINIMUM, GAS_RESERVE_PERCENTAGE } from '@/constants/paymaster';
 import type { SupportedChainId } from '@/constants/chains';
 import type { SmartWallet } from '@/wallet/base/wallets/SmartWallet';
 import type {
@@ -94,7 +97,7 @@ export abstract class BaseProtocol {
    */
   abstract withdraw(
     vaultInfo: VaultInfo,
-    amountInShares: string,
+    amount: string,
     smartWallet: SmartWallet,
     options?: { paymasterToken?: Address },
   ): Promise<VaultTxnResult>;
@@ -169,4 +172,94 @@ export abstract class BaseProtocol {
       args: [walletAddress, spenderAddress],
     });
   }
+
+  /**
+   * Validate gas reserve balance for operations using paymaster
+   * Ensures sufficient balance remains for gas payment when using paymaster with the same token
+   * @param tokenAddress Token address to check balance for
+   * @param walletAddress Wallet address to check
+   * @param tokenDecimals Number of decimals for the token
+   * @param operationAmount Optional: Amount for deposit operation (in token units). If provided, validates deposit; otherwise validates withdraw
+   * @throws Error if balance is insufficient for gas payment or operation
+   */
+  protected async validateGasReserve(
+    tokenAddress: Address,
+    walletAddress: Address,
+    tokenDecimals: number,
+    operationAmount?: bigint,
+  ): Promise<void> {
+    this.ensureInitialized();
+    const publicClient = this.chainManager!.getPublicClient(this.selectedChainId!);
+    const balance = await publicClient.readContract({
+      address: tokenAddress,
+      abi: erc20Abi,
+      functionName: 'balanceOf',
+      args: [walletAddress],
+    });
+
+    const gasReserve = this.calculateGasReserve(balance, tokenDecimals);
+
+    if (operationAmount !== undefined) {
+      // Check if operation amount exceeds available balance after gas reserve
+      const maxDepositAmount = balance > gasReserve ? balance - gasReserve : 0n;
+
+      if (operationAmount > maxDepositAmount) {
+        const maxDepositFormatted = formatUnits(maxDepositAmount, tokenDecimals);
+        throw new Error(
+          `Insufficient balance. Must reserve tokens for gas payment. Max deposit: ${maxDepositFormatted}`,
+        );
+      }
+    } else {
+      // Check if balance meets minimum gas reserve requirement
+      const minRequiredBalance = gasReserve;
+
+      if (balance < minRequiredBalance) {
+        const minRequiredFormatted = formatUnits(minRequiredBalance, tokenDecimals);
+        throw new Error(
+          `Insufficient wallet balance for gas payment. Wallet needs at least ${minRequiredFormatted} tokens to pay for gas before withdrawal.`,
+        );
+      }
+    }
+  }
+
+  /**
+   * Calculate gas reserve amount based on balance and token decimals
+   * Uses a more sophisticated calculation that considers token decimal places:
+   * - For tokens with low decimals (≤6): Uses a fixed minimum amount configured via
+   *   GAS_RESERVE_MINIMUM (e.g., 0.01 tokens), with at least 1 unit reserved
+   * - For tokens with higher decimals (>6): Uses GAS_RESERVE_PERCENTAGE% of balance
+   *   with a minimum of 1 unit
+   * @param balance Current token balance
+   * @param tokenDecimals Number of decimals for the token
+   * @returns Gas reserve amount in token units
+   */
+  protected calculateGasReserve(balance: bigint, tokenDecimals: number): bigint {
+    // For tokens with low decimals (e.g., WBTC with 8 decimals), use a fixed minimum
+    // This ensures sufficient gas coverage for high-value tokens
+    if (tokenDecimals <= 6) {
+      // Reserve 0.001 tokens (or 1 unit if that's larger)
+      const fixedReserve = parseUnits(GAS_RESERVE_MINIMUM, tokenDecimals);
+      const oneUnit = 1n;
+      return fixedReserve > oneUnit ? fixedReserve : oneUnit;
+    }
+
+    // For tokens with high decimals, use percentage-based approach
+    // Reserve 1% of balance with a minimum of 1 unit
+    const percentageReserve = (balance * BigInt(GAS_RESERVE_PERCENTAGE)) / 100n;
+    const oneUnit = 1n;
+    return percentageReserve > 0n ? percentageReserve : oneUnit;
+  }
+
+  /**
+   * Get the selected chain ID, ensuring it has been initialized
+   * @returns The selected chain ID
+   * @throws Error if `init()` has not been called or chain ID is not set
+   */
+  protected getSelectedChainId(): SupportedChainId {
+    this.ensureInitialized();
+    if (this.selectedChainId === undefined) {
+      throw new Error('Protocol chain ID not set. Ensure init() was called successfully.');
+    }
+    return this.selectedChainId;
+  }
 }

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

@@ -10,11 +10,8 @@ import type {
   VaultInfo,
   VaultBalance,
 } from '@mycelium-sdk/core/types/protocols/general';
-import type {
-  ProxyVaults,
-  ProxyBalance,
-  OperationCallDataType,
-} from '@mycelium-sdk/core/types/protocols/proxy';
+import type { ProxyVaults, ProxyBalance } from '@mycelium-sdk/core/types/protocols/proxy';
+import type { TransactionData } from '@mycelium-sdk/core/types/transaction';
 import { encodeFunctionData, erc20Abi, parseUnits, type Address, type Hash } from 'viem';
 
 // Mock viem functions
@@ -178,7 +175,7 @@ describe('ProxyProtocol integration tests', () => {
         BigInt('2000000000'), // Allowance is greater than deposit amount
       );
 
-      const mockOperationData: OperationCallDataType = {
+      const mockOperationData: TransactionData = {
         to: mockVaultInfo.vaultAddress,
         data: '0x1234' as `0x${string}`,
       };
@@ -220,7 +217,7 @@ describe('ProxyProtocol integration tests', () => {
       const mockApproveData = '0xhash123' as `0x${string}`;
       vi.mocked(encodeFunctionData).mockReturnValue(mockApproveData);
 
-      const mockOperationData: OperationCallDataType = {
+      const mockOperationData: TransactionData = {
         to: mockVaultInfo.vaultAddress,
         data: '0x1234' as `0x${string}`,
       };
@@ -289,7 +286,7 @@ describe('ProxyProtocol integration tests', () => {
         BigInt('2000000000'),
       );
 
-      const mockOperationData: OperationCallDataType = {
+      const mockOperationData: TransactionData = {
         to: mockVaultInfo.vaultAddress,
         data: '0x1234' as `0x${string}`,
       };
@@ -318,6 +315,63 @@ describe('ProxyProtocol integration tests', () => {
         }),
       );
     });
+
+    it('should deposit with paymaster token when paymaster token equals deposit token', async () => {
+      const paymasterToken = mockVaultInfo.tokenAddress;
+      const mockPublicClient = chainManager.getPublicClient(8453);
+
+      // Mock balance check (for gas reserve) and allowance check
+      vi.mocked(mockPublicClient.readContract as ReturnType<typeof vi.fn>)
+        .mockResolvedValueOnce(BigInt('2000000000'))
+        .mockResolvedValueOnce(BigInt('2000000000'));
+
+      const mockOperationData: TransactionData = {
+        to: mockVaultInfo.vaultAddress,
+        data: '0x1234' as `0x${string}`,
+      };
+
+      (apiClient.sendRequest as ReturnType<typeof vi.fn>)
+        .mockResolvedValueOnce({
+          success: true,
+          data: mockOperationData,
+        })
+        .mockResolvedValueOnce({
+          success: true,
+        });
+
+      (smartWallet.sendBatch as ReturnType<typeof vi.fn>).mockResolvedValue('0xhash123' as Hash);
+
+      const result = await proxyProtocol.deposit(mockVaultInfo, '1000', smartWallet, {
+        paymasterToken,
+      });
+
+      expect(mockPublicClient.readContract).toHaveBeenCalledWith({
+        address: mockVaultInfo.tokenAddress,
+        abi: erc20Abi,
+        functionName: 'balanceOf',
+        args: [expect.any(String)],
+      });
+
+      expect(smartWallet.sendBatch).toHaveBeenCalledWith([mockOperationData], 8453, {
+        paymasterToken,
+      });
+      expect(result.success).toBe(true);
+      expect(result.hash).toBe('0xhash123');
+    });
+
+    it('should throw error when deposit amount exceeds balance minus gas reserve', async () => {
+      const paymasterToken = mockVaultInfo.tokenAddress;
+      const mockPublicClient = chainManager.getPublicClient(8453);
+
+      // Mock balance that's too low for gas reserve
+      vi.mocked(mockPublicClient.readContract as ReturnType<typeof vi.fn>).mockResolvedValue(
+        BigInt('500000000'),
+      );
+
+      await expect(
+        proxyProtocol.deposit(mockVaultInfo, '1000', smartWallet, { paymasterToken }),
+      ).rejects.toThrow('Insufficient balance. Must reserve tokens for gas payment.');
+    });
   });
 
   describe('withdraw', () => {
@@ -335,7 +389,7 @@ describe('ProxyProtocol integration tests', () => {
 
       vi.mocked(smartWallet.getEarnBalances).mockResolvedValue(mockEarningBalances);
 
-      const mockOperationData: OperationCallDataType = {
+      const mockOperationData: TransactionData = {
         to: mockVaultInfo.vaultAddress,
         data: '0xhash123' as `0x${string}`,
       };
@@ -374,7 +428,7 @@ describe('ProxyProtocol integration tests', () => {
 
       vi.mocked(smartWallet.getEarnBalances).mockResolvedValue(mockEarningBalances);
 
-      const mockOperationData: OperationCallDataType = {
+      const mockOperationData: TransactionData = {
         to: mockVaultInfo.vaultAddress,
         data: '0xhash123' as `0x${string}`,
       };
@@ -454,7 +508,7 @@ describe('ProxyProtocol integration tests', () => {
 
       vi.mocked(smartWallet.getEarnBalances).mockResolvedValue(mockEarningBalances);
 
-      const mockOperationData: OperationCallDataType = {
+      const mockOperationData: TransactionData = {
         to: mockVaultInfo.vaultAddress,
         data: '0xhash123' as `0x${string}`,
       };
@@ -483,6 +537,78 @@ 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);
+
+      // Mock balance check for gas reserve
+      vi.mocked(mockPublicClient.readContract as ReturnType<typeof vi.fn>).mockResolvedValue(
+        BigInt('2000000000'),
+      );
+
+      const mockOperationData: TransactionData = {
+        to: mockVaultInfo.vaultAddress,
+        data: '0xhash123' as `0x${string}`,
+      };
+
+      (apiClient.sendRequest as ReturnType<typeof vi.fn>)
+        .mockResolvedValueOnce({
+          success: true,
+          data: mockOperationData,
+        })
+        .mockResolvedValueOnce({
+          success: true,
+        });
+
+      (smartWallet.send as ReturnType<typeof vi.fn>).mockResolvedValue('0xhash456' as Hash);
+
+      const result = await proxyProtocol.withdraw(mockVaultInfo, '500', smartWallet, {
+        paymasterToken,
+      });
+
+      // Verify balance check was performed
+      expect(mockPublicClient.readContract).toHaveBeenCalledWith({
+        address: mockVaultInfo.tokenAddress,
+        abi: erc20Abi,
+        functionName: 'balanceOf',
+        args: [expect.any(String)],
+      });
+
+      expect(smartWallet.send).toHaveBeenCalledWith(mockOperationData, 8453, { paymasterToken });
+      expect(result.success).toBe(true);
+      expect(result.hash).toBe('0xhash456');
+    });
+
+    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);
+
+      // Mock balance that's too low for gas reserve
+      vi.mocked(mockPublicClient.readContract as ReturnType<typeof vi.fn>).mockResolvedValue(
+        BigInt('0'),
+      );
+
+      await expect(
+        proxyProtocol.withdraw(mockVaultInfo, '500', smartWallet, { paymasterToken }),
+      ).rejects.toThrow('Insufficient wallet balance for gas payment');
+    });
   });
 
   describe('getBalances', () => {