roborock-cli 0.1.1__py3-none-any.whl

roborock_cli/connection.py

@@ -0,0 +1,202 @@
+"""
+连接管理模块
+
+负责：
+- 从 JSON 缓存加载 UserData
+- 创建 DeviceManager
+- 设备发现和模糊匹配
+- 连接生命周期管理
+"""
+
+import json
+import os
+import ssl
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import Any
+
+import aiohttp
+import certifi
+from roborock_cli._vendor.roborock.data import UserData
+from roborock_cli._vendor.roborock.devices.device_manager import create_device_manager, UserParams
+from roborock_cli._vendor.roborock.devices.file_cache import FileCache
+
+
+class AuthExpiredError(Exception):
+    """认证已过期，需要重新运行 setup_auth"""
+
+    def __init__(self):
+        super().__init__("认证已过期，请重新运行 `roborock-cli auth`")
+
+
+class NoDeviceError(Exception):
+    """未发现设备"""
+
+    def __init__(self):
+        super().__init__("未发现任何设备，请确认设备已绑定到账户")
+
+
+class AmbiguousDeviceError(Exception):
+    """设备名称匹配到多个设备"""
+
+    def __init__(self, name: str, matches: list[str]):
+        self.name = name
+        self.matches = matches
+        super().__init__(
+            f"设备名称 '{name}' 匹配到多个设备: {matches}，请使用更精确的名称"
+        )
+
+
+class DeviceNotFoundError(Exception):
+    """设备名称未匹配到任何设备"""
+
+    def __init__(self, name: str, available: list[str]):
+        self.name = name
+        self.available = available
+        super().__init__(
+            f"未找到设备 '{name}'，可用设备: {available}"
+        )
+
+
+class ConnectionManager:
+    """连接管理器"""
+
+    def __init__(self, auth_cache_path: Path, device_cache_path: Path):
+        self._auth_cache_path = auth_cache_path
+        self._device_cache_path = device_cache_path
+        self._manager: Any = None  # DeviceManager
+        self._cache: FileCache | None = None
+        self._user_params: UserParams | None = None
+        self._session: aiohttp.ClientSession | None = None
+
+    @classmethod
+    def from_env(cls) -> "ConnectionManager":
+        """从环境变量创建连接管理器"""
+        auth_cache = Path(
+            os.environ.get(
+                "ROBOROCK_AUTH_CACHE",
+                Path.home() / ".cache" / "roborock-cli" / "user_data.json"
+            )
+        )
+        device_cache = Path(
+            os.environ.get(
+                "ROBOROCK_DEVICE_CACHE",
+                Path.home() / ".cache" / "roborock-cli" / "device_cache.pkl"
+            )
+        )
+        return cls(auth_cache, device_cache)
+
+    async def ensure_connected(self) -> None:
+        """确保已连接，如果未连接则建立连接"""
+        if self._manager is not None:
+            return
+
+        # 1. 从 JSON 文件加载 UserData
+        if not self._auth_cache_path.exists():
+            raise AuthExpiredError()
+
+        try:
+            with open(self._auth_cache_path) as f:
+                data = json.load(f)
+            email = data["email"]
+            user_data = UserData.from_dict(data["user_data"])
+        except (json.JSONDecodeError, KeyError) as e:
+            raise AuthExpiredError() from e
+
+        # 2. 构建 UserParams（不传 base_url，让 SDK 自动发现）
+        self._user_params = UserParams(
+            username=email,
+            user_data=user_data,
+        )
+
+        # 3. 创建缓存目录
+        self._device_cache_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # 4. 创建 FileCache
+        self._cache = FileCache(self._device_cache_path)
+
+        # 5. 创建带 SSL 证书的 session（解决 macOS 证书问题）
+        ssl_ctx = ssl.create_default_context(cafile=certifi.where())
+        connector = aiohttp.TCPConnector(ssl=ssl_ctx)
+        self._session = aiohttp.ClientSession(connector=connector)
+
+        # 6. 创建 DeviceManager（内部已调用 discover_devices）
+        self._manager = await create_device_manager(
+            self._user_params,
+            cache=self._cache,
+            session=self._session,
+        )
+
+        # 7. 检查设备列表
+        devices = await self._manager.get_devices()
+        if not devices:
+            raise NoDeviceError()
+
+    async def list_devices(self) -> list[Any]:
+        """返回所有已发现设备"""
+        if self._manager is None:
+            await self.ensure_connected()
+        return await self._manager.get_devices()
+
+    async def get_device(self, device_name: str | None = None) -> Any:
+        """
+        获取指定设备
+
+        Args:
+            device_name: 设备名称（模糊匹配），为 None 时自动选择单设备
+
+        Returns:
+            RoborockDevice
+
+        Raises:
+            NoDeviceError: 无设备
+            AmbiguousDeviceError: 名称匹配到多个设备
+            DeviceNotFoundError: 名称未匹配到任何设备
+        """
+        devices = await self.list_devices()
+
+        # 单设备且未指定名称时，自动选择
+        if device_name is None:
+            if len(devices) == 1:
+                return devices[0]
+            else:
+                available = [d.name for d in devices]
+                raise ValueError(
+                    f"发现 {len(devices)} 个设备，请指定 device_name。可用设备: {available}"
+                )
+
+        # 模糊匹配
+        matches = []
+        for device in devices:
+            if device_name.lower() in device.name.lower():
+                matches.append(device)
+
+        if len(matches) == 0:
+            available = [d.name for d in devices]
+            raise DeviceNotFoundError(device_name, available)
+
+        if len(matches) > 1:
+            raise AmbiguousDeviceError(
+                device_name,
+                [d.name for d in matches]
+            )
+
+        return matches[0]
+
+    async def close(self) -> None:
+        """关闭连接，先 flush 缓存再关闭"""
+        if self._cache is not None:
+            await self._cache.flush()
+        if self._manager is not None:
+            await self._manager.close()
+        if self._session is not None:
+            await self._session.close()
+
+
+@asynccontextmanager
+async def lifespan(server: Any):
+    """MCP server lifespan 管理"""
+    conn = ConnectionManager.from_env()
+    await conn.ensure_connected()
+    yield {"conn": conn}
+    await conn.close()

