#!/usr/bin/env python # -*- coding: utf-8 -*- """ 脚本WebSocket API路由 提供实时日志推送和脚本状态监控 """ from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query from fastapi.responses import HTMLResponse import json from typing import Dict, Any from services.script_websocket_service import get_websocket_manager from utils.logger import get_logger logger = get_logger("routes.script_websocket_api") router = APIRouter() @router.websocket("/ws/script-logs") async def script_logs_websocket( websocket: WebSocket, client_id: str = Query(..., description="客户端ID"), client_type: str = Query(default="web", description="客户端类型") ): """脚本日志实时推送WebSocket连接""" websocket_manager = get_websocket_manager() connection_id = None try: # 建立连接 connection_id = await websocket_manager.connect(websocket, client_id, client_type) while True: # 接收客户端消息 data = await websocket.receive_text() message = json.loads(data) message_type = message.get("type") if message_type == "subscribe": # 订阅脚本日志 script_id = message.get("script_id") if script_id: success = await websocket_manager.subscribe_script(connection_id, script_id) if not success: await websocket.send_text(json.dumps({ "type": "error", "message": f"订阅脚本 {script_id} 失败" })) elif message_type == "unsubscribe": # 取消订阅脚本日志 script_id = message.get("script_id") if script_id: success = await websocket_manager.unsubscribe_script(connection_id, script_id) if not success: await websocket.send_text(json.dumps({ "type": "error", "message": f"取消订阅脚本 {script_id} 失败" })) elif message_type == "ping": # 心跳响应 await websocket.send_text(json.dumps({ "type": "pong", "timestamp": message.get("timestamp") })) elif message_type == "get_status": # 获取连接状态 status = await websocket_manager.get_connection_status() await websocket.send_text(json.dumps({ "type": "status", "data": status })) else: await websocket.send_text(json.dumps({ "type": "error", "message": f"未知消息类型: {message_type}" })) except WebSocketDisconnect: logger.info(f"WebSocket客户端断开连接: {client_id}") except json.JSONDecodeError: logger.error(f"WebSocket消息JSON解析失败: {data}") try: await websocket.send_text(json.dumps({ "type": "error", "message": "消息格式错误，请发送有效的JSON" })) except: pass except Exception as e: logger.error(f"WebSocket处理异常: {e}", exc_info=True) finally: # 清理连接 if connection_id: await websocket_manager.disconnect(connection_id) @router.get("/ws/test-page") async def get_websocket_test_page(): """获取WebSocket测试页面""" html = """ VWED 脚本日志 WebSocket 测试

VWED 脚本日志 WebSocket 测试

连接控制

未连接

订阅控制

实时日志

""" return HTMLResponse(content=html) @router.get("/ws/connections/status", summary="获取WebSocket连接状态") async def get_websocket_connections_status(): """获取当前WebSocket连接状态""" try: websocket_manager = get_websocket_manager() status = await websocket_manager.get_connection_status() return { "success": True, "data": status } except Exception as e: logger.error(f"获取WebSocket连接状态失败: {e}", exc_info=True) return { "success": False, "error": f"获取WebSocket连接状态失败: {str(e)}" } @router.post("/ws/broadcast/test", summary="测试广播消息") async def test_broadcast_message(request: Dict[str, Any]): """测试向指定脚本的订阅者广播消息""" try: script_id = request.get("script_id") message = request.get("message", "测试消息") level = request.get("level", "info") if not script_id: return {"success": False, "error": "缺少script_id参数"} websocket_manager = get_websocket_manager() await websocket_manager.broadcast_script_log(script_id, level.upper(), message) return { "success": True, "message": f"测试消息已广播到脚本 {script_id} 的订阅者" } except Exception as e: logger.error(f"测试广播消息失败: {e}", exc_info=True) return { "success": False, "error": f"测试广播消息失败: {str(e)}" }