112 lines
2.9 KiB
Python
112 lines
2.9 KiB
Python
"""Durable Cloud-side submission and Host governance policies."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
|
|
|
|
class TaskSubmissionPolicyError(PermissionError):
|
|
"""Raised when a human user's policy disallows a task submission."""
|
|
|
|
|
|
class GovernancePolicyConflictError(RuntimeError):
|
|
"""Raised when a policy write was based on an obsolete revision."""
|
|
|
|
|
|
class TokenBudgetExceededError(RuntimeError):
|
|
"""Raised before a Cloud-proxied call would exceed a Host's token budget."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class UserSubmissionPolicy:
|
|
user_id: str
|
|
revision: int
|
|
submission_enabled: bool
|
|
allowed_host_ids: tuple[str, ...] | None
|
|
allowed_device_targets: tuple[tuple[str, str], ...] | None
|
|
updated_at: datetime
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class HostGovernancePolicy:
|
|
host_id: str
|
|
revision: int
|
|
self_submission_enabled: bool
|
|
max_active_tasks: int | None
|
|
daily_token_budget: int | None
|
|
updated_at: datetime
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TokenReservation:
|
|
id: str
|
|
host_id: str
|
|
usage_day: str
|
|
reserved_tokens: int
|
|
task_id: str | None
|
|
attempt: int | None
|
|
created_at: datetime
|
|
expires_at: datetime
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TokenUsageEvent:
|
|
id: str
|
|
host_id: str
|
|
usage_day: str
|
|
task_id: str | None
|
|
attempt: int | None
|
|
provider: str
|
|
model: str
|
|
input_tokens: int | None
|
|
output_tokens: int | None
|
|
total_tokens: int
|
|
occurred_at: datetime
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TokenUsageSummary:
|
|
host_id: str
|
|
usage_day: str
|
|
daily_token_budget: int | None
|
|
used_tokens: int
|
|
reserved_tokens: int
|
|
|
|
@property
|
|
def remaining_tokens(self) -> int | None:
|
|
if self.daily_token_budget is None:
|
|
return None
|
|
return max(0, self.daily_token_budget - self.used_tokens - self.reserved_tokens)
|
|
|
|
|
|
def enforce_user_submission_policy(
|
|
policy: UserSubmissionPolicy | None,
|
|
*,
|
|
target_host_id: str | None,
|
|
target_device_id: str | None,
|
|
) -> None:
|
|
if policy is None:
|
|
return
|
|
if not policy.submission_enabled:
|
|
raise TaskSubmissionPolicyError("task submission is disabled for this user")
|
|
restricted = (
|
|
policy.allowed_host_ids is not None
|
|
or policy.allowed_device_targets is not None
|
|
)
|
|
if not restricted:
|
|
return
|
|
if target_host_id is None:
|
|
raise TaskSubmissionPolicyError("an explicit permitted target is required")
|
|
if (
|
|
policy.allowed_host_ids is not None
|
|
and target_host_id not in policy.allowed_host_ids
|
|
):
|
|
raise TaskSubmissionPolicyError("target host is not permitted")
|
|
if policy.allowed_device_targets is not None:
|
|
if target_device_id is None or (
|
|
target_host_id,
|
|
target_device_id,
|
|
) not in policy.allowed_device_targets:
|
|
raise TaskSubmissionPolicyError("target device is not permitted")
|