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