Merge pull request #71 from chainstacklabs/feature/cleanup-ata-support

Integrate ATA closing into trader

src/cleanup/manager.py

@@ -0,0 +1,98 @@
+from solders.pubkey import Pubkey
+from spl.token.instructions import BurnParams, CloseAccountParams, burn, close_account
+
+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
+
+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:
+            client: Solana RPC client
+            wallet: Wallet for signing transactions
+        """
+        self.client = client
+        self.wallet = wallet
+        self.priority_fee_manager = priority_fee_manager
+        self.use_priority_fee = use_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()
+
+        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:
+                logger.info(f"ATA {ata} does not exist or already closed.")
+                return
+
+            balance = await self.client.get_token_account_balance(ata)
+            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,
+                        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=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(
+                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=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/cleanup/modes.py

@@ -0,0 +1,37 @@
+from cleanup.manager import AccountCleanupManager
+from config import CLEANUP_MODE, CLEANUP_WITH_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, priority_fee_manager):
+    if should_cleanup_after_failure():
+        logger.info("[Cleanup] Triggered by failed buy transaction.")
+        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):
+    if should_cleanup_after_sell():
+        logger.info("[Cleanup] Triggered after token sell.")
+        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):
+    if should_cleanup_post_session():
+        logger.info("[Cleanup] Triggered post trading session.")
+        manager = AccountCleanupManager(client, wallet, priority_fee_manager, CLEANUP_WITH_PRIORITY_FEE)
+        for mint in mints:
+            await manager.cleanup_ata(mint)

src/cli.py

@@ -8,7 +8,7 @@ import asyncio
 import os
 import sys
 
-import config as config
+import config
 from trading.trader import PumpTrader
 from utils.logger import get_logger, setup_file_logging
 
@@ -109,7 +109,7 @@ async def main() -> None:
     except KeyboardInterrupt:
         logger.info("Trading stopped by user")
     except Exception as e:
-        logger.error(f"Trading stopped due to error: {str(e)}")
+        logger.error(f"Trading stopped due to error: {e!s}")
     finally:
         try:
             await trader.solana_client.close()
@@ -117,7 +117,7 @@ async def main() -> None:
             pass
 
 
-def sync_main():
+def sync_main() -> None:
     asyncio.run(main())
 
 

src/config.py

@@ -1,72 +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
+# Control trade execution: amount of SOL per trade and acceptable price deviation
 BUY_AMOUNT: int | float = 0.000_001  # Amount of SOL to spend when buying
-BUY_SLIPPAGE: float = 0.4  # 40% slippage tolerance for buying
-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
-ENABLE_DYNAMIC_PRIORITY_FEE: bool = False  # Enable dynamic priority fee calculation
-ENABLE_FIXED_PRIORITY_FEE: bool = True  # Enable fixed priority fee
-FIXED_PRIORITY_FEE: int = 2_000  # Fixed priority fee in microlamports
-EXTRA_PRIORITY_FEE: float = (
-    0.0  # Percentage increase applied to priority fee (0.1 = 10%)
-)
-HARD_CAP_PRIOR_FEE: int = (
-    200_000  # Maximum allowed priority fee in microlamports (hard cap)
-)
+# Priority fee configuration
+# Manage transaction speed and cost on the Solana network
+ENABLE_DYNAMIC_PRIORITY_FEE: bool = False  # Adaptive fee calculation
+ENABLE_FIXED_PRIORITY_FEE: bool = True  # Use consistent, predictable fee
+FIXED_PRIORITY_FEE: int = 2_000  # Base fee in microlamports
+EXTRA_PRIORITY_FEE: float = 0.0  # Percentage increase on priority fee (0.1 = 10%)
+HARD_CAP_PRIOR_FEE: int = 200_000  # Maximum allowable fee to prevent excessive spending in microlamports
 
 
 # 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
-MAX_RETRIES: int = 10  # Number of retries for transaction sending
-# TODO: waiting times will be replaced with retries to shorten delays
-WAIT_TIME_AFTER_CREATION: int | float = (
-    15  # Time to wait after token creation (in seconds)
-    # Too short a delay may cause the RPC node to be unaware of the bonding curve account
-)
-WAIT_TIME_AFTER_BUY: int | float = (
-    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
-)
+# Retry and timeout settings
+# Control bot resilience and transaction handling
+MAX_RETRIES: int = 10  # Number of attempts for transaction submission
+
+# Waiting periods in seconds between actions (TODO: to be replaced with retry mechanism)
+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_BEFORE_NEW_TOKEN: int | float = 15  # Pause between token trades
 
 
-# Maximum age (in seconds) for a token to be considered "fresh" and eligible for processing.
-# This threshold is checked before processing starts - tokens older than this are skipped
-# since they likely contain outdated information from the websocket stream
-MAX_TOKEN_AGE: int | float = 0.1
+# Token and account management
+# Control token processing and account cleanup strategies
+MAX_TOKEN_AGE: int | float = 0.1  # Maximum token age in seconds for processing
+
+# 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 = "disabled"
+CLEANUP_FORCE_CLOSE_WITH_BURN: bool = False  # 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
 
 
-# Node provider configuration
-# 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
-MAX_RPS: int = 25  # TODO: not implemented. Max RPS to avoid rate limit errors
+# Node provider configuration (TODO: to be implemented)
+# Manage RPC node interaction to prevent rate limiting
+MAX_RPS: int = 25  # Maximum requests per second
 
 
-def validate_priority_fee_config() -> None:
-    """Validate priority fee configuration values."""
-    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")
+def validate_configuration() -> None:
+    """
+    Comprehensive validation of bot configuration.
+    
+    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}")
 
 
-# Validate config on import
-validate_priority_fee_config()
+# Validate configuration on import
+validate_configuration()

src/trading/trader.py

@@ -8,7 +8,14 @@ import json
 import os
 from datetime import datetime
 
-import config as config
+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
@@ -137,10 +146,17 @@ 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()
 
+        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)
@@ -153,7 +169,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."""
@@ -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}"
@@ -254,7 +275,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.