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:
+            if balance > 0 and CLEANUP_FORCE_CLOSE_WITH_BURN:
-                logger.info(f"⚠️ Burning {balance} tokens from ATA {ata} (mint: {mint})...")
+                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, priority_fee_manager):
-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)
+        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, priority_fee_manager):
-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)
+        manager = AccountCleanupManager(client, wallet, priority_fee_manager, CLEANUP_WITH_PRIORITY_FEE)
         for mint in mints:
             await manager.cleanup_ata(mint)

src/config.py

@@ -1,87 +1,103 @@
 """
 Configuration for the pump.fun trading bot.
+
+This file defines comprehensive parameters and settings for the trading bot.
+Carefully review and adjust values to match your trading strategy and risk tolerance.
 """
 
 # Trading parameters
-BUY_AMOUNT: int | float = 0.000_001  # Amount of SOL to spend when buying
+# Control trade execution: amount of SOL per trade and acceptable price deviation
-BUY_SLIPPAGE: float = 0.4  # 40% slippage tolerance for buying
+BUY_AMOUNT: int | float = 0.000_001  # Minimal SOL amount to prevent dust transactions
-SELL_SLIPPAGE: float = 0.4  # 40% slippage tolerance for selling
+BUY_SLIPPAGE: float = 0.4  # Maximum acceptable price deviation (0.4 = 40%)
+SELL_SLIPPAGE: float = 0.4  # Consistent slippage tolerance to maintain trading strategy
 
 
-# Configuration for priority fee settings
+# Priority fee configuration
-ENABLE_DYNAMIC_PRIORITY_FEE: bool = False  # Enable dynamic priority fee calculation
+# Manage transaction speed and cost on the Solana network
-ENABLE_FIXED_PRIORITY_FEE: bool = True  # Enable fixed priority fee
+ENABLE_DYNAMIC_PRIORITY_FEE: bool = False  # Adaptive fee calculation
-FIXED_PRIORITY_FEE: int = 2_000  # Fixed priority fee in microlamports
+ENABLE_FIXED_PRIORITY_FEE: bool = True  # Use consistent, predictable fee
-EXTRA_PRIORITY_FEE: float = (
+FIXED_PRIORITY_FEE: int = 2_000  # Base fee in microlamports
-    0.0  # Percentage increase applied to priority fee (0.1 = 10%)
+EXTRA_PRIORITY_FEE: float = 0.0  # Percentage increase on base priority fee (0.1 = 10%)
-)
+HARD_CAP_PRIOR_FEE: int = 200_000  # Maximum allowable fee to prevent excessive spending in microlamports
-HARD_CAP_PRIOR_FEE: int = (
-    200_000  # Maximum allowed priority fee in microlamports (hard cap)
-)
 
 
 # Listener configuration
