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
@@ -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())