awoooi/apps/api/tests/test_gemini_cost_guard.py

from __future__ import annotations

import json
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from src.services import ai_control as ai_control_module
from src.services.ai_providers.gemini import GeminiProvider
from src.services.ai_rate_limiter import (
    _FINALIZE_GENERATION_LUA,
    _RESERVE_GENERATION_LUA,
    _WRITE_BLOCKED_GENERATION_RECEIPT_LUA,
    AIGenerationReservation,
    AIRateLimiter,
    get_gemini_pricing_policy,
)


@pytest.fixture(autouse=True)
def _gemini_runtime_control_enabled(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("ENABLE_GEMINI", "true")
    monkeypatch.setattr(
        ai_control_module,
        "is_provider_disabled",
        AsyncMock(return_value=False),
    )


def _reservation(*, allowed: bool = True, reason: str = "allowed") -> AIGenerationReservation:
    pricing = get_gemini_pricing_policy("gemini-2.5-flash-lite")
    assert pricing is not None
    return AIGenerationReservation(
        provider="gemini",
        receipt_id="receipt-guard-1",
        allowed=allowed,
        reason=reason,
        estimated_prompt_tokens=4,
        estimated_completion_tokens=2048,
        estimated_cost_usd=pricing.cost_usd(4, 2048),
        date="2026-07-14",
        model=pricing.model,
        pricing_source=pricing.source,
        pricing_version=pricing.version,
        pricing_checked_at=pricing.checked_at,
        input_usd_per_million=pricing.input_usd_per_million,
        output_usd_per_million=pricing.output_usd_per_million,
    )


class _ReserveRedis:
    def __init__(self, result: list[object] | None = None, error: Exception | None = None) -> None:
        self.result = result or [1, b"allowed"]
        self.error = error
        self.eval_calls: list[tuple[object, ...]] = []
        self.hsets: list[tuple[str, dict[str, object]]] = []
        self.expirations: list[tuple[str, int]] = []
        self.set_calls: list[tuple[str, object, bool | None, int | None]] = []

    async def eval(self, *args: object) -> list[object]:
        self.eval_calls.append(args)
        if self.error:
            raise self.error
        if args[0] == _WRITE_BLOCKED_GENERATION_RECEIPT_LUA:
            receipt_key = str(args[2])
            receipt_id = str(args[4])
            mapping = json.loads(str(args[5]))
            self.hsets.append((receipt_key, mapping))
            if int(args[6]) == 1:
                self.set_calls.append((str(args[3]), receipt_id, True, None))
            return [1, b"written", receipt_id]
        return self.result

    async def set(
        self,
        key: str,
        value: object,
        *,
        nx: bool | None = None,
        ex: int | None = None,
    ) -> bool:
        self.set_calls.append((key, value, nx, ex))
        return True

    def pipeline(self, transaction: bool = True) -> _ReceiptPipeline:
        return _ReceiptPipeline(self, transaction=transaction)


class _ReceiptPipeline:
    def __init__(self, redis: _ReserveRedis, *, transaction: bool) -> None:
        self.redis = redis
        self.transaction = transaction

    def set(
        self,
        key: str,
        value: object,
        *,
        nx: bool | None = None,
        ex: int | None = None,
    ) -> _ReceiptPipeline:
        self.redis.set_calls.append((key, value, nx, ex))
        return self

    def hset(
        self,
        key: str,
        *,
        mapping: dict[str, object],
    ) -> _ReceiptPipeline:
        self.redis.hsets.append((key, mapping))
        return self

    def expire(self, key: str, ttl: int) -> _ReceiptPipeline:
        self.redis.expirations.append((key, ttl))
        return self

    async def execute(self) -> list[int]:
        return [1, 1]


class _UsageRedis:
    async def get(self, _key: str) -> None:
        return None


@pytest.mark.asyncio
async def test_gemini_reservation_is_one_atomic_eval_with_all_hard_limits(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    redis = _ReserveRedis()
    limiter = AIRateLimiter()
    limiter._redis = redis
    monkeypatch.setattr(limiter, "_get_today", lambda: "2026-07-14")
    monkeypatch.setattr(limiter, "_seconds_until_tomorrow", lambda: 3600)

    reservation = await limiter.reserve_generation(
        "gemini",
        "alert diagnosis",
        context={
            "trace_id": "trace-1",
            "run_id": "run-1",
            "work_item_id": "work-1",
            "secret_value": "must-not-be-stored",
        },
    )

    assert reservation.allowed is True
    assert len(redis.eval_calls) == 1
    call = redis.eval_calls[0]
    assert call[0] == _RESERVE_GENERATION_LUA
    assert call[1] == 11
    serialized = " ".join(str(value) for value in call)
    assert "must-not-be-stored" not in serialized
    assert "daily_request_limit" in _RESERVE_GENERATION_LUA
    assert "daily_token_limit" in _RESERVE_GENERATION_LUA
    assert "daily_cost_limit" in _RESERVE_GENERATION_LUA
    assert "total_cost_limit" in _RESERVE_GENERATION_LUA
    assert "reserved_cost_micro_usd" in _RESERVE_GENERATION_LUA
    assert "actual_cost_micro_usd" in _RESERVE_GENERATION_LUA
    assert "pricing_source" in _RESERVE_GENERATION_LUA
    assert "pricing_version" in _RESERVE_GENERATION_LUA
    assert "pricing_checked_at" in _RESERVE_GENERATION_LUA
    assert "'EXPIRE', KEYS[8]" not in _RESERVE_GENERATION_LUA


def test_paid_reservation_identity_and_receipts_have_no_ttl_reopen_path() -> None:
    assert "redis.call('SET', KEYS[10], ARGV[10], 'NX')" in (
        _RESERVE_GENERATION_LUA
    )
    assert "redis.call('EXPIRE', KEYS[9]" not in _RESERVE_GENERATION_LUA
    assert "redis.call('EXPIRE', KEYS[10]" not in _RESERVE_GENERATION_LUA
    assert "redis.call('EXPIRE', KEYS[7]" not in _FINALIZE_GENERATION_LUA
    assert "EXPIRE" not in _WRITE_BLOCKED_GENERATION_RECEIPT_LUA


@pytest.mark.asyncio
async def test_gemini_unknown_model_writes_durable_blocked_no_write_receipt() -> None:
    redis = _ReserveRedis()
    limiter = AIRateLimiter()
    limiter._redis = redis

    reservation = await limiter.reserve_generation(
        "gemini",
        "alert diagnosis",
        model="gemini-unknown",
        context={
            "trace_id": "trace-1",
            "run_id": "run-1",
            "work_item_id": "work-1",
            "traceparent": "must-not-be-accepted",
            "secret": "must-not-be-stored",
        },
    )

    assert reservation.allowed is False
    assert reservation.reason == "pricing_policy_missing"
    assert reservation.model == "gemini-unknown"
    assert len(redis.eval_calls) == 1
    assert redis.eval_calls[0][0] == _WRITE_BLOCKED_GENERATION_RECEIPT_LUA
    assert len(redis.hsets) == 1
    receipt_key, receipt = redis.hsets[0]
    assert receipt_key.endswith(reservation.receipt_id)
    assert receipt["status"] == "blocked"
    assert receipt["terminal_state"] == "blocked_no_write"
    assert receipt["trace_id"] == "trace-1"
    assert receipt["run_id"] == "run-1"
    assert receipt["work_item_id"] == "work-1"
    assert receipt["estimated_total_tokens"] == "0"
    assert receipt["reserved_total_tokens"] == "0"
    assert receipt["reserved_cost_micro_usd"] == "0"
    serialized = repr(receipt)
    assert "traceparent" not in serialized
    assert "must-not-be-accepted" not in serialized
    assert "must-not-be-stored" not in serialized


@pytest.mark.asyncio
async def test_gemini_missing_correlation_identity_blocks_before_reservation() -> None:
    redis = _ReserveRedis()
    limiter = AIRateLimiter()
    limiter._redis = redis

    reservation = await limiter.reserve_generation(
        "gemini",
        "alert diagnosis",
        context={"trace_id": "trace-1", "run_id": "run-1"},
    )

    assert reservation.allowed is False
    assert reservation.reason == "missing_correlation_identity"
    assert len(redis.eval_calls) == 1
    assert redis.eval_calls[0][0] == _WRITE_BLOCKED_GENERATION_RECEIPT_LUA
    assert redis.hsets[0][1]["reason"] == "missing_correlation_identity"


def test_gemini_pricing_policy_is_exact_model_standard_text_contract() -> None:
    pricing = get_gemini_pricing_policy("gemini-2.5-flash-lite")

    assert pricing is not None
    assert pricing.input_usd_per_million == 0.10
    assert pricing.output_usd_per_million == 0.40
    assert pricing.checked_at == "2026-07-14"
    assert "ai.google.dev/gemini-api/docs/pricing" in pricing.source
    assert get_gemini_pricing_policy("gemini-2.0-flash") is None


def test_receipt_correlation_accepts_only_three_bounded_canonical_ids() -> None:
    limiter = AIRateLimiter()

    assert limiter._safe_receipt_context(
        {
            "trace_id": "trace:valid/1",
            "run_id": "run with spaces",
            "work_item_id": "work-1",
            "traceparent": "must-not-be-used",
            "incident_id": "must-not-be-used",
        }
    ) == ("trace:valid/1", "", "work-1")
    assert limiter._safe_receipt_context(
        {
            "trace_id": "x" * 161,
            "run_id": "run-1",
            "work_item_id": "work\nitem",
        }
    ) == ("", "run-1", "")


@pytest.mark.asyncio
async def test_gemini_usage_readback_publishes_limits_and_pricing_policy() -> None:
    limiter = AIRateLimiter()
    limiter._redis = _UsageRedis()

    stats = await limiter.get_usage_stats("gemini")

    assert stats["rpm"]["limit"] == 10
    assert stats["daily_requests"]["limit"] == 500
    assert stats["daily_tokens"]["limit"] == 100_000
    assert stats["total_cost_usd"]["limit"] == 5.0
    assert stats["total_cost_usd"]["alert_threshold"] == 4.0
    assert stats["pricing_policy"] == {
        "model": "gemini-2.5-flash-lite",
        "supported": True,
        "source": (
            "https://ai.google.dev/gemini-api/docs/pricing"
            "#gemini-2.5-flash-lite"
        ),
        "version": "google-ai-standard-text-2026-07-14",
        "checked_at": "2026-07-14",
        "input_usd_per_million": 0.1,
        "output_usd_per_million": 0.4,
    }


@pytest.mark.asyncio
async def test_gemini_reservation_fails_closed_when_redis_is_unavailable(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    limiter = AIRateLimiter()
    limiter._redis = _ReserveRedis(error=RuntimeError("offline"))
    monkeypatch.setattr(limiter, "_get_today", lambda: "2026-07-14")

    reservation = await limiter.reserve_generation(
        "gemini",
        "alert diagnosis",
        context={
            "trace_id": "trace-redis-down",
            "run_id": "run-redis-down",
            "work_item_id": "work-redis-down",
        },
    )

    assert reservation.allowed is False
    assert reservation.reason == "cost_guard_unavailable"


@pytest.mark.asyncio
async def test_gemini_finalize_uses_atomic_receipt_transition(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    redis = _ReserveRedis(result=[1, b"succeeded", 5, 1])

    async def _get(_key: str) -> bytes:
        return b"0.1"

    redis.get = _get  # type: ignore[attr-defined]
    limiter = AIRateLimiter()
    limiter._redis = redis
    monkeypatch.setattr(limiter, "_seconds_until_tomorrow", lambda: 3600)

    finalized = await limiter.finalize_generation(
        _reservation(),
        success=True,
        prompt_tokens=3,
        completion_tokens=2,
    )

    assert finalized is True
    assert redis.eval_calls[0][0] == _FINALIZE_GENERATION_LUA
    assert redis.eval_calls[0][1] == 9


@pytest.mark.asyncio
async def test_cost_warning_includes_pending_reservation_exposure() -> None:
    limiter = AIRateLimiter()
    limiter.get_usage_stats = AsyncMock(  # type: ignore[method-assign]
        return_value={
            "total_cost_usd": {
                "current": 3.9,
                "reserved": 0.2,
                "alert_threshold": 4.0,
            }
        }
    )
    limiter._send_cost_warning = AsyncMock()  # type: ignore[method-assign]

    await limiter._warn_for_cost_exposure("gemini")

    limiter._send_cost_warning.assert_awaited_once_with("gemini", 4.1, 4.0)


@pytest.mark.asyncio
async def test_gemini_provider_does_not_call_api_when_cost_guard_blocks() -> None:
    client = MagicMock()
    client.is_closed = False
    client.post = AsyncMock()
    limiter = MagicMock()
    limiter.reserve_generation = AsyncMock(
        return_value=_reservation(allowed=False, reason="daily_cost_limit")
    )

    provider = GeminiProvider()
    provider._http_client = client
    with patch(
        "src.services.ai_providers.gemini.settings",
        SimpleNamespace(
            GEMINI_API_KEY="configured",
            GEMINI_MODEL="gemini-2.5-flash-lite",
            GEMINI_MAX_OUTPUT_TOKENS=2048,
        ),
    ), patch(
        "src.services.ai_rate_limiter.get_ai_rate_limiter",
        return_value=limiter,
    ):
        result = await provider.analyze(
            "alert diagnosis",
            context={
                "trace_id": "trace-cost-block",
                "run_id": "run-cost-block",
                "work_item_id": "work-cost-block",
            },
        )

    assert result.success is False
    assert result.error == "cost_guard_blocked:daily_cost_limit"
    client.post.assert_not_awaited()


@pytest.mark.asyncio
async def test_gemini_provider_does_not_reserve_or_call_when_runtime_disabled(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    client = MagicMock()
    client.is_closed = False
    client.post = AsyncMock()
    limiter = MagicMock()
    limiter.reserve_generation = AsyncMock()
    monkeypatch.setattr(
        ai_control_module,
        "is_provider_disabled",
        AsyncMock(return_value=True),
    )
    provider = GeminiProvider()
    provider._http_client = client

    with patch(
        "src.services.ai_providers.gemini.settings",
        SimpleNamespace(GEMINI_API_KEY="configured-test-value"),
    ), patch(
        "src.services.ai_rate_limiter.get_ai_rate_limiter",
        return_value=limiter,
    ):
        result = await provider.analyze(
            "alert diagnosis",
            context={
                "trace_id": "trace-disabled",
                "run_id": "run-disabled",
                "work_item_id": "work-disabled",
            },
        )

    assert result.success is False
    assert result.error == "gemini_disabled_by_runtime_control"
    limiter.reserve_generation.assert_not_awaited()
    client.post.assert_not_awaited()


@pytest.mark.asyncio
async def test_gemini_provider_failure_finalizes_estimate_charged_receipt() -> None:
    client = MagicMock()
    client.is_closed = False
    client.post = AsyncMock(side_effect=RuntimeError("provider unavailable"))
    registry = MagicMock()
    registry.get_provider_options.return_value = {"temperature": 0.1}
    limiter = MagicMock()
    limiter.reserve_generation = AsyncMock(return_value=_reservation())
    limiter.finalize_generation = AsyncMock(return_value=True)

    provider = GeminiProvider()
    provider._http_client = client
    with patch(
        "src.services.ai_providers.gemini.settings",
        SimpleNamespace(
            GEMINI_API_KEY="configured",
            GEMINI_MODEL="gemini-2.5-flash-lite",
            GEMINI_MAX_OUTPUT_TOKENS=2048,
        ),
    ), patch(
        "src.services.ai_providers.gemini.get_model_registry",
        return_value=registry,
    ), patch(
        "src.services.ai_rate_limiter.get_ai_rate_limiter",
        return_value=limiter,
    ):
        result = await provider.analyze(
            "alert diagnosis",
            context={
                "trace_id": "trace-failure",
                "run_id": "run-failure",
                "work_item_id": "work-failure",
            },
        )

    assert result.success is False
    limiter.finalize_generation.assert_awaited_once_with(
        _reservation(),
        success=False,
        error_code="RuntimeError",
    )


@pytest.mark.asyncio
async def test_gemini_success_without_usage_metadata_is_not_zero_false_green() -> None:
    response = MagicMock()
    response.raise_for_status.return_value = None
    response.json.return_value = {
        "candidates": [{"content": {"parts": [{"text": '{"ok": true}'}]}}],
    }
    client = MagicMock()
    client.is_closed = False
    client.post = AsyncMock(return_value=response)
    registry = MagicMock()
    registry.get_provider_options.return_value = {"temperature": 0.1}
    limiter = MagicMock()
    limiter.reserve_generation = AsyncMock(return_value=_reservation())
    limiter.finalize_generation = AsyncMock(return_value=True)

    provider = GeminiProvider()
    provider._http_client = client
    with patch(
        "src.services.ai_providers.gemini.settings",
        SimpleNamespace(
            GEMINI_API_KEY="configured",
            GEMINI_MODEL="gemini-2.5-flash-lite",
            GEMINI_MAX_OUTPUT_TOKENS=2048,
        ),
    ), patch(
        "src.services.ai_providers.gemini.get_model_registry",
        return_value=registry,
    ), patch(
        "src.services.ai_rate_limiter.get_ai_rate_limiter",
        return_value=limiter,
    ):
        result = await provider.analyze(
            "alert diagnosis",
            context={
                "trace_id": "trace-no-usage",
                "run_id": "run-no-usage",
                "work_item_id": "work-no-usage",
            },
        )

    assert result.success is True
    assert result.tokens > 0
    assert result.cost_usd > 0
    limiter.finalize_generation.assert_awaited_once_with(
        _reservation(),
        success=True,
        prompt_tokens=0,
        completion_tokens=0,
    )