""" bot.py ====== Discord bot entry point. In WEB_ENABLED mode, also starts a FastAPI/Uvicorn web server in the same asyncio event loop via asyncio.gather. Jobs are submitted directly to ComfyUI — no internal SerialJobQueue. """ from __future__ import annotations import asyncio import json import logging from datetime import datetime, timezone from pathlib import Path import discord from discord.ext import commands from comfy_client import ComfyClient from config import BotConfig, COMMAND_PREFIX import generation_db from input_image_db import init_db, get_all_images from status_monitor import StatusMonitor from workflow_manager import WorkflowManager from workflow_state import WorkflowStateManager from commands import register_all_commands, CustomHelpCommand from commands.input_images import PersistentSetInputView from commands.server import autostart_comfy try: from dotenv import load_dotenv load_dotenv() except Exception: pass logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class _UvicornWSNoiseFilter(logging.Filter): """Suppress per-connection WebSocket rejection noise from Uvicorn. When an unknown path (e.g. /waapi probes) hits the server as a WebSocket upgrade, Uvicorn logs three lines at INFO: - uvicorn.access: '... "WebSocket /waapi" 403' - uvicorn.error: 'connection rejected (403 Forbidden)' - uvicorn.error: 'connection closed' These are expected and handled gracefully; we just don't want them filling the log. """ def filter(self, record: logging.LogRecord) -> bool: msg = record.getMessage() if '"WebSocket ' in msg and ' 403' in msg: return False if msg.startswith("connection rejected") or msg == "connection closed": return False return True _PROJECT_ROOT = Path(__file__).resolve().parent # --------------------------------------------------------------------------- # Bot setup # --------------------------------------------------------------------------- def get_prefix(bot, message): """Dynamic command prefix getter.""" msg = message.content.lower() if msg.startswith(COMMAND_PREFIX): return COMMAND_PREFIX return COMMAND_PREFIX intents = discord.Intents.default() intents.message_content = True intents.guilds = True bot = commands.Bot( command_prefix=get_prefix, intents=intents, help_command=CustomHelpCommand(), ) # --------------------------------------------------------------------------- # Event handlers # --------------------------------------------------------------------------- @bot.event async def on_ready() -> None: logger.info("Logged in as %s (ID: %s)", bot.user, bot.user.id) if not hasattr(bot, "start_time"): bot.start_time = datetime.now(timezone.utc) cfg = getattr(bot, "config", None) if cfg: for row in get_all_images(): view = PersistentSetInputView(bot, cfg, row["id"]) bot.add_view(view, message_id=row["bot_reply_id"]) asyncio.create_task(autostart_comfy(cfg)) if not hasattr(bot, "status_monitor"): log_ch = getattr(getattr(bot, "config", None), "log_channel_id", None) if log_ch: bot.status_monitor = StatusMonitor(bot, log_ch) if hasattr(bot, "status_monitor"): await bot.status_monitor.start() @bot.event async def on_disconnect() -> None: logger.info("Discord connection closed") @bot.event async def on_resumed() -> None: logger.info("Discord session resumed") # --------------------------------------------------------------------------- # Startup helpers # --------------------------------------------------------------------------- async def create_comfy(config: BotConfig) -> ComfyClient: state_manager = WorkflowStateManager(state_file="current-workflow-changes.json") workflow_manager = WorkflowManager() return ComfyClient( server_address=config.comfy_server, workflow_manager=workflow_manager, state_manager=state_manager, logger=logger, history_limit=config.comfy_history_limit, output_path=config.comfy_output_path, ) def _try_autoload_last_workflow(client: ComfyClient) -> None: """Re-load the last used workflow from the workflows/ folder on startup.""" last_wf = client.state_manager.get_last_workflow_file() if not last_wf: return wf_path = _PROJECT_ROOT / "workflows" / last_wf if not wf_path.exists(): logger.warning("Last workflow file not found: %s", wf_path) return try: with open(wf_path, "r", encoding="utf-8") as f: workflow = json.load(f) # Restore template without clearing overrides on restart client.workflow_manager.set_workflow_template(workflow) logger.info("Auto-loaded last workflow: %s", last_wf) except Exception as exc: logger.error("Failed to auto-load workflow %s: %s", last_wf, exc) # --------------------------------------------------------------------------- # Main entry point # --------------------------------------------------------------------------- async def main() -> None: try: config: BotConfig = BotConfig.from_env() logger.info("Configuration loaded") except RuntimeError as exc: logger.error("Configuration error: %s", exc) return bot.comfy = await create_comfy(config) bot.config = config # Auto-load last workflow (restores template without clearing overrides) _try_autoload_last_workflow(bot.comfy) # Fallback: WORKFLOW_FILE env var if not bot.comfy.get_workflow_template() and config.workflow_file: try: bot.comfy.load_workflow_from_file(config.workflow_file) logger.info("Loaded workflow from %s", config.workflow_file) except Exception as exc: logger.error("Failed to load workflow %s: %s", config.workflow_file, exc) from user_state_registry import UserStateRegistry bot.user_registry = UserStateRegistry( settings_dir=_PROJECT_ROOT / "user_settings", default_workflow=bot.comfy.get_workflow_template(), ) init_db() generation_db.init_db(_PROJECT_ROOT / "generation_history.db") register_all_commands(bot, config) logger.info("All commands registered") coroutines = [bot.start(config.discord_bot_token)] if config.web_enabled: try: import uvicorn from web.deps import set_bot from web.app import create_app set_bot(bot) fastapi_app = create_app() uvi_config = uvicorn.Config( fastapi_app, host=config.web_host, port=config.web_port, log_level="info", loop="none", # use existing event loop ) uvi_server = uvicorn.Server(uvi_config) _ws_filter = _UvicornWSNoiseFilter() logging.getLogger("uvicorn.access").addFilter(_ws_filter) logging.getLogger("uvicorn.error").addFilter(_ws_filter) coroutines.append(uvi_server.serve()) logger.info( "Web UI enabled at http://%s:%d", config.web_host, config.web_port ) except ImportError: logger.warning( "uvicorn or fastapi not installed — web UI disabled. " "pip install fastapi uvicorn[standard]" ) try: await asyncio.gather(*coroutines) finally: if hasattr(bot, "status_monitor"): await bot.status_monitor.stop() if hasattr(bot, "comfy"): await bot.comfy.close() if __name__ == "__main__": try: asyncio.run(main()) except KeyboardInterrupt: logger.info("Received interrupt, shutting down…")