DataOps-platform/tests/agent/test_wp09_critical_remediation.py

from __future__ import annotations

import uuid

import pytest

from app.core.llm.runtime_server import (
    RuntimeServerContext,
    RuntimeServerError,
    ServerGovernedInvocationService,
)


class _Repository:
    def __init__(self):
        self.calls = []

    def grant_context(self, *, agent_uid, grant_uid, actor_uid):
        if agent_uid != "agent-1" or grant_uid != "grant-1" or actor_uid != "owner-1":
            raise RuntimeServerError("server_scope_denied")
        return {
            "tenant_id": "tenant-1", "principal_id": "owner-1",
            "business_domain_uid": str(uuid.UUID(int=1)), "environment": "test",
            "provider": "controlled-provider", "model": "controlled-model-v1",
            "prompt_version": "prompt-v1", "generation": "generation-1",
            "interface_type": "mcp", "tool_name": "knowledge.search", "action": "read",
            "risk_level": "low",
        }

    def authorize_claim(self, request):
        self.calls.append(request)
        return {"decision": "authorized", "invocation_id": "invocation-1"}


def test_server_context_is_the_only_scope_source_and_payload_scope_is_rejected():
    repository = _Repository()
    service = ServerGovernedInvocationService(repository)
    context = RuntimeServerContext(agent_uid="agent-1", actor_uid="owner-1", roles=frozenset({"operator"}))
    result = service.authorize(context, {
        "grant_uid": "grant-1", "idempotency_key": "request-1", "input_text": "查询定义",
        "evidence_refs": [], "estimated_tokens": 1, "estimated_cost_micros": "1.00",
        "requested_time_ms": 1,
    })
    assert result["decision"] == "authorized"
    assert repository.calls[0]["tenant_id"] == "tenant-1"
    for field in ("tenant_id", "principal_id", "business_domain_uid", "environment", "provider", "model", "generation"):
        with pytest.raises(RuntimeServerError, match="payload_scope_forbidden"):
            service.authorize(context, {
                "grant_uid": "grant-1", "idempotency_key": f"request-{field}", "input_text": "查询定义",
                "evidence_refs": [], "estimated_tokens": 1, "estimated_cost_micros": "1.00",
                "requested_time_ms": 1, field: "forged",
            })


def test_server_path_rejects_bool_zero_negative_nan_and_ambiguous_budget_units():
    service = ServerGovernedInvocationService(_Repository())
    context = RuntimeServerContext(agent_uid="agent-1", actor_uid="owner-1", roles=frozenset({"operator"}))
    base = {"grant_uid": "grant-1", "idempotency_key": "request-1", "input_text": "查询定义", "evidence_refs": [], "estimated_tokens": 1, "estimated_cost_micros": "1.00", "requested_time_ms": 1}
    for field, value in (("estimated_tokens", True), ("estimated_tokens", 0), ("estimated_tokens", -1), ("estimated_cost_micros", "NaN"), ("estimated_cost_micros", "Infinity"), ("requested_time_ms", 0)):
        with pytest.raises(RuntimeServerError, match="budget_value_invalid"):
            service.authorize(context, {**base, field: value})


def test_server_path_requires_approval_for_high_risk_and_forwards_only_the_bound_task_uid():
    class _HighRiskRepository(_Repository):
        def grant_context(self, **kwargs):
            result = super().grant_context(**kwargs)
            result["action"] = "execute"
            result["risk_level"] = "high"
            return result

    service = ServerGovernedInvocationService(_HighRiskRepository())
    context = RuntimeServerContext(agent_uid="agent-1", actor_uid="owner-1", roles=frozenset({"operator"}))
    payload = {
        "grant_uid": "grant-1", "idempotency_key": "request-1", "input_text": "查询定义",
        "evidence_refs": [], "estimated_tokens": 1, "estimated_cost_micros": "1",
        "requested_time_ms": 1,
    }
    with pytest.raises(RuntimeServerError, match="human_approval_required"):
        service.authorize(context, {
            **payload,
        })
    result = service.authorize(context, {
        **payload, "approval_task_uid": "00000000-0000-0000-0000-000000000001",
    })
    assert result["decision"] == "authorized"
    assert service.repository.calls[-1]["approval_task_uid"].endswith("0001")