117 lines
4.2 KiB
Python
117 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import TYPE_CHECKING
|
|
|
|
from cloud.internal_api.models import DeviceSnapshotModel, HeartbeatResponse
|
|
from core.errors import DeviceRuntimeError
|
|
from device.manager import DeviceManager
|
|
from host_agent.client import HostAgentClient
|
|
from host_agent.config import HostAgentConfig
|
|
from host_agent.policy_cache import HostPolicyCacheStore
|
|
from host_agent.status import AgentStatusTracker
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Awaitable, Callable
|
|
|
|
from host_agent.mcp_lock import McpBusyTracker
|
|
|
|
|
|
def build_device_snapshot(manager: DeviceManager) -> list[DeviceSnapshotModel]:
|
|
return [
|
|
DeviceSnapshotModel(
|
|
device_id=device.id,
|
|
driver_type=device.driver_type,
|
|
status="idle" if device.status == "busy" else 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,
|
|
status_tracker: AgentStatusTracker | None = None,
|
|
on_sync: Callable[[int], None] | None = None,
|
|
policy_cache: HostPolicyCacheStore | None = None,
|
|
on_policy_sync: Callable[[int], None] | None = None,
|
|
mcp_busy_tracker: McpBusyTracker | None = None,
|
|
) -> None:
|
|
self.manager = manager
|
|
self.client = client
|
|
self.config = config
|
|
self.address = address
|
|
self._sleep = sleep
|
|
self.status_tracker = status_tracker
|
|
self.on_sync = on_sync
|
|
self.policy_cache = policy_cache
|
|
self.on_policy_sync = on_policy_sync
|
|
self.mcp_busy_tracker = mcp_busy_tracker
|
|
self.policy = policy_cache.load() if policy_cache is not None else None
|
|
self.policy_revision = self.policy.revision if self.policy is not None else 0
|
|
if self.status_tracker is not None:
|
|
self.status_tracker.mark_host_policy(self.policy)
|
|
|
|
async def sync_once(self) -> HeartbeatResponse:
|
|
snapshot = build_device_snapshot(self.manager)
|
|
mcp_busy_ids = (
|
|
self.mcp_busy_tracker.busy_device_ids()
|
|
if self.mcp_busy_tracker is not None
|
|
else []
|
|
)
|
|
response = await self.client.heartbeat(
|
|
snapshot,
|
|
address=self.address,
|
|
policy_revision=self.policy_revision,
|
|
mcp_busy_device_ids=mcp_busy_ids,
|
|
)
|
|
self.policy_revision = response.policy_revision
|
|
if response.policy is not None:
|
|
self.policy = response.policy
|
|
if self.policy_cache is not None:
|
|
self.policy_cache.save(response.policy)
|
|
if self.on_policy_sync is not None:
|
|
self.on_policy_sync(response.policy.revision)
|
|
elif response.policy_revision == 0:
|
|
self.policy = None
|
|
if self.policy_cache is not None:
|
|
self.policy_cache.clear()
|
|
if self.status_tracker is not None:
|
|
self.status_tracker.mark_host_policy(self.policy)
|
|
if self.status_tracker is not None:
|
|
self.status_tracker.mark_heartbeat(ok=True, device_count=len(snapshot))
|
|
if self.on_sync is not None:
|
|
try:
|
|
self.on_sync(len(snapshot))
|
|
except Exception:
|
|
pass
|
|
return response
|
|
|
|
async def run(self, stop: asyncio.Event) -> None:
|
|
self.connect_devices()
|
|
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
|
|
|
|
def connect_devices(self) -> None:
|
|
for device in self.manager.list_devices():
|
|
if device.status != "idle":
|
|
continue
|
|
try:
|
|
self.manager.connect(device.id)
|
|
except DeviceRuntimeError:
|
|
continue
|