q792602257
bd6b7e64e2
Tests / Test passed: 926
# Conflicts: # packages/cloud-platform/cloud/schema.py
581 lines
15 KiB
Python
581 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timedelta
|
|
from typing import TYPE_CHECKING, Any, Literal, Protocol
|
|
|
|
if TYPE_CHECKING:
|
|
from cloud.plugins import PluginManifest
|
|
from cloud.pool import HostRegistration, PooledDevice
|
|
from cloud.scheduler import ScheduledTask, ScheduledTaskStatus
|
|
from cloud.user_auth import (
|
|
AuthAuditEvent,
|
|
AuthenticatedUserSession,
|
|
LoginThrottle,
|
|
UserAccount,
|
|
UserSession,
|
|
)
|
|
from cloud.governance import (
|
|
HostGovernancePolicy,
|
|
TokenReservation,
|
|
TokenUsageSummary,
|
|
TokenUsageEvent,
|
|
UserSubmissionPolicy,
|
|
)
|
|
from cloud.llm_providers import (
|
|
LlmProviderProfile,
|
|
LlmProviderSettings,
|
|
ProviderType,
|
|
)
|
|
|
|
|
|
AttemptStatus = Literal["assigned", "dispatched", "done", "failed", "expired"]
|
|
TerminalTaskStatus = Literal["done", "failed", "cancelled"]
|
|
ResultRecordStatus = Literal["recorded", "already_recorded", "conflict"]
|
|
LeaseRenewalStatus = Literal["renewed", "not_found", "conflict", "expired"]
|
|
CancellationRequestStatus = Literal[
|
|
"requested", "already_terminal", "already_requested", "not_found"
|
|
]
|
|
|
|
|
|
class HostEnrollmentConflictError(RuntimeError):
|
|
"""Raised when a Host enrollment cannot preserve its existing binding."""
|
|
|
|
|
|
class EnrollmentTokenConflictError(HostEnrollmentConflictError):
|
|
"""Raised when a one-time enrollment token is already bound elsewhere."""
|
|
|
|
|
|
class DeviceEnrollmentConflictError(RuntimeError):
|
|
"""Raised when a local device enrollment conflicts with stored identity."""
|
|
|
|
|
|
class UserConflictError(RuntimeError):
|
|
"""Raised when a user operation violates a durable identity invariant."""
|
|
|
|
|
|
class LastAdministratorConflictError(UserConflictError):
|
|
"""Raised when a write would remove the last enabled administrator."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class HostEnrollment:
|
|
host_id: str
|
|
agent_instance_id: str
|
|
display_name: str | None
|
|
enrolled_at: datetime
|
|
revoked_at: datetime | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DeviceEnrollment:
|
|
device_id: str
|
|
host_id: str
|
|
local_device_id: str
|
|
driver_type: str
|
|
name: str | None
|
|
capability_tags: list[str]
|
|
enrolled_at: datetime
|
|
revoked_at: datetime | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TaskAttemptRecord:
|
|
task_id: str
|
|
attempt: int
|
|
lease_id: str
|
|
host_id: str
|
|
device_id: str
|
|
status: AttemptStatus
|
|
lease_expires_at: datetime
|
|
created_at: datetime
|
|
completed_at: datetime | None = None
|
|
failure_reason: str | None = None
|
|
terminal_result: dict[str, Any] | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LeaseRenewalResult:
|
|
"""Outcome of a lease renewal, including whether cancellation is pending.
|
|
|
|
``cancel_requested`` reflects the task's durable ``cancel_requested_at``
|
|
column at renewal time regardless of ``status`` — callers only act on it
|
|
when ``status == "renewed"``.
|
|
"""
|
|
|
|
status: LeaseRenewalStatus
|
|
cancel_requested: bool = False
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LeasedAssignment:
|
|
task_id: str
|
|
attempt: int
|
|
lease_id: str
|
|
lease_expires_at: datetime
|
|
host_id: str
|
|
device_id: str
|
|
goal: str | None
|
|
workflow_definition_id: str | None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AssignmentProgressSnapshot:
|
|
"""Latest in-progress step snapshot for an active assignment.
|
|
|
|
Mirrors the wire ``TaskProgressModel`` fields; carried through the
|
|
repository layer without depending on Pydantic models.
|
|
|
|
``None`` on the ``renew_lease`` parameter means "no progress reported
|
|
for this renewal" (leave any previously stored progress untouched).
|
|
"""
|
|
|
|
step_index: int
|
|
step_status: str
|
|
summary: str
|
|
updated_at: datetime
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PlannerDecisionRecord:
|
|
"""One persisted planner decision row (D7/D8)."""
|
|
|
|
id: str
|
|
host_id: str
|
|
task_id: str
|
|
attempt: int
|
|
step_index: int
|
|
system_prompt: str
|
|
user_prompt: str
|
|
tool_name: str
|
|
arguments_json: str
|
|
created_at: datetime
|
|
rationale: str | None = None
|
|
thinking: str | None = None
|
|
purpose: str | None = None
|
|
expected_outcome: str | None = None
|
|
|
|
|
|
class CloudRepository(Protocol):
|
|
"""Persistence port for cloud state and atomic scheduling operations."""
|
|
|
|
def enroll_host(
|
|
self,
|
|
*,
|
|
host_id: str,
|
|
agent_instance_id: str,
|
|
credential_digest: str,
|
|
enrollment_token_digest: str | None,
|
|
display_name: str | None,
|
|
enrolled_at: datetime,
|
|
) -> HostEnrollment: ...
|
|
|
|
def authenticate_enrolled_host(
|
|
self,
|
|
credential_digest: str,
|
|
) -> str | None: ...
|
|
|
|
def revoke_enrolled_host(
|
|
self,
|
|
host_id: str,
|
|
*,
|
|
revoked_at: datetime,
|
|
) -> bool: ...
|
|
|
|
def is_enrollment_managed_host(self, host_id: str) -> bool: ...
|
|
|
|
def enroll_device(
|
|
self,
|
|
*,
|
|
device_id: str,
|
|
host_id: str,
|
|
local_device_id: str,
|
|
driver_type: str,
|
|
name: str | None,
|
|
capability_tags: list[str],
|
|
enrolled_at: datetime,
|
|
) -> DeviceEnrollment: ...
|
|
|
|
def get_device_enrollment(
|
|
self,
|
|
device_id: str,
|
|
) -> DeviceEnrollment | None: ...
|
|
|
|
def list_device_enrollments(
|
|
self,
|
|
host_id: str,
|
|
) -> list[DeviceEnrollment]: ...
|
|
|
|
def upsert_host(
|
|
self,
|
|
host_id: str,
|
|
*,
|
|
address: str | None,
|
|
last_seen_at: datetime,
|
|
planner_transport: Literal["direct", "cloud"] = "direct",
|
|
) -> None: ...
|
|
|
|
def replace_host_devices(
|
|
self,
|
|
host_id: str,
|
|
devices: list[PooledDevice],
|
|
*,
|
|
allow_device_takeover: bool = False,
|
|
) -> None: ...
|
|
|
|
def list_hosts(self) -> list[HostRegistration]: ...
|
|
|
|
def get_host(self, host_id: str) -> HostRegistration | None: ...
|
|
|
|
def list_devices(self) -> list[PooledDevice]: ...
|
|
|
|
def get_device(self, device_id: str) -> PooledDevice | None: ...
|
|
|
|
def enqueue_task(self, task: ScheduledTask) -> None: ...
|
|
|
|
def list_queued_tasks(self) -> list[ScheduledTask]: ...
|
|
|
|
def list_tasks(
|
|
self,
|
|
*,
|
|
status: ScheduledTaskStatus | None = None,
|
|
limit: int = 50,
|
|
offset: int = 0,
|
|
) -> list[ScheduledTask]: ...
|
|
|
|
def count_tasks(self, status: ScheduledTaskStatus | None = None) -> int: ...
|
|
|
|
def get_task(self, task_id: str) -> ScheduledTask | None: ...
|
|
|
|
def update_task(
|
|
self,
|
|
task_id: str,
|
|
*,
|
|
status: str | None = None,
|
|
assigned_device_id: str | None = None,
|
|
assigned_host_id: str | None = None,
|
|
) -> None: ...
|
|
|
|
def count_queued_tasks(self) -> int: ...
|
|
|
|
def save_plugin(self, manifest: PluginManifest, *, wired: bool) -> None: ...
|
|
|
|
def list_plugins(self) -> list[tuple[PluginManifest, bool]]: ...
|
|
|
|
def get_plugin(self, name: str) -> tuple[PluginManifest, bool] | None: ...
|
|
|
|
def create_user(self, user: UserAccount) -> UserAccount: ...
|
|
|
|
def get_user(self, user_id: str) -> UserAccount | None: ...
|
|
|
|
def get_user_by_normalized_username(
|
|
self,
|
|
username_normalized: str,
|
|
) -> UserAccount | None: ...
|
|
|
|
def list_users(self, *, limit: int, offset: int) -> list[UserAccount]: ...
|
|
|
|
def update_user(
|
|
self,
|
|
user_id: str,
|
|
*,
|
|
display_name: str | None = None,
|
|
role: str | None = None,
|
|
enabled: bool | None = None,
|
|
updated_at: datetime,
|
|
) -> UserAccount: ...
|
|
|
|
def rehash_user_password(
|
|
self,
|
|
user_id: str,
|
|
*,
|
|
password_hash: str,
|
|
updated_at: datetime,
|
|
) -> UserAccount: ...
|
|
|
|
def update_user_password(
|
|
self,
|
|
user_id: str,
|
|
*,
|
|
password_hash: str,
|
|
must_change_password: bool,
|
|
updated_at: datetime,
|
|
revoke_sessions: bool,
|
|
) -> UserAccount: ...
|
|
|
|
def mark_user_login(self, user_id: str, *, now: datetime) -> UserAccount: ...
|
|
|
|
def create_user_session(self, session: UserSession) -> None: ...
|
|
|
|
def get_authenticated_user_session(
|
|
self,
|
|
token_digest: str,
|
|
*,
|
|
now: datetime,
|
|
) -> AuthenticatedUserSession | None: ...
|
|
|
|
def touch_user_session(
|
|
self,
|
|
session_id: str,
|
|
*,
|
|
last_seen_at: datetime,
|
|
idle_expires_at: datetime,
|
|
) -> UserSession: ...
|
|
|
|
def revoke_user_session(self, session_id: str, *, revoked_at: datetime) -> bool: ...
|
|
|
|
def revoke_user_sessions(self, user_id: str, *, revoked_at: datetime) -> int: ...
|
|
|
|
def get_login_throttle(
|
|
self,
|
|
username_normalized: str,
|
|
client_bucket: str,
|
|
) -> LoginThrottle | None: ...
|
|
|
|
def record_login_failure(
|
|
self,
|
|
*,
|
|
username_normalized: str,
|
|
client_bucket: str,
|
|
now: datetime,
|
|
failure_limit: int,
|
|
failure_window: timedelta,
|
|
block_duration: timedelta,
|
|
) -> LoginThrottle: ...
|
|
|
|
def clear_login_throttle(
|
|
self, username_normalized: str, client_bucket: str
|
|
) -> None: ...
|
|
|
|
def record_auth_audit(self, event: AuthAuditEvent) -> None: ...
|
|
|
|
def cleanup_auth_state(self, *, now: datetime, limit: int) -> int: ...
|
|
|
|
def get_user_submission_policy(
|
|
self, user_id: str
|
|
) -> UserSubmissionPolicy | None: ...
|
|
|
|
def upsert_user_submission_policy(
|
|
self,
|
|
*,
|
|
user_id: str,
|
|
submission_enabled: bool,
|
|
allowed_host_ids: tuple[str, ...] | None,
|
|
allowed_device_targets: tuple[tuple[str, str], ...] | None,
|
|
expected_revision: int | None,
|
|
updated_at: datetime,
|
|
) -> UserSubmissionPolicy: ...
|
|
|
|
def get_host_governance_policy(
|
|
self, host_id: str
|
|
) -> HostGovernancePolicy | None: ...
|
|
|
|
def upsert_host_governance_policy(
|
|
self,
|
|
*,
|
|
host_id: str,
|
|
self_submission_enabled: bool,
|
|
max_active_tasks: int | None,
|
|
daily_token_budget: int | None,
|
|
expected_revision: int | None,
|
|
updated_at: datetime,
|
|
) -> HostGovernancePolicy: ...
|
|
|
|
def count_active_tasks_for_host(self, host_id: str) -> int: ...
|
|
|
|
def reserve_host_token_budget(
|
|
self,
|
|
*,
|
|
reservation_id: str,
|
|
host_id: str,
|
|
usage_day: str,
|
|
reserved_tokens: int,
|
|
task_id: str | None,
|
|
attempt: int | None,
|
|
created_at: datetime,
|
|
expires_at: datetime,
|
|
) -> TokenReservation | None: ...
|
|
|
|
def settle_host_token_reservation(
|
|
self,
|
|
*,
|
|
reservation_id: str,
|
|
event_id: str,
|
|
provider: str,
|
|
model: str,
|
|
input_tokens: int | None,
|
|
output_tokens: int | None,
|
|
total_tokens: int,
|
|
occurred_at: datetime,
|
|
) -> TokenUsageEvent | None: ...
|
|
|
|
def cleanup_expired_token_reservations(
|
|
self, *, now: datetime, limit: int
|
|
) -> int: ...
|
|
|
|
def get_host_token_usage_summary(
|
|
self,
|
|
*,
|
|
host_id: str,
|
|
usage_day: str,
|
|
now: datetime,
|
|
) -> TokenUsageSummary: ...
|
|
|
|
def list_host_token_usage_events(
|
|
self,
|
|
*,
|
|
host_id: str,
|
|
limit: int,
|
|
offset: int,
|
|
) -> list[TokenUsageEvent]: ...
|
|
|
|
def get_llm_provider_settings(self) -> LlmProviderSettings: ...
|
|
|
|
def list_llm_provider_profiles(self) -> list[LlmProviderProfile]: ...
|
|
|
|
def get_llm_provider_profile(
|
|
self, profile_id: str
|
|
) -> LlmProviderProfile | None: ...
|
|
|
|
def create_llm_provider_profile(
|
|
self, profile: LlmProviderProfile
|
|
) -> LlmProviderProfile: ...
|
|
|
|
def update_llm_provider_profile(
|
|
self,
|
|
profile_id: str,
|
|
*,
|
|
name: str,
|
|
name_normalized: str,
|
|
provider_type: ProviderType,
|
|
model: str,
|
|
base_url: str | None,
|
|
timeout_seconds: float,
|
|
api_key_ciphertext: str,
|
|
key_last_rotated_at: datetime,
|
|
enabled: bool,
|
|
expected_revision: int | None,
|
|
updated_at: datetime,
|
|
) -> LlmProviderProfile: ...
|
|
|
|
def activate_llm_provider_profile(
|
|
self,
|
|
profile_id: str,
|
|
*,
|
|
expected_settings_revision: int | None,
|
|
updated_at: datetime,
|
|
) -> LlmProviderSettings: ...
|
|
|
|
def delete_llm_provider_profile(
|
|
self,
|
|
profile_id: str,
|
|
*,
|
|
expected_revision: int | None,
|
|
) -> None: ...
|
|
|
|
def list_reserved_device_ids(self, *, now: datetime) -> set[str]: ...
|
|
|
|
def assign_task(
|
|
self,
|
|
*,
|
|
task_id: str,
|
|
host_id: str,
|
|
device_id: str,
|
|
lease_id: str,
|
|
lease_expires_at: datetime,
|
|
now: datetime,
|
|
) -> LeasedAssignment | None: ...
|
|
|
|
def claim_assignment(
|
|
self,
|
|
*,
|
|
host_id: str,
|
|
now: datetime,
|
|
) -> LeasedAssignment | None: ...
|
|
|
|
def renew_lease(
|
|
self,
|
|
*,
|
|
task_id: str,
|
|
attempt: int,
|
|
lease_id: str,
|
|
host_id: str,
|
|
lease_expires_at: datetime,
|
|
now: datetime,
|
|
progress: AssignmentProgressSnapshot | None = None,
|
|
) -> LeaseRenewalResult: ...
|
|
|
|
def record_task_result(
|
|
self,
|
|
*,
|
|
task_id: str,
|
|
attempt: int,
|
|
lease_id: str,
|
|
host_id: str,
|
|
status: TerminalTaskStatus,
|
|
failure_reason: str | None,
|
|
terminal_result: dict[str, Any] | None,
|
|
completed_at: datetime,
|
|
) -> ResultRecordStatus: ...
|
|
|
|
def request_task_cancellation(
|
|
self,
|
|
task_id: str,
|
|
*,
|
|
requested_at: datetime,
|
|
) -> CancellationRequestStatus: ...
|
|
|
|
def reap_expired_leases(
|
|
self,
|
|
*,
|
|
now: datetime,
|
|
max_attempts: int,
|
|
) -> list[str]: ...
|
|
|
|
def list_task_attempts(self, task_id: str) -> list[TaskAttemptRecord]: ...
|
|
|
|
def record_planner_decision(
|
|
self,
|
|
*,
|
|
host_id: str,
|
|
task_id: str,
|
|
attempt: int,
|
|
system_prompt: str,
|
|
user_prompt: str,
|
|
tool_name: str,
|
|
arguments_json: str,
|
|
now: datetime,
|
|
rationale: str | None = None,
|
|
thinking: str | None = None,
|
|
purpose: str | None = None,
|
|
expected_outcome: str | None = None,
|
|
) -> int:
|
|
"""Insert one planner-decision log row, returning the assigned step_index.
|
|
|
|
``step_index`` is assigned as ``max(step_index) + 1`` scoped to
|
|
``(task_id, attempt)`` within the same transaction.
|
|
"""
|
|
|
|
def prune_planner_decision_log(
|
|
self,
|
|
*,
|
|
now: datetime,
|
|
prune_after_terminal_seconds: int,
|
|
) -> int:
|
|
"""Delete decision-log rows for tasks terminal more than the window ago.
|
|
|
|
Returns the count of deleted rows.
|
|
"""
|
|
|
|
def list_planner_decisions(
|
|
self,
|
|
*,
|
|
task_id: str,
|
|
attempt: int,
|
|
) -> list[PlannerDecisionRecord]:
|
|
"""Return all decision-log rows for a (task_id, attempt), ordered by step."""
|
|
|
|
def health_check(self) -> None: ...
|
|
|
|
def close(self) -> None: ...
|