Merge pull request #77 from chainstacklabs/feature/idempotent-ata-and-merge-burn-close

Optimize token handling: idempotent ATA for buying and combined burn & close instructions

src/cleanup/manager.py

@@ -51,6 +51,8 @@ class AccountCleanupManager:
                 return
 
             balance = await self.client.get_token_account_balance(ata)
+            instructions = []
+
             if balance > 0 and CLEANUP_FORCE_CLOSE_WITH_BURN:
                 logger.info(f"Burning {balance} tokens from ATA {ata} (mint: {mint})...")
                 burn_ix = burn(
@@ -62,13 +64,8 @@ class AccountCleanupManager:
                         program_id=SystemAddresses.TOKEN_PROGRAM,
                     )
                 )
-                await self.client.build_and_send_transaction(
-                    [burn_ix],
-                    self.wallet.keypair,
-                    skip_preflight=True,
-                    priority_fee=priority_fee,
-                )
-                logger.info(f"✅ Burned successfully from ATA {ata}")
+                instructions.append(burn_ix)
+
             elif balance > 0:
                 logger.info(
                     f"Skipping ATA {ata} with non-zero balance ({balance} tokens) "
@@ -76,6 +73,7 @@ class AccountCleanupManager:
                 )
                 return
 
+            # Include close account instruction
             logger.info(f"Closing ATA: {ata}")
             close_ix = close_account(
                 CloseAccountParams(
@@ -85,14 +83,18 @@ class AccountCleanupManager:
                     program_id=SystemAddresses.TOKEN_PROGRAM,
                 )
             )
-            tx_sig = await self.client.build_and_send_transaction(
-                [close_ix],
-                self.wallet.keypair,
-                skip_preflight=True,
-                priority_fee=priority_fee,
-            )
-            await self.client.confirm_transaction(tx_sig)
-            logger.info(f"Closed successfully: {ata}")
+            instructions.append(close_ix)
+
+            # Send both burn and close instructions in the same transaction
+            if instructions:
+                tx_sig = await self.client.build_and_send_transaction(
+                    instructions,
+                    self.wallet.keypair,
+                    skip_preflight=True,
+                    priority_fee=priority_fee,
+                )
+                await self.client.confirm_transaction(tx_sig)
+                logger.info(f"Closed successfully: {ata}")
 
         except Exception as e:
             logger.warning(f"Cleanup failed for ATA {ata}: {e!s}")

src/trading/buyer.py

@@ -7,7 +7,7 @@ from typing import Final
 
 from solders.instruction import AccountMeta, Instruction
 from solders.pubkey import Pubkey
-from spl.token.instructions import create_associated_token_account
+from spl.token.instructions import create_idempotent_associated_token_account
 
 from core.client import SolanaClient
 from core.curve import BondingCurveManager
@@ -93,10 +93,6 @@ class TokenBuyer(Trader):
                 token_info.mint
             )
 
-            await self._ensure_associated_token_account(
-                token_info.mint, associated_token_account
-            )
-
             tx_signature = await self._send_buy_transaction(
                 token_info,
                 associated_token_account,
@@ -124,51 +120,6 @@ class TokenBuyer(Trader):
             logger.error(f"Buy operation failed: {e!s}")
             return TradeResult(success=False, error_message=str(e))
 
-    async def _ensure_associated_token_account(
-        self, mint: Pubkey, associated_token_account: Pubkey
-    ) -> None:
-        """Ensure associated token account exists, else create it.
-
-        Args:
-            mint: Token mint
-            associated_token_account: Associated token account address
-        """
-        try:
-            solana_client = await self.client.get_client()
-            account_info = await solana_client.get_account_info(
-                associated_token_account, encoding="base64"
-            )
-
-            if account_info.value is None:
-                logger.info(f"Creating associated token account for {mint}...")
-
-                create_ata_ix = create_associated_token_account(
-                    payer=self.wallet.pubkey, owner=self.wallet.pubkey, mint=mint
-                )
-
-                tx_sig = await self.client.build_and_send_transaction(
-                    [create_ata_ix],
-                    self.wallet.keypair,
-                    skip_preflight=True,
-                    max_retries=self.max_retries,
-                    priority_fee=await self.priority_fee_manager.calculate_priority_fee(
-                        [mint, SystemAddresses.PROGRAM, SystemAddresses.TOKEN_PROGRAM]
-                    ),
-                )
-
-                await self.client.confirm_transaction(tx_sig)
-                logger.info(
-                    f"Associated token account created: {associated_token_account}"
-                )
-            else:
-                logger.info(
-                    f"Associated token account already exists: {associated_token_account}"
-                )
-
-        except Exception as e:
-            logger.error(f"Error creating associated token account: {e!s}")
-            raise
-
     async def _send_buy_transaction(
         self,
         token_info: TokenInfo,
@@ -225,6 +176,14 @@ class TokenBuyer(Trader):
             ),
         ]
 
+        # Prepare idempotent create ATA instruction: it will not fail if ATA already exists
+        idempotent_ata_ix = create_idempotent_associated_token_account(
+            self.wallet.pubkey,
+            self.wallet.pubkey,
+            token_info.mint,
+            SystemAddresses.TOKEN_PROGRAM
+        )
+
         # Prepare buy instruction data
         token_amount_raw = int(token_amount * 10**TOKEN_DECIMALS)
         data = (
@@ -236,7 +195,7 @@ class TokenBuyer(Trader):
 
         try:
             return await self.client.build_and_send_transaction(
-                [buy_ix],
+                [idempotent_ata_ix, buy_ix],
                 self.wallet.keypair,
                 skip_preflight=True,
                 max_retries=self.max_retries,