agentic-mobile-control/packages/cloud-platform/cloud/sdk/governance_api.py

from __future__ import annotations

from uuid import uuid4

from fastapi import APIRouter, HTTPException, Query, Request, status

from cloud.auth import AuthProvider, GOVERNANCE_ADMIN_SCOPE, GOVERNANCE_READ_SCOPE
from cloud.governance import GovernancePolicyConflictError
from cloud.observability import current_correlation_id
from cloud.sdk.models import (
    HostGovernancePolicyRequest,
    HostGovernancePolicyResponse,
    HostTokenUsageSummaryResponse,
    TokenUsageEventResponse,
    UserSubmissionPolicyRequest,
    UserSubmissionPolicyResponse,
)
from cloud.user_auth import AuthAuditEvent
from core.models import utc_now


def create_governance_router(*, repository, auth_provider: AuthProvider) -> APIRouter:
    router = APIRouter(prefix="/v1", tags=["cloud-governance"])

    def authorize(request: Request, required_scope: str):
        principal = auth_provider.authenticate(request)
        if principal is None:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="unauthorized",
                headers={"WWW-Authenticate": "Bearer"},
            )
        if principal.must_change_password or not principal.has_scope(required_scope):
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail=f"missing required scope: {required_scope}",
            )
        return principal

    @router.get(
        "/users/{user_id}/submission-policy",
        response_model=UserSubmissionPolicyResponse,
    )
    def get_user_policy(user_id: str, request: Request) -> UserSubmissionPolicyResponse:
        authorize(request, GOVERNANCE_READ_SCOPE)
        policy = repository.get_user_submission_policy(user_id)
        if policy is None:
            raise HTTPException(status_code=404, detail="submission policy not found")
        return UserSubmissionPolicyResponse(
            user_id=policy.user_id,
            revision=policy.revision,
            submission_enabled=policy.submission_enabled,
            allowed_host_ids=list(policy.allowed_host_ids)
            if policy.allowed_host_ids is not None
            else None,
            allowed_device_targets=[
                {"host_id": host_id, "device_id": device_id}
                for host_id, device_id in policy.allowed_device_targets or ()
            ]
            if policy.allowed_device_targets is not None
            else None,
            updated_at=policy.updated_at,
        )

    @router.put(
        "/users/{user_id}/submission-policy",
        response_model=UserSubmissionPolicyResponse,
    )
    def put_user_policy(
        user_id: str,
        payload: UserSubmissionPolicyRequest,
        request: Request,
    ) -> UserSubmissionPolicyResponse:
        principal = authorize(request, GOVERNANCE_ADMIN_SCOPE)
        try:
            policy = repository.upsert_user_submission_policy(
                user_id=user_id,
                submission_enabled=payload.submission_enabled,
                allowed_host_ids=_unique_strings(payload.allowed_host_ids),
                allowed_device_targets=(
                    tuple(
                        (item.host_id, item.device_id)
                        for item in payload.allowed_device_targets
                    )
                    if payload.allowed_device_targets is not None
                    else None
                ),
                expected_revision=payload.expected_revision,
                updated_at=utc_now(),
            )
        except KeyError as exc:
            raise HTTPException(status_code=404, detail="user not found") from exc
        except GovernancePolicyConflictError as exc:
            raise HTTPException(status_code=409, detail=str(exc)) from exc
        _audit(repository, principal.id, user_id, "user_submission_policy_update")
        return UserSubmissionPolicyResponse(
            user_id=policy.user_id,
            revision=policy.revision,
            submission_enabled=policy.submission_enabled,
            allowed_host_ids=list(policy.allowed_host_ids)
            if policy.allowed_host_ids is not None
            else None,
            allowed_device_targets=[
                {"host_id": host_id, "device_id": device_id}
                for host_id, device_id in policy.allowed_device_targets or ()
            ]
            if policy.allowed_device_targets is not None
            else None,
            updated_at=policy.updated_at,
        )

    @router.get(
        "/hosts/{host_id}/governance-policy",
        response_model=HostGovernancePolicyResponse,
    )
    def get_host_policy(host_id: str, request: Request) -> HostGovernancePolicyResponse:
        authorize(request, GOVERNANCE_READ_SCOPE)
        policy = repository.get_host_governance_policy(host_id)
        if policy is None:
            raise HTTPException(status_code=404, detail="Host policy not found")
        return _host_response(policy)

    @router.get(
        "/hosts/{host_id}/token-usage",
        response_model=HostTokenUsageSummaryResponse,
    )
    def get_host_token_usage(
        host_id: str, request: Request
    ) -> HostTokenUsageSummaryResponse:
        authorize(request, GOVERNANCE_READ_SCOPE)
        now = utc_now()
        summary = repository.get_host_token_usage_summary(
            host_id=host_id, usage_day=now.date().isoformat(), now=now
        )
        return HostTokenUsageSummaryResponse(
            host_id=summary.host_id,
            usage_day=summary.usage_day,
            daily_token_budget=summary.daily_token_budget,
            used_tokens=summary.used_tokens,
            reserved_tokens=summary.reserved_tokens,
            remaining_tokens=summary.remaining_tokens,
        )

    @router.get(
        "/hosts/{host_id}/token-usage-events",
        response_model=list[TokenUsageEventResponse],
    )
    def list_host_token_usage_events(
        host_id: str,
        request: Request,
        limit: int = Query(default=50, ge=1, le=100),
        offset: int = Query(default=0, ge=0),
    ) -> list[TokenUsageEventResponse]:
        authorize(request, GOVERNANCE_READ_SCOPE)
        return [
            TokenUsageEventResponse(
                id=event.id,
                host_id=event.host_id,
                usage_day=event.usage_day,
                task_id=event.task_id,
                attempt=event.attempt,
                provider=event.provider,
                model=event.model,
                input_tokens=event.input_tokens,
                output_tokens=event.output_tokens,
                total_tokens=event.total_tokens,
                occurred_at=event.occurred_at,
            )
            for event in repository.list_host_token_usage_events(
                host_id=host_id, limit=limit, offset=offset
            )
        ]

    @router.put(
        "/hosts/{host_id}/governance-policy",
        response_model=HostGovernancePolicyResponse,
    )
    def put_host_policy(
        host_id: str,
        payload: HostGovernancePolicyRequest,
        request: Request,
    ) -> HostGovernancePolicyResponse:
        principal = authorize(request, GOVERNANCE_ADMIN_SCOPE)
        try:
            policy = repository.upsert_host_governance_policy(
                host_id=host_id,
                self_submission_enabled=payload.self_submission_enabled,
                max_active_tasks=payload.max_active_tasks,
                daily_token_budget=payload.daily_token_budget,
                expected_revision=payload.expected_revision,
                updated_at=utc_now(),
            )
        except KeyError as exc:
            raise HTTPException(status_code=404, detail="Host not found") from exc
        except GovernancePolicyConflictError as exc:
            raise HTTPException(status_code=409, detail=str(exc)) from exc
        _audit(repository, principal.id, host_id, "host_governance_policy_update")
        return _host_response(policy)

    return router


def _unique_strings(values: list[str] | None) -> tuple[str, ...] | None:
    if values is None:
        return None
    return tuple(dict.fromkeys(value for value in values if value))


def _host_response(policy) -> HostGovernancePolicyResponse:
    return HostGovernancePolicyResponse(
        host_id=policy.host_id,
        revision=policy.revision,
        self_submission_enabled=policy.self_submission_enabled,
        max_active_tasks=policy.max_active_tasks,
        daily_token_budget=policy.daily_token_budget,
        updated_at=policy.updated_at,
    )


def _audit(repository, actor_id: str, target: str, action: str) -> None:
    repository.record_auth_audit(
        AuthAuditEvent(
            id=uuid4().hex,
            occurred_at=utc_now(),
            actor_principal_id=actor_id,
            target_user_id=target if action.startswith("user_") else None,
            action=action,
            outcome="success",
            correlation_id=current_correlation_id(),
            metadata={"target": target},
        )
    )