feat: added cleanup manager and modes

src/cleanup/manager.py

@@ -0,0 +1,83 @@
+from solders.pubkey import Pubkey
+from spl.token.instructions import BurnParams, CloseAccountParams, burn, close_account
+
+from config import CLEANUP_WITHOUT_PRIORITY_FEE
+from core.client import SolanaClient
+from core.pubkeys import SystemAddresses
+from core.wallet import Wallet
+from utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class AccountCleanupManager:
+    """Handles safe cleanup of token accounts (ATA) after trading sessions."""
+
+    def __init__(
+        self,
+        client: SolanaClient,
+        wallet: Wallet,
+    ):
+        """
+        Args:
+            client: Solana RPC client
+            wallet: Wallet for signing transactions
+        """
+        self.client = client
+        self.wallet = wallet
+        self.use_priority_fee = not CLEANUP_WITHOUT_PRIORITY_FEE
+
+    async def cleanup_ata(self, mint: Pubkey) -> None:
+        """
+        Attempt to burn any remaining tokens and close the ATA.
+        Skips if account doesn't exist or is already empty/closed.
+        """
+        ata = self.wallet.get_associated_token_address(mint)
+        solana_client = await self.client.get_client()
+
+        try:
+            info = await solana_client.get_account_info(ata)
+            if not info.value:
+                logger.info(f"ATA {ata} does not exist or already closed.")
+                return
+
+            balance = await self.client.get_token_account_balance(ata)
+            if balance > 0:
+                logger.info(f"⚠️ Burning {balance} tokens from ATA {ata} (mint: {mint})...")
+                burn_ix = burn(
+                    BurnParams(
+                        account=ata,
+                        mint=mint,
+                        owner=self.wallet.pubkey,
+                        amount=balance,
+                        program_id=SystemAddresses.TOKEN_PROGRAM,
+                    )
+                )
+                await self.client.build_and_send_transaction(
+                    [burn_ix],
+                    self.wallet.keypair,
+                    skip_preflight=True,
+                    priority_fee=None if not self.use_priority_fee else 0,
+                )
+                logger.info(f"✅ Burned successfully from ATA {ata}")
+
+            logger.info(f"Closing ATA: {ata}")
+            close_ix = close_account(
+                CloseAccountParams(
+                    account=ata,
+                    dest=self.wallet.pubkey,
+                    owner=self.wallet.pubkey,
+                    program_id=SystemAddresses.TOKEN_PROGRAM,
+                )
+            )
+            tx_sig = await self.client.build_and_send_transaction(
+                [close_ix],
+                self.wallet.keypair,
+                skip_preflight=True,
+                priority_fee=None if not self.use_priority_fee else 0,
+            )
+            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/cleanup/modes.py

@@ -0,0 +1,39 @@
+from cleanup.manager import AccountCleanupManager
+from config import CLEANUP_MODE, CLEANUP_WITHOUT_PRIORITY_FEE
+from utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+def should_cleanup_after_failure() -> bool:
+    return CLEANUP_MODE == "on_fail"
+
+
+def should_cleanup_after_sell() -> bool:
+    return CLEANUP_MODE == "after_sell"
+
+
+def should_cleanup_post_session() -> bool:
+    return CLEANUP_MODE == "post_session"
+
+
+async def handle_cleanup_after_failure(client, wallet, mint):
+    if should_cleanup_after_failure():
+        logger.info("[Cleanup] Triggered by failed buy transaction.")
+        manager = AccountCleanupManager(client, wallet, use_priority_fee=not CLEANUP_WITHOUT_PRIORITY_FEE)
+        await manager.cleanup_ata(mint)
+
+
+async def handle_cleanup_after_sell(client, wallet, mint):
+    if should_cleanup_after_sell():
+        logger.info("[Cleanup] Triggered after token sell.")
+        manager = AccountCleanupManager(client, wallet, use_priority_fee=not CLEANUP_WITHOUT_PRIORITY_FEE)
+        await manager.cleanup_ata(mint)
+
+
+async def handle_cleanup_post_session(client, wallet, mints):
+    if should_cleanup_post_session():
+        logger.info("[Cleanup] Triggered post trading session.")
+        manager = AccountCleanupManager(client, wallet, use_priority_fee=not CLEANUP_WITHOUT_PRIORITY_FEE)
+        for mint in mints:
+            await manager.cleanup_ata(mint)

src/trading/trader.py

@@ -8,7 +8,7 @@ import json
 import os
 from datetime import datetime
 
-import config as config
+import config
 from core.client import SolanaClient
 from core.curve import BondingCurveManager
 from core.priority_fee.manager import PriorityFeeManager
@@ -137,7 +137,7 @@ class PumpTrader:
             )
 
         except Exception as e:
-            logger.error(f"Trading stopped due to error: {str(e)}")
+            logger.error(f"Trading stopped due to error: {e!s}")
             processor_task.cancel()
             await self.solana_client.close()
 
@@ -153,7 +153,7 @@ class PumpTrader:
         self.token_timestamps[token_key] = asyncio.get_event_loop().time()
 
         await self.token_queue.put(token_info)
-        # logger.info(f"Queued new token: {token_info.symbol} ({token_info.mint})")
+        logger.info(f"Queued new token: {token_info.symbol} ({token_info.mint})")
 
     async def _process_token_queue(self, marry_mode: bool, yolo_mode: bool) -> None:
         """Continuously process tokens from the queue, only if they're fresh."""
@@ -254,7 +254,7 @@ class PumpTrader:
             await asyncio.sleep(config.WAIT_TIME_BEFORE_NEW_TOKEN)
 
         except Exception as e:
-            logger.error(f"Error handling token {token_info.symbol}: {str(e)}")
+            logger.error(f"Error handling token {token_info.symbol}: {e!s}")
 
     async def _save_token_info(self, token_info: TokenInfo) -> None:
         """Save token information to a file.