workphone-sdk/sdk/app/services/registry_cluster.py

"""
D6 多服务器注册中心：登记 peer SDK 节点，聚合各节点 /api/v3/devices 发现结果。

环境变量（可选）:
- WORKPHONE_SERVER_ID: 本机节点 ID（默认 hostname）
- WORKPHONE_PUBLIC_BASE_URL: 本机对外 Base，如 http://192.168.1.10:8899
- WORKPHONE_REGISTRY_TOKEN: 登记/心跳校验令牌（不设则开放）
- WORKPHONE_REGISTRY_PEERS: 静态 peer 列表，逗号分隔 Base URL（只读聚合，无需对方登记）
"""

from __future__ import annotations

import logging
import os
import time
from typing import Any, Dict, List, Optional

import httpx

logger = logging.getLogger(__name__)


def _local_server_id() -> str:
    return (os.getenv("WORKPHONE_SERVER_ID") or os.getenv("HOSTNAME") or "sdk-local").strip()


def _local_public_base() -> str:
    return (os.getenv("WORKPHONE_PUBLIC_BASE_URL") or "http://127.0.0.1:8899").rstrip("/")


def _registry_token_expected() -> str:
    return (os.getenv("WORKPHONE_REGISTRY_TOKEN") or "").strip()


def _static_peer_urls() -> List[str]:
    raw = os.getenv("WORKPHONE_REGISTRY_PEERS", "")
    return [u.strip().rstrip("/") for u in raw.split(",") if u.strip()]


def _check_token(token: Optional[str]) -> None:
    exp = _registry_token_expected()
    if exp and (token or "").strip() != exp:
        raise PermissionError("WORKPHONE_REGISTRY_TOKEN 校验失败")


class RegistryClusterService:
    """内存登记 peer；生产可换 Redis/Mongo（当前满足 D6 契约与联调）。"""

    def __init__(self) -> None:
        self._nodes: Dict[str, Dict[str, Any]] = {}

    async def register_peer(
        self,
        node_id: str,
        base_url: str,
        label: str = "",
        token: Optional[str] = None,
        version: str = "",
    ) -> Dict[str, Any]:
        _check_token(token)
        base_url = base_url.rstrip("/")
        self._nodes[node_id] = {
            "node_id": node_id,
            "base_url": base_url,
            "label": label or node_id,
            "version": version or "",
            "last_heartbeat": time.time(),
        }
        logger.info("registry: 登记节点 %s -> %s", node_id, base_url)
        return {"registered": True, "node_id": node_id}

    async def heartbeat_peer(
        self,
        node_id: str,
        token: Optional[str] = None,
        device_count: int = 0,
        version: str = "",
    ) -> Dict[str, Any]:
        _check_token(token)
        if node_id not in self._nodes:
            raise ValueError(f"未知节点: {node_id}，请先 register")
        n = self._nodes[node_id]
        n["last_heartbeat"] = time.time()
        n["device_count"] = device_count
        if version:
            n["version"] = version
        return {"ok": True}

    def list_peers(self, max_age_sec: float = 180.0) -> List[Dict[str, Any]]:
        now = time.time()
        out: List[Dict[str, Any]] = []
        for n in self._nodes.values():
            if now - float(n.get("last_heartbeat", 0)) <= max_age_sec:
                out.append(dict(n))
        return sorted(out, key=lambda x: x["node_id"])

    async def unregister_peer(self, node_id: str, token: Optional[str] = None) -> Dict[str, Any]:
        _check_token(token)
        self._nodes.pop(node_id, None)
        return {"removed": True, "node_id": node_id}

    @staticmethod
    async def fetch_devices_from_base(base_url: str) -> List[Dict[str, Any]]:
        url = f"{base_url.rstrip('/')}/api/v3/devices"
        try:
            async with httpx.AsyncClient(timeout=12.0) as client:
                r = await client.get(url)
                r.raise_for_status()
                j = r.json()
                data = j.get("data")
                if isinstance(data, list):
                    return data
                if isinstance(data, dict) and isinstance(data.get("devices"), list):
                    return data["devices"]
                return []
        except Exception as e:
            logger.warning("registry: 拉取 %s 失败: %s", url, e)
            return []

    async def aggregate_fleet(self) -> Dict[str, Any]:
        from services.device_fleet import list_merged_local_devices

        sid = _local_server_id()
        local_base = _local_public_base()
        local_list = await list_merged_local_devices()

        merged: List[Dict[str, Any]] = []
        for d in local_list:
            merged.append(
                {
                    **d,
                    "home_server_id": sid,
                    "home_base_url": local_base,
                    "_source": "local",
                }
            )

        seen_bases = {local_base}

        for peer in self.list_peers():
            b = peer["base_url"]
            if b in seen_bases:
                continue
            seen_bases.add(b)
            remote = await self.fetch_devices_from_base(b)
            for d in remote:
                merged.append(
                    {
                        **d,
                        "home_server_id": peer["node_id"],
                        "home_base_url": b,
                        "_source": "remote_registry",
                    }
                )

        for b in _static_peer_urls():
            if b in seen_bases:
                continue
            seen_bases.add(b)
            remote = await self.fetch_devices_from_base(b)
            nid = b.replace("http://", "").replace("https://", "").replace("/", "_")
            for d in remote:
                merged.append(
                    {
                        **d,
                        "home_server_id": f"static:{nid}",
                        "home_base_url": b,
                        "_source": "remote_static",
                    }
                )

        return {
            "local_server_id": sid,
            "local_public_base_url": local_base,
            "registered_peers": self.list_peers(),
            "devices": merged,
            "device_count": len(merged),
        }


registry_cluster = RegistryClusterService()