feat(cloud): add edge host enrollment

This commit is contained in:
2026-07-13 13:54:16 +08:00
parent cd56facbbf
commit e61dcca801
40 changed files with 2302 additions and 48 deletions
+41 -7
View File
@@ -8,10 +8,12 @@ from cloud.internal_api.models import AssignmentModel
from device.manager import DeviceManager
from driver.registry import build_driver_factory
from host_agent.assignment import AssignmentExecutor
from host_agent.client import HostAgentClient
from host_agent.client import HostAgentClient, HostAgentEnrollmentClient
from host_agent.config import HostAgentConfig, load_host_agent_config
from host_agent.enrollment import resolve_host_identity
from host_agent.execution import create_execution_factories
from host_agent.heartbeat import HeartbeatSynchronizer
from host_agent.identity import HostIdentityStore
from host_agent.lease import ActiveAssignmentRunner
from host_agent.processor import AssignmentProcessingResult, AssignmentProcessor
from storage.device_config import DeviceConfigStore
@@ -97,9 +99,29 @@ def create_application(
config: HostAgentConfig | None = None,
manager: DeviceManager | None = None,
device_config_store: DeviceConfigStore | None = None,
identity_store: HostIdentityStore | None = None,
enrollment_client: HostAgentEnrollmentClient | None = None,
) -> HostAgentApplication:
resolved_config = config or load_host_agent_config()
resolved_manager = manager or _configured_device_manager(device_config_store)
startup_config = config or load_host_agent_config()
config_store = device_config_store or DeviceConfigStore()
owned_enrollment_client = enrollment_client is None
bootstrap_client = enrollment_client or HostAgentEnrollmentClient(startup_config)
try:
resolved_config = resolve_host_identity(
startup_config,
identity_store=identity_store
or HostIdentityStore(startup_config.identity_path),
client=bootstrap_client,
)
bootstrap_client.config = resolved_config
resolved_manager = manager or _configured_device_manager(
config_store,
config=resolved_config,
enrollment_client=bootstrap_client,
)
finally:
if owned_enrollment_client:
bootstrap_client.close()
client = HostAgentClient(resolved_config)
heartbeat = HeartbeatSynchronizer(resolved_manager, client, resolved_config)
executor = AssignmentExecutor(create_execution_factories(resolved_manager))
@@ -112,13 +134,25 @@ def create_application(
def _configured_device_manager(
config_store: DeviceConfigStore | None,
config_store: DeviceConfigStore,
*,
config: HostAgentConfig,
enrollment_client: HostAgentEnrollmentClient,
) -> DeviceManager:
manager = DeviceManager()
store = config_store or DeviceConfigStore()
for device in store.list():
for device in config_store.list():
runtime_device_id = device["device_id"]
if config.enrollment_managed:
enrollment = enrollment_client.enroll_device(
local_device_id=device["device_id"],
driver_type=device["driver_type"],
name=device["name"],
capability_tags=[],
)
runtime_device_id = enrollment.device_id
config_store.set_cloud_device_id(device["device_id"], runtime_device_id)
manager.register_device(
device["device_id"],
runtime_device_id,
build_driver_factory(
device["driver_type"],
device["connection_info"],
+111 -13
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
import asyncio
import time
from collections.abc import Awaitable, Callable
from typing import Any
@@ -9,8 +10,10 @@ import httpx
from cloud.internal_api.models import (
AssignmentModel,
ClaimResponse,
DeviceEnrollmentResponse,
DeviceSnapshotModel,
HeartbeatResponse,
HostEnrollmentResponse,
LeaseRenewalResponse,
TerminalResultResponse,
)
@@ -30,6 +33,99 @@ class StaleLeaseError(HostAgentAPIError):
pass
class HostAgentEnrollmentClient:
def __init__(
self,
config: HostAgentConfig,
*,
http_client: httpx.Client | None = None,
sleep: Callable[[float], None] = time.sleep,
) -> None:
self.config = config
self._sleep = sleep
self._owns_client = http_client is None
self._client = http_client or httpx.Client(base_url=config.control_plane_url)
def enroll_host(
self,
*,
agent_instance_id: str,
host_token: str,
display_name: str | None,
) -> HostEnrollmentResponse:
if not self.config.enrollment_token:
raise HostAgentAPIError(0, "Host enrollment token is unavailable")
response = self._request(
"POST",
"/internal/v1/enrollments",
token=self.config.enrollment_token,
json={
"agent_instance_id": agent_instance_id,
"host_token": host_token,
"display_name": display_name,
},
)
return HostEnrollmentResponse.model_validate(response.json())
def enroll_device(
self,
*,
local_device_id: str,
driver_type: str,
name: str | None,
capability_tags: list[str],
) -> DeviceEnrollmentResponse:
if not self.config.host_id or not self.config.token:
raise HostAgentAPIError(0, "Host identity is unresolved")
response = self._request(
"POST",
f"/internal/v1/hosts/{self.config.host_id}/devices/enroll",
token=self.config.token,
json={
"local_device_id": local_device_id,
"driver_type": driver_type,
"name": name,
"capability_tags": list(capability_tags),
},
)
return DeviceEnrollmentResponse.model_validate(response.json())
def close(self) -> None:
if self._owns_client:
self._client.close()
def _request(
self,
method: str,
path: str,
*,
token: str,
json: dict[str, Any],
) -> httpx.Response:
backoff = self.config.retry_backoff_seconds
for attempt in range(1, self.config.max_retry_attempts + 1):
try:
response = self._client.request(
method,
path,
json=json,
headers={"Authorization": f"Bearer {token}"},
)
except httpx.TransportError:
if attempt == self.config.max_retry_attempts:
raise
else:
if response.status_code < 500:
if response.is_success:
return response
_raise_api_error(response)
if attempt == self.config.max_retry_attempts:
_raise_api_error(response)
self._sleep(backoff)
backoff = min(backoff * 2, self.config.max_retry_backoff_seconds)
raise AssertionError("retry loop exited unexpectedly")
class HostAgentClient:
def __init__(
self,
@@ -149,22 +245,24 @@ class HostAgentClient:
if response.status_code < 500:
if response.is_success:
return response
self._raise_api_error(response)
_raise_api_error(response, stale_lease=True)
if attempt == self.config.max_retry_attempts:
self._raise_api_error(response)
_raise_api_error(response, stale_lease=True)
await self._sleep(backoff)
backoff = min(backoff * 2, self.config.max_retry_backoff_seconds)
raise AssertionError("retry loop exited unexpectedly")
@staticmethod
def _raise_api_error(response: httpx.Response) -> None:
try:
payload = response.json()
except ValueError:
payload = {}
detail = payload.get("detail") or payload.get("code") or "request rejected"
error_type = (
StaleLeaseError if response.status_code == 409 else HostAgentAPIError
)
raise error_type(response.status_code, str(detail))
def _raise_api_error(response: httpx.Response, *, stale_lease: bool = False) -> None:
try:
payload = response.json()
except ValueError:
payload = {}
detail = payload.get("detail") or payload.get("code") or "request rejected"
error_type = (
StaleLeaseError
if stale_lease and response.status_code == 409
else HostAgentAPIError
)
raise error_type(response.status_code, str(detail))
+25 -7
View File
@@ -2,7 +2,8 @@ from __future__ import annotations
import os
from collections.abc import Mapping
from dataclasses import dataclass
from dataclasses import dataclass, field
from pathlib import Path
from urllib.parse import urlparse
@@ -13,8 +14,12 @@ class HostAgentConfigurationError(ValueError):
@dataclass(frozen=True)
class HostAgentConfig:
control_plane_url: str
host_id: str
token: str
host_id: str = ""
token: str = field(default="", repr=False)
enrollment_token: str = field(default="", repr=False)
identity_path: Path = Path("tasks/host_identity.json")
enrollment_managed: bool = False
display_name: str | None = None
heartbeat_interval_seconds: float = 30.0
poll_timeout_seconds: float = 20.0
retry_backoff_seconds: float = 1.0
@@ -41,16 +46,29 @@ def load_host_agent_config(
)
host_id = values.get("HOST_AGENT_HOST_ID", "").strip()
if not host_id:
raise HostAgentConfigurationError("HOST_AGENT_HOST_ID is required")
token = values.get("HOST_AGENT_TOKEN", "").strip()
if not token:
raise HostAgentConfigurationError("HOST_AGENT_TOKEN is required")
if bool(host_id) != bool(token):
raise HostAgentConfigurationError(
"HOST_AGENT_HOST_ID and HOST_AGENT_TOKEN must be configured together"
)
enrollment_token = values.get("HOST_AGENT_ENROLLMENT_TOKEN", "").strip()
identity_path = Path(
values.get("HOST_AGENT_IDENTITY_PATH", "tasks/host_identity.json").strip()
)
if not host_id and not enrollment_token and not identity_path.is_file():
raise HostAgentConfigurationError(
"explicit Host credentials, an enrollment token, or existing identity state "
"is required"
)
config = HostAgentConfig(
control_plane_url=control_plane_url,
host_id=host_id,
token=token,
enrollment_token=enrollment_token,
identity_path=identity_path,
enrollment_managed=not bool(host_id),
display_name=values.get("HOST_AGENT_DISPLAY_NAME") or None,
heartbeat_interval_seconds=_positive_float(
values,
"HOST_AGENT_HEARTBEAT_INTERVAL_SECONDS",
@@ -0,0 +1,35 @@
from __future__ import annotations
from dataclasses import replace
from host_agent.client import HostAgentEnrollmentClient
from host_agent.config import HostAgentConfig, HostAgentConfigurationError
from host_agent.identity import HostIdentityStore
def resolve_host_identity(
config: HostAgentConfig,
*,
identity_store: HostIdentityStore,
client: HostAgentEnrollmentClient,
) -> HostAgentConfig:
if config.host_id and config.token:
return config
state = identity_store.load_or_create()
if state.host_id is None:
if not config.enrollment_token:
raise HostAgentConfigurationError(
"HOST_AGENT_ENROLLMENT_TOKEN is required to complete enrollment"
)
response = client.enroll_host(
agent_instance_id=state.agent_instance_id,
host_token=state.token,
display_name=config.display_name,
)
state = identity_store.complete(state, response.host_id)
return replace(
config,
host_id=state.host_id or "",
token=state.token,
enrollment_managed=True,
)
@@ -0,0 +1,101 @@
from __future__ import annotations
import json
import os
from dataclasses import dataclass, field, replace
from pathlib import Path
from secrets import token_urlsafe
from uuid import uuid4
class HostIdentityStateError(RuntimeError):
"""Raised when persisted Host enrollment identity is missing or invalid."""
@dataclass(frozen=True)
class HostIdentityState:
agent_instance_id: str
token: str = field(repr=False)
host_id: str | None = None
class HostIdentityStore:
def __init__(self, path: str | Path) -> None:
self.path = Path(path)
def load(self) -> HostIdentityState | None:
if not self.path.exists():
return None
try:
payload = json.loads(self.path.read_text(encoding="utf-8"))
except (OSError, ValueError, json.JSONDecodeError) as exc:
raise HostIdentityStateError("Host identity state is unreadable") from exc
if not isinstance(payload, dict):
raise HostIdentityStateError("Host identity state must be an object")
agent_instance_id = payload.get("agent_instance_id")
token = payload.get("token")
host_id = payload.get("host_id")
if (
not isinstance(agent_instance_id, str)
or not agent_instance_id
or not isinstance(token, str)
or len(token) < 32
or (host_id is not None and (not isinstance(host_id, str) or not host_id))
):
raise HostIdentityStateError("Host identity state is invalid")
return HostIdentityState(
agent_instance_id=agent_instance_id,
token=token,
host_id=host_id,
)
def load_or_create(self) -> HostIdentityState:
existing = self.load()
if existing is not None:
return existing
state = HostIdentityState(
agent_instance_id=f"agent-{uuid4().hex}",
token=token_urlsafe(48),
)
self._write(state)
return state
def complete(self, state: HostIdentityState, host_id: str) -> HostIdentityState:
if not host_id:
raise ValueError("host_id must not be empty")
current = self.load()
if current is not None and (
current.agent_instance_id != state.agent_instance_id
or current.token != state.token
):
raise HostIdentityStateError("Host identity changed during enrollment")
completed = replace(state, host_id=host_id)
self._write(completed)
return completed
def _write(self, state: HostIdentityState) -> None:
self.path.parent.mkdir(parents=True, exist_ok=True)
temporary = self.path.with_name(f".{self.path.name}.{uuid4().hex}.tmp")
payload = {
"agent_instance_id": state.agent_instance_id,
"token": state.token,
"host_id": state.host_id,
}
try:
temporary.write_text(
json.dumps(payload, ensure_ascii=True, indent=2) + "\n",
encoding="utf-8",
)
_restrict_permissions(temporary)
os.replace(temporary, self.path)
_restrict_permissions(self.path)
finally:
if temporary.exists():
temporary.unlink()
def _restrict_permissions(path: Path) -> None:
try:
path.chmod(0o600)
except OSError:
return