feat: implement treasury stream cancellation and recovery

apps/onchain/contracts/treasury/src/lib.rs

@@ -12,11 +12,9 @@ use multisig::{
     propose as multisig_propose, replace_config as multisig_replace_config, sign as multisig_sign,
 };
 use reentrancy_guard::{acquire as acquire_reentrancy, release as release_reentrancy};
-use soroban_sdk::{contract, contractimpl, token, Address, Env, Vec};
+use soroban_sdk::{contract, contractimpl, token, Address, Env, String, Vec};
 use storage::{DataKey, StreamData, LEDGER_BUMP, LEDGER_THRESHOLD};
 
-// Re-exports so tests (and external clients) can construct / inspect the
-// multisig types without depending on the internal module layout.
 pub use storage::{
     MultisigConfig, Proposal, ProposalAction, ProposalStatus, Signer, MAX_SIGNERS,
     PROPOSAL_TTL_SECS,
@@ -139,7 +137,6 @@ impl TreasuryContract {
         request_id: soroban_sdk::BytesN<32>,
     ) -> Result<(), TreasuryError> {
         Self::with_reentrancy_guard(&env, || {
-            // Idempotency check
             if idempotency_guard::claim_request(&env, &request_id).is_err() {
                 return Err(TreasuryError::AlreadyExecuted);
             }
@@ -185,7 +182,6 @@ impl TreasuryContract {
                 LEDGER_BUMP,
             );
 
-            // Transfer tokens from admin to treasury
             let token_client = token::TokenClient::new(&env, &token_addr);
             token_client.transfer(&admin, env.current_contract_address(), &amount);
 
@@ -271,26 +267,21 @@ impl TreasuryContract {
                 .get(&old_key)
                 .ok_or(TreasuryError::StreamNotFound)?;
 
-            // Preserve the claimed amount and total amount, only change beneficiary
             let claimed_amount = stream.claimed_amount;
             let remaining_amount = stream.total_amount - claimed_amount;
 
             if claimed_amount == 0 {
-                // No claims yet: preserve vesting schedule, just change beneficiary
                 stream.beneficiary = new_beneficiary.clone();
             } else {
-                // Partial claims made: reset to immediate vesting of remaining amount
                 stream.claimed_amount = 0;
                 stream.total_amount = remaining_amount;
                 stream.start_time = env.ledger().timestamp();
                 stream.duration = 0;
                 stream.beneficiary = new_beneficiary.clone();
             }
 
-            // Remove old stream entry
             env.storage().persistent().remove(&old_key);
 
-            // Create new stream entry with updated beneficiary
             let new_key = DataKey::Stream(new_beneficiary.clone());
             env.storage().persistent().set(&new_key, &stream);
             env.storage()
@@ -309,9 +300,7 @@ impl TreasuryContract {
         })
     }
 
-    /// Multisig-gated admin rotation. The executor must be a signer consuming
-    /// an approved `SetAdmin` proposal. The proposal is marked Executed and
-    /// cannot be replayed. Unauthorized callers cannot bypass the multisig.
+    /// Multisig-gated admin rotation.
     pub fn set_admin_via_multisig(
         env: Env,
         executor: Address,
@@ -323,9 +312,7 @@ impl TreasuryContract {
         Ok(())
     }
 
-    /// Multisig-gated beneficiary rotation. The executor must be a signer
-    /// consuming an approved `RotateBeneficiary` proposal. Preserves all the
-    /// existing claim/vesting logic of `rotate_beneficiary`.
+    /// Multisig-gated beneficiary rotation.
     pub fn rotate_beneficiary_via_multisig(
         env: Env,
         executor: Address,
@@ -385,9 +372,7 @@ impl TreasuryContract {
         })
     }
 
-    /// Replace the multisig config (signers + threshold). Bootstrap-only:
-    /// gated to an approved `SetAdmin` proposal whose side-effect is the
-    /// signer-set swap. This lets the multisig rotate itself.
+    /// Replace the multisig config (signers + threshold).
     pub fn set_multisig_config(
         env: Env,
         executor: Address,
@@ -424,6 +409,119 @@ impl TreasuryContract {
             .get(&DataKey::Token)
             .ok_or(TreasuryError::NotInitialized)
     }
+
+    // ============================================
+    // CANCELLATION & RECOVERY FUNCTIONS
+    // ============================================
+
+    /// Cancel an active stream and return (total unlocked, refundable) amounts.
+    /// `total_unlocked` includes any amount already claimed prior to cancellation.
+    pub fn cancel_stream(
+        env: Env,
+        admin: Address,
+        beneficiary: Address,
+    ) -> Result<(i128, i128), TreasuryError> {
+        let stored_admin: Address = env
+            .storage()
+            .instance()
+            .get(&DataKey::Admin)
+            .ok_or(TreasuryError::NotInitialized)?;
+
+        if admin != stored_admin {
+            return Err(TreasuryError::Unauthorized);
+        }
+        admin.require_auth();
+
+        let token_address: Address = env
+            .storage()
+            .instance()
+            .get(&DataKey::Token)
+            .ok_or(TreasuryError::NotInitialized)?;
+
+        let token_client = token::TokenClient::new(&env, &token_address);
+        let contract_address = env.current_contract_address();
+
+        let stream_key = DataKey::Stream(beneficiary.clone());
+        let stream: StreamData = env
+            .storage()
+            .persistent()
+            .get(&stream_key)
+            .ok_or(TreasuryError::StreamNotFound)?;
+
+        let current_time = env.ledger().timestamp();
+        // `calculate_unlocked` returns the amount unlocked *since the last claim*
+        // (it subtracts claimed_amount internally). To report the total unlocked
+        // to date, add back what's already been claimed.
+        let newly_unlocked = Self::calculate_unlocked(current_time, &stream);
+        let remaining = stream.total_amount - stream.claimed_amount;
+        let refundable = remaining - newly_unlocked;
+        let total_unlocked = stream.claimed_amount + newly_unlocked;
+
+        if refundable > 0 {
+            token_client.transfer(&contract_address, &beneficiary, &refundable);
+        }
+
+        #[allow(deprecated)]
+        env.events().publish(
+            (String::from_str(&env, "stream_cancelled"),),
+            (&beneficiary, total_unlocked, refundable, current_time),
+        );
+
+        env.storage().persistent().remove(&stream_key);
+
+        Ok((total_unlocked, refundable))
+    }
+
+    /// Emergency stop - refund full remaining amount regardless of vesting
+    pub fn emergency_stop(
+        env: Env,
+        admin: Address,
+        beneficiary: Address,
+        reason: String,
+    ) -> Result<i128, TreasuryError> {
+        let stored_admin: Address = env
+            .storage()
+            .instance()
+            .get(&DataKey::Admin)
+            .ok_or(TreasuryError::NotInitialized)?;
+
+        if admin != stored_admin {
+            return Err(TreasuryError::Unauthorized);
+        }
+        admin.require_auth();
+
+        let token_address: Address = env
+            .storage()
+            .instance()
+            .get(&DataKey::Token)
+            .ok_or(TreasuryError::NotInitialized)?;
+
+        let token_client = token::TokenClient::new(&env, &token_address);
+        let contract_address = env.current_contract_address();
+
+        let stream_key = DataKey::Stream(beneficiary.clone());
+        let stream: StreamData = env
+            .storage()
+            .persistent()
+            .get(&stream_key)
+            .ok_or(TreasuryError::StreamNotFound)?;
+
+        let full_refund = stream.total_amount - stream.claimed_amount;
+
+        if full_refund > 0 {
+            token_client.transfer(&contract_address, &beneficiary, &full_refund);
+        }
+
+        #[allow(deprecated)]
+        env.events().publish(
+            (String::from_str(&env, "emergency_stop"),),
+            (&beneficiary, reason, full_refund),
+        );
+
+        env.storage().persistent().remove(&stream_key);
+
+        Ok(full_refund)
+    }
 }
 
 #[cfg(test)]

apps/onchain/contracts/treasury/src/test.rs

@@ -1,6 +1,6 @@
 use super::*;
 use soroban_sdk::testutils::{Address as _, Ledger};
-use soroban_sdk::{token, vec, Address, BytesN, Env};
+use soroban_sdk::{token, vec, Address, BytesN, Env, String};
 
 fn request_id(env: &Env) -> BytesN<32> {
     BytesN::from_array(env, &[0; 32])
@@ -417,6 +417,153 @@ fn test_rotate_beneficiary_stream_not_found() {
     );
 }
 
+// ── Cancellation & emergency recovery tests ──────────────────
+
+#[test]
+fn test_cancel_stream_partial_claim() {
+    let env = Env::default();
+    env.mock_all_auths();
+
+    let admin = Address::generate(&env);
+    let beneficiary = Address::generate(&env);
+    let token_admin = Address::generate(&env);
+
+    let token_id = env.register_stellar_asset_contract_v2(token_admin.clone());
+    let token_admin_client = token::StellarAssetClient::new(&env, &token_id.address());
+
+    let treasury_id = env.register(TreasuryContract, ());
+    let treasury_client = TreasuryContractClient::new(&env, &treasury_id);
+
+    treasury_client.initialize(&admin, &token_id.address());
+
+    let amount = 1000i128;
+    token_admin_client.mint(&admin, &amount);
+
+    let start_time = 1000u64;
+    let duration = 1000u64;
+    env.ledger().set_timestamp(start_time);
+
+    treasury_client.allocate_budget(
+        &admin,
+        &beneficiary,
+        &amount,
+        &start_time,
+        &duration,
+        &request_id(&env),
+    );
+
+    env.ledger().set_timestamp(start_time + 500);
+
+    let claimed = treasury_client.claim(&beneficiary);
+    assert_eq!(claimed, 500);
+
+    let (claimed_total, refunded) = treasury_client.cancel_stream(&admin, &beneficiary);
+    assert_eq!(claimed_total, 500);
+    assert_eq!(refunded, 500);
+}
+
+#[test]
+fn test_cancel_stream_immediate() {
+    let env = Env::default();
+    env.mock_all_auths();
+
+    let admin = Address::generate(&env);
+    let beneficiary = Address::generate(&env);
+    let token_admin = Address::generate(&env);
+
+    let token_id = env.register_stellar_asset_contract_v2(token_admin.clone());
+    let token_admin_client = token::StellarAssetClient::new(&env, &token_id.address());
+
+    let treasury_id = env.register(TreasuryContract, ());
+    let treasury_client = TreasuryContractClient::new(&env, &treasury_id);
+
+    treasury_client.initialize(&admin, &token_id.address());
+
+    let amount = 1000i128;
+    token_admin_client.mint(&admin, &amount);
+
+    let start_time = 1000u64;
+    let duration = 1000u64;
+    env.ledger().set_timestamp(start_time);
+
+    treasury_client.allocate_budget(
+        &admin,
+        &beneficiary,
+        &amount,
+        &start_time,
+        &duration,
+        &request_id(&env),
+    );
+
+    let (claimed, refunded) = treasury_client.cancel_stream(&admin, &beneficiary);
+    assert_eq!(claimed, 0);
+    assert_eq!(refunded, 1000);
+}
+
+#[test]
+fn test_emergency_stop_full_refund() {
+    let env = Env::default();
+    env.mock_all_auths();
+
+    let admin = Address::generate(&env);
+    let beneficiary = Address::generate(&env);
+    let token_admin = Address::generate(&env);
+
+    let token_id = env.register_stellar_asset_contract_v2(token_admin.clone());
+    let token_admin_client = token::StellarAssetClient::new(&env, &token_id.address());
+
+    let treasury_id = env.register(TreasuryContract, ());
+    let treasury_client = TreasuryContractClient::new(&env, &treasury_id);
+
+    treasury_client.initialize(&admin, &token_id.address());
+
+    let amount = 1000i128;
+    token_admin_client.mint(&admin, &amount);
+
+    let start_time = 1000u64;
+    let duration = 1000u64;
+    env.ledger().set_timestamp(start_time);
+
+    treasury_client.allocate_budget(
+        &admin,
+        &beneficiary,
+        &amount,
+        &start_time,
+        &duration,
+        &request_id(&env),
+    );
+
+    env.ledger().set_timestamp(start_time + 500);
+
+    let refunded = treasury_client.emergency_stop(
+        &admin,
+        &beneficiary,
+        &String::from_str(&env, "Security breach"),
+    );
+
+    assert_eq!(refunded, 1000);
+}
+
+#[test]
+fn test_cancel_nonexistent_stream() {
+    let env = Env::default();
+    env.mock_all_auths();
+
+    let admin = Address::generate(&env);
+    let beneficiary = Address::generate(&env);
+    let token_admin = Address::generate(&env);
+
+    let token_id = env.register_stellar_asset_contract_v2(token_admin.clone());
+
+    let treasury_id = env.register(TreasuryContract, ());
+    let treasury_client = TreasuryContractClient::new(&env, &treasury_id);
+
+    treasury_client.initialize(&admin, &token_id.address());
+
+    let result = treasury_client.try_cancel_stream(&admin, &beneficiary);
+    assert!(result.is_err());
+}
+
 // ── Issue #864: Multisig propose/execute lifecycle tests ─────────
 
 /// `set_admin_via_multisig` is gated: an outsider cannot execute it.