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
@@ -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