feat(cloud): remove static credentials and add host console
Tests / Test No test results found

This commit is contained in:
2026-07-13 19:45:53 +08:00
parent efeb3eb926
commit c162c2501b
61 changed files with 3118 additions and 1221 deletions
+115 -22
View File
@@ -4,18 +4,25 @@ import asyncio
from contextlib import suppress
from dataclasses import dataclass
import uvicorn
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, HostAgentEnrollmentClient
from host_agent.config import HostAgentConfig, load_host_agent_config
from host_agent.devices import register_local_device
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.history import ConsoleHistoryStore
from host_agent.identity import HostIdentityStore
from host_agent.lease import ActiveAssignmentRunner
from host_agent.local_account import LocalAccountStore
from host_agent.processor import AssignmentProcessingResult, AssignmentProcessor
from host_agent.status import AgentStatusTracker
from host_agent.web.app import create_console_app
from host_agent.web.auth import SessionManager
from storage.device_config import DeviceConfigStore
@@ -24,6 +31,8 @@ class HostAgentApplication:
client: HostAgentClient
heartbeat: HeartbeatSynchronizer
processor: AssignmentProcessor
console_server: uvicorn.Server | None = None
console_enrollment_client: HostAgentEnrollmentClient | None = None
def run(self) -> None:
asyncio.run(self.run_async())
@@ -32,6 +41,11 @@ class HostAgentApplication:
stop_requested = stop or asyncio.Event()
heartbeat_stop = asyncio.Event()
heartbeat_task = asyncio.create_task(self.heartbeat.run(heartbeat_stop))
console_task = (
asyncio.create_task(self.console_server.serve())
if self.console_server is not None
else None
)
active_processing: asyncio.Task[AssignmentProcessingResult] | None = None
try:
while not stop_requested.is_set():
@@ -60,11 +74,18 @@ class HostAgentApplication:
with suppress(Exception):
await asyncio.shield(active_processing)
heartbeat_stop.set()
if self.console_server is not None:
self.console_server.should_exit = True
try:
await asyncio.gather(heartbeat_task, return_exceptions=True)
with suppress(Exception):
await self.heartbeat.sync_once()
if console_task is not None:
with suppress(asyncio.CancelledError):
await asyncio.gather(console_task, return_exceptions=True)
finally:
if self.console_enrollment_client is not None:
self.console_enrollment_client.close()
await self.client.aclose()
async def _claim_until_stopped(
@@ -104,13 +125,15 @@ def create_application(
) -> HostAgentApplication:
startup_config = config or load_host_agent_config()
config_store = device_config_store or DeviceConfigStore()
resolved_identity_store = identity_store or HostIdentityStore(
startup_config.identity_path
)
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),
identity_store=resolved_identity_store,
client=bootstrap_client,
)
bootstrap_client.config = resolved_config
@@ -123,13 +146,93 @@ def create_application(
if owned_enrollment_client:
bootstrap_client.close()
client = HostAgentClient(resolved_config)
heartbeat = HeartbeatSynchronizer(resolved_manager, client, resolved_config)
history_store: ConsoleHistoryStore | None = None
status_tracker: AgentStatusTracker | None = None
console_server: uvicorn.Server | None = None
console_enrollment_client: HostAgentEnrollmentClient | None = None
if resolved_config.console_enabled:
history_store = ConsoleHistoryStore(
resolved_config.identity_path.parent / "host_console_history.sqlite3",
limit=resolved_config.console_history_limit,
)
status_tracker = AgentStatusTracker()
if resolved_config.enrollment_managed:
console_enrollment_client = HostAgentEnrollmentClient(resolved_config)
console_app = create_console_app(
config=resolved_config,
manager=resolved_manager,
config_store=config_store,
local_account_store=LocalAccountStore(resolved_config.local_account_path),
identity_store=resolved_identity_store,
history_store=history_store,
status_tracker=status_tracker,
session_manager=SessionManager(
ttl_seconds=resolved_config.console_session_ttl_seconds
),
enrollment_client=console_enrollment_client,
)
console_server = uvicorn.Server(
uvicorn.Config(
console_app,
host=resolved_config.console_bind_host,
port=resolved_config.console_port,
log_level="warning",
)
)
heartbeat = HeartbeatSynchronizer(
resolved_manager,
client,
resolved_config,
status_tracker=status_tracker,
on_sync=(
(
lambda device_count: history_store.record_heartbeat(
device_count=device_count
)
)
if history_store is not None
else None
),
)
executor = AssignmentExecutor(create_execution_factories(resolved_manager))
active_runner = ActiveAssignmentRunner(client, executor)
processor = AssignmentProcessor(
client,
active_runner,
status_tracker=status_tracker,
on_result=(
(
lambda assignment, result: _record_assignment_history(
history_store, assignment, result
)
)
if history_store is not None
else None
),
)
return HostAgentApplication(
client=client,
heartbeat=heartbeat,
processor=AssignmentProcessor(client, active_runner),
processor=processor,
console_server=console_server,
console_enrollment_client=console_enrollment_client,
)
def _record_assignment_history(
history_store: ConsoleHistoryStore,
assignment: AssignmentModel,
result: AssignmentProcessingResult,
) -> None:
status = "done" if result.execution.status == "done" else "failed"
history_store.record_assignment(
task_id=assignment.task_id,
attempt=assignment.attempt,
status=status,
failure_reason=result.execution.failure_reason if status == "failed" else None,
device_id=assignment.device_id,
)
@@ -141,24 +244,14 @@ def _configured_device_manager(
) -> DeviceManager:
manager = DeviceManager()
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(
runtime_device_id,
build_driver_factory(
device["driver_type"],
device["connection_info"],
),
name=device["name"],
register_local_device(
config_store,
manager,
device_id=device["device_id"],
driver_type=device["driver_type"],
connection_info=device["connection_info"],
name=device["name"],
config=config,
enrollment_client=enrollment_client,
)
return manager
+1 -1
View File
@@ -56,7 +56,7 @@ class HostAgentEnrollmentClient:
response = self._request(
"POST",
"/internal/v1/enrollments",
token=self.config.enrollment_token or None,
token=None,
json={
"agent_instance_id": agent_instance_id,
"host_token": host_token,
+48 -12
View File
@@ -11,12 +11,14 @@ class HostAgentConfigurationError(ValueError):
"""Raised when Host Agent process configuration is invalid."""
_LOOPBACK_BIND_HOSTS = frozenset({"127.0.0.1", "localhost", "::1"})
@dataclass(frozen=True)
class HostAgentConfig:
control_plane_url: 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")
local_account_path: Path = Path("tasks/host_local_account.json")
enrollment_managed: bool = False
@@ -26,6 +28,12 @@ class HostAgentConfig:
retry_backoff_seconds: float = 1.0
max_retry_backoff_seconds: float = 30.0
max_retry_attempts: int = 5
console_enabled: bool = False
console_bind_host: str = "127.0.0.1"
console_port: int = 8765
console_allow_non_loopback: bool = False
console_session_ttl_seconds: float = 43200.0
console_history_limit: int = 200
def load_host_agent_config(
@@ -46,13 +54,6 @@ def load_host_agent_config(
"HOST_AGENT_CONTROL_PLANE_URL must be an HTTP(S) URL"
)
host_id = values.get("HOST_AGENT_HOST_ID", "").strip()
token = values.get("HOST_AGENT_TOKEN", "").strip()
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()
)
@@ -64,12 +65,9 @@ def load_host_agent_config(
config = HostAgentConfig(
control_plane_url=control_plane_url,
host_id=host_id,
token=token,
enrollment_token=enrollment_token,
identity_path=identity_path,
local_account_path=local_account_path,
enrollment_managed=not bool(host_id),
enrollment_managed=True,
display_name=values.get("HOST_AGENT_DISPLAY_NAME") or None,
heartbeat_interval_seconds=_positive_float(
values,
@@ -96,11 +94,38 @@ def load_host_agent_config(
"HOST_AGENT_MAX_RETRY_ATTEMPTS",
5,
),
console_enabled=_truthy(values, "HOST_AGENT_CONSOLE_ENABLED", False),
console_bind_host=values.get(
"HOST_AGENT_CONSOLE_BIND_HOST", "127.0.0.1"
).strip(),
console_port=_positive_int(values, "HOST_AGENT_CONSOLE_PORT", 8765),
console_allow_non_loopback=_truthy(
values, "HOST_AGENT_CONSOLE_ALLOW_NON_LOOPBACK", False
),
console_session_ttl_seconds=_positive_float(
values,
"HOST_AGENT_CONSOLE_SESSION_TTL_SECONDS",
43200.0,
),
console_history_limit=_positive_int(
values,
"HOST_AGENT_CONSOLE_HISTORY_LIMIT",
200,
),
)
if config.max_retry_backoff_seconds < config.retry_backoff_seconds:
raise HostAgentConfigurationError(
"maximum retry backoff must not be less than initial backoff"
)
if config.console_enabled and (
config.console_bind_host not in _LOOPBACK_BIND_HOSTS
and not config.console_allow_non_loopback
):
raise HostAgentConfigurationError(
"HOST_AGENT_CONSOLE_BIND_HOST must be a loopback address "
f"({', '.join(sorted(_LOOPBACK_BIND_HOSTS))}) unless "
"HOST_AGENT_CONSOLE_ALLOW_NON_LOOPBACK is set"
)
return config
@@ -136,3 +161,14 @@ def _positive_int(
if value <= 0:
raise HostAgentConfigurationError(f"{name} must be greater than zero")
return value
def _truthy(
values: Mapping[str, str],
name: str,
default: bool,
) -> bool:
raw_value = values.get(name)
if raw_value is None:
return default
return raw_value.strip().lower() in {"true", "1"}
@@ -0,0 +1,66 @@
from __future__ import annotations
from typing import Any
from device.manager import DeviceManager
from driver.registry import build_driver_factory
from host_agent.client import HostAgentEnrollmentClient
from host_agent.config import HostAgentConfig
from storage.device_config import DeviceConfigStore
def register_local_device(
config_store: DeviceConfigStore,
manager: DeviceManager,
*,
device_id: str,
driver_type: str,
connection_info: dict[str, Any],
name: str | None,
config: HostAgentConfig,
enrollment_client: HostAgentEnrollmentClient | None,
) -> None:
previous = config_store.get(device_id)
previous_runtime_id = (
previous["cloud_device_id"] or previous["device_id"] if previous else None
)
config_store.add(
device_id=device_id,
name=name,
driver_type=driver_type,
connection_info=connection_info,
)
runtime_device_id = device_id
if config.enrollment_managed:
assert enrollment_client is not None
enrollment = enrollment_client.enroll_device(
local_device_id=device_id,
driver_type=driver_type,
name=name,
capability_tags=[],
)
runtime_device_id = enrollment.device_id
config_store.set_cloud_device_id(device_id, runtime_device_id)
if previous_runtime_id is not None and previous_runtime_id != runtime_device_id:
manager.unregister_device(previous_runtime_id)
manager.register_device(
runtime_device_id,
build_driver_factory(driver_type, connection_info),
name=name,
driver_type=driver_type,
connection_info=connection_info,
)
def unregister_local_device(
config_store: DeviceConfigStore,
manager: DeviceManager,
*,
device_id: str,
) -> None:
record = config_store.get(device_id)
if record is None:
return
runtime_device_id = record["cloud_device_id"] or record["device_id"]
manager.unregister_device(runtime_device_id)
config_store.remove(device_id)
+16 -2
View File
@@ -8,6 +8,7 @@ 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.status import AgentStatusTracker
if TYPE_CHECKING:
from collections.abc import Awaitable, Callable
@@ -34,18 +35,31 @@ class HeartbeatSynchronizer:
*,
address: str | None = None,
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
status_tracker: AgentStatusTracker | None = None,
on_sync: Callable[[int], None] | 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
async def sync_once(self) -> HeartbeatResponse:
return await self.client.heartbeat(
build_device_snapshot(self.manager),
snapshot = build_device_snapshot(self.manager)
response = await self.client.heartbeat(
snapshot,
address=self.address,
)
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()
@@ -0,0 +1,105 @@
from __future__ import annotations
import json
import sqlite3
from collections.abc import Callable
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
class ConsoleHistoryStore:
def __init__(
self,
db_path: str | Path = "tasks/host_console_history.sqlite3",
*,
limit: int = 200,
now: Callable[[], datetime] | None = None,
) -> None:
self.db_path = Path(db_path)
self.limit = limit
self._now = now or (lambda: datetime.now(UTC))
self.db_path.parent.mkdir(parents=True, exist_ok=True)
self._ensure_schema()
def record_assignment(
self,
*,
task_id: str,
attempt: int,
status: str,
failure_reason: str | None,
device_id: str,
) -> None:
summary = f"{task_id} attempt {attempt} on {device_id}: {status}"
if failure_reason:
summary = f"{summary} ({failure_reason})"
detail = {
"task_id": task_id,
"attempt": attempt,
"status": status,
"failure_reason": failure_reason,
"device_id": device_id,
}
self._insert("assignment", summary, detail)
def record_heartbeat(self, *, device_count: int) -> None:
summary = f"heartbeat: {device_count} devices"
detail = {"device_count": device_count}
self._insert("heartbeat", summary, detail)
def list_recent(self, limit: int | None = None) -> list[dict[str, Any]]:
effective_limit = limit if limit is not None else self.limit
with self._connect() as connection:
rows = connection.execute(
"select kind, occurred_at, summary, detail_json"
" from history_entries order by id desc limit ?",
(effective_limit,),
).fetchall()
return [
{
"kind": row["kind"],
"occurred_at": row["occurred_at"],
"summary": row["summary"],
"detail": json.loads(row["detail_json"]),
}
for row in rows
]
def _insert(self, kind: str, summary: str, detail: dict[str, Any]) -> None:
occurred_at = self._now().isoformat()
with self._connect() as connection:
connection.execute(
"""
insert into history_entries (kind, occurred_at, summary, detail_json)
values (?, ?, ?, ?)
""",
(kind, occurred_at, summary, json.dumps(detail, ensure_ascii=False)),
)
connection.execute(
"""
delete from history_entries where id not in (
select id from history_entries order by id desc limit ?
)
""",
(self.limit,),
)
def _ensure_schema(self) -> None:
with self._connect() as connection:
connection.execute(
"""
create table if not exists history_entries (
id integer primary key autoincrement,
kind text not null,
occurred_at text not null,
summary text not null,
detail_json text not null
)
"""
)
def _connect(self) -> sqlite3.Connection:
connection = sqlite3.connect(self.db_path)
connection.row_factory = sqlite3.Row
return connection
+36 -14
View File
@@ -1,11 +1,15 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import Literal, Protocol
from typing import TYPE_CHECKING, Literal, Protocol
from cloud.internal_api.models import AssignmentModel
from host_agent.assignment import AssignmentExecutionResult
from host_agent.client import HostAgentClient
from host_agent.status import AgentStatusTracker
if TYPE_CHECKING:
from collections.abc import Callable
class ActiveAssignmentExecutor(Protocol):
@@ -25,24 +29,42 @@ class AssignmentProcessor:
self,
client: HostAgentClient,
active_executor: ActiveAssignmentExecutor,
*,
status_tracker: AgentStatusTracker | None = None,
on_result: Callable[[AssignmentModel, AssignmentProcessingResult], None]
| None = None,
) -> None:
self.client = client
self.active_executor = active_executor
self.status_tracker = status_tracker
self.on_result = on_result
def request_stop(self) -> None:
self.active_executor.request_stop()
async def process(self, assignment: AssignmentModel) -> AssignmentProcessingResult:
execution = await self.active_executor.run(assignment)
status = "done" if execution.status == "done" else "failed"
failure_reason = execution.failure_reason if status == "failed" else None
response = await self.client.report_result(
assignment,
status=status,
failure_reason=failure_reason,
result=dict(execution.metadata),
)
return AssignmentProcessingResult(
execution=execution,
report_status=response.status,
)
if self.status_tracker is not None:
self.status_tracker.mark_assignment_started(assignment)
try:
execution = await self.active_executor.run(assignment)
status = "done" if execution.status == "done" else "failed"
failure_reason = execution.failure_reason if status == "failed" else None
response = await self.client.report_result(
assignment,
status=status,
failure_reason=failure_reason,
result=dict(execution.metadata),
)
result = AssignmentProcessingResult(
execution=execution,
report_status=response.status,
)
finally:
if self.status_tracker is not None:
self.status_tracker.mark_assignment_finished()
if self.on_result is not None:
try:
self.on_result(assignment, result)
except Exception:
pass
return result
@@ -0,0 +1,84 @@
from __future__ import annotations
import threading
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any
from cloud.internal_api.models import AssignmentModel
if TYPE_CHECKING:
from collections.abc import Callable
@dataclass(frozen=True)
class _CurrentAssignment:
task_id: str
device_id: str
goal: str | None
workflow_definition_id: str | None
started_at: datetime
@dataclass(frozen=True)
class _LastHeartbeat:
ok: bool
device_count: int
at: datetime
class AgentStatusTracker:
def __init__(self, *, now: Callable[[], datetime] | None = None) -> None:
self._now = now or (lambda: datetime.now(UTC))
self._lock = threading.Lock()
self._current_assignment: _CurrentAssignment | None = None
self._last_heartbeat: _LastHeartbeat | None = None
def mark_assignment_started(self, assignment: AssignmentModel) -> None:
with self._lock:
self._current_assignment = _CurrentAssignment(
task_id=assignment.task_id,
device_id=assignment.device_id,
goal=assignment.goal,
workflow_definition_id=assignment.workflow_definition_id,
started_at=self._now(),
)
def mark_assignment_finished(self) -> None:
with self._lock:
self._current_assignment = None
def mark_heartbeat(self, *, ok: bool, device_count: int) -> None:
with self._lock:
self._last_heartbeat = _LastHeartbeat(
ok=ok,
device_count=device_count,
at=self._now(),
)
def snapshot(self) -> dict[str, Any]:
with self._lock:
current_assignment = self._current_assignment
last_heartbeat = self._last_heartbeat
return {
"current_assignment": (
{
"task_id": current_assignment.task_id,
"device_id": current_assignment.device_id,
"goal": current_assignment.goal,
"workflow_definition_id": current_assignment.workflow_definition_id,
"started_at": current_assignment.started_at.isoformat(),
}
if current_assignment is not None
else None
),
"last_heartbeat": (
{
"ok": last_heartbeat.ok,
"device_count": last_heartbeat.device_count,
"at": last_heartbeat.at.isoformat(),
}
if last_heartbeat is not None
else None
),
}
@@ -0,0 +1,535 @@
from __future__ import annotations
import asyncio
import json
from html import escape as _escape
from typing import Any
from fastapi import Depends, FastAPI, HTTPException, Request
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response
from device.manager import DeviceManager
from host_agent.client import HostAgentEnrollmentClient
from host_agent.config import HostAgentConfig
from host_agent.devices import register_local_device, unregister_local_device
from host_agent.history import ConsoleHistoryStore
from host_agent.identity import HostIdentityState, HostIdentityStore
from host_agent.local_account import LocalAccountState, LocalAccountStore
from host_agent.status import AgentStatusTracker
from host_agent.web.auth import (
SessionManager,
SessionState,
attempt_login,
change_password,
)
from storage.device_config import DeviceConfigStore
SESSION_COOKIE_NAME = "host_console_session"
CSRF_HEADER_NAME = "X-CSRF-Token"
CSRF_FORM_FIELD = "csrf_token"
_LOOPBACK_BIND_HOSTS = frozenset({"127.0.0.1", "localhost", "::1"})
_CSS = """
body { font-family: system-ui, sans-serif; margin: 0; background: #f5f5f5; color: #222; }
header { background: #20303f; color: #fff; padding: 0.75rem 1.5rem; }
header nav { display: inline; margin-left: 1.5rem; }
header nav a, header nav form { display: inline-block; margin-right: 1rem; }
header a { color: #fff; text-decoration: none; }
header button { background: none; border: none; color: #fff; text-decoration: underline; cursor: pointer; padding: 0; font: inherit; }
main { padding: 1.5rem; max-width: 960px; margin: 0 auto; }
table { border-collapse: collapse; width: 100%; margin-bottom: 1rem; background: #fff; }
th, td { border: 1px solid #ccc; padding: 0.4rem 0.6rem; text-align: left; }
form.inline { display: inline; margin: 0; }
.error { color: #b00020; }
.notice { color: #1b5e20; }
"""
def escape(value: object) -> str:
if value is None:
return ""
return _escape(str(value), quote=True)
def _chrome(title: str, body_html: str, *, session: SessionState | None) -> str:
nav = ""
if session is not None:
nav = f"""
<nav>
<a href="/">Status</a>
<a href="/devices">Devices</a>
<a href="/account">Account</a>
<a href="/history">History</a>
<form class="inline" method="post" action="/logout">
<input type="hidden" name="{CSRF_FORM_FIELD}" value="{escape(session.csrf_token)}">
<button type="submit">Logout</button>
</form>
</nav>
"""
return f"""<!doctype html>
<html>
<head>
<meta charset="utf-8">
<title>{escape(title)}</title>
<style>{_CSS}</style>
</head>
<body>
<header>
<strong>Host Agent Console</strong>
{nav}
</header>
<main>
{body_html}
</main>
</body>
</html>"""
def _login_page(
*, account: LocalAccountState | None, error: str | None = None
) -> HTMLResponse:
if account is None:
body = """
<h1>Login</h1>
<p>No local account exists yet. Run <code>device-host-agent setup</code>
on this machine to create one before logging in to the console.</p>
"""
return HTMLResponse(_chrome("Login", body, session=None))
error_html = f'<p class="error">{escape(error)}</p>' if error else ""
body = f"""
<h1>Login</h1>
{error_html}
<form method="post" action="/login">
<label>Username <input type="text" name="username" required></label><br>
<label>Password <input type="password" name="password" required></label><br>
<button type="submit">Log in</button>
</form>
"""
return HTMLResponse(_chrome("Login", body, session=None))
def _dashboard_body(
*,
identity: HostIdentityState | None,
snapshot: dict[str, Any],
devices: list[Any],
config: HostAgentConfig,
) -> str:
heartbeat = snapshot.get("last_heartbeat")
assignment = snapshot.get("current_assignment")
heartbeat_text = (
f"{'ok' if heartbeat['ok'] else 'failed'} at {heartbeat['at']} "
f"({heartbeat['device_count']} devices)"
if heartbeat
else "never"
)
assignment_text = (
f"{assignment['task_id']} on {assignment['device_id']} "
f"(started {assignment['started_at']})"
if assignment
else "none"
)
device_rows = "".join(
f"<tr><td>{escape(device.id)}</td><td>{escape(device.name or '')}</td>"
f"<td>{escape(device.driver_type)}</td><td>{escape(device.status)}</td></tr>"
for device in devices
)
return f"""
<h1>Status</h1>
<section>
<h2>Enrollment</h2>
<p>Host ID: {escape(identity.host_id if identity else None) or "not enrolled"}</p>
<p>Agent instance ID: {escape(identity.agent_instance_id if identity else None) or "unknown"}</p>
<p>Control plane: {escape(config.control_plane_url)}</p>
</section>
<section>
<h2>Heartbeat</h2>
<p id="last-heartbeat">{escape(heartbeat_text)}</p>
</section>
<section>
<h2>Current assignment</h2>
<p id="current-assignment">{escape(assignment_text)}</p>
</section>
<section>
<h2>Devices</h2>
<table>
<thead><tr><th>ID</th><th>Name</th><th>Driver</th><th>Status</th></tr></thead>
<tbody id="device-status-body">{device_rows}</tbody>
</table>
</section>
<script>
(function () {{
function render(data) {{
var hb = data.status.last_heartbeat;
document.getElementById("last-heartbeat").textContent = hb
? (hb.ok ? "ok" : "failed") + " at " + hb.at + " (" + hb.device_count + " devices)"
: "never";
var current = data.status.current_assignment;
document.getElementById("current-assignment").textContent = current
? current.task_id + " on " + current.device_id + " (started " + current.started_at + ")"
: "none";
var body = document.getElementById("device-status-body");
body.innerHTML = "";
data.devices.forEach(function (device) {{
var row = document.createElement("tr");
["id", "name", "driver_type", "status"].forEach(function (key) {{
var cell = document.createElement("td");
cell.textContent = device[key] || "";
row.appendChild(cell);
}});
body.appendChild(row);
}});
}}
function poll() {{
fetch("/api/status", {{ credentials: "same-origin" }})
.then(function (response) {{ return response.ok ? response.json() : null; }})
.then(function (data) {{ if (data) render(data); }})
.catch(function () {{}});
}}
setInterval(poll, 5000);
}})();
</script>
"""
def _devices_body(
*,
devices: list[dict[str, Any]],
csrf_token: str,
edit_record: dict[str, Any] | None,
error: str | None,
) -> str:
error_html = f'<p class="error">{escape(error)}</p>' if error else ""
rows = "".join(
f"""
<tr>
<td>{escape(device["device_id"])}</td>
<td>{escape(device["name"] or "")}</td>
<td>{escape(device["driver_type"])}</td>
<td>{escape(device["cloud_device_id"] or "")}</td>
<td>
<a href="/devices?edit={escape(device["device_id"])}">Edit</a>
<form class="inline" method="post" action="/devices/remove">
<input type="hidden" name="{CSRF_FORM_FIELD}" value="{escape(csrf_token)}">
<input type="hidden" name="device_id" value="{escape(device["device_id"])}">
<button type="submit">Remove</button>
</form>
</td>
</tr>
"""
for device in devices
)
form_device_id = escape(edit_record["device_id"]) if edit_record else ""
form_name = escape(edit_record["name"] or "") if edit_record else ""
form_driver_type = escape(edit_record["driver_type"]) if edit_record else "wda"
form_connection_info = (
escape(json.dumps(edit_record["connection_info"])) if edit_record else "{}"
)
return f"""
<h1>Devices</h1>
{error_html}
<table>
<thead><tr><th>ID</th><th>Name</th><th>Driver</th><th>Cloud ID</th><th></th></tr></thead>
<tbody>{rows}</tbody>
</table>
<h2>{"Edit device" if edit_record else "Add device"}</h2>
<form method="post" action="/devices/save">
<input type="hidden" name="{CSRF_FORM_FIELD}" value="{escape(csrf_token)}">
<label>Device ID <input type="text" name="device_id" value="{form_device_id}" required></label><br>
<label>Name <input type="text" name="name" value="{form_name}"></label><br>
<label>Driver type <input type="text" name="driver_type" value="{form_driver_type}" required></label><br>
<label>Connection info (JSON)<br>
<textarea name="connection_info" rows="3" cols="50">{form_connection_info}</textarea>
</label><br>
<button type="submit">Save</button>
</form>
"""
def _account_body(*, csrf_token: str, message: str | None, error: str | None) -> str:
message_html = f'<p class="notice">{escape(message)}</p>' if message else ""
error_html = f'<p class="error">{escape(error)}</p>' if error else ""
return f"""
<h1>Account</h1>
{message_html}
{error_html}
<form method="post" action="/account">
<input type="hidden" name="{CSRF_FORM_FIELD}" value="{escape(csrf_token)}">
<label>Current password <input type="password" name="current_password" required></label><br>
<label>New password <input type="password" name="new_password" required></label><br>
<label>Confirm new password <input type="password" name="confirm_password" required></label><br>
<button type="submit">Change password</button>
</form>
"""
def _history_body(entries: list[dict[str, Any]]) -> str:
rows = "".join(
f"<tr><td>{escape(entry['occurred_at'])}</td><td>{escape(entry['kind'])}</td>"
f"<td>{escape(entry['summary'])}</td></tr>"
for entry in entries
)
return f"""
<h1>History</h1>
<table>
<thead><tr><th>Time</th><th>Kind</th><th>Summary</th></tr></thead>
<tbody>{rows}</tbody>
</table>
"""
def create_console_app(
*,
config: HostAgentConfig,
manager: DeviceManager,
config_store: DeviceConfigStore,
local_account_store: LocalAccountStore,
identity_store: HostIdentityStore,
history_store: ConsoleHistoryStore,
status_tracker: AgentStatusTracker,
session_manager: SessionManager,
enrollment_client: HostAgentEnrollmentClient | None,
) -> FastAPI:
app = FastAPI(title="Host Agent Console")
cookie_secure = config.console_bind_host not in _LOOPBACK_BIND_HOSTS
def _session_token(request: Request) -> str | None:
return request.cookies.get(SESSION_COOKIE_NAME)
def require_session(request: Request) -> SessionState:
token = _session_token(request)
state = session_manager.validate(token) if token else None
if state is None:
raise HTTPException(status_code=303, headers={"Location": "/login"})
return state
async def require_csrf(
request: Request,
session: SessionState = Depends(require_session),
) -> SessionState:
provided = request.headers.get(CSRF_HEADER_NAME)
if not provided:
form = await request.form()
raw = form.get(CSRF_FORM_FIELD)
provided = str(raw) if raw is not None else None
token = _session_token(request)
if (
not token
or not provided
or not session_manager.validate_csrf(token, provided)
):
raise HTTPException(status_code=403, detail="invalid CSRF token")
return session
@app.get("/login", response_class=HTMLResponse)
async def login_page() -> HTMLResponse:
account = await asyncio.to_thread(local_account_store.load)
return _login_page(account=account)
@app.post("/login")
async def login_submit(request: Request) -> Response:
account = await asyncio.to_thread(local_account_store.load)
if account is None:
return _login_page(account=None)
form = await request.form()
username = str(form.get("username", ""))
password = str(form.get("password", ""))
ok = await asyncio.to_thread(
attempt_login, local_account_store, username=username, password=password
)
if not ok:
return _login_page(account=account, error="Invalid username or password.")
session_token, _ = session_manager.create_session(username)
response = RedirectResponse(url="/", status_code=303)
response.set_cookie(
key=SESSION_COOKIE_NAME,
value=session_token,
httponly=True,
samesite="strict",
secure=cookie_secure,
path="/",
)
return response
@app.post("/logout")
async def logout(
request: Request,
session: SessionState = Depends(require_csrf),
) -> Response:
token = _session_token(request)
if token:
session_manager.invalidate(token)
response = RedirectResponse(url="/login", status_code=303)
response.delete_cookie(key=SESSION_COOKIE_NAME, path="/")
return response
@app.get("/", response_class=HTMLResponse)
async def dashboard(
session: SessionState = Depends(require_session),
) -> HTMLResponse:
identity = await asyncio.to_thread(identity_store.load)
snapshot = status_tracker.snapshot()
devices = manager.list_devices()
body = _dashboard_body(
identity=identity, snapshot=snapshot, devices=devices, config=config
)
return HTMLResponse(_chrome("Status", body, session=session))
@app.get("/api/status")
async def api_status(
session: SessionState = Depends(require_session),
) -> JSONResponse:
snapshot = status_tracker.snapshot()
devices = [
{
"id": device.id,
"name": device.name,
"driver_type": device.driver_type,
"status": device.status,
}
for device in manager.list_devices()
]
return JSONResponse({"status": snapshot, "devices": devices})
@app.get("/devices", response_class=HTMLResponse)
async def devices_page(
request: Request,
session: SessionState = Depends(require_session),
) -> HTMLResponse:
devices = await asyncio.to_thread(config_store.list)
edit_id = request.query_params.get("edit")
edit_record = (
await asyncio.to_thread(config_store.get, edit_id) if edit_id else None
)
body = _devices_body(
devices=devices,
csrf_token=session.csrf_token,
edit_record=edit_record,
error=None,
)
return HTMLResponse(_chrome("Devices", body, session=session))
@app.post("/devices/save")
async def devices_save(
request: Request,
session: SessionState = Depends(require_csrf),
) -> Response:
form = await request.form()
device_id = str(form.get("device_id", "")).strip()
driver_type = str(form.get("driver_type", "")).strip()
name = str(form.get("name", "")).strip() or None
connection_info_raw = str(form.get("connection_info", "") or "{}")
error: str | None = None
connection_info: dict[str, Any] = {}
if not device_id or not driver_type:
error = "Device ID and driver type are required."
else:
try:
parsed = json.loads(connection_info_raw)
except ValueError:
error = "Connection info must be valid JSON."
else:
if not isinstance(parsed, dict):
error = "Connection info must be a JSON object."
else:
connection_info = parsed
if error is None:
try:
await asyncio.to_thread(
register_local_device,
config_store,
manager,
device_id=device_id,
driver_type=driver_type,
connection_info=connection_info,
name=name,
config=config,
enrollment_client=enrollment_client,
)
except ValueError as exc:
error = str(exc)
if error is not None:
devices = await asyncio.to_thread(config_store.list)
body = _devices_body(
devices=devices,
csrf_token=session.csrf_token,
edit_record=None,
error=error,
)
return HTMLResponse(
_chrome("Devices", body, session=session), status_code=400
)
return RedirectResponse(url="/devices", status_code=303)
@app.post("/devices/remove")
async def devices_remove(
request: Request,
session: SessionState = Depends(require_csrf),
) -> Response:
form = await request.form()
device_id = str(form.get("device_id", "")).strip()
if device_id:
await asyncio.to_thread(
unregister_local_device, config_store, manager, device_id=device_id
)
return RedirectResponse(url="/devices", status_code=303)
@app.get("/account", response_class=HTMLResponse)
async def account_page(
session: SessionState = Depends(require_session),
) -> HTMLResponse:
body = _account_body(csrf_token=session.csrf_token, message=None, error=None)
return HTMLResponse(_chrome("Account", body, session=session))
@app.post("/account", response_class=HTMLResponse)
async def account_submit(
request: Request,
session: SessionState = Depends(require_csrf),
) -> HTMLResponse:
form = await request.form()
current_password = str(form.get("current_password", ""))
new_password = str(form.get("new_password", ""))
confirm_password = str(form.get("confirm_password", ""))
if not new_password or new_password != confirm_password:
body = _account_body(
csrf_token=session.csrf_token,
message=None,
error="New password and confirmation must match.",
)
return HTMLResponse(
_chrome("Account", body, session=session), status_code=400
)
ok = await asyncio.to_thread(
change_password,
local_account_store,
current_password=current_password,
new_password=new_password,
)
if not ok:
body = _account_body(
csrf_token=session.csrf_token,
message=None,
error="Current password is incorrect.",
)
return HTMLResponse(
_chrome("Account", body, session=session), status_code=400
)
body = _account_body(
csrf_token=session.csrf_token,
message="Password updated.",
error=None,
)
return HTMLResponse(_chrome("Account", body, session=session))
@app.get("/history", response_class=HTMLResponse)
async def history_page(
session: SessionState = Depends(require_session),
) -> HTMLResponse:
entries = await asyncio.to_thread(history_store.list_recent)
body = _history_body(entries)
return HTMLResponse(_chrome("History", body, session=session))
return app
@@ -0,0 +1,89 @@
from __future__ import annotations
import hmac
import secrets
from collections.abc import Callable
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from threading import Lock
from host_agent.local_account import LocalAccountStore
SESSION_TOKEN_BYTES = 32
CSRF_TOKEN_BYTES = 32
@dataclass(frozen=True)
class SessionState:
username: str
csrf_token: str
expires_at: datetime
class SessionManager:
def __init__(
self,
*,
ttl_seconds: float,
now: Callable[[], datetime] | None = None,
) -> None:
self._ttl = timedelta(seconds=ttl_seconds)
self._now = now or (lambda: datetime.now(UTC))
self._lock = Lock()
self._sessions: dict[str, SessionState] = {}
def create_session(self, username: str) -> tuple[str, str]:
session_token = secrets.token_urlsafe(SESSION_TOKEN_BYTES)
csrf_token = secrets.token_urlsafe(CSRF_TOKEN_BYTES)
state = SessionState(
username=username,
csrf_token=csrf_token,
expires_at=self._now() + self._ttl,
)
with self._lock:
self._sessions[session_token] = state
return session_token, csrf_token
def validate(self, session_token: str) -> SessionState | None:
now = self._now()
with self._lock:
state = self._sessions.get(session_token)
if state is None:
return None
if state.expires_at <= now:
del self._sessions[session_token]
return None
renewed = SessionState(
username=state.username,
csrf_token=state.csrf_token,
expires_at=now + self._ttl,
)
self._sessions[session_token] = renewed
return renewed
def validate_csrf(self, session_token: str, csrf_token: str) -> bool:
state = self.validate(session_token)
if state is None:
return False
return hmac.compare_digest(state.csrf_token, csrf_token)
def invalidate(self, session_token: str) -> None:
with self._lock:
self._sessions.pop(session_token, None)
def attempt_login(store: LocalAccountStore, *, username: str, password: str) -> bool:
account = store.load()
if account is None or account.username != username:
return False
return store.verify(account, password)
def change_password(
store: LocalAccountStore, *, current_password: str, new_password: str
) -> bool:
account = store.load()
if account is None or not store.verify(account, current_password):
return False
store.create(account.username, new_password)
return True
+2
View File
@@ -6,7 +6,9 @@ requires-python = ">=3.14"
dependencies = [
"device-agent-runtime==0.1.0",
"device-cloud-platform==0.1.0",
"fastapi>=0.115.0",
"httpx>=0.27.0",
"uvicorn[standard]>=0.30.0",
]
[project.scripts]
+104 -1
View File
@@ -1,9 +1,12 @@
from __future__ import annotations
import asyncio
import socket
from contextlib import suppress
from datetime import UTC, datetime, timedelta
import httpx
from cloud.internal_api.models import (
AssignmentModel,
DeviceEnrollmentResponse,
@@ -16,6 +19,12 @@ from host_agent.identity import HostIdentityStore
from storage.device_config import DeviceConfigStore
def _free_loopback_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
probe.bind(("127.0.0.1", 0))
return probe.getsockname()[1]
def _config() -> HostAgentConfig:
return HostAgentConfig(
control_plane_url="https://control.example",
@@ -88,7 +97,6 @@ def test_create_application_enrolls_host_and_devices_before_managed_startup(
def __init__(self) -> None:
self.config = HostAgentConfig(
control_plane_url="https://control.example",
enrollment_token="one-time-token",
enrollment_managed=True,
)
@@ -352,3 +360,98 @@ def test_main_task_cancellation_waits_for_active_work_shutdown() -> None:
assert events == ["work-finished", "final-heartbeat", "closed"]
asyncio.run(scenario())
def test_console_enabled_serves_http_and_shuts_down_cleanly(tmp_path) -> None:
async def scenario() -> None:
port = _free_loopback_port()
config = HostAgentConfig(
control_plane_url="https://control.example",
host_id="host-a",
token="secret",
identity_path=tmp_path / "host_identity.json",
local_account_path=tmp_path / "host_local_account.json",
console_enabled=True,
console_bind_host="127.0.0.1",
console_port=port,
)
application = create_application(config=config, manager=DeviceManager())
assert application.console_server is not None
claim_started = asyncio.Event()
claim_cancelled = asyncio.Event()
events: list[str] = []
class BlockingClient:
async def claim(self):
claim_started.set()
try:
await asyncio.Event().wait()
except asyncio.CancelledError:
claim_cancelled.set()
raise
async def aclose(self):
events.append("closed")
class RecordingHeartbeat:
async def run(self, stop):
await stop.wait()
async def sync_once(self):
events.append("final-heartbeat")
class IdleProcessor:
async def process(self, assignment):
raise AssertionError("no assignment expected")
def request_stop(self):
events.append("stop-work")
application.client = BlockingClient() # type: ignore[assignment]
application.heartbeat = RecordingHeartbeat() # type: ignore[assignment]
application.processor = IdleProcessor() # type: ignore[assignment]
stop = asyncio.Event()
running = asyncio.create_task(application.run_async(stop))
await claim_started.wait()
response: httpx.Response | None = None
async with httpx.AsyncClient() as http_client:
loop = asyncio.get_running_loop()
deadline = loop.time() + 5
while loop.time() < deadline:
try:
response = await http_client.get(
f"http://127.0.0.1:{port}/login", timeout=0.5
)
except httpx.TransportError:
await asyncio.sleep(0.05)
continue
break
assert response is not None
assert response.status_code == 200
assert "Login" in response.text
stop.set()
await asyncio.wait_for(running, timeout=5)
assert claim_cancelled.is_set()
assert events == ["stop-work", "final-heartbeat", "closed"]
assert application.console_server.should_exit is True
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
probe.settimeout(0.5)
with suppress(ConnectionRefusedError, OSError):
probe.connect(("127.0.0.1", port))
raise AssertionError("console socket should be closed after shutdown")
asyncio.run(scenario())
def test_console_disabled_by_default_opens_no_socket(tmp_path, monkeypatch) -> None:
monkeypatch.chdir(tmp_path)
application = create_application(config=_config(), manager=DeviceManager())
assert application.console_server is None
asyncio.run(application.client.aclose())
+2 -37
View File
@@ -172,7 +172,7 @@ def test_result_report_retries_identical_payload_after_response_loss() -> None:
assert payloads[0]["failure_reason"] == "planner unavailable"
def test_bootstrap_client_retries_identical_enrollment_and_enrolls_device() -> None:
def test_bootstrap_client_directly_enrolls_and_enrolls_device() -> None:
requests: list[httpx.Request] = []
host_attempts = 0
@@ -189,7 +189,6 @@ def test_bootstrap_client_retries_identical_enrollment_and_enrolls_device() -> N
config = _config(
host_id="",
token="",
enrollment_token="one-time-token",
enrollment_managed=True,
)
with httpx.Client(
@@ -209,7 +208,6 @@ def test_bootstrap_client_retries_identical_enrollment_and_enrolls_device() -> N
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(
@@ -223,38 +221,5 @@ def test_bootstrap_client_retries_identical_enrollment_and_enrolls_device() -> N
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))
def test_self_service_enrollment_sends_no_authorization_header() -> None:
requests: list[httpx.Request] = []
def handler(request: httpx.Request) -> httpx.Response:
requests.append(request)
return httpx.Response(201, json={"host_id": "host-cloud-a"})
config = _config(
host_id="",
token="",
enrollment_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="operator",
)
assert host.host_id == "host-cloud-a"
assert len(requests) == 1
assert "authorization" not in requests[0].headers
assert requests[2].headers["authorization"] == ("Bearer host-token-" + ("x" * 40))
+109 -30
View File
@@ -11,24 +11,16 @@ from host_agent.config import (
)
BASE_ENV = {
"HOST_AGENT_HOST_ID": "host-a",
"HOST_AGENT_TOKEN": "secret",
}
def test_load_host_agent_config_uses_managed_cloud_default() -> None:
assert load_host_agent_config(BASE_ENV) == HostAgentConfig(
assert load_host_agent_config({}) == HostAgentConfig(
control_plane_url="https://amcp.home.jerryyan.top",
host_id="host-a",
token="secret",
enrollment_managed=True,
)
def test_load_host_agent_config_parses_poll_and_retry_values() -> None:
config = load_host_agent_config(
{
**BASE_ENV,
"HOST_AGENT_CONTROL_PLANE_URL": "https://cloud.example/v1/",
"HOST_AGENT_HEARTBEAT_INTERVAL_SECONDS": "10",
"HOST_AGENT_POLL_TIMEOUT_SECONDS": "15",
@@ -42,12 +34,11 @@ 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:
def test_load_host_agent_config_uses_direct_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",
}
@@ -55,30 +46,30 @@ def test_load_host_agent_config_supports_managed_enrollment(tmp_path) -> None:
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")
def test_static_credential_environment_variables_are_ignored() -> None:
config = load_host_agent_config(
{
"HOST_AGENT_HOST_ID": "legacy-host",
"HOST_AGENT_TOKEN": "legacy-token",
"HOST_AGENT_ENROLLMENT_TOKEN": "legacy-enrollment-token",
}
)
config = load_host_agent_config({"HOST_AGENT_IDENTITY_PATH": str(identity_path)})
assert config.identity_path == Path(identity_path)
assert config.host_id == ""
assert config.token == ""
assert config.enrollment_managed is True
def test_fresh_install_with_no_token_is_valid_and_defaults_local_account_path() -> None:
config = load_host_agent_config({"HOST_AGENT_CONTROL_PLANE_URL": "https://cloud.example"})
def test_fresh_install_defaults_local_account_path() -> None:
config = load_host_agent_config(
{"HOST_AGENT_CONTROL_PLANE_URL": "https://cloud.example"}
)
assert config.host_id == ""
assert config.enrollment_token == ""
assert config.enrollment_managed is True
assert config.local_account_path == Path("tasks/host_local_account.json")
@@ -98,9 +89,6 @@ def test_local_account_path_can_be_overridden(tmp_path) -> None:
@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"},
{
@@ -113,4 +101,95 @@ def test_load_host_agent_config_rejects_invalid_values(
overrides: dict[str, str],
) -> None:
with pytest.raises(HostAgentConfigurationError):
load_host_agent_config({**BASE_ENV, **overrides})
load_host_agent_config(overrides)
def test_console_defaults_are_disabled_and_do_not_trigger_validation() -> None:
config = load_host_agent_config({})
assert config.console_enabled is False
assert config.console_bind_host == "127.0.0.1"
assert config.console_port == 8765
assert config.console_allow_non_loopback is False
assert config.console_session_ttl_seconds == 43200.0
assert config.console_history_limit == 200
def test_console_enabled_with_default_loopback_bind_passes() -> None:
config = load_host_agent_config({"HOST_AGENT_CONSOLE_ENABLED": "true"})
assert config.console_enabled is True
assert config.console_bind_host == "127.0.0.1"
@pytest.mark.parametrize("bind_host", ["127.0.0.1", "localhost", "::1"])
def test_console_enabled_with_loopback_bind_host_passes(bind_host: str) -> None:
config = load_host_agent_config(
{
"HOST_AGENT_CONSOLE_ENABLED": "true",
"HOST_AGENT_CONSOLE_BIND_HOST": bind_host,
}
)
assert config.console_bind_host == bind_host
def test_console_enabled_with_non_loopback_bind_without_opt_in_raises() -> None:
with pytest.raises(HostAgentConfigurationError):
load_host_agent_config(
{
"HOST_AGENT_CONSOLE_ENABLED": "true",
"HOST_AGENT_CONSOLE_BIND_HOST": "0.0.0.0",
}
)
def test_console_enabled_with_non_loopback_bind_with_opt_in_succeeds() -> None:
config = load_host_agent_config(
{
"HOST_AGENT_CONSOLE_ENABLED": "true",
"HOST_AGENT_CONSOLE_BIND_HOST": "0.0.0.0",
"HOST_AGENT_CONSOLE_ALLOW_NON_LOOPBACK": "true",
}
)
assert config.console_bind_host == "0.0.0.0"
assert config.console_allow_non_loopback is True
def test_console_disabled_with_non_loopback_bind_does_not_raise() -> None:
config = load_host_agent_config({"HOST_AGENT_CONSOLE_BIND_HOST": "0.0.0.0"})
assert config.console_enabled is False
assert config.console_bind_host == "0.0.0.0"
def test_console_env_vars_parse_numeric_and_bool_fields() -> None:
config = load_host_agent_config(
{
"HOST_AGENT_CONSOLE_ENABLED": "1",
"HOST_AGENT_CONSOLE_PORT": "9001",
"HOST_AGENT_CONSOLE_SESSION_TTL_SECONDS": "3600",
"HOST_AGENT_CONSOLE_HISTORY_LIMIT": "50",
}
)
assert config.console_enabled is True
assert config.console_port == 9001
assert config.console_session_ttl_seconds == 3600
assert config.console_history_limit == 50
@pytest.mark.parametrize(
"overrides",
[
{"HOST_AGENT_CONSOLE_PORT": "0"},
{"HOST_AGENT_CONSOLE_SESSION_TTL_SECONDS": "-1"},
{"HOST_AGENT_CONSOLE_HISTORY_LIMIT": "0"},
],
)
def test_console_numeric_fields_reject_invalid_values(
overrides: dict[str, str],
) -> None:
with pytest.raises(HostAgentConfigurationError):
load_host_agent_config(overrides)
@@ -0,0 +1,171 @@
from __future__ import annotations
from cloud.internal_api.models import DeviceEnrollmentResponse
from device.manager import DeviceManager
from host_agent.config import HostAgentConfig
from host_agent.devices import register_local_device, unregister_local_device
from storage.device_config import DeviceConfigStore
def _config(*, enrollment_managed: bool) -> HostAgentConfig:
return HostAgentConfig(
control_plane_url="https://control.example",
host_id="host-a",
token="secret",
enrollment_managed=enrollment_managed,
)
class RecordingEnrollmentClient:
def __init__(self, device_id: str) -> None:
self.device_id = device_id
self.calls: list[dict[str, object]] = []
def enroll_device(self, **payload):
self.calls.append(payload)
return DeviceEnrollmentResponse(device_id=self.device_id)
def test_register_local_device_enrollment_managed_registers_under_cloud_id(
tmp_path,
) -> None:
store = DeviceConfigStore(tmp_path / "devices.sqlite3")
manager = DeviceManager()
enrollment_client = RecordingEnrollmentClient("device-cloud-a")
register_local_device(
store,
manager,
device_id="local-device-a",
driver_type="wda",
connection_info={"server_url": "http://127.0.0.1:4723"},
name="Lab iPhone",
config=_config(enrollment_managed=True),
enrollment_client=enrollment_client, # type: ignore[arg-type]
)
assert enrollment_client.calls == [
{
"local_device_id": "local-device-a",
"driver_type": "wda",
"name": "Lab iPhone",
"capability_tags": [],
}
]
assert [device.id for device in manager.list_devices()] == ["device-cloud-a"]
assert store.get("local-device-a")["cloud_device_id"] == "device-cloud-a"
def test_register_local_device_not_enrollment_managed_registers_under_local_id(
tmp_path,
) -> None:
store = DeviceConfigStore(tmp_path / "devices.sqlite3")
manager = DeviceManager()
enrollment_client = RecordingEnrollmentClient("device-cloud-a")
register_local_device(
store,
manager,
device_id="local-device-a",
driver_type="wda",
connection_info={"server_url": "http://127.0.0.1:4723"},
name="Lab iPhone",
config=_config(enrollment_managed=False),
enrollment_client=enrollment_client, # type: ignore[arg-type]
)
assert enrollment_client.calls == []
assert [device.id for device in manager.list_devices()] == ["local-device-a"]
assert store.get("local-device-a")["cloud_device_id"] is None
def test_register_local_device_reregister_under_new_cloud_id_replaces_prior_entry(
tmp_path,
) -> None:
store = DeviceConfigStore(tmp_path / "devices.sqlite3")
manager = DeviceManager()
config = _config(enrollment_managed=True)
register_local_device(
store,
manager,
device_id="local-device-a",
driver_type="wda",
connection_info={"server_url": "http://127.0.0.1:4723"},
name="Lab iPhone",
config=config,
enrollment_client=RecordingEnrollmentClient( # type: ignore[arg-type]
"device-cloud-a"
),
)
assert [device.id for device in manager.list_devices()] == ["device-cloud-a"]
register_local_device(
store,
manager,
device_id="local-device-a",
driver_type="wda",
connection_info={"server_url": "http://127.0.0.1:5000"},
name="Lab iPhone (moved)",
config=config,
enrollment_client=RecordingEnrollmentClient( # type: ignore[arg-type]
"device-cloud-b"
),
)
devices = manager.list_devices()
assert [device.id for device in devices] == ["device-cloud-b"]
assert devices[0].name == "Lab iPhone (moved)"
assert devices[0].connection_info == {"server_url": "http://127.0.0.1:5000"}
assert store.get("local-device-a")["cloud_device_id"] == "device-cloud-b"
def test_unregister_local_device_removes_device_with_cloud_id(tmp_path) -> None:
store = DeviceConfigStore(tmp_path / "devices.sqlite3")
manager = DeviceManager()
register_local_device(
store,
manager,
device_id="local-device-a",
driver_type="wda",
connection_info={},
name=None,
config=_config(enrollment_managed=True),
enrollment_client=RecordingEnrollmentClient( # type: ignore[arg-type]
"device-cloud-a"
),
)
unregister_local_device(store, manager, device_id="local-device-a")
assert manager.list_devices() == []
assert store.get("local-device-a") is None
def test_unregister_local_device_removes_device_without_cloud_id(tmp_path) -> None:
store = DeviceConfigStore(tmp_path / "devices.sqlite3")
manager = DeviceManager()
register_local_device(
store,
manager,
device_id="local-device-a",
driver_type="wda",
connection_info={},
name=None,
config=_config(enrollment_managed=False),
enrollment_client=None,
)
unregister_local_device(store, manager, device_id="local-device-a")
assert manager.list_devices() == []
assert store.get("local-device-a") is None
def test_unregister_local_device_is_a_no_op_for_unknown_device(tmp_path) -> None:
store = DeviceConfigStore(tmp_path / "devices.sqlite3")
manager = DeviceManager()
unregister_local_device(store, manager, device_id="unknown-device")
assert manager.list_devices() == []
+58 -31
View File
@@ -10,12 +10,12 @@ import httpx
import pytest
from fastapi.testclient import TestClient
from cloud.auth import BearerCredential
from cloud.auth import digest_token
from cloud.control_config import CloudControlConfig
from cloud.sdk.client import CloudClient
from cloud.scheduler import TaskConstraints
from cloud_api.app import create_app
from core.models import Scene
from core.models import Scene, utc_now
from device.manager import DeviceManager
from driver.base import Driver
from host_agent.assignment import AssignmentExecutionResult, AssignmentExecutor
@@ -79,23 +79,6 @@ class FakeDriver(Driver):
return None
def _credential(host_id: str) -> BearerCredential:
return BearerCredential(
principal_id=f"agent-{host_id}",
token=f"token-{host_id}",
scopes=frozenset(),
host_id=host_id,
)
def _public_credential() -> BearerCredential:
return BearerCredential(
principal_id="sdk",
token="public-token",
scopes=frozenset({"tasks:submit", "tasks:read"}),
)
def _config(host_id: str) -> HostAgentConfig:
return HostAgentConfig(
control_plane_url="http://control.test",
@@ -120,10 +103,18 @@ async def _control_plane(
scheduler_interval_seconds=60,
lease_reaper_interval_seconds=60,
lease_duration_seconds=lease_duration_seconds,
credentials=tuple(_credential(host_id) for host_id in host_ids),
)
)
async with app.router.lifespan_context(app):
for host_id in host_ids:
app.state.cloud_services.repository.enroll_host(
host_id=host_id,
agent_instance_id=f"agent-{host_id}",
credential_digest=digest_token(f"token-{host_id}"),
enrollment_token_digest=None,
display_name=host_id,
enrolled_at=utc_now(),
)
yield app
@@ -162,12 +153,22 @@ class _RecordingTransport(httpx.AsyncBaseTransport):
async def _sync_fake_device(
app,
client: HostAgentClient,
host_id: str,
device_id: str,
*,
driver_type: str = "wda",
) -> FakeDriver:
app.state.cloud_services.repository.enroll_device(
device_id=device_id,
host_id=host_id,
local_device_id=device_id,
driver_type=driver_type,
name=device_id,
capability_tags=[],
enrolled_at=utc_now(),
)
driver = FakeDriver()
manager = DeviceManager()
manager.register_device(
@@ -248,7 +249,7 @@ def test_one_host_executes_assignment_through_outbound_protocol(tmp_path) -> Non
paths: list[str] = []
async with _control_plane(tmp_path / "one-host.sqlite3", "host-a") as app:
async with _host_client(app, "host-a", request_paths=paths) as client:
await _sync_fake_device(client, "host-a", "device-a")
await _sync_fake_device(app, client, "host-a", "device-a")
task_id = app.state.cloud_services.scheduler.submit(
goal="open settings"
)
@@ -275,7 +276,7 @@ def test_nat_style_host_requires_only_outbound_requests(tmp_path) -> None:
paths: list[str] = []
async with _control_plane(tmp_path / "outbound-only.sqlite3", "host-a") as app:
async with _host_client(app, "host-a", request_paths=paths) as client:
await _sync_fake_device(client, "host-a", "device-a")
await _sync_fake_device(app, client, "host-a", "device-a")
assert await client.claim() is None
assert paths == [
@@ -297,8 +298,9 @@ def test_multiple_hosts_claim_only_their_matching_devices(tmp_path) -> None:
_host_client(app, "host-a") as client_a,
_host_client(app, "host-b") as client_b,
):
await _sync_fake_device(client_a, "host-a", "device-a")
await _sync_fake_device(app, client_a, "host-a", "device-a")
await _sync_fake_device(
app,
client_b,
"host-b",
"device-b",
@@ -331,7 +333,7 @@ def test_control_plane_restart_preserves_dispatched_assignment(tmp_path) -> None
assignment = None
async with _control_plane(database_path, "host-a") as first_app:
async with _host_client(first_app, "host-a") as client:
await _sync_fake_device(client, "host-a", "device-a")
await _sync_fake_device(first_app, client, "host-a", "device-a")
task_id = first_app.state.cloud_services.scheduler.submit(goal="resume")
first_app.state.cloud_services.scheduler.assign()
assignment = await client.claim()
@@ -351,7 +353,7 @@ def test_host_agent_restart_reuses_active_lease(tmp_path) -> None:
async def scenario() -> None:
async with _control_plane(tmp_path / "agent-restart.sqlite3", "host-a") as app:
async with _host_client(app, "host-a") as first_client:
await _sync_fake_device(first_client, "host-a", "device-a")
await _sync_fake_device(app, first_client, "host-a", "device-a")
task_id = app.state.cloud_services.scheduler.submit(goal="resume host")
app.state.cloud_services.scheduler.assign()
assignment = await first_client.claim()
@@ -376,7 +378,7 @@ def test_lease_loss_rejects_stale_host_result(tmp_path) -> None:
lease_duration_seconds=0.2,
) as app:
async with _host_client(app, "host-a") as client:
await _sync_fake_device(client, "host-a", "device-a")
await _sync_fake_device(app, client, "host-a", "device-a")
task_id = app.state.cloud_services.scheduler.submit(goal="expire")
app.state.cloud_services.scheduler.assign()
assignment = await client.claim()
@@ -405,7 +407,6 @@ def test_public_sdk_reports_fake_device_success_and_runtime_failure(tmp_path) ->
scheduler_interval_seconds=60,
lease_reaper_interval_seconds=60,
lease_duration_seconds=30,
credentials=(_public_credential(), _credential("host-a")),
)
)
driver = FakeDriver()
@@ -413,10 +414,38 @@ def test_public_sdk_reports_fake_device_success_and_runtime_failure(tmp_path) ->
manager.register_device("device-a", lambda: driver, status="idle")
with TestClient(app) as http_client:
app.state.cloud_services.user_auth_service.create_user(
username="operator",
display_name="Operator",
role="admin",
password="correct-horse-battery-staple",
must_change_password=False,
)
login = http_client.post(
"/v1/auth/login",
json={"username": "operator", "password": "correct-horse-battery-staple"},
)
assert login.status_code == 200
app.state.cloud_services.repository.enroll_host(
host_id="host-a",
agent_instance_id="agent-host-a",
credential_digest=digest_token("token-host-a"),
enrollment_token_digest=None,
display_name="host-a",
enrolled_at=utc_now(),
)
app.state.cloud_services.repository.enroll_device(
device_id="device-a",
host_id="host-a",
local_device_id="device-a",
driver_type="wda",
name="device-a",
capability_tags=[],
enrolled_at=utc_now(),
)
cloud_client = CloudClient(
"http://testserver",
http_client=http_client,
token="public-token",
)
async def scenario() -> None:
@@ -428,9 +457,7 @@ def test_public_sdk_reports_fake_device_success_and_runtime_failure(tmp_path) ->
heartbeat.connect_devices()
await heartbeat.sync_once()
successful_task_id = cloud_client.submit_task(goal="tap screen")[
"task_id"
]
successful_task_id = cloud_client.submit_task(goal="tap screen")["task_id"]
app.state.cloud_services.scheduler.assign()
successful_assignment = await host_client.claim()
assert successful_assignment is not None
@@ -31,12 +31,12 @@ class _RejectingEnrollmentClient:
raise HostAgentAPIError(401, "unauthorized")
def test_fresh_install_with_no_token_self_enrolls(tmp_path) -> None:
def test_fresh_install_directly_enrolls(tmp_path) -> None:
identity_store = HostIdentityStore(tmp_path / "identity.json")
client = _RecordingEnrollmentClient()
resolved = resolve_host_identity(
_config(enrollment_token=""),
_config(),
identity_store=identity_store,
client=client,
)
@@ -52,7 +52,7 @@ def test_self_service_rejection_propagates_as_api_error(tmp_path) -> None:
try:
resolve_host_identity(
_config(enrollment_token=""),
_config(),
identity_store=identity_store,
client=client,
)
@@ -63,19 +63,6 @@ def test_self_service_rejection_propagates_as_api_error(tmp_path) -> None:
assert identity_store.load().host_id is None
def test_configured_enrollment_token_still_used_when_present(tmp_path) -> None:
identity_store = HostIdentityStore(tmp_path / "identity.json")
client = _RecordingEnrollmentClient()
resolve_host_identity(
_config(enrollment_token="one-time-token"),
identity_store=identity_store,
client=client,
)
assert client.calls[0]["agent_instance_id"]
def test_existing_cached_identity_skips_enrollment(tmp_path) -> None:
identity_store = HostIdentityStore(tmp_path / "identity.json")
identity_store.complete(identity_store.load_or_create(), "host-cloud-a")
@@ -85,7 +72,7 @@ def test_existing_cached_identity_skips_enrollment(tmp_path) -> None:
raise AssertionError("cached identity must skip enrollment")
resolved = resolve_host_identity(
_config(enrollment_token=""),
_config(),
identity_store=identity_store,
client=ExplodingClient(), # type: ignore[arg-type]
)
@@ -7,6 +7,7 @@ from cloud.internal_api.models import HeartbeatResponse
from device.manager import DeviceManager
from host_agent.config import HostAgentConfig
from host_agent.heartbeat import HeartbeatSynchronizer, build_device_snapshot
from host_agent.status import AgentStatusTracker
class ConnectableDriver:
@@ -81,3 +82,44 @@ def test_heartbeat_synchronizer_runs_at_configured_interval_until_stopped() -> N
asyncio.run(scenario())
assert calls == [["device-a"], ["device-a"], ["device-a"]]
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):
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())
@@ -0,0 +1,69 @@
from __future__ import annotations
from datetime import UTC, datetime
from host_agent.history import ConsoleHistoryStore
def test_history_store_records_assignment_and_heartbeat_newest_first(tmp_path) -> None:
store = ConsoleHistoryStore(tmp_path / "history.sqlite3")
store.record_assignment(
task_id="task-a",
attempt=1,
status="done",
failure_reason=None,
device_id="device-a",
)
store.record_heartbeat(device_count=2)
entries = store.list_recent()
assert len(entries) == 2
assert entries[0] == {
"kind": "heartbeat",
"occurred_at": entries[0]["occurred_at"],
"summary": "heartbeat: 2 devices",
"detail": {"device_count": 2},
}
assert entries[1] == {
"kind": "assignment",
"occurred_at": entries[1]["occurred_at"],
"summary": "task-a attempt 1 on device-a: done",
"detail": {
"task_id": "task-a",
"attempt": 1,
"status": "done",
"failure_reason": None,
"device_id": "device-a",
},
}
def test_history_store_prunes_oldest_entries_beyond_limit(tmp_path) -> None:
store = ConsoleHistoryStore(tmp_path / "history.sqlite3", limit=3)
for index in range(5):
store.record_heartbeat(device_count=index)
entries = store.list_recent()
assert len(entries) == 3
assert [entry["detail"]["device_count"] for entry in entries] == [4, 3, 2]
def test_history_store_returns_empty_list_when_no_entries(tmp_path) -> None:
store = ConsoleHistoryStore(tmp_path / "history.sqlite3")
assert store.list_recent() == []
def test_history_store_uses_injected_now_for_occurred_at(tmp_path) -> None:
fixed_now = datetime(2026, 1, 1, 12, 0, 0, tzinfo=UTC)
store = ConsoleHistoryStore(tmp_path / "history.sqlite3", now=lambda: fixed_now)
store.record_heartbeat(device_count=1)
entries = store.list_recent()
assert entries[0]["occurred_at"] == fixed_now.isoformat()
@@ -6,6 +6,7 @@ from datetime import UTC, datetime
from cloud.internal_api.models import AssignmentModel, TerminalResultResponse
from host_agent.assignment import AssignmentExecutionResult
from host_agent.processor import AssignmentProcessor
from host_agent.status import AgentStatusTracker
def _assignment() -> AssignmentModel:
@@ -85,3 +86,69 @@ def test_processor_preserves_runtime_failure_reason() -> None:
}
asyncio.run(scenario())
def test_status_tracker_sees_started_then_finished_even_on_raise() -> None:
async def scenario() -> None:
tracker = AgentStatusTracker()
snapshots: list[dict[str, object]] = []
class RaisingExecutor:
async def run(self, assignment):
snapshots.append(tracker.snapshot())
raise RuntimeError("executor exploded")
class RecordingClient:
async def report_result(self, assignment, **kwargs):
return TerminalResultResponse(status="recorded")
processor = AssignmentProcessor(
RecordingClient(), # type: ignore[arg-type]
RaisingExecutor(),
status_tracker=tracker,
)
try:
await processor.process(_assignment())
except RuntimeError:
pass
assert snapshots[0]["current_assignment"] is not None
assert snapshots[0]["current_assignment"]["task_id"] == "task-a"
assert tracker.snapshot()["current_assignment"] is None
asyncio.run(scenario())
def test_on_result_receives_assignment_and_result_and_swallows_exceptions() -> None:
async def scenario() -> None:
received: list[tuple[object, object]] = []
class SuccessfulExecutor:
async def run(self, assignment):
return AssignmentExecutionResult(
status="done",
failure_reason=None,
metadata={},
)
class RecordingClient:
async def report_result(self, assignment, **kwargs):
return TerminalResultResponse(status="recorded")
def on_result(assignment, result) -> None:
received.append((assignment, result))
raise RuntimeError("history recording exploded")
assignment = _assignment()
processor = AssignmentProcessor(
RecordingClient(), # type: ignore[arg-type]
SuccessfulExecutor(),
on_result=on_result,
)
result = await processor.process(assignment)
assert received == [(assignment, result)]
asyncio.run(scenario())
@@ -0,0 +1,66 @@
from __future__ import annotations
from datetime import UTC, datetime
from cloud.internal_api.models import AssignmentModel
from host_agent.status import AgentStatusTracker
def _assignment() -> AssignmentModel:
return AssignmentModel(
task_id="task-a",
attempt=1,
lease_id="lease-a",
lease_expires_at=datetime(2026, 7, 12, tzinfo=UTC),
host_id="host-a",
device_id="device-a",
goal="open settings",
workflow_definition_id="workflow-a",
)
def test_mark_assignment_started_reflected_in_snapshot() -> None:
ticks = iter([datetime(2026, 7, 13, 9, 0, 0, tzinfo=UTC)])
tracker = AgentStatusTracker(now=lambda: next(ticks))
tracker.mark_assignment_started(_assignment())
snapshot = tracker.snapshot()
assert snapshot["current_assignment"] == {
"task_id": "task-a",
"device_id": "device-a",
"goal": "open settings",
"workflow_definition_id": "workflow-a",
"started_at": "2026-07-13T09:00:00+00:00",
}
def test_mark_assignment_finished_clears_current_assignment() -> None:
tracker = AgentStatusTracker(now=lambda: datetime(2026, 7, 13, 9, 0, 0, tzinfo=UTC))
tracker.mark_assignment_started(_assignment())
tracker.mark_assignment_finished()
assert tracker.snapshot()["current_assignment"] is None
def test_mark_heartbeat_reflected_in_snapshot() -> None:
tracker = AgentStatusTracker(now=lambda: datetime(2026, 7, 13, 9, 5, 0, tzinfo=UTC))
tracker.mark_heartbeat(ok=True, device_count=3)
snapshot = tracker.snapshot()
assert snapshot["last_heartbeat"] == {
"ok": True,
"device_count": 3,
"at": "2026-07-13T09:05:00+00:00",
}
def test_snapshot_defaults_to_no_assignment_or_heartbeat() -> None:
tracker = AgentStatusTracker(now=lambda: datetime(2026, 7, 13, 9, 0, 0, tzinfo=UTC))
snapshot = tracker.snapshot()
assert snapshot["current_assignment"] is None
assert snapshot["last_heartbeat"] is None
@@ -0,0 +1,262 @@
from __future__ import annotations
import re
from fastapi.testclient import TestClient
from device.manager import DeviceManager
from host_agent.config import HostAgentConfig
from host_agent.history import ConsoleHistoryStore
from host_agent.identity import HostIdentityStore
from host_agent.local_account import LocalAccountStore
from host_agent.status import AgentStatusTracker
from host_agent.web.app import SESSION_COOKIE_NAME, create_console_app
from host_agent.web.auth import SessionManager
from storage.device_config import DeviceConfigStore
CSRF_PATTERN = re.compile(r'name="csrf_token" value="([^"]+)"')
def _build_client(tmp_path, *, create_account: bool = True) -> tuple[TestClient, dict]:
config = HostAgentConfig(
control_plane_url="https://control.example",
host_id="host-a",
token="secret",
enrollment_managed=False,
console_session_ttl_seconds=3600.0,
)
manager = DeviceManager()
config_store = DeviceConfigStore(tmp_path / "devices.sqlite3")
local_account_store = LocalAccountStore(tmp_path / "host_local_account.json")
if create_account:
local_account_store.create("operator", "correct horse battery staple")
identity_store = HostIdentityStore(tmp_path / "host_identity.json")
history_store = ConsoleHistoryStore(tmp_path / "history.sqlite3")
status_tracker = AgentStatusTracker()
session_manager = SessionManager(ttl_seconds=3600.0)
app = create_console_app(
config=config,
manager=manager,
config_store=config_store,
local_account_store=local_account_store,
identity_store=identity_store,
history_store=history_store,
status_tracker=status_tracker,
session_manager=session_manager,
enrollment_client=None,
)
client = TestClient(app)
context = {
"manager": manager,
"config_store": config_store,
"local_account_store": local_account_store,
"history_store": history_store,
"session_manager": session_manager,
}
return client, context
def _login(
client: TestClient,
*,
username: str = "operator",
password: str = "correct horse battery staple",
) -> str:
response = client.post("/login", data={"username": username, "password": password})
assert response.status_code == 200
match = CSRF_PATTERN.search(response.text)
assert match is not None
return match.group(1)
def test_unauthenticated_get_root_redirects_to_login(tmp_path) -> None:
client, _ = _build_client(tmp_path)
response = client.get("/", follow_redirects=False)
assert response.status_code == 303
assert response.headers["location"] == "/login"
def test_login_with_no_account_shows_setup_message_and_rejects_post(tmp_path) -> None:
client, context = _build_client(tmp_path, create_account=False)
get_response = client.get("/login")
assert "device-host-agent setup" in get_response.text
post_response = client.post(
"/login", data={"username": "operator", "password": "anything"}
)
assert "device-host-agent setup" in post_response.text
assert SESSION_COOKIE_NAME not in client.cookies
def test_login_with_wrong_password_fails_and_sets_no_cookie(tmp_path) -> None:
client, _ = _build_client(tmp_path)
response = client.post("/login", data={"username": "operator", "password": "wrong"})
assert response.status_code == 200
assert "Invalid username or password" in response.text
assert SESSION_COOKIE_NAME not in client.cookies
def test_login_with_correct_password_sets_cookie_and_dashboard_succeeds(
tmp_path,
) -> None:
client, _ = _build_client(tmp_path)
response = client.post(
"/login",
data={"username": "operator", "password": "correct horse battery staple"},
)
assert response.status_code == 200
assert SESSION_COOKIE_NAME in client.cookies
assert "Status" in response.text
assert "control.example" in response.text
def test_mutating_post_without_csrf_token_is_rejected_and_makes_no_change(
tmp_path,
) -> None:
client, context = _build_client(tmp_path)
_login(client)
response = client.post(
"/devices/save",
data={
"device_id": "device-a",
"driver_type": "wda",
"name": "Lab iPhone",
"connection_info": "{}",
},
)
assert response.status_code == 403
assert context["config_store"].get("device-a") is None
assert context["manager"].list_devices() == []
def test_add_device_appears_in_devices_page_and_manager(tmp_path) -> None:
client, context = _build_client(tmp_path)
csrf_token = _login(client)
response = client.post(
"/devices/save",
data={
"device_id": "device-a",
"driver_type": "wda",
"name": "Lab iPhone",
"connection_info": "{}",
"csrf_token": csrf_token,
},
)
assert response.status_code == 200
assert "device-a" in response.text
assert context["config_store"].get("device-a") is not None
assert [device.id for device in context["manager"].list_devices()] == ["device-a"]
def test_remove_device_unregisters_from_manager(tmp_path) -> None:
client, context = _build_client(tmp_path)
csrf_token = _login(client)
client.post(
"/devices/save",
data={
"device_id": "device-a",
"driver_type": "wda",
"name": "Lab iPhone",
"connection_info": "{}",
"csrf_token": csrf_token,
},
)
assert context["config_store"].get("device-a") is not None
response = client.post(
"/devices/remove",
data={"device_id": "device-a", "csrf_token": csrf_token},
)
assert response.status_code == 200
assert context["config_store"].get("device-a") is None
assert context["manager"].list_devices() == []
def test_change_password_wrong_current_fails_old_password_still_works(tmp_path) -> None:
client, context = _build_client(tmp_path)
csrf_token = _login(client)
response = client.post(
"/account",
data={
"current_password": "wrong",
"new_password": "new password",
"confirm_password": "new password",
"csrf_token": csrf_token,
},
)
assert response.status_code == 400
assert "incorrect" in response.text
account_store: LocalAccountStore = context["local_account_store"]
account = account_store.load()
assert account is not None
assert account_store.verify(account, "correct horse battery staple") is True
def test_change_password_correct_succeeds_old_password_no_longer_works(
tmp_path,
) -> None:
client, context = _build_client(tmp_path)
csrf_token = _login(client)
response = client.post(
"/account",
data={
"current_password": "correct horse battery staple",
"new_password": "new password",
"confirm_password": "new password",
"csrf_token": csrf_token,
},
)
assert response.status_code == 200
assert "Password updated" in response.text
account_store: LocalAccountStore = context["local_account_store"]
account = account_store.load()
assert account is not None
assert account_store.verify(account, "new password") is True
assert account_store.verify(account, "correct horse battery staple") is False
def test_history_reflects_entries_written_beforehand(tmp_path) -> None:
client, context = _build_client(tmp_path)
context["history_store"].record_heartbeat(device_count=3)
_login(client)
response = client.get("/history")
assert response.status_code == 200
assert "heartbeat: 3 devices" in response.text
def test_logout_invalidates_session_so_subsequent_request_redirects_to_login(
tmp_path,
) -> None:
client, _ = _build_client(tmp_path)
csrf_token = _login(client)
logout_response = client.post(
"/logout", data={"csrf_token": csrf_token}, follow_redirects=False
)
assert logout_response.status_code == 303
assert logout_response.headers["location"] == "/login"
response = client.get("/", follow_redirects=False)
assert response.status_code == 303
assert response.headers["location"] == "/login"
@@ -0,0 +1,162 @@
from __future__ import annotations
from datetime import UTC, datetime, timedelta
from host_agent.local_account import LocalAccountStore
from host_agent.web.auth import SessionManager, attempt_login, change_password
class FakeClock:
def __init__(self, start: datetime) -> None:
self.current = start
def __call__(self) -> datetime:
return self.current
def advance(self, seconds: float) -> None:
self.current += timedelta(seconds=seconds)
def test_create_session_then_validate_returns_username() -> None:
clock = FakeClock(datetime(2026, 1, 1, tzinfo=UTC))
manager = SessionManager(ttl_seconds=60, now=clock)
session_token, csrf_token = manager.create_session("operator")
state = manager.validate(session_token)
assert state is not None
assert state.username == "operator"
assert state.csrf_token == csrf_token
def test_validate_unknown_token_returns_none() -> None:
manager = SessionManager(ttl_seconds=60)
assert manager.validate("does-not-exist") is None
def test_validate_after_ttl_elapsed_returns_none() -> None:
clock = FakeClock(datetime(2026, 1, 1, tzinfo=UTC))
manager = SessionManager(ttl_seconds=60, now=clock)
session_token, _ = manager.create_session("operator")
clock.advance(61)
assert manager.validate(session_token) is None
def test_validate_before_expiry_slides_expiry_forward() -> None:
clock = FakeClock(datetime(2026, 1, 1, tzinfo=UTC))
manager = SessionManager(ttl_seconds=60, now=clock)
session_token, _ = manager.create_session("operator")
clock.advance(30)
first = manager.validate(session_token)
assert first is not None
clock.advance(30)
second = manager.validate(session_token)
assert second is not None
assert second.expires_at > first.expires_at
def test_validate_csrf_true_for_right_token() -> None:
manager = SessionManager(ttl_seconds=60)
session_token, csrf_token = manager.create_session("operator")
assert manager.validate_csrf(session_token, csrf_token) is True
def test_validate_csrf_false_for_wrong_token() -> None:
manager = SessionManager(ttl_seconds=60)
session_token, _ = manager.create_session("operator")
assert manager.validate_csrf(session_token, "wrong-token") is False
def test_validate_csrf_false_for_invalid_session() -> None:
manager = SessionManager(ttl_seconds=60)
assert manager.validate_csrf("does-not-exist", "anything") is False
def test_invalidate_makes_subsequent_validate_return_none() -> None:
manager = SessionManager(ttl_seconds=60)
session_token, _ = manager.create_session("operator")
manager.invalidate(session_token)
assert manager.validate(session_token) is None
def test_attempt_login_true_for_correct_credentials(tmp_path) -> None:
store = LocalAccountStore(tmp_path / "host_local_account.json")
store.create("operator", "correct horse battery staple")
assert (
attempt_login(
store, username="operator", password="correct horse battery staple"
)
is True
)
def test_attempt_login_false_for_wrong_password(tmp_path) -> None:
store = LocalAccountStore(tmp_path / "host_local_account.json")
store.create("operator", "correct horse battery staple")
assert attempt_login(store, username="operator", password="wrong") is False
def test_attempt_login_false_when_no_account_exists(tmp_path) -> None:
store = LocalAccountStore(tmp_path / "host_local_account.json")
assert attempt_login(store, username="operator", password="anything") is False
def test_attempt_login_false_for_wrong_username(tmp_path) -> None:
store = LocalAccountStore(tmp_path / "host_local_account.json")
store.create("operator", "correct horse battery staple")
assert (
attempt_login(
store, username="someone-else", password="correct horse battery staple"
)
is False
)
def test_change_password_succeeds_and_rotates_credential(tmp_path) -> None:
store = LocalAccountStore(tmp_path / "host_local_account.json")
store.create("operator", "old password")
assert (
change_password(
store, current_password="old password", new_password="new password"
)
is True
)
assert attempt_login(store, username="operator", password="new password") is True
assert attempt_login(store, username="operator", password="old password") is False
def test_change_password_fails_with_wrong_current_password(tmp_path) -> None:
store = LocalAccountStore(tmp_path / "host_local_account.json")
store.create("operator", "old password")
assert (
change_password(store, current_password="wrong", new_password="new password")
is False
)
assert attempt_login(store, username="operator", password="old password") is True
assert attempt_login(store, username="operator", password="new password") is False
def test_change_password_fails_when_no_account_exists(tmp_path) -> None:
store = LocalAccountStore(tmp_path / "host_local_account.json")
assert (
change_password(store, current_password="anything", new_password="new password")
is False
)