feat(host-protocol): renew and complete leases

openspec/changes/cloud-control-plane-integration/tasks.md

@@ -38,7 +38,7 @@
 - [x] 5.1 Add versioned `/internal/v1` request/response models for heartbeat snapshots, long-poll claim, lease renewal, and terminal result reporting.
 - [x] 5.2 Add atomic heartbeat/snapshot validation and device ownership-conflict handling before delegating to `DevicePool`.
 - [x] 5.3 Add long-poll assignment delivery that returns at most one claimed task and produces a normal empty timeout response.
-- [ ] 5.4 Add lease-renewal and idempotent terminal-result endpoints with typed stale-lease conflicts.
+- [x] 5.4 Add lease-renewal and idempotent terminal-result endpoints with typed stale-lease conflicts.
 - [ ] 5.5 Add internal API integration tests for multi-host isolation, duplicate device ids, timeout behavior, stale attempts, and repeated result reports.
 
 ## 6. Cloud Control Plane Composition

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

@@ -2,10 +2,12 @@ from __future__ import annotations
 
 import asyncio
 from collections.abc import Awaitable, Callable
+from datetime import timedelta
 from time import monotonic
 from typing import TYPE_CHECKING
 
 from fastapi import APIRouter, HTTPException, Request, status
+from fastapi.responses import JSONResponse
 
 from cloud.auth import (
     AuthProvider,
@@ -17,6 +19,11 @@ from cloud.internal_api.models import (
     ClaimResponse,
     HeartbeatRequest,
     HeartbeatResponse,
+    LeaseRenewalRequest,
+    LeaseRenewalResponse,
+    StaleLeaseConflict,
+    TerminalResultRequest,
+    TerminalResultResponse,
 )
 from core.models import Device, utc_now
 
@@ -30,10 +37,13 @@ def create_internal_router(
     auth_provider: AuthProvider,
     version_prefix: str = "/internal/v1",
     claim_poll_interval_seconds: float = 0.1,
+    lease_duration_seconds: float = 60.0,
     sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
 ) -> APIRouter:
     if claim_poll_interval_seconds <= 0:
         raise ValueError("claim_poll_interval_seconds must be greater than zero")
+    if lease_duration_seconds <= 0:
+        raise ValueError("lease_duration_seconds must be greater than zero")
     router = APIRouter(prefix=version_prefix, tags=["host-agent"])
 
     def authorize_host(request: Request, host_id: str) -> None:
@@ -117,9 +127,102 @@ def create_internal_router(
                 return ClaimResponse(timed_out=True)
             await sleep(min(claim_poll_interval_seconds, remaining))
 
+    @router.post(
+        "/hosts/{host_id}/assignments/{task_id}/renew",
+        response_model=LeaseRenewalResponse,
+        responses={status.HTTP_409_CONFLICT: {"model": StaleLeaseConflict}},
+    )
+    def renew_assignment(
+        host_id: str,
+        task_id: str,
+        payload: LeaseRenewalRequest,
+        request: Request,
+    ):
+        authorize_host(request, host_id)
+        _validate_assignment_identity(
+            host_id=host_id,
+            task_id=task_id,
+            payload_host_id=payload.host_id,
+            payload_task_id=payload.task_id,
+        )
+        now = utc_now()
+        lease_expires_at = now + timedelta(seconds=lease_duration_seconds)
+        renewal_status = pool.store.renew_lease(
+            task_id=task_id,
+            attempt=payload.attempt,
+            lease_id=payload.lease_id,
+            host_id=host_id,
+            lease_expires_at=lease_expires_at,
+            now=now,
+        )
+        if renewal_status == "not_found":
+            raise HTTPException(
+                status_code=status.HTTP_404_NOT_FOUND,
+                detail="assignment not found",
+            )
+        if renewal_status != "renewed":
+            return _stale_lease_conflict("assignment lease is stale or expired")
+        return LeaseRenewalResponse(
+            status="renewed",
+            lease_expires_at=lease_expires_at,
+        )
+
+    @router.post(
+        "/hosts/{host_id}/assignments/{task_id}/result",
+        response_model=TerminalResultResponse,
+        responses={status.HTTP_409_CONFLICT: {"model": StaleLeaseConflict}},
+    )
+    def report_result(
+        host_id: str,
+        task_id: str,
+        payload: TerminalResultRequest,
+        request: Request,
+    ):
+        authorize_host(request, host_id)
+        _validate_assignment_identity(
+            host_id=host_id,
+            task_id=task_id,
+            payload_host_id=payload.host_id,
+            payload_task_id=payload.task_id,
+        )
+        result_status = pool.store.record_task_result(
+            task_id=task_id,
+            attempt=payload.attempt,
+            lease_id=payload.lease_id,
+            host_id=host_id,
+            status=payload.status,
+            failure_reason=payload.failure_reason,
+            terminal_result=payload.result,
+            completed_at=utc_now(),
+        )
+        if result_status == "conflict":
+            return _stale_lease_conflict("assignment lease is stale or superseded")
+        return TerminalResultResponse(status=result_status)
+
     return router
 
 
+def _validate_assignment_identity(
+    *,
+    host_id: str,
+    task_id: str,
+    payload_host_id: str,
+    payload_task_id: str,
+) -> None:
+    if payload_host_id != host_id or payload_task_id != task_id:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
+            detail="assignment identity must match the request path",
+        )
+
+
+def _stale_lease_conflict(detail: str) -> JSONResponse:
+    return JSONResponse(
+        status_code=status.HTTP_409_CONFLICT,
+        content=StaleLeaseConflict(detail=detail).model_dump(),
+    )
+
+
 def _validate_snapshot(
     pool: DevicePool,
     *,

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

@@ -70,3 +70,8 @@ class TerminalResultRequest(BaseModel):
 
 class TerminalResultResponse(BaseModel):
     status: Literal["recorded", "already_recorded"]
+
+
+class StaleLeaseConflict(BaseModel):
+    code: Literal["stale_lease"] = "stale_lease"
+    detail: str

tests/test_host_agent_internal_api.py

@@ -204,3 +204,136 @@ def test_empty_long_poll_timeout_is_normal_response(tmp_path) -> None:
 
     assert response.status_code == 200
     assert response.json() == {"assignment": None, "timed_out": True}
+
+
+def _seed_active_assignment(pool: DevicePool) -> datetime:
+    now = datetime.now(UTC)
+    pool.store.upsert_host("host-a", address=None, last_seen_at=now)
+    pool.store.replace_host_devices(
+        "host-a",
+        [
+            PooledDevice(
+                device_id="active-device",
+                host_id="host-a",
+                driver_type="wda",
+                status="idle",
+                synced_at=now,
+            )
+        ],
+    )
+    pool.store.enqueue_task(
+        ScheduledTask(
+            id="active-task",
+            goal="execute assignment",
+            workflow_definition_id=None,
+            constraints=TaskConstraints(),
+            created_at=now,
+        )
+    )
+    pool.store.assign_task(
+        task_id="active-task",
+        host_id="host-a",
+        device_id="active-device",
+        lease_id="active-lease",
+        lease_expires_at=now + timedelta(minutes=1),
+        now=now,
+    )
+    pool.store.claim_assignment(host_id="host-a", now=now)
+    return now
+
+
+def test_lease_renewal_extends_active_assignment(tmp_path) -> None:
+    client, pool = _build_client(tmp_path)
+    original_time = _seed_active_assignment(pool)
+
+    response = client.post(
+        "/internal/v1/hosts/host-a/assignments/active-task/renew",
+        headers={"Authorization": "Bearer token-a"},
+        json={
+            "host_id": "host-a",
+            "task_id": "active-task",
+            "attempt": 1,
+            "lease_id": "active-lease",
+        },
+    )
+
+    assert response.status_code == 200
+    assert response.json()["status"] == "renewed"
+    renewed_expiry = datetime.fromisoformat(response.json()["lease_expires_at"])
+    assert renewed_expiry > original_time + timedelta(seconds=30)
+
+
+def test_stale_renewal_returns_typed_conflict(tmp_path) -> None:
+    client, pool = _build_client(tmp_path)
+    _seed_active_assignment(pool)
+
+    response = client.post(
+        "/internal/v1/hosts/host-a/assignments/active-task/renew",
+        headers={"Authorization": "Bearer token-a"},
+        json={
+            "host_id": "host-a",
+            "task_id": "active-task",
+            "attempt": 1,
+            "lease_id": "stale-lease",
+        },
+    )
+
+    assert response.status_code == 409
+    assert response.json()["code"] == "stale_lease"
+
+
+def test_terminal_result_is_idempotent_through_internal_api(tmp_path) -> None:
+    client, pool = _build_client(tmp_path)
+    _seed_active_assignment(pool)
+    payload = {
+        "host_id": "host-a",
+        "task_id": "active-task",
+        "attempt": 1,
+        "lease_id": "active-lease",
+        "status": "done",
+        "result": {"steps": 4},
+    }
+
+    first = client.post(
+        "/internal/v1/hosts/host-a/assignments/active-task/result",
+        headers={"Authorization": "Bearer token-a"},
+        json=payload,
+    )
+    repeated = client.post(
+        "/internal/v1/hosts/host-a/assignments/active-task/result",
+        headers={"Authorization": "Bearer token-a"},
+        json=payload,
+    )
+
+    assert first.status_code == 200
+    assert first.json()["status"] == "recorded"
+    assert repeated.status_code == 200
+    assert repeated.json()["status"] == "already_recorded"
+    assert pool.store.get_task("active-task").status == "done"  # type: ignore[union-attr]
+
+
+def test_conflicting_repeated_result_returns_stale_lease_conflict(tmp_path) -> None:
+    client, pool = _build_client(tmp_path)
+    _seed_active_assignment(pool)
+    base_payload = {
+        "host_id": "host-a",
+        "task_id": "active-task",
+        "attempt": 1,
+        "lease_id": "active-lease",
+        "status": "done",
+        "result": {"steps": 4},
+    }
+    client.post(
+        "/internal/v1/hosts/host-a/assignments/active-task/result",
+        headers={"Authorization": "Bearer token-a"},
+        json=base_payload,
+    )
+
+    response = client.post(
+        "/internal/v1/hosts/host-a/assignments/active-task/result",
+        headers={"Authorization": "Bearer token-a"},
+        json={**base_payload, "status": "failed", "failure_reason": "late failure"},
+    )
+
+    assert response.status_code == 409
+    assert response.json()["code"] == "stale_lease"