Merge branch 'worktree-task-cancellation': task cancellation feature
Tests / Test passed: 926

# Conflicts:
#	packages/cloud-platform/cloud/schema.py
This commit is contained in:
2026-07-15 19:39:28 +08:00
49 changed files with 2374 additions and 63 deletions
@@ -193,6 +193,9 @@ def make_task_detail_context(
task: dict[str, Any] | None = None,
task_rows: list[tuple[str, Any]] | None = None,
timeline_steps: list[dict[str, Any]] | None = None,
can_cancel: bool = False,
cancel_notice: str | None = None,
cancel_error: str | None = None,
) -> dict[str, Any]:
if task is None:
task = {
@@ -227,7 +230,11 @@ def make_task_detail_context(
return {
"title": "Task task-001",
"session": session,
"csrf_token": session.csrf_token,
"task": task,
"task_rows": task_rows,
"timeline_steps": timeline_steps,
"can_cancel": can_cancel,
"cancel_notice": cancel_notice,
"cancel_error": cancel_error,
}
@@ -77,6 +77,43 @@ def test_task_detail_renders(env, sample_session) -> None:
assert "<h2>Timeline</h2>" in html
def test_task_detail_shows_cancel_button_for_non_terminal_task(
env, sample_session
) -> None:
html = env.get_template("task_detail.html").render(
**make_task_detail_context(sample_session, can_cancel=True)
)
assert 'action="/tasks/task-001/cancel"' in html
assert "Cancel task" in html
def test_task_detail_hides_cancel_button_for_terminal_task(
env, sample_session
) -> None:
html = env.get_template("task_detail.html").render(
**make_task_detail_context(sample_session, can_cancel=False)
)
assert 'action="/tasks/task-001/cancel"' not in html
def test_task_detail_renders_cancel_notice_and_error(env, sample_session) -> None:
html = env.get_template("task_detail.html").render(
**make_task_detail_context(
sample_session,
cancel_notice="Cancellation requested. It may take a moment to take effect.",
)
)
assert 'id="cancel-notice"' in html
html = env.get_template("task_detail.html").render(
**make_task_detail_context(
sample_session,
cancel_error="Failed to request cancellation. Try again.",
)
)
assert 'id="cancel-error"' in html
# ---------------------------------------------------------------------------
# 6.4 XSS-probe tests (parametrised over templates with operator-influenced
# string fields set to <script>alert(1)</script>)
@@ -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]] = []
+57 -7
View File
@@ -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()
+164 -1
View File
@@ -3,10 +3,12 @@ from __future__ import annotations
import re
from collections.abc import Awaitable, Callable
from datetime import UTC, datetime
from typing import Any
from fastapi.testclient import TestClient
from cloud.internal_api.models import AssignmentModel
from core.models import Task
from device.manager import DeviceManager
from host_agent.client import HostAgentAPIError, HostTaskSubmissionUnknownError
from host_agent.config import HostAgentConfig
@@ -22,6 +24,7 @@ from storage.task_metadata import TaskMetadataStore
CSRF_PATTERN = re.compile(r'name="csrf_token" value="([^"]+)"')
TaskSubmissionCallable = Callable[[str, str | None], Awaitable[str]]
TaskCancellationCallable = Callable[[str], Awaitable[Any]]
def _build_client(
@@ -29,6 +32,7 @@ def _build_client(
*,
create_account: bool = True,
submit_self_task: TaskSubmissionCallable | None = None,
cancel_task: TaskCancellationCallable | None = None,
include_metadata_store: bool = True,
) -> tuple[TestClient, dict]:
config = HostAgentConfig(
@@ -62,6 +66,7 @@ def _build_client(
session_manager=session_manager,
enrollment_client=None,
submit_self_task=submit_self_task,
cancel_task=cancel_task,
metadata_store=metadata_store,
)
client = TestClient(app)
@@ -795,4 +800,162 @@ def test_submitted_redirect_does_not_include_goal_text(tmp_path) -> None:
assert "bearer-token-deadbeef" not in response.headers["location"]
follow = client.get(response.headers["location"])
assert secret_goal not in follow.text
assert "bearer-token-deadbeef" not in follow.text
def _make_cancellation_recorder(
*,
raise_api_error: HostAgentAPIError | None = None,
) -> tuple[TaskCancellationCallable, dict]:
captured: dict = {}
async def cancel(task_id: str) -> None:
captured["task_id"] = task_id
if raise_api_error is not None:
raise raise_api_error
return cancel, captured
def _seed_local_task(
metadata_store: TaskMetadataStore,
*,
status: str = "running",
source_task_id: str | None = "cloud-task-1",
) -> str:
task = Task(goal="open settings", device_id="dev-1", status=status)
metadata_store.create_task(
task, source_task_id=source_task_id, source_attempt=1
)
return task.id
def test_task_detail_page_shows_cancel_button_for_non_terminal_task(
tmp_path,
) -> None:
cancel, _ = _make_cancellation_recorder()
client, context = _build_client(tmp_path, cancel_task=cancel)
execution_id = _seed_local_task(context["metadata_store"], status="running")
_login(client)
response = client.get(f"/tasks/{execution_id}")
assert response.status_code == 200
assert f'action="/tasks/{execution_id}/cancel"' in response.text
assert "Cancel task" in response.text
def test_task_detail_page_hides_cancel_button_for_terminal_task(tmp_path) -> None:
cancel, _ = _make_cancellation_recorder()
client, context = _build_client(tmp_path, cancel_task=cancel)
execution_id = _seed_local_task(context["metadata_store"], status="completed")
_login(client)
response = client.get(f"/tasks/{execution_id}")
assert response.status_code == 200
assert f'action="/tasks/{execution_id}/cancel"' not in response.text
def test_task_detail_page_hides_cancel_button_when_client_unavailable(
tmp_path,
) -> None:
client, context = _build_client(tmp_path, cancel_task=None)
execution_id = _seed_local_task(context["metadata_store"], status="running")
_login(client)
response = client.get(f"/tasks/{execution_id}")
assert response.status_code == 200
assert f'action="/tasks/{execution_id}/cancel"' not in response.text
def test_cancel_task_success_calls_client_with_cloud_task_id_and_redirects(
tmp_path,
) -> None:
cancel, captured = _make_cancellation_recorder()
client, context = _build_client(tmp_path, cancel_task=cancel)
execution_id = _seed_local_task(
context["metadata_store"], status="running", source_task_id="cloud-task-99"
)
csrf_token = _login(client)
response = client.post(
f"/tasks/{execution_id}/cancel",
data={"csrf_token": csrf_token},
follow_redirects=False,
)
assert response.status_code == 303
assert response.headers["location"] == f"/tasks/{execution_id}?cancelled=1"
assert captured == {"task_id": "cloud-task-99"}
follow = client.get(response.headers["location"])
assert "cancel-notice" in follow.text
def test_cancel_task_client_error_redirects_with_cancel_error(tmp_path) -> None:
cancel, _ = _make_cancellation_recorder(
raise_api_error=HostAgentAPIError(502, "control plane unavailable")
)
client, context = _build_client(tmp_path, cancel_task=cancel)
execution_id = _seed_local_task(context["metadata_store"], status="running")
csrf_token = _login(client)
response = client.post(
f"/tasks/{execution_id}/cancel",
data={"csrf_token": csrf_token},
follow_redirects=False,
)
assert response.status_code == 303
assert response.headers["location"] == f"/tasks/{execution_id}?cancel_error=1"
follow = client.get(response.headers["location"])
assert "cancel-error" in follow.text
def test_cancel_task_unknown_execution_id_returns_404(tmp_path) -> None:
cancel, captured = _make_cancellation_recorder()
client, _ = _build_client(tmp_path, cancel_task=cancel)
csrf_token = _login(client)
response = client.post(
"/tasks/does-not-exist/cancel",
data={"csrf_token": csrf_token},
)
assert response.status_code == 404
assert captured == {}
def test_unauthenticated_cancel_redirects_to_login_without_calling_client(
tmp_path,
) -> None:
cancel, captured = _make_cancellation_recorder()
client, context = _build_client(tmp_path, cancel_task=cancel)
execution_id = _seed_local_task(context["metadata_store"], status="running")
response = client.post(
f"/tasks/{execution_id}/cancel",
data={},
follow_redirects=False,
)
assert response.status_code == 303
assert response.headers["location"] == "/login"
assert captured == {}
def test_cancel_task_without_csrf_token_is_rejected(tmp_path) -> None:
cancel, captured = _make_cancellation_recorder()
client, context = _build_client(tmp_path, cancel_task=cancel)
execution_id = _seed_local_task(context["metadata_store"], status="running")
_login(client)
response = client.post(
f"/tasks/{execution_id}/cancel",
data={"csrf_token": "wrong-token"},
)
assert response.status_code == 403
assert captured == {}