feat(host-agent): shut down gracefully

2026-07-12 19:47:50 +08:00
parent 91f08509a7
commit b1303569e3
6 changed files with 360 additions and 9 deletions
@@ -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),
+    )
@@ -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(
@@ -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"
@@ -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())
@@ -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())
@@ -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