feat(cloud): add durable cancellation support to task repository
- Add nullable cancel_requested_at column (migration 0012) - Widen ScheduledTaskStatus/TerminalTaskStatus to include cancelled - Add CancellationRequestStatus + request_task_cancellation() to CloudRepository protocol and SQLAlchemy implementation - renew_lease() now returns LeaseRenewalResult, surfacing whether cancellation is pending, instead of a bare status string - reap_expired_leases() resolves pending-cancellation tasks to cancelled instead of requeuing/failing them - record_task_result() accepts cancelled and clears cancel_requested_at on any terminal write Note: internal_api/api.py's renew_assignment route still compares renew_lease()'s return value against a bare string; it will be updated in the next task (Internal Host<->Cloud protocol) to consume LeaseRenewalResult and populate the new cancel_requested wire field.
This commit is contained in:
@@ -967,7 +967,8 @@ def test_active_lease_renews_for_owning_host(database_url: str) -> None:
|
||||
now=now + timedelta(seconds=30),
|
||||
)
|
||||
|
||||
assert status == "renewed"
|
||||
assert status.status == "renewed"
|
||||
assert status.cancel_requested is False
|
||||
task = database.repository.get_task(task_id)
|
||||
assert task is not None
|
||||
assert task.lease_expires_at == renewed_expiry
|
||||
@@ -1031,7 +1032,7 @@ def test_stale_or_foreign_lease_renewal_conflicts(
|
||||
now=now + timedelta(seconds=30),
|
||||
)
|
||||
|
||||
assert status == "conflict"
|
||||
assert status.status == "conflict"
|
||||
task = database.repository.get_task(task_id)
|
||||
assert task is not None
|
||||
assert task.lease_expires_at == initial_expiry
|
||||
@@ -1078,7 +1079,7 @@ def test_expired_or_missing_lease_cannot_be_renewed(database_url: str) -> None:
|
||||
host_id=host_id,
|
||||
lease_expires_at=now + timedelta(minutes=2),
|
||||
now=now + timedelta(seconds=2),
|
||||
)
|
||||
).status
|
||||
== "expired"
|
||||
)
|
||||
assert (
|
||||
@@ -1089,7 +1090,7 @@ def test_expired_or_missing_lease_cannot_be_renewed(database_url: str) -> None:
|
||||
host_id=host_id,
|
||||
lease_expires_at=now + timedelta(minutes=2),
|
||||
now=now,
|
||||
)
|
||||
).status
|
||||
== "not_found"
|
||||
)
|
||||
finally:
|
||||
@@ -1810,3 +1811,251 @@ def test_record_planner_decision_stores_null_rationale_and_thinking(
|
||||
assert decisions[0].thinking is None
|
||||
finally:
|
||||
database.close()
|
||||
|
||||
|
||||
def test_cancel_queued_task_is_immediate(database_url: str) -> None:
|
||||
database = CloudDatabase(database_url)
|
||||
task_id = _unique_id("cancel-queued-task")
|
||||
now = datetime(2026, 7, 15, 8, 0, tzinfo=UTC)
|
||||
|
||||
try:
|
||||
database.repository.enqueue_task(
|
||||
ScheduledTask(
|
||||
id=task_id,
|
||||
goal="cancel before assignment",
|
||||
workflow_definition_id=None,
|
||||
constraints=TaskConstraints(),
|
||||
created_at=now,
|
||||
)
|
||||
)
|
||||
|
||||
status = database.repository.request_task_cancellation(
|
||||
task_id, requested_at=now
|
||||
)
|
||||
|
||||
assert status == "requested"
|
||||
task = database.repository.get_task(task_id)
|
||||
assert task is not None
|
||||
assert task.status == "cancelled"
|
||||
assert task.cancel_requested_at is None
|
||||
finally:
|
||||
database.close()
|
||||
|
||||
|
||||
def test_cancel_request_on_assigned_task_is_durable(database_url: str) -> None:
|
||||
database = CloudDatabase(database_url)
|
||||
host_id = _unique_id("cancel-durable-host")
|
||||
device_id = _unique_id("cancel-durable-device")
|
||||
task_id = _unique_id("cancel-durable-task")
|
||||
now = datetime(2026, 7, 15, 8, 0, tzinfo=UTC)
|
||||
|
||||
try:
|
||||
database.repository.upsert_host(host_id, address=None, last_seen_at=now)
|
||||
database.repository.replace_host_devices(
|
||||
host_id,
|
||||
[_device(device_id, host_id)],
|
||||
)
|
||||
database.repository.enqueue_task(
|
||||
ScheduledTask(
|
||||
id=task_id,
|
||||
goal="cancel while running",
|
||||
workflow_definition_id=None,
|
||||
constraints=TaskConstraints(),
|
||||
created_at=now,
|
||||
)
|
||||
)
|
||||
database.repository.assign_task(
|
||||
task_id=task_id,
|
||||
host_id=host_id,
|
||||
device_id=device_id,
|
||||
lease_id="durable-cancel-lease",
|
||||
lease_expires_at=now + timedelta(minutes=5),
|
||||
now=now,
|
||||
)
|
||||
|
||||
status = database.repository.request_task_cancellation(
|
||||
task_id, requested_at=now
|
||||
)
|
||||
assert status == "requested"
|
||||
|
||||
# Simulate a process restart by re-fetching the task from a fresh read.
|
||||
task = database.repository.get_task(task_id)
|
||||
assert task is not None
|
||||
assert task.status == "assigned"
|
||||
assert task.cancel_requested_at == now
|
||||
|
||||
repeat_status = database.repository.request_task_cancellation(
|
||||
task_id, requested_at=now + timedelta(seconds=5)
|
||||
)
|
||||
assert repeat_status == "already_requested"
|
||||
task = database.repository.get_task(task_id)
|
||||
assert task is not None
|
||||
assert task.cancel_requested_at == now
|
||||
finally:
|
||||
database.close()
|
||||
|
||||
|
||||
def test_renew_lease_reports_pending_cancellation(database_url: str) -> None:
|
||||
database = CloudDatabase(database_url)
|
||||
host_id = _unique_id("cancel-renew-host")
|
||||
device_id = _unique_id("cancel-renew-device")
|
||||
task_id = _unique_id("cancel-renew-task")
|
||||
now = datetime(2026, 7, 15, 8, 0, tzinfo=UTC)
|
||||
|
||||
try:
|
||||
database.repository.upsert_host(host_id, address=None, last_seen_at=now)
|
||||
database.repository.replace_host_devices(
|
||||
host_id,
|
||||
[_device(device_id, host_id)],
|
||||
)
|
||||
database.repository.enqueue_task(
|
||||
ScheduledTask(
|
||||
id=task_id,
|
||||
goal="report pending cancellation on renewal",
|
||||
workflow_definition_id=None,
|
||||
constraints=TaskConstraints(),
|
||||
created_at=now,
|
||||
)
|
||||
)
|
||||
database.repository.assign_task(
|
||||
task_id=task_id,
|
||||
host_id=host_id,
|
||||
device_id=device_id,
|
||||
lease_id="renew-cancel-lease",
|
||||
lease_expires_at=now + timedelta(minutes=5),
|
||||
now=now,
|
||||
)
|
||||
database.repository.request_task_cancellation(task_id, requested_at=now)
|
||||
|
||||
result = database.repository.renew_lease(
|
||||
task_id=task_id,
|
||||
attempt=1,
|
||||
lease_id="renew-cancel-lease",
|
||||
host_id=host_id,
|
||||
lease_expires_at=now + timedelta(minutes=10),
|
||||
now=now + timedelta(seconds=30),
|
||||
)
|
||||
|
||||
assert result.status == "renewed"
|
||||
assert result.cancel_requested is True
|
||||
finally:
|
||||
database.close()
|
||||
|
||||
|
||||
def test_expired_lease_with_pending_cancellation_resolves_to_cancelled(
|
||||
database_url: str,
|
||||
) -> None:
|
||||
database = CloudDatabase(database_url)
|
||||
host_id = _unique_id("cancel-expiry-host")
|
||||
device_id = _unique_id("cancel-expiry-device")
|
||||
task_id = _unique_id("cancel-expiry-task")
|
||||
now = datetime(2026, 7, 15, 8, 0, tzinfo=UTC)
|
||||
expired_at = now + timedelta(seconds=10)
|
||||
reaped_at = expired_at + timedelta(seconds=1)
|
||||
|
||||
try:
|
||||
database.repository.upsert_host(host_id, address=None, last_seen_at=now)
|
||||
database.repository.replace_host_devices(
|
||||
host_id,
|
||||
[_device(device_id, host_id)],
|
||||
)
|
||||
database.repository.enqueue_task(
|
||||
ScheduledTask(
|
||||
id=task_id,
|
||||
goal="cancelled task must not be requeued",
|
||||
workflow_definition_id=None,
|
||||
constraints=TaskConstraints(),
|
||||
created_at=now,
|
||||
)
|
||||
)
|
||||
database.repository.assign_task(
|
||||
task_id=task_id,
|
||||
host_id=host_id,
|
||||
device_id=device_id,
|
||||
lease_id="cancel-then-expire-lease",
|
||||
lease_expires_at=expired_at,
|
||||
now=now,
|
||||
)
|
||||
database.repository.request_task_cancellation(task_id, requested_at=now)
|
||||
|
||||
reaped_task_ids = database.repository.reap_expired_leases(
|
||||
now=reaped_at,
|
||||
# A high attempt limit proves cancellation takes priority over retry.
|
||||
max_attempts=5,
|
||||
)
|
||||
assert task_id in reaped_task_ids
|
||||
|
||||
task = database.repository.get_task(task_id)
|
||||
assert task is not None
|
||||
assert task.status == "cancelled"
|
||||
assert task.cancel_requested_at is None
|
||||
assert task.assigned_host_id is None
|
||||
assert task.assigned_device_id is None
|
||||
finally:
|
||||
database.close()
|
||||
|
||||
|
||||
def test_cancel_request_rejected_on_terminal_task(database_url: str) -> None:
|
||||
database = CloudDatabase(database_url)
|
||||
host_id = _unique_id("cancel-terminal-host")
|
||||
device_id = _unique_id("cancel-terminal-device")
|
||||
task_id = _unique_id("cancel-terminal-task")
|
||||
now = datetime(2026, 7, 15, 8, 0, tzinfo=UTC)
|
||||
|
||||
try:
|
||||
database.repository.upsert_host(host_id, address=None, last_seen_at=now)
|
||||
database.repository.replace_host_devices(
|
||||
host_id,
|
||||
[_device(device_id, host_id)],
|
||||
)
|
||||
database.repository.enqueue_task(
|
||||
ScheduledTask(
|
||||
id=task_id,
|
||||
goal="already finished",
|
||||
workflow_definition_id=None,
|
||||
constraints=TaskConstraints(),
|
||||
created_at=now,
|
||||
)
|
||||
)
|
||||
database.repository.assign_task(
|
||||
task_id=task_id,
|
||||
host_id=host_id,
|
||||
device_id=device_id,
|
||||
lease_id="terminal-lease",
|
||||
lease_expires_at=now + timedelta(minutes=5),
|
||||
now=now,
|
||||
)
|
||||
database.repository.record_task_result(
|
||||
task_id=task_id,
|
||||
attempt=1,
|
||||
lease_id="terminal-lease",
|
||||
host_id=host_id,
|
||||
status="done",
|
||||
failure_reason=None,
|
||||
terminal_result={"ok": True},
|
||||
completed_at=now + timedelta(seconds=10),
|
||||
)
|
||||
|
||||
status = database.repository.request_task_cancellation(
|
||||
task_id, requested_at=now + timedelta(seconds=20)
|
||||
)
|
||||
|
||||
assert status == "already_terminal"
|
||||
task = database.repository.get_task(task_id)
|
||||
assert task is not None
|
||||
assert task.status == "done"
|
||||
finally:
|
||||
database.close()
|
||||
|
||||
|
||||
def test_cancel_request_unknown_task_not_found(database_url: str) -> None:
|
||||
database = CloudDatabase(database_url)
|
||||
now = datetime(2026, 7, 15, 8, 0, tzinfo=UTC)
|
||||
|
||||
try:
|
||||
status = database.repository.request_task_cancellation(
|
||||
_unique_id("missing-cancel-task"), requested_at=now
|
||||
)
|
||||
assert status == "not_found"
|
||||
finally:
|
||||
database.close()
|
||||
|
||||
@@ -90,7 +90,8 @@ def test_renew_lease_writes_progress_on_success(tmp_path) -> None:
|
||||
now=now + timedelta(seconds=5),
|
||||
progress=progress,
|
||||
)
|
||||
assert result == "renewed"
|
||||
assert result.status == "renewed"
|
||||
assert result.cancel_requested is False
|
||||
|
||||
task = database.repository.get_task(task_id)
|
||||
assert task is not None
|
||||
|
||||
Reference in New Issue
Block a user