Skip to content

Commit ac5b6fc

Browse files
haritkapadiaharitkapadia
committed
feat(sdk-coin-sol): conditionally create ATA for jito staking
Ticket: SC-3118
1 parent 39bfb65 commit ac5b6fc

File tree

9 files changed

+169
-114
lines changed

9 files changed

+169
-114
lines changed

‎examples/ts/sol/stake-jito.ts‎

Lines changed: 65 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -3,106 +3,110 @@
33
*
44
* Copyright 2025, BitGo, Inc. All Rights Reserved.
55
*/
6-
import { BitGoAPI } from '@bitgo/sdk-api'
7-
import { TransactionBuilderFactory, Tsol } from '@bitgo/sdk-coin-sol'
8-
import { coins } from '@bitgo/statics'
9-
import { Connection, PublicKey, clusterApiUrl, Transaction, Keypair, LAMPORTS_PER_SOL } from "@solana/web3.js"
10-
import { getStakePoolAccount, updateStakePool } from '@solana/spl-stake-pool'
6+
import { SolStakingTypeEnum } from '@bitgo/public-types';
7+
import { BitGoAPI } from '@bitgo/sdk-api';
8+
import { TransactionBuilderFactory, Tsol } from '@bitgo/sdk-coin-sol';
9+
import { coins } from '@bitgo/statics';
10+
import { Connection, PublicKey, clusterApiUrl, Transaction, Keypair, LAMPORTS_PER_SOL } from '@solana/web3.js';
11+
import { getStakePoolAccount, updateStakePool } from '@solana/spl-stake-pool';
12+
import { getAssociatedTokenAddressSync } from '@solana/spl-token';
1113
import * as bs58 from 'bs58';
1214

13-
require('dotenv').config({ path: '../../.env' })
15+
require('dotenv').config({ path: '../../.env' });
1416

15-
const AMOUNT_LAMPORTS = 1000
16-
const JITO_STAKE_POOL_ADDRESS = 'Jito4APyf642JPZPx3hGc6WWJ8zPKtRbRs4P815Awbb'
17-
const NETWORK = 'devnet'
17+
const AMOUNT_LAMPORTS = 1000;
18+
const JITO_STAKE_POOL_ADDRESS = 'Jito4APyf642JPZPx3hGc6WWJ8zPKtRbRs4P815Awbb';
19+
const NETWORK = 'devnet';
1820

1921
const bitgo = new BitGoAPI({
2022
accessToken: process.env.TESTNET_ACCESS_TOKEN,
2123
env: 'test',
22-
})
23-
const coin = coins.get("tsol")
24-
bitgo.register(coin.name, Tsol.createInstance)
24+
});
25+
const coin = coins.get('tsol');
26+
bitgo.register(coin.name, Tsol.createInstance);
2527

