feat: Implement Phase 3 (Queue & Permissions) with DM safety fix
This commit is contained in:
AEmotionStudio
+4
-2
@@ -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:
|
||||
|
||||
@@ -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))
|
||||
@@ -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))
|
||||
@@ -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)
|
||||
@@ -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")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user