roborock_cli/helpers.py

@@ -0,0 +1,71 @@
+"""
+共享辅助函数，供 server.py 和 cli.py 使用
+"""
+
+import logging
+from enum import Enum
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+async def safe_refresh(trait, name: str) -> bool:
+    """安全刷新单个 trait，失败不影响其他"""
+    try:
+        await trait.refresh()
+        return True
+    except Exception as e:
+        logger.warning("Trait %s refresh 失败: %s", name, e)
+        return False
+
+
+def serialize_value(value: Any) -> Any:
+    """将 Enum 等特殊类型转为可序列化值"""
+    if value is None:
+        return None
+    if isinstance(value, Enum):
+        return value.name
+    if isinstance(value, list):
+        return [serialize_value(v) for v in value]
+    return value
+
+
+def field(value: Any, description: str) -> dict[str, Any]:
+    """包装字段值和描述"""
+    return {"value": value, "description": description}
+
+
+# StatusV2 白名单：仅输出有意义的字段
+# state/error_code 改在 computed properties 中输出（避免与 state_name/error_code_name 重复）
+STATUS_FIELDS: dict[str, str] = {
+    # 核心状态
+    "battery": "电池电量 (%)",
+    "charge_status": "充电状态 (0=未充电, 1=充电中, 2=充电完成)",
+    # 清洁进度
+    "clean_percent": "清洁任务完成百分比 (%)",
+    "in_cleaning": "清洁状态 (complete/global_clean_not_complete/zone_clean_not_complete/segment_clean_not_complete)",
+    "in_returning": "是否在返回途中 (0=否, 1=是)",
+    "avoid_count": "当前任务避障次数",
+    "repeat": "当前任务重复遍数",
+    # 地图
+    "map_present": "是否有地图 (0=无, 1=有)",
+    # 基座
+    "dock_type": "基座类型 (no_dock/auto_empty_dock/empty_wash_fill_dock/p10_pro_dock 等)",
+    "dock_error_status": "基座错误状态 (ok/duct_blockage/water_empty/waste_water_tank_full 等)",
+    "dust_collection_status": "集尘状态 (0=未集尘, 1=集尘中)",
+    "auto_dust_collection": "自动集尘开关 (0=关闭, 1=开启)",
+    # 洗拖布/烘干
+    "wash_status": "洗拖布状态 (0=空闲, 1=洗涤中)",
+    "wash_phase": "洗拖布阶段 (0=空闲)",
+    "wash_ready": "是否准备好洗拖布 (0=否, 1=是)",
+    "dry_status": "热风烘干状态 (0=关闭, 1=烘干中)",
+    # 水箱
+    "water_box_status": "水箱安装状态 (0=未安装, 1=已安装)",
+    "water_box_carriage_status": "水箱托架状态 (0=正常)",
+    "water_shortage_status": "缺水状态 (0=正常)",
+    # 其他
+    "is_locating": "是否正在定位 (0=否, 1=是)",
+    "lock_status": "童锁状态 (0=未锁定, 1=锁定)",
+    "mop_forbidden_enable": "拖地禁区是否启用 (0=关闭, 1=启用)",
+    "collision_avoid_status": "避障功能 (0=关闭, 1=开启)",
+}