feat(host-protocol): long poll assignments

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.
 

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
 
 

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}