e6261532f7
## 主要修复 - 修复Creo软件运行状态检测失败问题 - 添加完整的软件停止功能支持 - 改进多进程软件的进程管理逻辑 ## 技术改进 - 更新软件配置支持多进程名称检测 - 优化进程停止逻辑，增加超时配置 - 新增 stop_software WebSocket消息类型 - 完善错误处理和日志记录 ## 配置更新 - configs/software_config.yaml: 支持进程名称列表和停止超时 - 添加Revit 2017配置支持 ## 文档更新 - README.md: 更新软件配置说明和API列表 - frontend-api-docs.md: 添加停止软件API文档 - CHECKPOINT.md: 记录修复进展和解决方案 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
126 lines
4.2 KiB
Python
126 lines
4.2 KiB
Python
"""
|
||
WebSocket连接管理器
|
||
管理所有WebSocket连接，提供消息广播和连接状态监控功能
|
||
"""
|
||
from typing import Dict, List, Optional
|
||
from fastapi import WebSocket, WebSocketDisconnect
|
||
import json
|
||
import asyncio
|
||
import logging
|
||
from enum import Enum
|
||
from datetime import datetime
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class DateTimeEncoder(json.JSONEncoder):
|
||
"""自定义JSON编码器，处理datetime对象"""
|
||
def default(self, obj):
|
||
if isinstance(obj, datetime):
|
||
return obj.isoformat()
|
||
return super().default(obj)
|
||
|
||
|
||
class MessageType(str, Enum):
|
||
"""消息类型枚举"""
|
||
SOFTWARE_STATUS = "software_status"
|
||
PROCESS_UPDATE = "process_update"
|
||
ERROR = "error"
|
||
INFO = "info"
|
||
HEARTBEAT = "heartbeat"
|
||
# 软件操作相关消息类型
|
||
TASK_UPDATE = "task_update"
|
||
SOFTWARE_STARTED = "software_started"
|
||
SOFTWARE_START_FAILED = "software_start_failed"
|
||
SOFTWARE_LIST_UPDATE = "software_list_update"
|
||
# 日志相关消息类型
|
||
LOG_OPERATION = "log_operation"
|
||
LOG_RECORDED = "log_recorded"
|
||
|
||
|
||
class WebSocketManager:
|
||
"""WebSocket连接管理器"""
|
||
|
||
def __init__(self):
|
||
# 存储活跃连接
|
||
self.active_connections: Dict[str, WebSocket] = {}
|
||
# 连接对应的用户信息
|
||
self.connection_users: Dict[str, str] = {}
|
||
|
||
async def connect(self, websocket: WebSocket, client_id: str, user_id: str = None):
|
||
"""接受WebSocket连接"""
|
||
await websocket.accept()
|
||
self.active_connections[client_id] = websocket
|
||
if user_id:
|
||
self.connection_users[client_id] = user_id
|
||
logger.info(f"WebSocket客户端 {client_id} 已连接，用户: {user_id}")
|
||
|
||
# 发送连接成功消息
|
||
await self.send_personal_message({
|
||
"type": MessageType.INFO,
|
||
"message": "WebSocket连接已建立",
|
||
"timestamp": self._get_timestamp()
|
||
}, client_id)
|
||
|
||
def disconnect(self, client_id: str):
|
||
"""断开WebSocket连接"""
|
||
if client_id in self.active_connections:
|
||
del self.active_connections[client_id]
|
||
if client_id in self.connection_users:
|
||
del self.connection_users[client_id]
|
||
logger.info(f"WebSocket客户端 {client_id} 已断开连接")
|
||
|
||
async def send_personal_message(self, message: dict, client_id: str):
|
||
"""发送消息给指定客户端"""
|
||
if client_id in self.active_connections:
|
||
try:
|
||
websocket = self.active_connections[client_id]
|
||
await websocket.send_text(json.dumps(message, ensure_ascii=False, cls=DateTimeEncoder))
|
||
except Exception as e:
|
||
logger.error(f"发送消息给客户端 {client_id} 失败: {e}")
|
||
self.disconnect(client_id)
|
||
|
||
async def broadcast(self, message: dict):
|
||
"""广播消息给所有连接的客户端"""
|
||
if not self.active_connections:
|
||
return
|
||
|
||
# 并发发送消息
|
||
tasks = []
|
||
for client_id in list(self.active_connections.keys()):
|
||
task = self.send_personal_message(message, client_id)
|
||
tasks.append(task)
|
||
|
||
if tasks:
|
||
await asyncio.gather(*tasks, return_exceptions=True)
|
||
|
||
async def broadcast_to_user(self, message: dict, user_id: str):
|
||
"""广播消息给指定用户的所有连接"""
|
||
target_clients = [
|
||
client_id for client_id, uid in self.connection_users.items()
|
||
if uid == user_id
|
||
]
|
||
|
||
tasks = []
|
||
for client_id in target_clients:
|
||
task = self.send_personal_message(message, client_id)
|
||
tasks.append(task)
|
||
|
||
if tasks:
|
||
await asyncio.gather(*tasks, return_exceptions=True)
|
||
|
||
def get_active_connections_count(self) -> int:
|
||
"""获取活跃连接数"""
|
||
return len(self.active_connections)
|
||
|
||
def get_connected_users(self) -> List[str]:
|
||
"""获取已连接的用户列表"""
|
||
return list(set(self.connection_users.values()))
|
||
|
||
def _get_timestamp(self) -> str:
|
||
"""获取当前时间戳"""
|
||
return datetime.now().isoformat()
|
||
|
||
|
||
# 全局WebSocket管理器实例
|
||
websocket_manager = WebSocketManager() |