from __future__ import annotations import asyncio from datetime import UTC, datetime from cloud.internal_api.models import HeartbeatResponse from cloud.internal_api.models import HostGovernancePolicyModel from device.manager import DeviceManager from host_agent.config import HostAgentConfig from host_agent.heartbeat import HeartbeatSynchronizer, build_device_snapshot from host_agent.mcp_lock import McpBusyTracker from host_agent.policy_cache import HostPolicyCacheStore from host_agent.status import AgentStatusTracker class ConnectableDriver: def connect(self) -> None: return None def screenshot(self) -> bytes: return b"ok" def _config() -> HostAgentConfig: return HostAgentConfig( control_plane_url="https://control.example", host_id="host-a", token="secret", heartbeat_interval_seconds=0.01, ) def test_build_device_snapshot_copies_complete_non_secret_state() -> None: manager = DeviceManager() manager.register_device( "device-b", lambda: object(), # type: ignore[arg-type,return-value] driver_type="appium", connection_info={"password": "must-not-leave-host"}, status="busy", capability_tags=["android", "physical"], ) manager.register_device( "device-a", lambda: object(), # type: ignore[arg-type,return-value] status="offline", ) snapshot = build_device_snapshot(manager) assert [device.device_id for device in snapshot] == ["device-a", "device-b"] assert snapshot[1].driver_type == "appium" assert snapshot[1].status == "idle" assert snapshot[1].capability_tags == ["android", "physical"] assert "must-not-leave-host" not in repr(snapshot) def test_heartbeat_synchronizer_runs_at_configured_interval_until_stopped() -> None: manager = DeviceManager() manager.register_device( "device-a", lambda: ConnectableDriver(), # type: ignore[arg-type,return-value] ) calls: list[list[str]] = [] class FakeClient: async def heartbeat( self, devices, *, address=None, policy_revision=0, **kwargs ): calls.append([device.device_id for device in devices]) return HeartbeatResponse( host_id="host-a", accepted_devices=len(devices), received_at=datetime.now(UTC), ) async def scenario() -> None: stop = asyncio.Event() synchronizer = HeartbeatSynchronizer( manager, FakeClient(), # type: ignore[arg-type] _config(), ) task = asyncio.create_task(synchronizer.run(stop)) while len(calls) < 3: await asyncio.sleep(0.005) stop.set() await task asyncio.run(scenario()) assert calls == [["device-a"], ["device-a"], ["device-a"]] assert manager.status("device-a") == "busy" def test_offline_device_is_retried_on_next_heartbeat_cycle() -> None: class FlakyDriver(ConnectableDriver): def __init__(self, fail: bool) -> None: self.fail = fail def screenshot(self) -> bytes: if self.fail: raise RuntimeError("WDA disconnected") return b"ok" instances: list[FlakyDriver] = [] def factory() -> FlakyDriver: driver = FlakyDriver(not instances) instances.append(driver) return driver manager = DeviceManager() manager.register_device("device-a", factory) # type: ignore[arg-type] manager.connect("device-a") sync = HeartbeatSynchronizer(manager, object(), _config()) # type: ignore[arg-type] sync.probe_connected_devices() assert manager.status("device-a") == "offline" sync.connect_devices() assert manager.status("device-a") == "busy" def test_sync_once_notifies_status_tracker_and_on_sync_with_device_count() -> None: manager = DeviceManager() manager.register_device( "device-a", lambda: ConnectableDriver(), # type: ignore[arg-type,return-value] ) manager.register_device( "device-b", lambda: ConnectableDriver(), # type: ignore[arg-type,return-value] ) class FakeClient: async def heartbeat( self, devices, *, address=None, policy_revision=0, **kwargs ): return HeartbeatResponse( host_id="host-a", accepted_devices=len(devices), received_at=datetime.now(UTC), ) async def scenario() -> None: tracker = AgentStatusTracker() on_sync_calls: list[int] = [] synchronizer = HeartbeatSynchronizer( manager, FakeClient(), # type: ignore[arg-type] _config(), status_tracker=tracker, on_sync=on_sync_calls.append, ) await synchronizer.sync_once() assert on_sync_calls == [2] last_heartbeat = tracker.snapshot()["last_heartbeat"] assert last_heartbeat is not None assert last_heartbeat["ok"] is True assert last_heartbeat["device_count"] == 2 asyncio.run(scenario()) def test_heartbeat_caches_safe_host_policy_and_reuses_its_revision(tmp_path) -> None: manager = DeviceManager() cache = HostPolicyCacheStore(tmp_path / "host_policy.json") revisions: list[int] = [] class UpdatingClient: async def heartbeat( self, devices, *, address=None, policy_revision=0, **kwargs ): revisions.append(policy_revision) return HeartbeatResponse( host_id="host-a", accepted_devices=len(devices), received_at=datetime.now(UTC), policy_revision=4, policy=HostGovernancePolicyModel( revision=4, self_submission_enabled=False, max_active_tasks=2, daily_token_budget=900, ), ) async def scenario() -> None: tracker = AgentStatusTracker() synchronizer = HeartbeatSynchronizer( manager, UpdatingClient(), # type: ignore[arg-type] _config(), policy_cache=cache, status_tracker=tracker, ) await synchronizer.sync_once() assert tracker.snapshot()["host_policy"] == { "revision": 4, "self_submission_enabled": False, "max_active_tasks": 2, "daily_token_budget": 900, } restarted = HeartbeatSynchronizer( manager, UpdatingClient(), # type: ignore[arg-type] _config(), policy_cache=cache, ) assert restarted.policy_revision == 4 asyncio.run(scenario()) assert revisions == [0] assert '"token":' not in (tmp_path / "host_policy.json").read_text(encoding="utf-8") def test_sync_once_passes_mcp_busy_device_ids_to_client() -> None: """When mcp_busy_tracker has a lease, sync_once relays the device_ids.""" manager = DeviceManager() tracker = McpBusyTracker() assert tracker.acquire("phone-1", "sess-a") last_kwargs: dict[str, object] = {} class FakeClient: async def heartbeat( self, devices, *, address=None, policy_revision=0, **kwargs ): last_kwargs.update(kwargs) return HeartbeatResponse( host_id="host-a", accepted_devices=len(devices), received_at=datetime.now(UTC), ) async def scenario() -> None: sync = HeartbeatSynchronizer( manager, FakeClient(), # type: ignore[arg-type] _config(), mcp_busy_tracker=tracker, ) await sync.sync_once() asyncio.run(scenario()) assert last_kwargs.get("mcp_busy_device_ids") == ["phone-1"] def test_sync_once_passes_empty_when_tracker_is_none() -> None: """Default: no tracker → no busy device ids forwarded.""" manager = DeviceManager() last_kwargs: dict[str, object] = {} class FakeClient: async def heartbeat( self, devices, *, address=None, policy_revision=0, **kwargs ): last_kwargs.update(kwargs) return HeartbeatResponse( host_id="host-a", accepted_devices=len(devices), received_at=datetime.now(UTC), ) async def scenario() -> None: sync = HeartbeatSynchronizer( manager, FakeClient(), # type: ignore[arg-type] _config(), ) await sync.sync_once() asyncio.run(scenario()) assert not last_kwargs.get("mcp_busy_device_ids")