DataOps-platform/app/core/llm/runtime_server.py

"""Server-derived invocation contract for the hardened WP09 runtime path."""

from __future__ import annotations

import hashlib
from dataclasses import dataclass
from decimal import Decimal, InvalidOperation
from typing import Any


class RuntimeServerError(ValueError):
    pass


@dataclass(frozen=True)
class RuntimeServerContext:
    agent_uid: str
    actor_uid: str
    roles: frozenset[str]


_PAYLOAD_FIELDS = frozenset(
    {
        "grant_uid", "idempotency_key", "input_text", "evidence_refs",
        "estimated_tokens", "estimated_cost_micros", "requested_time_ms",
    }
)
_OPTIONAL_APPROVAL_FIELD = "approval_task_uid"
_SCOPE_FIELDS = frozenset(
    {
        "tenant_id", "principal_id", "business_domain_uid", "environment", "provider",
        "model", "prompt_version", "generation", "interface_type", "tool_name", "action",
        "risk_level", "agent_uid",
    }
)


def _positive_integer(value: Any) -> int:
    if isinstance(value, bool) or not isinstance(value, int) or not 0 < value <= 1_000_000:
        raise RuntimeServerError("budget_value_invalid")
    return value


def _cost_micros(value: Any) -> int:
    if isinstance(value, bool) or not isinstance(value, str) or len(value) > 32:
        raise RuntimeServerError("budget_value_invalid")
    try:
        decimal = Decimal(value)
    except (InvalidOperation, ValueError) as error:
        raise RuntimeServerError("budget_value_invalid") from error
    if not decimal.is_finite() or decimal <= 0 or decimal != decimal.to_integral_value():
        raise RuntimeServerError("budget_value_invalid")
    result = int(decimal)
    if result > 1_000_000_000:
        raise RuntimeServerError("budget_value_invalid")
    return result


class ServerGovernedInvocationService:
    """Reject client-supplied scope, deriving every authority field server-side."""

    def __init__(self, repository):
        self.repository = repository

    def authorize(self, context: RuntimeServerContext, payload: Any) -> dict[str, Any]:
        if not isinstance(payload, dict):
            raise RuntimeServerError("payload_invalid")
        if set(payload) & _SCOPE_FIELDS:
            raise RuntimeServerError("payload_scope_forbidden")
        if set(payload) not in {_PAYLOAD_FIELDS, _PAYLOAD_FIELDS | {_OPTIONAL_APPROVAL_FIELD}}:
            raise RuntimeServerError("payload_schema_invalid")
        if not context.agent_uid or not context.actor_uid or not context.roles:
            raise RuntimeServerError("server_identity_invalid")
        grant_uid = payload["grant_uid"]
        if not isinstance(grant_uid, str) or not grant_uid:
            raise RuntimeServerError("payload_invalid")
        idempotency_key = payload["idempotency_key"]
        if not isinstance(idempotency_key, str) or not idempotency_key or len(idempotency_key) > 160:
            raise RuntimeServerError("payload_invalid")
        input_text = payload["input_text"]
        if not isinstance(input_text, str) or not input_text.strip() or len(input_text) > 8000:
            raise RuntimeServerError("payload_invalid")
        evidence_refs = payload["evidence_refs"]
        if not isinstance(evidence_refs, list) or len(evidence_refs) > 20:
            raise RuntimeServerError("payload_invalid")
        derived = self.repository.grant_context(
            agent_uid=context.agent_uid, grant_uid=grant_uid, actor_uid=context.actor_uid
        )
        requires_approval = derived.get("risk_level") != "low" or derived.get("action") == "execute"
        approval_task_uid = payload.get(_OPTIONAL_APPROVAL_FIELD)
        if requires_approval and (not isinstance(approval_task_uid, str) or len(approval_task_uid) != 36):
            raise RuntimeServerError("human_approval_required")
        if not requires_approval and approval_task_uid is not None:
            raise RuntimeServerError("approval_not_applicable")
        request = {
            **derived,
            "agent_uid": context.agent_uid,
            "grant_uid": grant_uid,
            "idempotency_key": idempotency_key,
            "request_digest": hashlib.sha256(
                (input_text.strip() + "|" + repr(evidence_refs)).encode("utf-8")
            ).hexdigest(),
            "input_hash": hashlib.sha256(input_text.strip().encode("utf-8")).hexdigest(),
            "evidence_digests": [],
            "estimated_tokens": _positive_integer(payload["estimated_tokens"]),
            "estimated_cost_micros": _cost_micros(payload["estimated_cost_micros"]),
            "requested_time_ms": _positive_integer(payload["requested_time_ms"]),
        }
        if requires_approval:
            request["approval_task_uid"] = approval_task_uid
        return self.repository.authorize_claim(request)