feat(sdk-coin-sol): conditionally create ATA for jito staking

examples/ts/sol/stake-jito.ts

@@ -3,106 +3,110 @@
  *
  * Copyright 2025, BitGo, Inc.  All Rights Reserved.
  */
-import { BitGoAPI } from '@bitgo/sdk-api'
-import { TransactionBuilderFactory, Tsol } from '@bitgo/sdk-coin-sol'
-import { coins } from '@bitgo/statics'
-import { Connection, PublicKey, clusterApiUrl, Transaction, Keypair, LAMPORTS_PER_SOL } from "@solana/web3.js"
-import { getStakePoolAccount, updateStakePool } from '@solana/spl-stake-pool'
+import { SolStakingTypeEnum } from '@bitgo/public-types';
+import { BitGoAPI } from '@bitgo/sdk-api';
+import { TransactionBuilderFactory, Tsol } from '@bitgo/sdk-coin-sol';
+import { coins } from '@bitgo/statics';
+import { Connection, PublicKey, clusterApiUrl, Transaction, Keypair, LAMPORTS_PER_SOL } from '@solana/web3.js';
+import { getStakePoolAccount, updateStakePool } from '@solana/spl-stake-pool';
+import { getAssociatedTokenAddressSync } from '@solana/spl-token';
 import * as bs58 from 'bs58';
 
-require('dotenv').config({ path: '../../.env' })
+require('dotenv').config({ path: '../../.env' });
 
