feat(sdk-coin-sui): add fundsInAddressBalance support to StakingBuilder (AddStake)

modules/sdk-coin-sui/src/lib/stakingBuilder.ts

@@ -18,6 +18,7 @@ import {
   MoveCallTransaction,
   Inputs,
 } from './mystenlab/builder';
+import BigNumber from 'bignumber.js';
 import {
   ADD_STAKE_FUN_NAME,
   SUI_SYSTEM_ADDRESS,
@@ -31,6 +32,7 @@ import { MAX_COMMAND_ARGS, MAX_GAS_OBJECTS } from './constants';
 export class StakingBuilder extends TransactionBuilder<StakingProgrammableTransaction> {
   protected _addStakeTx: RequestAddStake[];
   protected _withdrawDelegation: RequestWithdrawStakedSui;
+  protected _fundsInAddressBalance: BigNumber = new BigNumber(0);
 
   constructor(_coinConfig: Readonly<CoinConfig>) {
     super(_coinConfig);
@@ -108,6 +110,18 @@ export class StakingBuilder extends TransactionBuilder<StakingProgrammableTransa
     return this;
   }
 
+  /**
+   * Set the amount of SUI held in the sender's address balance system.
+   * When set, the PTB will redeem this amount as a Coin<SUI> and merge it
+   * into the gas coin before splitting the stake amount.
+   *
+   * @param {string} amount - amount in MIST held in address balance
+   */
+  fundsInAddressBalance(amount: string): this {
+    this._fundsInAddressBalance = new BigNumber(amount);
+    return this;
+  }
+
   /** @inheritdoc */
   protected fromImplementation(rawTransaction: string): Transaction<StakingProgrammableTransaction> {
     const tx = new StakingTransaction(this._coinConfig);
@@ -151,6 +165,21 @@ export class StakingBuilder extends TransactionBuilder<StakingProgrammableTransa
     this.sender(txData.sender);
     this.gasData(txData.gasData);
 
+    if (txData.expiration && !('None' in txData.expiration)) {
+      this._expiration = txData.expiration;
+    }
+
+    // Restore fundsInAddressBalance from BalanceWithdrawal input if present.
+    const withdrawalInput = (tx.suiTransaction?.tx?.inputs as any[])?.find(
+      (input: any) =>
+        (input !== null && typeof input === 'object' && 'BalanceWithdrawal' in input) ||
+        (input?.value !== null && typeof input?.value === 'object' && 'BalanceWithdrawal' in (input.value ?? {}))
+    );
+    if (withdrawalInput) {
+      const bw = withdrawalInput.BalanceWithdrawal ?? withdrawalInput.value?.BalanceWithdrawal;
+      this._fundsInAddressBalance = new BigNumber(String(bw.reservation?.MaxAmountU64 ?? bw.amount));
+    }
+
     const requests = utils.getStakeRequests(tx.suiTransaction.tx);
     this.stake(requests);
   }
@@ -200,6 +229,18 @@ export class StakingBuilder extends TransactionBuilder<StakingProgrammableTransa
           }
         }
 
+        // When the sender has funds in the address balance system, redeem them as a
+        // Coin<SUI> and merge into the gas coin so SplitCoins can draw from the full
+        // available balance (coin objects + address balance).
+        if (this._fundsInAddressBalance.gt(0)) {
+          const [addrCoin] = programmableTxBuilder.moveCall({
+            target: '0x2::coin::redeem_funds',
+            typeArguments: ['0x2::sui::SUI'],
+            arguments: [programmableTxBuilder.withdrawal({ amount: BigInt(this._fundsInAddressBalance.toFixed()) })],
+          });
+          programmableTxBuilder.mergeCoins(programmableTxBuilder.gas, [addrCoin]);
+        }
+
         // Create a new coin with staking balance, based on the coins used as gas payment.
         this._addStakeTx.forEach((req) => {
           const coin = programmableTxBuilder.splitCoins(programmableTxBuilder.gas, [
@@ -242,6 +283,8 @@ export class StakingBuilder extends TransactionBuilder<StakingProgrammableTransa
         ...this._gasData,
         payment: this._gasData.payment.slice(0, MAX_GAS_OBJECTS - 1),
       },
+      expiration: this._expiration,
+      fundsInAddressBalance: this._fundsInAddressBalance.gt(0) ? this._fundsInAddressBalance.toFixed() : undefined,
     };
   }
 }

