Add bridge pending status and platform restart

Author: greeeen-dev

shinobu/__main__.py

@@ -34,13 +34,13 @@
 from shinobu.cli import secrets as secrets_cli, installer as installer_cli
 from shinobu.runtime.secrets.encryptor import EncryptedData
 
-# Manifest path
-manifest_path = os.path.join(os.path.dirname(__file__), "manifest.json")
-
 # Prevent attacks via import
 if __name__ != "__main__":
     raise RuntimeError("Bootscript should not be imported!")
 
+# Manifest path
+manifest_path = os.path.join(os.path.dirname(__file__), "manifest.json")
+
 # Create argument parser
 parser = argparse.ArgumentParser(
     prog="shinobu",

shinobu/beacon/cogs/config.py

@@ -0,0 +1,49 @@
+"""
+Shinobu - Converse from anywhere, anytime.
+Copyright (C) 2026-present  Green (@greeeen-dev)
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU Affero General Public License as
+published by the Free Software Foundation, either version 3 of the
+License, or (at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+
+import discord
+from discord.ext import commands, bridge
+from shinobu.runtime.models import shinobu_cog
+from shinobu.beacon.protocol import beacon
+
+class BeaconConfig(shinobu_cog.ShinobuCog):
+    def __init__(self, bot):
+        # Register cog metadata
+        super().__init__(
+            bot,
+            shinobu_metadata=shinobu_cog.ShinobuCogMetadata(
+                name="Config",
+                description="A module containing configuration commands.",
+                emoji="\U0001F4E1",
+                visible_in_help=True
+            )
+        )
+
+        # Get Beacon
+        self._beacon: beacon.Beacon = self.bot.shared_objects.get("beacon")
+
+    @bridge.bridge_group(name="config")
+    async def config_universal(self, ctx):
+        # Universal command group.
+        pass
+
+def get_cog_type():
+    return BeaconConfig
+
+def setup(bot):
+    bot.add_cog(BeaconConfig(bot))

shinobu/beacon/discord/driver.py

@@ -751,6 +751,7 @@ async def _edit(self, message: beacon_message.BeaconMessage, content: beacon_mes
 
     async def _delete(self, message: beacon_message.BeaconMessage):
         channel = self.bot.get_channel(int(message.channel.id))
+        partial_message: discord.PartialMessage = discord.PartialMessage(channel=channel, id=int(message.id))
 
         if message.webhook_id:
             # Ensure webhook is in cache
@@ -760,12 +761,14 @@ async def _delete(self, message: beacon_message.BeaconMessage):
             webhook_obj: discord.Webhook = self._webhooks.get_webhook(message.webhook_id)
 
             # Delete message
-            await webhook_obj.delete_message(message_id=int(message.id))
+            try:
+                await webhook_obj.delete_message(message_id=int(message.id))
+            except RuntimeError:
+                # Use regular delete
+                await partial_message.delete()
         else:
-            message_obj = await channel.fetch_message(int(message.id))
-
             # Delete message
-            await message_obj.delete()
+            await partial_message.delete()
 
     async def _purge(self, messages: list[beacon_message.BeaconMessage]):
         channel: discord.TextChannel = self.bot.get_channel(int(messages[0].channel.id))

shinobu/beacon/discord/parent.py

@@ -190,41 +190,27 @@ async def _get_attachment_data(attachment: discord.Attachment) -> beacon_file.Be
             spoiler=spoiler
         )
 
-    @commands.Cog.listener()
-    async def on_message(self, message: discord.Message):
-        origin_driver: beacon_driver.BeaconDriver = self._beacon.drivers.get_driver("discord")
-
+    async def handle_edit(self, message: discord.Message):
         # noinspection DuplicatedCode
-        if message.content.startswith(self.bot.command_prefix):
-            # Assume this is a text command
-            return
+        origin_driver: beacon_driver.BeaconDriver = self._beacon.drivers.get_driver("discord")
 
-        if message.author.id == self.bot.user.id:
-            # Do not self-bridge
+        # Get the BeaconMessage object for the message
+        message_obj: beacon_message.BeaconMessage = self._beacon.messages.get_message(str(message.id))
+        if not message_obj:
+            # We can't edit messages that aren't cached
             return
 
+        # Did we bridge this message?
         if message.webhook_id:
-            # Check if the webhook was ours (to prevent a self-bridge)
-            webhook: discord.Webhook | None = origin_driver.webhooks.get_webhook(str(message.webhook_id))
-
-            if not webhook:
-                try:
-                    webhook = await self.bot.fetch_webhook(message.webhook_id)
-                    origin_driver.webhooks.store_webhook(str(webhook.id), webhook)
-                except discord.HTTPException:
-                    pass
-
-            if webhook:
-                # Does the bot own the webhook?
-                if webhook.user.id == self.bot.user.id:
-                    # Do not self-bridge
-                    return
+            if message_obj.author.id != str(message.webhook_id):
+                # We probably did
+                return
 
         # Convert guild data to server.BeaconServer
-        # noinspection DuplicatedCode
         server: beacon_server.BeaconServer = origin_driver.get_server(str(message.guild.id))
 
         # Convert author data to member.BeaconMember
+        # noinspection DuplicatedCode
         author: beacon_member.BeaconMember = origin_driver.get_member(server, str(message.author.id))
 
         # Convert channel data to channel.BeaconChannel
@@ -236,15 +222,14 @@ async def on_message(self, message: discord.Message):
         # Get Space
         space: beacon_space.BeaconSpace = self._beacon.spaces.get_space_for_channel(channel)
 
-        # noinspection DuplicatedCode
-        if not space:
-            # We can't bridge
-            return
-
         # Get the ID of the webhook to use
         membership: beacon_space.BeaconSpaceMember = space.get_member(server)
         webhook_id = membership.webhook_id
 
+        if not space:
+            # We can't bridge
+            return
+
         # Convert message data to message.BeaconMessageContent
         content: beacon_message.BeaconMessageContent = await self._to_beacon_content(message)
 
@@ -261,26 +246,23 @@ async def on_message(self, message: discord.Message):
         if preliminary_block:
             return
 
-        # Send message!
+        # Edit the message!
         try:
-            await self._beacon.send(
-                author=author,
-                space=space,
-                content=content,
-                webhook_id=webhook_id
+            await self._beacon.edit(
+                message=message_obj,
+                content=content
             )
         except beacon.BeaconPlatformDisabled:
             pass
 
-    @commands.Cog.listener()
-    async def on_message_edit(self, _, message: discord.Message):
+    async def handle_delete(self, message: discord.Message):
         # noinspection DuplicatedCode
         origin_driver: beacon_driver.BeaconDriver = self._beacon.drivers.get_driver("discord")
 
         # Get the BeaconMessage object for the message
         message_obj: beacon_message.BeaconMessage = self._beacon.messages.get_message(str(message.id))
         if not message_obj:
-            # We can't edit messages that aren't cached
+            # We can't remove messages that aren't cached
             return
 
         # Did we bridge this message?
@@ -292,11 +274,8 @@ async def on_message_edit(self, _, message: discord.Message):
         # Convert guild data to server.BeaconServer
         server: beacon_server.BeaconServer = origin_driver.get_server(str(message.guild.id))
 
-        # Convert author data to member.BeaconMember
-        # noinspection DuplicatedCode
-        author: beacon_member.BeaconMember = origin_driver.get_member(server, str(message.author.id))
-
         # Convert channel data to channel.BeaconChannel
+        # noinspection DuplicatedCode
         channel: beacon_channel.BeaconChannel = origin_driver.get_channel(server, str(message.channel.id))
         if not channel:
             # We can't bridge
@@ -305,61 +284,54 @@ async def on_message_edit(self, _, message: discord.Message):
         # Get Space
         space: beacon_space.BeaconSpace = self._beacon.spaces.get_space_for_channel(channel)
 
-        # Get the ID of the webhook to use
-        membership: beacon_space.BeaconSpaceMember = space.get_member(server)
-        webhook_id = membership.webhook_id
-
         if not space:
-            # We can't bridge
-            return
-
-        # Convert message data to message.BeaconMessageContent
-        content: beacon_message.BeaconMessageContent = await self._to_beacon_content(message)
-
-        # Run preliminary checks
-        preliminary_block: beacon.BeaconMessageBlockedReason | None = await self._beacon.can_send(
-            author=author,
-            space=space,
-            content=content,
-            webhook_id=webhook_id,
-            skip_filter=True
-        )
-
-        # TODO: Add returning the block reason.
-        if preliminary_block:
+            # We can't bridge deletes, even if it was sent in the Space by the server
             return
 
-        # Edit the message!
+        # Delete the message!
         try:
-            await self._beacon.edit(
-                message=message_obj,
-                content=content
-            )
+            await self._beacon.delete(message=message_obj)
         except beacon.BeaconPlatformDisabled:
             pass
 
     @commands.Cog.listener()
-    async def on_message_delete(self, message: discord.Message):
-        # noinspection DuplicatedCode
+    async def on_message(self, message: discord.Message):
         origin_driver: beacon_driver.BeaconDriver = self._beacon.drivers.get_driver("discord")
 
-        # Get the BeaconMessage object for the message
-        message_obj: beacon_message.BeaconMessage = self._beacon.messages.get_message(str(message.id))
-        if not message_obj:
-            # We can't remove messages that aren't cached
+        # noinspection DuplicatedCode
+        if message.content.startswith(self.bot.command_prefix):
+            # Assume this is a text command
+            return
+
+        if message.author.id == self.bot.user.id:
+            # Do not self-bridge
             return
 
-        # Did we bridge this message?
         if message.webhook_id:
-            if message_obj.author.id != str(message.webhook_id):
-                # We probably did
-                return
+            # Check if the webhook was ours (to prevent a self-bridge)
+            webhook: discord.Webhook | None = origin_driver.webhooks.get_webhook(str(message.webhook_id))
+
+            if not webhook:
+                try:
+                    webhook = await self.bot.fetch_webhook(message.webhook_id)
+                    origin_driver.webhooks.store_webhook(str(webhook.id), webhook)
+                except discord.HTTPException:
+                    pass
+
+            if webhook:
+                # Does the bot own the webhook?
+                if webhook.user.id == self.bot.user.id:
+                    # Do not self-bridge
+                    return
 
         # Convert guild data to server.BeaconServer
+        # noinspection DuplicatedCode
         server: beacon_server.BeaconServer = origin_driver.get_server(str(message.guild.id))
 
+        # Convert author data to member.BeaconMember
+        author: beacon_member.BeaconMember = origin_driver.get_member(server, str(message.author.id))
+
         # Convert channel data to channel.BeaconChannel
-        # noinspection DuplicatedCode
         channel: beacon_channel.BeaconChannel = origin_driver.get_channel(server, str(message.channel.id))
         if not channel:
             # We can't bridge
@@ -368,16 +340,62 @@ async def on_message_delete(self, message: discord.Message):
         # Get Space
         space: beacon_space.BeaconSpace = self._beacon.spaces.get_space_for_channel(channel)
 
+        # noinspection DuplicatedCode
         if not space:
-            # We can't bridge deletes, even if it was sent in the Space by the server
+            # We can't bridge
             return
 
-        # Delete the message!
+        # Get the ID of the webhook to use
+        membership: beacon_space.BeaconSpaceMember = space.get_member(server)
+        webhook_id = membership.webhook_id
+
+        # Convert message data to message.BeaconMessageContent
+        content: beacon_message.BeaconMessageContent = await self._to_beacon_content(message)
+
+        # Run preliminary checks
+        preliminary_block: beacon.BeaconMessageBlockedReason | None = await self._beacon.can_send(
+            author=author,
+            space=space,
+            content=content,
+            webhook_id=webhook_id,
+            skip_filter=True
+        )
+
+        # TODO: Add returning the block reason.
+        if preliminary_block:
+            return
+
+        # Send message!
         try:
-            await self._beacon.delete(message=message_obj)
+            await self._beacon.send(
+                author=author,
+                space=space,
+                content=content,
+                webhook_id=webhook_id
+            )
         except beacon.BeaconPlatformDisabled:
             pass
 
+    @commands.Cog.listener()
+    async def on_message_edit(self, _, message: discord.Message):
+        # Check if message is pending
+        if self._beacon.is_pending(str(message.id)):
+            # Add callback
+            self._beacon.add_callback(str(message.id), self.handle_edit, [message])
+        else:
+            # Run directly
+            await self.handle_edit(message)
+
+    @commands.Cog.listener()
+    async def on_message_delete(self, message: discord.Message):
+        # Check if message is pending
+        if self._beacon.is_pending(str(message.id)):
+            # Add callback
+            self._beacon.add_callback(str(message.id), self.handle_delete, [message])
+        else:
+            # Run directly
+            await self.handle_delete(message)
+
     @commands.Cog.listener()
     async def on_bulk_message_delete(self, messages: list[discord.Message]):
         # noinspection DuplicatedCode