feat(bot): token2022 support in cleanup flow

Author: smypmsa

src/cleanup/manager.py

@@ -34,12 +34,19 @@ def __init__(
         self.use_priority_fee = use_priority_fee
         self.close_with_force_burn = force_burn
 
-    async def cleanup_ata(self, mint: Pubkey) -> None:
+    async def cleanup_ata(self, mint: Pubkey, token_program_id: Pubkey | None = None) -> None:
         """
         Attempt to burn any remaining tokens and close the ATA.
         Skips if account doesn't exist or is already empty/closed.
+
+        Args:
+            mint: Token mint address
+            token_program_id: Token program (TOKEN or TOKEN_2022). Defaults to TOKEN_2022_PROGRAM
         """
-        ata = self.wallet.get_associated_token_address(mint)
+        if token_program_id is None:
+            token_program_id = SystemAddresses.TOKEN_2022_PROGRAM
+
+        ata = self.wallet.get_associated_token_address(mint, token_program_id)
         solana_client = await self.client.get_client()
 
         priority_fee = (
@@ -70,7 +77,7 @@ async def cleanup_ata(self, mint: Pubkey) -> None:
                         mint=mint,
                         owner=self.wallet.pubkey,
                         amount=balance,
-                        program_id=SystemAddresses.TOKEN_PROGRAM,
+                        program_id=token_program_id,
                     )
                 )
                 instructions.append(burn_ix)
@@ -89,7 +96,7 @@ async def cleanup_ata(self, mint: Pubkey) -> None:
                     account=ata,
                     dest=self.wallet.pubkey,
                     owner=self.wallet.pubkey,
-                    program_id=SystemAddresses.TOKEN_PROGRAM,
+                    program_id=token_program_id,
                 )
             )
             instructions.append(close_ix)

src/cleanup/modes.py

@@ -20,6 +20,7 @@ async def handle_cleanup_after_failure(
     client,
     wallet,
     mint,
+    token_program_id,
     priority_fee_manager,
     cleanup_mode,
     cleanup_with_prior_fee,
@@ -30,13 +31,14 @@ async def handle_cleanup_after_failure(
         manager = AccountCleanupManager(
             client, wallet, priority_fee_manager, cleanup_with_prior_fee, force_burn
         )
-        await manager.cleanup_ata(mint)
+        await manager.cleanup_ata(mint, token_program_id)
 
 
 async def handle_cleanup_after_sell(
     client,
     wallet,
     mint,
+    token_program_id,
     priority_fee_manager,
     cleanup_mode,
     cleanup_with_prior_fee,
@@ -47,13 +49,14 @@ async def handle_cleanup_after_sell(
         manager = AccountCleanupManager(
             client, wallet, priority_fee_manager, cleanup_with_prior_fee, force_burn
         )
-        await manager.cleanup_ata(mint)
+        await manager.cleanup_ata(mint, token_program_id)
 
 
 async def handle_cleanup_post_session(
     client,
     wallet,
     mints,
+    token_program_ids,
     priority_fee_manager,
     cleanup_mode,
     cleanup_with_prior_fee,
@@ -64,5 +67,5 @@ async def handle_cleanup_post_session(
         manager = AccountCleanupManager(
             client, wallet, priority_fee_manager, cleanup_with_prior_fee, force_burn
         )
-        for mint in mints:
-            await manager.cleanup_ata(mint)
+        for mint, token_program_id in zip(mints, token_program_ids):
+            await manager.cleanup_ata(mint, token_program_id)

src/core/wallet.py

@@ -7,6 +7,8 @@
 from solders.pubkey import Pubkey
 from spl.token.instructions import get_associated_token_address
 
+from core.pubkeys import SystemAddresses
+
 
 class Wallet:
     """Manages a Solana wallet for trading operations."""
@@ -30,16 +32,21 @@ def keypair(self) -> Keypair:
         """Get the keypair for signing transactions."""
         return self._keypair
 
-    def get_associated_token_address(self, mint: Pubkey) -> Pubkey:
+    def get_associated_token_address(
+        self, mint: Pubkey, token_program_id: Pubkey | None = None
+    ) -> Pubkey:
         """Get the associated token account address for a mint.
 
         Args:
             mint: Token mint address
+            token_program_id: Token program (TOKEN or TOKEN_2022). Defaults to TOKEN_2022_PROGRAM
 
         Returns:
             Associated token account address
         """
-        return get_associated_token_address(self.pubkey, mint)
+        if token_program_id is None:
+            token_program_id = SystemAddresses.TOKEN_2022_PROGRAM
+        return get_associated_token_address(self.pubkey, mint, token_program_id)
 
     @staticmethod
     def _load_keypair(private_key: str) -> Keypair:

src/trading/universal_trader.py

@@ -193,6 +193,7 @@ def __init__(
 
         # State tracking
         self.traded_mints: set[Pubkey] = set()
+        self.traded_token_programs: dict[str, Pubkey] = {}  # Maps mint (as string) to token_program_id
         self.token_queue: asyncio.Queue = asyncio.Queue()
         self.processing: bool = False
         self.processed_tokens: set[str] = set()
@@ -325,10 +326,17 @@ async def _cleanup_resources(self) -> None:
         if self.traded_mints:
             try:
                 logger.info(f"Cleaning up {len(self.traded_mints)} traded token(s)...")
+                # Build parallel lists of mints and token_program_ids
+                mints_list = list(self.traded_mints)
+                token_program_ids = [
+                    self.traded_token_programs.get(str(mint))
+                    for mint in mints_list
+                ]
                 await handle_cleanup_post_session(
                     self.solana_client,
                     self.wallet,
-                    list(self.traded_mints),
+                    mints_list,
+                    token_program_ids,
                     self.priority_fee_manager,
                     self.cleanup_mode,
                     self.cleanup_with_priority_fee,
@@ -447,6 +455,10 @@ async def _handle_successful_buy(
             buy_result.tx_signature,
         )
         self.traded_mints.add(token_info.mint)
+        # Track token program for cleanup
+        mint_str = str(token_info.mint)
+        if token_info.token_program_id:
+            self.traded_token_programs[mint_str] = token_info.token_program_id
 
         # Choose exit strategy
         if not self.marry_mode:
@@ -469,6 +481,7 @@ async def _handle_failed_buy(
             self.solana_client,
             self.wallet,
             token_info.mint,
+            token_info.token_program_id,
             self.priority_fee_manager,
             self.cleanup_mode,
             self.cleanup_with_priority_fee,
@@ -521,6 +534,7 @@ async def _handle_time_based_exit(self, token_info: TokenInfo) -> None:
                 self.solana_client,
                 self.wallet,
                 token_info.mint,
+                token_info.token_program_id,
                 self.priority_fee_manager,
                 self.cleanup_mode,
                 self.cleanup_with_priority_fee,
@@ -590,6 +604,7 @@ async def _monitor_position_until_exit(
                             self.solana_client,
                             self.wallet,
                             token_info.mint,
+                            token_info.token_program_id,
                             self.priority_fee_manager,
                             self.cleanup_mode,
                             self.cleanup_with_priority_fee,