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,
    PAID_RUN_BUDGET_KEY,
    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 _ReceiptReadRedis:
    async def hgetall(self, _key: str) -> dict[bytes, bytes]:
        return {
            b"provider": b"gemini",
            b"receipt_id": b"receipt-readback-1",
            b"status": b"succeeded",
            b"terminal_state": b"finalized_accounted",
            b"success": b"1",
            b"trace_id": b"trace-readback",
            b"run_id": b"run-readback",
            b"work_item_id": b"AIA-SRE-013",
            b"model": b"gemini-2.5-flash-lite",
            b"pricing_version": b"pricing-v1",
            b"usage_source": b"provider_metadata",
            b"accounted_tokens": b"120",
            b"accounted_cost_micro_usd": b"42",
            b"run_budget_required": b"1",
            b"run_budget_id": b"paid-run-readback",
            b"run_budget_max_micro_usd": b"250000",
            b"run_budget_reserved_cost_micro_usd": b"0",
            b"reserved_prompt_tokens": b"0",
            b"reserved_completion_tokens": b"0",
            b"reserved_total_tokens": b"0",
            b"reserved_cost_micro_usd": b"0",
        }


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

    receipt = await limiter.get_generation_receipt_readback(
        "gemini",
        "receipt-readback-1",
        trace_id="trace-readback",
        run_id="run-readback",
        work_item_id="AIA-SRE-013",
    )

    assert receipt["found"] is True
    assert receipt["correlation_verified"] is True
    assert receipt["finalized_accounted"] is True
    assert receipt["reservation_released"] is True
    assert receipt["accounted_tokens"] == 120
    assert receipt["accounted_cost_usd"] == 0.000042
    assert receipt["accounted_cost_micro_usd"] == 42
    assert receipt["run_budget_required"] is True
    assert receipt["raw_prompt_persisted"] is False
    assert "trace_id" not in receipt


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


class _RunBudgetReadRedis:
    async def hgetall(self, _key: str) -> dict[bytes, bytes]:
        return {
            b"schema_version": b"ai_paid_run_budget_receipt_v1",
            b"run_budget_id": b"paid-run-6b098d09eeef680a5b119ca0b4d31713",
            b"run_id": b"run-readback",
            b"work_item_id": b"AIA-SRE-013",
            b"status": b"finalized",
            b"max_cost_micro_usd": b"250000",
            b"reserved_cost_micro_usd": b"0",
            b"accounted_cost_micro_usd": b"2000",
            b"reservation_count": b"2",
            b"finalized_count": b"2",
        }


@pytest.mark.asyncio
async def test_paid_run_budget_readback_is_allowlisted_and_exactly_correlated() -> None:
    limiter = AIRateLimiter()
    run_id = "run-readback"
    work_item_id = "AIA-SRE-013"
    budget_hash = limiter._identity_hash(run_id, work_item_id)
    budget_id = f"paid-run-{budget_hash[:32]}"
    redis = _RunBudgetReadRedis()

    async def _hgetall(_key: str) -> dict[bytes, bytes]:
        payload = await _RunBudgetReadRedis().hgetall(_key)
        payload[b"run_budget_id"] = budget_id.encode()
        return payload

    redis.hgetall = _hgetall  # type: ignore[method-assign]
    limiter._redis = redis

    receipt = await limiter.get_paid_run_budget_readback(
        run_id=run_id,
        work_item_id=work_item_id,
        expected_max_cost_micro_usd=250_000,
    )

    assert receipt == {
        "schema_version": "ai_paid_run_budget_receipt_v1",
        "run_budget_id": budget_id,
        "found": True,
        "status": "finalized",
        "max_cost_micro_usd": 250_000,
        "reserved_cost_micro_usd": 0,
        "accounted_cost_micro_usd": 2000,
        "accounted_cost_usd": 0.002,
        "reservation_count": 2,
        "finalized_count": 2,
        "correlation_verified": True,
        "cap_verified": True,
        "terminal_verified": True,
        "raw_prompt_persisted": False,
        "raw_response_persisted": False,
        "secret_value_present": False,
    }


@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] == 12
    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 "run_cost_limit" in _RESERVE_GENERATION_LUA
    assert "'EXPIRE', KEYS[8]" not in _RESERVE_GENERATION_LUA


