Freeze unstaking coins during cooldown; close account on full unstake

.github/workflows/verifiable-build.yml

@@ -5,6 +5,12 @@ on:
   pull_request:
   workflow_dispatch:
 
+# Cancel an in-progress run when a newer one starts for the same ref+event,
+# so pushes don't pile up stale runs competing for runners.
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}-${{ github.event_name }}
+  cancel-in-progress: true
+
 env:
   SOLANA_VERSION: v2.0.25
 
@@ -22,9 +28,15 @@ jobs:
       - name: Generate keypair
         run: solana-keygen new --no-bip39-passphrase
 
+      - name: Install protobuf compiler (solana-program-test build dep)
+        run: sudo apt-get update && sudo apt-get install -y protobuf-compiler
+
       - name: Build program
         run: cargo build-sbf --workspace
 
+      - name: Run Rust tests (unit + legacy-migration integration)
+        run: cargo test -p chiefstaker
+
       - name: Start test validator with upgradeable program
         run: |
           solana config set --url http://localhost:8899 --keypair ~/.config/solana/id.json

README.md

@@ -63,16 +63,16 @@ Rewards can be deposited directly via instruction or sent to the pool PDA (e.g.,
 |---|-------------|-------------|
 | 0 | `InitializePool` | Create a new staking pool for a Token 2022 mint |
 | 1 | `Stake` | Stake tokens into the pool |
-| 2 | `Unstake` | Unstake tokens (direct, when no cooldown) |
+| 2 | `Unstake` | Unstake tokens (direct, when no cooldown); full unstake closes the account |
 | 3 | `ClaimRewards` | Claim accumulated SOL rewards |
 | 4 | `DepositRewards` | Deposit SOL rewards into the pool |
 | 5 | `SyncPool` | Rebase pool math to prevent overflow |
 | 6 | `SyncRewards` | Sync SOL sent directly to the pool PDA |
 | 7 | `UpdatePoolSettings` | Set min stake, lock duration, cooldown (authority only) |
 | 8 | `TransferAuthority` | Transfer or renounce pool authority |
-| 9 | `RequestUnstake` | Start unstake cooldown (tokens keep earning) |
-| 10 | `CompleteUnstake` | Finish unstake after cooldown elapsed |
-| 11 | `CancelUnstakeRequest` | Cancel a pending unstake request |
+| 9 | `RequestUnstake` | Start unstake cooldown (requested coins stop earning; earned rewards settled) |
+| 10 | `CompleteUnstake` | Finish unstake after cooldown; full unstake closes the account |
+| 11 | `CancelUnstakeRequest` | Cancel a pending request, restoring coins to the active position |
 | 12 | `CloseStakeAccount` | Close zero-balance stake account to reclaim rent |
 | 13 | ~~`FixTotalRewardDebt`~~ | Deprecated (no-op, returns error) |
 | 14 | `SetPoolMetadata` | Set pool name, tags, and URL (permissionless) |

idl.json

@@ -123,7 +123,8 @@
       "docs": [
         "Unstake tokens from the pool (direct unstake).",
         "Only available when pool has no cooldown configured.",
-        "Claims pending rewards and redistributes stranded rewards."
+        "Claims pending rewards and redistributes stranded rewards.",
+        "A full unstake closes the user stake account."
       ],
       "accounts": [
         {
@@ -159,6 +160,18 @@
         {
           "name": "tokenProgram",
           "docs": ["Token program (SPL Token or Token 2022)"]
+        },
+        {
+          "name": "systemProgram",
+          "optional": true,
+          "address": "11111111111111111111111111111111",
+          "docs": ["System program (optional, for legacy account reallocation)"]
+        },
+        {
+          "name": "metadata",
+          "writable": true,
+          "optional": true,
+          "docs": ["Metadata PDA (optional, to decrement member_count on full-unstake close)"]
         }
       ],
       "args": [
@@ -332,7 +345,10 @@
       "discriminator": [9],
       "docs": [
         "Request unstake - starts cooldown period.",
-        "Tokens remain staked during cooldown. Only one pending request at a time."
+        "The requested coins stop earning rewards immediately (removed from pool",
+        "weight) and their already-earned rewards are settled and paid out now.",
+        "Only the token transfer is deferred to completeUnstake.",
+        "Only one pending request at a time."
       ],
       "accounts": [
         {
@@ -347,8 +363,15 @@
         },
         {
           "name": "owner",
+          "writable": true,
           "signer": true,
-          "docs": ["User / owner"]
+          "docs": ["User / owner (receives settled reward SOL)"]
+        },
+        {
+          "name": "systemProgram",
+          "optional": true,
+          "address": "11111111111111111111111111111111",
+          "docs": ["System program (optional, for legacy account reallocation)"]
         }
       ],
       "args": [
@@ -363,7 +386,8 @@
       "discriminator": [10],
       "docs": [
         "Complete unstake after cooldown has elapsed.",
-        "Claims pending rewards and redistributes stranded rewards."
+        "Delivers the frozen tokens from the vault to the user (reward settlement",
+        "already happened at requestUnstake). A full unstake closes the account."
       ],
       "accounts": [
         {
@@ -399,6 +423,18 @@
         {
           "name": "tokenProgram",
           "docs": ["Token program (SPL Token or Token 2022)"]
+        },
+        {
+          "name": "systemProgram",
+          "optional": true,
+          "address": "11111111111111111111111111111111",
+          "docs": ["System program (optional)"]
+        },
+        {
+          "name": "metadata",
+          "writable": true,
+          "optional": true,
+          "docs": ["Metadata PDA (optional, to decrement member_count on close)"]
         }
       ],
       "args": []
@@ -407,12 +443,14 @@
       "name": "cancelUnstakeRequest",
       "discriminator": [11],
       "docs": [
-        "Cancel a pending unstake request.",
-        "Tokens remain staked."
+        "Cancel a pending unstake request, restoring the frozen coins to the",
+        "active staking position (re-grants weight; re-added tokens get a fresh",
+        "reward snapshot so they earn no cooldown-period rewards)."
       ],
       "accounts": [
         {
           "name": "pool",
+          "writable": true,
           "docs": ["Pool account"]
         },
         {
@@ -424,6 +462,12 @@
           "name": "owner",
           "signer": true,
           "docs": ["User / owner"]
+        },
+        {
+          "name": "systemProgram",
+          "optional": true,
+          "address": "11111111111111111111111111111111",
+          "docs": ["System program (optional, for legacy account reallocation)"]
         }
       ],
       "args": []
@@ -835,7 +879,7 @@
       "docs": [
         "User stake account.",
         "PDA seeds: [\"stake\", pool, owner].",
-        "Size: 177 bytes (legacy accounts at 153 or 161 bytes are lazily reallocated)."
+        "Size: 178 bytes (legacy accounts at 153, 161, or 177 bytes are lazily reallocated)."
       ],
       "type": {
         "kind": "struct",
@@ -911,6 +955,11 @@
             "name": "claimedRewardsWad",
             "type": "u128",
             "docs": ["Cumulative WAD-scaled rewards already paid for the current position. Used for frequency-independent claims: pending = full_entitlement - claimedRewardsWad. Reset to 0 on stake/unstake. Defaults to 0 for legacy accounts."]
+          },
+          {
+            "name": "unstakeRequestSettled",
+            "type": "u8",
+            "docs": ["1 when the pending unstake request was created under the settle-at-request flow (frozen coins already removed from the pool). 0 = no request, or a legacy in-flight request from before this upgrade. Cancel/Complete branch on this to avoid double-restoring or double-removing legacy requests. Defaults to 0 for legacy accounts."]
           }
         ]
       }
@@ -1169,6 +1218,11 @@
       "code": 6034,
       "name": "RewardDebtExceedsBound",
       "msg": "New total_reward_debt exceeds maximum accumulated rewards"
+    },
+    {
+      "code": 6035,
+      "name": "ResidualRewardsPending",
+      "msg": "Claim residual rewards before cancelling a full unstake request"
     }
   ]
 }

