From 1a231ee0ce9e859ddbb5911bc8b65da217c9bc17 Mon Sep 17 00:00:00 2001
From: Thoralf-M <46689931+Thoralf-M@users.noreply.github.com>
Date: Mon, 29 Apr 2024 14:57:33 +0200
Subject: [PATCH] Fix staking amount (#2238)

Co-authored-by: Thibault Martinez <thibault@iota.org>
---
 .../api/block_builder/transaction_builder/transition.rs  | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/sdk/src/client/api/block_builder/transaction_builder/transition.rs b/sdk/src/client/api/block_builder/transaction_builder/transition.rs
index 8889001432..f7cb4eef02 100644
--- a/sdk/src/client/api/block_builder/transaction_builder/transition.rs
+++ b/sdk/src/client/api/block_builder/transaction_builder/transition.rs
@@ -69,6 +69,8 @@ impl TransactionBuilder {
             .cloned()
             .collect::<Vec<_>>();
 
+        let mut new_amount = None;
+
         if let Some(change) = self.transitions.as_ref().and_then(|t| t.accounts.get(&account_id)) {
             match change {
                 AccountChange::BeginStaking {
@@ -83,6 +85,7 @@ impl TransactionBuilder {
                         self.protocol_parameters
                             .past_bounded_slot(self.latest_slot_commitment_id),
                     );
+                    new_amount = Some(*staked_amount);
                     features.push(
                         StakingFeature::new(
                             *staked_amount,
@@ -101,6 +104,7 @@ impl TransactionBuilder {
                             .protocol_parameters
                             .future_bounded_epoch(self.latest_slot_commitment_id);
                         let staking_feature = feature.as_staking();
+                        new_amount = Some(staking_feature.staked_amount());
                         // Just extend the end epoch if it's still possible
                         if future_bounded_epoch <= staking_feature.end_epoch() {
                             *feature = StakingFeature::new(
@@ -162,11 +166,14 @@ impl TransactionBuilder {
         }
 
         let mut builder = AccountOutputBuilder::from(input)
-            .with_minimum_amount(self.protocol_parameters.storage_score_parameters())
             .with_mana(0)
             .with_account_id(account_id)
             .with_foundry_counter(u32::max(highest_foundry_serial_number, input.foundry_counter()))
             .with_features(features);
+        match new_amount {
+            Some(amount) => builder = builder.with_amount(amount),
+            None => builder = builder.with_minimum_amount(self.protocol_parameters.storage_score_parameters()),
+        }
 
         // Block issuers cannot move their mana elsewhere.
         if input.is_block_issuer() {