2628
async function main() {
27-
const account = getAccount()
28-
const connection = new Connection(clusterApiUrl(NETWORK), 'confirmed')
29-
const recentBlockhash = await connection.getLatestBlockhash()
30-
const stakePoolAccount = await getStakePoolAccount(connection, new PublicKey(JITO_STAKE_POOL_ADDRESS))
31-
29+
const account = getAccount();
30+
const connection = new Connection(clusterApiUrl(NETWORK), 'confirmed');
31+
const recentBlockhash = await connection.getLatestBlockhash();
32+
const stakePoolAccount = await getStakePoolAccount(connection, new PublicKey(JITO_STAKE_POOL_ADDRESS));
33+
const associatedTokenAddress = getAssociatedTokenAddressSync(
34+
stakePoolAccount.account.data.poolMint,
35+
account.publicKey
36+
);
37+
const associatedTokenAccountExists = !!(await connection.getAccountInfo(associatedTokenAddress));
3238

3339
// Account should have sufficient balance
34-
const accountBalance = await connection.getBalance(account.publicKey)
40+
const accountBalance = await connection.getBalance(account.publicKey);
3541
if (accountBalance < 0.1 * LAMPORTS_PER_SOL) {
36-
console.info(`Your account balance is ${accountBalance / LAMPORTS_PER_SOL} SOL, requesting airdrop`)
37-
const sig = await connection.requestAirdrop(account.publicKey, 2 * LAMPORTS_PER_SOL)
38-
await connection.confirmTransaction(sig)
39-
console.info(`Airdrop successful: ${sig}`)
42+
console.info(`Your account balance is ${accountBalance / LAMPORTS_PER_SOL} SOL, requesting airdrop`);
43+
const sig = await connection.requestAirdrop(account.publicKey, 2 * LAMPORTS_PER_SOL);
44+
await connection.confirmTransaction(sig);
45+
console.info(`Airdrop successful: ${sig}`);
4046
}
4147

4248
// Stake pool should be up to date
43-
const epochInfo = await connection.getEpochInfo()
49+
const epochInfo = await connection.getEpochInfo();
4450
if (stakePoolAccount.account.data.lastUpdateEpoch.ltn(epochInfo.epoch)) {
45-
console.info('Stake pool is out of date.')
46-
const usp = await updateStakePool(connection, stakePoolAccount)
47-
const tx = new Transaction()
48-
tx.add(...usp.updateListInstructions, ...usp.finalInstructions)
49-
const signer = Keypair.fromSecretKey(account.secretKeyArray)
50-
const sig = await connection.sendTransaction(tx, [signer])
51-
await connection.confirmTransaction(sig)
52-
console.info(`Stake pool updated: ${sig}`)
51+
console.info('Stake pool is out of date.');
52+
const usp = await updateStakePool(connection, stakePoolAccount);
53+
const tx = new Transaction();
54+
tx.add(...usp.updateListInstructions, ...usp.finalInstructions);
55+
const signer = Keypair.fromSecretKey(account.secretKeyArray);
56+
const sig = await connection.sendTransaction(tx, [signer]);
57+
await connection.confirmTransaction(sig);
58+
console.info(`Stake pool updated: ${sig}`);
5359
}
5460

5561
// Use BitGoAPI to build depositSol instruction
56-
const txBuilder = new TransactionBuilderFactory(coin).getStakingActivateBuilder()
62+
const txBuilder = new TransactionBuilderFactory(coin).getStakingActivateBuilder();
5763
txBuilder
5864
.amount(`${AMOUNT_LAMPORTS}`)
5965
.sender(account.publicKey.toBase58())
6066
.stakingAddress(JITO_STAKE_POOL_ADDRESS)
6167
.validator(JITO_STAKE_POOL_ADDRESS)
62-
.stakingTypeParams({
63-
type: 'JITO',
68+
.stakingType(SolStakingTypeEnum.JITO)
69+
.extraParams({
6470
stakePoolData: {
65-
managerFeeAccount: stakePoolAccount.account.data.managerFeeAccount.toString(),
66-
poolMint: stakePoolAccount.account.data.poolMint.toString(),
67-
reserveStake: stakePoolAccount.account.data.toString(),
68-
}
71+
managerFeeAccount: stakePoolAccount.account.data.managerFeeAccount.toBase58(),
72+
poolMint: stakePoolAccount.account.data.poolMint.toBase58(),
73+
reserveStake: stakePoolAccount.account.data.reserveStake.toBase58(),
74+
},
75+
createAssociatedTokenAccount: !associatedTokenAccountExists,
6976
})
70-
.nonce(recentBlockhash.blockhash)
71-
txBuilder.sign({ key: account.secretKey })
72-
const tx = await txBuilder.build()
73-
const serializedTx = tx.toBroadcastFormat()
74-
console.info(`Transaction JSON:\n${JSON.stringify(tx.toJson(), undefined, 2)}`)
77+
.nonce(recentBlockhash.blockhash);
78+
txBuilder.sign({ key: account.secretKey });
79+
const tx = await txBuilder.build();
80+
const serializedTx = tx.toBroadcastFormat();
81+
console.info(`Transaction JSON:\n${JSON.stringify(tx.toJson(), undefined, 2)}`);
7582

7683
// Send transaction
7784
try {
78-
const sig = await connection.sendRawTransaction(Buffer.from(serializedTx, 'base64'))
79-
await connection.confirmTransaction(sig)
80-
console.log(`${AMOUNT_LAMPORTS / LAMPORTS_PER_SOL} SOL deposited`, sig)
85+
const sig = await connection.sendRawTransaction(Buffer.from(serializedTx, 'base64'));
86+
await connection.confirmTransaction(sig);
87+
console.log(`${AMOUNT_LAMPORTS / LAMPORTS_PER_SOL} SOL deposited`, sig);
8188
} catch (e) {
82-
console.log('Error sending transaction')
83-
console.error(e)
84-
if (e.transactionMessage === 'Transaction simulation failed: Error processing Instruction 0: Provided owner is not allowed') {
85-
console.error('If you successfully staked JitoSOL once, you cannot stake again.')
86-
}
89+
console.log('Error sending transaction');
90+
console.error(e);
8791
}
8892
}
8993

9094
const getAccount = () => {
91-
const publicKey = process.env.ACCOUNT_PUBLIC_KEY
92-
const secretKey = process.env.ACCOUNT_SECRET_KEY
95+
const publicKey = process.env.ACCOUNT_PUBLIC_KEY;
96+
const secretKey = process.env.ACCOUNT_SECRET_KEY;
9397
if (publicKey === undefined || secretKey === undefined) {
94-
const { publicKey, secretKey } = Keypair.generate()
95-
console.log('# Here is a new account to save into your .env file.')
96-
console.log(`ACCOUNT_PUBLIC_KEY=${publicKey.toBase58()}`)
97-
console.log(`ACCOUNT_SECRET_KEY=${bs58.encode(secretKey)}`)
98-
throw new Error("Missing account information")
98+
const { publicKey, secretKey } = Keypair.generate();
99+
console.log('# Here is a new account to save into your .env file.');
100+
console.log(`ACCOUNT_PUBLIC_KEY=${publicKey.toBase58()}`);
101+
console.log(`ACCOUNT_SECRET_KEY=${bs58.encode(secretKey)}`);
102+
throw new Error('Missing account information');
99103
}
100104

101105
return {
102106
publicKey: new PublicKey(publicKey),
103107
secretKey,
104108
secretKeyArray: new Uint8Array(bs58.decode(secretKey)),
105-
}
106-
}
109+
};
110+
};
107111

