fix Auto-Claim via Keeper Bot

Author: Eniola3321

contracts/vesting_contracts/src/lib.rs

@@ -18,6 +18,7 @@ const INITIAL_SUPPLY: Symbol = Symbol::new(&"INITIAL_SUPPLY");
 const ADMIN_BALANCE: Symbol = Symbol::new(&"ADMIN_BALANCE");
 const ADMIN_ADDRESS: Symbol = Symbol::new(&"ADMIN_ADDRESS");
 const PROPOSED_ADMIN: Symbol = Symbol::new(&"PROPOSED_ADMIN");
+const KEEPER_FEES: Symbol = Symbol::new(&"KEEPER_FEES");
 
 // Vault structure with lazy initialization
 #[contracttype]
@@ -28,6 +29,7 @@ pub struct Vault {
     pub released_amount: i128,
     pub start_time: u64,
     pub end_time: u64,
+    pub keeper_fee: i128, // Fee paid to anyone who triggers auto_claim
     pub is_initialized: bool, // Lazy initialization flag
     pub is_irrevocable: bool, // Security flag to prevent admin withdrawal
 }
@@ -38,6 +40,7 @@ pub struct BatchCreateData {
     pub amounts: Vec<i128>,
     pub start_times: Vec<u64>,
     pub end_times: Vec<u64>,
+    pub keeper_fees: Vec<i128>,
 }
 
 #[contracttype]
@@ -116,7 +119,7 @@ impl VestingContract {
     }
 
     // Full initialization - writes all metadata immediately
