feat(cloud): enforce host governance budgets

packages/cloud-platform/cloud/sql_repository.py

@@ -24,6 +24,8 @@ from cloud.db_models import (
     UserRow,
     UserSessionRow,
     UserSubmissionPolicyRow,
+    TokenReservationRow,
+    TokenUsageEventRow,
 )
 from cloud.observability import current_correlation_id
 from core.models import utc_now
@@ -824,12 +826,21 @@ class SQLAlchemyCloudRepository:
         allowed_host_ids: tuple[str, ...] | None,
         allowed_device_targets: tuple[tuple[str, str], ...] | None,
         updated_at: datetime,
+        expected_revision: int | None = None,
     ) -> Any:
+        from cloud.governance import GovernancePolicyConflictError
+
         with self._sessions.begin() as session:
             if session.get(UserRow, user_id) is None:
                 raise KeyError(f"unknown user {user_id!r}")
-            row = session.get(UserSubmissionPolicyRow, user_id)
+            row = session.get(
+                UserSubmissionPolicyRow,
+                user_id,
+                with_for_update=self.engine.dialect.name == "postgresql",
+            )
             if row is None:
+                if expected_revision not in {None, 0}:
+                    raise GovernancePolicyConflictError("submission policy revision changed")
                 row = UserSubmissionPolicyRow(
                     user_id=user_id,
                     revision=1,
@@ -842,6 +853,11 @@ class SQLAlchemyCloudRepository:
                 )
                 session.add(row)
             else:
+                if (
+                    expected_revision is not None
+                    and expected_revision != row.revision
+                ):
+                    raise GovernancePolicyConflictError("submission policy revision changed")
                 row.revision += 1
                 row.submission_enabled = 1 if submission_enabled else 0
                 row.allowed_host_ids_json = _dump_optional_list(allowed_host_ids)
@@ -865,12 +881,21 @@ class SQLAlchemyCloudRepository:
         max_active_tasks: int | None,
         daily_token_budget: int | None,
         updated_at: datetime,
+        expected_revision: int | None = None,
     ) -> Any:
+        from cloud.governance import GovernancePolicyConflictError
+
         with self._sessions.begin() as session:
             if session.get(HostRow, host_id) is None:
                 raise KeyError(f"unknown host {host_id!r}")
-            row = session.get(HostGovernancePolicyRow, host_id)
+            row = session.get(
+                HostGovernancePolicyRow,
+                host_id,
+                with_for_update=self.engine.dialect.name == "postgresql",
+            )
             if row is None:
+                if expected_revision not in {None, 0}:
+                    raise GovernancePolicyConflictError("Host policy revision changed")
                 row = HostGovernancePolicyRow(
                     host_id=host_id,
                     revision=1,
@@ -881,6 +906,11 @@ class SQLAlchemyCloudRepository:
                 )
                 session.add(row)
             else:
+                if (
+                    expected_revision is not None
+                    and expected_revision != row.revision
+                ):
+                    raise GovernancePolicyConflictError("Host policy revision changed")
                 row.revision += 1
                 row.self_submission_enabled = 1 if self_submission_enabled else 0
                 row.max_active_tasks = max_active_tasks
@@ -889,6 +919,133 @@ class SQLAlchemyCloudRepository:
             session.flush()
             return _host_governance_policy_from_row(row)
 
