feat(host-agent): stop assignment execution collaboratively on cancellation

- LeaseGuard gains an is_cancellation convenience property
- ActiveAssignmentRunner marks the lease lost with a cancellation
  reason when a renewal response reports cancel_requested
- AssignmentExecutor threads stop_reason through to TaskRunner/
  WorkflowRunner and maps a cancellation-flavored stop to
  AssignmentExecutionResult.status = "cancelled" instead of "failed"
- AssignmentProcessor forwards a three-way done/cancelled/failed
  status when reporting the terminal result
- Add/extend tests across lease, assignment, processor, and client
This commit is contained in:
2026-07-15 18:28:21 +08:00
parent 8a0d48eada
commit d3024b4810
9 changed files with 233 additions and 22 deletions
@@ -9,6 +9,7 @@ from core.models import Task
from host_agent.execution import ExecutionFactories from host_agent.execution import ExecutionFactories
from host_agent.planner_context import bind_planner_execution_context from host_agent.planner_context import bind_planner_execution_context
from host_agent.progress import TaskProgressHolder, TaskProgressSnapshot from host_agent.progress import TaskProgressHolder, TaskProgressSnapshot
from runtime.task import is_cancellation_reason
@dataclass(frozen=True) @dataclass(frozen=True)
@@ -32,18 +33,24 @@ class AssignmentExecutor:
assignment: AssignmentModel, assignment: AssignmentModel,
*, *,
should_stop: Callable[[], bool] | None = None, should_stop: Callable[[], bool] | None = None,
stop_reason: Callable[[], str | None] | None = None,
) -> AssignmentExecutionResult: ) -> AssignmentExecutionResult:
self._progress.clear() self._progress.clear()
with bind_planner_execution_context(assignment): with bind_planner_execution_context(assignment):
if should_stop is not None and should_stop(): if should_stop is not None and should_stop():
reason = stop_reason() if stop_reason is not None else None
return AssignmentExecutionResult( return AssignmentExecutionResult(
status="failed", status="cancelled" if is_cancellation_reason(reason) else "failed",
failure_reason="execution interrupted", failure_reason=reason or "execution interrupted",
) )
if assignment.workflow_definition_id is not None: if assignment.workflow_definition_id is not None:
return self._execute_workflow(assignment, should_stop=should_stop) return self._execute_workflow(
assignment, should_stop=should_stop, stop_reason=stop_reason
)
if assignment.goal is not None: if assignment.goal is not None:
return self._execute_goal(assignment, should_stop=should_stop) return self._execute_goal(
assignment, should_stop=should_stop, stop_reason=stop_reason
)
return AssignmentExecutionResult( return AssignmentExecutionResult(
status="failed", status="failed",
failure_reason="assignment has neither goal nor workflow definition", failure_reason="assignment has neither goal nor workflow definition",
@@ -54,6 +61,7 @@ class AssignmentExecutor:
assignment: AssignmentModel, assignment: AssignmentModel,
*, *,
should_stop: Callable[[], bool] | None, should_stop: Callable[[], bool] | None,
stop_reason: Callable[[], str | None] | None,
) -> AssignmentExecutionResult: ) -> AssignmentExecutionResult:
task = Task(goal=assignment.goal or "", device_id=assignment.device_id) task = Task(goal=assignment.goal or "", device_id=assignment.device_id)
if self.factories.metadata_store is not None: if self.factories.metadata_store is not None:
@@ -67,9 +75,11 @@ class AssignmentExecutor:
if should_stop is None: if should_stop is None:
completed = runner.run(task) completed = runner.run(task)
else: else:
completed = runner.run(task, should_stop=should_stop) completed = runner.run(
task, should_stop=should_stop, stop_reason=stop_reason
)
return AssignmentExecutionResult( return AssignmentExecutionResult(
status="done" if completed.status == "completed" else "failed", status=_terminal_status(completed.status),
failure_reason=completed.failure_reason, failure_reason=completed.failure_reason,
metadata={ metadata={
"runtime_task_id": completed.id, "runtime_task_id": completed.id,
@@ -82,6 +92,7 @@ class AssignmentExecutor:
assignment: AssignmentModel, assignment: AssignmentModel,
*, *,
should_stop: Callable[[], bool] | None, should_stop: Callable[[], bool] | None,
stop_reason: Callable[[], str | None] | None,
) -> AssignmentExecutionResult: ) -> AssignmentExecutionResult:
definition_id = assignment.workflow_definition_id or "" definition_id = assignment.workflow_definition_id or ""
definition = self.factories.workflow_store.get_definition(definition_id) definition = self.factories.workflow_store.get_definition(definition_id)
@@ -98,9 +109,10 @@ class AssignmentExecutor:
definition, definition,
device_id=assignment.device_id, device_id=assignment.device_id,
should_stop=should_stop, should_stop=should_stop,
stop_reason=stop_reason,
) )
return AssignmentExecutionResult( return AssignmentExecutionResult(
status="done" if run.status == "completed" else "failed", status=_terminal_status(run.status),
failure_reason=( failure_reason=(
None if run.status == "completed" else f"workflow ended as {run.status}" None if run.status == "completed" else f"workflow ended as {run.status}"
), ),
@@ -109,3 +121,11 @@ class AssignmentExecutor:
"workflow_status": run.status, "workflow_status": run.status,
}, },
) )
def _terminal_status(runtime_status: str) -> str:
if runtime_status == "completed":
return "done"
if runtime_status == "cancelled":
return "cancelled"
return "failed"
@@ -12,6 +12,7 @@ from cloud.internal_api.models import AssignmentModel
from host_agent.assignment import AssignmentExecutionResult from host_agent.assignment import AssignmentExecutionResult
from host_agent.client import HostAgentAPIError, HostAgentClient, StaleLeaseError from host_agent.client import HostAgentAPIError, HostAgentClient, StaleLeaseError
from host_agent.progress import TaskProgressSnapshot from host_agent.progress import TaskProgressSnapshot
from runtime.task import is_cancellation_reason
class InterruptibleAssignmentExecutor(Protocol): class InterruptibleAssignmentExecutor(Protocol):
@@ -20,6 +21,7 @@ class InterruptibleAssignmentExecutor(Protocol):
assignment: AssignmentModel, assignment: AssignmentModel,
*, *,
should_stop: Callable[[], bool] | None = None, should_stop: Callable[[], bool] | None = None,
stop_reason: Callable[[], str | None] | None = None,
) -> AssignmentExecutionResult: ... ) -> AssignmentExecutionResult: ...
def latest_progress(self) -> TaskProgressSnapshot | None: ... def latest_progress(self) -> TaskProgressSnapshot | None: ...
@@ -36,6 +38,10 @@ class LeaseGuard:
with self._lock: with self._lock:
return self._reason return self._reason
@property
def is_cancellation(self) -> bool:
return is_cancellation_reason(self.reason)
def is_lost(self) -> bool: def is_lost(self) -> bool:
return self._lost.is_set() return self._lost.is_set()
@@ -69,6 +75,7 @@ class ActiveAssignmentRunner:
self.executor.execute, self.executor.execute,
assignment, assignment,
should_stop=lambda: guard.is_lost() or self._stop_requested.is_set(), should_stop=lambda: guard.is_lost() or self._stop_requested.is_set(),
stop_reason=lambda: guard.reason,
) )
) )
renewal = asyncio.create_task( renewal = asyncio.create_task(
@@ -108,4 +115,7 @@ class ActiveAssignmentRunner:
guard.mark_lost("lease renewal failed after transport retries") guard.mark_lost("lease renewal failed after transport retries")
return return
else: else:
if response.cancel_requested:
guard.mark_lost("cancellation requested by control plane")
return
lease_expires_at = response.lease_expires_at lease_expires_at = response.lease_expires_at
@@ -47,8 +47,11 @@ class AssignmentProcessor:
self.status_tracker.mark_assignment_started(assignment) self.status_tracker.mark_assignment_started(assignment)
try: try:
execution = await self.active_executor.run(assignment) execution = await self.active_executor.run(assignment)
status = "done" if execution.status == "done" else "failed" if execution.status in {"done", "cancelled"}:
failure_reason = execution.failure_reason if status == "failed" else None status = execution.status
else:
status = "failed"
failure_reason = execution.failure_reason if status != "done" else None
response = await self.client.report_result( response = await self.client.report_result(
assignment, assignment,
status=status, status=status,
@@ -78,6 +78,39 @@ def test_goal_assignment_preserves_runtime_failure_reason() -> None:
assert result.failure_reason == "planner unavailable" assert result.failure_reason == "planner unavailable"
def test_goal_assignment_maps_cancellation_stop_to_cancelled_status() -> None:
# First should_stop() call is Executor.execute()'s pre-flight check (must pass
# through so the runner is actually invoked); the runner's own loop then stops.
calls = {"count": 0}
def should_stop() -> bool:
calls["count"] += 1
return calls["count"] > 1
class FakeTaskRunner:
def run(self, task: Task, *, should_stop=None, stop_reason=None) -> Task:
assert should_stop is not None and should_stop()
assert stop_reason is not None
task.status = "cancelled"
task.failure_reason = stop_reason()
return task
factories = ExecutionFactories(
task_runner_factory=lambda: FakeTaskRunner(), # type: ignore[arg-type,return-value]
workflow_runner_factory=lambda: object(), # type: ignore[arg-type,return-value]
workflow_store=object(), # type: ignore[arg-type]
)
result = AssignmentExecutor(factories).execute(
_assignment(),
should_stop=should_stop,
stop_reason=lambda: "cancellation requested by control plane",
)
assert result.status == "cancelled"
assert result.failure_reason == "cancellation requested by control plane"
def test_workflow_assignment_loads_and_executes_definition() -> None: def test_workflow_assignment_loads_and_executes_definition() -> None:
definition = object() definition = object()
calls: list[tuple[object, str]] = [] calls: list[tuple[object, str]] = []
@@ -109,6 +142,44 @@ def test_workflow_assignment_loads_and_executes_definition() -> None:
} }
def test_workflow_assignment_maps_cancellation_stop_to_cancelled_status() -> None:
calls = {"count": 0}
def should_stop() -> bool:
calls["count"] += 1
return calls["count"] > 1
class FakeWorkflowStore:
def get_definition(self, definition_id: str):
return object() if definition_id == "workflow-a" else None
class FakeWorkflowRunner:
def run(self, loaded_definition, device_id: str, *, should_stop=None, stop_reason=None):
assert should_stop is not None and should_stop()
assert stop_reason is not None
return SimpleNamespace(
id="run-a", status="cancelled", failure_reason=stop_reason()
)
factories = ExecutionFactories(
task_runner_factory=lambda: object(), # type: ignore[arg-type,return-value]
workflow_runner_factory=lambda: FakeWorkflowRunner(), # type: ignore[arg-type,return-value]
workflow_store=FakeWorkflowStore(), # type: ignore[arg-type]
)
result = AssignmentExecutor(factories).execute(
_assignment(goal=None, workflow_definition_id="workflow-a"),
should_stop=should_stop,
stop_reason=lambda: "cancellation requested by control plane",
)
assert result.status == "cancelled"
assert result.metadata == {
"workflow_run_id": "run-a",
"workflow_status": "cancelled",
}
def test_unknown_workflow_fails_without_running() -> None: def test_unknown_workflow_fails_without_running() -> None:
class FakeWorkflowStore: class FakeWorkflowStore:
def get_definition(self, definition_id: str): def get_definition(self, definition_id: str):
@@ -140,6 +140,30 @@ def test_stale_lease_response_raises_typed_error_without_retry() -> None:
assert attempts == 1 assert attempts == 1
def test_renew_deserializes_cancel_requested_flag() -> None:
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(
200,
json={
"status": "renewed",
"lease_expires_at": "2026-07-12T00:05:00Z",
"cancel_requested": True,
},
)
async def scenario() -> None:
async with httpx.AsyncClient(
transport=httpx.MockTransport(handler),
base_url="https://control.example",
) as http_client:
client = HostAgentClient(_config(), http_client=http_client)
response = await client.renew(_assignment())
assert response.status == "renewed"
assert response.cancel_requested is True
asyncio.run(scenario())
def test_result_report_retries_identical_payload_after_response_loss() -> None: def test_result_report_retries_identical_payload_after_response_loss() -> None:
payloads: list[dict[str, object]] = [] payloads: list[dict[str, object]] = []
+57 -7
View File
@@ -29,7 +29,7 @@ def test_lease_renews_while_execution_is_active() -> None:
renewed = asyncio.Event() renewed = asyncio.Event()
class BlockingExecutor: class BlockingExecutor:
def execute(self, assignment, *, should_stop=None): def execute(self, assignment, *, should_stop=None, stop_reason=None):
execution_started.set() execution_started.set()
release_execution.wait(timeout=2) release_execution.wait(timeout=2)
return AssignmentExecutionResult(status="done") return AssignmentExecutionResult(status="done")
@@ -61,13 +61,13 @@ def test_lease_renews_while_execution_is_active() -> None:
asyncio.run(scenario()) asyncio.run(scenario())
def test_stale_lease_stops_later_interruptible_actions() -> None: def test_cancel_requested_renewal_stops_execution_with_cancelled_status() -> None:
async def scenario() -> None: async def scenario() -> None:
first_action_started = Event() first_action_started = Event()
actions: list[str] = [] actions: list[str] = []
class CooperativeExecutor: class CooperativeExecutor:
def execute(self, assignment, *, should_stop=None): def execute(self, assignment, *, should_stop=None, stop_reason=None):
assert should_stop is not None assert should_stop is not None
actions.append("first") actions.append("first")
first_action_started.set() first_action_started.set()
@@ -76,9 +76,59 @@ def test_stale_lease_stops_later_interruptible_actions() -> None:
Event().wait(0.001) Event().wait(0.001)
if not should_stop(): if not should_stop():
actions.append("second") actions.append("second")
reason = stop_reason() if stop_reason is not None else None
return AssignmentExecutionResult( return AssignmentExecutionResult(
status="failed", status="cancelled" if reason and "cancel" in reason else "failed",
failure_reason="execution interrupted", failure_reason=reason,
)
def latest_progress(self):
return None
class CancellingClient:
async def renew(self, assignment, *, progress=None):
assert await asyncio.to_thread(first_action_started.wait, 1)
return LeaseRenewalResponse(
status="renewed",
lease_expires_at=datetime.now(UTC) + timedelta(seconds=30),
cancel_requested=True,
)
result = await asyncio.wait_for(
ActiveAssignmentRunner(
CancellingClient(), # type: ignore[arg-type]
CooperativeExecutor(),
).run(_assignment()),
timeout=1,
)
assert result.status == "cancelled"
assert result.failure_reason == "cancellation requested by control plane"
assert actions == ["first"]
asyncio.run(scenario())
def test_stale_lease_stops_later_interruptible_actions() -> None:
async def scenario() -> None:
first_action_started = Event()
actions: list[str] = []
class CooperativeExecutor:
def execute(self, assignment, *, should_stop=None, stop_reason=None):
assert should_stop is not None
actions.append("first")
first_action_started.set()
assert first_action_started.wait(timeout=1)
while not should_stop():
Event().wait(0.001)
if not should_stop():
actions.append("second")
reason = stop_reason() if stop_reason is not None else None
assert reason == "lease rejected by control plane"
return AssignmentExecutionResult(
status="cancelled" if reason and "cancel" in reason else "failed",
failure_reason=reason,
) )
def latest_progress(self): def latest_progress(self):
@@ -108,7 +158,7 @@ def test_renewal_loop_exits_when_execution_finishes() -> None:
renew_calls = 0 renew_calls = 0
class ImmediateExecutor: class ImmediateExecutor:
def execute(self, assignment, *, should_stop=None): def execute(self, assignment, *, should_stop=None, stop_reason=None):
return AssignmentExecutionResult(status="done") return AssignmentExecutionResult(status="done")
def latest_progress(self): def latest_progress(self):
@@ -142,7 +192,7 @@ def test_shutdown_request_stops_active_execution_cooperatively() -> None:
execution_started = Event() execution_started = Event()
class CooperativeExecutor: class CooperativeExecutor:
def execute(self, assignment, *, should_stop=None): def execute(self, assignment, *, should_stop=None, stop_reason=None):
assert should_stop is not None assert should_stop is not None
execution_started.set() execution_started.set()
while not should_stop(): while not should_stop():
@@ -88,6 +88,38 @@ def test_processor_preserves_runtime_failure_reason() -> None:
asyncio.run(scenario()) asyncio.run(scenario())
def test_processor_reports_cancelled_status_with_reason() -> None:
async def scenario() -> None:
reports: list[dict[str, object]] = []
class CancelledExecutor:
async def run(self, assignment):
return AssignmentExecutionResult(
status="cancelled",
failure_reason="cancellation requested by control plane",
metadata={"runtime_status": "cancelled"},
)
class RecordingClient:
async def report_result(self, assignment, **kwargs):
reports.append(kwargs)
return TerminalResultResponse(status="recorded")
result = await AssignmentProcessor(
RecordingClient(), # type: ignore[arg-type]
CancelledExecutor(),
).process(_assignment())
assert result.report_status == "recorded"
assert reports[0] == {
"status": "cancelled",
"failure_reason": "cancellation requested by control plane",
"result": {"runtime_status": "cancelled"},
}
asyncio.run(scenario())
def test_status_tracker_sees_started_then_finished_even_on_raise() -> None: def test_status_tracker_sees_started_then_finished_even_on_raise() -> None:
async def scenario() -> None: async def scenario() -> None:
tracker = AgentStatusTracker() tracker = AgentStatusTracker()
+6 -6
View File
@@ -29,12 +29,12 @@
## 4. Host Agent collaborative stop ## 4. Host Agent collaborative stop
- [ ] 4.1 Extend `LeaseGuard` (`host_agent/lease.py`) with a `reason` attribute already implied by `mark_lost(reason)` — confirm it's readable, add a `is_cancellation` helper or convention (e.g. reason string prefix) to distinguish cancellation from other lost-lease reasons. - [x] 4.1 Extend `LeaseGuard` (`host_agent/lease.py`) with a `reason` attribute already implied by `mark_lost(reason)` — confirm it's readable, add a `is_cancellation` helper or convention (e.g. reason string prefix) to distinguish cancellation from other lost-lease reasons.
- [ ] 4.2 Update `ActiveAssignmentRunner._renew_while_running()` (`host_agent/lease.py`): when `response.cancel_requested` is true, call `guard.mark_lost("cancellation requested by control plane")` instead of continuing the renewal loop. - [x] 4.2 Update `ActiveAssignmentRunner._renew_while_running()` (`host_agent/lease.py`): when `response.cancel_requested` is true, call `guard.mark_lost("cancellation requested by control plane")` instead of continuing the renewal loop.
- [ ] 4.3 Update `client.py`'s `renew()` to ensure `LeaseRenewalResponse.cancel_requested` deserializes correctly (should be automatic via Pydantic model update, but add a test). - [x] 4.3 Update `client.py`'s `renew()` to ensure `LeaseRenewalResponse.cancel_requested` deserializes correctly (should be automatic via Pydantic model update, but add a test).
- [ ] 4.4 Update `AssignmentExecutor._execute_goal()` and `_execute_workflow()` (`host_agent/assignment.py`) to map a cancellation-flavored stop to `AssignmentExecutionResult.status = "cancelled"` (new value alongside `"done"`/`"failed"`), reading the underlying `Task`/`WorkflowRun` status (`"cancelled"`) instead of collapsing it to `"failed"`. - [x] 4.4 Update `AssignmentExecutor._execute_goal()` and `_execute_workflow()` (`host_agent/assignment.py`) to map a cancellation-flavored stop to `AssignmentExecutionResult.status = "cancelled"` (new value alongside `"done"`/`"failed"`), reading the underlying `Task`/`WorkflowRun` status (`"cancelled"`) instead of collapsing it to `"failed"`.
- [ ] 4.5 Update `AssignmentProcessor.process()` (`host_agent/processor.py`) so its `status = "done" if execution.status == "done" else "failed"` mapping becomes a three-way mapping that preserves `"cancelled"`, and `report_result` is called with `status="cancelled"` in that case. - [x] 4.5 Update `AssignmentProcessor.process()` (`host_agent/processor.py`) so its `status = "done" if execution.status == "done" else "failed"` mapping becomes a three-way mapping that preserves `"cancelled"`, and `report_result` is called with `status="cancelled"` in that case.
- [ ] 4.6 Add/extend Host Agent tests covering: a renewal response with `cancel_requested=True` stops the active `should_stop`-driven loop; the resulting `AssignmentExecutionResult.status` and reported terminal status are `"cancelled"`; a lease-loss stop unrelated to cancellation still reports `"failed"`. - [x] 4.6 Add/extend Host Agent tests covering: a renewal response with `cancel_requested=True` stops the active `should_stop`-driven loop; the resulting `AssignmentExecutionResult.status` and reported terminal status are `"cancelled"`; a lease-loss stop unrelated to cancellation still reports `"failed"`.
## 5. Public SDK endpoint ## 5. Public SDK endpoint
@@ -58,6 +58,7 @@ class _FakeExecutor:
assignment: AssignmentModel, assignment: AssignmentModel,
*, *,
should_stop: Any | None = None, should_stop: Any | None = None,
stop_reason: Any | None = None,
) -> Any: ) -> Any:
from host_agent.assignment import AssignmentExecutionResult from host_agent.assignment import AssignmentExecutionResult