From 735ab11d4d711d9dd7a2063e75ec772399d37696 Mon Sep 17 00:00:00 2001 From: smypmsa Date: Thu, 27 Mar 2025 15:19:52 +0000 Subject: [PATCH] feat: added cleanup manager and modes --- src/cleanup/manager.py | 83 ++++++++++++++++++++++++++++++++++++++++++ src/cleanup/modes.py | 39 ++++++++++++++++++++ src/trading/trader.py | 8 ++-- 3 files changed, 126 insertions(+), 4 deletions(-) diff --git a/src/cleanup/manager.py b/src/cleanup/manager.py index e69de29..eb00e53 100644 --- a/src/cleanup/manager.py +++ b/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}") diff --git a/src/cleanup/modes.py b/src/cleanup/modes.py index e69de29..26cf079 100644 --- a/src/cleanup/modes.py +++ b/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) diff --git a/src/trading/trader.py b/src/trading/trader.py index 6204217..470b125 100644 --- a/src/trading/trader.py +++ b/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.