From a75981466ff64afac788ca4efd969e6a93120996 Mon Sep 17 00:00:00 2001
From: dongcool <95206510+dongcool@users.noreply.github.com>
Date: Fri, 8 May 2026 10:15:39 +0800
Subject: [PATCH] fix reward callback failure unlock

---
 contracts/linear/src/epoch_actions.rs          | 14 +++++++++++++-
 contracts/linear/src/validator_pool.rs         |  4 ++++
 .../linear/epoch-action-failure.ava.ts         | 18 ++++++++++++++++++
 3 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/contracts/linear/src/epoch_actions.rs b/contracts/linear/src/epoch_actions.rs
index 90d75007..98b6a6fa 100644
--- a/contracts/linear/src/epoch_actions.rs
+++ b/contracts/linear/src/epoch_actions.rs
@@ -378,13 +378,25 @@ impl LiquidStakingContract {
     pub fn validator_get_balance_callback(
         &mut self,
         validator_id: AccountId,
-        #[callback] total_balance: U128,
+        #[callback_result] result: Result<U128, PromiseError>,
     ) {
         let mut validator = self
             .validator_pool
             .get_validator(&validator_id)
             .expect(ERR_VALIDATOR_NOT_EXIST);
 
+        let total_balance = match result {
+            Ok(total_balance) => total_balance,
+            Err(_) => {
+                validator.on_refresh_total_balance_failed(&mut self.validator_pool);
+                log!(
+                    "Failed to refresh total balance from validator {}",
+                    validator_id
+                );
+                return;
+            }
+        };
+
         let new_balance = total_balance.0;
         let rewards = new_balance - validator.total_balance();
         Event::EpochUpdateRewards {
diff --git a/contracts/linear/src/validator_pool.rs b/contracts/linear/src/validator_pool.rs
index f0f6635c..702eca03 100644
--- a/contracts/linear/src/validator_pool.rs
+++ b/contracts/linear/src/validator_pool.rs
@@ -1049,6 +1049,10 @@ impl Validator {
         pool.save_validator(self);
     }
 
+    pub fn on_refresh_total_balance_failed(&mut self, pool: &mut ValidatorPool) {
+        self.post_execution(pool);
+    }
+
     /// Due to shares calculation and rounding of staking pool contract,
     /// the amount of staked and unstaked balance might be a little bit
     /// different than we requested.
diff --git a/tests/__tests__/linear/epoch-action-failure.ava.ts b/tests/__tests__/linear/epoch-action-failure.ava.ts
index 8f0436c2..704a3638 100644
--- a/tests/__tests__/linear/epoch-action-failure.ava.ts
+++ b/tests/__tests__/linear/epoch-action-failure.ava.ts
@@ -441,4 +441,22 @@ test('get balance failure', async (t) => {
 
   // balance should not change
   await assertValidator(v1, '60', '0');
+
+  await v1.call(v1, 'set_panic', {
+    panic: false,
+  });
+
+  // update reward should be retryable after the failed callback releases executing
+  await owner.call(
+    contract,
+    'epoch_update_rewards',
+    {
+      validator_id: v1.accountId,
+    },
+    {
+      gas: Gas.parse('200 Tgas'),
+    },
+  );
+
+  await assertValidator(v1, '61', '0');
 });