fix(artillery, init): refactor ledger balance management and add PKP balance checks

Author: Ansonhkg

e2e/artillery/src/init.ts

@@ -6,63 +6,54 @@ import {
 import { createLitClient } from '@lit-protocol/lit-client';
 import * as NetworkManager from '../../../e2e/src/helper/NetworkManager';
 import { getOrCreatePkp } from '../../../e2e/src/helper/pkp-utils';
+import { printAligned } from '../../../e2e/src/helper/utils';
 import '../../src/helper/supressLogs';
 import * as AccountManager from '../src/AccountManager';
 import * as StateManager from './StateManager';
 
 const _network = process.env['NETWORK'];
 
-const AUTO_TOP_UP_FLAG = '--auto-topup';
-const args = process.argv.slice(2);
-
 // CONFIGURATIONS
 const REJECT_BALANCE_THRESHOLD = 0;
 const LEDGER_MINIMUM_BALANCE = 10000;
-const AUTO_TOP_UP_ENABLED = args.includes(AUTO_TOP_UP_FLAG);
-const AUTO_TOP_UP_INTERVAL = 10_000;
-const AUTO_TOP_UP_THRESHOLD = LEDGER_MINIMUM_BALANCE;
 
 if (Number.isNaN(LEDGER_MINIMUM_BALANCE) || LEDGER_MINIMUM_BALANCE < 0) {
   throw new Error('❌ LEDGER_MINIMUM_BALANCE must be a non-negative number');
 }
 
-const ensureLedgerThreshold = async ({
-  paymentManager,
-  accountAddress,
+const ensureLedgerBalance = async ({
+  label,
+  balanceFetcher,
   minimumBalance,
+  topUp,
 }: {
-  paymentManager: Awaited<
-    ReturnType<typeof AccountManager.getAccountDetails>
-  >['paymentManager'];
-  accountAddress: `0x${string}`;
+  label: string;
+  balanceFetcher: () => Promise<{ availableBalance: string }>;
   minimumBalance: number;
+  topUp: (difference: number) => Promise<void>;
 }) => {
-  const { availableBalance } = await paymentManager.getBalance({
-    userAddress: accountAddress,
-  });
+  const { availableBalance } = await balanceFetcher();
 
   const currentAvailable = Number(availableBalance);
 
   if (currentAvailable >= minimumBalance) {
+    console.log(
+      `✅ ${label} ledger balance healthy (${currentAvailable} ETH, threshold ${minimumBalance} ETH)`
+    );
     return currentAvailable;
   }
 
-  const diff = minimumBalance - currentAvailable;
+  const difference = minimumBalance - currentAvailable;
 
   console.log(
-    `🚨 Live Master Account Ledger Balance (${currentAvailable}) is below threshold (${minimumBalance}). Depositing ${difference} ETH.`
+    `🚨 ${label} ledger balance (${currentAvailable} ETH) is below threshold (${minimumBalance} ETH). Depositing ${difference} ETH.`
   );
 
-  await paymentManager.deposit({
-    amountInEth: diff.toString(),
-  });
+  await topUp(difference);
 
-  const { availableBalance: postTopUpBalance } =
-    await paymentManager.getBalance({
-      userAddress: accountAddress,
-    });
+  const { availableBalance: postTopUpBalance } = await balanceFetcher();
 
-  console.log('✅ New Live Master Account Payment Balance:', postTopUpBalance);
+  console.log(`✅ ${label} ledger balance after top-up: ${postTopUpBalance} ETH`);
 
   return Number(postTopUpBalance);
 };
@@ -94,10 +85,18 @@ const ensureLedgerThreshold = async ({
     );
   }
 
-  await ensureLedgerThreshold({
-    paymentManager: masterAccountDetails.paymentManager,
-    accountAddress: masterAccount.address,
+  await ensureLedgerBalance({
+    label: 'Master Account',
+    balanceFetcher: () =>
+      masterAccountDetails.paymentManager.getBalance({
+        userAddress: masterAccount.address,
+      }),
     minimumBalance: LEDGER_MINIMUM_BALANCE,
+    topUp: async (difference) => {
+      await masterAccountDetails.paymentManager.deposit({
+        amountInEth: difference.toString(),
+      });
+    },
   });
 
   // 3. Authenticate the master account and store the auth data
@@ -130,6 +129,58 @@ const ensureLedgerThreshold = async ({
 
   console.log('✅ Master Account PKP:', masterAccountPkp);
 
+  const pkpEthAddress = masterAccountPkp?.ethAddress;
+
+  if (!pkpEthAddress) {
+    throw new Error('❌ Master Account PKP is missing an ethAddress');
+  }
+
+  const pkpLedgerBalance = await masterAccountDetails.paymentManager.getBalance(
+    {
+      userAddress: pkpEthAddress,
+    }
+  );
+
+  console.log('\n========== Master Account PKP Details ==========');
+
+  const pkpStatus =
+    Number(pkpLedgerBalance.availableBalance) < 0
+      ? {
+          label: '🚨 Status:',
+          value: `Negative balance (debt): ${pkpLedgerBalance.availableBalance}`,
+        }
+      : { label: '', value: '' };
+
+  printAligned(
+    [
+      { label: '🔑 PKP ETH Address:', value: pkpEthAddress },
+      {
+        label: '💳 Ledger Total Balance:',
+        value: pkpLedgerBalance.totalBalance,
+      },
+      {
+        label: '💳 Ledger Available Balance:',
+        value: pkpLedgerBalance.availableBalance,
+      },
+      pkpStatus,
+    ].filter((item) => item.label)
+  );
+
+  await ensureLedgerBalance({
+    label: 'Master Account PKP',
+    balanceFetcher: () =>
+      masterAccountDetails.paymentManager.getBalance({
+        userAddress: pkpEthAddress,
+      }),
+    minimumBalance: LEDGER_MINIMUM_BALANCE,
+    topUp: async (difference) => {
+      await masterAccountDetails.paymentManager.depositForUser({
+        userAddress: pkpEthAddress,
+        amountInEth: difference.toString(),
+      });
+    },
+  });
+
   // create pkp auth context
   // const masterAccountPkpAuthContext = await authManager.createPkpAuthContext({
   //   authData: masterAccountAuthData,
@@ -175,38 +226,5 @@ const ensureLedgerThreshold = async ({
 
   // console.log('✅ PKP Sign Test Result:', res);
 
-  if (AUTO_TOP_UP_ENABLED) {
-    console.log(
-      `\n✅ Auto top-up enabled. Monitoring every ${AUTO_TOP_UP_INTERVAL}ms with threshold ${AUTO_TOP_UP_THRESHOLD} ETH. Press Ctrl+C to exit.`
-    );
-
-    let isTopUpInProgress = false;
-
-    const poll = async () => {
-      if (isTopUpInProgress) {
-        return;
-      }
-
-      isTopUpInProgress = true;
-
-      try {
-        await ensureLedgerThreshold({
-          paymentManager: masterAccountDetails.paymentManager,
-          accountAddress: masterAccount.address,
-          minimumBalance: AUTO_TOP_UP_THRESHOLD,
-        });
-      } catch (error) {
-        console.error('❌ Auto top-up check failed:', error);
-      } finally {
-        isTopUpInProgress = false;
-      }
-    };
-
-    await poll();
-    setInterval(() => {
-      void poll();
-    }, AUTO_TOP_UP_INTERVAL);
-  } else {
-    process.exit();
-  }
+  process.exit();
 })();