@ignitionfi/fogo-stake-pool 1.0.2 → 1.0.3

package/dist/instructions.d.ts

@@ -178,8 +178,8 @@ 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;
+    /** Reserve stake account for rent funding */
+    reserveStake: PublicKey;
     poolTokensIn: number;
     minimumLamportsOut: number;
     /** Seed used to derive the user stake PDA */
@@ -344,7 +344,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.0.3",
   "description": "Fogo Stake Pool SDK",
   "contributors": [
     "Anza Maintainers <maintainers@anza.xyz>",

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

@@ -47,6 +47,7 @@ import {
   lamportsToSol,
   newStakeAccount,
   prepareWithdrawAccounts,
+  PrepareWithdrawPrefetchedData,
   solToLamports,
   ValidatorAccount,
 } from './utils'
@@ -1034,41 +1035,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 +1104,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 +1149,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 +1202,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 +1217,7 @@ export async function withdrawStakeWithSession(
         poolMint: stakePool.poolMint,
         tokenProgramId: stakePool.tokenProgramId,
         programSigner,
-        payer,
+        reserveStake: stakePool.reserveStake,
         poolTokensIn: withdrawAccount.poolAmount.toNumber(),
         minimumLamportsOut,
         userStakeSeed,
@@ -1197,6 +1230,7 @@ export async function withdrawStakeWithSession(
     instructions,
     stakeAccountPubkeys,
     userStakeSeeds,
+    remainingPoolTokens: partialRemainingAmount ? lamportsToSol(partialRemainingAmount) : 0,
   }
 }
 

package/src/instructions.ts

@@ -426,8 +426,8 @@ 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
+  /** Reserve stake account for rent funding */
+  reserveStake: PublicKey
   poolTokensIn: number
   minimumLamportsOut: number
   /** Seed used to derive the user stake PDA */
@@ -1222,7 +1222,7 @@ export class StakePoolInstruction {
 
   /**
    * 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 {
     const type = STAKE_POOL_INSTRUCTION_LAYOUTS.WithdrawStakeWithSession
@@ -1238,17 +1238,19 @@ export class StakePoolInstruction {
       { pubkey: params.withdrawAuthority, isSigner: false, isWritable: false },
       { pubkey: params.stakeToSplit, isSigner: false, isWritable: true },
       { pubkey: params.stakeToReceive, isSigner: false, isWritable: true },
-      { pubkey: params.sessionSigner, isSigner: true, isWritable: false }, // user_stake_authority_info (signer_or_session)
-      { pubkey: params.sessionSigner, isSigner: false, isWritable: false }, // user_transfer_authority_info (not used in session path)
+      { pubkey: params.sessionSigner, isSigner: true, isWritable: false }, // user_stake_authority_info
+      { pubkey: params.sessionSigner, isSigner: false, isWritable: false }, // user_transfer_authority_info (unused in session path)
       { pubkey: params.burnFromPool, isSigner: false, isWritable: true },
       { pubkey: params.managerFeeAccount, isSigner: false, isWritable: true },
       { pubkey: params.poolMint, isSigner: false, isWritable: true },
       { pubkey: SYSVAR_CLOCK_PUBKEY, isSigner: false, isWritable: false },
       { pubkey: params.tokenProgramId, isSigner: false, isWritable: false },
       { pubkey: StakeProgram.programId, isSigner: false, isWritable: false },
+      // Session-specific accounts
       { pubkey: params.programSigner, isSigner: false, isWritable: false },
       { pubkey: SystemProgram.programId, isSigner: false, isWritable: false },
-      { pubkey: params.payer, isSigner: true, isWritable: true },
+      { pubkey: params.reserveStake, isSigner: false, isWritable: true },
+      { pubkey: SYSVAR_STAKE_HISTORY_PUBKEY, isSigner: false, isWritable: false },
     ]
 
     return new TransactionInstruction({

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
 }