-const AMOUNT_LAMPORTS = 1000
-const JITO_STAKE_POOL_ADDRESS = 'Jito4APyf642JPZPx3hGc6WWJ8zPKtRbRs4P815Awbb'
-const NETWORK = 'devnet'
+const AMOUNT_LAMPORTS = 1000;
+const JITO_STAKE_POOL_ADDRESS = 'Jito4APyf642JPZPx3hGc6WWJ8zPKtRbRs4P815Awbb';
+const NETWORK = 'devnet';
 
 const bitgo = new BitGoAPI({
   accessToken: process.env.TESTNET_ACCESS_TOKEN,
   env: 'test',
-})
-const coin = coins.get("tsol")
-bitgo.register(coin.name, Tsol.createInstance)
+});
+const coin = coins.get('tsol');
+bitgo.register(coin.name, Tsol.createInstance);
 
 async function main() {
-  const account = getAccount()
-  const connection = new Connection(clusterApiUrl(NETWORK), 'confirmed')
-  const recentBlockhash = await connection.getLatestBlockhash()
-  const stakePoolAccount = await getStakePoolAccount(connection, new PublicKey(JITO_STAKE_POOL_ADDRESS))
-
+  const account = getAccount();
+  const connection = new Connection(clusterApiUrl(NETWORK), 'confirmed');
+  const recentBlockhash = await connection.getLatestBlockhash();
+  const stakePoolAccount = await getStakePoolAccount(connection, new PublicKey(JITO_STAKE_POOL_ADDRESS));
+  const associatedTokenAddress = getAssociatedTokenAddressSync(
+    stakePoolAccount.account.data.poolMint,
+    account.publicKey
+  );
+  const associatedTokenAccountExists = !!(await connection.getAccountInfo(associatedTokenAddress));
 
   // Account should have sufficient balance
-  const accountBalance = await connection.getBalance(account.publicKey)
+  const accountBalance = await connection.getBalance(account.publicKey);
   if (accountBalance < 0.1 * LAMPORTS_PER_SOL) {
-    console.info(`Your account balance is ${accountBalance / LAMPORTS_PER_SOL} SOL, requesting airdrop`)
-    const sig = await connection.requestAirdrop(account.publicKey, 2 * LAMPORTS_PER_SOL)
-    await connection.confirmTransaction(sig)
-    console.info(`Airdrop successful: ${sig}`)
+    console.info(`Your account balance is ${accountBalance / LAMPORTS_PER_SOL} SOL, requesting airdrop`);
+    const sig = await connection.requestAirdrop(account.publicKey, 2 * LAMPORTS_PER_SOL);
+    await connection.confirmTransaction(sig);
+    console.info(`Airdrop successful: ${sig}`);
   }
 
   // Stake pool should be up to date
-  const epochInfo = await connection.getEpochInfo()
+  const epochInfo = await connection.getEpochInfo();
   if (stakePoolAccount.account.data.lastUpdateEpoch.ltn(epochInfo.epoch)) {
-    console.info('Stake pool is out of date.')
-    const usp = await updateStakePool(connection, stakePoolAccount)
-    const tx = new Transaction()
-    tx.add(...usp.updateListInstructions, ...usp.finalInstructions)
-    const signer = Keypair.fromSecretKey(account.secretKeyArray)
-    const sig = await connection.sendTransaction(tx, [signer])
-    await connection.confirmTransaction(sig)
-    console.info(`Stake pool updated: ${sig}`)
+    console.info('Stake pool is out of date.');
+    const usp = await updateStakePool(connection, stakePoolAccount);
+    const tx = new Transaction();
+    tx.add(...usp.updateListInstructions, ...usp.finalInstructions);
+    const signer = Keypair.fromSecretKey(account.secretKeyArray);
+    const sig = await connection.sendTransaction(tx, [signer]);
+    await connection.confirmTransaction(sig);
+    console.info(`Stake pool updated: ${sig}`);
   }
 
   // Use BitGoAPI to build depositSol instruction
-  const txBuilder = new TransactionBuilderFactory(coin).getStakingActivateBuilder()
+  const txBuilder = new TransactionBuilderFactory(coin).getStakingActivateBuilder();
   txBuilder
     .amount(`${AMOUNT_LAMPORTS}`)
     .sender(account.publicKey.toBase58())
     .stakingAddress(JITO_STAKE_POOL_ADDRESS)
     .validator(JITO_STAKE_POOL_ADDRESS)
-    .stakingTypeParams({
-      type: 'JITO',
+    .stakingType(SolStakingTypeEnum.JITO)
+    .extraParams({
       stakePoolData: {
-        managerFeeAccount: stakePoolAccount.account.data.managerFeeAccount.toString(),
-        poolMint: stakePoolAccount.account.data.poolMint.toString(),
-        reserveStake: stakePoolAccount.account.data.toString(),
-      }
+        managerFeeAccount: stakePoolAccount.account.data.managerFeeAccount.toBase58(),
+        poolMint: stakePoolAccount.account.data.poolMint.toBase58(),
+        reserveStake: stakePoolAccount.account.data.reserveStake.toBase58(),
+      },
+      createAssociatedTokenAccount: !associatedTokenAccountExists,
     })
-    .nonce(recentBlockhash.blockhash)
-  txBuilder.sign({ key: account.secretKey })
-  const tx = await txBuilder.build()
-  const serializedTx = tx.toBroadcastFormat()
-  console.info(`Transaction JSON:\n${JSON.stringify(tx.toJson(), undefined, 2)}`)
+    .nonce(recentBlockhash.blockhash);
+  txBuilder.sign({ key: account.secretKey });
+  const tx = await txBuilder.build();
+  const serializedTx = tx.toBroadcastFormat();
+  console.info(`Transaction JSON:\n${JSON.stringify(tx.toJson(), undefined, 2)}`);
 
   // Send transaction
   try {
-    const sig = await connection.sendRawTransaction(Buffer.from(serializedTx, 'base64'))
-    await connection.confirmTransaction(sig)
-    console.log(`${AMOUNT_LAMPORTS / LAMPORTS_PER_SOL} SOL deposited`, sig)
+    const sig = await connection.sendRawTransaction(Buffer.from(serializedTx, 'base64'));
+    await connection.confirmTransaction(sig);
+    console.log(`${AMOUNT_LAMPORTS / LAMPORTS_PER_SOL} SOL deposited`, sig);
   } catch (e) {
-    console.log('Error sending transaction')
-    console.error(e)
-    if (e.transactionMessage === 'Transaction simulation failed: Error processing Instruction 0: Provided owner is not allowed') {
-      console.error('If you successfully staked JitoSOL once, you cannot stake again.')
-    }
+    console.log('Error sending transaction');
+    console.error(e);
   }
 }
 
 const getAccount = () => {
-  const publicKey = process.env.ACCOUNT_PUBLIC_KEY
-  const secretKey = process.env.ACCOUNT_SECRET_KEY
+  const publicKey = process.env.ACCOUNT_PUBLIC_KEY;
+  const secretKey = process.env.ACCOUNT_SECRET_KEY;
   if (publicKey === undefined || secretKey === undefined) {
-    const { publicKey, secretKey } = Keypair.generate()
-    console.log('# Here is a new account to save into your .env file.')
-    console.log(`ACCOUNT_PUBLIC_KEY=${publicKey.toBase58()}`)
-    console.log(`ACCOUNT_SECRET_KEY=${bs58.encode(secretKey)}`)
-    throw new Error("Missing account information")
+    const { publicKey, secretKey } = Keypair.generate();
+    console.log('# Here is a new account to save into your .env file.');
+    console.log(`ACCOUNT_PUBLIC_KEY=${publicKey.toBase58()}`);
+    console.log(`ACCOUNT_SECRET_KEY=${bs58.encode(secretKey)}`);
+    throw new Error('Missing account information');
   }
 
   return {
     publicKey: new PublicKey(publicKey),
     secretKey,
     secretKeyArray: new Uint8Array(bs58.decode(secretKey)),
-  }
-}
+  };
+};
 
-main().catch((e) => console.error(e))
+main().catch((e) => console.error(e));

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

