Anson/naga artillery fix init and product 3

e2e/artillery/src/init.ts

@@ -9,12 +9,57 @@ import { createLitClient } from '@lit-protocol/lit-client';
 import { getOrCreatePkp } from '../../../e2e/src/helper/pkp-utils';
 import * as NetworkManager from '../../../e2e/src/helper/NetworkManager';
 import * as AccountManager from '../src/AccountManager';
+import { printAligned } from '../../../e2e/src/helper/utils';
 
 const _network = process.env['NETWORK'];
 
 // CONFIGURATIONS
 const REJECT_BALANCE_THRESHOLD = 0;
-const LEDGER_MINIMUM_BALANCE = 10000;
+const MASTER_LEDGER_MINIMUM_BALANCE = 3_000;
+const PKP_LEDGER_MINIMUM_BALANCE = 3_000;
+
+if (MASTER_LEDGER_MINIMUM_BALANCE < 0 || PKP_LEDGER_MINIMUM_BALANCE < 0) {
+  throw new Error('❌ Ledger minimum balances must be non-negative numbers');
+}
+
+const ensureLedgerBalance = async ({
+  label,
+  balanceFetcher,
+  minimumBalance,
+  topUp,
+}: {
+  label: string;
+  balanceFetcher: () => Promise<{ availableBalance: string }>;
+  minimumBalance: number;
+  topUp: (difference: number) => Promise<void>;
+}) => {
+  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 difference = minimumBalance - currentAvailable;
+
+  console.log(
+    `🚨 ${label} ledger balance (${currentAvailable} ETH) is below threshold (${minimumBalance} ETH). Depositing ${difference} ETH.`
+  );
+
+  await topUp(difference);
+
+  const { availableBalance: postTopUpBalance } = await balanceFetcher();
+
+  console.log(
+    `✅ ${label} ledger balance after top-up: ${postTopUpBalance} ETH`
+  );
+
+  return Number(postTopUpBalance);
+};
 
 (async () => {
   // -- Start
@@ -43,32 +88,19 @@ const LEDGER_MINIMUM_BALANCE = 10000;
     );
   }
 
