feat(host-agent): renew active assignment leases

apps/device-host-agent/host_agent/assignment.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections.abc import Callable
 from dataclasses import dataclass, field
 from typing import Any
 
@@ -19,11 +20,21 @@ class AssignmentExecutor:
     def __init__(self, factories: ExecutionFactories) -> None:
         self.factories = factories
 
-    def execute(self, assignment: AssignmentModel) -> AssignmentExecutionResult:
+    def execute(
+        self,
+        assignment: AssignmentModel,
+        *,
+        should_stop: Callable[[], bool] | None = None,
+    ) -> AssignmentExecutionResult:
+        if should_stop is not None and should_stop():
+            return AssignmentExecutionResult(
+                status="failed",
+                failure_reason="execution interrupted",
+            )
         if assignment.workflow_definition_id is not None:
-            return self._execute_workflow(assignment)
+            return self._execute_workflow(assignment, should_stop=should_stop)
         if assignment.goal is not None:
-            return self._execute_goal(assignment)
+            return self._execute_goal(assignment, should_stop=should_stop)
         return AssignmentExecutionResult(
             status="failed",
             failure_reason="assignment has neither goal nor workflow definition",
@@ -32,9 +43,15 @@ class AssignmentExecutor:
     def _execute_goal(
         self,
         assignment: AssignmentModel,
+        *,
+        should_stop: Callable[[], bool] | None,
     ) -> AssignmentExecutionResult:
         task = Task(goal=assignment.goal or "", device_id=assignment.device_id)
-        completed = self.factories.task_runner_factory().run(task)
+        runner = self.factories.task_runner_factory()
+        if should_stop is None:
+            completed = runner.run(task)
+        else:
+            completed = runner.run(task, should_stop=should_stop)
         return AssignmentExecutionResult(
             status="done" if completed.status == "completed" else "failed",
             failure_reason=completed.failure_reason,
@@ -47,6 +64,8 @@ class AssignmentExecutor:
     def _execute_workflow(
         self,
         assignment: AssignmentModel,
+        *,
+        should_stop: Callable[[], bool] | None,
     ) -> AssignmentExecutionResult:
         definition_id = assignment.workflow_definition_id or ""
         definition = self.factories.workflow_store.get_definition(definition_id)
@@ -55,10 +74,15 @@ class AssignmentExecutor:
                 status="failed",
                 failure_reason=f"unknown workflow definition {definition_id!r}",
             )
-        run = self.factories.workflow_runner_factory().run(
-            definition,
-            device_id=assignment.device_id,
-        )
+        runner = self.factories.workflow_runner_factory()
+        if should_stop is None:
+            run = runner.run(definition, device_id=assignment.device_id)
+        else:
+            run = runner.run(
+                definition,
+                device_id=assignment.device_id,
+                should_stop=should_stop,
+            )
         return AssignmentExecutionResult(
             status="done" if run.status == "completed" else "failed",
             failure_reason=(