Fix TON STON.fi quote after reserve and jetton swap

crates/gem_ton/src/provider/preload.rs

@@ -5,7 +5,7 @@ use num_bigint::BigInt;
 use primitives::FeeOption;
 use primitives::{
     AssetSubtype, FeePriority, FeeRate, GasPriceType, TransactionFee, TransactionInputType, TransactionLoadData, TransactionLoadInput, TransactionLoadMetadata,
-    TransactionPreloadInput,
+    TransactionPreloadInput, swap::SwapQuoteDataType,
 };
 use std::collections::HashMap;
 use std::error::Error;
@@ -17,32 +17,29 @@ const TON_BASE_FEE: u64 = 10_000_000;
 const JETTON_ACCOUNT_FEE_EXISTING: u64 = 100_000_000;
 const JETTON_ACCOUNT_FEE_EXISTING_WITH_MEMO: u64 = 60_000_000;
 const JETTON_ACCOUNT_CREATION: u64 = 200_000_000;
+const SWAP_NATIVE_RESERVE: u64 = 310_000_000;
 
 pub fn calculate_transaction_fee(input: &TransactionLoadInput, recipient_token_address: Option<String>) -> TransactionFee {
     let base_fee = BigInt::from(TON_BASE_FEE);
     let mut options = HashMap::new();
 
     let fee = match &input.input_type {
-        TransactionInputType::Transfer(asset) | TransactionInputType::TransferNft(asset, _) | TransactionInputType::Account(asset, _) => match asset.id.token_subtype() {
-            AssetSubtype::NATIVE => base_fee.clone(),
-            AssetSubtype::TOKEN => {
-                let jetton_fee = if recipient_token_address.is_some() {
-                    if input.memo.is_some() {
-                        BigInt::from(JETTON_ACCOUNT_FEE_EXISTING_WITH_MEMO)
-                    } else {
-                        BigInt::from(JETTON_ACCOUNT_FEE_EXISTING)
-                    }
-                } else {
-                    BigInt::from(JETTON_ACCOUNT_CREATION)
-                };
-                options.insert(FeeOption::TokenAccountCreation, jetton_fee.clone());
+        TransactionInputType::Transfer(asset) | TransactionInputType::TransferNft(asset, _) | TransactionInputType::Account(asset, _) => {
+            transfer_fee(asset.id.token_subtype(), input.memo.as_deref(), recipient_token_address.as_deref(), &base_fee, &mut options)
+        }
+        TransactionInputType::Swap(from_asset, _, swap_data) => match &swap_data.data.data_type {
+            SwapQuoteDataType::Contract => {
+                options.insert(FeeOption::TokenAccountCreation, BigInt::from(SWAP_NATIVE_RESERVE));
                 base_fee
             }
+            SwapQuoteDataType::Transfer => transfer_fee(
+                from_asset.id.token_subtype(),
+                input.memo.as_deref(),
+                recipient_token_address.as_deref(),
+                &base_fee,
+                &mut options,
+            ),
         },
-        TransactionInputType::Swap(_, _, _) => {
-            options.insert(FeeOption::TokenAccountCreation, BigInt::from(JETTON_ACCOUNT_CREATION));
-            base_fee
-        }
         TransactionInputType::TokenApprove(_, _) => base_fee.clone(),
         TransactionInputType::Generic(_, _, _) => base_fee.clone(),
         TransactionInputType::Perpetual(_, _) => base_fee.clone(),
@@ -52,6 +49,25 @@ pub fn calculate_transaction_fee(input: &TransactionLoadInput, recipient_token_a
     TransactionFee::new_gas_price_type(GasPriceType::regular(fee.clone()), fee.clone(), BigInt::from(1), options)
 }
 
+fn transfer_fee(asset_subtype: AssetSubtype, memo: Option<&str>, recipient_token_address: Option<&str>, base_fee: &BigInt, options: &mut HashMap<FeeOption, BigInt>) -> BigInt {
+    match asset_subtype {
+        AssetSubtype::NATIVE => base_fee.clone(),
+        AssetSubtype::TOKEN => {
+            let jetton_fee = if recipient_token_address.is_some() {
+                if memo.is_some() {
+                    BigInt::from(JETTON_ACCOUNT_FEE_EXISTING_WITH_MEMO)
+                } else {
+                    BigInt::from(JETTON_ACCOUNT_FEE_EXISTING)
+                }
+            } else {
+                BigInt::from(JETTON_ACCOUNT_CREATION)
+            };
+            options.insert(FeeOption::TokenAccountCreation, jetton_fee);
+            base_fee.clone()
+        }
+    }
+}
+
 #[async_trait]
 impl<C: Client> ChainTransactionLoad for TonClient<C> {
     async fn get_transaction_preload(&self, input: TransactionPreloadInput) -> Result<TransactionLoadMetadata, Box<dyn Error + Sync + Send>> {
@@ -64,13 +80,19 @@ impl<C: Client> ChainTransactionLoad for TonClient<C> {
                 let token_id = asset.id.token_id.as_ref().ok_or("Missing token ID for jetton transaction")?;
                 let jetton_token_id = base64_to_hex_address(token_id).ok_or("Invalid jetton token ID")?.to_uppercase();
 
-                let (sender_jetton_wallets, recipient_jetton_wallets) = futures::try_join!(
-                    self.get_jetton_wallets(input.sender_address.clone()),
-                    self.get_jetton_wallets(input.destination_address.clone()), //TODO: different destination address for swaps
-                )?;
+                let sender_wallets = self.get_jetton_wallets(input.sender_address.clone());
+                let recipient_wallets = async {
+                    match get_recipient_jetton_wallet(&input) {
+                        Some(address) => self.get_jetton_wallets(address.to_string()).await.map(Some),
+                        None => Ok(None),
+                    }
+                };
+                let (sender_jetton_wallets, recipient_jetton_wallets) = futures::future::try_join(sender_wallets, recipient_wallets).await?;
 
                 let sender_jetton_wallet_address = sender_jetton_wallets.jetton_wallets.iter().find(|wallet| wallet.jetton == jetton_token_id);
-                let recipient_jetton_wallet_address = recipient_jetton_wallets.jetton_wallets.iter().find(|wallet| wallet.jetton == jetton_token_id);
+                let recipient_jetton_wallet_address = recipient_jetton_wallets
+                    .as_ref()
+                    .and_then(|wallets| wallets.jetton_wallets.iter().find(|wallet| wallet.jetton == jetton_token_id));
 
                 Ok(TransactionLoadMetadata::Ton {
                     sender_token_address: sender_jetton_wallet_address.map(|x| x.address.clone()),
@@ -99,11 +121,21 @@ impl<C: Client> ChainTransactionLoad for TonClient<C> {
     }
 }
 
+fn get_recipient_jetton_wallet(input: &TransactionPreloadInput) -> Option<&str> {
+    match &input.input_type {
+        TransactionInputType::Swap(_, _, swap_data) => match &swap_data.data.data_type {
+            SwapQuoteDataType::Transfer => Some(&swap_data.data.to),
+            SwapQuoteDataType::Contract => None,
+        },
+        _ => Some(&input.destination_address),
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use num_bigint::BigInt;
-    use primitives::{Asset, AssetId, AssetType, Chain, GasPriceType};
+    use primitives::{Asset, AssetId, AssetType, Chain, GasPriceType, SwapProvider, TransactionPreloadInput, swap::SwapData};
 
     fn create_input(asset_type: AssetType, memo: Option<String>) -> TransactionLoadInput {
         let (token_id, name, symbol, decimals) = match asset_type {
@@ -183,4 +215,94 @@ mod tests {
         assert_eq!(fee.fee, BigInt::from(TON_BASE_FEE + JETTON_ACCOUNT_CREATION));
         assert_eq!(fee.options.get(&FeeOption::TokenAccountCreation), Some(&BigInt::from(JETTON_ACCOUNT_CREATION)));
     }
+
+    #[test]
+    fn test_swap_contract_native_fee_includes_native_reserve() {
+        let swap_data = SwapData::mock_contract(SwapProvider::StonfiV2, "400000000", "1000000", "710000000");
+        let input = TransactionLoadInput {
+            input_type: TransactionInputType::Swap(Asset::from_chain(Chain::Ton), Asset::from_chain(Chain::Ton), swap_data),
+            value: "400000000".to_string(),
+            ..create_input(AssetType::NATIVE, None)
+        };
+
+        let fee = calculate_transaction_fee(&input, None);
+
+        assert_eq!(fee.fee, BigInt::from(320000000u64));
+        assert_eq!(fee.options.get(&FeeOption::TokenAccountCreation), Some(&BigInt::from(310000000u64)));
+    }
+
+    #[test]
+    fn test_swap_contract_jetton_fee_includes_native_reserve() {
+        let from_asset = Asset::mock_ton_usdt();
+        let swap_data = SwapData::mock_contract(SwapProvider::StonfiV2, "2000000", "400000000", "300000000");
+        let input = TransactionLoadInput {
+            input_type: TransactionInputType::Swap(from_asset, Asset::from_chain(Chain::Ton), swap_data),
+            value: "2000000".to_string(),
+            ..create_input(AssetType::JETTON, None)
+        };
+
+        let fee = calculate_transaction_fee(&input, None);
+
+        assert_eq!(fee.fee, BigInt::from(TON_BASE_FEE + SWAP_NATIVE_RESERVE));
+        assert_eq!(fee.options.get(&FeeOption::TokenAccountCreation), Some(&BigInt::from(SWAP_NATIVE_RESERVE)));
+    }
+
+    #[test]
+    fn test_swap_transfer_native_fee_uses_transfer_fee() {
+        let swap_data = SwapData::mock_transfer(SwapProvider::NearIntents, "400000000", "1000000", "ton_deposit_address");
+        let input = TransactionLoadInput {
+            input_type: TransactionInputType::Swap(Asset::from_chain(Chain::Ton), Asset::from_chain(Chain::Near), swap_data),
+            value: "400000000".to_string(),
+            ..create_input(AssetType::NATIVE, None)
+        };
+
+        let fee = calculate_transaction_fee(&input, None);
+
+        assert_eq!(fee.fee, BigInt::from(TON_BASE_FEE));
+        assert_eq!(fee.options.len(), 0);
+    }
+
+    #[test]
+    fn test_swap_transfer_jetton_fee_uses_token_transfer_fee() {
+        let swap_data = SwapData::mock_transfer(SwapProvider::NearIntents, "2000000", "400000000", "ton_deposit_address");
+        let input = TransactionLoadInput {
+            input_type: TransactionInputType::Swap(Asset::mock_ton_usdt(), Asset::from_chain(Chain::Near), swap_data),
+            value: "2000000".to_string(),
+            ..create_input(AssetType::JETTON, None)
+        };
+
+        let fee = calculate_transaction_fee(&input, None);
+
+        assert_eq!(fee.fee, BigInt::from(TON_BASE_FEE + JETTON_ACCOUNT_CREATION));
+        assert_eq!(fee.options.get(&FeeOption::TokenAccountCreation), Some(&BigInt::from(JETTON_ACCOUNT_CREATION)));
+    }
+
+    #[test]
+    fn test_get_recipient_jetton_wallet() {
+        let transfer = TransactionPreloadInput {
+            input_type: TransactionInputType::Transfer(Asset::mock_ton_usdt()),
+            sender_address: "sender".to_string(),
+            destination_address: "recipient".to_string(),
+        };
+        assert_eq!(get_recipient_jetton_wallet(&transfer), Some("recipient"));
+
+        let swap_data = SwapData::mock_transfer(SwapProvider::NearIntents, "2000000", "400000000", "ton_deposit_address");
+        let transfer_swap = TransactionPreloadInput {
+            input_type: TransactionInputType::Swap(Asset::mock_ton_usdt(), Asset::from_chain(Chain::Ethereum), swap_data),
+            sender_address: "sender".to_string(),
+            destination_address: "0xrecipient".to_string(),
+        };
+        assert_eq!(get_recipient_jetton_wallet(&transfer_swap), Some("ton_deposit_address"));
+
+        let contract_swap = TransactionPreloadInput {
+            input_type: TransactionInputType::Swap(
+                Asset::mock_ton_usdt(),
+                Asset::from_chain(Chain::Ton),
+                SwapData::mock_contract(SwapProvider::StonfiV2, "2000000", "400000000", "300000000"),
+            ),
+            sender_address: "sender".to_string(),
+            destination_address: "recipient".to_string(),
+        };
+        assert_eq!(get_recipient_jetton_wallet(&contract_swap), None);
+    }
 }

crates/gem_ton/src/rpc/client.rs

@@ -11,6 +11,7 @@ use crate::models::{
     SimpleJettonBalance, WalletInfo,
 };
 
+#[derive(Debug)]
 pub struct TonClient<C: Client> {
     pub client: C,
 }

crates/gem_ton/src/signer/transaction/sign.rs

@@ -175,11 +175,8 @@ mod tests {
         swap_data.data.value = "241000000".to_string();
         swap_data.data.data = mock_cell();
         swap_data.data.gas_limit = None;
-        let input = SignerInput::mock_with_input_type(
+        let input = SignerInput::mock_ton(
             TransactionInputType::Swap(Asset::from_chain(Chain::Ton), Asset::from_chain(Chain::Ton), swap_data),
-            "",
-            "",
-            "0",
             TransactionLoadMetadata::mock_ton(1),
         );
 

crates/primitives/src/testkit/asset_mock.rs

@@ -1,6 +1,6 @@
 use crate::{
     Asset, AssetId, AssetType, Chain,
-    asset_constants::{ETHEREUM_USDC_ASSET_ID, SOLANA_USDC_ASSET_ID},
+    asset_constants::{ETHEREUM_USDC_ASSET_ID, SOLANA_USDC_ASSET_ID, TON_USDT_ASSET_ID},
 };
 
 impl Asset {
@@ -20,6 +20,10 @@ impl Asset {
         Asset::new(ETHEREUM_USDC_ASSET_ID.clone(), "USD Coin".to_string(), "USDC".to_string(), 6, AssetType::ERC20)
     }
 
+    pub fn mock_ton_usdt() -> Self {
+        Asset::new(TON_USDT_ASSET_ID.clone(), "Tether USD".to_string(), "USDT".to_string(), 6, AssetType::JETTON)
+    }
+
     pub fn mock_eth() -> Self {
         Asset::from_chain(Chain::Ethereum)
     }

crates/primitives/src/testkit/swap_mock.rs

@@ -66,6 +66,25 @@ impl SwapData {
             data: SwapQuoteData::mock(),
         }
     }
+
+    pub fn mock_contract(provider: SwapProvider, from_value: &str, to_value: &str, value: &str) -> Self {
+        let swap_data = Self::mock_with_values(provider, from_value, to_value);
+        SwapData {
+            data: SwapQuoteData {
+                value: value.to_string(),
+                ..swap_data.data
+            },
+            ..swap_data
+        }
+    }
+
+    pub fn mock_transfer(provider: SwapProvider, from_value: &str, to_value: &str, to: &str) -> Self {
+        let swap_data = Self::mock_with_values(provider, from_value, to_value);
+        SwapData {
+            data: SwapQuoteData::new_tranfer(to.to_string(), from_value.to_string(), None),
+            ..swap_data
+        }
+    }
 }
 
 impl SwapQuote {

crates/primitives/src/testkit/transaction_load_input_mock.rs

@@ -97,6 +97,22 @@ impl SignerInput {
             TransactionFee::default(),
         )
     }
+
+    pub fn mock_ton(input_type: TransactionInputType, metadata: TransactionLoadMetadata) -> Self {
+        SignerInput::new(
+            TransactionLoadInput {
+                input_type,
+                sender_address: "UQBY1cVPu4SIr36q0M3HWcqPb_efyVVRBsEzmwN-wKQDR6zg".to_string(),
+                destination_address: "UQBY1cVPu4SIr36q0M3HWcqPb_efyVVRBsEzmwN-wKQDR6zg".to_string(),
+                value: "10000".to_string(),
+                gas_price: GasPriceType::regular(0),
+                memo: None,
+                is_max_value: false,
+                metadata,
+            },
+            TransactionFee::default(),
+        )
+    }
 }
 
 impl TransactionLoadInput {

crates/swapper/src/chainflip/provider.rs

@@ -20,7 +20,7 @@ use crate::{
     amount_to_value,
     approval::{check_approval_erc20, get_swap_gas_limit_with_approval},
     cross_chain::VaultAddresses,
-    fees::{DEFAULT_CHAINFLIP_FEE_BPS, apply_slippage_in_bp, resolve_max_quote_value},
+    fees::{DEFAULT_CHAINFLIP_FEE_BPS, apply_slippage_in_bp, quote_value_after_reserve_by_chain},
     solana::DEFAULT_SWAP_GAS_LIMIT,
 };
 use primitives::{ChainType, chain::Chain, swap::QuoteAsset};
@@ -70,7 +70,7 @@ where
     }
 
     fn get_quote_value(request: &QuoteRequest) -> Result<String, SwapperError> {
-        let value = resolve_max_quote_value(request)?;
+        let value = quote_value_after_reserve_by_chain(request)?;
         if !request.options.use_max_amount || !request.from_asset.asset_id().is_native() {
             return Ok(value);
         }

crates/swapper/src/fees/mod.rs

@@ -3,7 +3,7 @@ mod reserve;
 mod slippage;
 
 pub use referral::{ReferralFee, ReferralFees, default_referral_fees};
-pub use reserve::{RESERVED_NATIVE_FEES, reserved_tx_fees, resolve_max_quote_value};
+pub use reserve::{RESERVED_NATIVE_FEES, quote_value_after_reserve, quote_value_after_reserve_by_chain, reserved_tx_fees};
 pub use slippage::{BasisPointConvert, apply_slippage_in_bp};
 
 pub const DEFAULT_SWAP_FEE_BPS: u32 = 50;

crates/swapper/src/fees/reserve.rs

@@ -45,13 +45,10 @@ pub fn reserved_tx_fees(chain: Chain) -> Option<&'static str> {
     RESERVED_NATIVE_FEES.get(&chain).copied()
 }
 
-pub fn resolve_max_quote_value(request: &QuoteRequest) -> Result<String, SwapperError> {
+pub fn quote_value_after_reserve(request: &QuoteRequest, reserved: &str) -> Result<String, SwapperError> {
     if !request.options.use_max_amount || !request.from_asset.asset_id().is_native() {
         return Ok(request.value.clone());
     }
-    let Some(reserved) = reserved_tx_fees(request.from_asset.chain()) else {
-        return Ok(request.value.clone());
-    };
     let reserved_fee = U256::from_str(reserved).map_err(|_| SwapperError::ComputeQuoteError(format!("invalid reserved fee: {reserved}")))?;
     let amount = U256::from_str(&request.value).map_err(|_| SwapperError::ComputeQuoteError(format!("invalid amount: {}", request.value)))?;
     if amount <= reserved_fee {
@@ -61,3 +58,10 @@ pub fn resolve_max_quote_value(request: &QuoteRequest) -> Result<String, Swapper
     }
     Ok((amount - reserved_fee).to_string())
 }
+
+pub fn quote_value_after_reserve_by_chain(request: &QuoteRequest) -> Result<String, SwapperError> {
+    let Some(reserved) = reserved_tx_fees(request.from_asset.chain()) else {
+        return Ok(request.value.clone());
+    };
+    quote_value_after_reserve(request, reserved)
+}