diff --git a/apps/device-host-agent/host_agent/assignment.py b/apps/device-host-agent/host_agent/assignment.py index 448d202..42d3146 100644 --- a/apps/device-host-agent/host_agent/assignment.py +++ b/apps/device-host-agent/host_agent/assignment.py @@ -9,6 +9,7 @@ from core.models import Task from host_agent.execution import ExecutionFactories from host_agent.planner_context import bind_planner_execution_context from host_agent.progress import TaskProgressHolder, TaskProgressSnapshot +from runtime.task import is_cancellation_reason @dataclass(frozen=True) @@ -32,18 +33,24 @@ class AssignmentExecutor: assignment: AssignmentModel, *, should_stop: Callable[[], bool] | None = None, + stop_reason: Callable[[], str | None] | None = None, ) -> AssignmentExecutionResult: self._progress.clear() with bind_planner_execution_context(assignment): if should_stop is not None and should_stop(): + reason = stop_reason() if stop_reason is not None else None return AssignmentExecutionResult( - status="failed", - failure_reason="execution interrupted", + status="cancelled" if is_cancellation_reason(reason) else "failed", + failure_reason=reason or "execution interrupted", ) 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: - return self._execute_goal(assignment, should_stop=should_stop) + return self._execute_goal( + assignment, should_stop=should_stop, stop_reason=stop_reason + ) return AssignmentExecutionResult( status="failed", failure_reason="assignment has neither goal nor workflow definition", @@ -54,6 +61,7 @@ class AssignmentExecutor: assignment: AssignmentModel, *, should_stop: Callable[[], bool] | None, + stop_reason: Callable[[], str | None] | None, ) -> AssignmentExecutionResult: task = Task(goal=assignment.goal or "", device_id=assignment.device_id) if self.factories.metadata_store is not None: @@ -67,9 +75,11 @@ class AssignmentExecutor: if should_stop is None: completed = runner.run(task) else: - completed = runner.run(task, should_stop=should_stop) + completed = runner.run( + task, should_stop=should_stop, stop_reason=stop_reason + ) return AssignmentExecutionResult( - status="done" if completed.status == "completed" else "failed", + status=_terminal_status(completed.status), failure_reason=completed.failure_reason, metadata={ "runtime_task_id": completed.id, @@ -82,6 +92,7 @@ class AssignmentExecutor: assignment: AssignmentModel, *, should_stop: Callable[[], bool] | None, + stop_reason: Callable[[], str | None] | None, ) -> AssignmentExecutionResult: definition_id = assignment.workflow_definition_id or "" definition = self.factories.workflow_store.get_definition(definition_id) @@ -98,9 +109,10 @@ class AssignmentExecutor: definition, device_id=assignment.device_id, should_stop=should_stop, + stop_reason=stop_reason, ) return AssignmentExecutionResult( - status="done" if run.status == "completed" else "failed", + status=_terminal_status(run.status), failure_reason=( None if run.status == "completed" else f"workflow ended as {run.status}" ), @@ -109,3 +121,11 @@ class AssignmentExecutor: "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" diff --git a/apps/device-host-agent/host_agent/lease.py b/apps/device-host-agent/host_agent/lease.py index 29e9e1c..b8bd95c 100644 --- a/apps/device-host-agent/host_agent/lease.py +++ b/apps/device-host-agent/host_agent/lease.py @@ -12,6 +12,7 @@ from cloud.internal_api.models import AssignmentModel from host_agent.assignment import AssignmentExecutionResult from host_agent.client import HostAgentAPIError, HostAgentClient, StaleLeaseError from host_agent.progress import TaskProgressSnapshot +from runtime.task import is_cancellation_reason class InterruptibleAssignmentExecutor(Protocol): @@ -20,6 +21,7 @@ class InterruptibleAssignmentExecutor(Protocol): assignment: AssignmentModel, *, should_stop: Callable[[], bool] | None = None, + stop_reason: Callable[[], str | None] | None = None, ) -> AssignmentExecutionResult: ... def latest_progress(self) -> TaskProgressSnapshot | None: ... @@ -36,6 +38,10 @@ class LeaseGuard: with self._lock: return self._reason + @property + def is_cancellation(self) -> bool: + return is_cancellation_reason(self.reason) + def is_lost(self) -> bool: return self._lost.is_set() @@ -69,6 +75,7 @@ class ActiveAssignmentRunner: self.executor.execute, assignment, should_stop=lambda: guard.is_lost() or self._stop_requested.is_set(), + stop_reason=lambda: guard.reason, ) ) renewal = asyncio.create_task( @@ -108,4 +115,7 @@ class ActiveAssignmentRunner: guard.mark_lost("lease renewal failed after transport retries") return else: + if response.cancel_requested: + guard.mark_lost("cancellation requested by control plane") + return lease_expires_at = response.lease_expires_at diff --git a/apps/device-host-agent/host_agent/processor.py b/apps/device-host-agent/host_agent/processor.py index d17f4e2..5e068c7 100644 --- a/apps/device-host-agent/host_agent/processor.py +++ b/apps/device-host-agent/host_agent/processor.py @@ -47,8 +47,11 @@ class AssignmentProcessor: self.status_tracker.mark_assignment_started(assignment) try: execution = await self.active_executor.run(assignment) - status = "done" if execution.status == "done" else "failed" - failure_reason = execution.failure_reason if status == "failed" else None + if execution.status in {"done", "cancelled"}: + status = execution.status + else: + status = "failed" + failure_reason = execution.failure_reason if status != "done" else None response = await self.client.report_result( assignment, status=status, diff --git a/apps/device-host-agent/tests/test_assignment.py b/apps/device-host-agent/tests/test_assignment.py index b6b6162..9818ce2 100644 --- a/apps/device-host-agent/tests/test_assignment.py +++ b/apps/device-host-agent/tests/test_assignment.py @@ -78,6 +78,39 @@ def test_goal_assignment_preserves_runtime_failure_reason() -> None: 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: definition = object() 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: class FakeWorkflowStore: def get_definition(self, definition_id: str): diff --git a/apps/device-host-agent/tests/test_client.py b/apps/device-host-agent/tests/test_client.py index 1f84ead..e8c70bd 100644 --- a/apps/device-host-agent/tests/test_client.py +++ b/apps/device-host-agent/tests/test_client.py @@ -140,6 +140,30 @@ def test_stale_lease_response_raises_typed_error_without_retry() -> None: 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: payloads: list[dict[str, object]] = [] diff --git a/apps/device-host-agent/tests/test_lease.py b/apps/device-host-agent/tests/test_lease.py index e21cc45..03f2aa6 100644 --- a/apps/device-host-agent/tests/test_lease.py +++ b/apps/device-host-agent/tests/test_lease.py @@ -29,7 +29,7 @@ def test_lease_renews_while_execution_is_active() -> None: renewed = asyncio.Event() class BlockingExecutor: - def execute(self, assignment, *, should_stop=None): + def execute(self, assignment, *, should_stop=None, stop_reason=None): execution_started.set() release_execution.wait(timeout=2) return AssignmentExecutionResult(status="done") @@ -61,13 +61,13 @@ def test_lease_renews_while_execution_is_active() -> None: 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: first_action_started = Event() actions: list[str] = [] 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 actions.append("first") first_action_started.set() @@ -76,9 +76,59 @@ def test_stale_lease_stops_later_interruptible_actions() -> None: Event().wait(0.001) if not should_stop(): actions.append("second") + reason = stop_reason() if stop_reason is not None else None return AssignmentExecutionResult( - status="failed", - failure_reason="execution interrupted", + status="cancelled" if reason and "cancel" in reason else "failed", + 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): @@ -108,7 +158,7 @@ def test_renewal_loop_exits_when_execution_finishes() -> None: renew_calls = 0 class ImmediateExecutor: - def execute(self, assignment, *, should_stop=None): + def execute(self, assignment, *, should_stop=None, stop_reason=None): return AssignmentExecutionResult(status="done") def latest_progress(self): @@ -142,7 +192,7 @@ def test_shutdown_request_stops_active_execution_cooperatively() -> None: execution_started = Event() 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 execution_started.set() while not should_stop(): diff --git a/apps/device-host-agent/tests/test_processor.py b/apps/device-host-agent/tests/test_processor.py index c0f8be2..c92ad6b 100644 --- a/apps/device-host-agent/tests/test_processor.py +++ b/apps/device-host-agent/tests/test_processor.py @@ -88,6 +88,38 @@ def test_processor_preserves_runtime_failure_reason() -> None: 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: async def scenario() -> None: tracker = AgentStatusTracker() diff --git a/openspec/changes/task-cancellation/tasks.md b/openspec/changes/task-cancellation/tasks.md index 96c3c1a..6803e7e 100644 --- a/openspec/changes/task-cancellation/tasks.md +++ b/openspec/changes/task-cancellation/tasks.md @@ -29,12 +29,12 @@ ## 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. -- [ ] 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). -- [ ] 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. -- [ ] 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.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.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.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.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.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.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 diff --git a/tests/test_host_agent_progress_reporting.py b/tests/test_host_agent_progress_reporting.py index ce114f6..acaef96 100644 --- a/tests/test_host_agent_progress_reporting.py +++ b/tests/test_host_agent_progress_reporting.py @@ -58,6 +58,7 @@ class _FakeExecutor: assignment: AssignmentModel, *, should_stop: Any | None = None, + stop_reason: Any | None = None, ) -> Any: from host_agent.assignment import AssignmentExecutionResult