programs/chiefstaker/Cargo.toml

@@ -28,3 +28,4 @@ solana-security-txt = "1.1"
 solana-program-test = "2.0"
 solana-sdk = "2.0"
 tokio = { version = "1", features = ["full"] }
+spl-token-2022 = { version = "5.0", features = ["no-entrypoint"] }

programs/chiefstaker/src/error.rs

@@ -108,6 +108,9 @@ pub enum StakingError {
 
     #[error("New total_reward_debt exceeds maximum accumulated rewards")]
     RewardDebtExceedsBound,
+
+    #[error("Claim residual rewards before cancelling a full unstake request")]
+    ResidualRewardsPending,
 }
 
 impl From<StakingError> for ProgramError {

programs/chiefstaker/src/instructions/cancel_unstake.rs

@@ -10,15 +10,24 @@ use solana_program::{
 
 use crate::{
     error::StakingError,
+    math::{wad_mul, U256, WAD},
     state::{StakingPool, UserStake},
 };
 
-/// Cancel a pending unstake request
+/// Cancel a pending unstake request, restoring the frozen coins to the active
+/// staking position so they earn rewards again.
+///
+/// RequestUnstake removed the frozen coins from the pool's reward accounting. This
+/// adds them back: total_staked and sum_stake_exp are restored (re-granting the
+/// weight at the original maturity, since exp_start_factor is unchanged), and the
+/// re-added tokens receive a fresh reward snapshot so they do not retroactively
+/// claim rewards distributed during the cooldown (mirrors add-stake).
 ///
 /// Accounts:
-/// 0. `[]` Pool account
+/// 0. `[writable]` Pool account
 /// 1. `[writable]` User stake account
 /// 2. `[signer]` User/owner
+/// 3. `[]` System program (optional, for legacy account reallocation)
 pub fn process_cancel_unstake_request(
     program_id: &Pubkey,
     accounts: &[AccountInfo],
@@ -38,7 +47,7 @@ pub fn process_cancel_unstake_request(
     if pool_info.owner != program_id {
         return Err(StakingError::InvalidAccountOwner.into());
     }
-    let pool = StakingPool::try_from_slice(&pool_info.try_borrow_data()?)?;
+    let mut pool = StakingPool::try_from_slice(&pool_info.try_borrow_data()?)?;
     if !pool.is_initialized() {
         return Err(StakingError::NotInitialized.into());
     }
@@ -49,6 +58,11 @@ pub fn process_cancel_unstake_request(
         return Err(StakingError::InvalidPDA.into());
     }
 
+    // Check if pool needs rebasing (we are about to grow sum_stake_exp)
+    if pool.get_sum_stake_exp().needs_rebase() {
+        return Err(StakingError::PoolRequiresSync.into());
+    }
+
     // Realloc legacy accounts to current size (payer = user)
     // System program is optional trailing account, only needed for legacy accounts
     let system_program_info = account_info_iter.next();
@@ -78,25 +92,94 @@ pub fn process_cancel_unstake_request(
         return Err(StakingError::InvalidPDA.into());
     }
 
-    // Lazily adjust exp_start_factor if pool has been rebased
-    user_stake.sync_to_pool(&pool)?;
-
     // Check there is a pending request
     if !user_stake.has_pending_unstake_request() {
         return Err(StakingError::NoPendingUnstakeRequest.into());
     }
 
-    let cancelled_amount = user_stake.unstake_request_amount;
+    // Lazily adjust exp_start_factor if pool has been rebased
+    user_stake.sync_to_pool(&pool)?;
+
+    let frozen = user_stake.unstake_request_amount;
+
+    if user_stake.unstake_request_settled == 1 {
+        // New-style request: RequestUnstake removed the frozen coins from the pool
+        // and settled their rewards. Restore them to the active position.
+
+        // Edge case: a full unstake request (amount == 0) that still has unpaid
+        // residual rewards stores those in reward_debt. Reactivating the position
+        // would overwrite that storage with a snapshot, losing the residual.
+        // Require the user to claim it first (ClaimRewards works on amount==0).
+        if user_stake.amount == 0 && user_stake.reward_debt / WAD > 0 {
+            return Err(StakingError::ResidualRewardsPending.into());
+        }
+
+        let frozen_wad = (frozen as u128)
+            .checked_mul(WAD)
+            .ok_or(StakingError::MathOverflow)?;
+
+        // sum_stake_exp += frozen * exp_start_factor (restores weight / maturity)
+        let contribution = wad_mul(frozen_wad, user_stake.exp_start_factor)?;
+        let new_sum = pool
+            .get_sum_stake_exp()
+            .checked_add(U256::from_u128(contribution))
+            .ok_or(StakingError::MathOverflow)?;
+        pool.set_sum_stake_exp(new_sum);
+
+        // total_staked += frozen
+        pool.total_staked = pool
+            .total_staked
+            .checked_add(frozen as u128)
+            .ok_or(StakingError::MathOverflow)?;
+
+        // Fresh reward snapshot for the re-added tokens (no retroactive rewards)
+        let new_token_debt = wad_mul(frozen_wad, pool.acc_reward_per_weighted_share)?;
+
+        if user_stake.amount == 0 {
+            // Was a full request (no residual remains): start clean.
+            user_stake.reward_debt = new_token_debt;
+            user_stake.claimed_rewards_wad = 0;
+        } else {
+            // Partial request: keep the active remainder's pending intact and add a
+            // fresh debt snapshot for the re-added tokens (mirrors add-stake).
+            user_stake.reward_debt = user_stake
+                .reward_debt
+                .checked_add(new_token_debt)
+                .ok_or(StakingError::MathOverflow)?;
+        }
+
+        user_stake.amount = user_stake
+            .amount
+            .checked_add(frozen)
+            .ok_or(StakingError::MathOverflow)?;
+
+        pool.total_reward_debt = pool
+            .total_reward_debt
+            .checked_add(new_token_debt)
+            .ok_or(StakingError::MathOverflow)?;
+    } else {
+        // Legacy in-flight request (created before this upgrade): the coins were
+        // never removed from the pool's accounting, so there is nothing to restore
+        // — adding them back here would credit the user stake they never lost.
+        // Just clear the request, matching the pre-upgrade cancel behavior.
+    }
 
     // Clear the request fields
     user_stake.unstake_request_amount = 0;
     user_stake.unstake_request_time = 0;
+    user_stake.unstake_request_settled = 0;
 
-    // Save user stake
-    let mut stake_data = user_stake_info.try_borrow_mut_data()?;
-    user_stake.serialize(&mut &mut stake_data[..])?;
+    // Save pool and user stake
+    {
+        let mut pool_data = pool_info.try_borrow_mut_data()?;
+        pool.serialize(&mut &mut pool_data[..])?;
+    }
+    {
+        let mut stake_data = user_stake_info.try_borrow_mut_data()?;
+        user_stake.serialize(&mut &mut stake_data[..])?;
+    }
 
-    msg!("Cancelled unstake request for {} tokens", cancelled_amount);
+    msg!("Cancelled unstake request for {} tokens", frozen);
 
     Ok(())
 }