+    def count_active_tasks_for_host(self, host_id: str) -> int:
+        with self._sessions() as session:
+            count = session.scalar(
+                select(func.count())
+                .select_from(ScheduledTaskRow)
+                .where(
+                    ScheduledTaskRow.assigned_host_id == host_id,
+                    ScheduledTaskRow.status.in_(("assigned", "dispatched")),
+                )
+            )
+        return int(count or 0)
+
+    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,
+    ) -> Any | None:
+        from cloud.governance import TokenBudgetExceededError
+
+        with self._sessions.begin() as session:
+            statement = select(HostGovernancePolicyRow).where(
+                HostGovernancePolicyRow.host_id == host_id
+            )
+            if self.engine.dialect.name == "postgresql":
+                statement = statement.with_for_update()
+            policy = session.scalars(statement).first()
+            if policy is None or policy.daily_token_budget is None:
+                return None
+            used = session.scalar(
+                select(func.coalesce(func.sum(TokenUsageEventRow.total_tokens), 0)).where(
+                    TokenUsageEventRow.host_id == host_id,
+                    TokenUsageEventRow.usage_day == usage_day,
+                )
+            )
+            reserved = session.scalar(
+                select(func.coalesce(func.sum(TokenReservationRow.reserved_tokens), 0)).where(
+                    TokenReservationRow.host_id == host_id,
+                    TokenReservationRow.usage_day == usage_day,
+                    TokenReservationRow.expires_at > _iso(created_at),
+                )
+            )
+            if int(used or 0) + int(reserved or 0) + reserved_tokens > policy.daily_token_budget:
+                raise TokenBudgetExceededError("Host daily token budget is exhausted")
+            row = TokenReservationRow(
+                id=reservation_id, host_id=host_id, usage_day=usage_day,
+                reserved_tokens=reserved_tokens, task_id=task_id, attempt=attempt,
+                created_at=_iso(created_at), expires_at=_iso(expires_at),
+            )
+            session.add(row)
+            session.flush()
+            return _token_reservation_from_row(row)
+
+    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,
+    ) -> Any | None:
+        with self._sessions.begin() as session:
+            row = session.get(
+                TokenReservationRow, reservation_id,
+                with_for_update=self.engine.dialect.name == "postgresql",
+            )
+            if row is None:
+                return None
+            event = TokenUsageEventRow(
+                id=event_id, host_id=row.host_id, usage_day=row.usage_day,
+                task_id=row.task_id, attempt=row.attempt, provider=provider, model=model,
+                input_tokens=input_tokens, output_tokens=output_tokens,
+                total_tokens=total_tokens, occurred_at=_iso(occurred_at),
+            )
+            session.add(event)
+            session.delete(row)
+            session.flush()
+            return _token_usage_event_from_row(event)
+
+    def cleanup_expired_token_reservations(self, *, now: datetime, limit: int) -> int:
+        with self._sessions.begin() as session:
+            rows = session.scalars(
+                select(TokenReservationRow)
+                .where(TokenReservationRow.expires_at <= _iso(now))
+                .order_by(TokenReservationRow.expires_at)
+                .limit(limit)
+            ).all()
+            for row in rows:
+                session.delete(row)
+            return len(rows)
+
+    def get_host_token_usage_summary(
+        self, *, host_id: str, usage_day: str, now: datetime,
+    ) -> Any:
+        from cloud.governance import TokenUsageSummary
+
+        with self._sessions() as session:
+            policy = session.get(HostGovernancePolicyRow, host_id)
+            used = session.scalar(
+                select(func.coalesce(func.sum(TokenUsageEventRow.total_tokens), 0)).where(
+                    TokenUsageEventRow.host_id == host_id,
+                    TokenUsageEventRow.usage_day == usage_day,
+                )
+            )
+            reserved = session.scalar(
+                select(func.coalesce(func.sum(TokenReservationRow.reserved_tokens), 0)).where(
+                    TokenReservationRow.host_id == host_id,
+                    TokenReservationRow.usage_day == usage_day,
+                    TokenReservationRow.expires_at > _iso(now),
+                )
+            )
+        return TokenUsageSummary(
+            host_id=host_id, usage_day=usage_day,
+            daily_token_budget=(policy.daily_token_budget if policy else None),
+            used_tokens=int(used or 0), reserved_tokens=int(reserved or 0),
+        )
+
     def list_reserved_device_ids(self, *, now: datetime) -> set[str]:
         with self._sessions() as session:
             device_ids = session.scalars(
@@ -930,6 +1087,24 @@ class SQLAlchemyCloudRepository:
             if device is None or device.status != "idle":
                 return None
 
+            policy_statement = select(HostGovernancePolicyRow).where(
+                HostGovernancePolicyRow.host_id == host_id
+            )
+            if self.engine.dialect.name == "postgresql":
+                policy_statement = policy_statement.with_for_update()
+            policy = session.scalars(policy_statement).first()
+            if policy is not None and policy.max_active_tasks is not None:
+                active_count = session.scalar(
+                    select(func.count())
+                    .select_from(ScheduledTaskRow)
+                    .where(
+                        ScheduledTaskRow.assigned_host_id == host_id,
+                        ScheduledTaskRow.status.in_(("assigned", "dispatched")),
+                    )
+                )
+                if int(active_count or 0) >= policy.max_active_tasks:
+                    return None
+
             active_reservation = session.scalar(
                 select(ScheduledTaskRow.id)
                 .where(
@@ -1376,6 +1551,28 @@ def _host_governance_policy_from_row(row: HostGovernancePolicyRow) -> Any:
     )
 
 
+def _token_reservation_from_row(row: TokenReservationRow) -> Any:
+    from cloud.governance import TokenReservation
+
+    return TokenReservation(
+        id=row.id, host_id=row.host_id, usage_day=row.usage_day,
+        reserved_tokens=row.reserved_tokens, task_id=row.task_id, attempt=row.attempt,
+        created_at=_parse_dt(row.created_at) or utc_now(),
+        expires_at=_parse_dt(row.expires_at) or utc_now(),
+    )
+
+
+def _token_usage_event_from_row(row: TokenUsageEventRow) -> Any:
+    from cloud.governance import TokenUsageEvent
+
+    return TokenUsageEvent(
+        id=row.id, host_id=row.host_id, usage_day=row.usage_day,
+        task_id=row.task_id, attempt=row.attempt, provider=row.provider, model=row.model,
+        input_tokens=row.input_tokens, output_tokens=row.output_tokens,
+        total_tokens=row.total_tokens, occurred_at=_parse_dt(row.occurred_at) or utc_now(),
+    )
+
+
 def _task_attempt_from_row(row: TaskAttemptRow) -> Any:
     from cloud.repository import TaskAttemptRecord