Initial commit — ComfyUI Discord bot + web UI
Full source for the-third-rev: Discord bot (discord.py), FastAPI web UI (React/TS/Vite/Tailwind), ComfyUI integration, generation history DB, preset manager, workflow inspector, and all supporting modules. Excluded from tracking: .env, invite_tokens.json, *.db (SQLite), current-workflow-changes.json, user_settings/, presets/, logs/, web-static/ (build output), frontend/node_modules/. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
61
web/routers/ws_router.py
Normal file
61
web/routers/ws_router.py
Normal file
@@ -0,0 +1,61 @@
|
||||
"""WebSocket /ws?token=<jwt> — real-time event stream"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
||||
|
||||
from web.auth import verify_ws_token
|
||||
from web.ws_bus import get_bus
|
||||
|
||||
router = APIRouter()
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@router.websocket("/ws")
|
||||
async def websocket_endpoint(websocket: WebSocket, token: str = ""):
|
||||
"""
|
||||
Authenticate via JWT query param or ttb_session cookie, then stream events from WSBus.
|
||||
|
||||
Events common to all users: status_snapshot, queue_update, node_executing, server_state
|
||||
Events private to submitter: generation_complete, generation_error
|
||||
"""
|
||||
payload = verify_ws_token(token)
|
||||
if payload is None:
|
||||
# Fallback: browsers send cookies automatically with WebSocket connections
|
||||
cookie_token = websocket.cookies.get("ttb_session", "")
|
||||
payload = verify_ws_token(cookie_token)
|
||||
if payload is None:
|
||||
await websocket.close(code=4001, reason="Unauthorized")
|
||||
return
|
||||
|
||||
user_label: str = payload.get("sub", "anonymous")
|
||||
bus = get_bus()
|
||||
queue = bus.subscribe(user_label)
|
||||
await websocket.accept()
|
||||
logger.info("WS connected: user=%s", user_label)
|
||||
|
||||
try:
|
||||
while True:
|
||||
# Wait for an event from the bus
|
||||
try:
|
||||
frame = await asyncio.wait_for(queue.get(), timeout=30.0)
|
||||
except asyncio.TimeoutError:
|
||||
# Send a keepalive ping
|
||||
try:
|
||||
await websocket.send_text('{"type":"ping"}')
|
||||
except Exception:
|
||||
break
|
||||
continue
|
||||
|
||||
try:
|
||||
await websocket.send_text(frame)
|
||||
except Exception:
|
||||
break
|
||||
|
||||
except WebSocketDisconnect:
|
||||
pass
|
||||
finally:
|
||||
bus.unsubscribe(user_label, queue)
|
||||
logger.info("WS disconnected: user=%s", user_label)
|
||||
Reference in New Issue
Block a user