feat: cleanup manager integrated into trader, enhanced config

src/cleanup/manager.py

@@ -1,8 +1,9 @@
 from solders.pubkey import Pubkey
 from spl.token.instructions import BurnParams, CloseAccountParams, burn, close_account
 
-from config import CLEANUP_WITHOUT_PRIORITY_FEE
+from config import CLEANUP_FORCE_CLOSE_WITH_BURN
 from core.client import SolanaClient
+from core.priority_fee.manager import PriorityFeeManager
 from core.pubkeys import SystemAddresses
 from core.wallet import Wallet
 from utils.logger import get_logger
@@ -12,11 +13,12 @@ logger = get_logger(__name__)
 
 class AccountCleanupManager:
     """Handles safe cleanup of token accounts (ATA) after trading sessions."""
-
     def __init__(
         self,
         client: SolanaClient,
         wallet: Wallet,
+        priority_fee_manager: PriorityFeeManager,
+        use_priority_fee: bool = False,
     ):
         """
         Args:
@@ -25,7 +27,8 @@ class AccountCleanupManager:
         """
         self.client = client
         self.wallet = wallet
-        self.use_priority_fee = not CLEANUP_WITHOUT_PRIORITY_FEE
+        self.priority_fee_manager = priority_fee_manager
+        self.use_priority_fee = use_priority_fee
 
     async def cleanup_ata(self, mint: Pubkey) -> None:
         """
@@ -35,6 +38,12 @@ class AccountCleanupManager:
         ata = self.wallet.get_associated_token_address(mint)
         solana_client = await self.client.get_client()
 
+        priority_fee = (
+            await self.priority_fee_manager.calculate_priority_fee([ata])
+            if self.use_priority_fee
+            else None
+        )
+
         try:
             info = await solana_client.get_account_info(ata)
             if not info.value:
@@ -42,8 +51,8 @@ class AccountCleanupManager:
                 return
 
             balance = await self.client.get_token_account_balance(ata)
-            if balance > 0:
-                logger.info(f"⚠️ Burning {balance} tokens from ATA {ata} (mint: {mint})...")
+            if balance > 0 and CLEANUP_FORCE_CLOSE_WITH_BURN:
+                logger.info(f"Burning {balance} tokens from ATA {ata} (mint: {mint})...")
                 burn_ix = burn(
                     BurnParams(
                         account=ata,
@@ -57,9 +66,15 @@ class AccountCleanupManager:
                     [burn_ix],
                     self.wallet.keypair,
                     skip_preflight=True,
-                    priority_fee=None if not self.use_priority_fee else 0,
+                    priority_fee=priority_fee,
                 )
                 logger.info(f"✅ Burned successfully from ATA {ata}")
+            elif balance > 0:
+                logger.info(
+                    f"Skipping ATA {ata} with non-zero balance ({balance} tokens) "
+                    f"because CLEANUP_FORCE_CLOSE_WITH_BURN is disabled."
+                )
+                return
 
             logger.info(f"Closing ATA: {ata}")
             close_ix = close_account(
@@ -74,10 +89,10 @@ class AccountCleanupManager:
                 [close_ix],
                 self.wallet.keypair,
                 skip_preflight=True,
-                priority_fee=None if not self.use_priority_fee else 0,
+                priority_fee=priority_fee,
             )
             await self.client.confirm_transaction(tx_sig)
-            logger.info(f"✅ Closed successfully: {ata}")
+            logger.info(f"Closed successfully: {ata}")
 
         except Exception as e:
-            logger.warning(f"⚠️ Cleanup failed for ATA {ata}: {e!s}")
+            logger.warning(f"Cleanup failed for ATA {ata}: {e!s}")

src/cleanup/modes.py

@@ -1,5 +1,5 @@
 from cleanup.manager import AccountCleanupManager
-from config import CLEANUP_MODE, CLEANUP_WITHOUT_PRIORITY_FEE
+from config import CLEANUP_MODE, CLEANUP_WITH_PRIORITY_FEE
 from utils.logger import get_logger
 
 logger = get_logger(__name__)
@@ -17,23 +17,21 @@ def should_cleanup_post_session() -> bool:
     return CLEANUP_MODE == "post_session"
 
 
-async def handle_cleanup_after_failure(client, wallet, mint):
+async def handle_cleanup_after_failure(client, wallet, mint, priority_fee_manager):
     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)
+        manager = AccountCleanupManager(client, wallet, priority_fee_manager, CLEANUP_WITH_PRIORITY_FEE)
         await manager.cleanup_ata(mint)
 
-
-async def handle_cleanup_after_sell(client, wallet, mint):
+async def handle_cleanup_after_sell(client, wallet, mint, priority_fee_manager):
     if should_cleanup_after_sell():
         logger.info("[Cleanup] Triggered after token sell.")
-        manager = AccountCleanupManager(client, wallet, use_priority_fee=not CLEANUP_WITHOUT_PRIORITY_FEE)
+        manager = AccountCleanupManager(client, wallet, priority_fee_manager, CLEANUP_WITH_PRIORITY_FEE)
         await manager.cleanup_ata(mint)
 
-
-async def handle_cleanup_post_session(client, wallet, mints):
+async def handle_cleanup_post_session(client, wallet, mints, priority_fee_manager):
     if should_cleanup_post_session():
         logger.info("[Cleanup] Triggered post trading session.")
-        manager = AccountCleanupManager(client, wallet, use_priority_fee=not CLEANUP_WITHOUT_PRIORITY_FEE)
+        manager = AccountCleanupManager(client, wallet, priority_fee_manager, CLEANUP_WITH_PRIORITY_FEE)
         for mint in mints:
             await manager.cleanup_ata(mint)