275 lines
8.8 KiB
Python
275 lines
8.8 KiB
Python
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,
|
|
HostAuthorizationError,
|
|
)
|
|
from cloud.internal_api.models import (
|
|
AssignmentModel,
|
|
ClaimRequest,
|
|
ClaimResponse,
|
|
HeartbeatRequest,
|
|
HeartbeatResponse,
|
|
LeaseRenewalRequest,
|
|
LeaseRenewalResponse,
|
|
StaleLeaseConflict,
|
|
TerminalResultRequest,
|
|
TerminalResultResponse,
|
|
)
|
|
from core.models import Device, utc_now
|
|
|
|
if TYPE_CHECKING:
|
|
from cloud.pool import DevicePool
|
|
|
|
|
|
def create_internal_router(
|
|
*,
|
|
pool: DevicePool,
|
|
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:
|
|
principal = auth_provider.authenticate(request)
|
|
if principal is None:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_401_UNAUTHORIZED,
|
|
detail="unauthorized",
|
|
headers={"WWW-Authenticate": "Bearer"},
|
|
)
|
|
try:
|
|
principal.require_host(host_id)
|
|
except HostAuthorizationError as exc:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN,
|
|
detail=str(exc),
|
|
) from exc
|
|
|
|
@router.put(
|
|
"/hosts/{host_id}/heartbeat",
|
|
response_model=HeartbeatResponse,
|
|
)
|
|
def heartbeat(
|
|
host_id: str,
|
|
payload: HeartbeatRequest,
|
|
request: Request,
|
|
) -> HeartbeatResponse:
|
|
authorize_host(request, host_id)
|
|
allow_device_takeover = _validate_snapshot(
|
|
pool,
|
|
host_id=host_id,
|
|
payload=payload,
|
|
)
|
|
devices = [
|
|
Device(
|
|
id=device.device_id,
|
|
driver_type=device.driver_type,
|
|
status=device.status,
|
|
capability_tags=list(device.capability_tags),
|
|
)
|
|
for device in payload.devices
|
|
]
|
|
pool.sync_host_devices(
|
|
host_id,
|
|
devices,
|
|
address=payload.address,
|
|
allow_device_takeover=allow_device_takeover,
|
|
)
|
|
return HeartbeatResponse(
|
|
host_id=host_id,
|
|
accepted_devices=len(devices),
|
|
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))
|
|
|
|
@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,
|
|
*,
|
|
host_id: str,
|
|
payload: HeartbeatRequest,
|
|
) -> bool:
|
|
if payload.host_id != host_id:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
|
|
detail="heartbeat host_id must match the request path",
|
|
)
|
|
device_ids = [device.device_id for device in payload.devices]
|
|
if len(device_ids) != len(set(device_ids)):
|
|
raise HTTPException(
|
|
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
|
|
detail="heartbeat snapshot contains duplicate device ids",
|
|
)
|
|
|
|
now = utc_now()
|
|
hosts = {host.host_id: host for host in pool.store.list_hosts()}
|
|
conflicts: list[str] = []
|
|
stale_owner_found = False
|
|
requested_ids = set(device_ids)
|
|
for device in pool.store.list_devices():
|
|
if device.device_id not in requested_ids or device.host_id == host_id:
|
|
continue
|
|
owner = hosts.get(device.host_id)
|
|
if owner is None:
|
|
continue
|
|
age_seconds = (now - owner.last_seen_at).total_seconds()
|
|
if age_seconds <= pool.config.stale_after_seconds:
|
|
conflicts.append(device.device_id)
|
|
else:
|
|
stale_owner_found = True
|
|
if conflicts:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_409_CONFLICT,
|
|
detail=f"device ownership conflict: {sorted(set(conflicts))}",
|
|
)
|
|
return stale_owner_found
|