@ignitionfi/fogo-stake-pool 1.0.2 → 1.1.0

package/dist/instructions.d.ts

@@ -1,6 +1,15 @@
 import * as BufferLayout from '@solana/buffer-layout';
 import { PublicKey, TransactionInstruction } from '@solana/web3.js';
+import BN from 'bn.js';
 import { InstructionType } from './utils';
+/**
+ * Type for amounts that can be converted to BN.
+ * Accepts number, bigint, BN, or string representations.
+ *
+ * IMPORTANT: For values > 9,007,199,254,740,991 (Number.MAX_SAFE_INTEGER),
+ * you MUST use bigint, BN, or string to avoid precision loss.
+ */
+export type AmountInput = number | bigint | BN | string;
 /**
  * An enumeration of valid StakePoolInstructionType's
  */
@@ -60,8 +69,8 @@ export type DecreaseValidatorStakeParams = {
     validatorList: PublicKey;
     validatorStake: PublicKey;
     transientStake: PublicKey;
-    lamports: number;
-    transientStakeSeed: number;
+    lamports: AmountInput;
+    transientStakeSeed: AmountInput;
 };
 export interface DecreaseValidatorStakeWithReserveParams extends DecreaseValidatorStakeParams {
     reserveStake: PublicKey;
@@ -69,7 +78,7 @@ export interface DecreaseValidatorStakeWithReserveParams extends DecreaseValidat
 export interface DecreaseAdditionalValidatorStakeParams extends DecreaseValidatorStakeParams {
     reserveStake: PublicKey;
     ephemeralStake: PublicKey;
-    ephemeralStakeSeed: number;
+    ephemeralStakeSeed: AmountInput;
 }
 /**
  * (Staker only) Increase stake on a validator from the reserve account.
@@ -84,12 +93,12 @@ export type IncreaseValidatorStakeParams = {
     transientStake: PublicKey;
     validatorStake: PublicKey;
     validatorVote: PublicKey;
-    lamports: number;
-    transientStakeSeed: number;
+    lamports: AmountInput;
+    transientStakeSeed: AmountInput;
 };
 export interface IncreaseAdditionalValidatorStakeParams extends IncreaseValidatorStakeParams {
     ephemeralStake: PublicKey;
-    ephemeralStakeSeed: number;
+    ephemeralStakeSeed: AmountInput;
 }
 /**
  * Deposits a stake account into the pool in exchange for pool tokens
@@ -123,7 +132,7 @@ export type WithdrawStakeParams = {
     sourcePoolAccount: PublicKey;
     managerFeeAccount: PublicKey;
     poolMint: PublicKey;
-    poolTokens: number;
+    poolTokens: AmountInput;
 };
 /**
  * Withdraw sol instruction params
@@ -139,7 +148,7 @@ export type WithdrawSolParams = {
     solWithdrawAuthority?: PublicKey | undefined;
     managerFeeAccount: PublicKey;
     poolMint: PublicKey;
-    poolTokens: number;
+    poolTokens: AmountInput;
 };
 /**
  * Withdraw WSOL with session instruction params
@@ -158,8 +167,8 @@ export type WithdrawWsolWithSessionParams = {
     wsolMint: PublicKey;
     programSigner: PublicKey;
     userWallet: PublicKey;
-    poolTokensIn: number;
-    minimumLamportsOut: number;
+    poolTokensIn: AmountInput;
+    minimumLamportsOut: AmountInput;
     solWithdrawAuthority?: PublicKey;
 };
 export type WithdrawStakeWithSessionParams = {
@@ -178,12 +187,12 @@ export type WithdrawStakeWithSessionParams = {
     tokenProgramId: PublicKey;
     /** The program signer PDA derived from PROGRAM_SIGNER_SEED */
     programSigner: PublicKey;
-    /** The payer for stake account rent (typically the paymaster) */
-    payer: PublicKey;
-    poolTokensIn: number;
-    minimumLamportsOut: number;
+    /** Reserve stake account for rent funding */
+    reserveStake: PublicKey;
+    poolTokensIn: AmountInput;
+    minimumLamportsOut: AmountInput;
     /** Seed used to derive the user stake PDA */
-    userStakeSeed: number;
+    userStakeSeed: AmountInput;
 };
 export type WithdrawFromStakeAccountWithSessionParams = {
     programId: PublicKey;
@@ -194,9 +203,9 @@ export type WithdrawFromStakeAccountWithSessionParams = {
     /** The session signer (user or session) */
     sessionSigner: PublicKey;
     /** Seed used to derive the user stake PDA */
-    userStakeSeed: number;
-    /** Lamports to withdraw (use Number.MAX_SAFE_INTEGER for full withdrawal) */
-    lamports: number;
+    userStakeSeed: AmountInput;
+    /** Lamports to withdraw (use BigInt(Number.MAX_SAFE_INTEGER) for full withdrawal) */
+    lamports: AmountInput;
 };
 /**
  * Deposit SOL directly into the pool's reserve account. The output is a "pool" token
@@ -213,7 +222,7 @@ export type DepositSolParams = {
     managerFeeAccount: PublicKey;
     referralPoolAccount: PublicKey;
     poolMint: PublicKey;
-    lamports: number;
+    lamports: AmountInput;
 };
 export type CreateTokenMetadataParams = {
     programId?: PublicKey | undefined;
@@ -325,8 +334,8 @@ export declare class StakePoolInstruction {
         tokenProgramId: PublicKey;
         programId: PublicKey;
         userWallet: PublicKey;
-        lamportsIn: number;
-        minimumPoolTokensOut: number;
+        lamportsIn: AmountInput;
+        minimumPoolTokensOut: AmountInput;
         payer?: PublicKey;
     }): TransactionInstruction;
     /**
@@ -344,7 +353,7 @@ export declare class StakePoolInstruction {
     static withdrawWsolWithSession(params: WithdrawWsolWithSessionParams): TransactionInstruction;
     /**
      * Creates a transaction instruction to withdraw stake from a stake pool using a Fogo session.
-     * The stake account is created as a PDA and rent is paid by the payer (typically paymaster).
+     * The stake account is created as a PDA and rent is funded from the reserve.
      */
     static withdrawStakeWithSession(params: WithdrawStakeWithSessionParams): TransactionInstruction;
     /**

package/dist/utils/stake.d.ts

@@ -12,12 +12,29 @@ export declare function getValidatorListAccount(connection: Connection, pubkey:
     };
 }>;
 export interface ValidatorAccount {
-    type: 'preferred' | 'active' | 'transient' | 'reserve';
+    type: 'preferred' | 'active' | 'transient';
     voteAddress?: PublicKey | undefined;
     stakeAddress: PublicKey;
     lamports: BN;
 }
-export declare function prepareWithdrawAccounts(connection: Connection, stakePool: StakePool, stakePoolAddress: PublicKey, amount: BN, compareFn?: (a: ValidatorAccount, b: ValidatorAccount) => number, skipFee?: boolean): Promise<WithdrawAccount[]>;
+export interface PrepareWithdrawResult {
+    withdrawAccounts: WithdrawAccount[];
+    /** Pool tokens that will be withdrawn via delayed unstake */
+    delayedAmount: BN;
+    /** Pool tokens remaining that need instant unstake */
+    remainingAmount: BN;
+}
+/** Pre-fetched data to avoid duplicate RPC calls */
+export interface PrepareWithdrawPrefetchedData {
+    /** Raw validator list account data from getAccountInfo */
+    validatorListData: Buffer | null;
+    /** Rent exemption for stake accounts in lamports */
+    minBalanceForRentExemption: number;
+    /** Minimum stake delegation in lamports */
+    stakeMinimumDelegation: number;
+}
+export declare function prepareWithdrawAccounts(connection: Connection, stakePool: StakePool, stakePoolAddress: PublicKey, amount: BN, compareFn?: (a: ValidatorAccount, b: ValidatorAccount) => number, skipFee?: boolean, allowPartial?: boolean, prefetchedData?: PrepareWithdrawPrefetchedData): Promise<WithdrawAccount[]>;
+export declare function prepareWithdrawAccounts(connection: Connection, stakePool: StakePool, stakePoolAddress: PublicKey, amount: BN, compareFn: ((a: ValidatorAccount, b: ValidatorAccount) => number) | undefined, skipFee: boolean | undefined, allowPartial: true, prefetchedData?: PrepareWithdrawPrefetchedData): Promise<PrepareWithdrawResult>;
 /**
  * Calculate the pool tokens that should be minted for a deposit of `stakeLamports`
  */

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@ignitionfi/fogo-stake-pool",
-  "version": "1.0.2",
+  "version": "1.1.0",
   "description": "Fogo Stake Pool SDK",
   "contributors": [
     "Anza Maintainers <maintainers@anza.xyz>",

package/src/codecs.ts

@@ -1,3 +1,4 @@
+import * as BufferLayout from '@solana/buffer-layout'
 import { PublicKey } from '@solana/web3.js'
 import BN from 'bn.js'
 import { blob, Layout as LayoutCls, offset, seq, struct, u8, u32 } from 'buffer-layout'
@@ -15,18 +16,22 @@ export interface Layout<T> {
   replicate: (name: string) => this
 }
 
-class BNLayout extends LayoutCls<BN> {
-  blob: Layout<Buffer>
+/**
+ * BN-based layout for data decoding using buffer-layout.
+ * Used for decoding on-chain account data (StakePool, ValidatorList, etc.)
+ */
+class BNDataLayout extends LayoutCls<BN> {
+  blobLayout: Layout<Buffer>
   signed: boolean
 
   constructor(span: number, signed: boolean, property?: string) {
     super(span, property)
-    this.blob = blob(span)
+    this.blobLayout = blob(span)
     this.signed = signed
   }
 
   decode(b: Buffer, offset = 0) {
-    const num = new BN(this.blob.decode(b, offset), 10, 'le')
+    const num = new BN(this.blobLayout.decode(b, offset), 10, 'le')
     if (this.signed) {
       return num.fromTwos(this.span * 8).clone()
     }
@@ -37,12 +42,60 @@ class BNLayout extends LayoutCls<BN> {
     if (this.signed) {
       src = src.toTwos(this.span * 8)
     }
-    return this.blob.encode(src.toArrayLike(Buffer, 'le', this.span), b, offset)
+    return this.blobLayout.encode(src.toArrayLike(Buffer, 'le', this.span), b, offset)
   }
 }
 
+/**
+ * Creates a u64 layout for data decoding (account layouts).
+ * Used in StakePoolLayout, ValidatorListLayout, etc.
+ */
 export function u64(property?: string): Layout<BN> {
-  return new BNLayout(8, false, property)
+  return new BNDataLayout(8, false, property)
+}
+
+/**
+ * BN-based layout for 64-bit unsigned integers using @solana/buffer-layout.
+ * Used for encoding instruction data with support for values > MAX_SAFE_INTEGER.
+ */
+class BNInstructionLayout extends BufferLayout.Layout<BN> {
+  blobLayout: BufferLayout.Blob
+  signed: boolean
+
+  constructor(span: number, signed: boolean, property?: string) {
+    super(span, property)
+    this.blobLayout = BufferLayout.blob(span)
+    this.signed = signed
+  }
+
+  decode(b: Uint8Array, offset = 0): BN {
+    const num = new BN(this.blobLayout.decode(b, offset), 10, 'le')
+    if (this.signed) {
+      return num.fromTwos(this.span * 8).clone()
+    }
+    return num
+  }
+
+  encode(src: BN, b: Uint8Array, offset = 0): number {
+    if (this.signed) {
+      src = src.toTwos(this.span * 8)
+    }
+    return this.blobLayout.encode(src.toArrayLike(Buffer, 'le', this.span), b, offset)
+  }
+
+  getSpan(_b?: Uint8Array, _offset?: number): number {
+    return this.span
+  }
+}
+
+/**
+ * Creates a u64 layout for instruction encoding.
+ * Properly handles BN values larger than Number.MAX_SAFE_INTEGER.
+ * Compatible with @solana/buffer-layout.struct().
+ */
+// eslint-disable-next-line ts/no-explicit-any
+export function u64Instruction(property?: string): any {
+  return new BNInstructionLayout(8, false, property)
 }
 
 class WrappedLayout<T, U> extends LayoutCls<U> {

package/src/constants.ts

@@ -1,5 +1,5 @@
 import { Buffer } from 'node:buffer'
-import { LAMPORTS_PER_SOL, PublicKey } from '@solana/web3.js'
+import { PublicKey } from '@solana/web3.js'
 
 // Public key that identifies the metadata program.
 export const METADATA_PROGRAM_ID = new PublicKey('metaqbxxUerdq28cj1RbAWkYQm3ybzjb6a8bt518x1s')

package/src/index.ts

@@ -24,7 +24,7 @@ import {
   MINIMUM_ACTIVE_STAKE,
   STAKE_POOL_PROGRAM_ID,
 } from './constants'
-import { StakePoolInstruction } from './instructions'
+import { AmountInput, StakePoolInstruction } from './instructions'
 import {
   StakeAccount,
   StakePool,
@@ -47,6 +47,7 @@ import {
   lamportsToSol,
   newStakeAccount,
   prepareWithdrawAccounts,
+  PrepareWithdrawPrefetchedData,
   solToLamports,
   ValidatorAccount,
 } from './utils'
@@ -303,8 +304,8 @@ export async function depositWsolWithSession(
   stakePoolAddress: PublicKey,
   signerOrSession: PublicKey,
   userPubkey: PublicKey,
-  lamports: number,
-  minimumPoolTokensOut: number = 0,
+  lamports: AmountInput,
+  minimumPoolTokensOut: AmountInput = 0,
   destinationTokenAccount?: PublicKey,
   referrerTokenAccount?: PublicKey,
   depositAuthority?: PublicKey,
@@ -323,13 +324,22 @@ export async function depositWsolWithSession(
       'confirmed',
     )
     const wsolBalance = tokenAccountInfo
-      ? parseInt(tokenAccountInfo.value.amount)
-      : 0
-
-    if (wsolBalance < lamports) {
+      ? BigInt(tokenAccountInfo.value.amount)
+      : BigInt(0)
+
+    // Convert lamports to BigInt for comparison
+    const lamportsBigInt = typeof lamports === 'bigint'
+      ? lamports
+      : typeof lamports === 'string'
+        ? BigInt(lamports)
+        : BN.isBN(lamports)
+          ? BigInt(lamports.toString())
+          : BigInt(lamports)
+
+    if (wsolBalance < lamportsBigInt) {
       throw new Error(
         `Not enough WSOL to deposit into pool. Maximum deposit amount is ${lamportsToSol(
-          wsolBalance,
+          Number(wsolBalance),
         )} WSOL.`,
       )
     }
@@ -401,13 +411,23 @@ export async function depositSol(
   connection: Connection,
   stakePoolAddress: PublicKey,
   from: PublicKey,
-  lamports: number,
+  lamports: AmountInput,
   destinationTokenAccount?: PublicKey,
   referrerTokenAccount?: PublicKey,
   depositAuthority?: PublicKey,
 ) {
   const fromBalance = await connection.getBalance(from, 'confirmed')
-  if (fromBalance < lamports) {
+
+  // Convert lamports to BigInt for comparison
+  const lamportsBigInt = typeof lamports === 'bigint'
+    ? lamports
+    : typeof lamports === 'string'
+      ? BigInt(lamports)
+      : BN.isBN(lamports)
+        ? BigInt(lamports.toString())
+        : BigInt(lamports)
+
+  if (BigInt(fromBalance) < lamportsBigInt) {
     throw new Error(
       `Not enough SOL to deposit into pool. Maximum deposit amount is ${lamportsToSol(
         fromBalance,
@@ -432,7 +452,7 @@ export async function depositSol(
     SystemProgram.transfer({
       fromPubkey: from,
       toPubkey: userSolTransfer.publicKey,
-      lamports,
+      lamports: lamportsBigInt,
     }),
   )
 
@@ -1034,41 +1054,57 @@ export async function getUserStakeAccounts(
  * Withdraws stake from a stake pool using a Fogo session.
  *
  * The on-chain program creates stake account PDAs. The rent for these accounts
- * is paid by the payer (typically the paymaster), not deducted from the user's withdrawal.
+ * is funded from the reserve stake.
  *
  * @param connection - Solana connection
  * @param stakePoolAddress - The stake pool to withdraw from
  * @param signerOrSession - The session signer public key
  * @param userPubkey - User's wallet (used for PDA derivation and token ownership)
- * @param payer - Payer for stake account rent (typically paymaster)
  * @param amount - Amount of pool tokens to withdraw
  * @param userStakeSeedStart - Starting seed for user stake PDA derivation (default: 0)
  * @param useReserve - Whether to withdraw from reserve (default: false)
  * @param voteAccountAddress - Optional specific validator to withdraw from
  * @param minimumLamportsOut - Minimum lamports to receive (slippage protection)
  * @param validatorComparator - Optional comparator for validator selection
+ * @param allowPartial - If true, returns partial results instead of throwing when not enough stake available
  */
 export async function withdrawStakeWithSession(
   connection: Connection,
   stakePoolAddress: PublicKey,
   signerOrSession: PublicKey,
   userPubkey: PublicKey,
-  payer: PublicKey,
   amount: number,
   userStakeSeedStart: number = 0,
   useReserve = false,
   voteAccountAddress?: PublicKey,
   minimumLamportsOut: number = 0,
   validatorComparator?: (_a: ValidatorAccount, _b: ValidatorAccount) => number,
+  allowPartial = false,
 ) {
-  const stakePoolAccount = await getStakePoolAccount(connection, stakePoolAddress)
   const stakePoolProgramId = getStakePoolProgramId(connection.rpcEndpoint)
+
+  // First fetch: get stake pool to know other account addresses
+  const stakePoolAccount = await getStakePoolAccount(connection, stakePoolAddress)
   const stakePool = stakePoolAccount.account.data
   const poolTokens = solToLamports(amount)
   const poolAmount = new BN(poolTokens)
 
   const poolTokenAccount = getAssociatedTokenAddressSync(stakePool.poolMint, userPubkey)
-  const tokenAccount = await getAccount(connection, poolTokenAccount)
+
+  // Second fetch: get ALL remaining data in parallel
+  const [tokenAccount, stakeAccountRentExemption, validatorListAcc, stakeMinimumDelegationResponse] = await Promise.all([
+    getAccount(connection, poolTokenAccount),
+    connection.getMinimumBalanceForRentExemption(StakeProgram.space),
+    connection.getAccountInfo(stakePool.validatorList),
+    connection.getStakeMinimumDelegation(),
+  ])
+
+  // Pre-fetch data to avoid duplicate RPC calls in prepareWithdrawAccounts
+  const prefetchedData: PrepareWithdrawPrefetchedData = {
+    validatorListData: validatorListAcc?.data ?? null,
+    minBalanceForRentExemption: stakeAccountRentExemption,
+    stakeMinimumDelegation: Number(stakeMinimumDelegationResponse.value),
+  }
 
   if (tokenAccount.amount < poolTokens) {
     throw new Error(
@@ -1087,10 +1123,9 @@ export async function withdrawStakeWithSession(
     stakePoolAddress,
   )
 
-  const stakeAccountRentExemption = await connection.getMinimumBalanceForRentExemption(StakeProgram.space)
-
   // Determine which stake accounts to withdraw from
   const withdrawAccounts: WithdrawAccount[] = []
+  let partialRemainingAmount: BN | undefined
 
   if (useReserve) {
     withdrawAccounts.push({
@@ -1133,16 +1168,33 @@ export async function withdrawStakeWithSession(
     })
   } else {
     // Get the list of accounts to withdraw from automatically
-    withdrawAccounts.push(
-      ...(await prepareWithdrawAccounts(
+    if (allowPartial) {
+      const result = await prepareWithdrawAccounts(
         connection,
         stakePool,
         stakePoolAddress,
         poolAmount,
         validatorComparator,
         poolTokenAccount.equals(stakePool.managerFeeAccount),
-      )),
-    )
+        true,
+        prefetchedData,
+      )
+      withdrawAccounts.push(...result.withdrawAccounts)
+      partialRemainingAmount = result.remainingAmount
+    } else {
+      withdrawAccounts.push(
+        ...(await prepareWithdrawAccounts(
+          connection,
+          stakePool,
+          stakePoolAddress,
+          poolAmount,
+          validatorComparator,
+          poolTokenAccount.equals(stakePool.managerFeeAccount),
+          undefined,
+          prefetchedData,
+        )),
+      )
+    }
   }
 
   const instructions: TransactionInstruction[] = []
@@ -1169,7 +1221,7 @@ export async function withdrawStakeWithSession(
     stakeAccountPubkeys.push(stakeReceiverPubkey)
     userStakeSeeds.push(userStakeSeed)
 
-    // The on-chain program creates the stake account PDA and rent is paid by payer.
+    // The on-chain program creates the stake account PDA and rent is funded from reserve.
     instructions.push(
       StakePoolInstruction.withdrawStakeWithSession({
         programId: stakePoolProgramId,
@@ -1184,7 +1236,7 @@ export async function withdrawStakeWithSession(
         poolMint: stakePool.poolMint,
         tokenProgramId: stakePool.tokenProgramId,
         programSigner,
-        payer,
+        reserveStake: stakePool.reserveStake,
         poolTokensIn: withdrawAccount.poolAmount.toNumber(),
         minimumLamportsOut,
         userStakeSeed,
@@ -1197,6 +1249,7 @@ export async function withdrawStakeWithSession(
     instructions,
     stakeAccountPubkeys,
     userStakeSeeds,
+    remainingPoolTokens: partialRemainingAmount ? lamportsToSol(partialRemainingAmount) : 0,
   }
 }
 
@@ -1329,7 +1382,7 @@ export async function increaseValidatorStake(
   connection: Connection,
   stakePoolAddress: PublicKey,
   validatorVote: PublicKey,
-  lamports: number,
+  lamports: AmountInput,
   ephemeralStakeSeed?: number,
 ) {
   const stakePool = await getStakePoolAccount(connection, stakePoolAddress)
@@ -1427,7 +1480,7 @@ export async function decreaseValidatorStake(
   connection: Connection,
   stakePoolAddress: PublicKey,
   validatorVote: PublicKey,
-  lamports: number,
+  lamports: AmountInput,
   ephemeralStakeSeed?: number,
 ) {
   const stakePool = await getStakePoolAccount(connection, stakePoolAddress)