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

custom_bot_manager.py

@@ -45,6 +45,7 @@ class CustomBot(commands.Bot):
         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."""
@@ -59,6 +60,26 @@ 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,
@@ -151,22 +172,49 @@ def run_custom_bot_in_thread(user_id: str, token: str) -> Tuple[bool, str]:
 
     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}"
     )
@@ -197,29 +245,51 @@ def stop_custom_bot(user_id: str) -> Tuple[bool, str]:
     # 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)
+    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}"
 
@@ -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
-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 =============

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: