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