@@ -115,6 +115,11 @@ export const marinadeStakingActivateInstructionsIndexes = {
 
 /** Const to check the order of the Jito Staking Activate instructions when decode */
 export const jitoStakingActivateInstructionsIndexes = {
+  DepositSol: 0,
+} as const;
+
+/** Const to check the order of the Jito Staking Activate instructions when decode */
+export const jitoStakingActivateWithATAInstructionsIndexes = {
   InitializeAssociatedTokenAccount: 0,
   DepositSol: 1,
 } as const;

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

@@ -123,6 +123,7 @@ export interface Approve {
 
 export interface JitoStakingActivateParams {
   stakePoolData: DepositSolStakePoolData;
+  createAssociatedTokenAccount?: boolean;
 }
 
 export type StakingActivateExtraParams = JitoStakingActivateParams;

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

@@ -330,6 +330,7 @@ type StakingInstructions = {
   create?: CreateAccountParams;
   initialize?: InitializeStakeParams;
   delegate?: DelegateStakeParams;
+  hasAtaInit?: boolean;
 };
 
 type JitoStakingInstructions = StakingInstructions & {
@@ -419,6 +420,7 @@ function parseStakingActivateInstructions(
         break;
 
       case ValidInstructionTypesEnum.InitializeAssociatedTokenAccount:
+        stakingInstructions.hasAtaInit = true;
         instructionData.push({
           type: InstructionBuilderTypes.CreateAssociatedTokenAccount,
           params: {
@@ -441,7 +443,7 @@ function parseStakingActivateInstructions(
   switch (stakingType) {
     case SolStakingTypeEnum.JITO: {
       assert(isJitoStakingInstructions(stakingInstructions));
-      const { depositSol } = stakingInstructions;
+      const { depositSol, hasAtaInit } = stakingInstructions;
       stakingActivate = {
         type: InstructionBuilderTypes.StakingActivate,
         params: {
@@ -456,6 +458,7 @@ function parseStakingActivateInstructions(
               poolMint: depositSol.poolMint.toString(),
               reserveStake: depositSol.reserveStake.toString(),
             },
+            createAssociatedTokenAccount: !!hasAtaInit,
           },
         },
       };

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

@@ -115,7 +115,8 @@ export type DepositSolStakePoolData = Pick<StakePoolData, 'poolMint' | 'reserveS
  */
 export function depositSolInstructions(
   params: DepositSolInstructionsParams,
-  stakePool: DepositSolStakePoolData
+  stakePool: DepositSolStakePoolData,
+  createAssociatedTokenAccount: boolean
 ): TransactionInstruction[] {
   const { stakePoolAddress, from, lamports } = params;
   const poolMint = new PublicKey(stakePool.poolMint);
@@ -124,11 +125,15 @@ export function depositSolInstructions(
 
   // findWithdrawAuthorityProgramAddress
   const withdrawAuthority = findWithdrawAuthorityProgramAddressSync(STAKE_POOL_PROGRAM_ID, stakePoolAddress);
-
   const associatedAddress = getAssociatedTokenAddressSync(poolMint, from);
 
-  return [
-    createAssociatedTokenAccountInstruction(from, associatedAddress, from, poolMint),
+  const instructions: TransactionInstruction[] = [];
+
+  if (createAssociatedTokenAccount) {
+    instructions.push(createAssociatedTokenAccountInstruction(from, associatedAddress, from, poolMint));
+  }
+
+  instructions.push(
     StakePoolInstruction.depositSol({
       stakePool: stakePoolAddress,
       reserveStake,
@@ -139,8 +144,10 @@ export function depositSolInstructions(
       poolMint,
       lamports: Number(lamports),
       withdrawAuthority,
-    }),
-  ];
+    })
+  );
+
+  return instructions;
 }
 
 function parseKey(key: AccountMeta, template: { isSigner: boolean; isWritable: boolean }): PublicKey {

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

@@ -296,7 +296,8 @@ function stakingInitializeInstruction(data: StakingActivate): TransactionInstruc
           from: fromPubkey,
           lamports: BigInt(amount),
         },
-        extraParams.stakePoolData
+        extraParams.stakePoolData,
+        !!extraParams.createAssociatedTokenAccount
       );
       tx.add(...instructions);
       break;

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

@@ -51,6 +51,7 @@ import {
   walletInitInstructionIndexes,
   jitoStakingActivateInstructionsIndexes,
   jitoStakingDeactivateInstructionsIndexes,
+  jitoStakingActivateWithATAInstructionsIndexes,
 } from './constants';
 import { ValidInstructionTypes } from './iface';
 import { STAKE_POOL_INSTRUCTION_LAYOUTS, STAKE_POOL_PROGRAM_ID } from '@solana/spl-stake-pool';
@@ -327,6 +328,7 @@ export function getTransactionType(transaction: SolTransaction): TransactionType
   } else if (
     matchTransactionTypeByInstructionsOrder(instructions, marinadeStakingActivateInstructionsIndexes) ||
     matchTransactionTypeByInstructionsOrder(instructions, jitoStakingActivateInstructionsIndexes) ||
+    matchTransactionTypeByInstructionsOrder(instructions, jitoStakingActivateWithATAInstructionsIndexes) ||
     matchTransactionTypeByInstructionsOrder(instructions, stakingActivateInstructionsIndexes)
   ) {
     return TransactionType.StakingActivate;