From a24efc1043765c8f48405c9f80c66370cc8c598a Mon Sep 17 00:00:00 2001 From: Jerry Yan <792602257@qq.com> Date: Sun, 12 Jul 2026 18:15:04 +0800 Subject: [PATCH] feat(host-protocol): renew and complete leases --- .../cloud-control-plane-integration/tasks.md | 2 +- .../cloud-platform/cloud/internal_api/api.py | 103 ++++++++++++++ .../cloud/internal_api/models.py | 5 + tests/test_host_agent_internal_api.py | 133 ++++++++++++++++++ 4 files changed, 242 insertions(+), 1 deletion(-) diff --git a/openspec/changes/cloud-control-plane-integration/tasks.md b/openspec/changes/cloud-control-plane-integration/tasks.md index 59feb68..8b25e62 100644 --- a/openspec/changes/cloud-control-plane-integration/tasks.md +++ b/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 diff --git a/packages/cloud-platform/cloud/internal_api/api.py b/packages/cloud-platform/cloud/internal_api/api.py index 3027ee7..5b70d91 100644 --- a/packages/cloud-platform/cloud/internal_api/api.py +++ b/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, *, diff --git a/packages/cloud-platform/cloud/internal_api/models.py b/packages/cloud-platform/cloud/internal_api/models.py index 768118c..a76946c 100644 --- a/packages/cloud-platform/cloud/internal_api/models.py +++ b/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 diff --git a/tests/test_host_agent_internal_api.py b/tests/test_host_agent_internal_api.py index 89bc243..66757e6 100644 --- a/tests/test_host_agent_internal_api.py +++ b/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"