Add restart and shutdown commands

Author: greeeen-dev

shinobu/__main__.py

@@ -21,10 +21,12 @@
 import time
 import copy
 import argparse
+import traceback
 import ujson as json
 import orjson
 import getpass
 import uuid
+import asyncio
 import discord
 from discord.ext import commands
 from dotenv import load_dotenv
@@ -54,6 +56,10 @@
 launch_secrets_cli: bool = launch_args.secrets
 launch_installer_cli: bool = launch_args.installer
 
+# Store restart options
+restart_message_id: int | None = None
+restart_message_channel_id: int | None = None
+
 # Create TokenStore variable (do not initialize yet)
 tokenstore: manager.TokenStore | None = None
 
@@ -557,7 +563,11 @@ def load_extension(self, package: str, *args, **kwargs):
 def start_bot() -> bool:
     """Starts Shinobu!"""
 
-    global tokenstore, secrets_authority, raw_encryptor, extension_map, password
+    global tokenstore, secrets_authority, raw_encryptor, extension_map, password, restart_message_id,\
+           restart_message_channel_id
+
+    # Ensure asyncio event loop is always fresh
+    asyncio.set_event_loop(asyncio.new_event_loop())
 
     # Regenerate bootscript-level objects
     tokenstore = manager.TokenStore(
@@ -582,6 +592,8 @@ def start_bot() -> bool:
 
     # Create Shinobu bot instance
     bot: runtime.ShinobuBot = runtime.ShinobuBot(command_prefix="sh!", intents=intents, manifest=manifest_path)
+    bot.restart_message_id = restart_message_id
+    bot.restart_message_channel_id = restart_message_channel_id
     bot.setup_entitlements_loader(ModuleLoader(bot))
     bot.load_builtins()
 
@@ -593,11 +605,26 @@ def start_bot() -> bool:
         bot.run(tokenstore.retrieve("TOKEN"))
     except KeyboardInterrupt:
         pass
+    except RuntimeError:
+        # Assume session was closed (i.e. bot restarted)
+        pass
     except:
+        traceback.print_exc()
         print("Bot has crashed.")
         should_restart = bot.should_restart
 
+    # Run cleanup
     bot.cleanup()
+
+    # Set restart state and channel
+    if bot.requested_restart:
+        should_restart = True
+        restart_message_id = bot.restart_message_id
+        restart_message_channel_id = bot.restart_message_channel_id
+
+    # Clear bot memory
+    del bot
+
     return should_restart
 
 def start_secrets_cli():

shinobu/runtime/modules/admin.py

@@ -22,6 +22,7 @@
 import time
 import traceback
 import base64
+import discord
 from contextlib import redirect_stdout
 from discord.ext import commands
 from shinobu.runtime.models import shinobu_cog
@@ -121,5 +122,22 @@ async def eval(self, ctx: commands.Context, *, body: str):
             else:
                 await ctx.send(f':white_check_mark: Evaluation completed in `{exec_time}` seconds.\n```\n{value}\n```')
 
+    @commands.command(name="shutdown", description="Shuts the bot down.", aliases=["poweroff"])
+    @commands.is_owner()
+    async def shutdown(self, ctx: commands.Context):
+        await ctx.reply("shutting down :zzz:")
+
+        # Shut down bot
+        await self.bot.close()
+
+    @commands.command(name="restart", description="Restarts the bot.", aliases=["reboot"])
+    @commands.is_owner()
+    async def restart(self, ctx: commands.Context):
+        message: discord.Message = await ctx.reply("restarting :arrows_counterclockwise:")
+
+        # Request restart then shut down bot
+        self.bot.request_restart(message=message)
+        await self.bot.close()
+
 def setup(bot):
     bot.add_cog(Admin(bot))

shinobu/runtime/modules/events.py

@@ -20,7 +20,6 @@
 import traceback
 import discord
 from discord.ext import commands, bridge
-from discord.ext.bridge import BridgeExtContext
 from shinobu.runtime.models.colors import Colors
 from shinobu.runtime.utils import check_slash
 
@@ -39,9 +38,18 @@ def __init__(self, bot, **kwargs):
 
     @commands.Cog.listener()
     async def on_ready(self):
-        print("Bot is ready, woohoo!")
+        print("Shinobu Runtime is ready! :3")
         print(f"Logged in as {self.bot.user.name}#{self.bot.user.discriminator} ({self.bot.user.id})")
 
+        # Handle restart
+        if self.bot.restart_message_id:
+            channel: discord.TextChannel | None = self.bot.get_channel(self.bot.restart_message_channel_id)
+            if not channel:
+                return
+
+            message: discord.PartialMessage = discord.PartialMessage(channel=channel, id=self.bot.restart_message_id)
+            await message.edit(content="bot restarted! :white_check_mark:")
+
     @commands.Cog.listener()
     async def on_bridge_command_error(self, ctx: bridge.BridgeApplicationContext | bridge.BridgeExtContext, error):
         is_slash: bool = check_slash.is_slash(ctx)

shinobu/runtime/runtime.py

@@ -18,6 +18,7 @@
 
 import tomllib
 import ujson as json
+import discord
 from discord.ext import bridge
 
 class ShinobuErrorManager:
@@ -50,7 +51,12 @@ def __init__(self, *args, **kwargs):
         self.__errors: ShinobuErrorManager = ShinobuErrorManager()
         self.__cog_entitlements_loader = None
         self._cleanups = {}
-        self._should_restart: bool = False
+
+        # Restart state
+        self._should_restart: bool = False # Restart on crash
+        self._requested_restart: bool = False # Instance owner requested restart
+        self.restart_message_id: int | None = None
+        self.restart_message_channel_id: int | None = None
 
         # Load configs
         self._config: dict = {}
@@ -121,6 +127,10 @@ def config(self) -> dict:
     def should_restart(self) -> bool:
         return self._should_restart
 
+    @property
+    def requested_restart(self) -> bool:
+        return self._requested_restart
+
     @property
     def shared_objects(self) -> ShinobuSharedObjects:
         return self.__shared_objects
@@ -132,3 +142,11 @@ def errors(self) -> ShinobuErrorManager:
     @property
     def cog_entitlements_loader(self):
         return self.__cog_entitlements_loader
+
+    def request_restart(self, message: discord.Message | None = None):
+        """Sets bot to restart. This does not restart the bot on its own."""
+        self._requested_restart = True
+
+        if message:
+            self.restart_message_id = message.id
+            self.restart_message_channel_id = message.channel.id