hospital-front-agent/agent/ws_client.py

# -*- coding: utf-8 -*-
"""Agent WebSocket 客户端。

- 主动出站连接 B 端（穿透白名单，无需 A 端开放入站端口）
- 断线指数退避重连（1s → 2s → 4s → ... 最大 60s）
- 周期心跳：拉取主程序 /api/admin/health 并转发给 B 端
- 收到下行命令交给 CommandDispatcher 处理

上行消息格式：{"type":"HEARTBEAT","deviceId":"...","ts":...,"payload":{...}}
"""

import json
import logging
import threading
import time

import websocket  # websocket-client

log = logging.getLogger(__name__)


class AgentClient:
    """WebSocket 长连接客户端（连接管理 + 心跳 + 重连）。"""

    def __init__(self, config, dispatcher):
        self.config = config
        self.dispatcher = dispatcher

        self._ws = None                      # 当前 WebSocket 连接
        self._connected = False              # 是否已连接
        self._reconnect_attempts = 0         # 连续重连次数（成功后清零）
        self._lock = threading.Lock()        # 保护 _ws/_connected 的并发访问
        self._heartbeat_stop = threading.Event()
        self._threads: list[threading.Thread] = []

    # ================= 生命周期 =================

    def start(self):
        """启动：连接 + 心跳线程（首次心跳立即发送，避免网关空闲超时在首个周期内掐断连接）。"""
        self._connect()
        t = threading.Thread(target=self._heartbeat_loop, name="heartbeat", daemon=True)
        t.start()
        self._threads.append(t)

    def stop(self):
        """停止并释放资源。"""
        self._heartbeat_stop.set()
        with self._lock:
            if self._ws is not None:
                try:
                    self._ws.close()
                except Exception:
                    pass
                self._ws = None
            self._connected = False

    # ================= 连接管理 =================

    def _connect(self):
        """建立 WebSocket 连接（阻塞直到连接断开，内部自行调度重连）。"""
        url = self.config.server_url
        log.info("连接管理服务器：%s", url)

        headers = [
            f"X-Device-Id: {self.config.device_id}",
            f"X-Device-Token: {self.config.device_token}",
        ]

        def on_open(ws):
            self._connected = True
            self._reconnect_attempts = 0
            log.info("已连接管理服务器")
            self.send("REGISTER", {
                "deviceId": self.config.device_id,
                "version": "1.0.0",
            })

        def on_message(ws, text):
            try:
                node = json.loads(text)
                msg_type = node.get("type")
                payload = node.get("payload")
                self.dispatcher.dispatch(msg_type, payload)
            except Exception as e:
                log.warning("命令处理异常：%s", e)

        def on_close(ws, code, reason):
            self._connected = False
            log.info("连接关闭：code=%s, reason=%s", code, reason)
            self._schedule_reconnect()

        def on_error(ws, error):
            self._connected = False
            log.warning("连接失败：%s，将重连", error)
            # websocket-client 出错后必定回调 on_close，重连由 on_close 统一调度
            # 此处不重复调度，避免双重重连

        ws = websocket.WebSocketApp(
            url,
            header=headers,
            on_open=on_open,
            on_message=on_message,
            on_close=on_close,
            on_error=on_error,
        )
        with self._lock:
            self._ws = ws
        # 阻塞运行（跳过协议层 ping，应用层心跳已足够；关掉默认 ping 避免双重心跳）
        ws.run_forever(ping_interval=25, ping_timeout=10)

    def _schedule_reconnect(self):
        """指数退避重连：1s 起，翻倍至 60s 封顶。"""
        self._reconnect_attempts += 1
        attempt = self._reconnect_attempts
        delay = min(1000 * (2 ** min(attempt - 1, 6)), 60_000) / 1000.0
        log.info("%.0fs 后进行第 %d 次重连", delay, attempt)
        time.sleep(delay)
        if not self._heartbeat_stop.is_set():
            self._connect()

    # ================= 心跳 =================

    def _heartbeat_loop(self):
        """周期心跳：拉主程序健康数据转发。"""
        while not self._heartbeat_stop.wait(self.config.heartbeat_sec):
            try:
                self._heartbeat_once()
            except Exception as e:
                log.warning("心跳发送异常：%s", e)

    def _heartbeat_once(self):
        main_health = self.dispatcher.main_app_client.health()
        self.send("HEARTBEAT", {
            "deviceId": self.config.device_id,
            "mainApp": main_health,
        })

    # ================= 上行消息 =================

    def send(self, msg_type: str, payload):
        """发送上行消息（未连接时丢弃并记日志）。"""
        with self._lock:
            ws = self._ws
            connected = self._connected
        if ws is None or not connected:
            log.debug("未连接，丢弃上行消息：%s", msg_type)
            return
        envelope = {
            "type": msg_type,
            "deviceId": self.config.device_id,
            "ts": int(time.time()),
            "payload": payload,
        }
        try:
            ws.send(json.dumps(envelope, ensure_ascii=False))
        except Exception as e:
            log.warning("上行消息发送失败（%s）：%s", msg_type, e)

    def is_connected(self) -> bool:
        """当前是否已连接（供健康状态展示）。"""
        return self._connected