Merge commit 'a0579a338b2c42e1a2706f7a42550a60e8611a22'

src/grpc/yellow_stone.rs

@@ -480,6 +480,30 @@ impl YellowstoneGrpc {
             return Ok(());
         }
 
+        let mut buy_instruction_events = vec![];
+        let mut sell_instruction_events = vec![];
+
+        if let Some(versioned_tx) = trade_raw.transaction.decode() {
+            let signature = versioned_tx.signatures[0].to_string();
+            let instructions: Vec<crate::common::pumpswap::logs_data::PumpSwapInstruction> =
+                crate::common::pumpswap::logs_filters::LogFilter::parse_pumpswap_compiled_instruction(versioned_tx).unwrap();
+            for instruction in instructions {
+                match instruction {
+                    crate::common::pumpswap::logs_data::PumpSwapInstruction::Buy(mut e) => {
+                        e.signature = signature.clone();
+                        e.slot = slot;
+                        buy_instruction_events.push(e);
+                    }
+                    crate::common::pumpswap::logs_data::PumpSwapInstruction::Sell(mut e) => {
+                        e.signature = signature.clone();
+                        e.slot = slot;
+                        sell_instruction_events.push(e);
+                    }
+                    _ => {}
+                }
+            }
+        }
+
         // 获取日志
         let logs = if let solana_transaction_status::option_serializer::OptionSerializer::Some(logs) = &meta.log_messages {
             logs
@@ -497,10 +521,26 @@ impl YellowstoneGrpc {
                 crate::common::pumpswap::logs_events::PumpSwapEvent::Buy(e) => {
                     e.signature = transaction_pretty.signature.to_string();
                     e.slot = slot;
+                    if let Some(ie) = buy_instruction_events.iter().find(|ie| ie.signature == e.signature && ie.slot == e.slot && ie.pool == e.pool && ie.user == e.user) {
+                        e.base_mint = ie.base_mint;
+                        e.quote_mint = ie.quote_mint;
+                        e.pool_base_token_account = ie.pool_base_token_account;
+                        e.pool_quote_token_account = ie.pool_quote_token_account;
+                        e.coin_creator_vault_ata = ie.coin_creator_vault_ata;
+                        e.coin_creator_vault_authority = ie.coin_creator_vault_authority;
+                    }
                 },
                 crate::common::pumpswap::logs_events::PumpSwapEvent::Sell(e) => {
                     e.signature = transaction_pretty.signature.to_string();
                     e.slot = slot;
+                    if let Some(ie) = sell_instruction_events.iter().find(|ie| ie.signature == e.signature && ie.slot == e.slot && ie.pool == e.pool && ie.user == e.user) {
+                        e.base_mint = ie.base_mint;
+                        e.quote_mint = ie.quote_mint;
+                        e.pool_base_token_account = ie.pool_base_token_account;
+                        e.pool_quote_token_account = ie.pool_quote_token_account;
+                        e.coin_creator_vault_ata = ie.coin_creator_vault_ata;
+                        e.coin_creator_vault_authority = ie.coin_creator_vault_authority;
+                    }
                 },
                 crate::common::pumpswap::logs_events::PumpSwapEvent::CreatePool(e) => {
                     e.signature = transaction_pretty.signature.to_string();

src/lib.rs

@@ -267,6 +267,7 @@ impl PumpFun {
                 None,
                 None,
                 None,
+                true,
             ).await
         } else {
             Err(anyhow::anyhow!("Unsupported trade platform: {}", trade_platform))
@@ -350,6 +351,7 @@ impl PumpFun {
                 protocol_params.pool_quote_token_account.clone(),
                 protocol_params.user_base_token_account.clone(),
                 protocol_params.user_quote_token_account.clone(),
+                protocol_params.auto_handle_wsol,
             ).await
         } else {
             return Err(anyhow::anyhow!("Invalid protocol params for PumpFun"));
@@ -403,6 +405,7 @@ impl PumpFun {
                 None,
                 None,
                 None,
+                true,
             ).await
         } else {
             Err(anyhow::anyhow!("Unsupported trade platform: {}", trade_platform))

src/pumpswap/buy.rs

@@ -27,6 +27,7 @@ pub async fn buy(
     pool_quote_token_account: Option<Pubkey>,
     user_base_token_account: Option<Pubkey>,
     user_quote_token_account: Option<Pubkey>,
+    auto_handle_wsol: bool,
 ) -> Result<(), anyhow::Error> {
     // 创建执行器
     let executor = TradeFactory::create_executor(Protocol::PumpSwap);
@@ -37,6 +38,7 @@ pub async fn buy(
         pool_quote_token_account: pool_quote_token_account,
         user_base_token_account: user_base_token_account,
         user_quote_token_account: user_quote_token_account,
+        auto_handle_wsol: auto_handle_wsol,
     });
     // 创建买入参数
     let buy_params = BuyParams {
@@ -75,6 +77,7 @@ pub async fn buy_with_tip(
     pool_quote_token_account: Option<Pubkey>,
     user_base_token_account: Option<Pubkey>,
     user_quote_token_account: Option<Pubkey>,
+    auto_handle_wsol: bool,
 ) -> Result<(), anyhow::Error> {
     // 创建执行器
     let executor = TradeFactory::create_executor(Protocol::PumpSwap);
@@ -85,6 +88,7 @@ pub async fn buy_with_tip(
         pool_quote_token_account: pool_quote_token_account,
         user_base_token_account: user_base_token_account,
         user_quote_token_account: user_quote_token_account,
+        auto_handle_wsol: auto_handle_wsol,
     });
     // 创建买入参数
     let buy_params = BuyParams {

src/pumpswap/sell.rs

@@ -34,6 +34,7 @@ pub async fn sell(
         pool_quote_token_account,
         user_base_token_account,
         user_quote_token_account,
+        auto_handle_wsol: true,
     });
     // 创建卖出参数
     let sell_params = SellParams {
@@ -163,6 +164,7 @@ pub async fn sell_with_tip(
         pool_quote_token_account,
         user_base_token_account,
         user_quote_token_account,
+        auto_handle_wsol: true,
     });
     // 创建卖出参数
     let sell_params = SellParams {

src/trading/core/params.rs

@@ -109,6 +109,7 @@ pub struct PumpSwapParams {
     pub pool_quote_token_account: Option<Pubkey>,
     pub user_base_token_account: Option<Pubkey>,
     pub user_quote_token_account: Option<Pubkey>,
+    pub auto_handle_wsol: bool,
 }
 
 impl ProtocolParams for PumpSwapParams {

src/trading/protocols/pumpswap.rs

@@ -58,6 +58,7 @@ impl InstructionBuilder for PumpSwapInstructionBuilder {
                     *pool_quote_token_account,
                     *user_base_token_account,
                     *user_quote_token_account,
+                    protocol_params.auto_handle_wsol,
                 )
                 .await
             }
@@ -148,6 +149,7 @@ impl PumpSwapInstructionBuilder {
             pool_quote_token_account,
             user_base_token_account,
             user_quote_token_account,
+            true,
         )
         .await
     }
@@ -210,6 +212,7 @@ impl PumpSwapInstructionBuilder {
         pool_quote_token_account: Pubkey,
         user_base_token_account: Pubkey,
         user_quote_token_account: Pubkey,
+        auto_handle_wsol: bool,
     ) -> Result<Vec<Instruction>> {
         if params.rpc.is_none() {
             return Err(anyhow!("RPC is not set"));
@@ -226,6 +229,36 @@ impl PumpSwapInstructionBuilder {
 
         let mut instructions = vec![];
 
+        if auto_handle_wsol {
+            // 插入wsol
+            instructions.push(
+                // 创建wSOL ATA账户，如果不存在
+                create_associated_token_account_idempotent(
+                    &params.payer.pubkey(),
+                    &params.payer.pubkey(),
+                    &accounts::WSOL_TOKEN_ACCOUNT,
+                    &accounts::TOKEN_PROGRAM,
+                ),
+            );
+            instructions.push(
+                // 将SOL转入wSOL ATA账户
+                solana_sdk::system_instruction::transfer(
+                    &params.payer.pubkey(),
+                    &user_quote_token_account,
+                    max_sol_amount,
+                ),
+            );
+
+            // 同步wSOL余额
+            instructions.push(
+                spl_token::instruction::sync_native(
+                    &accounts::TOKEN_PROGRAM,
+                    &user_quote_token_account,
+                )
+                .unwrap(),
+            );
+        }
+
         // 创建用户的基础代币账户
         instructions.push(create_associated_token_account_idempotent(
             &params.payer.pubkey(),
@@ -275,6 +308,20 @@ impl PumpSwapInstructionBuilder {
             data,
         });
 
+        if auto_handle_wsol {
+            // 关闭wSOL ATA账户，回收租金
+            instructions.push(
+                spl_token::instruction::close_account(
+                    &accounts::TOKEN_PROGRAM,
+                    &user_quote_token_account,
+                    &params.payer.pubkey(),
+                    &params.payer.pubkey(),
+                    &[],
+                )
+                .unwrap(),
+            );
+        }
+
         Ok(instructions)
     }