# Copyright 2026 Logan Fick # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """The simple nonversation Discord bot. Provides the Crabstero subclass that owns the full bot lifecycle: database connection, cog loading, ingestion worker pool, and graceful shutdown. """ import asyncio import logging import discord from discord import app_commands from discord.ext import commands from crabstero import __version__ as crabstero_version from crabstero.database import Database from crabstero.listeners import interaction, message, server_events from crabstero.tasks.ingestion import ingest_channel logger = logging.getLogger(__name__) class Crabstero(commands.Bot): """Central bot subclass that owns all lifecycle state. The database is opened in setup_hook and closed in close(). Background ingestion is handled by a bounded queue and a fixed worker pool. """ def __init__( self, token: str, database_path: str, ingestion_workers: int = 4, ingest_only: bool = False, ) -> None: """Configure intents, store configuration, and prepare ingestion queue state. :param token: The Discord bot token. :param database_path: The file path to the SQLite database. :param ingestion_workers: The number of concurrent ingestion workers. :param ingest_only: When True, the bot only ingests data and never responds. """ intents = discord.Intents.default() intents.guilds = True intents.guild_messages = True intents.message_content = True super().__init__( command_prefix=[], intents=intents, max_messages=None, # Disables the message cache. ) self._token = token self._database_path = database_path self._ingestion_worker_count = ingestion_workers self._ingest_only = ingest_only self._ingestion_queue: asyncio.Queue[ discord.TextChannel | discord.VoiceChannel ] = asyncio.Queue() self._ingestion_workers: list[asyncio.Task[None]] = [] self.db: Database repo_url = "https://git.logal.dev/LogalDeveloper/Crabstero" self.http.user_agent = f"DiscordBot ({repo_url}, {crabstero_version})" async def setup_hook(self) -> None: """Open the database, start ingestion workers, and load all cogs.""" self.db = await Database.connect(self._database_path) self._start_ingestion_workers() if not self._ingest_only: await interaction.setup(self) await message.setup(self) await server_events.setup(self) if not self._ingest_only: # Only sync slash commands if the registered commands # differ from local definitions. local_commands = { cmd.name: cmd.description for cmd in self.tree.get_commands() if isinstance(cmd, (app_commands.Command, app_commands.Group)) } try: remote_commands = { cmd.name: cmd.description for cmd in await self.tree.fetch_commands() } except discord.HTTPException: remote_commands = {} if local_commands != remote_commands: logger.info("Slash command tree has changed, syncing with Discord.") await self.tree.sync() async def on_ready(self) -> None: """Log that the bot has started successfully.""" logger.info("Crabstero started!") async def start(self, token: str = "", *, reconnect: bool = True) -> None: """Start the bot using the stored token by default. :param token: Optional token override. Falls back to the stored token if empty. :param reconnect: Whether to automatically reconnect on disconnect. """ await super().start(token or self._token, reconnect=reconnect) async def close(self) -> None: """Cancel ingestion workers, close the database, and then the bot connection.""" if self.is_closed(): return logger.info("Shutting down Crabstero...") for worker in self._ingestion_workers: worker.cancel() await asyncio.gather(*self._ingestion_workers, return_exceptions=True) self._ingestion_workers.clear() if hasattr(self, "db"): await self.db.close() await super().close() def queue_channel_for_ingestion( self, channel: discord.TextChannel | discord.VoiceChannel ) -> None: """Enqueue a single channel for background message history ingestion. :param channel: The channel to enqueue. """ self._ingestion_queue.put_nowait(channel) def _start_ingestion_workers(self) -> None: """Spawn the fixed pool of ingestion worker tasks.""" for _ in range(self._ingestion_worker_count): task = asyncio.create_task(self._ingestion_worker()) self._ingestion_workers.append(task) async def _ingestion_worker(self) -> None: """Loop forever pulling channels from the ingestion queue and ingesting them.""" while True: channel = await self._ingestion_queue.get() try: await ingest_channel(channel, self.db) finally: self._ingestion_queue.task_done()