1ed3c9ec4b
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>
62 lines
1.9 KiB
Python
62 lines
1.9 KiB
Python
"""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)
|