SPL Stake Pool extract yield to stake account. Test complete

packages/sdks/common/src/types/spl_beam.ts

@@ -669,7 +669,10 @@ export type SplBeam = {
         {
           "name": "newStakeAccount",
           "isMut": true,
-          "isSigner": false
+          "isSigner": false,
+          "docs": [
+            "The uninitialized new stake account. Will be initialised by CPI to the SPL StakePool program."
+          ]
         },
         {
           "name": "vaultAuthority",
@@ -711,6 +714,11 @@ export type SplBeam = {
           "isMut": false,
           "isSigner": false
         },
+        {
+          "name": "sysvarStakeHistory",
+          "isMut": false,
+          "isSigner": false
+        },
         {
           "name": "sunriseProgram",
           "isMut": false,
@@ -815,6 +823,16 @@ export type SplBeam = {
       "code": 6002,
       "name": "Unimplemented",
       "msg": "This feature is unimplemented for this beam"
+    },
+    {
+      "code": 6003,
+      "name": "YieldStakeAccountNotCooledDown",
+      "msg": "The yield stake account cannot yet be claimed"
+    },
+    {
+      "code": 6004,
+      "name": "InsufficientYieldToExtract",
+      "msg": "The yield being extracted is insufficient to cover the rent of the stake account"
     }
   ]
 };
@@ -1490,7 +1508,10 @@ export const IDL: SplBeam = {
         {
           "name": "newStakeAccount",
           "isMut": true,
-          "isSigner": false
+          "isSigner": false,
+          "docs": [
+            "The uninitialized new stake account. Will be initialised by CPI to the SPL StakePool program."
+          ]
         },
         {
           "name": "vaultAuthority",
@@ -1532,6 +1553,11 @@ export const IDL: SplBeam = {
           "isMut": false,
           "isSigner": false
         },
+        {
+          "name": "sysvarStakeHistory",
+          "isMut": false,
+          "isSigner": false
+        },
         {
           "name": "sunriseProgram",
           "isMut": false,
@@ -1636,6 +1662,16 @@ export const IDL: SplBeam = {
       "code": 6002,
       "name": "Unimplemented",
       "msg": "This feature is unimplemented for this beam"
+    },
+    {
+      "code": 6003,
+      "name": "YieldStakeAccountNotCooledDown",
+      "msg": "The yield stake account cannot yet be claimed"
+    },
+    {
+      "code": 6004,
+      "name": "InsufficientYieldToExtract",
+      "msg": "The yield being extracted is insufficient to cover the rent of the stake account"
     }
   ]
 };

packages/sdks/spl/src/index.ts

@@ -449,6 +449,7 @@ export class SplClient extends BeamInterface<SplBeam.SplBeam, StateAccount> {
       managerFeeAccount: this.spl.stakePoolState.managerFeeAccount,
       sysvarClock: SYSVAR_CLOCK_PUBKEY,
       nativeStakeProgram: StakeProgram.programId,
+      sysvarStakeHistory: SYSVAR_STAKE_HISTORY_PUBKEY,
       sunriseProgram: this.sunrise.program.programId,
       splStakePoolProgram: SPL_STAKE_POOL_PROGRAM_ID,
       systemProgram: SystemProgram.programId,
@@ -462,6 +463,30 @@ export class SplClient extends BeamInterface<SplBeam.SplBeam, StateAccount> {
     return new Transaction().add(instruction);
   }
 
+  public async claimExtractedYieldStakeAcccount(): Promise<Transaction> {
+    const [stakeAccount] = Utils.deriveExtractYieldStakeAccount(
+      this.program.programId,
+      this.stateAddress,
+    );
+    const accounts = {
+      state: this.stateAddress,
+      sunriseState: this.state.sunriseState,
+      yieldAccount: this.sunrise.state.yieldAccount,
+      stakeAccount,
+      vaultAuthority: this.vaultAuthority[0],
+      sysvarClock: SYSVAR_CLOCK_PUBKEY,
+      sysvarStakeHistory: SYSVAR_STAKE_HISTORY_PUBKEY,
+      nativeStakeProgram: StakeProgram.programId,
+      systemProgram: SystemProgram.programId,
+    };
+    const instruction = await this.program.methods
+      .claimExtractedYieldStakeAccount()
+      .accounts(accounts)
+      .instruction();
+
+    return new Transaction().add(instruction);
+  }
+
   /**
    * Return a transaction to redeem a ticket received from ordering a withdrawal.
    * This is not a supported feature for SPL beams and will throw an error.

packages/tests/src/functional/beams/spl-stake-pool.test.ts

@@ -6,13 +6,12 @@ import { SplClient } from "@sunrisestake/beams-spl";
 import { Keypair, LAMPORTS_PER_SOL, PublicKey } from "@solana/web3.js";
 import BN from "bn.js";
 import {
-  expectStakeAccountBalance,
+  expectSolBalance,
   expectTokenBalance,
   fund,
   registerSunriseState,
   sendAndConfirmTransaction,
   tokenAccountBalance,
-  waitForNextEpoch,
 } from "../../utils.js";
 import { provider, staker, stakerIdentity } from "../setup.js";
 import { expect } from "chai";
@@ -218,35 +217,15 @@ describe("SPL stake pool beam", () => {
     const expectedFee = burnAmount
       .mul(beamClient.spl.stakePoolState.stakeWithdrawalFee.numerator)
       .div(beamClient.spl.stakePoolState.stakeWithdrawalFee.denominator);
-    const expectedStakeAmount = burnAmount.sub(expectedFee);
+    const expectedExtractedYield = burnAmount.sub(expectedFee);
 
-    // there is no yield yet, but we have created a stake account for the yield
-    await expectStakeAccountBalance(
-      beamClient.provider,
-      beamClient.yieldStakeAccount,
-      expectedStakeAmount,
-      1,
-    );
-  });
-
-  it.skip("can claim stake account into yield account after cooldown", async () => {
-    // wait an epoch for the stake account to cool down
-    await waitForNextEpoch(beamClient.provider);
-
-    // TODO enable
-    // await sendAndConfirmTransaction(
-    //   // anyone can extract yield to the yield account, but let's use the staker provider (rather than the admin provider) for this test
-    //   // to show that it doesn't have to be an admin
-    //   stakerIdentity,
-    //   // TODO move this and extract yield into "management" object to make it clear that these
-    //   // do not need to be executed by end-users
-    //   await beamClient.claimExtractedYieldStakeAcccount(),
-    // );
-
-    await expectTokenBalance(
+    await expectSolBalance(
       beamClient.provider,
       beamClient.sunrise.state.yieldAccount,
-      burnAmount,
+      expectedExtractedYield,
+      // the calculation appears to be slightly inaccurate at present, but in our favour,
+      // so we can leave this as a low priority TODO to improve the accuracy
+      3000,
     );
   });
 });

packages/tests/src/utils.ts

@@ -109,8 +109,15 @@ export const expectStakerSolBalance = async (
   provider: AnchorProvider,
   expectedAmount: number | BN,
   tolerance = 0, // Allow for a tolerance as the balance depends on the fees which are unstable at the beginning of a test validator
+) => expectSolBalance(provider, provider.publicKey, expectedAmount, tolerance);
+
+export const expectSolBalance = async (
+  provider: AnchorProvider,
+  address = provider.publicKey,
+  expectedAmount: number | BN,
+  tolerance = 0, // Allow for a tolerance as the balance depends on the fees which are unstable at the beginning of a test validator
 ) => {
-  const actualAmount = await solBalance(provider);
+  const actualAmount = await solBalance(provider, address);
   expectAmount(actualAmount, expectedAmount, tolerance);
 };
 
@@ -140,10 +147,14 @@ export const waitForNextEpoch = async (provider: AnchorProvider) => {
   const startSlot = startingEpoch.slotIndex;
   let subscriptionId = 0;
 
+  log("Waiting for epoch", nextEpoch);
+
   await new Promise((resolve) => {
     subscriptionId = provider.connection.onSlotChange((slotInfo) => {
+      log("slot", slotInfo.slot, "startSlot", startSlot);
       if (slotInfo.slot % SLOTS_IN_EPOCH === 1 && slotInfo.slot > startSlot) {
         void provider.connection.getEpochInfo().then((currentEpoch) => {
+          log("currentEpoch", currentEpoch);
           if (currentEpoch.epoch === nextEpoch) {
             resolve(slotInfo.slot);
           }

programs/spl-beam/src/cpi_interface/spl.rs

@@ -1,7 +1,8 @@
+use crate::constants::STAKE_ACCOUNT_SIZE;
 use crate::cpi_interface::stake_pool::StakePool;
 use crate::seeds::*;
 use crate::state::State;
-use crate::{ExtractYield, WithdrawStake};
+use crate::{ExtractYield, SplBeamError, WithdrawStake};
 use anchor_lang::{
     prelude::*,
     solana_program::program::{invoke, invoke_signed},
@@ -222,9 +223,22 @@ pub fn extract_stake(accounts: &ExtractStakeAccount, lamports: u64) -> Result<()
     let state_address = accounts.state.key();
     let seeds = &[state_address.as_ref(), VAULT_AUTHORITY, bump][..];
 
+    let stake_account_rent = Rent::get()?.minimum_balance(STAKE_ACCOUNT_SIZE);
+    let mut total_extractable_lamports = lamports.saturating_sub(stake_account_rent);
+
+    if total_extractable_lamports > accounts.stake_to_split.lamports() {
+        total_extractable_lamports = accounts.stake_to_split.lamports();
+        msg!("Limiting the extraction to the amount of lamports in the reserve stake account of the pool: {}", total_extractable_lamports);
+    }
+    if total_extractable_lamports == 0 {
+        return Err(SplBeamError::InsufficientYieldToExtract.into());
+    }
+
     let pool = &accounts.stake_pool;
-    let pool_tokens =
-        crate::utils::pool_tokens_from_lamports(&pool.clone().into_inner(), lamports)?;
+    let pool_tokens = crate::utils::pool_tokens_from_lamports(
+        &pool.clone().into_inner(),
+        total_extractable_lamports,
+    )?;
 
     invoke_signed(
         &spl_stake_pool::instruction::withdraw_stake(

programs/spl-beam/src/cpi_interface/stake_account.rs

@@ -1,5 +1,12 @@
-use anchor_lang::prelude::Pubkey;
+use crate::seeds::VAULT_AUTHORITY;
+use crate::state::State;
+use crate::ExtractYield;
+use anchor_lang::prelude::*;
+use anchor_lang::solana_program::clock::Epoch;
+use anchor_lang::solana_program::program::invoke_signed;
 use anchor_lang::solana_program::stake::program::ID;
+use anchor_lang::Key;
+use anchor_spl::token::spl_token::solana_program;
 use borsh::BorshDeserialize;
 use spl_stake_pool::solana_program::stake::state::StakeStateV2;
 use std::ops::Deref;
@@ -9,19 +16,19 @@ use std::ops::Deref;
 #[derive(Clone)]
 pub struct StakeAccount(StakeStateV2);
 
-impl anchor_lang::AccountDeserialize for StakeAccount {
-    fn try_deserialize(buf: &mut &[u8]) -> anchor_lang::Result<Self> {
+impl AccountDeserialize for StakeAccount {
+    fn try_deserialize(buf: &mut &[u8]) -> Result<Self> {
         Self::try_deserialize_unchecked(buf)
     }
 
-    fn try_deserialize_unchecked(buf: &mut &[u8]) -> anchor_lang::Result<Self> {
+    fn try_deserialize_unchecked(buf: &mut &[u8]) -> Result<Self> {
         StakeStateV2::deserialize(buf).map(Self).map_err(Into::into)
     }
 }
 
-impl anchor_lang::AccountSerialize for StakeAccount {}
+impl AccountSerialize for StakeAccount {}
 
-impl anchor_lang::Owner for StakeAccount {
+impl Owner for StakeAccount {
     fn owner() -> Pubkey {
         ID
     }
@@ -34,3 +41,70 @@ impl Deref for StakeAccount {
         &self.0
     }
 }
+
+impl StakeAccount {
+    pub fn can_be_withdrawn(&self, current_epoch: &Epoch) -> bool {
+        match self.0 {
+            StakeStateV2::Stake(_, stake, _) => {
+                stake.delegation.deactivation_epoch <= *current_epoch
+            }
+            StakeStateV2::Initialized(_) => true,
+            _ => false,
+        }
+    }
+}
+
+pub struct ClaimStakeAccount<'info> {
+    pub state: Box<Account<'info, State>>,
+    pub stake_account: Account<'info, StakeAccount>,
+    pub withdrawer: AccountInfo<'info>,
+    pub to: AccountInfo<'info>,
+    pub native_stake_program: AccountInfo<'info>,
+    pub sysvar_clock: AccountInfo<'info>,
+    pub sysvar_stake_history: AccountInfo<'info>,
+}
+impl<'a> From<ExtractYield<'a>> for ClaimStakeAccount<'a> {
+    /// Convert the ExtractYield beam instruction accounts to the ClaimStakeAccount accounts
+    fn from(extract_yield: ExtractYield<'a>) -> Self {
+        Self {
+            state: extract_yield.state,
+            stake_account: extract_yield.new_stake_account,
+            withdrawer: extract_yield.vault_authority.to_account_info(),
+            to: extract_yield.yield_account.to_account_info(),
+            native_stake_program: extract_yield.native_stake_program.to_account_info(),
+            sysvar_clock: extract_yield.sysvar_clock.to_account_info(),
+            sysvar_stake_history: extract_yield.sysvar_stake_history.to_account_info(),
+        }
+    }
+}
+impl<'a> From<&ExtractYield<'a>> for ClaimStakeAccount<'a> {
+    fn from(extract_yield: &ExtractYield<'a>) -> Self {
+        extract_yield.to_owned().into()
+    }
+}
+
+pub fn claim_stake_account(accounts: &ClaimStakeAccount, lamports: u64) -> Result<()> {
+    let bump = &[accounts.state.vault_authority_bump][..];
+    let state_address = accounts.state.key();
+    let seeds = &[state_address.as_ref(), VAULT_AUTHORITY, bump][..];
+
+    invoke_signed(
+        &solana_program::stake::instruction::withdraw(
+            &accounts.stake_account.key(),
+            accounts.withdrawer.key,
+            accounts.to.key,
+            lamports,
+            None,
+        ),
+        &[
+            accounts.stake_account.to_account_info(),
+            accounts.to.clone(),
+            accounts.sysvar_clock.clone(),
+            accounts.sysvar_stake_history.clone(),
+            accounts.withdrawer.clone(),
+        ],
+        &[seeds],
+    )?;
+
+    Ok(())
+}