From d1fca116ba894a1f5f5e30099e2f247871dd5d34 Mon Sep 17 00:00:00 2001 From: Greg Zaitsev Date: Mon, 18 May 2026 16:47:40 -0400 Subject: [PATCH] Handle swap errors when taking fees in alpha --- pallets/transaction-fee/src/lib.rs | 51 +++++++++------ pallets/transaction-fee/src/tests/mod.rs | 82 +++++++++++++++++++++++- 2 files changed, 111 insertions(+), 22 deletions(-) diff --git a/pallets/transaction-fee/src/lib.rs b/pallets/transaction-fee/src/lib.rs index 72b27cdcfb..7b4e0d6d2d 100644 --- a/pallets/transaction-fee/src/lib.rs +++ b/pallets/transaction-fee/src/lib.rs @@ -3,6 +3,7 @@ // FRAME use frame_support::{ pallet_prelude::*, + storage::{TransactionOutcome, with_transaction}, traits::{ Imbalance, IsSubType, OnUnbalanced, fungible::{ @@ -16,8 +17,9 @@ use frame_support::{ // Runtime use sp_runtime::{ - Perbill, Saturating, + DispatchError, Perbill, Saturating, traits::{DispatchInfoOf, PostDispatchInfoOf}, + transaction_validity::{InvalidTransaction, TransactionValidityError}, }; // Pallets @@ -67,7 +69,7 @@ pub trait AlphaFeeHandler { coldkey: &AccountIdOf, alpha_vec: &[(AccountIdOf, NetUid)], tao_amount: TaoBalance, - ) -> (AlphaBalance, TaoBalance, NetUid); + ) -> Result<(AlphaBalance, TaoBalance, NetUid), TransactionValidityError>; fn get_all_netuids_for_coldkey_and_hotkey( coldkey: &AccountIdOf, hotkey: &AccountIdOf, @@ -155,9 +157,9 @@ where coldkey: &AccountIdOf, alpha_vec: &[(AccountIdOf, NetUid)], tao_amount: TaoBalance, - ) -> (AlphaBalance, TaoBalance, NetUid) { + ) -> Result<(AlphaBalance, TaoBalance, NetUid), TransactionValidityError> { if alpha_vec.len() != 1 { - return (0.into(), 0.into(), NetUid::ROOT); + return Ok((0.into(), 0.into(), NetUid::ROOT)); } if let Some((hotkey, netuid)) = alpha_vec.first() { @@ -176,26 +178,33 @@ where // Sell alpha_fee and burn received tao (ignore unstake_from_subnet return). if let Some(author) = T::author() { - let swap_result = pallet_subtensor::Pallet::::unstake_from_subnet( - hotkey, - coldkey, - &author, - *netuid, - alpha_fee, - 0.into(), - true, - ); - if let Ok(tao_amount) = swap_result { - (alpha_fee, tao_amount, *netuid) - } else { - (0.into(), 0.into(), NetUid::ROOT) - } + with_transaction( + || -> TransactionOutcome> { + match pallet_subtensor::Pallet::::unstake_from_subnet( + hotkey, + coldkey, + &author, + *netuid, + alpha_fee, + 0.into(), + true, + ) { + Ok(tao_amount) => TransactionOutcome::Commit(Ok(tao_amount)), + Err(err) => TransactionOutcome::Rollback(Err(err)), + } + }, + ) + .map(|tao_amount| (alpha_fee, tao_amount, *netuid)) + .map_err(|err| { + log::error!("Error withdrawing transaction fee in alpha: {err:?}"); + InvalidTransaction::Payment.into() + }) } else { // Fallback: no author => no fees (do nothing) - (0.into(), 0.into(), NetUid::ROOT) + Ok((0.into(), 0.into(), NetUid::ROOT)) } } else { - (0.into(), 0.into(), NetUid::ROOT) + Ok((0.into(), 0.into(), NetUid::ROOT)) } } @@ -343,7 +352,7 @@ where if !alpha_vec.is_empty() { let fee_u64: u64 = fee.saturated_into::(); let (alpha_fee, tao_amount, netuid) = - OU::withdraw_in_alpha(who, &alpha_vec, fee_u64.into()); + OU::withdraw_in_alpha(who, &alpha_vec, fee_u64.into())?; return Ok(Some(WithdrawnFee::Alpha((alpha_fee, tao_amount, netuid)))); } Err(InvalidTransaction::Payment.into()) diff --git a/pallets/transaction-fee/src/tests/mod.rs b/pallets/transaction-fee/src/tests/mod.rs index d54c63a63b..8203070d76 100644 --- a/pallets/transaction-fee/src/tests/mod.rs +++ b/pallets/transaction-fee/src/tests/mod.rs @@ -3,6 +3,7 @@ use crate::{AlphaFeeHandler, SubtensorTxFeeHandler, TransactionFeeHandler, Trans use approx::assert_abs_diff_eq; use frame_support::dispatch::GetDispatchInfo; use frame_support::pallet_prelude::Zero; +use frame_support::traits::Currency; use frame_support::{assert_err, assert_ok}; use pallet_subtensor_swap::AlphaSqrtPrice; use sp_runtime::{ @@ -187,7 +188,7 @@ fn test_rejects_multi_subnet_alpha_fee_deduction() { &alpha_vec, 1.into(), ), - (0.into(), 0.into(), NetUid::ROOT) + Ok((0.into(), 0.into(), NetUid::ROOT)) ); let alpha_after_0 = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( @@ -326,6 +327,85 @@ fn test_remove_stake_fees_alpha() { }); } +#[test] +fn test_alpha_fee_withdraw_failure_aborts_and_rolls_back() { + new_test_ext().execute_with(|| { + let stake_amount = TAO; + let unstake_amount = AlphaBalance::from(TAO / 50); + let sn = setup_subnets(1, 1); + let netuid = sn.subnets[0].netuid; + let hotkey = sn.hotkeys[0]; + + setup_stake(netuid, &sn.coldkey, &hotkey, stake_amount); + + let current_balance = Balances::free_balance(sn.coldkey); + remove_balance_from_coldkey_account( + &sn.coldkey, + current_balance - ExistentialDeposit::get(), + ); + + // Force the alpha-fee unstake to fail after AMM bookkeeping by draining + // the subnet account used by transfer_tao_from_subnet. + let subnet_account = SubtensorModule::get_subnet_account_id(netuid).unwrap(); + Balances::make_free_balance_be(&subnet_account, 0.into()); + + let block_builder = U256::from(MOCK_BLOCK_BUILDER); + let block_builder_balance_before = Balances::free_balance(block_builder); + let signer_balance_before = Balances::free_balance(sn.coldkey); + let alpha_before = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, + &sn.coldkey, + netuid, + ); + let subnet_alpha_in_before = SubnetAlphaIn::::get(netuid); + let subnet_alpha_out_before = SubnetAlphaOut::::get(netuid); + let subnet_tao_before = SubnetTAO::::get(netuid); + let total_stake_before = TotalStake::::get(); + let subnet_volume_before = SubnetVolume::::get(netuid); + + let call = RuntimeCall::SubtensorModule(pallet_subtensor::Call::remove_stake { + hotkey, + netuid, + amount_unstaked: unstake_amount, + }); + let info = call.get_dispatch_info(); + let ext = pallet_transaction_payment::ChargeTransactionPayment::::from(0.into()); + + let result = + ext.dispatch_transaction(RuntimeOrigin::signed(sn.coldkey).into(), call, &info, 0, 0); + + assert_eq!( + result.unwrap_err(), + TransactionValidityError::Invalid(InvalidTransaction::Payment) + ); + assert_eq!( + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &hotkey, + &sn.coldkey, + netuid, + ), + alpha_before + ); + assert_eq!(SubnetAlphaIn::::get(netuid), subnet_alpha_in_before); + assert_eq!(SubnetAlphaOut::::get(netuid), subnet_alpha_out_before); + assert_eq!(SubnetTAO::::get(netuid), subnet_tao_before); + assert_eq!(TotalStake::::get(), total_stake_before); + assert_eq!(SubnetVolume::::get(netuid), subnet_volume_before); + assert_eq!(Balances::free_balance(sn.coldkey), signer_balance_before); + assert_eq!( + Balances::free_balance(block_builder), + block_builder_balance_before + ); + + assert!(!System::events().iter().any(|event_record| { + matches!( + &event_record.event, + RuntimeEvent::SubtensorModule(SubtensorEvent::TransactionFeePaidWithAlpha { .. }) + ) + })); + }); +} + // Test that unstaking on root with no free balance results in charging fees from // staked amount //