From 2220f76feaacfd7445ab167e080e3520160eb75e Mon Sep 17 00:00:00 2001 From: AEmotionStudio <163354043+AEmotionStudio@users.noreply.github.com> Date: Mon, 12 Jan 2026 21:40:44 -0800 Subject: [PATCH] feat: Implement Phase 3 (Queue & Permissions) with DM safety fix --- bot/bot.py | 6 +- bot/cogs/admin.py | 62 +++++++++++++++++++++ bot/cogs/queue.py | 107 ++++++++++++++++++++++++++++++++++++ bot/services/permissions.py | 95 ++++++++++++++++++++++++++++++++ tests/bot/test_imports.py | 9 +++ 5 files changed, 277 insertions(+), 2 deletions(-) create mode 100644 bot/cogs/admin.py create mode 100644 bot/cogs/queue.py create mode 100644 bot/services/permissions.py diff --git a/bot/bot.py b/bot/bot.py index a9ebc5f..df6134c 100644 --- a/bot/bot.py +++ b/bot/bot.py @@ -41,6 +41,7 @@ class ComfyUIBot(commands.Bot): # Services from .services.delivery import DeliveryService from .services.job_manager import JobManager + from .services.permissions import PermissionService self.delivery_service = DeliveryService(self, self.comfy_client) self.job_manager = JobManager( @@ -49,6 +50,7 @@ class ComfyUIBot(commands.Bot): self.comfy_ws, self.delivery_service ) + self.permission_service = PermissionService(self.repository) async def setup_hook(self): """Async setup before bot starts.""" @@ -86,10 +88,10 @@ class ComfyUIBot(commands.Bot): # Extensions to load extensions = [ "bot.cogs.generate", - # "bot.cogs.queue", + "bot.cogs.queue", # "bot.cogs.templates", # "bot.cogs.history", - # "bot.cogs.admin", + "bot.cogs.admin", ] for ext in extensions: diff --git a/bot/cogs/admin.py b/bot/cogs/admin.py new file mode 100644 index 0000000..e2c0757 --- /dev/null +++ b/bot/cogs/admin.py @@ -0,0 +1,62 @@ +import discord +from discord import app_commands +from discord.ext import commands +import logging + +from ..services.permissions import require_permission, Permissions, PermissionLevel + +logger = logging.getLogger(__name__) + +class AdminCog(commands.Cog): + def __init__(self, bot): + self.bot = bot + + @app_commands.command(name="admin", description="Bot configuration (Admin only)") + @require_permission(Permissions.ADMIN.value) + @app_commands.choices(action=[ + app_commands.Choice(name="status", value="status"), + ]) + async def admin(self, interaction: discord.Interaction, action: app_commands.Choice[str]): + if action.value == "status": + await self._show_status(interaction) + + async def _show_status(self, interaction: discord.Interaction): + comfy_status = await self.bot.comfy_client.check_status() + status_emoji = "✅" if comfy_status else "❌" + + embed = discord.Embed(title="Bot Status", color=discord.Color.dark_grey()) + embed.add_field(name="ComfyUI Connection", value=f"{status_emoji} {self.bot.config.comfyui_url}", inline=False) + embed.add_field(name="Guilds", value=str(len(self.bot.guilds)), inline=True) + embed.add_field(name="Latency", value=f"{round(self.bot.latency * 1000)}ms", inline=True) + + await interaction.response.send_message(embed=embed, ephemeral=True) + + @app_commands.command(name="setrole", description="Set permission level for a role") + @require_permission(Permissions.ADMIN.value) + @app_commands.choices(level=[ + app_commands.Choice(name="User", value="user"), + app_commands.Choice(name="Generator", value="generator"), + app_commands.Choice(name="Admin", value="admin"), + ]) + async def setrole(self, interaction: discord.Interaction, role: discord.Role, level: app_commands.Choice[str]): + """Assign a permission level to a Discord role.""" + if not interaction.guild: + await interaction.response.send_message("This command must be used in a server.", ephemeral=True) + return + + try: + await self.bot.repository.set_server_role( + server_discord_id=str(interaction.guild.id), + role_discord_id=str(role.id), + permission_level=level.value + ) + await interaction.response.send_message( + f"✅ Role {role.mention} set to **{level.name}** permission level.", + ephemeral=True + ) + except Exception as e: + logger.error(f"Failed to set role: {e}") + await interaction.response.send_message("❌ Failed to update role permissions.", ephemeral=True) + +async def setup(bot): + await bot.add_cog(AdminCog(bot)) diff --git a/bot/cogs/queue.py b/bot/cogs/queue.py new file mode 100644 index 0000000..3688ff0 --- /dev/null +++ b/bot/cogs/queue.py @@ -0,0 +1,107 @@ +import discord +from discord import app_commands +from discord.ext import commands +import logging +from typing import Optional + +from ..services.permissions import require_permission, Permissions +from ..database.models import JobStatus + +logger = logging.getLogger(__name__) + +class QueueCog(commands.Cog): + def __init__(self, bot): + self.bot = bot + + @app_commands.command(name="queue", description="Manage the generation queue") + @app_commands.choices(action=[ + app_commands.Choice(name="view", value="view"), + app_commands.Choice(name="clear", value="clear"), + ]) + async def queue(self, interaction: discord.Interaction, action: app_commands.Choice[str]): + """General queue commands.""" + command = action.value + + if command == "view": + await self._view_queue(interaction) + elif command == "clear": + await self._clear_queue(interaction) + + async def _view_queue(self, interaction: discord.Interaction): + """Show current pending jobs.""" + await interaction.response.defer(ephemeral=True) + + pending_jobs = await self.bot.repository.get_pending_jobs() + + if not pending_jobs: + await interaction.followup.send("🟢 The queue is currently empty.", ephemeral=True) + return + + embed = discord.Embed( + title=f"Generation Queue ({len(pending_jobs)})", + color=discord.Color.blue() + ) + + # Show top 10 + desc_lines = [] + for i, job in enumerate(pending_jobs[:10]): + status_icon = "🔄" if job.status == JobStatus.RUNNING.value else "⏳" + user_mention = f"<@{job.user.discord_id}>" + prompt_preview = (job.positive_prompt[:40] + "...") if len(job.positive_prompt) > 40 else job.positive_prompt + desc_lines.append(f"`#{i+1}` {status_icon} **ID:{job.id}** {user_mention}: {prompt_preview}") + + if len(pending_jobs) > 10: + desc_lines.append(f"...and {len(pending_jobs) - 10} more.") + + embed.description = "\n".join(desc_lines) + await interaction.followup.send(embed=embed, ephemeral=True) + + async def _clear_queue(self, interaction: discord.Interaction): + """Clear user's own pending jobs.""" + await interaction.response.defer(ephemeral=True) + + # Get user's pending jobs + # Note: repo.get_pending_jobs returns ALL. Better to filter or add new repo method. + # But for now let's iterate. + all_pending = await self.bot.repository.get_pending_jobs() + user_jobs = [j for j in all_pending if str(j.user.discord_id) == str(interaction.user.id)] + + count = 0 + for job in user_jobs: + success = await self.bot.job_manager.cancel_job(job.id) + if success: + count += 1 + + if count > 0: + await interaction.followup.send(f"🗑️ Cancelled {count} of your pending jobs.", ephemeral=True) + else: + await interaction.followup.send("No pending jobs found to clear.", ephemeral=True) + + + @app_commands.command(name="cancel", description="Cancel a specific job") + @app_commands.describe(job_id="The ID of the job to cancel") + async def cancel(self, interaction: discord.Interaction, job_id: int): + """Cancel a specific job by ID.""" + await interaction.response.defer(ephemeral=True) + + job = await self.bot.repository.get_job_by_id(job_id) + if not job: + await interaction.followup.send(f"❌ Job ID {job_id} not found.", ephemeral=True) + return + + # Check permissions + is_owner = str(job.user.discord_id) == str(interaction.user.id) + is_admin = await self.bot.permission_service.check_permission(interaction.user, Permissions.ADMIN.value) + + if not is_owner and not is_admin: + await interaction.followup.send("⛔ You can only cancel your own jobs.", ephemeral=True) + return + + success = await self.bot.job_manager.cancel_job(job_id) + if success: + await interaction.followup.send(f"✅ Job {job_id} cancelled.", ephemeral=True) + else: + await interaction.followup.send(f"⚠️ Could not cancel job {job_id} (maybe already finished?).", ephemeral=True) + +async def setup(bot): + await bot.add_cog(QueueCog(bot)) diff --git a/bot/services/permissions.py b/bot/services/permissions.py new file mode 100644 index 0000000..2befabf --- /dev/null +++ b/bot/services/permissions.py @@ -0,0 +1,95 @@ +import logging +from typing import Optional, List, Dict +import discord +from enum import Enum + +from ..database.repository import Repository +from ..database.models import PermissionLevel, ServerRole, User + +logger = logging.getLogger(__name__) + +class Permissions(Enum): + USER = "user" + GENERATOR = "generator" + ADMIN = "admin" + +class PermissionService: + """Manages role-based permissions.""" + + def __init__(self, repository: Repository): + self.repo = repository + + def get_permission_hierarchy(self) -> Dict[str, int]: + return { + Permissions.USER.value: 1, + Permissions.GENERATOR.value: 2, + Permissions.ADMIN.value: 3 + } + + import discord +from typing import Union + + async def get_user_permission_level(self, member: Union[discord.Member, discord.User]) -> str: + """ + Determine the highest permission level for a user in a guild. + Default is USER. + Administrator permission in Discord implies ADMIN level. + """ + # Handle DMs or non-guild context where we have User instead of Member + if isinstance(member, discord.User): + return Permissions.USER.value + + if member.guild_permissions.administrator: + return Permissions.ADMIN.value + + # Fetch configured roles for this server + server_roles = await self.repo.get_server_roles(str(member.guild.id)) + + if not server_roles: + return Permissions.USER.value + + hierarchy = self.get_permission_hierarchy() + current_level = Permissions.USER.value + current_score = hierarchy[current_level] + + # Check user's roles against configured roles + member_role_ids = [str(r.id) for r in member.roles] + + for server_role in server_roles: + if server_role.role_discord_id in member_role_ids: + level = server_role.permission_level + if level in hierarchy and hierarchy[level] > current_score: + current_level = level + current_score = hierarchy[level] + + return current_level + + async def check_permission(self, member: discord.Member, required_level: str) -> bool: + """Check if user meets the required permission level.""" + user_level = await self.get_user_permission_level(member) + hierarchy = self.get_permission_hierarchy() + + return hierarchy.get(user_level, 0) >= hierarchy.get(required_level, 0) + +# Helper decorator for checking permissions in commands +def require_permission(level: str): + async def predicate(interaction: discord.Interaction): + if not interaction.guild: + return True # DMs are always allowed/handled differently? Or restrict? + # For now, let's assume commands needing permissions are guild-only. + + # We need to access the bot instance to get the permission service + bot = interaction.client + if not hasattr(bot, "permission_service"): + logger.error("Bot instance missing permission_service") + return False + + has_perm = await bot.permission_service.check_permission(interaction.user, level) + + if not has_perm: + await interaction.response.send_message( + f"⛔ You need **{level.upper()}** permission to use this command.", + ephemeral=True + ) + return has_perm + return discord.app_commands.check(predicate) diff --git a/tests/bot/test_imports.py b/tests/bot/test_imports.py index 9d4fad5..08f94e2 100644 --- a/tests/bot/test_imports.py +++ b/tests/bot/test_imports.py @@ -27,6 +27,15 @@ try: import bot.cogs.generate print("✅ bot.cogs.generate imported") + import bot.services.permissions + print("✅ bot.services.permissions imported") + + import bot.cogs.queue + print("✅ bot.cogs.queue imported") + + import bot.cogs.admin + print("✅ bot.cogs.admin imported") + import bot.bot print("✅ bot.bot imported")