modules/sdk-coin-sui/src/lib/stakingTransaction.ts

@@ -74,7 +74,8 @@ export class StakingTransaction extends Transaction<StakingProgrammableTransacti
         ...tx.gasData,
         payment: [...tx.gasData.payment, ...this.getInputGasPaymentObjectsFromTx(tx.tx)],
       },
-      expiration: { None: null },
+      expiration: tx.expiration ?? { None: null },
+      fundsInAddressBalance: tx.fundsInAddressBalance,
     };
   }
 
@@ -191,7 +192,15 @@ export class StakingTransaction extends Transaction<StakingProgrammableTransacti
           return Inputs.Pure(amount, BCS.U64);
         }
       }
-      if (input.kind === 'Input' && (input.value.hasOwnProperty('Object') || input.value.hasOwnProperty('Pure'))) {
+      if (input.hasOwnProperty('BalanceWithdrawal')) {
+        return input;
+      }
+      if (
+        input.kind === 'Input' &&
+        (input.value.hasOwnProperty('Object') ||
+          input.value.hasOwnProperty('Pure') ||
+          input.value.hasOwnProperty('BalanceWithdrawal'))
+      ) {
         return input.value;
       }
 
@@ -206,7 +215,7 @@ export class StakingTransaction extends Transaction<StakingProgrammableTransacti
 
     return {
       sender: this._suiTransaction.sender,
-      expiration: { None: null },
+      expiration: this._suiTransaction.expiration ?? { None: null },
       gasData: {
         ...this._suiTransaction.gasData,
         payment: this._suiTransaction.gasData.payment.slice(0, MAX_GAS_OBJECTS - 1),

modules/sdk-coin-sui/src/lib/utils.ts

@@ -383,7 +383,10 @@ export class Utils implements BaseUtils {
         const amountInputIdx = ((transaction as SplitCoinsTransaction).amounts[0] as TransactionBlockInput).index;
         amounts.push(utils.getAmount(tx.inputs[amountInputIdx] as TransactionBlockInput));
       }
-      if (transaction.kind === 'MoveCall') {
+      if (
+        transaction.kind === 'MoveCall' &&
+        (transaction as MoveCallTransaction).target.endsWith('::sui_system::request_add_stake')
+      ) {
         const validatorAddressInputIdx = ((transaction as MoveCallTransaction).arguments[2] as TransactionBlockInput)
           .index;
         const validatorAddress = utils.getAddress(tx.inputs[validatorAddressInputIdx] as TransactionBlockInput);
@@ -406,7 +409,10 @@ export class Utils implements BaseUtils {
         const amountInputIdx = ((transaction as SplitCoinsTransaction).amounts[0] as TransactionBlockInput).index;
         amounts.push(utils.getAmount(tx.inputs[amountInputIdx] as TransactionBlockInput));
       }
-      if (transaction.kind === 'MoveCall') {
+      if (
+        transaction.kind === 'MoveCall' &&
+        (transaction as MoveCallTransaction).target.endsWith('::staking::stake_with_pool')
+      ) {
         const validatorAddressInputIdx = ((transaction as MoveCallTransaction).arguments[2] as TransactionBlockInput)
           .index;
         const validatorAddress = utils.getAddress(tx.inputs[validatorAddressInputIdx] as TransactionBlockInput);

modules/sdk-coin-sui/src/lib/walrusStakingBuilder.ts

@@ -11,7 +11,7 @@ import { Transaction } from './transaction';
 import utils from './utils';
 import assert from 'assert';
 import { TransferTransaction } from './transferTransaction';
-import { StakingTransaction } from './stakingTransaction';
+import { WalrusStakingTransaction } from './walrusStakingTransaction';
 import {
   Inputs,
   MoveCallTransaction,
@@ -30,7 +30,7 @@ export class WalrusStakingBuilder extends TransactionBuilder<WalrusStakingProgra
 
   constructor(_coinConfig: Readonly<CoinConfig>) {
     super(_coinConfig);
-    this._transaction = new StakingTransaction(_coinConfig);
+    this._transaction = new WalrusStakingTransaction(_coinConfig);
 
     // TODO improve mainnet vs. testnet configuration
     this.walrusConfig = _coinConfig.network.type === NetworkType.MAINNET ? WALRUS_PROD_CONFIG : WALRUS_TESTNET_CONFIG;
@@ -120,7 +120,7 @@ export class WalrusStakingBuilder extends TransactionBuilder<WalrusStakingProgra
 
   /** @inheritdoc */
   protected fromImplementation(rawTransaction: string): Transaction<WalrusStakingProgrammableTransaction> {
-    const tx = new StakingTransaction(this._coinConfig);
+    const tx = new WalrusStakingTransaction(this._coinConfig);
     this.validateRawTransaction(rawTransaction);
     tx.fromRawTransaction(rawTransaction);
     this.initBuilder(tx);
@@ -147,7 +147,7 @@ export class WalrusStakingBuilder extends TransactionBuilder<WalrusStakingProgra
   /**
    * Initialize the transaction builder fields using the decoded transaction data
    *
-   * @param {StakingTransaction} tx the transaction data
+   * @param {WalrusStakingTransaction} tx the transaction data
    */
   initBuilder(tx: Transaction<WalrusStakingProgrammableTransaction>): void {
     this._transaction = tx;

modules/sdk-coin-sui/test/unit/transactionBuilder/stakingBuilder.ts

@@ -98,6 +98,179 @@ describe('Sui Staking Builder', () => {
     });
   });
 
+  describe('fundsInAddressBalance', () => {
+    const FUNDS_IN_ADDRESS_BALANCE = '5000000000'; // 5 SUI in MIST
+
+    it('should build AddStake with mixed coin objects + address balance (redeem → merge → split → add_stake)', async function () {
+      // When fundsInAddressBalance > 0, the PTB must insert redeem_funds + mergeCoins into gas
+      // coin BEFORE the splitCoins + request_add_stake commands, so the full balance is available.
+      const txBuilder = factory.getStakingBuilder();
+      txBuilder.type(SuiTransactionType.AddStake);
+      txBuilder.sender(testData.sender.address);
+      txBuilder.stake([testData.requestAddStake]);
+      txBuilder.gasData(testData.gasData);
+      txBuilder.fundsInAddressBalance(FUNDS_IN_ADDRESS_BALANCE);
+
+      const tx = await txBuilder.build();
+      should.equal(tx.type, TransactionType.StakingAdd);
+
+      const suiTx = tx as SuiTransaction<StakingProgrammableTransaction>;
+
+      // Expected PTB command sequence:
+      //   0: MoveCall(redeem_funds)       — materialise Coin<SUI> from address balance
+      //   1: MergeCoins(gas, [addrCoin])  — merge address-balance coin into gas coin
+      //   2: SplitCoins(gas, [amount])    — split stake amount from gas coin
+      //   3: MoveCall(request_add_stake)  — stake to validator
+      const cmds = suiTx.suiTransaction.tx.transactions as any[];
+      cmds.length.should.equal(4, 'expected 4 commands: redeem, merge, split, add_stake');
+      cmds[0].kind.should.equal('MoveCall', 'command 0 must be MoveCall(redeem_funds)');
+      cmds[0].target.should.equal('0x2::coin::redeem_funds');
+      cmds[1].kind.should.equal('MergeCoins', 'command 1 must be MergeCoins(gas, [addrCoin])');
+      cmds[2].kind.should.equal('SplitCoins', 'command 2 must be SplitCoins(gas, [amount])');
+      cmds[3].kind.should.equal('MoveCall', 'command 3 must be MoveCall(request_add_stake)');
+      cmds[3].target.should.endWith('::sui_system::request_add_stake');
+
+      // fundsInAddressBalance must be persisted on the SuiTransaction
+      suiTx.suiTransaction.fundsInAddressBalance!.should.equal(FUNDS_IN_ADDRESS_BALANCE);
+
+      // getStakeRequests must still correctly identify the request despite the extra commands
+      const requests = utils.getStakeRequests(suiTx.suiTransaction.tx);
+      requests.length.should.equal(1);
+      requests[0].validatorAddress.should.equal(testData.requestAddStake.validatorAddress);
+      requests[0].amount.should.equal(testData.requestAddStake.amount);
+
+      // toJson must emit fundsInAddressBalance
+      const json = tx.toJson();
+      json.fundsInAddressBalance!.should.equal(FUNDS_IN_ADDRESS_BALANCE);
+
+      const rawTx = tx.toBroadcastFormat();
+      should.equal(utils.isValidRawTransaction(rawTx), true);
+    });
+
+    it('should round-trip: build → serialize → initBuilder restores fundsInAddressBalance', async function () {
+      const txBuilder = factory.getStakingBuilder();
+      txBuilder.type(SuiTransactionType.AddStake);
+      txBuilder.sender(testData.sender.address);
+      txBuilder.stake([testData.requestAddStake]);
+      txBuilder.gasData(testData.gasData);
+      txBuilder.fundsInAddressBalance(FUNDS_IN_ADDRESS_BALANCE);
+
+      const tx = await txBuilder.build();
+      const rawTx = tx.toBroadcastFormat();
+      should.equal(utils.isValidRawTransaction(rawTx), true);
+
+      // Deserialize and rebuild
+      const rebuilder = factory.from(rawTx);
+      rebuilder.addSignature({ pub: testData.sender.publicKey }, Buffer.from(testData.sender.signatureHex));
+      const rebuiltTx = await rebuilder.build();
+      rebuiltTx.toBroadcastFormat().should.equal(rawTx);
+
+      // initBuilder must have restored fundsInAddressBalance
+      const rebuiltSuiTx = rebuiltTx as SuiTransaction<StakingProgrammableTransaction>;
+      rebuiltSuiTx.suiTransaction.fundsInAddressBalance!.should.equal(FUNDS_IN_ADDRESS_BALANCE);
+
+      // stake requests must still be intact
+      const requests = utils.getStakeRequests(rebuiltSuiTx.suiTransaction.tx);
+      requests.length.should.equal(1);
+      requests[0].validatorAddress.should.equal(testData.requestAddStake.validatorAddress);
+    });
+
+    it('should build multiple stake requests with fundsInAddressBalance (multi-stake)', async function () {
+      const txBuilder = factory.getStakingBuilder();
+      txBuilder.type(SuiTransactionType.AddStake);
+      txBuilder.sender(testData.sender.address);
+      txBuilder.stake(testData.requestAddStakeMany);
+      txBuilder.gasData(testData.gasData);
+      txBuilder.fundsInAddressBalance(FUNDS_IN_ADDRESS_BALANCE);
+
+      const tx = await txBuilder.build();
+      should.equal(tx.type, TransactionType.StakingAdd);
+
+      const suiTx = tx as SuiTransaction<StakingProgrammableTransaction>;
+      const cmds = suiTx.suiTransaction.tx.transactions as any[];
+
+      // Command sequence: redeem_funds, mergeCoins, then for each stake: splitCoins + request_add_stake
+      cmds[0].kind.should.equal('MoveCall', 'command 0 must be redeem_funds');
+      cmds[0].target.should.equal('0x2::coin::redeem_funds');
+      cmds[1].kind.should.equal('MergeCoins', 'command 1 must be mergeCoins');
+
+      const numStakes = testData.requestAddStakeMany.length;
+      // After redeem+merge, each stake adds 2 commands: splitCoins + MoveCall
+      cmds.length.should.equal(2 + numStakes * 2);
+
+      const requests = utils.getStakeRequests(suiTx.suiTransaction.tx);
+      requests.length.should.equal(numStakes);
+      for (let i = 0; i < numStakes; i++) {
+        requests[i].validatorAddress.should.equal(testData.requestAddStakeMany[i].validatorAddress);
+        requests[i].amount.should.equal(testData.requestAddStakeMany[i].amount);
+      }
+    });
+
+    it('should not insert redeem_funds when fundsInAddressBalance is 0 (default behavior unchanged)', async function () {
+      const txBuilder = factory.getStakingBuilder();
+      txBuilder.type(SuiTransactionType.AddStake);
+      txBuilder.sender(testData.sender.address);
+      txBuilder.stake([testData.requestAddStake]);
+      txBuilder.gasData(testData.gasData);
+      // No fundsInAddressBalance call — default is 0
+
+      const tx = await txBuilder.build();
+      const suiTx = tx as SuiTransaction<StakingProgrammableTransaction>;
+      const cmds = suiTx.suiTransaction.tx.transactions as any[];
+
+      // Original 2-command sequence: splitCoins + request_add_stake
+      cmds.length.should.equal(2);
+      cmds[0].kind.should.equal('SplitCoins');
+      cmds[1].kind.should.equal('MoveCall');
+      cmds[1].target.should.endWith('::sui_system::request_add_stake');
+
+      should.equal(suiTx.suiTransaction.fundsInAddressBalance, undefined);
+    });
+
+    it('should build Case 2 (addr-bal only, empty payment): redeem → merge → split → add_stake', async function () {
+      // Case 2: gasData.payment=[] — the stake amount comes entirely from address balance.
+      // SplitCoins(GasCoin) is unsafe with empty payment because GasCoin only carries up to
+      // gas-budget-worth of balance. The same redeem_funds → mergeCoins(gas) path is used,
+      // after which GasCoin holds the full address balance and SplitCoins works correctly.
+      const gasDataNoPayment = {
+        ...testData.gasDataWithoutGasPayment,
+        payment: [],
+      };
+      const txBuilder = factory.getStakingBuilder();
+      txBuilder.type(SuiTransactionType.AddStake);
+      txBuilder.sender(testData.sender.address);
+      txBuilder.stake([testData.requestAddStake]);
+      txBuilder.gasData(gasDataNoPayment);
+      txBuilder.fundsInAddressBalance(FUNDS_IN_ADDRESS_BALANCE);
+
+      const tx = await txBuilder.build();
+      should.equal(tx.type, TransactionType.StakingAdd);
+
+      const suiTx = tx as SuiTransaction<StakingProgrammableTransaction>;
+      suiTx.suiTransaction.gasData.payment.length.should.equal(0);
+
+      // Command sequence: redeem_funds, mergeCoins, splitCoins, request_add_stake
+      const cmds = suiTx.suiTransaction.tx.transactions as any[];
+      cmds.length.should.equal(4);
+      cmds[0].kind.should.equal('MoveCall');
+      cmds[0].target.should.equal('0x2::coin::redeem_funds');
+      cmds[1].kind.should.equal('MergeCoins');
+      cmds[2].kind.should.equal('SplitCoins');
+      cmds[3].kind.should.equal('MoveCall');
+      cmds[3].target.should.endWith('::sui_system::request_add_stake');
+
+      suiTx.suiTransaction.fundsInAddressBalance!.should.equal(FUNDS_IN_ADDRESS_BALANCE);
+
+      const rawTx = tx.toBroadcastFormat();
+      should.equal(utils.isValidRawTransaction(rawTx), true);
+
+      // Round-trip
+      const rebuilder = factory.from(rawTx);
+      const rebuiltTx = await rebuilder.build();
+      rebuiltTx.toBroadcastFormat().should.equal(rawTx);
+    });
+  });
+
   describe('Fail', () => {
     it('should fail for invalid sender', async function () {
       const builder = factory.getStakingBuilder();