mirror of
https://github.com/TheFunny/ArisuAutoSweeper
synced 2025-12-16 19:55:12 +00:00
115 lines
3.6 KiB
Python
115 lines
3.6 KiB
Python
"""
|
|
WebSocket handler for real-time log streaming
|
|
"""
|
|
import asyncio
|
|
from typing import Set
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
|
from module.webui.process_manager import ProcessManager
|
|
from module.logger import logger
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
class ConnectionManager:
|
|
"""Manage WebSocket connections"""
|
|
|
|
def __init__(self):
|
|
self.active_connections: Set[WebSocket] = set()
|
|
|
|
async def connect(self, websocket: WebSocket):
|
|
await websocket.accept()
|
|
self.active_connections.add(websocket)
|
|
|
|
def disconnect(self, websocket: WebSocket):
|
|
self.active_connections.discard(websocket)
|
|
|
|
async def broadcast(self, message: str):
|
|
for connection in self.active_connections.copy():
|
|
try:
|
|
await connection.send_text(message)
|
|
except Exception:
|
|
self.disconnect(connection)
|
|
|
|
|
|
manager = ConnectionManager()
|
|
|
|
|
|
@router.websocket("/logs/{instance_name}")
|
|
async def websocket_logs(websocket: WebSocket, instance_name: str):
|
|
"""WebSocket endpoint for streaming logs"""
|
|
await manager.connect(websocket)
|
|
|
|
try:
|
|
alas = ProcessManager.get_manager(instance_name)
|
|
|
|
# Send initial connection message
|
|
await websocket.send_json({
|
|
"type": "connected",
|
|
"instance": instance_name
|
|
})
|
|
|
|
# Keep connection alive and send log updates
|
|
while True:
|
|
try:
|
|
# Check if process is alive
|
|
if hasattr(alas, 'alive') and alas.alive:
|
|
await websocket.send_json({
|
|
"type": "status",
|
|
"alive": True,
|
|
"state": alas.state
|
|
})
|
|
|
|
# Wait a bit before next update
|
|
await asyncio.sleep(1)
|
|
|
|
# Check if client sent any message (to keep connection alive)
|
|
try:
|
|
data = await asyncio.wait_for(websocket.receive_text(), timeout=0.1)
|
|
if data == "ping":
|
|
await websocket.send_text("pong")
|
|
except asyncio.TimeoutError:
|
|
pass
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in WebSocket loop: {e}")
|
|
break
|
|
|
|
except WebSocketDisconnect:
|
|
logger.info(f"WebSocket disconnected for {instance_name}")
|
|
except Exception as e:
|
|
logger.exception(f"WebSocket error: {e}")
|
|
finally:
|
|
manager.disconnect(websocket)
|
|
|
|
|
|
@router.websocket("/system")
|
|
async def websocket_system(websocket: WebSocket):
|
|
"""WebSocket endpoint for system-wide updates"""
|
|
await manager.connect(websocket)
|
|
|
|
try:
|
|
await websocket.send_json({
|
|
"type": "connected",
|
|
"message": "System WebSocket connected"
|
|
})
|
|
|
|
# Keep connection alive
|
|
while True:
|
|
try:
|
|
data = await asyncio.wait_for(websocket.receive_text(), timeout=10)
|
|
if data == "ping":
|
|
await websocket.send_text("pong")
|
|
except asyncio.TimeoutError:
|
|
# Send heartbeat
|
|
await websocket.send_json({"type": "heartbeat"})
|
|
except Exception as e:
|
|
logger.error(f"Error in system WebSocket: {e}")
|
|
break
|
|
|
|
except WebSocketDisconnect:
|
|
logger.info("System WebSocket disconnected")
|
|
except Exception as e:
|
|
logger.exception(f"System WebSocket error: {e}")
|
|
finally:
|
|
manager.disconnect(websocket)
|