feat(cloud): attribute planner token usage
This commit is contained in:
@@ -373,6 +373,7 @@ def create_internal_router(
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
|
||||
detail="planner-decision host_id must match the request path",
|
||||
)
|
||||
_validate_planner_context(pool, host_id=host_id, payload=payload)
|
||||
|
||||
screenshot: bytes | None = None
|
||||
if payload.screenshot_base64 is not None:
|
||||
@@ -401,8 +402,8 @@ def create_internal_router(
|
||||
host_id=host_id,
|
||||
usage_day=now.date().isoformat(),
|
||||
reserved_tokens=planner_token_reservation_ceiling,
|
||||
task_id=None,
|
||||
attempt=None,
|
||||
task_id=payload.task_id,
|
||||
attempt=payload.attempt,
|
||||
created_at=now,
|
||||
expires_at=now + timedelta(seconds=planner_token_reservation_ttl_seconds),
|
||||
)
|
||||
@@ -485,6 +486,28 @@ def _validate_assignment_identity(
|
||||
)
|
||||
|
||||
|
||||
def _validate_planner_context(pool, *, host_id: str, payload: PlannerDecisionRequest) -> None:
|
||||
context_values = (payload.task_id, payload.attempt, payload.lease_id)
|
||||
if not any(value is not None for value in context_values):
|
||||
return
|
||||
if any(value is None for value in context_values):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
|
||||
detail="planner context requires task_id, attempt, and lease_id together",
|
||||
)
|
||||
attempts = pool.store.list_task_attempts(payload.task_id or "")
|
||||
if not any(
|
||||
attempt.attempt == payload.attempt
|
||||
and attempt.host_id == host_id
|
||||
and attempt.lease_id == payload.lease_id
|
||||
for attempt in attempts
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
|
||||
detail="planner context does not match a Host assignment",
|
||||
)
|
||||
|
||||
|
||||
def _stale_lease_conflict(detail: str) -> JSONResponse:
|
||||
return JSONResponse(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
|
||||
Reference in New Issue
Block a user