-    pub fn create_vault_full(env: Env, owner: Address, amount: i128, start_time: u64, end_time: u64) -> u64 {
+    pub fn create_vault_full(env: Env, owner: Address, amount: i128, start_time: u64, end_time: u64, keeper_fee: i128) -> u64 {
         Self::require_admin(&env);
 
         // Get next vault ID
@@ -137,6 +140,7 @@ impl VestingContract {
             released_amount: 0,
             start_time,
             end_time,
+            keeper_fee,
             is_initialized: true, // Mark as fully initialized
             is_irrevocable: false, // Default to revocable
         };
@@ -170,7 +174,7 @@ impl VestingContract {
     }
 
     // Lazy initialization - writes minimal data initially
-    pub fn create_vault_lazy(env: Env, owner: Address, amount: i128, start_time: u64, end_time: u64) -> u64 {
+    pub fn create_vault_lazy(env: Env, owner: Address, amount: i128, start_time: u64, end_time: u64, keeper_fee: i128) -> u64 {
         Self::require_admin(&env);
 
         // Get next vault ID
@@ -191,6 +195,7 @@ impl VestingContract {
             released_amount: 0,
             start_time,
             end_time,
+            keeper_fee,
             is_initialized: false, // Mark as lazy initialized
             is_irrevocable: false, // Default to revocable
         };
@@ -232,6 +237,7 @@ impl VestingContract {
                     released_amount: 0,
                     start_time: 0,
                     end_time: 0,
+                    keeper_fee: 0,
                     is_initialized: false,
                     is_irrevocable: false,
                 }
@@ -395,6 +401,7 @@ impl VestingContract {
                 released_amount: 0,
                 start_time: batch_data.start_times.get(i).unwrap(),
                 end_time: batch_data.end_times.get(i).unwrap(),
+                keeper_fee: batch_data.keeper_fees.get(i).unwrap(),
                 is_initialized: false, // Lazy initialization
                 is_irrevocable: false, // Default to revocable
             };
@@ -448,6 +455,7 @@ impl VestingContract {
                 released_amount: 0,
                 start_time: batch_data.start_times.get(i).unwrap(),
                 end_time: batch_data.end_times.get(i).unwrap(),
+                keeper_fee: batch_data.keeper_fees.get(i).unwrap(),
                 is_initialized: true, // Full initialization
                 is_irrevocable: false, // Default to revocable
             };
@@ -496,6 +504,7 @@ impl VestingContract {
                     released_amount: 0,
                     start_time: 0,
                     end_time: 0,
+                    keeper_fee: 0,
                     is_initialized: false,
                     is_irrevocable: false,
                 }
@@ -529,6 +538,7 @@ impl VestingContract {
                         released_amount: 0,
                         start_time: 0,
                         end_time: 0,
+                        keeper_fee: 0,
                         is_initialized: false,
                         is_irrevocable: false,
                     }
@@ -683,4 +693,79 @@ impl VestingContract {
         let sum = total_locked + total_claimed + admin_balance;
         sum == initial_supply
     }
+
+    // --- New Auto-Claim Logic ---
+
+    // Calculate currently claimable tokens based on linear vesting
+    pub fn get_claimable_amount(env: Env, vault_id: u64) -> i128 {
+        let vault: Vault = env.storage().instance()
+            .get(&VAULT_DATA, &vault_id)
+            .unwrap_or_else(|| panic!("Vault not found"));
+
+        let now = env.ledger().timestamp();
+        
+        if now <= vault.start_time {
+            return 0;
+        }
+
+        let elapsed = if now >= vault.end_time {
+            vault.end_time - vault.start_time
+        } else {
+            now - vault.start_time
+        };
+
+        let total_duration = vault.end_time - vault.start_time;
+        
+        let vested = if total_duration > 0 {
+            // Use i128 for calculation to prevent overflow then back to i128
+            (vault.total_amount * elapsed as i128) / total_duration as i128
+        } else {
+            vault.total_amount
+        };
+
+        if vested > vault.released_amount {
+            vested - vault.released_amount
+        } else {
+            0
+        }
+    }
+
+    // Auto-claim function that anyone can call.
+    // Tokens go to beneficiary, but keeper can get a tip.
+    pub fn auto_claim(env: Env, vault_id: u64, keeper: Address) {
+        let mut vault: Vault = env.storage().instance()
+            .get(&VAULT_DATA, &vault_id)
+            .unwrap_or_else(|| panic!("Vault not found"));
+
+        require!(vault.is_initialized, "Vault not initialized");
+
+        let claimable = Self::get_claimable_amount(env.clone(), vault_id);
+        
+        // Ensure there's enough to cover the fee and something left for beneficiary
+        require!(claimable > vault.keeper_fee, "Insufficient claimable tokens to cover fee");
+
+        let beneficiary_amount = claimable - vault.keeper_fee;
+        
+        // Update vault
+        vault.released_amount += claimable;
+        env.storage().instance().set(&VAULT_DATA, &vault_id, &vault);
+
+        // Update keeper fees
+        let mut fees: Map<Address, i128> = env.storage().instance().get(&KEEPER_FEES).unwrap_or(Map::new(&env));
+        let current_fees = fees.get(keeper.clone()).unwrap_or(0);
+        fees.set(keeper.clone(), current_fees + vault.keeper_fee);
+        env.storage().instance().set(&KEEPER_FEES, &fees);
+
+        // Emit KeeperClaim event
+        env.events().publish(
+            (Symbol::new(&env, "KeeperClaim"), vault_id, keeper),
+            (beneficiary_amount, vault.keeper_fee)
+        );
+    }
+
+    // Get accumulated fees for a keeper
+    pub fn get_keeper_fee(env: Env, keeper: Address) -> i128 {
+        let fees: Map<Address, i128> = env.storage().instance().get(&KEEPER_FEES).unwrap_or(Map::new(&env));
+        fees.get(keeper).unwrap_or(0)
+    }
 }

contracts/vesting_contracts/src/test.rs

@@ -104,12 +104,12 @@ fn test_admin_access_control() {
     });
 
     let result = std::panic::catch_unwind(|| {
-        client.create_vault_full(&vault_owner, &1000i128, &100u64, &200u64);
+        client.create_vault_full(&vault_owner, &1000i128, &100u64, &200u64, &0i128);
     });
     assert!(result.is_err());
 
     let result = std::panic::catch_unwind(|| {
-        client.create_vault_lazy(&vault_owner, &1000i128, &100u64, &200u64);
+        client.create_vault_lazy(&vault_owner, &1000i128, &100u64, &200u64, &0i128);
     });
     assert!(result.is_err());
 
@@ -118,10 +118,10 @@ fn test_admin_access_control() {
         env.current_contract_address().set(&admin);
     });
 
-    let vault_id = client.create_vault_full(&vault_owner, &1000i128, &100u64, &200u64);
+    let vault_id = client.create_vault_full(&vault_owner, &1000i128, &100u64, &200u64, &0i128);
     assert_eq!(vault_id, 1);
 
-    let vault_id2 = client.create_vault_lazy(&vault_owner, &500i128, &150u64, &250u64);
+    let vault_id2 = client.create_vault_lazy(&vault_owner, &500i128, &150u64, &250u64, &0i128);
     assert_eq!(vault_id2, 2);
 }
 
