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.
This commit is contained in:
2026-07-15 17:52:33 +08:00
parent 947434b65a
commit 88189770ff
5 changed files with 199 additions and 17 deletions
+5 -5
View File
@@ -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
+16 -6
View File
@@ -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
+70
View File
@@ -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,
+76
View File
@@ -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",
+32 -6
View File
@@ -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,