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
+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
)