@pytest.mark.asyncio
async def test_paid_canary_cap_is_part_of_the_same_atomic_reservation_eval(
    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",
        "sanitized canary",
        context={
            "trace_id": "trace-paid-canary",
            "run_id": "run-paid-canary",
            "work_item_id": "AIA-SRE-013",
            "synthetic_validation": True,
            "paid_canary_max_cost_usd": 0.25,
        },
    )

    assert reservation.allowed is True
    assert reservation.run_budget_id.startswith("paid-run-")
    assert reservation.run_budget_max_micro_usd == 250_000
    assert len(redis.eval_calls) == 1
    call = redis.eval_calls[0]
    assert call[0] == _RESERVE_GENERATION_LUA
    assert call[1] == 12
    assert call[13] == PAID_RUN_BUDGET_KEY.format(budget_id=reservation.run_budget_id)
    assert call[-3:] == (1, reservation.run_budget_id, 250_000)


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("cap_marker", "expected_reason"),
    [
        (None, "run_cost_cap_missing"),
        (True, "run_cost_cap_invalid"),
        (0, "run_cost_cap_invalid"),
        (-1, "run_cost_cap_invalid"),
        (0.0000001, "run_cost_cap_invalid"),
        (0.250001, "run_cost_cap_invalid"),
        ("NaN", "run_cost_cap_invalid"),
        ("Infinity", "run_cost_cap_invalid"),
        ("not-a-number", "run_cost_cap_invalid"),
    ],
)
async def test_paid_canary_missing_or_invalid_cap_fails_closed_before_reservation(
    cap_marker: object,
    expected_reason: str,
) -> None:
    redis = _ReserveRedis()
    limiter = AIRateLimiter()
    limiter._redis = redis
    context = {
        "trace_id": "trace-invalid-cap",
        "run_id": f"run-invalid-cap-{expected_reason}",
        "work_item_id": "AIA-SRE-013",
        "synthetic_validation": True,
    }
    if cap_marker is not None:
        context["paid_canary_max_cost_usd"] = cap_marker

    reservation = await limiter.reserve_generation(
        "gemini",
        "sanitized canary",
        context=context,
    )

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


@pytest.mark.asyncio
async def test_unrelated_synthetic_call_does_not_inherit_aia_sre_013_run_cap(
    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",
        "unrelated synthetic fixture",
        context={
            "trace_id": "trace-unrelated-synthetic",
            "run_id": "run-unrelated-synthetic",
            "work_item_id": "OTHER-SYNTHETIC-WORK",
            "synthetic_validation": True,
        },
    )

    assert reservation.allowed is True
    assert reservation.run_budget_id == ""
    assert reservation.run_budget_max_micro_usd == 0


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

    reservation = await limiter.reserve_generation(
        "gemini",
        "canary missing exact marker",
        context={
            "trace_id": "trace-canary-marker-missing",
            "run_id": "run-canary-marker-missing",
            "work_item_id": "AIA-SRE-013",
            "paid_canary_max_cost_usd": 0.25,
        },
    )

    assert reservation.allowed is False
    assert reservation.reason == "run_cost_cap_context_invalid"
    assert redis.eval_calls[0][0] == _WRITE_BLOCKED_GENERATION_RECEIPT_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[12]" not in _RESERVE_GENERATION_LUA
    assert "redis.call('EXPIRE', KEYS[7]" not in _FINALIZE_GENERATION_LUA
    assert "redis.call('EXPIRE', KEYS[10]" 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] == 10


@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,
            GEMINI_EXECUTION_MODE="production",
            GEMINI_CANARY_PERCENT=0,
        ),
    ), 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",
                "data_classification": "sanitized",
            },
        )

    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()
    disabled_readback = AsyncMock(return_value=True)
    monkeypatch.setattr(
        ai_control_module,
        "is_provider_disabled",
        disabled_readback,
    )
    provider = GeminiProvider()
    provider._http_client = client

    with patch(
        "src.services.ai_providers.gemini.settings",
        SimpleNamespace(
            GEMINI_API_KEY="configured-test-value",
            GEMINI_EXECUTION_MODE="production",
            GEMINI_CANARY_PERCENT=0,
        ),
    ), 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",
                "data_classification": "sanitized",
            },
        )

    assert result.success is False
    assert result.error == "gemini_disabled_by_runtime_control"
    disabled_readback.assert_awaited_once_with(
        "gemini",
        run_id="run-disabled",
    )
    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,
            GEMINI_EXECUTION_MODE="production",
            GEMINI_CANARY_PERCENT=0,
        ),
    ), 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",
                "data_classification": "sanitized",
            },
        )

    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,
            GEMINI_EXECUTION_MODE="production",
            GEMINI_CANARY_PERCENT=0,
        ),
    ), 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",
                "data_classification": "sanitized",
            },
        )

    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,
    )