108-
main().catch((e) => console.error(e))
112+
main().catch((e) => console.error(e));

‎modules/sdk-coin-sol/src/lib/constants.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,11 @@ export const marinadeStakingActivateInstructionsIndexes = {
115115

116116
/** Const to check the order of the Jito Staking Activate instructions when decode */
117117
export const jitoStakingActivateInstructionsIndexes = {
118+
DepositSol: 0,
119+
} as const;
120+
121+
/** Const to check the order of the Jito Staking Activate instructions when decode */
122+
export const jitoStakingActivateWithATAInstructionsIndexes = {
118123
InitializeAssociatedTokenAccount: 0,
119124
DepositSol: 1,
120125
} as const;

‎modules/sdk-coin-sol/src/lib/iface.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,7 @@ export interface Approve {
123123

124124
export interface JitoStakingActivateParams {
125125
stakePoolData: DepositSolStakePoolData;
126+
createAssociatedTokenAccount?: boolean;
126127
}
127128

128129
export type StakingActivateExtraParams = JitoStakingActivateParams;

‎modules/sdk-coin-sol/src/lib/instructionParamsFactory.ts‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -330,6 +330,7 @@ type StakingInstructions = {
330330
create?: CreateAccountParams;
331331
initialize?: InitializeStakeParams;
332332
delegate?: DelegateStakeParams;
333+
hasAtaInit?: boolean;
333334
};
334335

335336
type JitoStakingInstructions = StakingInstructions & {
@@ -419,6 +420,7 @@ function parseStakingActivateInstructions(
419420
break;
420421

421422
case ValidInstructionTypesEnum.InitializeAssociatedTokenAccount:
423+
stakingInstructions.hasAtaInit = true;
422424
instructionData.push({
423425
type: InstructionBuilderTypes.CreateAssociatedTokenAccount,
424426
params: {
@@ -441,7 +443,7 @@ function parseStakingActivateInstructions(
441443
switch (stakingType) {
442444
case SolStakingTypeEnum.JITO: {
443445
assert(isJitoStakingInstructions(stakingInstructions));
444-
const { depositSol } = stakingInstructions;
446+
const { depositSol, hasAtaInit } = stakingInstructions;
445447
stakingActivate = {
446448
type: InstructionBuilderTypes.StakingActivate,
447449
params: {
@@ -456,6 +458,7 @@ function parseStakingActivateInstructions(
456458
poolMint: depositSol.poolMint.toString(),
457459
reserveStake: depositSol.reserveStake.toString(),
458460
},
461+
createAssociatedTokenAccount: !!hasAtaInit,
459462
},
460463
},
461464
};

‎modules/sdk-coin-sol/src/lib/jitoStakePoolOperations.ts‎

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,8 @@ export type DepositSolStakePoolData = Pick<StakePoolData, 'poolMint' | 'reserveS
115115
*/
116116
export function depositSolInstructions(
117117
params: DepositSolInstructionsParams,
118-
stakePool: DepositSolStakePoolData
118+
stakePool: DepositSolStakePoolData,
119+
createAssociatedTokenAccount: boolean
119120
): TransactionInstruction[] {
120121
const { stakePoolAddress, from, lamports } = params;
121122
const poolMint = new PublicKey(stakePool.poolMint);
@@ -124,11 +125,15 @@ export function depositSolInstructions(
124125

125126
// findWithdrawAuthorityProgramAddress
126127
const withdrawAuthority = findWithdrawAuthorityProgramAddressSync(STAKE_POOL_PROGRAM_ID, stakePoolAddress);
127-
128128
const associatedAddress = getAssociatedTokenAddressSync(poolMint, from);
129129

130-
return [
131-
createAssociatedTokenAccountInstruction(from, associatedAddress, from, poolMint),
130+
const instructions: TransactionInstruction[] = [];
131+
132+
if (createAssociatedTokenAccount) {
133+
instructions.push(createAssociatedTokenAccountInstruction(from, associatedAddress, from, poolMint));
134+
}
135+
136+
instructions.push(
132137
StakePoolInstruction.depositSol({
133138
stakePool: stakePoolAddress,
134139
reserveStake,
@@ -139,8 +144,10 @@ export function depositSolInstructions(
139144
poolMint,
140145
lamports: Number(lamports),
141146
withdrawAuthority,
142-
}),
143-
];
147+
})
148+
);
149+
150+
return instructions;
144151
}
145152

146153
function parseKey(key: AccountMeta, template: { isSigner: boolean; isWritable: boolean }): PublicKey {

‎modules/sdk-coin-sol/src/lib/solInstructionFactory.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -296,7 +296,8 @@ function stakingInitializeInstruction(data: StakingActivate): TransactionInstruc
296296
from: fromPubkey,
297297
lamports: BigInt(amount),
298298
},
299-
extraParams.stakePoolData
299+
extraParams.stakePoolData,
300+
!!extraParams.createAssociatedTokenAccount
300301
);
301302
tx.add(...instructions);
302303
break;

‎modules/sdk-coin-sol/src/lib/utils.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ import {
5151
walletInitInstructionIndexes,
5252
jitoStakingActivateInstructionsIndexes,
5353
jitoStakingDeactivateInstructionsIndexes,
54+
jitoStakingActivateWithATAInstructionsIndexes,
5455
} from './constants';
5556
import { ValidInstructionTypes } from './iface';
5657
import { STAKE_POOL_INSTRUCTION_LAYOUTS, STAKE_POOL_PROGRAM_ID } from '@solana/spl-stake-pool';
@@ -327,6 +328,7 @@ export function getTransactionType(transaction: SolTransaction): TransactionType
327328
} else if (
328329
matchTransactionTypeByInstructionsOrder(instructions, marinadeStakingActivateInstructionsIndexes) ||
329330
matchTransactionTypeByInstructionsOrder(instructions, jitoStakingActivateInstructionsIndexes) ||
331+
matchTransactionTypeByInstructionsOrder(instructions, jitoStakingActivateWithATAInstructionsIndexes) ||
330332
matchTransactionTypeByInstructionsOrder(instructions, stakingActivateInstructionsIndexes)
331333
) {
332334
return TransactionType.StakingActivate;

0 commit comments

Comments
 (0)