Merge branch 'worktree-task-cancellation': task cancellation feature

# Conflicts:
#	packages/cloud-platform/cloud/schema.py

packages/cloud-platform/cloud/internal_api/api.py

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

packages/cloud-platform/cloud/internal_api/models.py

@@ -95,6 +95,7 @@ class LeaseRenewalRequest(BaseModel):
 class LeaseRenewalResponse(BaseModel):
     status: Literal["renewed"]
     lease_expires_at: datetime
+    cancel_requested: bool = False
 
 
 class TerminalResultRequest(BaseModel):
@@ -102,7 +103,7 @@ class TerminalResultRequest(BaseModel):
     task_id: str = Field(min_length=1)
     attempt: int = Field(ge=1)
     lease_id: str = Field(min_length=1)
-    status: Literal["done", "failed"]
+    status: Literal["done", "failed", "cancelled"]
     failure_reason: str | None = None
     result: dict[str, Any] | None = None
 
@@ -121,6 +122,11 @@ class HostTaskSubmissionResponse(BaseModel):
     task_id: str
 
 
+class HostTaskCancellationResponse(BaseModel):
+    task_id: str
+    status: str
+
+
 class StaleLeaseConflict(BaseModel):
     code: Literal["stale_lease"] = "stale_lease"
     detail: str