Skip to content

feat(sdk-coin-sui): add fundsInAddressBalance to WalrusStakingBuilder #9340

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
bitgo-ai-agent-dev wants to merge 1 commit into
base: master
Choose a base branch
from
+295 −34
Draft

feat(sdk-coin-sui): add fundsInAddressBalance to WalrusStakingBuilder #9340

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
108 changes: 97 additions & 11 deletions modules/sdk-coin-sui/src/lib/walrusStakingBuilder.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,10 +21,12 @@ import {
import { MAX_GAS_OBJECTS } from './constants';
import { WALRUS_PROD_CONFIG, WALRUS_TESTNET_CONFIG } from './resources/walrusConfig';
import { SuiObjectRef } from './mystenlab/types';
import BigNumber from 'bignumber.js';

export class WalrusStakingBuilder extends TransactionBuilder<WalrusStakingProgrammableTransaction> {
protected _stakeWithPoolTx: RequestWalrusStakeWithPool[];
protected _inputObjects: SuiObjectRef[];
protected _fundsInAddressBalance: BigNumber = new BigNumber(0);

private walrusConfig: any; // TODO improve

Expand Down Expand Up @@ -108,6 +110,19 @@ export class WalrusStakingBuilder extends TransactionBuilder<WalrusStakingProgra
return this;
}

/**
* Set the amount of WAL held in the sender's address balance system.
* When set, the PTB will redeem this amount as a Coin<WAL> before
* merging/splitting and staking. WAL principal can come from coin
* objects, address balance, or a mix.
*
* @param {string} amount - amount in WAL base units held in address balance
*/
fundsInAddressBalance(amount: string): this {
this._fundsInAddressBalance = new BigNumber(amount);
return this;
}

private validateInputObjects(inputObjects: SuiObjectRef[]): void {
assert(
inputObjects && inputObjects.length > 0,
Expand Down Expand Up @@ -161,11 +176,27 @@ export class WalrusStakingBuilder extends TransactionBuilder<WalrusStakingProgra
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.getWalrusStakeWithPoolRequests(tx.suiTransaction.tx);
this.stake(requests);

assert(txData.inputObjects);
this.inputObjects(txData.inputObjects);
if (txData.inputObjects && txData.inputObjects.length > 0) {
this.inputObjects(txData.inputObjects);
}
}

/**
Expand All @@ -180,7 +211,17 @@ export class WalrusStakingBuilder extends TransactionBuilder<WalrusStakingProgra
});
assert(this._gasData, new BuildTransactionError('gasData is required before building'));
this.validateGasData(this._gasData);
this.validateInputObjects(this._inputObjects);
Comment thread
benkelcher-bitgo marked this conversation as resolved.
const hasInputObjects = this._inputObjects && this._inputObjects.length > 0;
const hasAddrBal = this._fundsInAddressBalance && this._fundsInAddressBalance.gt(0);
assert(
hasInputObjects || hasAddrBal,
new BuildTransactionError('either inputObjects or fundsInAddressBalance is required before building')
);
if (hasInputObjects) {
this._inputObjects.forEach((inputObject) => {
this.validateSuiObjectRef(inputObject, 'input object');
});
}
}

/**
Expand All @@ -194,20 +235,54 @@ export class WalrusStakingBuilder extends TransactionBuilder<WalrusStakingProgra

const programmableTxBuilder = new ProgrammingTransactionBlockBuilder();
switch (this._type) {
case SuiTransactionType.WalrusStakeWithPool:
const inputObjects = this._inputObjects.map((token) => programmableTxBuilder.object(Inputs.ObjectRef(token)));
const mergedObject = inputObjects.shift() as TransactionArgument;
case SuiTransactionType.WalrusStakeWithPool: {
const walCoinType = `${this.walrusConfig.WAL_PKG_ID}::${this.walrusConfig.WAL_COIN_MODULE_NAME}::${this.walrusConfig.WAL_COIN_NAME}`;

// Step 1: consolidate any provided Coin<WAL> objects into a single coin.
let baseWalCoin: TransactionArgument | undefined;
if (this._inputObjects && this._inputObjects.length > 0) {
const inputObjects = this._inputObjects.map((token) => programmableTxBuilder.object(Inputs.ObjectRef(token)));
baseWalCoin = inputObjects.shift() as TransactionArgument;
if (inputObjects.length > 0) {
programmableTxBuilder.mergeCoins(baseWalCoin, inputObjects);
}
}

// Step 2: redeem address-balance funds as a Coin<WAL> if requested.
let redeemedCoin: TransactionArgument | undefined;
if (this._fundsInAddressBalance.gt(0)) {
const [coin] = programmableTxBuilder.moveCall({
target: '0x2::coin::redeem_funds',
typeArguments: [walCoinType],
arguments: [
programmableTxBuilder.withdrawal({
amount: BigInt(this._fundsInAddressBalance.toFixed()),
type: walCoinType,
}),
],
});
redeemedCoin = coin;
}

if (inputObjects.length > 0) {
programmableTxBuilder.mergeCoins(mergedObject, inputObjects);
// Step 3: combine into a single definite walCoin.
let walCoin: TransactionArgument;
if (baseWalCoin !== undefined && redeemedCoin !== undefined) {
programmableTxBuilder.mergeCoins(baseWalCoin, [redeemedCoin]);
walCoin = baseWalCoin;
} else if (baseWalCoin !== undefined) {
walCoin = baseWalCoin;
} else if (redeemedCoin !== undefined) {
walCoin = redeemedCoin;
} else {
// Unreachable: validateTransactionFields ensures at least one source.
throw new BuildTransactionError('either inputObjects or fundsInAddressBalance is required before building');
}

// Create a new coin with staking balance, based on the coins used as gas payment.
// Create a new coin with staking balance and stake each request.
const stakedWals = this._stakeWithPoolTx.map((req) => {
const splitObject = programmableTxBuilder.splitCoins(mergedObject, [
const splitObject = programmableTxBuilder.splitCoins(walCoin, [
programmableTxBuilder.pure(Number(req.amount)),
]);
// Stake the split coin to a specific validator address.
return programmableTxBuilder.moveCall({
target: `${this.walrusConfig.WALRUS_PKG_ID}::${this.walrusConfig.WALRUS_STAKING_MODULE_NAME}::${this.walrusConfig.WALRUS_STAKE_WITH_POOL_FUN_NAME}`,
arguments: [
Expand All @@ -218,8 +293,17 @@ export class WalrusStakingBuilder extends TransactionBuilder<WalrusStakingProgra
} as unknown as MoveCallTransaction);
});

// Transfer staked WAL objects back to sender.
programmableTxBuilder.transferObjects(stakedWals, programmableTxBuilder.object(this._sender));

// Coin<WAL> cannot be merged into the SUI gas coin — transfer any
// residual WAL coin back to sender to avoid a no-drop violation.
if (redeemedCoin !== undefined) {
programmableTxBuilder.transferObjects([walCoin], programmableTxBuilder.object(this._sender));
}

break;
}
default:
throw new InvalidTransactionError(`unsupported target method`);
}
Expand All @@ -236,6 +320,8 @@ export class WalrusStakingBuilder extends TransactionBuilder<WalrusStakingProgra
...this._gasData,
payment: this._gasData.payment.slice(0, MAX_GAS_OBJECTS - 1),
},
expiration: this._expiration,
fundsInAddressBalance: this._fundsInAddressBalance.gt(0) ? this._fundsInAddressBalance.toFixed() : undefined,
};
}
}
66 changes: 44 additions & 22 deletions modules/sdk-coin-sui/src/lib/walrusStakingTransaction.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,13 +68,15 @@ export class WalrusStakingTransaction extends Transaction<WalrusStakingProgramma
}

const tx = this._suiTransaction;
const inputObjects = this.getInputObjectsFromTx(tx.tx);
return {
id: this._id,
sender: tx.sender,
kind: { ProgrammableTransaction: tx.tx },
gasData: tx.gasData,
expiration: { None: null },
inputObjects: this.getInputObjectsFromTx(tx.tx),
expiration: tx.expiration ?? { None: null },
...(inputObjects.length > 0 && { inputObjects }),
fundsInAddressBalance: tx.fundsInAddressBalance,
};
}

Expand Down Expand Up @@ -191,7 +193,15 @@ export class WalrusStakingTransaction extends Transaction<WalrusStakingProgramma
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;
}

Expand All @@ -206,7 +216,7 @@ export class WalrusStakingTransaction extends Transaction<WalrusStakingProgramma

return {
sender: this._suiTransaction.sender,
expiration: { None: null },
expiration: this._suiTransaction.expiration ?? { None: null },
gasData: this._suiTransaction.gasData,
kind: {
ProgrammableTransaction: programmableTx,
Expand Down Expand Up @@ -255,28 +265,40 @@ export class WalrusStakingTransaction extends Transaction<WalrusStakingProgramma

private getInputObjectsFromTx(tx: WalrusStakingProgrammableTransaction): SuiObjectRef[] {
const inputs = tx.inputs;
const transaction = tx.transactions[0] as SuiTransactionBlockType;

let args: TransactionArgument[] = [];
if (transaction.kind === 'MergeCoins') {
const { destination, sources } = transaction;
args = [destination, ...sources];
} else if (transaction.kind === 'SplitCoins') {
args = [transaction.coin];
}

// Scan all transactions for MergeCoins/SplitCoins that reference ImmOrOwned
// coin object inputs. When fundsInAddressBalance is used, a MoveCall
// (redeem_funds) may appear before the first merge/split, so we cannot
// assume transactions[0] is the relevant one.
const inputObjects: SuiObjectRef[] = [];
args.forEach((arg) => {
if (arg.kind === 'Input') {
let input = inputs[arg.index];
if ('value' in input) {
input = input.value;
}
if ('Object' in input && isImmOrOwnedObj(input.Object)) {
inputObjects.push(input.Object.ImmOrOwned);
for (const txn of tx.transactions as SuiTransactionBlockType[]) {
let args: TransactionArgument[] = [];
if (txn.kind === 'MergeCoins') {
const { destination, sources } = txn;
args = [destination, ...sources];
} else if (txn.kind === 'SplitCoins') {
args = [txn.coin];
} else {
continue;
}

args.forEach((arg) => {
if (arg.kind === 'Input') {
let input = inputs[arg.index];
if ('value' in input) {
input = input.value;
}
if ('Object' in input && isImmOrOwnedObj(input.Object)) {
inputObjects.push(input.Object.ImmOrOwned);
}
}
});

// Stop after the first merge/split that yields coin object refs.
if (inputObjects.length > 0) {
break;
}
});
}

return inputObjects;
}
Expand Down
Loading
Loading