diff --git a/apps/device-host-agent/host_agent/client.py b/apps/device-host-agent/host_agent/client.py index 76bf447..42a9f10 100644 --- a/apps/device-host-agent/host_agent/client.py +++ b/apps/device-host-agent/host_agent/client.py @@ -168,17 +168,21 @@ class HostAgentClient: *, address: str | None = None, policy_revision: int = 0, + mcp_busy_device_ids: list[str] | None = None, ) -> HeartbeatResponse: + payload: dict[str, Any] = { + "host_id": self.config.host_id, + "address": address, + "devices": [device.model_dump(mode="json") for device in devices], + "policy_revision": policy_revision, + "planner_transport": self.config.ai_planner_transport, + } + if mcp_busy_device_ids: + payload["mcp_busy_device_ids"] = list(mcp_busy_device_ids) response = await self._request( "PUT", f"/internal/v1/hosts/{self.config.host_id}/heartbeat", - json={ - "host_id": self.config.host_id, - "address": address, - "devices": [device.model_dump(mode="json") for device in devices], - "policy_revision": policy_revision, - "planner_transport": self.config.ai_planner_transport, - }, + json=payload, ) return HeartbeatResponse.model_validate(response.json()) diff --git a/apps/device-host-agent/host_agent/heartbeat.py b/apps/device-host-agent/host_agent/heartbeat.py index ecdf2eb..152229b 100644 --- a/apps/device-host-agent/host_agent/heartbeat.py +++ b/apps/device-host-agent/host_agent/heartbeat.py @@ -14,6 +14,8 @@ 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 [ @@ -40,6 +42,7 @@ class HeartbeatSynchronizer: 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 @@ -50,6 +53,7 @@ class HeartbeatSynchronizer: 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: @@ -57,10 +61,16 @@ class HeartbeatSynchronizer: 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: diff --git a/apps/device-host-agent/tests/test_client.py b/apps/device-host-agent/tests/test_client.py index e8c70bd..9ec0be7 100644 --- a/apps/device-host-agent/tests/test_client.py +++ b/apps/device-host-agent/tests/test_client.py @@ -353,6 +353,60 @@ def test_submit_self_task_does_not_duplicate_when_response_is_lost() -> None: assert attempts == 1 +def test_heartbeat_includes_mcp_busy_device_ids_in_payload() -> None: + """When mcp_busy_device_ids is passed, the client sends it in the request.""" + captured: list[dict[str, object]] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(json.loads(request.content)) + return httpx.Response( + 200, + json={ + "host_id": "host-a", + "accepted_devices": 0, + "received_at": "2026-07-12T00:00:00Z", + }, + ) + + async def scenario() -> None: + async with httpx.AsyncClient( + transport=httpx.MockTransport(handler), + base_url="https://control.example", + ) as http_client: + client = HostAgentClient(_config(), http_client=http_client) + await client.heartbeat([], mcp_busy_device_ids=["phone-1"]) + + asyncio.run(scenario()) + assert captured[0]["mcp_busy_device_ids"] == ["phone-1"] + + +def test_heartbeat_omits_mcp_busy_device_ids_when_empty() -> None: + """Empty list is omitted from the payload (backward compatible).""" + captured: list[dict[str, object]] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(json.loads(request.content)) + return httpx.Response( + 200, + json={ + "host_id": "host-a", + "accepted_devices": 0, + "received_at": "2026-07-12T00:00:00Z", + }, + ) + + async def scenario() -> None: + async with httpx.AsyncClient( + transport=httpx.MockTransport(handler), + base_url="https://control.example", + ) as http_client: + client = HostAgentClient(_config(), http_client=http_client) + await client.heartbeat([], mcp_busy_device_ids=[]) + + asyncio.run(scenario()) + assert "mcp_busy_device_ids" not in captured[0] + + def test_bootstrap_client_directly_enrolls_and_enrolls_device() -> None: requests: list[httpx.Request] = [] host_attempts = 0 diff --git a/apps/device-host-agent/tests/test_heartbeat.py b/apps/device-host-agent/tests/test_heartbeat.py index 72e75e3..d865de7 100644 --- a/apps/device-host-agent/tests/test_heartbeat.py +++ b/apps/device-host-agent/tests/test_heartbeat.py @@ -8,6 +8,7 @@ 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 @@ -60,7 +61,7 @@ def test_heartbeat_synchronizer_runs_at_configured_interval_until_stopped() -> N calls: list[list[str]] = [] class FakeClient: - async def heartbeat(self, devices, *, address=None, policy_revision=0): + 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", @@ -98,7 +99,7 @@ def test_sync_once_notifies_status_tracker_and_on_sync_with_device_count() -> No ) class FakeClient: - async def heartbeat(self, devices, *, address=None, policy_revision=0): + async def heartbeat(self, devices, *, address=None, policy_revision=0, **kwargs): return HeartbeatResponse( host_id="host-a", accepted_devices=len(devices), @@ -133,7 +134,7 @@ def test_heartbeat_caches_safe_host_policy_and_reuses_its_revision(tmp_path) -> revisions: list[int] = [] class UpdatingClient: - async def heartbeat(self, devices, *, address=None, policy_revision=0): + async def heartbeat(self, devices, *, address=None, policy_revision=0, **kwargs): revisions.append(policy_revision) return HeartbeatResponse( host_id="host-a", @@ -178,3 +179,58 @@ def test_heartbeat_caches_safe_host_policy_and_reuses_its_revision(tmp_path) -> 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")