feat(host-agent): renew active assignment leases

runtime/task.py

@@ -39,6 +39,7 @@ class TaskRunnerConfig:
 Observer = Callable[[str], Scene]
 ScreenshotProvider = Callable[[str], bytes]
 TaskSucceededHook = Callable[[str, str, Timeline], None]
+StopRequested = Callable[[], bool]
 
 
 class TaskRunner:
@@ -89,7 +90,12 @@ class TaskRunner:
         else:
             self.on_task_succeeded = None
 
-    def run(self, task: Task) -> Task:
+    def run(
+        self,
+        task: Task,
+        *,
+        should_stop: StopRequested | None = None,
+    ) -> Task:
         context = TaskContext(task_id=task.id, goal=task.goal)
         world_handle = self._start_world_view(task.id)
         if world_handle is not None:
@@ -97,6 +103,8 @@ class TaskRunner:
         self._update_task(task, status="running")
 
         for _ in range(self.config.max_steps):
+            if should_stop is not None and should_stop():
+                return self._interrupt_task(task)
             try:
                 scene = self.observer(task.device_id)
                 context.add_scene(scene)
@@ -119,6 +127,8 @@ class TaskRunner:
                 return self._complete_task(task)
 
             for step in steps:
+                if should_stop is not None and should_stop():
+                    return self._interrupt_task(task)
                 executable_step = self._step_for_device(step, task.device_id)
                 result = self.executor.execute(
                     executable_step,
@@ -143,6 +153,15 @@ class TaskRunner:
         )
         return task
 
+    def _interrupt_task(self, task: Task) -> Task:
+        self._update_task(
+            task,
+            status="failed",
+            completed=True,
+            failure_reason="execution interrupted",
+        )
+        return task
+
     def _start_world_view(self, task_id: str) -> TaskWorldView | None:
         """Start (or resume) this task's isolated `WorldModel` handle, if enabled.