-LISTENER_TYPE = "block"  # Options: "block" or "logs"
+# Choose method for detecting new tokens on the network
+# "logs": Recommended for more stable token detection
+# "blocks": Unstable method, potentially less reliable
+LISTENER_TYPE = "logs"
 
 
-# Retries and timeouts
+# Retry and timeout settings
-MAX_RETRIES: int = 10  # Number of retries for transaction sending
+# Control bot resilience and transaction handling
-# TODO: waiting times will be replaced with retries to shorten delays
+MAX_RETRIES: int = 10  # Number of attempts for transaction submission
-WAIT_TIME_AFTER_CREATION: int | float = (
+
-    15  # Time to wait after token creation (in seconds)
+# Waiting periods in seconds between actions (TODO: to be replaced with retry mechanism)
-    # Too short a delay may cause the RPC node to be unaware of the bonding curve account
+WAIT_TIME_AFTER_CREATION: int | float = 15  # Seconds to wait after token creation
-)
+WAIT_TIME_AFTER_BUY: int | float = 15  # Holding period after buy transaction
-WAIT_TIME_AFTER_BUY: int | float = (
+WAIT_TIME_BEFORE_NEW_TOKEN: int | float = 15  # Pause between token trades
-    15  # Time to wait after a buy transaction is confirmed (in seconds)
-    # Acts as a simple holding period
-    # Too short delay may cause the RPC node to be unaware of account balance
-)
-WAIT_TIME_BEFORE_NEW_TOKEN: int | float = (
-    5  # Time to wait after a sell transaction is confirmed (in seconds)
-    # Provides a pause between completed trades, can be set to 0
-)
 
 
-# Maximum age (in seconds) for a token to be considered "fresh" and eligible for processing.
+# Token and account management
-# This threshold is checked before processing starts - tokens older than this are skipped
+# Control token processing and account cleanup strategies
-# since they likely contain outdated information from the websocket stream
+MAX_TOKEN_AGE: int | float = 0.1  # Maximum token age in seconds for processing
-MAX_TOKEN_AGE: int | float = 0.1
+
+# Cleanup mode determines when to manage token accounts. Options:
+# "disabled": No cleanup will occur.
+# "on_fail": Only clean up if a buy transaction fails.
+# "after_sell": Clean up after selling, but only if the balance is zero.
+# "post_session": Clean up all empty accounts after a trading session ends.
+CLEANUP_MODE: str = "after_sell"
+CLEANUP_FORCE_CLOSE_WITH_BURN: bool = True  # Burn remaining tokens before closing account, else skip ATA with non-zero balances
+CLEANUP_WITH_PRIORITY_FEE: bool = False  # Use priority fees for cleanup transactions
 
 
-# Cleanup configuration
+# Node provider configuration (TODO: to be implemented)
-
+# Manage RPC node interaction to prevent rate limiting
-# CLEANUP_MODE controls when to attempt closing token accounts (ATA)
+MAX_RPS: int = 25  # Maximum requests per second
-# Options:
-#   - "disabled": no cleanup at all
-#   - "on_fail": close ATA only if a buy transaction fails
-#   - "after_sell": close ATA after selling, but only if token balance is zero
-#   - "post_session": clean up all empty ATA accounts at the end of trading session
-CLEANUP_MODE: str = "disabled"
-
-# If True, cleanup transactions will skip priority fees (cheaper, slower confirmation)
-# Set to False if you want faster confirmations (e.g. when racing for SOL reclaim)
-CLEANUP_WITHOUT_PRIORITY_FEE: bool = True
 
 
-# Node provider configuration
+def validate_configuration() -> None:
-# Tested with Chainstack nodes (https://console.chainstack.com), but you can use any node provider
+    """
-# You can get a trader node https://docs.chainstack.com/docs/solana-trader-nodes
+    Comprehensive validation of bot configuration.
-MAX_RPS: int = 25  # TODO: not implemented. Max RPS to avoid rate limit errors
+    
+    Checks:
+    - Type correctness
+    - Value ranges
+    - Logical consistency of settings
+    """
+    # Configuration validation checks
+    config_checks = [
+        # (value, type, min_value, max_value, error_message)
+        (BUY_AMOUNT, (int, float), 0, float('inf'), "BUY_AMOUNT must be a positive number"),
+        (BUY_SLIPPAGE, float, 0, 1, "BUY_SLIPPAGE must be between 0 and 1"),
+        (SELL_SLIPPAGE, float, 0, 1, "SELL_SLIPPAGE must be between 0 and 1"),
+        (FIXED_PRIORITY_FEE, int, 0, float('inf'), "FIXED_PRIORITY_FEE must be a non-negative integer"),
+        (EXTRA_PRIORITY_FEE, float, 0, 1, "EXTRA_PRIORITY_FEE must be between 0 and 1"),
+        (HARD_CAP_PRIOR_FEE, int, 0, float('inf'), "HARD_CAP_PRIOR_FEE must be a non-negative integer"),
+        (MAX_RETRIES, int, 0, 100, "MAX_RETRIES must be between 0 and 100")
+    ]
+
+    for value, expected_type, min_val, max_val, error_msg in config_checks:
+        if not isinstance(value, expected_type):
+            raise ValueError(f"Type error: {error_msg}")
+        
+        if isinstance(value, (int, float)) and not (min_val <= value <= max_val):
+            raise ValueError(f"Range error: {error_msg}")
+
+    # Logical consistency checks
+    if ENABLE_DYNAMIC_PRIORITY_FEE and ENABLE_FIXED_PRIORITY_FEE:
+        raise ValueError("Cannot enable both dynamic and fixed priority fees simultaneously")
+
+    # Validate listener type
+    if LISTENER_TYPE not in ["logs", "blocks"]:
+        raise ValueError("LISTENER_TYPE must be either 'logs' or 'blocks'")
+
+    # Validate cleanup mode
+    valid_cleanup_modes = ["disabled", "on_fail", "after_sell", "post_session"]
+    if CLEANUP_MODE not in valid_cleanup_modes:
+        raise ValueError(f"CLEANUP_MODE must be one of {valid_cleanup_modes}")
 
 
-def validate_priority_fee_config() -> None:
+# Validate configuration on import
-    """Validate priority fee configuration values."""
+validate_configuration()
-    if not isinstance(ENABLE_DYNAMIC_PRIORITY_FEE, bool):
-        raise ValueError("ENABLE_DYNAMIC_PRIORITY_FEE must be a boolean")
-    if not isinstance(ENABLE_FIXED_PRIORITY_FEE, bool):
-        raise ValueError("ENABLE_FIXED_PRIORITY_FEE must be a boolean")
-    if not isinstance(FIXED_PRIORITY_FEE, int) or FIXED_PRIORITY_FEE < 0:
-        raise ValueError("FIXED_PRIORITY_FEE must be a non-negative integer")
-    if not isinstance(EXTRA_PRIORITY_FEE, float) or EXTRA_PRIORITY_FEE < 0:
-        raise ValueError("EXTRA_PRIORITY_FEE must be a non-negative float")
-    if not isinstance(HARD_CAP_PRIOR_FEE, int) or HARD_CAP_PRIOR_FEE < 0:
-        raise ValueError("HARD_CAP_PRIOR_FEE must be a non-negative integer")
-
-
-# Validate config on import
-validate_priority_fee_config()

