691 lines
30 KiB
Python
691 lines
30 KiB
Python
"""
|
||
工作手机SDK v3.0 - WebSocket Hub
|
||
管理所有设备的WebSocket连接
|
||
"""
|
||
|
||
from fastapi import WebSocket
|
||
from typing import Dict, Optional, Callable, Any
|
||
import asyncio
|
||
import json
|
||
import uuid
|
||
import logging
|
||
from datetime import datetime
|
||
from config import settings
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 延迟导入避免循环依赖，在 handle_message 内使用
|
||
def _get_device_manager():
|
||
from services.device_manager import device_manager
|
||
return device_manager
|
||
|
||
|
||
async def _record_device_event(device_id: str, event_type: str, payload: dict) -> None:
|
||
"""把真实 WS 设备事件写入 Hook 事件总线，供存客宝/超管/AI 数字员工实时订阅。"""
|
||
try:
|
||
from services.hook_module_service import hook_module_service
|
||
|
||
await hook_module_service.add_event({
|
||
"source": "ws_agent",
|
||
"platform": "android",
|
||
"device_id": device_id,
|
||
"event_type": event_type,
|
||
"payload": payload,
|
||
})
|
||
except Exception as e:
|
||
logger.debug(f"设备事件落总线失败 [{device_id}] {event_type}: {e}")
|
||
|
||
|
||
class WebSocketHub:
|
||
"""WebSocket连接管理中心"""
|
||
|
||
def __init__(self):
|
||
# 设备连接: device_id -> WebSocket
|
||
self.connections: Dict[str, WebSocket] = {}
|
||
|
||
# 设备信息: device_id -> device_info
|
||
self.device_info: Dict[str, dict] = {}
|
||
|
||
# 项目设备映射: project_id -> [device_id, ...]
|
||
self.project_devices: Dict[str, list] = {}
|
||
|
||
# 命令响应: command_id -> asyncio.Future
|
||
self.pending_commands: Dict[str, asyncio.Future] = {}
|
||
# 设备在途命令计数（sweep 时跳过误踢）
|
||
self.pending_device_commands: Dict[str, int] = {}
|
||
|
||
# 命令结果回调: command_id -> result
|
||
self.command_results: Dict[str, dict] = {}
|
||
|
||
# 消息处理器
|
||
self.message_handlers: Dict[str, Callable] = {}
|
||
# 设备心跳配置: device_id -> interval_seconds
|
||
self.heartbeat_config: Dict[str, int] = {}
|
||
|
||
async def connect(self, websocket: WebSocket, device_id: str):
|
||
"""设备连接"""
|
||
await websocket.accept()
|
||
self.connections[device_id] = websocket
|
||
logger.info(f"✅ 设备连接: {device_id}")
|
||
|
||
async def disconnect(self, device_id: str):
|
||
"""设备断开"""
|
||
old_info = dict(self.device_info.get(device_id) or {})
|
||
# 从项目中移除
|
||
if device_id in self.device_info:
|
||
project_id = self.device_info[device_id].get("project_id")
|
||
if project_id and project_id in self.project_devices:
|
||
if device_id in self.project_devices[project_id]:
|
||
self.project_devices[project_id].remove(device_id)
|
||
del self.device_info[device_id]
|
||
|
||
if device_id in self.connections:
|
||
del self.connections[device_id]
|
||
|
||
logger.info(f"❌ 设备断开: {device_id}")
|
||
await _record_device_event(device_id, "device_disconnected", {
|
||
"project_id": old_info.get("project_id"),
|
||
"model": old_info.get("model"),
|
||
"last_heartbeat": old_info.get("last_heartbeat"),
|
||
})
|
||
|
||
def is_online(self, device_id: str) -> bool:
|
||
"""检查设备是否在线"""
|
||
return device_id in self.connections
|
||
|
||
async def handle_message(self, device_id: str, data: dict):
|
||
"""处理设备消息"""
|
||
msg_type = data.get("type")
|
||
|
||
if msg_type == "register":
|
||
# 设备注册（支持APP直接发送的格式）
|
||
# APP发送: {type: "register", device_id: xxx, project_id: xxx, ...}
|
||
# Agent发送: {type: "register", data: {...}}
|
||
|
||
device_data = data.get("data") or data
|
||
project_id = device_data.get("project_id", "")
|
||
heartbeat_interval = int(device_data.get("heartbeat_interval_seconds") or settings.WS_HEARTBEAT_INTERVAL)
|
||
heartbeat_interval = max(5, min(120, heartbeat_interval))
|
||
|
||
now_iso = datetime.now().isoformat()
|
||
self.device_info[device_id] = {
|
||
**device_data,
|
||
"device_id": device_id,
|
||
"project_id": project_id,
|
||
"status": "online",
|
||
"connected_at": now_iso,
|
||
"last_heartbeat": now_iso,
|
||
"heartbeat_interval_seconds": heartbeat_interval,
|
||
}
|
||
self.heartbeat_config[device_id] = heartbeat_interval
|
||
|
||
# 添加到项目设备映射
|
||
if project_id:
|
||
if project_id not in self.project_devices:
|
||
self.project_devices[project_id] = []
|
||
if device_id not in self.project_devices[project_id]:
|
||
self.project_devices[project_id].append(device_id)
|
||
logger.info(f"📱 设备注册: {device_id} -> 项目: {project_id}")
|
||
else:
|
||
logger.info(f"📱 设备注册: {device_id} (无项目绑定)")
|
||
|
||
# 先发确认（不被 MongoDB 阻塞）
|
||
await self.send_to_device(device_id, {
|
||
"type": "registered",
|
||
"success": True,
|
||
"device_id": device_id,
|
||
"project_id": project_id,
|
||
"heartbeat_interval_seconds": heartbeat_interval,
|
||
})
|
||
|
||
# 注册落库 fire-and-forget
|
||
_reg_data = {
|
||
"device_id": device_id, "project_id": project_id,
|
||
"status": "online", "connected_at": now_iso, "last_heartbeat": now_iso,
|
||
**{k: v for k, v in device_data.items() if k not in ("device_id",)}
|
||
}
|
||
async def _reg_persist():
|
||
try:
|
||
dm = _get_device_manager()
|
||
await dm.register_device(_reg_data)
|
||
except Exception as e:
|
||
logger.warning(f"设备落库失败: {e}")
|
||
asyncio.create_task(_reg_persist())
|
||
asyncio.create_task(_record_device_event(device_id, "device_registered", {
|
||
"project_id": device_data.get("project_id"),
|
||
"model": device_data.get("model"),
|
||
"agent_version": device_data.get("agent_version"),
|
||
"capabilities": device_data.get("capabilities") or [],
|
||
"frida_available": device_data.get("frida_available"),
|
||
"ai_brain_enabled": device_data.get("ai_brain_enabled"),
|
||
}))
|
||
|
||
elif msg_type == "heartbeat":
|
||
now_iso = datetime.now().isoformat()
|
||
if device_id in self.device_info:
|
||
self.device_info[device_id]["last_heartbeat"] = now_iso
|
||
hb_status = data.get("status") or {}
|
||
if hb_status:
|
||
self.device_info[device_id]["quick_status"] = hb_status
|
||
_sync_fields = [
|
||
"wechat_running", "wechat_foreground", "network_type",
|
||
"screen_on", "current_app", "battery_level",
|
||
"battery_charging", "memory_usage_pct", "storage_free_mb",
|
||
"uptime_sec", "cpu_usage_pct",
|
||
"connect_stage", # BIND-07 寻服阶段实时同步
|
||
]
|
||
for _f in _sync_fields:
|
||
if hb_status.get(_f) is not None:
|
||
self.device_info[device_id][_f] = hb_status[_f]
|
||
# 兼容旧 Agent（≤3.x/4.0 心跳用 "battery" 键）：统一落 battery_level
|
||
if hb_status.get("battery") is not None and hb_status.get("battery_level") is None:
|
||
self.device_info[device_id]["battery_level"] = hb_status["battery"]
|
||
self.device_info[device_id]["heartbeat_count"] = (
|
||
self.device_info[device_id].get("heartbeat_count", 0) + 1
|
||
)
|
||
guard_data = data.get("guard")
|
||
if guard_data:
|
||
self.device_info[device_id]["guard"] = guard_data
|
||
self._store_guard_events(device_id, guard_data)
|
||
# 先回 pong（不被 MongoDB 阻塞）；双 type 兼容旧 Agent
|
||
pong_payload: dict = {
|
||
"type": "pong",
|
||
"server_time": now_iso,
|
||
"heartbeat_ack": True,
|
||
}
|
||
from services.ai_heartbeat import ai_heartbeat
|
||
q = ai_heartbeat.get_task_queue_status(device_id)
|
||
if q.get("pending_count", 0) > 0:
|
||
pong_payload["pending_tasks"] = q["pending_count"]
|
||
await self.send_to_device(device_id, pong_payload)
|
||
|
||
# 心跳落库改为 fire-and-forget（MongoDB 挂了不影响 ACK）
|
||
async def _hb_persist():
|
||
try:
|
||
dm = _get_device_manager()
|
||
await dm.update_heartbeat(device_id)
|
||
except Exception as e:
|
||
logger.debug(f"心跳落库: {e}")
|
||
asyncio.create_task(_hb_persist())
|
||
|
||
elif msg_type in ("response", "result"):
|
||
# 命令响应
|
||
# - Python Agent: {type:"response", code, message, data}
|
||
# - Android App: {type:"result", success, message}
|
||
command_id = data.get("command_id")
|
||
if not command_id:
|
||
return
|
||
if msg_type == "result":
|
||
success = bool(data.get("success"))
|
||
msg = data.get("message") or ("success" if success else "failed")
|
||
raw_data = data.get("data") or {}
|
||
if not raw_data and msg:
|
||
raw_data = {"output": msg}
|
||
normalized = {
|
||
"type": "response",
|
||
"command_id": command_id,
|
||
"device_id": device_id,
|
||
"code": 200 if success else 500,
|
||
"message": msg,
|
||
"data": raw_data,
|
||
"timestamp": data.get("timestamp"),
|
||
}
|
||
data = normalized
|
||
if command_id in self.pending_commands:
|
||
future = self.pending_commands.pop(command_id)
|
||
cnt = self.pending_device_commands.get(device_id, 0) - 1
|
||
if cnt <= 0:
|
||
self.pending_device_commands.pop(device_id, None)
|
||
else:
|
||
self.pending_device_commands[device_id] = cnt
|
||
if device_id in self.device_info:
|
||
self.device_info[device_id].pop("command_grace_seconds", None)
|
||
if not future.done():
|
||
future.set_result(data)
|
||
|
||
elif msg_type == "status_report":
|
||
# 详细状态上报（设备端定期上报）
|
||
status_payload = data.get("data", data)
|
||
if device_id in self.device_info:
|
||
self.device_info[device_id]["last_status"] = status_payload
|
||
self.device_info[device_id]["last_heartbeat"] = datetime.now().isoformat()
|
||
# 与心跳同源：把电量/网络等关键状态提升到顶层 device_info，
|
||
# 让 fleet/status、设备详情等聚合端「各端统一」读到真实电量（否则只在 last_status 里）。
|
||
if isinstance(status_payload, dict):
|
||
_sync_fields = [
|
||
"wechat_running", "wechat_foreground", "network_type",
|
||
"screen_on", "current_app", "battery_level",
|
||
"battery_charging", "battery_status", "charging",
|
||
"memory_usage_pct", "storage_free_mb",
|
||
"uptime_sec", "cpu_usage_pct", "connect_stage",
|
||
]
|
||
for _f in _sync_fields:
|
||
if status_payload.get(_f) is not None:
|
||
self.device_info[device_id][_f] = status_payload[_f]
|
||
# 旧 Agent 用 "battery" 键：统一落 battery_level
|
||
if status_payload.get("battery") is not None and status_payload.get("battery_level") is None:
|
||
self.device_info[device_id]["battery_level"] = status_payload["battery"]
|
||
logger.debug(f"设备状态上报: {device_id}")
|
||
asyncio.create_task(_record_device_event(device_id, "device_status_report", status_payload))
|
||
|
||
elif msg_type == "event":
|
||
event_name = data.get("event", "unknown")
|
||
event_data = data.get("data", {})
|
||
logger.info(f"设备事件 [{device_id}] {event_name}: {event_data}")
|
||
asyncio.create_task(_record_device_event(device_id, event_name, event_data))
|
||
|
||
if event_name == "ai_brain_acted":
|
||
await self._handle_ai_brain_acted(device_id, event_data)
|
||
elif event_name == "offline_buffer_upload":
|
||
await self._handle_offline_buffer(device_id, event_data)
|
||
elif event_name == "agent_started":
|
||
if device_id in self.device_info:
|
||
self.device_info[device_id]["frida_available"] = event_data.get("frida_available", False)
|
||
self.device_info[device_id]["ai_brain_enabled"] = event_data.get("ai_brain_enabled", False)
|
||
|
||
handler = self.message_handlers.get("event")
|
||
if handler:
|
||
await handler(device_id, data)
|
||
|
||
elif msg_type == "device_request":
|
||
# 设备端请求服务端执行操作（拉配置、落库、通知业务等）
|
||
request_id = data.get("request_id", "")
|
||
action = data.get("action", "")
|
||
params = data.get("params", {})
|
||
logger.info(f"设备请求 [{device_id}] {action} request_id={request_id}")
|
||
ack = await self._handle_device_request(device_id, request_id, action, params)
|
||
if ack:
|
||
await self.send_to_device(device_id, ack)
|
||
handler = self.message_handlers.get("device_request")
|
||
if handler:
|
||
await handler(device_id, data)
|
||
|
||
else:
|
||
# 自定义处理器
|
||
handler = self.message_handlers.get(msg_type)
|
||
if handler:
|
||
await handler(device_id, data)
|
||
|
||
async def _handle_device_request(
|
||
self, device_id: str, request_id: str, action: str, params: dict
|
||
) -> Optional[dict]:
|
||
"""处理设备端 device_request，返回要下发给设备的 ack（可选）"""
|
||
if action == "get_config":
|
||
interval = self.heartbeat_config.get(device_id, settings.WS_HEARTBEAT_INTERVAL)
|
||
config_data = {
|
||
"heartbeat_interval": interval,
|
||
"heartbeat_interval_seconds": interval,
|
||
"ws_timeout_seconds": settings.WS_TIMEOUT,
|
||
}
|
||
ai_cfg = getattr(settings, "AI_BRAIN_CONFIG", None)
|
||
if ai_cfg:
|
||
config_data["ai_brain"] = ai_cfg
|
||
return {
|
||
"type": "device_request_ack",
|
||
"request_id": request_id,
|
||
"success": True,
|
||
"data": config_data,
|
||
}
|
||
if action == "log_result":
|
||
# 仅记录日志，可选落库由 message_handlers 扩展
|
||
logger.info(f"设备请求落库 [{device_id}] {params}")
|
||
return {
|
||
"type": "device_request_ack",
|
||
"request_id": request_id,
|
||
"success": True,
|
||
}
|
||
logger.warning(f"未知 device_request action: {action}")
|
||
return {
|
||
"type": "device_request_ack",
|
||
"request_id": request_id,
|
||
"success": False,
|
||
"error": f"unknown action: {action}",
|
||
}
|
||
|
||
# ================================================================
|
||
# AI Brain 服务端支持
|
||
# ================================================================
|
||
|
||
async def _handle_ai_brain_acted(self, device_id: str, event_data: dict):
|
||
"""处理 AI Brain 自主执行事件"""
|
||
results_count = event_data.get("results_count", 0)
|
||
reason = event_data.get("reason", "")
|
||
logger.info(f"🧠 AI Brain [{device_id}] 执行了 {results_count} 个操作: {reason}")
|
||
if device_id in self.device_info:
|
||
ai_stats = self.device_info[device_id].setdefault("ai_brain_stats", {
|
||
"total_acts": 0, "last_act": None,
|
||
})
|
||
ai_stats["total_acts"] += results_count
|
||
ai_stats["last_act"] = datetime.now().isoformat()
|
||
ai_stats["last_reason"] = reason
|
||
|
||
async def _handle_offline_buffer(self, device_id: str, event_data: dict):
|
||
"""处理设备离线期间缓冲的执行结果"""
|
||
count = event_data.get("count", 0)
|
||
results = event_data.get("results", [])
|
||
logger.info(f"📥 离线缓冲上传 [{device_id}]: {count} 条结果")
|
||
if device_id in self.device_info:
|
||
self.device_info[device_id].setdefault("offline_uploads", []).append({
|
||
"uploaded_at": datetime.now().isoformat(),
|
||
"count": count,
|
||
})
|
||
try:
|
||
dm = _get_device_manager()
|
||
await dm.log_offline_results(device_id, results)
|
||
except Exception as e:
|
||
logger.debug(f"离线缓冲落库: {e}")
|
||
|
||
async def push_ai_task(self, device_id: str, instruction: str, priority: int = 5) -> bool:
|
||
"""向设备推送 AI 任务"""
|
||
return await self.send_to_device(device_id, {
|
||
"type": "ai_task",
|
||
"data": {"instruction": instruction, "priority": priority},
|
||
})
|
||
|
||
async def push_standing_order(self, device_id: str, order: str) -> bool:
|
||
"""向设备推送常驻指令"""
|
||
return await self.send_to_device(device_id, {
|
||
"type": "standing_order",
|
||
"data": {"order": order},
|
||
})
|
||
|
||
def get_ai_brain_status(self, device_id: str) -> Optional[dict]:
|
||
"""获取设备 AI Brain 状态"""
|
||
info = self.device_info.get(device_id)
|
||
if not info:
|
||
return None
|
||
return {
|
||
"device_id": device_id,
|
||
"ai_brain_enabled": info.get("ai_brain_enabled", False),
|
||
"frida_available": info.get("frida_available", False),
|
||
"ai_brain_stats": info.get("ai_brain_stats"),
|
||
"offline_uploads": info.get("offline_uploads", [])[-5:],
|
||
"last_status": info.get("last_status", {}),
|
||
}
|
||
|
||
def _store_guard_events(self, device_id: str, guard_data: dict):
|
||
"""将心跳中的守护数据存入事件存储（供 API 查询）"""
|
||
try:
|
||
from routers.devices import _guard_event_store, _MAX_EVENTS_PER_DEVICE
|
||
import time
|
||
if device_id not in _guard_event_store:
|
||
_guard_event_store[device_id] = []
|
||
store = _guard_event_store[device_id]
|
||
|
||
popups = guard_data.get("popups")
|
||
if popups and popups.get("count", 0) > 0:
|
||
store.append({
|
||
"type": "popup_dismissed",
|
||
"dismissed": popups.get("dismissed", []),
|
||
"count": popups["count"],
|
||
"server_received_at": time.time(),
|
||
})
|
||
|
||
full_cycle = guard_data.get("full_cycle")
|
||
if full_cycle:
|
||
store.append({
|
||
"type": "guard_cycle",
|
||
"all_ok": full_cycle.get("all_ok", True),
|
||
"u2_alive": full_cycle.get("u2", {}).get("alive"),
|
||
"network_ok": full_cycle.get("network", {}).get("network"),
|
||
"screen_on": full_cycle.get("screen", {}).get("screen_on"),
|
||
"server_received_at": time.time(),
|
||
})
|
||
if not full_cycle.get("all_ok"):
|
||
store.append({
|
||
"type": "guard_alert",
|
||
"issues": full_cycle,
|
||
"server_received_at": time.time(),
|
||
})
|
||
|
||
if len(store) > _MAX_EVENTS_PER_DEVICE:
|
||
_guard_event_store[device_id] = store[-_MAX_EVENTS_PER_DEVICE:]
|
||
except Exception as e:
|
||
logger.debug(f"守护事件存储异常: {e}")
|
||
|
||
async def send_to_device(self, device_id: str, data: dict) -> bool:
|
||
"""发送消息到设备（失败重试，避免单次抖动即踢下线）"""
|
||
ws = self.connections.get(device_id)
|
||
if not ws:
|
||
return False
|
||
|
||
last_err = None
|
||
for attempt in range(3):
|
||
try:
|
||
await ws.send_json(data)
|
||
return True
|
||
except Exception as e:
|
||
last_err = e
|
||
logger.warning(f"发送失败 [{device_id}] attempt={attempt + 1}/3: {e}")
|
||
if attempt < 2:
|
||
await asyncio.sleep(0.4 * (attempt + 1))
|
||
|
||
logger.error(f"发送失败 [{device_id}] 已重试 3 次: {last_err}")
|
||
await self.disconnect(device_id)
|
||
return False
|
||
|
||
def get_heartbeat_status(self, device_id: str) -> Optional[dict]:
|
||
info = self.device_info.get(device_id)
|
||
if not info:
|
||
return None
|
||
last = info.get("last_heartbeat")
|
||
if not last:
|
||
return None
|
||
try:
|
||
dt = datetime.fromisoformat(last)
|
||
age = max(0, int((datetime.now() - dt).total_seconds()))
|
||
except Exception:
|
||
age = -1
|
||
interval = int(info.get("heartbeat_interval_seconds") or self.heartbeat_config.get(device_id, settings.WS_HEARTBEAT_INTERVAL))
|
||
return {
|
||
"device_id": device_id,
|
||
"online": self.is_online(device_id),
|
||
"last_heartbeat": last,
|
||
"heartbeat_interval_seconds": interval,
|
||
"heartbeat_age_seconds": age,
|
||
"stale": age >= (interval * 3) if age >= 0 else True,
|
||
}
|
||
|
||
async def set_heartbeat_interval(self, device_id: str, interval_seconds: int) -> bool:
|
||
if interval_seconds < 5 or interval_seconds > 120:
|
||
return False
|
||
self.heartbeat_config[device_id] = interval_seconds
|
||
if device_id in self.device_info:
|
||
self.device_info[device_id]["heartbeat_interval_seconds"] = interval_seconds
|
||
if self.is_online(device_id):
|
||
await self.send_to_device(device_id, {
|
||
"type": "config",
|
||
"heartbeat_interval_seconds": interval_seconds,
|
||
})
|
||
return True
|
||
|
||
async def sweep_stale_devices(self, timeout_seconds: int):
|
||
# 主动清理超时连接，避免设备长时间假在线
|
||
to_drop = []
|
||
now = datetime.now()
|
||
for device_id, info in list(self.device_info.items()):
|
||
if self.pending_device_commands.get(device_id, 0) > 0:
|
||
logger.debug(f"跳过超时下线（设备有在途命令）: {device_id}")
|
||
continue
|
||
last = info.get("last_heartbeat")
|
||
if not last:
|
||
continue
|
||
try:
|
||
age = (now - datetime.fromisoformat(last)).total_seconds()
|
||
except Exception:
|
||
age = timeout_seconds + 1
|
||
grace = int(info.get("command_grace_seconds") or 0)
|
||
effective_timeout = timeout_seconds + max(0, grace)
|
||
if age > effective_timeout:
|
||
to_drop.append(device_id)
|
||
for device_id in to_drop:
|
||
logger.warning(f"⏱️ 心跳超时下线: {device_id}")
|
||
await self.disconnect(device_id)
|
||
|
||
async def send_command(
|
||
self,
|
||
device_id: str,
|
||
command: dict,
|
||
timeout: int = 30
|
||
) -> dict:
|
||
"""发送命令并等待响应"""
|
||
|
||
if not self.is_online(device_id):
|
||
return {"code": 503, "message": "设备不在线"}
|
||
|
||
# 生成命令ID
|
||
command_id = str(uuid.uuid4())
|
||
command["command_id"] = command_id
|
||
command["timestamp"] = int(datetime.now().timestamp())
|
||
|
||
if device_id in self.device_info:
|
||
self.device_info[device_id]["last_command_at"] = datetime.now().isoformat()
|
||
# 长任务（微信登录等）期间延长心跳宽限期，避免误踢
|
||
timeout_hint = int(command.get("timeout") or command.get("data", {}).get("timeout") or 0)
|
||
grace = max(120, min(timeout_hint, 600))
|
||
self.device_info[device_id]["command_grace_seconds"] = grace
|
||
|
||
# Android App 协议兼容：execute 需要扁平 action/params
|
||
try:
|
||
info = self.device_info.get(device_id) or {}
|
||
platform = (info.get("platform") or "").lower()
|
||
if platform == "android" and command.get("type") == "execute":
|
||
payload = command.get("data") or {}
|
||
if isinstance(payload, dict):
|
||
action = payload.get("action")
|
||
params = payload.get("params") or {}
|
||
if action:
|
||
script = payload.get("script")
|
||
hook_only = payload.get("hook_only")
|
||
command.pop("data", None)
|
||
command["action"] = action
|
||
command["params"] = params
|
||
if script is not None:
|
||
command["script"] = script
|
||
if hook_only is not None:
|
||
command["hook_only"] = hook_only
|
||
except Exception:
|
||
pass
|
||
|
||
# 创建Future等待响应
|
||
future = asyncio.get_running_loop().create_future()
|
||
self.pending_commands[command_id] = future
|
||
self.pending_device_commands[device_id] = self.pending_device_commands.get(device_id, 0) + 1
|
||
|
||
try:
|
||
# 发送命令（保留 pending_commands 直至收到 response 或超时）
|
||
success = await self.send_to_device(device_id, command)
|
||
if not success:
|
||
self.pending_commands.pop(command_id, None)
|
||
cnt = self.pending_device_commands.get(device_id, 0) - 1
|
||
if cnt <= 0:
|
||
self.pending_device_commands.pop(device_id, None)
|
||
else:
|
||
self.pending_device_commands[device_id] = cnt
|
||
if not future.done():
|
||
future.cancel()
|
||
return {"code": 503, "message": "发送失败"}
|
||
|
||
result = await asyncio.wait_for(future, timeout=timeout)
|
||
return result
|
||
|
||
except asyncio.TimeoutError:
|
||
self.pending_commands.pop(command_id, None)
|
||
cnt = self.pending_device_commands.get(device_id, 0) - 1
|
||
if cnt <= 0:
|
||
self.pending_device_commands.pop(device_id, None)
|
||
else:
|
||
self.pending_device_commands[device_id] = cnt
|
||
if device_id in self.device_info:
|
||
self.device_info[device_id].pop("command_grace_seconds", None)
|
||
logger.warning(f"[ws_hub] 设备响应超时 device_id={device_id} command_id={command_id} timeout={timeout}s")
|
||
return {"code": 408, "message": "设备响应超时"}
|
||
|
||
except Exception as e:
|
||
self.pending_commands.pop(command_id, None)
|
||
cnt = self.pending_device_commands.get(device_id, 0) - 1
|
||
if cnt <= 0:
|
||
self.pending_device_commands.pop(device_id, None)
|
||
else:
|
||
self.pending_device_commands[device_id] = cnt
|
||
return {"code": 500, "message": str(e)}
|
||
|
||
def get_online_devices(self) -> list:
|
||
"""获取所有在线设备"""
|
||
return list(self.device_info.values())
|
||
|
||
def get_device_info(self, device_id: str) -> Optional[dict]:
|
||
"""获取设备信息"""
|
||
return self.device_info.get(device_id)
|
||
|
||
def get_project_devices(self, project_id: str) -> list:
|
||
"""获取项目下所有在线设备"""
|
||
device_ids = self.project_devices.get(project_id, [])
|
||
return [self.device_info[did] for did in device_ids if did in self.device_info]
|
||
|
||
def get_all_projects(self) -> dict:
|
||
"""获取所有项目及其设备数量"""
|
||
result = {}
|
||
for project_id, device_ids in self.project_devices.items():
|
||
online_count = sum(1 for did in device_ids if did in self.device_info)
|
||
result[project_id] = {
|
||
"project_id": project_id,
|
||
"total_devices": len(device_ids),
|
||
"online_devices": online_count
|
||
}
|
||
return result
|
||
|
||
async def broadcast_to_project(self, project_id: str, command: dict) -> dict:
|
||
"""向项目下所有设备广播命令"""
|
||
device_ids = self.project_devices.get(project_id, [])
|
||
results = {"success": 0, "failed": 0, "devices": []}
|
||
|
||
for device_id in device_ids:
|
||
if device_id in self.connections:
|
||
success = await self.send_to_device(device_id, command)
|
||
if success:
|
||
results["success"] += 1
|
||
else:
|
||
results["failed"] += 1
|
||
results["devices"].append({
|
||
"device_id": device_id,
|
||
"success": success
|
||
})
|
||
|
||
return results
|
||
|
||
async def execute_on_project(
|
||
self,
|
||
project_id: str,
|
||
command: dict,
|
||
timeout: int = 30
|
||
) -> list:
|
||
"""在项目下所有设备执行命令并等待结果"""
|
||
device_ids = self.project_devices.get(project_id, [])
|
||
tasks = []
|
||
|
||
for device_id in device_ids:
|
||
if device_id in self.connections:
|
||
task = self.send_command(device_id, command.copy(), timeout)
|
||
tasks.append((device_id, task))
|
||
|
||
results = []
|
||
for device_id, task in tasks:
|
||
try:
|
||
result = await task
|
||
results.append({
|
||
"device_id": device_id,
|
||
"result": result
|
||
})
|
||
except Exception as e:
|
||
results.append({
|
||
"device_id": device_id,
|
||
"error": str(e)
|
||
})
|
||
|
||
return results
|
||
|
||
|
||
# 全局实例
|
||
ws_hub = WebSocketHub()
|