diff --git a/custom_bot_manager.py b/custom_bot_manager.py index d1c2076..ed48843 100644 --- a/custom_bot_manager.py +++ b/custom_bot_manager.py @@ -40,16 +40,17 @@ DEFAULT_COGS = [ class CustomBot(commands.Bot): """Custom bot class with additional functionality for user-specific bots.""" - + def __init__(self, user_id: str, *args, **kwargs): super().__init__(*args, **kwargs) self.user_id = user_id self.owner_id = int(os.getenv('OWNER_USER_ID', '0')) - + self._cleanup_tasks = [] # Track cleanup tasks + async def setup_hook(self): """Called when the bot is first connected to Discord.""" log.info(f"Custom bot for user {self.user_id} is setting up...") - + # Load default cogs for cog in DEFAULT_COGS: try: @@ -59,48 +60,68 @@ class CustomBot(commands.Bot): log.error(f"Failed to load extension {cog} for custom bot {self.user_id}: {e}") traceback.print_exc() + async def close(self): + """Override close to ensure proper cleanup of all resources.""" + log.info(f"Closing custom bot for user {self.user_id}...") + + # Close all cogs that have aiohttp sessions + for cog_name, cog in self.cogs.items(): + try: + if hasattr(cog, 'session') and cog.session and not cog.session.closed: + await cog.session.close() + log.info(f"Closed aiohttp session for cog {cog_name} in custom bot {self.user_id}") + except Exception as e: + log.error(f"Error closing session for cog {cog_name} in custom bot {self.user_id}: {e}") + + # Wait a bit for sessions to close properly + await asyncio.sleep(0.1) + + # Call parent close + await super().close() + log.info(f"Custom bot for user {self.user_id} closed successfully") + async def create_custom_bot( - user_id: str, - token: str, + user_id: str, + token: str, prefix: str = "!", status_type: str = "listening", status_text: str = "!help" ) -> Tuple[bool, str]: """ Create a new custom bot instance for a user. - + Args: user_id: The Discord user ID who owns this bot token: The Discord bot token prefix: Command prefix for the bot status_type: Activity type (playing, listening, watching, competing) status_text: Status text to display - + Returns: Tuple of (success, message) """ # Check if a bot already exists for this user if user_id in custom_bots and custom_bot_status.get(user_id) == STATUS_RUNNING: return False, f"A bot is already running for user {user_id}. Stop it first." - + try: # Set up intents intents = discord.Intents.default() intents.message_content = True intents.members = True - + # Create bot instance bot = CustomBot( user_id=user_id, command_prefix=prefix, intents=intents ) - + # Set up events @bot.event async def on_ready(): log.info(f"Custom bot {bot.user.name} (ID: {bot.user.id}) for user {user_id} is ready!") - + # Set the bot's status activity_type = getattr(discord.ActivityType, status_type, discord.ActivityType.listening) await bot.change_presence( @@ -109,23 +130,23 @@ async def create_custom_bot( name=status_text ) ) - + # Update status custom_bot_status[user_id] = STATUS_RUNNING if user_id in custom_bot_errors: del custom_bot_errors[user_id] - + @bot.event async def on_error(event, *args, **kwargs): log.error(f"Error in custom bot for user {user_id} in event {event}: {sys.exc_info()[1]}") custom_bot_errors[user_id] = str(sys.exc_info()[1]) - + # Store the bot instance custom_bots[user_id] = bot custom_bot_status[user_id] = STATUS_STOPPED - + return True, f"Custom bot created for user {user_id}" - + except Exception as e: log.error(f"Error creating custom bot for user {user_id}: {e}") custom_bot_status[user_id] = STATUS_ERROR @@ -135,101 +156,150 @@ async def create_custom_bot( def run_custom_bot_in_thread(user_id: str, token: str) -> Tuple[bool, str]: """ Run a custom bot in a separate thread. - + Args: user_id: The Discord user ID who owns this bot token: The Discord bot token - + Returns: Tuple of (success, message) """ if user_id not in custom_bots: return False, f"No bot instance found for user {user_id}" - + if user_id in custom_bot_threads and custom_bot_threads[user_id].is_alive(): return False, f"Bot is already running for user {user_id}" - + bot = custom_bots[user_id] - - async def _run_bot(): + + def _run_bot_thread(): + """Run the bot in a new event loop within this thread.""" try: - await bot.start(token) - except discord.errors.LoginFailure: - log.error(f"Invalid token for custom bot (user {user_id})") - custom_bot_status[user_id] = STATUS_ERROR - custom_bot_errors[user_id] = "Invalid Discord bot token. Please check your token and try again." + # Create a new event loop for this thread + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + async def _run_bot(): + try: + await bot.start(token) + except discord.errors.LoginFailure: + log.error(f"Invalid token for custom bot (user {user_id})") + custom_bot_status[user_id] = STATUS_ERROR + custom_bot_errors[user_id] = "Invalid Discord bot token. Please check your token and try again." + except Exception as e: + log.error(f"Error running custom bot for user {user_id}: {e}") + custom_bot_status[user_id] = STATUS_ERROR + custom_bot_errors[user_id] = str(e) + finally: + # Ensure proper cleanup + if not bot.is_closed(): + try: + await bot.close() + except Exception as e: + log.error(f"Error closing bot during cleanup for user {user_id}: {e}") + + # Run the bot + loop.run_until_complete(_run_bot()) + except Exception as e: - log.error(f"Error running custom bot for user {user_id}: {e}") + log.error(f"Error in bot thread for user {user_id}: {e}") custom_bot_status[user_id] = STATUS_ERROR custom_bot_errors[user_id] = str(e) - + finally: + # Clean up the event loop + try: + loop.close() + except Exception as e: + log.error(f"Error closing event loop for user {user_id}: {e}") + # Create and start the thread - loop = asyncio.new_event_loop() thread = threading.Thread( - target=lambda: loop.run_until_complete(_run_bot()), + target=_run_bot_thread, daemon=True, name=f"custom-bot-{user_id}" ) thread.start() - + # Store the thread custom_bot_threads[user_id] = thread - + return True, f"Started custom bot for user {user_id}" def stop_custom_bot(user_id: str) -> Tuple[bool, str]: """ Stop a running custom bot. - + Args: user_id: The Discord user ID who owns this bot - + Returns: Tuple of (success, message) """ if user_id not in custom_bots: return False, f"No bot instance found for user {user_id}" - + if user_id not in custom_bot_threads or not custom_bot_threads[user_id].is_alive(): custom_bot_status[user_id] = STATUS_STOPPED return True, f"Bot was not running for user {user_id}" - + # Get the bot instance bot = custom_bots[user_id] - - # Close the bot (this will be done in a new thread to avoid blocking) - async def _close_bot(): + + def _close_bot_thread(): + """Close the bot in a proper event loop context.""" try: - await bot.close() - custom_bot_status[user_id] = STATUS_STOPPED + # Create a new event loop for this thread + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + async def _close_bot(): + try: + await bot.close() + custom_bot_status[user_id] = STATUS_STOPPED + log.info(f"Successfully closed custom bot for user {user_id}") + except Exception as e: + log.error(f"Error closing custom bot for user {user_id}: {e}") + custom_bot_status[user_id] = STATUS_ERROR + custom_bot_errors[user_id] = str(e) + + # Run the close operation + loop.run_until_complete(_close_bot()) + except Exception as e: - log.error(f"Error closing custom bot for user {user_id}: {e}") + log.error(f"Error in close thread for user {user_id}: {e}") custom_bot_status[user_id] = STATUS_ERROR custom_bot_errors[user_id] = str(e) - + finally: + # Clean up the event loop + try: + loop.close() + except Exception as e: + log.error(f"Error closing event loop in close thread for user {user_id}: {e}") + # Run the close operation in a new thread - loop = asyncio.new_event_loop() close_thread = threading.Thread( - target=lambda: loop.run_until_complete(_close_bot()), + target=_close_bot_thread, daemon=True, name=f"close-bot-{user_id}" ) close_thread.start() - + # Wait for the close thread to finish (with timeout) - close_thread.join(timeout=5.0) - - # The thread will be cleaned up when the bot is started again - + close_thread.join(timeout=10.0) # Increased timeout for proper cleanup + + # Clean up the thread reference + if user_id in custom_bot_threads: + del custom_bot_threads[user_id] + return True, f"Stopped custom bot for user {user_id}" def get_custom_bot_status(user_id: str) -> Dict: """ Get the status of a custom bot. - + Args: user_id: The Discord user ID who owns this bot - + Returns: Dict with status information """ @@ -240,15 +310,15 @@ def get_custom_bot_status(user_id: str) -> Dict: "error": None, "is_running": False } - + status = custom_bot_status.get(user_id, STATUS_STOPPED) error = custom_bot_errors.get(user_id) is_running = ( - user_id in custom_bot_threads and + user_id in custom_bot_threads and custom_bot_threads[user_id].is_alive() and status == STATUS_RUNNING ) - + return { "exists": True, "status": status, @@ -259,7 +329,7 @@ def get_custom_bot_status(user_id: str) -> Dict: def get_all_custom_bot_statuses() -> Dict[str, Dict]: """ Get the status of all custom bots. - + Returns: Dict mapping user_id to status information """ @@ -267,3 +337,47 @@ def get_all_custom_bot_statuses() -> Dict[str, Dict]: for user_id in custom_bots: result[user_id] = get_custom_bot_status(user_id) return result + +def list_custom_bots() -> List[Dict]: + """ + List all custom bot instances and their status. + + Returns: + List of dictionaries containing bot information + """ + bots = [] + for user_id in custom_bots.keys(): + bot_info = get_custom_bot_status(user_id) + bots.append(bot_info) + return bots + +def cleanup_all_custom_bots() -> None: + """ + Clean up all custom bot instances. Should be called when the main bot shuts down. + """ + log.info("Cleaning up all custom bots...") + + # Stop all running bots + for user_id in list(custom_bots.keys()): + try: + if custom_bot_status.get(user_id) == STATUS_RUNNING: + log.info(f"Stopping custom bot for user {user_id}") + stop_custom_bot(user_id) + except Exception as e: + log.error(f"Error stopping custom bot for user {user_id}: {e}") + + # Wait a bit for all bots to stop + import time + time.sleep(2) + + # Force cleanup any remaining threads + for user_id, thread in list(custom_bot_threads.items()): + if thread.is_alive(): + log.warning(f"Force terminating thread for custom bot {user_id}") + # Note: We can't force kill threads in Python, but we can clean up references + try: + del custom_bot_threads[user_id] + except KeyError: + pass + + log.info("Custom bot cleanup completed") diff --git a/discord_bot_sync_api.py b/discord_bot_sync_api.py index dca7f8f..5470846 100644 --- a/discord_bot_sync_api.py +++ b/discord_bot_sync_api.py @@ -19,10 +19,18 @@ from typing import Optional # Added for GurtCog type hint # --- Placeholder for GurtCog instance and bot instance --- # These need to be set by the script that starts the bot and API server -from gurt.cog import GurtCog # Import GurtCog for type hint and access -from cogs.mod_log_cog import ModLogCog # Import ModLogCog for type hint -gurt_cog_instance: Optional[GurtCog] = None -mod_log_cog_instance: Optional[ModLogCog] = None # Placeholder for ModLogCog +# Import GurtCog and ModLogCog conditionally to avoid dependency issues +try: + from gurt.cog import GurtCog # Import GurtCog for type hint and access + from cogs.mod_log_cog import ModLogCog # Import ModLogCog for type hint + gurt_cog_instance: Optional[GurtCog] = None + mod_log_cog_instance: Optional[ModLogCog] = None # Placeholder for ModLogCog +except ImportError as e: + print(f"Warning: Could not import GurtCog or ModLogCog: {e}") + # Use Any type as fallback + from typing import Any + gurt_cog_instance: Optional[Any] = None + mod_log_cog_instance: Optional[Any] = None bot_instance = None # Will be set to the Discord bot instance # ============= Models ============= diff --git a/main.py b/main.py index 1183910..6c8a38d 100644 --- a/main.py +++ b/main.py @@ -653,26 +653,11 @@ async def main(args): # Pass parsed args else: log.info("Flask server process was not running or already terminated.") - # Stop all custom bots + # Stop all custom bots using the improved cleanup function try: - log.info("Stopping all custom bots...") - # Get all running custom bots - bot_statuses = custom_bot_manager.get_all_custom_bot_statuses() - stopped_count = 0 - - for user_id, status in bot_statuses.items(): - if status.get('is_running', False): - log.info(f"Stopping custom bot for user {user_id}") - success, message = custom_bot_manager.stop_custom_bot(user_id) - if success: - stopped_count += 1 - log.info(f"Successfully stopped custom bot for user {user_id}") - else: - log.error(f"Failed to stop custom bot for user {user_id}: {message}") - - log.info(f"Stopped {stopped_count} custom bots") + custom_bot_manager.cleanup_all_custom_bots() except Exception as e: - log.exception(f"Error stopping custom bots: {e}") + log.exception(f"Error during custom bot cleanup: {e}") # Close database/cache pools if they were initialized if bot.pg_pool: