This commit is contained in:
@@ -25,8 +25,11 @@ from cloud.internal_api.models import (
|
||||
DeviceEnrollmentResponse,
|
||||
HeartbeatRequest,
|
||||
HeartbeatResponse,
|
||||
HostGovernancePolicyModel,
|
||||
HostEnrollmentRequest,
|
||||
HostEnrollmentResponse,
|
||||
HostTaskSubmissionRequest,
|
||||
HostTaskSubmissionResponse,
|
||||
LeaseRenewalRequest,
|
||||
LeaseRenewalResponse,
|
||||
PlannerDecisionError,
|
||||
@@ -47,6 +50,7 @@ from runtime.tool_specs import ToolSpec
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from cloud.pool import DevicePool
|
||||
from cloud.scheduler import TaskScheduler
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -60,6 +64,7 @@ def create_internal_router(
|
||||
lease_duration_seconds: float = 60.0,
|
||||
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
|
||||
planner_client_factory: Callable[[], ToolCallingClient] | None = None,
|
||||
scheduler: TaskScheduler | None = None,
|
||||
) -> APIRouter:
|
||||
if claim_poll_interval_seconds <= 0:
|
||||
raise ValueError("claim_poll_interval_seconds must be greater than zero")
|
||||
@@ -166,12 +171,71 @@ def create_internal_router(
|
||||
address=payload.address,
|
||||
allow_device_takeover=allow_device_takeover,
|
||||
)
|
||||
policy = pool.store.get_host_governance_policy(host_id)
|
||||
policy_revision = policy.revision if policy is not None else 0
|
||||
return HeartbeatResponse(
|
||||
host_id=host_id,
|
||||
accepted_devices=len(devices),
|
||||
received_at=utc_now(),
|
||||
policy_revision=policy_revision,
|
||||
policy=(
|
||||
HostGovernancePolicyModel(
|
||||
revision=policy.revision,
|
||||
self_submission_enabled=policy.self_submission_enabled,
|
||||
max_active_tasks=policy.max_active_tasks,
|
||||
daily_token_budget=policy.daily_token_budget,
|
||||
)
|
||||
if policy is not None and payload.policy_revision != policy.revision
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
@router.post(
|
||||
"/hosts/{host_id}/tasks",
|
||||
response_model=HostTaskSubmissionResponse,
|
||||
status_code=status.HTTP_201_CREATED,
|
||||
)
|
||||
def submit_host_task(
|
||||
host_id: str,
|
||||
payload: HostTaskSubmissionRequest,
|
||||
request: Request,
|
||||
) -> HostTaskSubmissionResponse:
|
||||
authorize_host(request, host_id)
|
||||
if payload.host_id != host_id:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
|
||||
detail="task host_id must match the request path",
|
||||
)
|
||||
if scheduler is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
||||
detail="task submission is unavailable",
|
||||
)
|
||||
policy = pool.store.get_host_governance_policy(host_id)
|
||||
if policy is not None and not policy.self_submission_enabled:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="Host self-submission is disabled",
|
||||
)
|
||||
if payload.device_id is not None and not any(
|
||||
device.host_id == host_id and device.device_id == payload.device_id
|
||||
for device in pool.list_devices()
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
|
||||
detail="target device is not owned by this Host",
|
||||
)
|
||||
from cloud.scheduler import TaskConstraints
|
||||
|
||||
task_id = scheduler.submit(
|
||||
goal=payload.goal,
|
||||
constraints=TaskConstraints(
|
||||
target_host_id=host_id,
|
||||
target_device_id=payload.device_id,
|
||||
),
|
||||
)
|
||||
return HostTaskSubmissionResponse(task_id=task_id)
|
||||
|
||||
@router.post(
|
||||
"/hosts/{host_id}/assignments/claim",
|
||||
response_model=ClaimResponse,
|
||||
@@ -285,6 +349,7 @@ def create_internal_router(
|
||||
@router.post(
|
||||
"/hosts/{host_id}/planner/decide",
|
||||
response_model=PlannerDecisionResponse,
|
||||
response_model_exclude_none=True,
|
||||
responses={
|
||||
status.HTTP_502_BAD_GATEWAY: {"model": PlannerDecisionError},
|
||||
},
|
||||
@@ -354,6 +419,9 @@ def create_internal_router(
|
||||
return PlannerDecisionResponse(
|
||||
tool_name=decision.tool_name,
|
||||
arguments=dict(decision.arguments),
|
||||
input_tokens=(decision.usage.input_tokens if decision.usage else None),
|
||||
output_tokens=(decision.usage.output_tokens if decision.usage else None),
|
||||
total_tokens=(decision.usage.total_tokens if decision.usage else None),
|
||||
)
|
||||
|
||||
return router
|
||||
|
||||
@@ -38,12 +38,22 @@ class HeartbeatRequest(BaseModel):
|
||||
host_id: str = Field(min_length=1)
|
||||
address: str | None = None
|
||||
devices: list[DeviceSnapshotModel] = Field(default_factory=list)
|
||||
policy_revision: int = Field(default=0, ge=0)
|
||||
|
||||
|
||||
class HostGovernancePolicyModel(BaseModel):
|
||||
revision: int = Field(ge=1)
|
||||
self_submission_enabled: bool
|
||||
max_active_tasks: int | None = None
|
||||
daily_token_budget: int | None = None
|
||||
|
||||
|
||||
class HeartbeatResponse(BaseModel):
|
||||
host_id: str
|
||||
accepted_devices: int
|
||||
received_at: datetime
|
||||
policy_revision: int = Field(default=0, ge=0)
|
||||
policy: HostGovernancePolicyModel | None = None
|
||||
|
||||
|
||||
class ClaimRequest(BaseModel):
|
||||
@@ -93,6 +103,16 @@ class TerminalResultResponse(BaseModel):
|
||||
status: Literal["recorded", "already_recorded"]
|
||||
|
||||
|
||||
class HostTaskSubmissionRequest(BaseModel):
|
||||
host_id: str = Field(min_length=1)
|
||||
goal: str = Field(min_length=1)
|
||||
device_id: str | None = Field(default=None, min_length=1)
|
||||
|
||||
|
||||
class HostTaskSubmissionResponse(BaseModel):
|
||||
task_id: str
|
||||
|
||||
|
||||
class StaleLeaseConflict(BaseModel):
|
||||
code: Literal["stale_lease"] = "stale_lease"
|
||||
detail: str
|
||||
@@ -116,6 +136,9 @@ class PlannerDecisionRequest(BaseModel):
|
||||
class PlannerDecisionResponse(BaseModel):
|
||||
tool_name: str
|
||||
arguments: dict[str, Any] = Field(default_factory=dict)
|
||||
input_tokens: int | None = Field(default=None, ge=0)
|
||||
output_tokens: int | None = Field(default=None, ge=0)
|
||||
total_tokens: int | None = Field(default=None, ge=0)
|
||||
|
||||
|
||||
class PlannerDecisionError(BaseModel):
|
||||
|
||||
Reference in New Issue
Block a user