feat(host-agent): renew active assignment leases

workflow/runner.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from collections.abc import Callable
-from dataclasses import replace
 from datetime import datetime
 from types import SimpleNamespace
 from typing import Any
@@ -33,6 +32,7 @@ TaskRunnerFactory = Callable[[], TaskRunner]
 SceneProvider = Callable[[], Scene | None]
 WorldStateProvider = Callable[[], object | None]
 SleepFunc = Callable[[float], None]
+StopRequested = Callable[[], bool]
 
 TERMINAL_STATUSES = {"completed", "failed", "cancelled"}
 
@@ -68,6 +68,8 @@ class WorkflowRunner:
         definition: WorkflowDefinition,
         device_id: str,
         initial_variables: dict[str, Any] | None = None,
+        *,
+        should_stop: StopRequested | None = None,
     ) -> WorkflowRun:
         self.store.save_definition(definition)
         run = self.store.create_run(
@@ -75,9 +77,14 @@ class WorkflowRunner:
             initial_variables or {},
             device_id=device_id,
         )
-        return self._drive(definition, run)
+        return self._drive(definition, run, should_stop=should_stop)
 
-    def resume(self, run_id: str) -> WorkflowRun:
+    def resume(
+        self,
+        run_id: str,
+        *,
+        should_stop: StopRequested | None = None,
+    ) -> WorkflowRun:
         run = self.store.get_run(run_id)
         if run is None:
             raise KeyError(f"unknown workflow run {run_id}")
@@ -86,15 +93,19 @@ class WorkflowRunner:
         definition = self.store.get_definition(run.definition_id)
         if definition is None:
             raise KeyError(f"unknown workflow definition {run.definition_id}")
-        return self._drive(definition, run)
+        return self._drive(definition, run, should_stop=should_stop)
 
     def _drive(
         self,
         definition: WorkflowDefinition,
         run: WorkflowRun,
+        *,
+        should_stop: StopRequested | None = None,
     ) -> WorkflowRun:
         executed = 0
         while run.status not in TERMINAL_STATUSES and run.current_step_id:
+            if should_stop is not None and should_stop():
+                return self._checkpoint(run, run.current_step_id, "cancelled")
             if self.step_limit is not None and executed >= self.step_limit:
                 return run
             step = definition.step_by_id(run.current_step_id)
@@ -109,7 +120,15 @@ class WorkflowRunner:
                 run = self._checkpoint(run, next_step_id, next_status)
                 continue
 
-            result, branch_next_step_id = self._execute_step(definition, run, step)
+            result, branch_next_step_id = self._execute_step(
+                definition,
+                run,
+                step,
+                should_stop=should_stop,
+            )
+            if should_stop is not None and should_stop():
+                self.store.append_step_result(run.id, result)
+                return self._checkpoint(run, run.current_step_id, "cancelled")
             next_status, next_step_id = self._resolve_outcome(
                 definition, step, result.success, branch_next_step_id
             )
@@ -164,13 +183,23 @@ class WorkflowRunner:
         definition: WorkflowDefinition,
         run: WorkflowRun,
         step: WorkflowStep,
+        *,
+        should_stop: StopRequested | None = None,
     ) -> tuple[WorkflowStepResult, str | None]:
         if isinstance(step, PlannedGoalStep):
-            return self._execute_planned_goal_step(run, step), None
+            return self._execute_planned_goal_step(
+                run,
+                step,
+                should_stop=should_stop,
+            ), None
         if isinstance(step, SkillInvocationStep):
             return self._execute_skill_invocation_step(step), None
         if isinstance(step, WaitForConditionStep):
-            return self._execute_wait_step(run, step), None
+            return self._execute_wait_step(
+                run,
+                step,
+                should_stop=should_stop,
+            ), None
         if isinstance(step, BranchStep):
             return self._execute_branch_step(run, step)
         return (
@@ -187,10 +216,15 @@ class WorkflowRunner:
         self,
         run: WorkflowRun,
         step: PlannedGoalStep,
+        *,
+        should_stop: StopRequested | None = None,
     ) -> WorkflowStepResult:
         task = Task(goal=step.goal, device_id=run.device_id or "")
         task_runner = self.task_runner_factory()
-        result_task = task_runner.run(task)
+        if should_stop is None:
+            result_task = task_runner.run(task)
+        else:
+            result_task = task_runner.run(task, should_stop=should_stop)
         success = result_task.status == "completed"
         return WorkflowStepResult(
             step_id=step.step_id,
@@ -238,11 +272,20 @@ class WorkflowRunner:
         self,
         run: WorkflowRun,
         step: WaitForConditionStep,
+        *,
+        should_stop: StopRequested | None = None,
     ) -> WorkflowStepResult:
         started_at = datetime.now().astimezone()
         timeout_seconds = step.timeout_seconds
         poll_interval_seconds = step.poll_interval_seconds
         while True:
+            if should_stop is not None and should_stop():
+                return WorkflowStepResult(
+                    step_id=step.step_id,
+                    kind=step.kind,
+                    success=False,
+                    detail={"reason": "execution interrupted"},
+                )
             try:
                 if self._condition_is_true(run, step.condition, started_at=started_at):
                     return WorkflowStepResult(