diff --git a/apps/device-host-agent/host_agent/app.py b/apps/device-host-agent/host_agent/app.py index fef486e..cc865bc 100644 --- a/apps/device-host-agent/host_agent/app.py +++ b/apps/device-host-agent/host_agent/app.py @@ -1,15 +1,108 @@ from __future__ import annotations +import asyncio +from contextlib import suppress from dataclasses import dataclass +from cloud.internal_api.models import AssignmentModel +from device.manager import DeviceManager +from host_agent.assignment import AssignmentExecutor +from host_agent.client import HostAgentClient +from host_agent.config import HostAgentConfig, load_host_agent_config +from host_agent.execution import create_execution_factories +from host_agent.heartbeat import HeartbeatSynchronizer +from host_agent.lease import ActiveAssignmentRunner +from host_agent.processor import AssignmentProcessingResult, AssignmentProcessor + @dataclass class HostAgentApplication: - """Process shell expanded by the Host Agent implementation tasks.""" + client: HostAgentClient + heartbeat: HeartbeatSynchronizer + processor: AssignmentProcessor def run(self) -> None: - return None + asyncio.run(self.run_async()) + + async def run_async(self, stop: asyncio.Event | None = None) -> None: + stop_requested = stop or asyncio.Event() + heartbeat_stop = asyncio.Event() + heartbeat_task = asyncio.create_task(self.heartbeat.run(heartbeat_stop)) + active_processing: asyncio.Task[AssignmentProcessingResult] | None = None + try: + while not stop_requested.is_set(): + assignment = await self._claim_until_stopped(stop_requested) + if assignment is None: + continue + active_processing = asyncio.create_task( + self.processor.process(assignment) + ) + stopped = asyncio.create_task(stop_requested.wait()) + done, _ = await asyncio.wait( + {active_processing, stopped}, + return_when=asyncio.FIRST_COMPLETED, + ) + if stopped in done: + self.processor.request_stop() + else: + stopped.cancel() + with suppress(asyncio.CancelledError): + await stopped + await asyncio.shield(active_processing) + active_processing = None + finally: + self.processor.request_stop() + if active_processing is not None: + with suppress(Exception): + await asyncio.shield(active_processing) + heartbeat_stop.set() + try: + await asyncio.gather(heartbeat_task, return_exceptions=True) + with suppress(Exception): + await self.heartbeat.sync_once() + finally: + await self.client.aclose() + + async def _claim_until_stopped( + self, + stop: asyncio.Event, + ) -> AssignmentModel | None: + claim = asyncio.create_task(self.client.claim()) + stopped = asyncio.create_task(stop.wait()) + try: + done, _ = await asyncio.wait( + {claim, stopped}, + return_when=asyncio.FIRST_COMPLETED, + ) + except asyncio.CancelledError: + claim.cancel() + stopped.cancel() + await asyncio.gather(claim, stopped, return_exceptions=True) + raise + if stopped in done: + claim.cancel() + with suppress(asyncio.CancelledError): + await claim + return None + stopped.cancel() + with suppress(asyncio.CancelledError): + await stopped + return await claim -def create_application() -> HostAgentApplication: - return HostAgentApplication() +def create_application( + *, + config: HostAgentConfig | None = None, + manager: DeviceManager | None = None, +) -> HostAgentApplication: + resolved_config = config or load_host_agent_config() + resolved_manager = manager or DeviceManager() + client = HostAgentClient(resolved_config) + heartbeat = HeartbeatSynchronizer(resolved_manager, client, resolved_config) + executor = AssignmentExecutor(create_execution_factories(resolved_manager)) + active_runner = ActiveAssignmentRunner(client, executor) + return HostAgentApplication( + client=client, + heartbeat=heartbeat, + processor=AssignmentProcessor(client, active_runner), + ) diff --git a/apps/device-host-agent/host_agent/lease.py b/apps/device-host-agent/host_agent/lease.py index beab2b5..daf719c 100644 --- a/apps/device-host-agent/host_agent/lease.py +++ b/apps/device-host-agent/host_agent/lease.py @@ -54,6 +54,10 @@ class ActiveAssignmentRunner: self.client = client self.executor = executor self._now = now or (lambda: datetime.now(UTC)) + self._stop_requested = Event() + + def request_stop(self) -> None: + self._stop_requested.set() async def run(self, assignment: AssignmentModel) -> AssignmentExecutionResult: guard = LeaseGuard() @@ -61,7 +65,8 @@ class ActiveAssignmentRunner: asyncio.to_thread( self.executor.execute, assignment, - should_stop=guard.is_lost, + should_stop=lambda: guard.is_lost() + or self._stop_requested.is_set(), ) ) renewal = asyncio.create_task( diff --git a/apps/device-host-agent/host_agent/processor.py b/apps/device-host-agent/host_agent/processor.py index dd94718..adb1e3e 100644 --- a/apps/device-host-agent/host_agent/processor.py +++ b/apps/device-host-agent/host_agent/processor.py @@ -11,6 +11,8 @@ from host_agent.client import HostAgentClient class ActiveAssignmentExecutor(Protocol): async def run(self, assignment: AssignmentModel) -> AssignmentExecutionResult: ... + def request_stop(self) -> None: ... + @dataclass(frozen=True) class AssignmentProcessingResult: @@ -27,6 +29,9 @@ class AssignmentProcessor: self.client = client self.active_executor = active_executor + 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" diff --git a/apps/device-host-agent/tests/test_app.py b/apps/device-host-agent/tests/test_app.py index d71bb48..decb740 100644 --- a/apps/device-host-agent/tests/test_app.py +++ b/apps/device-host-agent/tests/test_app.py @@ -1,10 +1,222 @@ from __future__ import annotations +import asyncio +from contextlib import suppress +from datetime import UTC, datetime, timedelta + +from cloud.internal_api.models import AssignmentModel +from device.manager import DeviceManager from host_agent.app import HostAgentApplication, create_application +from host_agent.config import HostAgentConfig -def test_create_application_returns_host_agent_process_shell() -> None: - application = create_application() +def _config() -> HostAgentConfig: + return HostAgentConfig( + control_plane_url="https://control.example", + host_id="host-a", + token="secret", + ) + + +def _assignment() -> AssignmentModel: + return AssignmentModel( + task_id="task-a", + attempt=1, + lease_id="lease-a", + lease_expires_at=datetime.now(UTC) + timedelta(seconds=30), + host_id="host-a", + device_id="device-a", + goal="open settings", + ) + + +def test_create_application_composes_host_agent_services(tmp_path, monkeypatch) -> None: + monkeypatch.chdir(tmp_path) + application = create_application(config=_config(), manager=DeviceManager()) assert isinstance(application, HostAgentApplication) - assert application.run() is None + asyncio.run(application.client.aclose()) + + +def test_shutdown_cancels_long_poll_and_sends_final_heartbeat() -> None: + async def scenario() -> 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") + + stop = asyncio.Event() + application = HostAgentApplication( + client=BlockingClient(), # type: ignore[arg-type] + heartbeat=RecordingHeartbeat(), # type: ignore[arg-type] + processor=IdleProcessor(), # type: ignore[arg-type] + ) + running = asyncio.create_task(application.run_async(stop)) + await claim_started.wait() + stop.set() + await asyncio.wait_for(running, timeout=1) + + assert claim_cancelled.is_set() + assert events == ["stop-work", "final-heartbeat", "closed"] + + asyncio.run(scenario()) + + +def test_shutdown_interrupts_active_work_before_final_heartbeat() -> None: + async def scenario() -> None: + processing_started = asyncio.Event() + processing_stopped = asyncio.Event() + events: list[str] = [] + claims = 0 + + class AssignedClient: + async def claim(self): + nonlocal claims + claims += 1 + if claims == 1: + return _assignment() + await asyncio.Event().wait() + + 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 CooperativeProcessor: + async def process(self, assignment): + processing_started.set() + await processing_stopped.wait() + events.append("work-finished") + + def request_stop(self): + events.append("stop-work") + processing_stopped.set() + + stop = asyncio.Event() + application = HostAgentApplication( + client=AssignedClient(), # type: ignore[arg-type] + heartbeat=RecordingHeartbeat(), # type: ignore[arg-type] + processor=CooperativeProcessor(), # type: ignore[arg-type] + ) + running = asyncio.create_task(application.run_async(stop)) + await processing_started.wait() + stop.set() + await asyncio.wait_for(running, timeout=1) + + assert events.index("work-finished") < events.index("final-heartbeat") + assert events[-1] == "closed" + + asyncio.run(scenario()) + + +def test_final_heartbeat_failure_does_not_prevent_client_close() -> None: + async def scenario() -> None: + closed = False + + class StoppedClient: + async def claim(self): + raise AssertionError("polling must not start") + + async def aclose(self): + nonlocal closed + closed = True + + class FailingHeartbeat: + async def run(self, stop): + await stop.wait() + + async def sync_once(self): + raise OSError("control plane unavailable") + + class IdleProcessor: + def request_stop(self): + return None + + stop = asyncio.Event() + stop.set() + application = HostAgentApplication( + client=StoppedClient(), # type: ignore[arg-type] + heartbeat=FailingHeartbeat(), # type: ignore[arg-type] + processor=IdleProcessor(), # type: ignore[arg-type] + ) + + await application.run_async(stop) + assert closed + + asyncio.run(scenario()) + + +def test_main_task_cancellation_waits_for_active_work_shutdown() -> None: + async def scenario() -> None: + processing_started = asyncio.Event() + processing_stopped = asyncio.Event() + events: list[str] = [] + + class AssignedClient: + async def claim(self): + return _assignment() + + 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 CooperativeProcessor: + async def process(self, assignment): + processing_started.set() + await processing_stopped.wait() + events.append("work-finished") + + def request_stop(self): + processing_stopped.set() + + application = HostAgentApplication( + client=AssignedClient(), # type: ignore[arg-type] + heartbeat=RecordingHeartbeat(), # type: ignore[arg-type] + processor=CooperativeProcessor(), # type: ignore[arg-type] + ) + running = asyncio.create_task(application.run_async()) + await processing_started.wait() + running.cancel() + + with suppress(asyncio.CancelledError): + await running + + assert events == ["work-finished", "final-heartbeat", "closed"] + + asyncio.run(scenario()) diff --git a/apps/device-host-agent/tests/test_lease.py b/apps/device-host-agent/tests/test_lease.py index d5f5560..9efcb66 100644 --- a/apps/device-host-agent/tests/test_lease.py +++ b/apps/device-host-agent/tests/test_lease.py @@ -126,3 +126,39 @@ def test_renewal_loop_exits_when_execution_finishes() -> None: assert renew_calls == 0 asyncio.run(scenario()) + + +def test_shutdown_request_stops_active_execution_cooperatively() -> None: + async def scenario() -> None: + execution_started = Event() + + class CooperativeExecutor: + def execute(self, assignment, *, should_stop=None): + assert should_stop is not None + execution_started.set() + while not should_stop(): + Event().wait(0.001) + return AssignmentExecutionResult( + status="failed", + failure_reason="execution interrupted", + ) + + class RenewingClient: + async def renew(self, assignment): + return LeaseRenewalResponse( + status="renewed", + lease_expires_at=datetime.now(UTC) + timedelta(seconds=30), + ) + + runner = ActiveAssignmentRunner( + RenewingClient(), # type: ignore[arg-type] + CooperativeExecutor(), + ) + running = asyncio.create_task(runner.run(_assignment(expires_in=30))) + assert await asyncio.to_thread(execution_started.wait, 1) + runner.request_stop() + + result = await asyncio.wait_for(running, timeout=1) + assert result.failure_reason == "execution interrupted" + + asyncio.run(scenario()) diff --git a/openspec/changes/cloud-control-plane-integration/tasks.md b/openspec/changes/cloud-control-plane-integration/tasks.md index 7e66ec6..cfa680d 100644 --- a/openspec/changes/cloud-control-plane-integration/tasks.md +++ b/openspec/changes/cloud-control-plane-integration/tasks.md @@ -58,7 +58,7 @@ - [x] 7.4 Execute goal assignments through the configured Runtime Planner/Executor and workflow assignments through the existing workflow runner. - [x] 7.5 Run lease renewal alongside active execution and stop further interruptible actions after confirmed lease loss. - [x] 7.6 Normalize and report successful/failed terminal outcomes, including Runtime failure reasons, with idempotent retries after response loss. -- [ ] 7.7 Implement graceful shutdown that stops polling, finishes or interrupts current work according to lease policy, and performs a final heartbeat when possible. +- [x] 7.7 Implement graceful shutdown that stops polling, finishes or interrupts current work according to lease policy, and performs a final heartbeat when possible. - [ ] 7.8 Add fake-driver end-to-end tests for one host, multiple hosts, NAT-style outbound-only operation, control-plane restart, Host Agent restart, and lease loss. ## 8. Public SDK And Operational Delivery