src/trading/trader.py

@@ -8,7 +8,14 @@ import json
 import os
 from datetime import datetime
 
+from solders.pubkey import Pubkey
+
 import config
+from cleanup.modes import (
+    handle_cleanup_after_failure,
+    handle_cleanup_after_sell,
+    handle_cleanup_post_session,
+)
 from core.client import SolanaClient
 from core.curve import BondingCurveManager
 from core.priority_fee.manager import PriorityFeeManager
@@ -96,6 +103,8 @@ class PumpTrader:
         self.max_retries = max_retries
         self.max_token_age = config.MAX_TOKEN_AGE
 
+        self.traded_mints: set[Pubkey] = set()
+
         # Token processing state
         self.token_queue = asyncio.Queue()
         self.processing = False
@@ -141,6 +150,13 @@ class PumpTrader:
             processor_task.cancel()
             await self.solana_client.close()
 
+        finally:
+            processor_task.cancel()
+            if self.traded_mints:
+                # Close ATA if enabled
+                await handle_cleanup_post_session(self.solana_client, self.wallet, list(self.traded_mints), self.priority_fee_manager)
+            await self.solana_client.close()
+
     async def _queue_token(self, token_info: TokenInfo) -> None:
         """Queue a token for processing if not already processed."""
         token_key = str(token_info.mint)
@@ -215,10 +231,13 @@ class PumpTrader:
                     buy_result.amount,  # type: ignore
                     buy_result.tx_signature,
                 )
+                self.traded_mints.add(token_info.mint)
             else:
                 logger.error(
                     f"Failed to buy {token_info.symbol}: {buy_result.error_message}"
                 )
+                # Close ATA if enabled
+                await handle_cleanup_after_failure(self.solana_client, self.wallet, token_info.mint, self.priority_fee_manager)
 
             # Sell token if not in marry mode
             if not marry_mode and buy_result.success:
@@ -239,6 +258,8 @@ class PumpTrader:
                         sell_result.amount,  # type: ignore
                         sell_result.tx_signature,
                     )
+                    # Close ATA if enabled
+                    await handle_cleanup_after_sell(self.solana_client, self.wallet, token_info.mint, self.priority_fee_manager)
                 else:
                     logger.error(
                         f"Failed to sell {token_info.symbol}: {sell_result.error_message}"