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
2026-07-15 18:28:21 +08:00
parent 8a0d48eada
commit d3024b4810
9 changed files with 233 additions and 22 deletions
@@ -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):
@@ -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]] = []

@@ -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():
@@ -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()