feat(host-agent): synchronize device snapshots

2026-07-12 18:48:17 +08:00
parent f6bc8f6b98
commit 6e59f2f3ca
3 changed files with 136 additions and 1 deletions
@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING
+
+from cloud.internal_api.models import DeviceSnapshotModel, HeartbeatResponse
+from device.manager import DeviceManager
+from host_agent.client import HostAgentClient
+from host_agent.config import HostAgentConfig
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+
+def build_device_snapshot(manager: DeviceManager) -> list[DeviceSnapshotModel]:
+    return [
+        DeviceSnapshotModel(
+            device_id=device.id,
+            driver_type=device.driver_type,
+            status=device.status,
+            capability_tags=list(device.capability_tags),
+        )
+        for device in sorted(manager.list_devices(), key=lambda item: item.id)
+    ]
+
+
+class HeartbeatSynchronizer:
+    def __init__(
+        self,
+        manager: DeviceManager,
+        client: HostAgentClient,
+        config: HostAgentConfig,
+        *,
+        address: str | None = None,
+        sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
+    ) -> None:
+        self.manager = manager
+        self.client = client
+        self.config = config
+        self.address = address
+        self._sleep = sleep
+
+    async def sync_once(self) -> HeartbeatResponse:
+        return await self.client.heartbeat(
+            build_device_snapshot(self.manager),
+            address=self.address,
+        )
+
+    async def run(self, stop: asyncio.Event) -> None:
+        while not stop.is_set():
+            await self.sync_once()
+            try:
+                await asyncio.wait_for(
+                    stop.wait(),
+                    timeout=self.config.heartbeat_interval_seconds,
+                )
+            except TimeoutError:
+                continue