1748cbf8d2
Unknown-path WS probes (e.g. /waapi) generate three INFO lines from uvicorn.access + uvicorn.error on every attempt. Install a logging.Filter on both loggers at startup to drop: - access log entries: "WebSocket <path>" 403 - error log entries: "connection rejected ..." / "connection closed" These are handled gracefully by _SPAStaticFiles; the logs add no value. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
244 lines
7.7 KiB
Python
244 lines
7.7 KiB
Python
"""
|
|
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…")
|