-  if (LEDGER_MINIMUM_BALANCE > Number(masterAccountDetails.ledgerBalance)) {
-    // find the difference between the minimum balance and the current balance
-    const difference =
-      LEDGER_MINIMUM_BALANCE - Number(masterAccountDetails.ledgerBalance);
-
-    console.log(
-      `🚨 Live Master Account Ledger Balance is less than LEDGER_MINIMUM_BALANCE: ${LEDGER_MINIMUM_BALANCE} ETH. Attempting to top up the difference of ${difference} ETH to the master account.`
-    );
-
-    // deposit the difference
-    console.log(
-      '\x1b[90m✅ Depositing the difference to Live Master Account Payment Manager...\x1b[0m'
-    );
-    await masterAccountDetails.paymentManager.deposit({
-      amountInEth: difference.toString(),
-    });
-
-    // print the new balance
-    const newBalance = await masterAccountDetails.paymentManager.getBalance({
-      userAddress: masterAccount.address,
-    });
-    console.log(
-      '✅ New Live Master Account Payment Balance:',
-      newBalance.availableBalance
-    );
-  }
+  await ensureLedgerBalance({
+    label: 'Master Account',
+    balanceFetcher: () =>
+      masterAccountDetails.paymentManager.getBalance({
+        userAddress: masterAccount.address,
+      }),
+    minimumBalance: MASTER_LEDGER_MINIMUM_BALANCE,
+    topUp: async (difference) => {
+      await masterAccountDetails.paymentManager.deposit({
+        amountInEth: difference.toString(),
+      });
+    },
+  });
 
   // 3. Authenticate the master account and store the auth data
   const masterAccountAuthData = await StateManager.getOrUpdate(
@@ -100,6 +132,58 @@ const LEDGER_MINIMUM_BALANCE = 10000;
 
   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: PKP_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,

packages/auth/src/lib/AuthManager/authAdapters/getPkpAuthContextAdapter.ts

@@ -136,9 +136,9 @@ export async function getPkpAuthContextAdapter(
   const nodeUrls = litClientCtx.getMaxPricesForNodeProduct({
     nodePrices: nodePrices,
     userMaxPrice: litClientCtx.getUserMaxPrice({
-      product: 'LIT_ACTION',
+      product: 'SIGN_SESSION_KEY',
     }),
-    productId: PRODUCT_IDS['LIT_ACTION'],
+    productId: PRODUCT_IDS['SIGN_SESSION_KEY'],
     numRequiredNodes: threshold,
   });
 

packages/networks/src/networks/vNaga/shared/managers/LitChainClient/apis/rawContractApis/pricing/getNodesForRequest.ts

@@ -15,6 +15,7 @@ import {
  * - DECRYPTION (0): Used for decryption operations
  * - SIGN (1): Used for signing operations
  * - LA (2): Used for Lit Actions execution
+ * - SIGN_SESSION_KEY (3): Used for sign session key operations
  */
 export const PRODUCT_IDS = {
   DECRYPTION: 0n, // For decryption operations

packages/networks/src/networks/vNaga/shared/managers/pricing-manager/getMaxPricesForNodeProduct.spec.ts

@@ -0,0 +1,60 @@
+import { PRODUCT_IDS } from '@lit-protocol/constants';
+
+import { getMaxPricesForNodeProduct } from './getMaxPricesForNodeProduct';
+
+describe('getMaxPricesForNodeProduct', () => {
+  it('uses the requested product column when ranking nodes', () => {
+    const nodePrices = [
+      {
+        url: 'https://node-a',
+        prices: [80n, 5n, 9n, 30n],
+      },
+      {
+        url: 'https://node-b',
+        prices: [70n, 4n, 8n, 10n],
+      },
+      {
+        url: 'https://node-c',
+        prices: [60n, 3n, 7n, 20n],
+      },
+    ];
+
+    // Log the incoming order to show the decryption column is already sorted lowest-first.
+    console.log(
+      'incoming order',
+      nodePrices.map(({ url, prices }) => ({
+        url,
+        decryptionPrice: prices[PRODUCT_IDS.DECRYPTION],
+        signPrice: prices[PRODUCT_IDS.SIGN],
+        litActionPrice: prices[PRODUCT_IDS.LIT_ACTION],
+        signSessionKeyPrice: prices[PRODUCT_IDS.SIGN_SESSION_KEY],
+      }))
+    );
+
+    // Call the helper exactly like the SDK does: ask for SIGN_SESSION_KEY prices,
+    // pass the raw price feed output, and cap the request at two nodes.
+    const result = getMaxPricesForNodeProduct({
+      nodePrices,
+      userMaxPrice: 100n,
+      productId: PRODUCT_IDS.SIGN_SESSION_KEY,
+      numRequiredNodes: 2,
+    });
+
+    console.log(
+      'selected nodes',
+      result.map(({ url, price }) => ({ url, price }))
+    );
+
+    // After sorting the nodes by the session-key column, the helper should
+    // return node-b (10) and node-c (20) even though the original array was
+    // ordered by the decryption price column.
+    expect(result).toHaveLength(2);
+    expect(result[0].url).toBe('https://node-b');
+    expect(result[1].url).toBe('https://node-c');
+
+    // Base prices are taken from the SIGN_SESSION_KEY column (10 and 20)
+    // with the excess (100 - 30 = 70) split evenly.
+    expect(result[0].price).toBe(10n + 35n);
+    expect(result[1].price).toBe(20n + 35n);
+  });
+});

packages/networks/src/networks/vNaga/shared/managers/pricing-manager/getMaxPricesForNodeProduct.ts

@@ -15,6 +15,13 @@ export interface MaxPricesForNodes {
  * Ensures the total cost does not exceed userMaxPrice.
  * Operates in the order of lowest priced node to highest.
  *
+ * Example:
+ * - Selected nodes have SIGN_SESSION_KEY prices of 10 and 20.
+ * - `userMaxPrice` is 100.
+ * - Base total = 10 + 20 = 30.
+ * - Excess = 100 - 30 = 70.
+ * - Each node receives 70 / 2 = 35 extra budget, yielding 45 and 55.
+ *
  * @param nodePrices - An object where keys are node addresses and values are arrays of prices for different action types.
  * @param userMaxPrice - The maximum price the user is willing to pay to execute the request.
  * @param productId - The ID of the product to determine which price to consider.
@@ -28,19 +35,33 @@ export function getMaxPricesForNodeProduct({
   productId,
   numRequiredNodes,
 }: MaxPricesForNodes): { url: string; price: bigint }[] {
+  // Always evaluate pricing using the product-specific column so we truly pick
+  // the cheapest validators for that product (the upstream feed is sorted by
+  // prices[0]/decryption price only).
+  const sortedNodes = [...nodePrices].sort((a, b) => {
+    const priceA = a.prices[productId];
+    const priceB = b.prices[productId];
+
+    if (priceA === priceB) {
+      return 0;
+    }
+
+    return priceA < priceB ? -1 : 1;
+  });
+
   // If we don't need all nodes to service the request, only use the cheapest `n` of them
   const nodesToConsider = numRequiredNodes
-    ? nodePrices.slice(0, numRequiredNodes)
-    : nodePrices;
+    ? sortedNodes.slice(0, numRequiredNodes)
+    : sortedNodes;
 
+  // Sum the unadjusted cost for the nodes we plan to use.
   let totalBaseCost = 0n;
-
-  // Calculate the base total cost without adjustments
   for (const { prices } of nodesToConsider) {
+    // Example: base total accumulates 10 + 20 = 30 for the two cheapest nodes.
     totalBaseCost += prices[productId];
   }
 
-  // Verify that we have a high enough userMaxPrice to fulfill the request
+  // Refuse to proceed if the caller's budget cannot even cover the base cost.
   if (totalBaseCost > userMaxPrice) {
     throw new MaxPriceTooLow(
       {
@@ -58,13 +79,16 @@ export function getMaxPricesForNodeProduct({
    * our request to fail if the price on some of the nodes is higher than we think it was, as long as it's not
    * drastically different than we expect it to be
    */
+  // Any remaining budget is spread across the participating nodes to
+  // provide cushion for minor pricing fluctuations. Example: 100 - 30 = 70.
   const excessBalance = userMaxPrice - totalBaseCost;
 
   // Map matching the keys from `nodePrices`, but w/ the per-node maxPrice computed based on `userMaxPrice`
   const maxPricesPerNode: { url: string; price: bigint }[] = [];
 
   for (const { url, prices } of nodesToConsider) {
-    // For now, we'll distribute the remaining balance equally across nodes
+    // Distribute the remaining budget evenly across nodes to form the max price.
+    // Example: each node receives 70 / 2 = 35, becoming 10+35 and 20+35.
     maxPricesPerNode.push({
       url,
       price: excessBalance

packages/networks/src/networks/vNaga/shared/managers/pricing-manager/schema.ts

@@ -4,7 +4,7 @@ import { PRODUCT_IDS } from './constants';
 
 export const PricingContextSchema = z
   .object({
-    product: z.enum(['DECRYPTION', 'SIGN', 'LIT_ACTION']),
+    product: z.enum(['DECRYPTION', 'SIGN', 'LIT_ACTION', 'SIGN_SESSION_KEY']),
     userMaxPrice: z.bigint().optional(),
     nodePrices: z.array(
       z.object({ url: z.string(), prices: z.array(z.bigint()) })