Gan, Jimmy
ddaf3c6cee
- Add JWT token auth to OPC WebSocket (unauthenticated → 403, per-user tracking) - Externalize Transmission RPC credentials to TRANSMISSION_USER/PASS env vars - Remove hardcoded SECRET_KEY fallback from dev compose - Rate-limit passkey register options endpoint at 5/minute - Add PDD (docs/improvement-plan.md) and sprint specs (specs/) Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
135 lines
4.4 KiB
Python
135 lines
4.4 KiB
Python
"""
|
|
OPC WebSocket Router - Real-time updates for tasks and agent executions
|
|
"""
|
|
|
|
import logging
|
|
from datetime import datetime
|
|
from typing import Any
|
|
|
|
from fastapi import APIRouter, HTTPException, Query, WebSocket, WebSocketDisconnect, status
|
|
|
|
import auth_service as auth
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
def serialize_datetime(obj: Any) -> Any:
|
|
"""Recursively serialize datetime objects to ISO format strings"""
|
|
if isinstance(obj, datetime):
|
|
return obj.isoformat()
|
|
elif isinstance(obj, dict):
|
|
return {key: serialize_datetime(value) for key, value in obj.items()}
|
|
elif isinstance(obj, list):
|
|
return [serialize_datetime(item) for item in obj]
|
|
else:
|
|
return obj
|
|
|
|
|
|
class ConnectionManager:
|
|
"""Manages WebSocket connections with per-user tracking"""
|
|
|
|
def __init__(self):
|
|
self.active_connections: dict[WebSocket, str] = {}
|
|
|
|
async def connect(self, websocket: WebSocket, username: str):
|
|
"""Accept and register a new connection"""
|
|
await websocket.accept()
|
|
self.active_connections[websocket] = username
|
|
logger.info(f"WebSocket connected: {username}. Total connections: {len(self.active_connections)}")
|
|
|
|
def disconnect(self, websocket: WebSocket):
|
|
"""Remove a connection"""
|
|
username = self.active_connections.pop(websocket, "unknown")
|
|
logger.info(f"WebSocket disconnected: {username}. Total connections: {len(self.active_connections)}")
|
|
|
|
async def broadcast(self, message: dict):
|
|
"""Broadcast message to all connected clients"""
|
|
disconnected = set()
|
|
for connection in self.active_connections:
|
|
try:
|
|
await connection.send_json(message)
|
|
except Exception as e:
|
|
logger.error(f"Failed to send message: {e}")
|
|
disconnected.add(connection)
|
|
|
|
# Clean up disconnected clients
|
|
for conn in disconnected:
|
|
self.disconnect(conn)
|
|
|
|
|
|
manager = ConnectionManager()
|
|
|
|
|
|
@router.websocket("/ws")
|
|
async def websocket_endpoint(websocket: WebSocket, token: str = Query(None)):
|
|
"""WebSocket endpoint for OPC real-time updates. Requires JWT token via ?token= query param."""
|
|
if not token:
|
|
await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="Missing authentication token")
|
|
return
|
|
|
|
try:
|
|
user = auth.user_from_access_token(token, include_auth_header=False)
|
|
except HTTPException:
|
|
await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="Invalid authentication token")
|
|
return
|
|
|
|
await manager.connect(websocket, user.username)
|
|
|
|
try:
|
|
while True:
|
|
# Keep connection alive and receive any client messages
|
|
data = await websocket.receive_text()
|
|
|
|
# Echo back for ping/pong
|
|
if data == "ping":
|
|
await websocket.send_text("pong")
|
|
|
|
except WebSocketDisconnect:
|
|
manager.disconnect(websocket)
|
|
except Exception as e:
|
|
logger.error(f"WebSocket error: {e}")
|
|
manager.disconnect(websocket)
|
|
|
|
|
|
async def broadcast_task_update(task_id: int, action: str, task_data: dict):
|
|
"""Broadcast task update to all connected clients"""
|
|
# Serialize all datetime objects recursively
|
|
serialized_data = serialize_datetime(task_data)
|
|
|
|
message = {
|
|
"type": "task_update",
|
|
"task_id": task_id,
|
|
"action": action, # created, updated, moved, deleted
|
|
"data": serialized_data,
|
|
"timestamp": datetime.utcnow().isoformat(),
|
|
}
|
|
await manager.broadcast(message)
|
|
|
|
|
|
async def broadcast_agent_execution(execution_id: int, status: str, execution_data: dict):
|
|
"""Broadcast agent execution update"""
|
|
# Serialize all datetime objects recursively
|
|
serialized_data = serialize_datetime(execution_data)
|
|
|
|
message = {
|
|
"type": "agent_execution",
|
|
"execution_id": execution_id,
|
|
"status": status, # pending, running, completed, failed, pending_approval
|
|
"data": serialized_data,
|
|
"timestamp": datetime.utcnow().isoformat(),
|
|
}
|
|
await manager.broadcast(message)
|
|
|
|
|
|
async def broadcast_agent_status(agent_id: str, status: str, message_text: str = ""):
|
|
"""Broadcast agent status update"""
|
|
message = {
|
|
"type": "agent_status",
|
|
"agent_id": agent_id,
|
|
"status": status, # idle, working, error
|
|
"message": message_text,
|
|
}
|
|
await manager.broadcast(message)
|