fix(engine)!: charge finalization fees against the state actually persisted

crates/engine/src/fees/fee_module.rs

@@ -7,7 +7,7 @@ use tari_template_lib::types::TemplateAddress;
 
 use super::FeeTable;
 use crate::{
-    runtime::{RuntimeEvent, RuntimeModule, RuntimeModuleError, StateTracker},
+    runtime::{ChargeableState, RuntimeEvent, RuntimeModule, RuntimeModuleError, StateTracker},
     state_store::StateReader,
 };
 
@@ -45,6 +45,113 @@ impl FeeModule {
         Ok((base_bytes, premium))
     }
 
+    /// Charges everything that is a function of the state being persisted, rather than of the work
+    /// done to produce it: storage, the template-publish premium, substate slots, the metering of
+    /// WASM and native execution, and the exhaust burn over the resulting total.
+    ///
+    /// Every charge here is *assigned*, not accumulated, so that running this again against a
+    /// different state replaces the result rather than doubling it. That is what lets a transaction
+    /// be gated on the cost of the state it asked to commit and then billed for the state that is
+    /// actually persisted — see [`RuntimeModule::on_before_persist`].
+    fn charge_finalization_fees<TStore: StateReader>(
+        &self,
+        state: &mut ChargeableState<'_, TStore>,
+    ) -> Result<(), RuntimeModuleError> {
+        let mut counter = ByteCounter::new();
+        let mut template_base_bytes = 0u64;
+        let mut template_premium = 0u64;
+        for substate in state.substates_to_persist().values() {
+            // A published template's binary is priced by the dedicated base + quadratic publish
+            // model, so keep it out of the flat per-byte storage tally. Accumulate the raw
+            // metrics here and apply the storage divisor once below.
+            if let SubstateValue::Template(template) = substate {
+                let (tpl_base_bytes, tpl_premium) = self.template_publish_metrics(template.binary.len())?;
+                template_base_bytes = template_base_bytes.checked_add(tpl_base_bytes).ok_or_else(|| {
+                    RuntimeModuleError::Overflow("Overflow accumulating template base bytes".to_string())
+                })?;
+                template_premium = template_premium.checked_add(tpl_premium).ok_or_else(|| {
+                    RuntimeModuleError::Overflow("Overflow accumulating template publish premium".to_string())
+                })?;
+                continue;
+            }
+            encode_into_writer(substate, &mut counter)?;
+        }
+
+        // Finalization persists the transaction receipt on top of the mutated substates. It carries
+        // the transaction's events, so leaving it out of the tally would make that payload — the
+        // largest caller-controlled contribution to permanent state after the substates themselves —
+        // free.
+        let receipt_bytes = state
+            .transaction_receipt_size()
+            .map_err(|e| RuntimeModuleError::Runtime(e.to_string()))?;
+        let total_storage = counter
+            .get()
+            .checked_add(receipt_bytes)
+            .ok_or_else(|| RuntimeModuleError::Overflow("Overflow accumulating storage bytes".to_string()))?;
+
+        let cost = self
+            .fee_table
+            .per_byte_storage_cost()
+            .checked_mul(total_storage as u64)
+            .ok_or_else(|| RuntimeModuleError::Overflow("Overflow calculating storage cost".to_string()))?;
+        let storage_cost = cost / self.fee_table.storage_cost_divisor();
+
+        let template_base_cost = self
+            .fee_table
+            .per_byte_storage_cost()
+            .checked_mul(template_base_bytes)
+            .ok_or_else(|| {
+                RuntimeModuleError::Overflow("Overflow calculating template base storage cost".to_string())
+            })? /
+            self.fee_table.storage_cost_divisor();
+        let template_publish_cost = template_base_cost
+            .checked_add(template_premium)
+            .ok_or_else(|| RuntimeModuleError::Overflow("Overflow calculating template publish cost".to_string()))?;
+
+        // The receipt occupies a slot of its own — it is always newly created, since it is addressed
+        // by a transaction id that can only be finalized once.
+        let new_substate_count = state
+            .count_newly_created_substates()
+            .map_err(|e| RuntimeModuleError::Runtime(e.to_string()))?
+            .saturating_add(1);
+        let create_cost = (new_substate_count as u64)
+            .checked_mul(self.fee_table.per_substate_create_cost())
+            .ok_or_else(|| RuntimeModuleError::Overflow("Overflow calculating substate create cost".to_string()))?;
+
+        // WASM execution: charge once against the transaction's accumulated points so the divisor
+        // rounds against the total. Per-call rounding would let a transaction split work into
+        // sub-divisor chunks and pay zero for any single one (each `points/divisor` is `0`), even
+        // though the summed work is non-trivial.
+        let units = state.fee_state().accumulated_wasm_points() / self.fee_table.wasm_points_cost_divisor();
+        let wasm_cost = units
+            .checked_mul(self.fee_table.per_wasm_point_cost())
+            .ok_or_else(|| RuntimeModuleError::Overflow("Overflow calculating WASM execution cost".to_string()))?;
+
+        // Native verification is priced in the same points and charged at the same rate, under its
+        // own source so the breakdown distinguishes crypto verification from template execution.
+        let native_units = state.fee_state().accumulated_native_points() / self.fee_table.wasm_points_cost_divisor();
+        let native_cost = native_units
+            .checked_mul(self.fee_table.per_wasm_point_cost())
+            .ok_or_else(|| RuntimeModuleError::Overflow("Overflow calculating native execution cost".to_string()))?;
+
+        let fee_state = state.fee_state_mut();
+        fee_state.set_charge(FeeSource::Storage, storage_cost);
+        fee_state.set_charge(FeeSource::TemplatePublish, template_publish_cost);
+        fee_state.set_charge(FeeSource::SubstateCreate, create_cost);
+        fee_state.set_charge(FeeSource::WasmExecution, wasm_cost);
+        fee_state.set_charge(FeeSource::NativeExecution, native_cost);
+
+        // Exhaust burn: charged on top of the execution fee accrued so far, so leaders receive the execution fee in
+        // full and the burn amount is destroyed separately. The rate is seeded onto the fee state at execution time
+        // for the execution epoch. Zeroed first so that the total it is taken over never includes a burn from an
+        // earlier pass over a different state.
+        fee_state.set_charge(FeeSource::ExhaustBurn, 0);
+        let burn = calculate_burn_amount(fee_state.total_charges(), fee_state.burn_rate_bps())?;
+        fee_state.set_charge(FeeSource::ExhaustBurn, burn);
+
+        Ok(())
+    }
+
     #[cfg(test)]
     fn template_publish_cost(&self, binary_len: usize) -> Result<u64, RuntimeModuleError> {
         let (base_bytes, premium) = self.template_publish_metrics(binary_len)?;
@@ -108,98 +215,11 @@ impl<TStore: StateReader> RuntimeModule<TStore> for FeeModule {
     }
 
     fn on_before_finalize(&self, track: &mut StateTracker<TStore>) -> Result<(), RuntimeModuleError> {
-        let (total_storage, template_base_bytes, template_premium) = track.with_substates_to_persist(|changes| {
-            let mut counter = ByteCounter::new();
-            let mut base_bytes = 0u64;
-            let mut premium = 0u64;
-            for substate in changes.values() {
-                // A published template's binary is priced by the dedicated base + quadratic publish
-                // model, so keep it out of the flat per-byte storage tally. Accumulate the raw
-                // metrics here and apply the storage divisor once below.
-                if let SubstateValue::Template(template) = substate {
-                    let (tpl_base_bytes, tpl_premium) = self.template_publish_metrics(template.binary.len())?;
-                    base_bytes = base_bytes.checked_add(tpl_base_bytes).ok_or_else(|| {
-                        RuntimeModuleError::Overflow("Overflow accumulating template base bytes".to_string())
-                    })?;
-                    premium = premium.checked_add(tpl_premium).ok_or_else(|| {
-                        RuntimeModuleError::Overflow("Overflow accumulating template publish premium".to_string())
-                    })?;
-                    continue;
-                }
-                encode_into_writer(substate, &mut counter)?;
-            }
-            Ok::<_, RuntimeModuleError>((counter.get(), base_bytes, premium))
-        })?;
-
-        // Finalization persists the transaction receipt on top of the mutated substates. It carries
-        // the transaction's events, so leaving it out of the tally would make that payload — the
-        // largest caller-controlled contribution to permanent state after the substates themselves —
-        // free.
-        let receipt_bytes = track
-            .transaction_receipt_size()
-            .map_err(|e| RuntimeModuleError::Runtime(e.to_string()))?;
-        let total_storage = total_storage
-            .checked_add(receipt_bytes)
-            .ok_or_else(|| RuntimeModuleError::Overflow("Overflow accumulating storage bytes".to_string()))?;
-
-        let cost = self
-            .fee_table
-            .per_byte_storage_cost()
-            .checked_mul(total_storage as u64)
-            .ok_or_else(|| RuntimeModuleError::Overflow("Overflow calculating storage cost".to_string()))?;
-        track.add_fee_charge(FeeSource::Storage, cost / self.fee_table.storage_cost_divisor());
-
-        let template_base_cost = self
-            .fee_table
-            .per_byte_storage_cost()
-            .checked_mul(template_base_bytes)
-            .ok_or_else(|| {
-                RuntimeModuleError::Overflow("Overflow calculating template base storage cost".to_string())
-            })? /
-            self.fee_table.storage_cost_divisor();
-        let template_publish_cost = template_base_cost
-            .checked_add(template_premium)
-            .ok_or_else(|| RuntimeModuleError::Overflow("Overflow calculating template publish cost".to_string()))?;
-        if template_publish_cost > 0 {
-            track.add_fee_charge(FeeSource::TemplatePublish, template_publish_cost);
-        }
-
-        // The receipt occupies a slot of its own — it is always newly created, since it is addressed
-        // by a transaction id that can only be finalized once.
-        let new_substate_count = track
-            .count_newly_created_substates()
-            .map_err(|e| RuntimeModuleError::Runtime(e.to_string()))?
-            .saturating_add(1);
-        let create_cost = (new_substate_count as u64)
-            .checked_mul(self.fee_table.per_substate_create_cost())
-            .ok_or_else(|| RuntimeModuleError::Overflow("Overflow calculating substate create cost".to_string()))?;
-        track.add_fee_charge(FeeSource::SubstateCreate, create_cost);
-
-        // WASM execution: charge once against the transaction's accumulated points so the divisor
-        // rounds against the total. Per-call rounding would let a transaction split work into
-        // sub-divisor chunks and pay zero for any single one (each `points/divisor` is `0`), even
-        // though the summed work is non-trivial.
-        let units = track.accumulated_wasm_points() / self.fee_table.wasm_points_cost_divisor();
-        let wasm_cost = units
-            .checked_mul(self.fee_table.per_wasm_point_cost())
-            .ok_or_else(|| RuntimeModuleError::Overflow("Overflow calculating WASM execution cost".to_string()))?;
-        track.add_fee_charge(FeeSource::WasmExecution, wasm_cost);
-
-        // Native verification is priced in the same points and charged at the same rate, under its
-        // own source so the breakdown distinguishes crypto verification from template execution.
-        let native_units = track.accumulated_native_points() / self.fee_table.wasm_points_cost_divisor();
-        let native_cost = native_units
-            .checked_mul(self.fee_table.per_wasm_point_cost())
-            .ok_or_else(|| RuntimeModuleError::Overflow("Overflow calculating native execution cost".to_string()))?;
-        track.add_fee_charge(FeeSource::NativeExecution, native_cost);
-
-        // Exhaust burn: charged on top of the execution fee accrued so far, so leaders receive the execution fee in
-        // full and the burn amount is destroyed separately. The rate is seeded onto the fee state at execution time
-        // for the execution epoch.
-        let burn = calculate_burn_amount(track.total_fee_charges(), track.fee_burn_rate_bps())?;
-        track.add_fee_charge(FeeSource::ExhaustBurn, burn);
+        self.charge_finalization_fees(&mut track.chargeable_state())
+    }
 
-        Ok(())
+    fn on_before_persist(&self, state: &mut ChargeableState<'_, TStore>) -> Result<(), RuntimeModuleError> {
+        self.charge_finalization_fees(state)
     }
 
     fn on_runtime_event(

crates/engine/src/runtime/fee_state.rs

@@ -93,6 +93,12 @@ impl FeeState {
         self.fee_charges.add(source, amount)
     }
 
+    /// Replaces the charge for `source`. Used by the fee module to recompute the finalization
+    /// charges once the state that will actually be persisted is known.
+    pub fn set_charge(&mut self, source: FeeSource, amount: u64) {
+        self.fee_charges.set(source, amount)
+    }
+
     pub fn accumulate_wasm_points(&mut self, points: u64) {
         self.accumulated_wasm_points = self.accumulated_wasm_points.saturating_add(points);
     }

crates/engine/src/runtime/impl.rs

@@ -164,7 +164,7 @@ use crate::{
         locking::{LockError, LockedSubstate},
         pay_fee::PayFee,
         scope::PushCallFrame,
-        tracker::StateTracker,
+        tracker::{FinalizedState, StateTracker},
     },
     state_store::StateReader,
     template::LoadedTemplate,
@@ -275,6 +275,29 @@ impl<TStore: StateReader + Clone + 'static, TTemplateProvider: TemplateProvider<
         Ok(())
     }
 
+    /// Settles the transaction into a result.
+    ///
+    /// The fee module charges twice here, and the order is the point of the split. It first charges
+    /// against the live state: those charges are what [`StateTracker::select_finalized_state`] tests
+    /// against the payments, so they decide whether the transaction commits or falls back to a
+    /// fee-intent commit. Once that is decided, it charges again against the state actually chosen,
+    /// which on a fee-intent commit holds only what the fee intent touched. A transaction is
+    /// therefore gated on the cost of the state it asked to commit, but pays for the state that is
+    /// really persisted.
+    fn finalize_with(&mut self, failure: Option<RejectReason>) -> Result<FinalizeResult, RuntimeError> {
+        self.invoke_modules_on_before_finalize()?;
+        let mut finalized = self.tracker.select_finalized_state(failure)?;
+        self.invoke_modules_on_before_persist(&mut finalized)?;
+        self.tracker.finalize(finalized)
+    }
+
+    fn invoke_modules_on_before_persist(&mut self, finalized: &mut FinalizedState<TStore>) -> Result<(), RuntimeError> {
+        for module in self.modules.iter() {
+            module.on_before_persist(&mut finalized.chargeable_state())?;
+        }
+        Ok(())
+    }
+
     fn invoke_modules_on_runtime_event(&mut self, event: RuntimeEvent) -> Result<(), RuntimeError> {
         for module in self.modules.iter() {
             module.on_runtime_event(&mut self.tracker, &event)?;
@@ -3370,16 +3393,12 @@ where
 
     fn finalize(&mut self) -> Result<FinalizeResult, RuntimeError> {
         self.invoke_modules_on_runtime_call("finalize")?;
-        // If the fee module is present, this will add substate storage fees
-        self.invoke_modules_on_before_finalize()?;
-        self.tracker.finalize(None)
+        self.finalize_with(None)
     }
 
     fn finalize_failure(&mut self, reason: RejectReason) -> Result<FinalizeResult, RuntimeError> {
         self.invoke_modules_on_runtime_call("finalize_failure")?;
-        // If the fee module is present, this will add substate storage fees
-        self.invoke_modules_on_before_finalize()?;
-        self.tracker.finalize(Some(reason))
+        self.finalize_with(Some(reason))
     }
 
     fn validate_finalized(&self) -> Result<(), RuntimeError> {

crates/engine/src/runtime/mod.rs

@@ -109,7 +109,8 @@ use tari_template_lib::{
         stealth::StealthTransferStatement,
     },
 };
-pub use tracker::StateTracker;
+pub use tracker::{FinalizedState, StateTracker};
+pub use working_state::ChargeableState;
 
 use crate::runtime::{locking::LockedSubstate, scope::PushCallFrame};
 

crates/engine/src/runtime/module.rs

@@ -3,7 +3,7 @@
 
 use tari_template_lib::types::TemplateAddress;
 
-use crate::runtime::StateTracker;
+use crate::runtime::{ChargeableState, StateTracker};
 
 pub trait RuntimeModule<TStore>: Send + Sync {
     fn on_initialize(&self, _track: &mut StateTracker<TStore>) -> Result<(), RuntimeModuleError> {
@@ -32,10 +32,21 @@ pub trait RuntimeModule<TStore>: Send + Sync {
         Ok(())
     }
 
+    /// Invoked at the start of finalization, against the live working state — before it is known
+    /// whether the transaction commits or falls back to a fee-intent commit. Charges added here
+    /// decide that outcome: they are what the paid-in-full check sees.
     fn on_before_finalize(&self, _track: &mut StateTracker<TStore>) -> Result<(), RuntimeModuleError> {
         Ok(())
     }
 
+    /// Invoked once the state that finalization will persist has been chosen, and before its fees
+    /// are settled. On a fee-intent commit that state is the fee checkpoint, not the live state
+    /// [`Self::on_before_finalize`] saw, so any charge that is a function of what gets persisted
+    /// must be recomputed here against `state`.
+    fn on_before_persist(&self, _state: &mut ChargeableState<'_, TStore>) -> Result<(), RuntimeModuleError> {
+        Ok(())
+    }
+
     fn on_runtime_event(
         &self,
         _track: &mut StateTracker<TStore>,