From 88189770ff1d6e40318aeb4fda147cb8ab9ac452 Mon Sep 17 00:00:00 2001 From: Jerry Yan <792602257@qq.com> Date: Wed, 15 Jul 2026 17:52:33 +0800 Subject: [PATCH] feat(runtime): add cancellation-aware stop_reason to TaskRunner and WorkflowRunner - TaskRunner.run() and WorkflowRunner.run()/resume() accept an optional stop_reason callable alongside should_stop, distinguishing a genuine cancellation from other stop conditions (e.g. lost lease). - is_cancellation_reason() shared helper added to runtime/task.py. - WorkflowRunner._stop_status() now branches cancelled/failed based on stop_reason, correcting a prior blanket cancelled-on-any-stop behavior that conflicted with the host-agent-protocol spec's requirement to distinguish cancellation from lease-loss stops. - Default behavior (stop_reason=None) is preserved exactly for both runners so existing callers/tests are unaffected. - Task 1 of openspec change task-cancellation. --- openspec/changes/task-cancellation/tasks.md | 10 +-- runtime/task.py | 22 ++++-- tests/test_task_loop.py | 70 +++++++++++++++++++ tests/test_workflow_runner.py | 76 +++++++++++++++++++++ workflow/runner.py | 38 +++++++++-- 5 files changed, 199 insertions(+), 17 deletions(-) diff --git a/openspec/changes/task-cancellation/tasks.md b/openspec/changes/task-cancellation/tasks.md index 9cfcbc6..b44c1bb 100644 --- a/openspec/changes/task-cancellation/tasks.md +++ b/openspec/changes/task-cancellation/tasks.md @@ -1,10 +1,10 @@ ## 1. Core and Runtime status plumbing -- [ ] 1.1 Confirm `core/models.py`'s `TaskStatus` already includes `"cancelled"` (it does); add `StepStatus` no change needed — verify no other literal needs widening. -- [ ] 1.2 Add an optional `stop_reason: Callable[[], str] | None` parameter to `TaskRunner.run()` (`runtime/task.py`), defaulting to `None`. -- [ ] 1.3 Update `TaskRunner._interrupt_task()` to accept the resolved reason string, set `status="cancelled"` when the reason indicates cancellation (e.g. contains `"cancel"`), else keep `status="failed"` as today, and record the reason as `failure_reason` in both cases. -- [ ] 1.4 Update `WorkflowRunner`'s stop branch (`workflow/runner.py`) to accept and pass through the same optional `stop_reason`, reusing its existing `"cancelled"` checkpoint call — confirm no behavior change needed since it already lands on `"cancelled"` for any stop; only wire the reason through for consistency/logging. -- [ ] 1.5 Add/extend unit tests in `tests/` for `TaskRunner` covering: cancellation-flavored stop → `status="cancelled"`; lease-loss-flavored stop → `status="failed"` (existing behavior preserved). +- [x] 1.1 Confirm `core/models.py`'s `TaskStatus` already includes `"cancelled"` (it does); add `StepStatus` no change needed — verify no other literal needs widening. +- [x] 1.2 Add an optional `stop_reason: Callable[[], str] | None` parameter to `TaskRunner.run()` (`runtime/task.py`), defaulting to `None`. +- [x] 1.3 Update `TaskRunner._interrupt_task()` to accept the resolved reason string, set `status="cancelled"` when the reason indicates cancellation (e.g. contains `"cancel"`), else keep `status="failed"` as today, and record the reason as `failure_reason` in both cases. +- [x] 1.4 Update `WorkflowRunner`'s stop branch (`workflow/runner.py`) to accept and pass through the same optional `stop_reason`, reusing its existing `"cancelled"` checkpoint call — confirm no behavior change needed since it already lands on `"cancelled"` for any stop; only wire the reason through for consistency/logging. (Correction during implementation: WorkflowRunner previously collapsed *every* stop, including lease-loss, into `"cancelled"`, which contradicts the host-agent-protocol spec's requirement to distinguish cancellation from lease-loss stops. `_stop_status()` now branches on `stop_reason` the same way `TaskRunner` does, while preserving the exact pre-existing default (`"cancelled"`) when no `stop_reason` is supplied.) +- [x] 1.5 Add/extend unit tests in `tests/` for `TaskRunner` covering: cancellation-flavored stop → `status="cancelled"`; lease-loss-flavored stop → `status="failed"` (existing behavior preserved). Also added matching `WorkflowRunner` coverage for the same `_stop_status` branching. ## 2. Cloud persistence: schema and repository diff --git a/runtime/task.py b/runtime/task.py index 5a3a867..3ed6885 100644 --- a/runtime/task.py +++ b/runtime/task.py @@ -40,9 +40,15 @@ Observer = Callable[[str], Scene] ScreenshotProvider = Callable[[str], bytes] TaskSucceededHook = Callable[[str, str, Timeline], None] StopRequested = Callable[[], bool] +StopReason = Callable[[], "str | None"] StepProgressCallback = Callable[[int, str, str], None] +def is_cancellation_reason(reason: str | None) -> bool: + """Distinguish an explicit cancellation stop from other stop reasons (e.g. lost lease).""" + return bool(reason) and "cancel" in reason.lower() + + class TaskRunner: def __init__( self, @@ -98,6 +104,7 @@ class TaskRunner: task: Task, *, should_stop: StopRequested | None = None, + stop_reason: StopReason | None = None, ) -> Task: if self.metadata_store: self.metadata_store.create_task(task) @@ -109,7 +116,7 @@ class TaskRunner: for _ in range(self.config.max_steps): if should_stop is not None and should_stop(): - return self._interrupt_task(task) + return self._interrupt_task(task, stop_reason) try: scene = self.observer(task.device_id) context.add_scene(scene) @@ -136,7 +143,7 @@ class TaskRunner: for step_index, step in enumerate(steps): if should_stop is not None and should_stop(): - return self._interrupt_task(task) + return self._interrupt_task(task, stop_reason) executable_step = self._step_for_device(step, task.device_id) if step_index == 0 and screenshot is not None: # Reuse the screenshot already captured for planning instead of @@ -194,13 +201,16 @@ class TaskRunner: ) return task - def _interrupt_task(self, task: Task) -> Task: - self._emit_step_progress(-1, "failed", "execution interrupted") + def _interrupt_task(self, task: Task, stop_reason: StopReason | None = None) -> Task: + reason = stop_reason() if stop_reason is not None else None + message = reason or "execution interrupted" + status = "cancelled" if is_cancellation_reason(reason) else "failed" + self._emit_step_progress(-1, "failed", message) self._update_task( task, - status="failed", + status=status, completed=True, - failure_reason="execution interrupted", + failure_reason=message, ) return task diff --git a/tests/test_task_loop.py b/tests/test_task_loop.py index 729dca0..ad60e1b 100644 --- a/tests/test_task_loop.py +++ b/tests/test_task_loop.py @@ -115,6 +115,76 @@ def test_task_runner_stops_before_the_next_planned_action() -> None: assert actions == ["tap"] +def test_task_runner_stop_reason_cancellation_yields_cancelled_status() -> None: + scene = Scene(width=10, height=20, elements=[]) + stop_requested = False + + def record_action(**kwargs): + nonlocal stop_requested + stop_requested = True + return {"ok": True} + + runner = TaskRunner( + planner=ScriptedPlanner( + [ + PlannedStep(action="tap", description="first", args={}), + PlannedStep(action="tap", description="second", args={}), + ] + ), + executor=Executor( + tools={"tap": record_action}, + config=ExecutorConfig(max_retries=1, backoff_seconds=0), + ), + config=TaskRunnerConfig(max_steps=5), + observer=lambda device_id: scene, + screenshot_provider=lambda device_id: PNG_10X20, + ) + + result = runner.run( + Task(goal="perform two actions", device_id="phone"), + should_stop=lambda: stop_requested, + stop_reason=lambda: "cancellation requested by control plane", + ) + + assert result.status == "cancelled" + assert result.failure_reason == "cancellation requested by control plane" + + +def test_task_runner_stop_reason_lease_loss_yields_failed_status() -> None: + scene = Scene(width=10, height=20, elements=[]) + stop_requested = False + + def record_action(**kwargs): + nonlocal stop_requested + stop_requested = True + return {"ok": True} + + runner = TaskRunner( + planner=ScriptedPlanner( + [ + PlannedStep(action="tap", description="first", args={}), + PlannedStep(action="tap", description="second", args={}), + ] + ), + executor=Executor( + tools={"tap": record_action}, + config=ExecutorConfig(max_retries=1, backoff_seconds=0), + ), + config=TaskRunnerConfig(max_steps=5), + observer=lambda device_id: scene, + screenshot_provider=lambda device_id: PNG_10X20, + ) + + result = runner.run( + Task(goal="perform two actions", device_id="phone"), + should_stop=lambda: stop_requested, + stop_reason=lambda: "lease rejected by control plane", + ) + + assert result.status == "failed" + assert result.failure_reason == "lease rejected by control plane" + + def test_task_runner_persists_action_evidence_and_raw_ocr(tmp_path) -> None: scene = Scene( width=10, diff --git a/tests/test_workflow_runner.py b/tests/test_workflow_runner.py index f4ecf21..3a46c7e 100644 --- a/tests/test_workflow_runner.py +++ b/tests/test_workflow_runner.py @@ -166,6 +166,82 @@ def test_workflow_runner_stops_before_the_next_step(tmp_path) -> None: assert calls == ["first"] +def test_workflow_runner_stop_reason_cancellation_yields_cancelled_status( + tmp_path, +) -> None: + stop_requested = False + calls: list[str] = [] + + class StoppingTaskRunner: + def run(self, task: Task, *, should_stop=None, stop_reason=None) -> Task: + nonlocal stop_requested + calls.append(task.goal) + stop_requested = True + task.status = "completed" + return task + + definition = WorkflowDefinition( + name="interruptible", + entry_step_id="first", + steps=[ + PlannedGoalStep("first", "first", next_step_id="second"), + PlannedGoalStep("second", "second"), + ], + ) + runner = WorkflowRunner( + _store(tmp_path), + task_runner_factory=lambda: StoppingTaskRunner(), # type: ignore[arg-type] + ) + + run = runner.run( + definition, + "phone", + should_stop=lambda: stop_requested, + stop_reason=lambda: "cancellation requested by control plane", + ) + + assert run.status == "cancelled" + assert calls == ["first"] + + +def test_workflow_runner_stop_reason_lease_loss_yields_failed_status( + tmp_path, +) -> None: + stop_requested = False + calls: list[str] = [] + + class StoppingTaskRunner: + def run(self, task: Task, *, should_stop=None, stop_reason=None) -> Task: + nonlocal stop_requested + calls.append(task.goal) + stop_requested = True + task.status = "completed" + return task + + definition = WorkflowDefinition( + name="interruptible", + entry_step_id="first", + steps=[ + PlannedGoalStep("first", "first", next_step_id="second"), + PlannedGoalStep("second", "second"), + ], + ) + runner = WorkflowRunner( + _store(tmp_path), + task_runner_factory=lambda: StoppingTaskRunner(), # type: ignore[arg-type] + ) + + run = runner.run( + definition, + "phone", + should_stop=lambda: stop_requested, + stop_reason=lambda: "lease rejected by control plane", + ) + + assert run.status == "failed" + assert calls == ["first"] + + def test_workflow_runner_failing_planned_goal_marks_run_failed(tmp_path) -> None: definition = WorkflowDefinition( name="fail", diff --git a/workflow/runner.py b/workflow/runner.py index 4b75d4f..5deabe9 100644 --- a/workflow/runner.py +++ b/workflow/runner.py @@ -7,7 +7,7 @@ from typing import Any from time import sleep from core.models import Scene, Task -from runtime.task import TaskRunner +from runtime.task import StopReason, TaskRunner, is_cancellation_reason from skills_learning.store import SkillStore, get_default_store from workflow.conditions import ( ConditionEvaluator, @@ -70,6 +70,7 @@ class WorkflowRunner: initial_variables: dict[str, Any] | None = None, *, should_stop: StopRequested | None = None, + stop_reason: StopReason | None = None, ) -> WorkflowRun: self.store.save_definition(definition) run = self.store.create_run( @@ -77,13 +78,14 @@ class WorkflowRunner: initial_variables or {}, device_id=device_id, ) - return self._drive(definition, run, should_stop=should_stop) + return self._drive(definition, run, should_stop=should_stop, stop_reason=stop_reason) def resume( self, run_id: str, *, should_stop: StopRequested | None = None, + stop_reason: StopReason | None = None, ) -> WorkflowRun: run = self.store.get_run(run_id) if run is None: @@ -93,7 +95,18 @@ 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, should_stop=should_stop) + return self._drive(definition, run, should_stop=should_stop, stop_reason=stop_reason) + + def _stop_status(self, stop_reason: StopReason | None) -> str: + """Resolve the terminal status for a should_stop-triggered stop. + + No `stop_reason` preserves the pre-existing default of `cancelled` for + any stop; a supplied reason distinguishes an explicit cancellation + from other stop conditions (e.g. lost lease), which resolve to `failed`. + """ + if stop_reason is None: + return "cancelled" + return "cancelled" if is_cancellation_reason(stop_reason()) else "failed" def _drive( self, @@ -101,11 +114,14 @@ class WorkflowRunner: run: WorkflowRun, *, should_stop: StopRequested | None = None, + stop_reason: StopReason | 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") + return self._checkpoint( + run, run.current_step_id, self._stop_status(stop_reason) + ) if self.step_limit is not None and executed >= self.step_limit: return run step = definition.step_by_id(run.current_step_id) @@ -125,10 +141,13 @@ class WorkflowRunner: run, step, should_stop=should_stop, + stop_reason=stop_reason, ) 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") + return self._checkpoint( + run, run.current_step_id, self._stop_status(stop_reason) + ) next_status, next_step_id = self._resolve_outcome( definition, step, result.success, branch_next_step_id ) @@ -185,12 +204,14 @@ class WorkflowRunner: step: WorkflowStep, *, should_stop: StopRequested | None = None, + stop_reason: StopReason | None = None, ) -> tuple[WorkflowStepResult, str | None]: if isinstance(step, PlannedGoalStep): return self._execute_planned_goal_step( run, step, should_stop=should_stop, + stop_reason=stop_reason, ), None if isinstance(step, SkillInvocationStep): return self._execute_skill_invocation_step(step), None @@ -218,13 +239,18 @@ class WorkflowRunner: step: PlannedGoalStep, *, should_stop: StopRequested | None = None, + stop_reason: StopReason | None = None, ) -> WorkflowStepResult: task = Task(goal=step.goal, device_id=run.device_id or "") task_runner = self.task_runner_factory() if should_stop is None: result_task = task_runner.run(task) - else: + elif stop_reason is None: result_task = task_runner.run(task, should_stop=should_stop) + else: + result_task = task_runner.run( + task, should_stop=should_stop, stop_reason=stop_reason + ) success = result_task.status == "completed" return WorkflowStepResult( step_id=step.step_id,