@@ -147,6 +147,7 @@ fn test_batch_operations_admin_control() {
         amounts: vec![&env, 1000i128, 2000i128],
         start_times: vec![&env, 100u64, 150u64],
         end_times: vec![&env, 200u64, 250u64],
+        keeper_fees: vec![&env, 0i128, 0i128],
     };
 
     // Test: Unauthorized user cannot create batch vaults
@@ -176,79 +177,62 @@ fn test_batch_operations_admin_control() {
 }
 
 #[test]
-
+fn test_auto_claim_functionality() {
     let env = Env::default();
     let contract_id = env.register(VestingContract, ());
     let client = VestingContractClient::new(&env, &contract_id);
 
-    // Create addresses for testing
     let admin = Address::generate(&env);
-    let vault_owner = Address::generate(&env);
-
-    let unauthorized_user = Address::generate(&env);
+    let beneficiary = Address::generate(&env);
+    let keeper = Address::generate(&env);
 
-    // Initialize contract with admin
     let initial_supply = 1000000i128;
     client.initialize(&admin, &initial_supply);
 
-
+    let amount = 1000i128;
+    let keeper_fee = 10i128;
+    let start_time = env.ledger().timestamp();
+    let end_time = start_time + 100; // 100 seconds duration
+    
+    // Create vault
     env.as_contract(&contract_id, || {
         env.current_contract_address().set(&admin);
     });
+    let vault_id = client.create_vault_full(&beneficiary, &amount, &start_time, &end_time, &keeper_fee);
 
-
-    env.as_contract(&contract_id, || {
-        env.current_contract_address().set(&unauthorized_user);
-    });
+    // Advance time by 50 seconds (halfway)
+    env.ledger().set_timestamp(start_time + 50);
 
-    let result = std::panic::catch_unwind(|| {
-
-    });
-    assert!(result.is_err());
-}
-
-#[test]
-
-    let env = Env::default();
-    let contract_id = env.register(VestingContract, ());
-    let client = VestingContractClient::new(&env, &contract_id);
+    // Check claimable amount
+    let claimable = client.get_claimable_amount(&vault_id);
+    assert!(claimable >= 500); // Linear vesting should be ~500
 
-    // Create addresses for testing
-    let admin = Address::generate(&env);
-    let vault_owner = Address::generate(&env);
-
+    // Perform auto-claim
+    client.auto_claim(&vault_id, &keeper);
 
-    // Initialize contract with admin
-    let initial_supply = 1000000i128;
-    client.initialize(&admin, &initial_supply);
+    // Verify vault state
+    let vault = client.get_vault(&vault_id);
+    assert_eq!(vault.released_amount, claimable);
 
-
-    env.as_contract(&contract_id, || {
-        env.current_contract_address().set(&admin);
-    });
+    // Verify keeper received fee
+    let accumulated_fee = client.get_keeper_fee(&keeper);
+    assert_eq!(accumulated_fee, keeper_fee);
 
-
+    // Test: Calling auto_claim again immediately should fail (no tokens to cover fee)
+    let result = std::panic::catch_unwind(|| {
+        client.auto_claim(&vault_id, &keeper);
     });
     assert!(result.is_err());
-}
-
-#[test]
-
-    let env = Env::default();
-    let contract_id = env.register(VestingContract, ());
-    let client = VestingContractClient::new(&env, &contract_id);
 
-    // Create addresses for testing
-    let admin = Address::generate(&env);
-    let vault_owner = Address::generate(&env);
-
-    
-    // Initialize contract with admin
-    let initial_supply = 1000000i128;
-    client.initialize(&admin, &initial_supply);
+    // Advance time to end
+    env.ledger().set_timestamp(end_time);
 
-
-    env.as_contract(&contract_id, || {
-        env.current_contract_address().set(&admin);
-    });
+    // Perform another auto-claim
+    let current_released = client.get_vault(&vault_id).released_amount;
+    let second_claimable = client.get_claimable_amount(&vault_id);
+    client.auto_claim(&vault_id, &keeper);
 
+    let vault_final = client.get_vault(&vault_id);
+    assert_eq!(vault_final.released_amount, current_released + second_claimable);
+    assert_eq!(client.get_keeper_fee(&keeper), keeper_fee * 2);
+}