feat(cloud): add edge host enrollment

apps/device-host-agent/tests/test_app.py

@@ -4,10 +4,15 @@ import asyncio
 from contextlib import suppress
 from datetime import UTC, datetime, timedelta
 
-from cloud.internal_api.models import AssignmentModel
+from cloud.internal_api.models import (
+    AssignmentModel,
+    DeviceEnrollmentResponse,
+    HostEnrollmentResponse,
+)
 from device.manager import DeviceManager
 from host_agent.app import HostAgentApplication, create_application
 from host_agent.config import HostAgentConfig
+from host_agent.identity import HostIdentityStore
 from storage.device_config import DeviceConfigStore
 
 
@@ -65,6 +70,106 @@ def test_create_application_loads_persisted_device_configuration(
     asyncio.run(application.client.aclose())
 
 
+def test_create_application_enrolls_host_and_devices_before_managed_startup(
+    tmp_path,
+    monkeypatch,
+) -> None:
+    monkeypatch.chdir(tmp_path)
+    store = DeviceConfigStore(tmp_path / "devices.sqlite3")
+    store.add(
+        device_id="local-device-a",
+        name="Lab iPhone",
+        driver_type="wda",
+        connection_info={"server_url": "http://127.0.0.1:4723"},
+    )
+    events: list[str] = []
+
+    class EnrollmentClient:
+        def __init__(self) -> None:
+            self.config = HostAgentConfig(
+                control_plane_url="https://control.example",
+                enrollment_token="one-time-token",
+                enrollment_managed=True,
+            )
+
+        def enroll_host(self, **payload):
+            events.append(f"host:{payload['agent_instance_id']}")
+            return HostEnrollmentResponse(host_id="host-cloud-a")
+
+        def enroll_device(self, **payload):
+            events.append(f"device:{payload['local_device_id']}")
+            return DeviceEnrollmentResponse(device_id="device-cloud-a")
+
+        def close(self):
+            raise AssertionError("injected client must not be closed")
+
+    identity_store = HostIdentityStore(tmp_path / "host_identity.json")
+    enrollment_client = EnrollmentClient()
+    application = create_application(
+        config=enrollment_client.config,
+        device_config_store=store,
+        identity_store=identity_store,
+        enrollment_client=enrollment_client,  # type: ignore[arg-type]
+    )
+
+    assert events[0].startswith("host:agent-")
+    assert events[1] == "device:local-device-a"
+    assert application.client.config.host_id == "host-cloud-a"
+    assert application.client.config.enrollment_managed is True
+    assert [device.id for device in application.heartbeat.manager.list_devices()] == [
+        "device-cloud-a"
+    ]
+    assert store.get("local-device-a")["cloud_device_id"] == "device-cloud-a"
+    assert identity_store.load().host_id == "host-cloud-a"
+    asyncio.run(application.client.aclose())
+
+
+def test_managed_restart_reuses_identity_and_recovers_device_mapping(
+    tmp_path,
+    monkeypatch,
+) -> None:
+    monkeypatch.chdir(tmp_path)
+    store = DeviceConfigStore(tmp_path / "devices.sqlite3")
+    store.add(
+        device_id="local-device-a",
+        driver_type="wda",
+        connection_info={},
+    )
+    identity_store = HostIdentityStore(tmp_path / "host_identity.json")
+    identity_store.complete(identity_store.load_or_create(), "host-cloud-a")
+    events: list[str] = []
+
+    class EnrollmentClient:
+        config = HostAgentConfig(
+            control_plane_url="https://control.example",
+            identity_path=tmp_path / "host_identity.json",
+            enrollment_managed=True,
+        )
+
+        def enroll_host(self, **payload):
+            raise AssertionError("completed identity must skip Host enrollment")
+
+        def enroll_device(self, **payload):
+            events.append(payload["local_device_id"])
+            return DeviceEnrollmentResponse(device_id="device-cloud-a")
+
+        def close(self):
+            return None
+
+    enrollment_client = EnrollmentClient()
+    application = create_application(
+        config=enrollment_client.config,
+        device_config_store=store,
+        identity_store=identity_store,
+        enrollment_client=enrollment_client,  # type: ignore[arg-type]
+    )
+
+    assert events == ["local-device-a"]
+    assert application.client.config.host_id == "host-cloud-a"
+    assert store.get("local-device-a")["cloud_device_id"] == "device-cloud-a"
+    asyncio.run(application.client.aclose())
+
+
 def test_shutdown_cancels_long_poll_and_sends_final_heartbeat() -> None:
     async def scenario() -> None:
         claim_started = asyncio.Event()

apps/device-host-agent/tests/test_client.py

@@ -8,7 +8,11 @@ import httpx
 import pytest
 
 from cloud.internal_api.models import AssignmentModel, DeviceSnapshotModel
-from host_agent.client import HostAgentClient, StaleLeaseError
+from host_agent.client import (
+    HostAgentClient,
+    HostAgentEnrollmentClient,
+    StaleLeaseError,
+)
 from host_agent.config import HostAgentConfig
 
 
@@ -166,3 +170,58 @@ def test_result_report_retries_identical_payload_after_response_loss() -> None:
     assert len(payloads) == 2
     assert payloads[0] == payloads[1]
     assert payloads[0]["failure_reason"] == "planner unavailable"
+
+
+def test_bootstrap_client_retries_identical_enrollment_and_enrolls_device() -> None:
+    requests: list[httpx.Request] = []
+    host_attempts = 0
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal host_attempts
+        requests.append(request)
+        if request.url.path == "/internal/v1/enrollments":
+            host_attempts += 1
+            if host_attempts == 1:
+                raise httpx.ReadError("response lost", request=request)
+            return httpx.Response(201, json={"host_id": "host-cloud-a"})
+        return httpx.Response(201, json={"device_id": "device-cloud-a"})
+
+    config = _config(
+        host_id="",
+        token="",
+        enrollment_token="one-time-token",
+        enrollment_managed=True,
+    )
+    with httpx.Client(
+        transport=httpx.MockTransport(handler),
+        base_url="https://control.example",
+    ) as http_client:
+        client = HostAgentEnrollmentClient(
+            config,
+            http_client=http_client,
+            sleep=lambda _delay: None,
+        )
+        host = client.enroll_host(
+            agent_instance_id="agent-instance-a",
+            host_token="host-token-" + ("x" * 40),
+            display_name="Edge Mac",
+        )
+        client.config = _config(
+            host_id=host.host_id,
+            token="host-token-" + ("x" * 40),
+            enrollment_token="one-time-token",
+            enrollment_managed=True,
+        )
+        device = client.enroll_device(
+            local_device_id="local-device-a",
+            driver_type="wda",
+            name="iPhone",
+            capability_tags=["ios"],
+        )
+
+    assert host.host_id == "host-cloud-a"
+    assert device.device_id == "device-cloud-a"
+    assert len(requests) == 3
+    assert requests[0].content == requests[1].content
+    assert requests[0].headers["authorization"] == "Bearer one-time-token"
+    assert requests[2].headers["authorization"] == ("Bearer host-token-" + ("x" * 40))

apps/device-host-agent/tests/test_config.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from pathlib import Path
+
 import pytest
 
 from host_agent.config import (
@@ -40,11 +42,44 @@ def test_load_host_agent_config_parses_poll_and_retry_values() -> None:
     assert config.max_retry_backoff_seconds == 20
 
 
+def test_load_host_agent_config_supports_managed_enrollment(tmp_path) -> None:
+    identity_path = tmp_path / "host_identity.json"
+    config = load_host_agent_config(
+        {
+            "HOST_AGENT_CONTROL_PLANE_URL": "https://cloud.example",
+            "HOST_AGENT_ENROLLMENT_TOKEN": "one-time-token",
+            "HOST_AGENT_IDENTITY_PATH": str(identity_path),
+            "HOST_AGENT_DISPLAY_NAME": "Edge Mac",
+        }
+    )
+
+    assert config.host_id == ""
+    assert config.token == ""
+    assert config.enrollment_token == "one-time-token"
+    assert config.identity_path == identity_path
+    assert config.enrollment_managed is True
+    assert config.display_name == "Edge Mac"
+    assert "one-time-token" not in repr(config)
+
+
+def test_existing_identity_state_allows_restart_without_enrollment_token(
+    tmp_path,
+) -> None:
+    identity_path = tmp_path / "host_identity.json"
+    identity_path.write_text("{}", encoding="utf-8")
+
+    config = load_host_agent_config({"HOST_AGENT_IDENTITY_PATH": str(identity_path)})
+
+    assert config.identity_path == Path(identity_path)
+    assert config.enrollment_managed is True
+
+
 @pytest.mark.parametrize(
     "overrides",
     [
         {"HOST_AGENT_HOST_ID": ""},
         {"HOST_AGENT_TOKEN": ""},
+        {"HOST_AGENT_HOST_ID": "host-a", "HOST_AGENT_TOKEN": ""},
         {"HOST_AGENT_CONTROL_PLANE_URL": "ftp://cloud.example"},
         {"HOST_AGENT_POLL_TIMEOUT_SECONDS": "0"},
         {

apps/device-host-agent/tests/test_identity.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import os
+
+import pytest
+
+from host_agent.identity import (
+    HostIdentityStateError,
+    HostIdentityStore,
+)
+
+
+def test_identity_store_persists_pending_and_completed_state(tmp_path) -> None:
+    path = tmp_path / "state" / "host_identity.json"
+    store = HostIdentityStore(path)
+
+    pending = store.load_or_create()
+    assert pending.host_id is None
+    assert len(pending.token) >= 32
+    assert "token=" not in repr(pending)
+    assert store.load() == pending
+
+    completed = store.complete(pending, "host-cloud-a")
+    assert completed.host_id == "host-cloud-a"
+    assert store.load() == completed
+    assert "host-cloud-a" in path.read_text(encoding="utf-8")
+    if os.name != "nt":
+        assert path.stat().st_mode & 0o777 == 0o600
+
+
+def test_identity_store_rejects_invalid_or_changed_state(tmp_path) -> None:
+    path = tmp_path / "host_identity.json"
+    path.write_text('{"agent_instance_id":"a"}', encoding="utf-8")
+    store = HostIdentityStore(path)
+
+    with pytest.raises(HostIdentityStateError):
+        store.load()
+
+    path.unlink()
+    pending = store.load_or_create()
+    path.write_text(
+        '{"agent_instance_id":"other","token":"' + ("x" * 40) + '"}',
+        encoding="utf-8",
+    )
+    with pytest.raises(HostIdentityStateError, match="changed"):
+        store.complete(pending, "host-a")