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
@@ -10,7 +10,7 @@ from time import monotonic
from typing import TYPE_CHECKING
from uuid import uuid4
from fastapi import APIRouter, HTTPException, Request, status
from fastapi import APIRouter, HTTPException, Request, Response, status
from fastapi.responses import JSONResponse
from cloud.auth import (
@@ -29,6 +29,7 @@ from cloud.internal_api.models import (
HostGovernancePolicyModel,
HostEnrollmentRequest,
HostEnrollmentResponse,
HostTaskCancellationResponse,
HostTaskSubmissionRequest,
HostTaskSubmissionResponse,
LeaseRenewalRequest,
@@ -250,6 +251,50 @@ def create_internal_router(
)
return HostTaskSubmissionResponse(task_id=task_id)
@router.post(
"/hosts/{host_id}/tasks/{task_id}/cancel",
response_model=HostTaskCancellationResponse,
responses={
status.HTTP_202_ACCEPTED: {"model": HostTaskCancellationResponse},
},
)
def cancel_host_task(
host_id: str,
task_id: str,
request: Request,
response: Response,
) -> HostTaskCancellationResponse:
authorize_host(request, host_id)
if scheduler is None:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="task cancellation is unavailable",
)
task = scheduler.store.get_task(task_id)
if task is None or task.constraints.target_host_id != host_id:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"task {task_id!r} not found",
)
result = scheduler.store.request_task_cancellation(
task_id, requested_at=utc_now()
)
if result == "not_found":
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail=f"task {task_id!r} not found",
)
if result == "already_terminal":
raise HTTPException(
status_code=status.HTTP_409_CONFLICT,
detail=f"task {task_id!r} has already reached a terminal state",
)
task = scheduler.store.get_task(task_id)
assert task is not None
if result == "requested" and task.status != "cancelled":
response.status_code = status.HTTP_202_ACCEPTED
return HostTaskCancellationResponse(task_id=task_id, status=task.status)
@router.post(
"/hosts/{host_id}/assignments/claim",
response_model=ClaimResponse,
@@ -316,7 +361,7 @@ def create_internal_router(
summary=payload.progress.summary[:500],
updated_at=now,
)
renewal_status = pool.store.renew_lease(
renewal = pool.store.renew_lease(
task_id=task_id,
attempt=payload.attempt,
lease_id=payload.lease_id,
@@ -325,16 +370,17 @@ def create_internal_router(
now=now,
progress=progress_snapshot,
)
if renewal_status == "not_found":
if renewal.status == "not_found":
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail="assignment not found",
)
if renewal_status != "renewed":
if renewal.status != "renewed":
return _stale_lease_conflict("assignment lease is stale or expired")
return LeaseRenewalResponse(
status="renewed",
lease_expires_at=lease_expires_at,
cancel_requested=renewal.cancel_requested,
)
@router.post(