feat(host-protocol): renew and complete leases
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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,
|
||||
*,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user