From c7f8d0c128046503d7d96f7718e08283c8ef30c1 Mon Sep 17 00:00:00 2001 From: Jerry Yan <792602257@qq.com> Date: Sun, 12 Jul 2026 18:12:33 +0800 Subject: [PATCH] feat(host-protocol): long poll assignments --- .../cloud-control-plane-integration/tasks.md | 2 +- .../cloud-platform/cloud/internal_api/api.py | 53 ++++++++++++++- tests/test_host_agent_internal_api.py | 68 +++++++++++++++++++ 3 files changed, 121 insertions(+), 2 deletions(-) diff --git a/openspec/changes/cloud-control-plane-integration/tasks.md b/openspec/changes/cloud-control-plane-integration/tasks.md index e1e4918..59feb68 100644 --- a/openspec/changes/cloud-control-plane-integration/tasks.md +++ b/openspec/changes/cloud-control-plane-integration/tasks.md @@ -37,7 +37,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`. -- [ ] 5.3 Add long-poll assignment delivery that returns at most one claimed task and produces a normal empty timeout response. +- [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. - [ ] 5.5 Add internal API integration tests for multi-host isolation, duplicate device ids, timeout behavior, stale attempts, and repeated result reports. diff --git a/packages/cloud-platform/cloud/internal_api/api.py b/packages/cloud-platform/cloud/internal_api/api.py index 7f45eb2..3027ee7 100644 --- a/packages/cloud-platform/cloud/internal_api/api.py +++ b/packages/cloud-platform/cloud/internal_api/api.py @@ -1,5 +1,8 @@ from __future__ import annotations +import asyncio +from collections.abc import Awaitable, Callable +from time import monotonic from typing import TYPE_CHECKING from fastapi import APIRouter, HTTPException, Request, status @@ -8,7 +11,13 @@ from cloud.auth import ( AuthProvider, HostAuthorizationError, ) -from cloud.internal_api.models import HeartbeatRequest, HeartbeatResponse +from cloud.internal_api.models import ( + AssignmentModel, + ClaimRequest, + ClaimResponse, + HeartbeatRequest, + HeartbeatResponse, +) from core.models import Device, utc_now if TYPE_CHECKING: @@ -20,7 +29,11 @@ def create_internal_router( pool: DevicePool, auth_provider: AuthProvider, version_prefix: str = "/internal/v1", + claim_poll_interval_seconds: float = 0.1, + 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") router = APIRouter(prefix=version_prefix, tags=["host-agent"]) def authorize_host(request: Request, host_id: str) -> None: @@ -66,6 +79,44 @@ def create_internal_router( received_at=utc_now(), ) + @router.post( + "/hosts/{host_id}/assignments/claim", + response_model=ClaimResponse, + ) + async def claim_assignment( + host_id: str, + payload: ClaimRequest, + request: Request, + ) -> ClaimResponse: + authorize_host(request, host_id) + if payload.host_id != host_id: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, + detail="claim host_id must match the request path", + ) + + deadline = monotonic() + payload.timeout_seconds + while True: + assignment = pool.store.claim_assignment(host_id=host_id, now=utc_now()) + if assignment is not None: + return ClaimResponse( + assignment=AssignmentModel( + task_id=assignment.task_id, + attempt=assignment.attempt, + lease_id=assignment.lease_id, + lease_expires_at=assignment.lease_expires_at, + host_id=assignment.host_id, + device_id=assignment.device_id, + goal=assignment.goal, + workflow_definition_id=assignment.workflow_definition_id, + ) + ) + + remaining = deadline - monotonic() + if remaining <= 0: + return ClaimResponse(timed_out=True) + await sleep(min(claim_poll_interval_seconds, remaining)) + return router diff --git a/tests/test_host_agent_internal_api.py b/tests/test_host_agent_internal_api.py index 23766eb..89bc243 100644 --- a/tests/test_host_agent_internal_api.py +++ b/tests/test_host_agent_internal_api.py @@ -1,6 +1,7 @@ from __future__ import annotations from datetime import UTC, datetime +from datetime import timedelta from fastapi import FastAPI from fastapi.testclient import TestClient @@ -9,6 +10,7 @@ from cloud.auth import BearerCredential, ConfiguredBearerAuthProvider from cloud.config import CloudConfig from cloud.internal_api.api import create_internal_router from cloud.pool import DevicePool, PooledDevice +from cloud.scheduler import ScheduledTask, TaskConstraints from cloud.store import CloudStore @@ -136,3 +138,69 @@ def test_host_token_cannot_submit_heartbeat_for_another_host(tmp_path) -> None: assert response.status_code == 403 assert pool.store.get_host("host-b") is None + + +def test_long_poll_claim_returns_at_most_one_owned_assignment(tmp_path) -> None: + client, pool = _build_client(tmp_path) + 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="claim-device", + host_id="host-a", + driver_type="wda", + status="idle", + synced_at=now, + ) + ], + ) + pool.store.enqueue_task( + ScheduledTask( + id="claim-task", + goal="open settings", + workflow_definition_id=None, + constraints=TaskConstraints(), + created_at=now, + ) + ) + pool.store.assign_task( + task_id="claim-task", + host_id="host-a", + device_id="claim-device", + lease_id="claim-lease", + lease_expires_at=now + timedelta(minutes=1), + now=now, + ) + + first = client.post( + "/internal/v1/hosts/host-a/assignments/claim", + headers={"Authorization": "Bearer token-a"}, + json={"host_id": "host-a", "timeout_seconds": 0}, + ) + second = client.post( + "/internal/v1/hosts/host-a/assignments/claim", + headers={"Authorization": "Bearer token-a"}, + json={"host_id": "host-a", "timeout_seconds": 0}, + ) + + assert first.status_code == 200 + assignment = first.json()["assignment"] + assert assignment["task_id"] == "claim-task" + assert assignment["lease_id"] == "claim-lease" + assert second.status_code == 200 + assert second.json() == {"assignment": None, "timed_out": True} + + +def test_empty_long_poll_timeout_is_normal_response(tmp_path) -> None: + client, _ = _build_client(tmp_path) + + response = client.post( + "/internal/v1/hosts/host-a/assignments/claim", + headers={"Authorization": "Bearer token-a"}, + json={"host_id": "host-a", "timeout_seconds": 0}, + ) + + assert response.status_code == 200 + assert response.json() == {"assignment": None, "timed_out": True}