@ignitionfi/fogo-stake-pool 1.0.2 → 1.1.0

package/src/utils/stake.ts

@@ -38,12 +38,50 @@ export async function getValidatorListAccount(connection: Connection, pubkey: Pu
 }
 
 export interface ValidatorAccount {
-  type: 'preferred' | 'active' | 'transient' | 'reserve'
+  type: 'preferred' | 'active' | 'transient'
   voteAddress?: PublicKey | undefined
   stakeAddress: PublicKey
   lamports: BN
 }
 
+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 async 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 async 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>
 export async function prepareWithdrawAccounts(
   connection: Connection,
   stakePool: StakePool,
@@ -51,26 +89,65 @@ export async function prepareWithdrawAccounts(
   amount: BN,
   compareFn?: (a: ValidatorAccount, b: ValidatorAccount) => number,
   skipFee?: boolean,
-): Promise<WithdrawAccount[]> {
+  allowPartial?: boolean,
+  prefetchedData?: PrepareWithdrawPrefetchedData,
+): Promise<WithdrawAccount[] | PrepareWithdrawResult> {
   const stakePoolProgramId = getStakePoolProgramId(connection.rpcEndpoint)
-  const validatorListAcc = await connection.getAccountInfo(stakePool.validatorList)
-  const validatorList = ValidatorListLayout.decode(validatorListAcc?.data) as ValidatorList
+
+  // Use prefetched data if available, otherwise fetch from RPC
+  let validatorListData: Buffer | null
+  let minBalanceForRentExemption: number
+  let stakeMinimumDelegation: number
+
+  if (prefetchedData) {
+    validatorListData = prefetchedData.validatorListData
+    minBalanceForRentExemption = prefetchedData.minBalanceForRentExemption
+    stakeMinimumDelegation = prefetchedData.stakeMinimumDelegation
+  } else {
+    const [validatorListAcc, rentExemption, stakeMinimumDelegationResponse] = await Promise.all([
+      connection.getAccountInfo(stakePool.validatorList),
+      connection.getMinimumBalanceForRentExemption(StakeProgram.space),
+      connection.getStakeMinimumDelegation(),
+    ])
+    validatorListData = validatorListAcc?.data ?? null
+    minBalanceForRentExemption = rentExemption
+    stakeMinimumDelegation = Number(stakeMinimumDelegationResponse.value)
+  }
+
+  if (!validatorListData) {
+    throw new Error('No staked funds available for delayed unstake. Use instant unstake instead.')
+  }
+
+  const validatorList = ValidatorListLayout.decode(validatorListData) as ValidatorList
 
   if (!validatorList?.validators || validatorList?.validators.length === 0) {
-    throw new Error('No accounts found')
+    throw new Error('No staked funds available for delayed unstake. Use instant unstake instead.')
   }
 
-  const minBalanceForRentExemption = await connection.getMinimumBalanceForRentExemption(
-    StakeProgram.space,
+  // minBalance = rent + max(stake_minimum_delegation, MINIMUM_ACTIVE_STAKE)
+  const minimumDelegation = Math.max(stakeMinimumDelegation, MINIMUM_ACTIVE_STAKE)
+  const minBalance = new BN(minBalanceForRentExemption + minimumDelegation)
+
+  // Threshold for has_active_stake check (ceiling division for lamports_per_pool_token)
+  const lamportsPerPoolToken = stakePool.totalLamports
+    .add(stakePool.poolTokenSupply)
+    .sub(new BN(1))
+    .div(stakePool.poolTokenSupply)
+  const minimumLamportsWithTolerance = minBalance.add(lamportsPerPoolToken)
+
+  const hasActiveStake = validatorList.validators.some(
+    v => v.status === ValidatorStakeInfoStatus.Active
+      && v.activeStakeLamports.gt(minimumLamportsWithTolerance),
   )
-  const minBalance = new BN(minBalanceForRentExemption + MINIMUM_ACTIVE_STAKE)
+  const hasTransientStake = validatorList.validators.some(
+    v => v.status === ValidatorStakeInfoStatus.Active
+      && v.transientStakeLamports.gt(minimumLamportsWithTolerance),
+  )
+
+  // ValidatorRemoval mode: no validator above threshold
+  const isValidatorRemovalMode = !hasActiveStake && !hasTransientStake
 
-  // First, collect all stake account addresses we need to check
-  const accountsToFetch: Array<{
-    type: 'preferred' | 'active' | 'transient'
-    voteAddress: PublicKey
-    stakeAddress: PublicKey
-  }> = []
+  let accounts: ValidatorAccount[] = []
 
   for (const validator of validatorList.validators) {
     if (validator.status !== ValidatorStakeInfoStatus.Active) {
@@ -83,63 +160,39 @@ export async function prepareWithdrawAccounts(
       stakePoolAddress,
     )
 
-    const isPreferred = stakePool?.preferredWithdrawValidatorVoteAddress?.equals(
-      validator.voteAccountAddress,
-    )
+    // ValidatorRemoval: full balance available; Normal: leave minBalance
+    const availableActiveLamports = isValidatorRemovalMode
+      ? validator.activeStakeLamports
+      : validator.activeStakeLamports.sub(minBalance)
 
-    // Add active stake account if validator list indicates it has stake
-    if (validator.activeStakeLamports.gt(new BN(0))) {
-      accountsToFetch.push({
+    if (availableActiveLamports.gt(new BN(0))) {
+      const isPreferred = stakePool?.preferredWithdrawValidatorVoteAddress?.equals(
+        validator.voteAccountAddress,
+      )
+      accounts.push({
         type: isPreferred ? 'preferred' : 'active',
         voteAddress: validator.voteAccountAddress,
         stakeAddress: stakeAccountAddress,
+        lamports: availableActiveLamports,
       })
     }
 
-    // Add transient stake account if validator list indicates it has stake
-    if (validator.transientStakeLamports.gt(new BN(0))) {
+    const availableTransientLamports = isValidatorRemovalMode
+      ? validator.transientStakeLamports
+      : validator.transientStakeLamports.sub(minBalance)
+
+    if (availableTransientLamports.gt(new BN(0))) {
       const transientStakeAccountAddress = await findTransientStakeProgramAddress(
         stakePoolProgramId,
         validator.voteAccountAddress,
         stakePoolAddress,
         validator.transientSeedSuffixStart,
       )
-      accountsToFetch.push({
+      accounts.push({
         type: 'transient',
         voteAddress: validator.voteAccountAddress,
         stakeAddress: transientStakeAccountAddress,
-      })
-    }
-  }
-
-  // Fetch all stake accounts + reserve in one batch call
-  const addressesToFetch = [
-    ...accountsToFetch.map(a => a.stakeAddress),
-    stakePool.reserveStake,
-  ]
-  const accountInfos = await connection.getMultipleAccountsInfo(addressesToFetch)
-
-  // Build accounts list using actual on-chain balances
-  let accounts: ValidatorAccount[] = []
-
-  for (let i = 0; i < accountsToFetch.length; i++) {
-    const { type, voteAddress, stakeAddress } = accountsToFetch[i]
-    const accountInfo = accountInfos[i]
-
-    if (!accountInfo) {
-      continue
-    }
-
-    // Use actual on-chain balance instead of validator list value
-    const actualLamports = new BN(accountInfo.lamports)
-    const availableLamports = actualLamports.sub(minBalance)
-
-    if (availableLamports.gt(new BN(0))) {
-      accounts.push({
-        type,
-        voteAddress,
-        stakeAddress,
-        lamports: availableLamports,
+        lamports: availableTransientLamports,
       })
     }
   }
@@ -147,17 +200,6 @@ export async function prepareWithdrawAccounts(
   // Sort from highest to lowest balance
   accounts = accounts.sort(compareFn || ((a, b) => b.lamports.sub(a.lamports).toNumber()))
 
-  // Add reserve stake using actual balance (last item in batch fetch)
-  const reserveAccountInfo = accountInfos[accountInfos.length - 1]
-  const reserveStakeBalance = new BN((reserveAccountInfo?.lamports ?? 0) - minBalanceForRentExemption)
-  if (reserveStakeBalance.gt(new BN(0))) {
-    accounts.push({
-      type: 'reserve',
-      stakeAddress: stakePool.reserveStake,
-      lamports: reserveStakeBalance,
-    })
-  }
-
   // Prepare the list of accounts to withdraw from
   const withdrawFrom: WithdrawAccount[] = []
   let remainingAmount = new BN(amount)
@@ -168,14 +210,10 @@ export async function prepareWithdrawAccounts(
     denominator: fee.denominator,
   }
 
-  for (const type of ['preferred', 'active', 'transient', 'reserve']) {
+  for (const type of ['preferred', 'active', 'transient']) {
     const filteredAccounts = accounts.filter(a => a.type === type)
 
     for (const { stakeAddress, voteAddress, lamports } of filteredAccounts) {
-      if (lamports.lte(minBalance) && type === 'transient') {
-        continue
-      }
-
       let availableForWithdrawal = calcPoolTokensForDeposit(stakePool, lamports)
 
       if (!skipFee && !inverseFee.numerator.isZero()) {
@@ -184,12 +222,17 @@ export async function prepareWithdrawAccounts(
           .div(inverseFee.numerator)
       }
 
+      // In ValidatorRemoval mode, must withdraw full validator balance (no partial)
+      // Skip if remaining amount is less than full validator balance
+      if (isValidatorRemovalMode && remainingAmount.lt(availableForWithdrawal)) {
+        continue
+      }
+
       const poolAmount = BN.min(availableForWithdrawal, remainingAmount)
       if (poolAmount.lte(new BN(0))) {
         continue
       }
 
-      // Those accounts will be withdrawn completely with `claim` instruction
       withdrawFrom.push({ stakeAddress, voteAddress, poolAmount })
       remainingAmount = remainingAmount.sub(poolAmount)
 
@@ -205,13 +248,28 @@ export async function prepareWithdrawAccounts(
 
   // Not enough stake to withdraw the specified amount
   if (remainingAmount.gt(new BN(0))) {
+    if (allowPartial) {
+      const delayedAmount = amount.sub(remainingAmount)
+      return {
+        withdrawAccounts: withdrawFrom,
+        delayedAmount,
+        remainingAmount,
+      }
+    }
+    const availableAmount = amount.sub(remainingAmount)
     throw new Error(
-      `No stake accounts found in this pool with enough balance to withdraw ${lamportsToSol(
-        amount,
-      )} pool tokens.`,
+      `Not enough staked funds for delayed unstake. Requested ${lamportsToSol(amount)} iFOGO, but only ${lamportsToSol(availableAmount)} available. Use instant unstake for the remaining amount.`,
     )
   }
 
+  if (allowPartial) {
+    return {
+      withdrawAccounts: withdrawFrom,
+      delayedAmount: amount,
+      remainingAmount: new BN(0),
+    }
+  }
+
   return withdrawFrom
 }