feat: Enhance custom bot management with improved cleanup and resource handling

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, 
+    user_id: str,
-    token: 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)
+            # Create a new event loop for this thread
-        except discord.errors.LoginFailure:
+            loop = asyncio.new_event_loop()
-            log.error(f"Invalid token for custom bot (user {user_id})")
+            asyncio.set_event_loop(loop)
-            custom_bot_status[user_id] = STATUS_ERROR
+
-            custom_bot_errors[user_id] = "Invalid Discord bot token. Please check your token and try again."
+            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)
+    def _close_bot_thread():
-    async def _close_bot():
+        """Close the bot in a proper event loop context."""
         try:
-            await bot.close()
+            # Create a new event loop for this thread
-            custom_bot_status[user_id] = STATUS_STOPPED
+            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)
+    close_thread.join(timeout=10.0)  # Increased timeout for proper cleanup
-    
+
-    # The thread will be cleaned up when the bot is started again
+    # 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")

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
+# Import GurtCog and ModLogCog conditionally to avoid dependency issues
-from cogs.mod_log_cog import ModLogCog # Import ModLogCog for type hint
+try:
-gurt_cog_instance: Optional[GurtCog] = None
+    from gurt.cog import GurtCog # Import GurtCog for type hint and access
-mod_log_cog_instance: Optional[ModLogCog] = None # Placeholder for ModLogCog
+    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 =============

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...")
+            custom_bot_manager.cleanup_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")
         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: