73 lines
2.6 KiB
Python
73 lines
2.6 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import uuid
|
|
from collections.abc import Callable
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
from cloud.internal_api.models import (
|
|
ClaimResponse,
|
|
HeartbeatResponse,
|
|
HostTaskCancellationResponse,
|
|
HostTaskSubmissionResponse,
|
|
LeaseRenewalResponse,
|
|
TerminalResultResponse,
|
|
)
|
|
from host_agent.progress import TaskProgressSnapshot
|
|
|
|
|
|
class LocalHostAgentClient:
|
|
"""In-process task broker used when the Host Agent runs without Cloud."""
|
|
|
|
def __init__(self, *, host_id: str, device_ids: Callable[[], list[str]]) -> None:
|
|
self.host_id = host_id
|
|
self._device_ids = device_ids
|
|
self._queue: asyncio.Queue[tuple[str, str, str | None]] = asyncio.Queue()
|
|
self._cancelled: set[str] = set()
|
|
|
|
async def submit_self_task(self, *, goal: str, device_id: str | None = None) -> HostTaskSubmissionResponse:
|
|
task_id = f"local-{uuid.uuid4().hex}"
|
|
await self._queue.put((task_id, goal, device_id))
|
|
return HostTaskSubmissionResponse(task_id=task_id)
|
|
|
|
async def claim(self):
|
|
task_id, goal, requested_device = await self._queue.get()
|
|
devices = self._device_ids()
|
|
device_id = requested_device or (devices[0] if devices else "")
|
|
if not device_id:
|
|
return None
|
|
from cloud.internal_api.models import AssignmentModel
|
|
|
|
return AssignmentModel(
|
|
task_id=task_id,
|
|
attempt=1,
|
|
lease_id=f"local-lease-{uuid.uuid4().hex}",
|
|
lease_expires_at=datetime.now(UTC) + timedelta(days=3650),
|
|
host_id=self.host_id,
|
|
device_id=device_id,
|
|
goal=goal,
|
|
)
|
|
|
|
async def heartbeat(self, *args, **kwargs) -> HeartbeatResponse:
|
|
return HeartbeatResponse(
|
|
host_id=self.host_id,
|
|
accepted_devices=len(self._device_ids()),
|
|
received_at=datetime.now(UTC),
|
|
)
|
|
|
|
async def renew(self, assignment, *, progress: TaskProgressSnapshot | None = None) -> LeaseRenewalResponse:
|
|
return LeaseRenewalResponse(
|
|
status="renewed",
|
|
lease_expires_at=datetime.now(UTC) + timedelta(days=3650),
|
|
)
|
|
|
|
async def report_result(self, assignment, *, status: str, failure_reason: str | None = None, result: dict | None = None) -> TerminalResultResponse:
|
|
return TerminalResultResponse(status="recorded")
|
|
|
|
async def cancel_task(self, task_id: str) -> HostTaskCancellationResponse:
|
|
self._cancelled.add(task_id)
|
|
return HostTaskCancellationResponse(task_id=task_id, status="cancel_requested")
|
|
|
|
async def aclose(self) -> None:
|
|
return None
|