diff --git a/smolrouter/database.py b/smolrouter/database.py index 75da32a..b2dd586 100644 --- a/smolrouter/database.py +++ b/smolrouter/database.py @@ -1081,7 +1081,14 @@ def select(): return [] @staticmethod - async def mark_invalid_by_hash(api_key_hash: str, provider_name: str) -> int: + async def mark_invalid_by_hash( + api_key_hash: str, + provider_name: str, + *, + reason: str = "operator", + status_code: int | None = None, + request_id: str | None = None, + ) -> int: """Mark all quota entries for a given key hash as invalid. Args: @@ -1091,7 +1098,13 @@ async def mark_invalid_by_hash(api_key_hash: str, provider_name: str) -> int: Returns: Number of quota entries marked as invalid """ - return await RedisApiKeyQuota.mark_invalid(api_key_hash, provider_name) + return await RedisApiKeyQuota.mark_invalid( + api_key_hash, + provider_name, + reason=reason, + status_code=status_code, + request_id=request_id, + ) @staticmethod async def mark_quota_exhausted(api_key: str, provider_id: str, model_name: str, error: str = None) -> None: diff --git a/smolrouter/google_genai_provider.py b/smolrouter/google_genai_provider.py index 3e8dc68..eb12fed 100644 --- a/smolrouter/google_genai_provider.py +++ b/smolrouter/google_genai_provider.py @@ -16,7 +16,7 @@ from types import SimpleNamespace from datetime import datetime, timedelta, timezone from typing import List, Dict, Any, Optional, Tuple, Callable -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from urllib.parse import urlparse from zoneinfo import ZoneInfo import httpx @@ -84,29 +84,15 @@ } GOOGLE_IMAGE_RESPONSE_FORMATS = {"url", "b64_json"} OPENAI_IMAGE_REQUEST_FIELDS = {"model", "prompt", "n", "size", "response_format", "extra_body", "user"} -GOOGLE_NON_INVALID_403_INDICATORS = ( - "paid plan", - "paid plans", - "upgrade your account", - "billing", - "billed user", - "not supported for predict", - "forbidden for project", - "model is not found", +GOOGLE_ERROR_DETAIL_TYPE_FIELD = "@type" +# Permanent key invalidation has a fleet-wide blast radius. Generic status codes +# and Google request-validation failures are deliberately not evidence of a bad key. +GOOGLE_INVALID_KEY_REASON_PATTERNS = ( + ("api_key_not_valid", re.compile(r"(?:^|[.:]\s*)api key (?:is )?not valid(?:[.!:].*)?$", re.IGNORECASE)), + ("api_key_invalid", re.compile(r"(?:^|[.:]\s*)invalid api key(?:[.!:].*)?$", re.IGNORECASE)), + ("api_key_expired", re.compile(r"(?:^|[.:]\s*)api key expired(?:[.!:].*)?$", re.IGNORECASE)), + ("api_key_revoked", re.compile(r"(?:^|[.:]\s*)api key revoked(?:[.!:].*)?$", re.IGNORECASE)), ) -GOOGLE_INVALID_KEY_INDICATORS = ( - "permission denied", - "permission_denied", - "permissiondenied", - "api key not valid", - "invalid api key", - "api_key_invalid", - "authentication failed", - "unauthorized", - "credentials are missing or invalid", - "api key expired", -) -GOOGLE_FALLBACK_INVALID_KEY_INDICATORS = GOOGLE_INVALID_KEY_INDICATORS + ("invalid_argument",) AUDIO_WAV_MIME_TYPE = "audio/wav" AUDIO_PCM_MIME_TYPE = "audio/pcm" AUDIO_MPEG_MIME_TYPE = "audio/mpeg" @@ -226,6 +212,8 @@ def __init__( api_key_index: Optional[int] = None, api_key_total: Optional[int] = None, proxy_used: Optional[str] = None, + status_code: Optional[int] = None, + error_type: str = "api_error", ): super().__init__(message) self.provider_id = provider_id @@ -234,9 +222,35 @@ def __init__( self.api_key_index = api_key_index self.api_key_total = api_key_total self.proxy_used = proxy_used + self.status_code = status_code + self.error_type = error_type self.retry_after_seconds: Optional[float] = None +@dataclass(frozen=True) +class GoogleErrorEvidence: + """Normalized upstream evidence used for routing state transitions. + + Only the provider's top-level message is eligible for the narrow text fallback + used to invalidate a key. Nested details are kept for quota parsing only. + """ + + status_code: Optional[int] + status: str = "" + message: str = "" + reason: str = "" + quota_id: str = "" + retry_delay_seconds: Optional[float] = None + + +class GoogleGenAIUpstreamError(RuntimeError): + """HTTP error whose parsed Google response is safe to classify internally.""" + + def __init__(self, evidence: GoogleErrorEvidence): + super().__init__(f"Google upstream HTTP {evidence.status_code or 'error'}") + self.evidence = evidence + + @dataclass class GoogleGenAIConfig(ProviderConfig): """Extended configuration for Google GenAI provider""" @@ -978,6 +992,7 @@ async def _update_api_key_stats( tokens: int = 0, error: Optional[str] = None, status_code: Optional[int] = None, + error_evidence: Optional[GoogleErrorEvidence] = None, ): """Update statistics for an API key + model combination after a request @@ -999,19 +1014,18 @@ async def _update_api_key_stats( await self._record_api_key_success(RedisApiKeyQuota, quota, api_key, model_name, tokens) else: error_message = error or "" - # Check error type and handle appropriately - # IMPORTANT: Check both error message AND status code (403 errors may have empty error string) - if self._is_invalid_key_error(error_message, status_code): - await self._record_invalid_api_key(ApiKeyQuota, api_key, error_message, status_code) + invalid_reason = self._invalid_key_reason(error_evidence) + if invalid_reason is not None: + await self._record_invalid_api_key(ApiKeyQuota, api_key, status_code, invalid_reason) # Check if this is a quota exhaustion error - elif self._is_quota_exhausted_error(error_message): - if self._is_per_day_quota_error(error_message): + elif self._is_quota_exhausted_error(error_message, error_evidence): + if self._is_per_day_quota_error(error_message, error_evidence): # Only explicit per-day (RPD) signals should hard-bench a key. await self._record_quota_exhaustion(ApiKeyQuota, quota, api_key, model_name, error_message) else: # Explicit RPM and ambiguous quota 429s are treated as transient cooldowns. - cooldown_seconds = self._quota_cooldown_seconds(error_message) + cooldown_seconds = self._quota_cooldown_seconds(error_message, error_evidence) await self._record_rate_limit_cooldown( ApiKeyQuota, quota, api_key, model_name, error_message, cooldown_seconds ) @@ -1053,20 +1067,21 @@ async def _record_api_key_success( quota.mark_request_success(tokens=tokens) async def _record_invalid_api_key( - self, quota_backend: Any, api_key: str, error_message: str, status_code: Optional[int] + self, quota_backend: Any, api_key: str, status_code: Optional[int], invalid_reason: str ) -> None: key_hash = quota_backend.hash_api_key(api_key) try: - await quota_backend.mark_invalid_by_hash(key_hash, self.config.name) - logger.error(f"🚫 API key {redact_secret(api_key)} MARKED INVALID (status={status_code}, error={error_message})") - - try: - self.config.api_keys.remove(api_key) - logger.warning( - f"🗑️ Removed API key {redact_secret(api_key)} from selection pool ({len(self.config.api_keys)} keys remaining)" - ) - except ValueError: - pass + await quota_backend.mark_invalid_by_hash( + key_hash, self.config.name, reason=invalid_reason, status_code=status_code + ) + # Redis is the shared selector authority. Leaving the configured list + # intact makes an operator recovery effective without a process restart. + logger.error( + "API key %s marked invalid (reason=%s, status=%s)", + redact_secret(api_key), + invalid_reason, + status_code, + ) except Exception: logger.exception("❌ Failed to mark API key as invalid") @@ -1179,35 +1194,33 @@ def _log_quota_status(self, api_key: str, model_name: str, quota: QuotaRecord) - model_limit, ) - def _is_invalid_key_error( - self, error_msg: Optional[str] = None, status_code: Optional[int] = None - ) -> bool: - """Check if error indicates invalid/expired API key - - Args: - error_msg: Error message text (may be None or empty) - status_code: HTTP status code + def _invalid_key_reason(self, evidence: Optional[GoogleErrorEvidence]) -> Optional[str]: + """Return explicit Google key-invalid evidence, never a status-only inference.""" + if evidence is None: + return None - Returns: - True if this error indicates an invalid/expired API key - """ - if status_code == 403: - if error_msg: - error_lower = error_msg.lower() - if any(indicator in error_lower for indicator in GOOGLE_NON_INVALID_403_INDICATORS): - return False - return any(indicator in error_lower for indicator in GOOGLE_INVALID_KEY_INDICATORS) - return True + if evidence.reason.upper() == "API_KEY_INVALID": + return "api_key_invalid" - # Check error message for known invalid key indicators - if error_msg: - error_lower = error_msg.lower() - return any(indicator in error_lower for indicator in GOOGLE_FALLBACK_INVALID_KEY_INDICATORS) + normalized_message = " ".join(evidence.message.strip().split()) + for reason, pattern in GOOGLE_INVALID_KEY_REASON_PATTERNS: + if pattern.search(normalized_message): + return reason + return None - return False + def _is_invalid_key_error( + self, error_msg: Optional[str] = None, status_code: Optional[int] = None + ) -> bool: + """Test a direct top-level Google message without status-only inference.""" + _ = status_code + return self._invalid_key_reason(GoogleErrorEvidence(status_code=None, message=error_msg or "")) is not None - def _is_quota_exhausted_error(self, error_msg: Optional[str] = None) -> bool: + def _is_quota_exhausted_error( + self, error_msg: Optional[str] = None, evidence: Optional[GoogleErrorEvidence] = None + ) -> bool: """Check if error indicates quota exhaustion""" + if evidence is not None and (evidence.status_code == 429 or evidence.status.upper() == "RESOURCE_EXHAUSTED"): + return True if not error_msg: return False @@ -1224,26 +1237,36 @@ def _is_quota_exhausted_error(self, error_msg: Optional[str] = None) -> bool: return any(indicator in error_lower for indicator in quota_indicators) - def _is_per_minute_quota_error(self, error_msg: Optional[str] = None) -> bool: + def _is_per_minute_quota_error( + self, error_msg: Optional[str] = None, evidence: Optional[GoogleErrorEvidence] = None + ) -> bool: """Check if a quota 429 is a per-minute (RPM) limit rather than per-day (RPD). Google returns the same quotaMetric (generate_content_free_tier_requests) for both RPM and RPD on the free tier; the only reliable discriminator is the quotaId, e.g. ``GenerateRequestsPerMinutePerProjectPerModel-FreeTier``. """ + if evidence is not None and "perminute" in evidence.quota_id.lower(): + return True if not error_msg: return False error_lower = error_msg.lower() return "perminute" in error_lower or "per minute" in error_lower - def _is_per_day_quota_error(self, error_msg: Optional[str] = None) -> bool: + def _is_per_day_quota_error( + self, error_msg: Optional[str] = None, evidence: Optional[GoogleErrorEvidence] = None + ) -> bool: """Check if a quota 429 is an explicit per-day (RPD) limit.""" + if evidence is not None and "perday" in evidence.quota_id.lower(): + return True if not error_msg: return False error_lower = error_msg.lower() return "perday" in error_lower or "per day" in error_lower or "requests per day" in error_lower - def _quota_cooldown_seconds(self, error_msg: Optional[str] = None) -> float: + def _quota_cooldown_seconds( + self, error_msg: Optional[str] = None, evidence: Optional[GoogleErrorEvidence] = None + ) -> float: """Return a transient cooldown (seconds) for a quota 429. - Explicit per-minute (RPM) → cooldown driven by Google's retryDelay (or default). @@ -1251,7 +1274,7 @@ def _quota_cooldown_seconds(self, error_msg: Optional[str] = None) -> float: enough evidence for all-day exhaustion. - Explicit per-day (RPD) is handled by the caller before reaching this helper. """ - retry_delay = self._extract_retry_delay(error_msg) + retry_delay = evidence.retry_delay_seconds if evidence is not None else self._extract_retry_delay(error_msg) base = retry_delay if retry_delay is not None else self.DEFAULT_RPM_COOLDOWN_SECONDS return base + self.RPM_COOLDOWN_BUFFER_SECONDS @@ -2478,23 +2501,148 @@ def _convert_imagen_response_to_openai_images(self, response: Dict[str, Any]) -> return openai_response @staticmethod - def _extract_image_error_message(response: httpx.Response) -> str: + def _parse_google_duration(value: Any) -> Optional[float]: + if isinstance(value, (int, float)): + return float(value) + if not isinstance(value, str): + return None + match = re.fullmatch(r"(\d+(?:\.\d+)?)s", value.strip()) + return float(match.group(1)) if match else None + + @staticmethod + def _google_error_payload(payload: Any) -> Dict[str, Any]: + if not isinstance(payload, dict): + return {} + error_payload = payload.get("error", payload) + return error_payload if isinstance(error_payload, dict) else {} + + @staticmethod + def _google_error_status_code(error_payload: Dict[str, Any], fallback_status: Optional[int]) -> Optional[int]: + status_code = error_payload.get("code", fallback_status) + return status_code if isinstance(status_code, int) else fallback_status + + @staticmethod + def _google_error_reason(detail: Dict[str, Any]) -> str: + reason = detail.get("reason") + if str(detail.get(GOOGLE_ERROR_DETAIL_TYPE_FIELD, "")).endswith("ErrorInfo") and isinstance(reason, str): + return reason + return "" + + @staticmethod + def _google_quota_id(detail: Dict[str, Any]) -> str: + if not str(detail.get(GOOGLE_ERROR_DETAIL_TYPE_FIELD, "")).endswith("QuotaFailure"): + return "" + violations = detail.get("violations") + if not isinstance(violations, list): + return "" + for violation in violations: + quota_id = violation.get("quotaId") if isinstance(violation, dict) else None + if isinstance(quota_id, str): + return quota_id + return "" + + @classmethod + def _google_error_detail_evidence(cls, details: Any) -> tuple[str, str, Optional[float]]: + if not isinstance(details, list): + return "", "", None + + reason = "" + quota_id = "" + retry_delay_seconds = None + for detail in details: + if not isinstance(detail, dict): + continue + reason = cls._google_error_reason(detail) or reason + quota_id = cls._google_quota_id(detail) or quota_id + if str(detail.get(GOOGLE_ERROR_DETAIL_TYPE_FIELD, "")).endswith("RetryInfo"): + retry_delay_seconds = cls._parse_google_duration(detail.get("retryDelay")) + return reason, quota_id, retry_delay_seconds + + @classmethod + def _google_error_evidence_from_payload( + cls, payload: Any, fallback_status: Optional[int] = None + ) -> GoogleErrorEvidence: + error_payload = cls._google_error_payload(payload) + if not error_payload: + return GoogleErrorEvidence(status_code=fallback_status) + + status = error_payload.get("status") + message = error_payload.get("message") + reason, quota_id, retry_delay_seconds = cls._google_error_detail_evidence(error_payload.get("details")) + + return GoogleErrorEvidence( + status_code=cls._google_error_status_code(error_payload, fallback_status), + status=status if isinstance(status, str) else "", + message=message if isinstance(message, str) else "", + reason=reason, + quota_id=quota_id, + retry_delay_seconds=retry_delay_seconds, + ) + + @staticmethod + def _google_exception_response(error: Exception) -> Any: + return getattr(error, "response", None) or getattr(error, "_response", None) + + @staticmethod + def _google_exception_status_code(error: Exception, response: Any) -> Optional[int]: + if isinstance(error, ResourceExhausted): + fallback_status = 429 + elif isinstance(error, PermissionDenied): + fallback_status = 403 + elif isinstance(error, InvalidArgument): + fallback_status = 400 + else: + fallback_status = None + + error_code = getattr(error, "code", None) + if isinstance(error_code, int): + fallback_status = error_code + response_status = getattr(response, "status_code", None) + return response_status if isinstance(response_status, int) else fallback_status + + @staticmethod + def _google_exception_payload(error: Exception, response: Any) -> Any: + payload = getattr(error, "details", None) + if payload: + return payload + if response is not None and callable(getattr(response, "json", None)): + try: + return response.json() + except (TypeError, ValueError): + pass + + errors = getattr(error, "errors", None) or getattr(error, "_errors", None) + if isinstance(errors, dict): + return errors + if isinstance(errors, (list, tuple)): + return next((item for item in errors if isinstance(item, dict)), None) + return None + + def _extract_google_error_evidence(self, error: Exception) -> GoogleErrorEvidence: + if isinstance(error, GoogleGenAIUpstreamError): + return error.evidence + + response = self._google_exception_response(error) + fallback_status = self._google_exception_status_code(error, response) + payload = self._google_exception_payload(error, response) + evidence = self._google_error_evidence_from_payload(payload, fallback_status) + if evidence.message or evidence.reason or evidence.quota_id: + return evidence + + # Google API exceptions expose the provider's top-level message separately. + message = getattr(error, "message", "") + return GoogleErrorEvidence( + status_code=fallback_status, + message=message if isinstance(message, str) else "", + retry_delay_seconds=self._extract_retry_after_seconds(error), + ) + + def _extract_image_error_evidence(self, response: httpx.Response) -> GoogleErrorEvidence: try: payload = response.json() except ValueError: payload = None - - if isinstance(payload, dict): - error_payload = payload.get("error") - if isinstance(error_payload, dict): - message = error_payload.get("message") - if message: - return f"{response.status_code} {message}" - - body_text = response.text.strip() - if body_text: - return f"{response.status_code} {body_text}" - return f"{response.status_code} {response.reason_phrase}" + return self._google_error_evidence_from_payload(payload, response.status_code) async def _initialize_request_context( self, @@ -2753,7 +2901,7 @@ async def _run_image_generation_request( ), ) if response.is_error: - raise RuntimeError(self._extract_image_error_message(response)) + raise GoogleGenAIUpstreamError(self._extract_image_error_evidence(response)) image_response = response.json() if self._is_imagen_generation_model(context.model_name): @@ -3127,7 +3275,12 @@ def _resolve_observation_state( return actual_key_suffix, actual_proxy, key_verified, proxy_verified def _build_request_error_from_context( - self, context: GoogleGenAICompletionContext, message: str + self, + context: GoogleGenAICompletionContext, + message: str, + *, + status_code: Optional[int] = None, + error_type: str = "api_error", ) -> GoogleGenAIRequestError: return GoogleGenAIRequestError( message, @@ -3137,23 +3290,39 @@ def _build_request_error_from_context( api_key_index=context.api_key_index, api_key_total=context.api_key_total, proxy_used=self._proxy_info_to_string(context.proxy_info), + status_code=status_code, + error_type=error_type, ) async def _record_completion_failure( self, context: GoogleGenAICompletionContext, - error_message: str, + error_evidence: GoogleErrorEvidence, status_code: Optional[int], ) -> None: if context.api_key and context.model_name: + error_message = self._quota_error_record(error_evidence) await self._update_api_key_stats( context.api_key, context.model_name, success=False, error=error_message, status_code=status_code, + error_evidence=error_evidence, ) + @staticmethod + def _quota_error_record(evidence: GoogleErrorEvidence) -> str: + """Persist bounded classification fields instead of an upstream response body.""" + fields = [f"status={evidence.status_code or 'unknown'}"] + if evidence.status: + fields.append(f"google_status={evidence.status}") + if evidence.reason: + fields.append(f"reason={evidence.reason}") + if evidence.quota_id: + fields.append(f"quota_id={evidence.quota_id}") + return "google_upstream " + " ".join(fields) + def _extract_retry_after_seconds(self, error: Exception) -> Optional[float]: retry_after_seconds = None error_details = getattr(error, "errors", None) @@ -3165,6 +3334,12 @@ def _extract_retry_after_seconds(self, error: Exception) -> Optional[float]: return retry_after_seconds def _extract_status_code_from_exception(self, error: Exception) -> Optional[int]: + structured_status = self._extract_google_error_evidence(error).status_code + if structured_status is not None: + return structured_status + + # Compatibility for locally generated validation errors. This helper does + # not decide key eligibility; selector state only consumes typed evidence. error_str = str(error).lower() if "403" in error_str or "permission" in error_str: return 403 @@ -3179,33 +3354,48 @@ def _extract_status_code_from_exception(self, error: Exception) -> Optional[int] async def _handle_completion_exception( self, context: GoogleGenAICompletionContext, error: Exception ) -> GoogleGenAIRequestError: + evidence = self._extract_google_error_evidence(error) + status_code = evidence.status_code + if status_code is None: + status_code = self._extract_status_code_from_exception(error) + if status_code is not None: + # The legacy fallback is status-only: never let opaque text affect + # key invalidation or daily-exhaustion classification. + evidence = replace(evidence, status_code=status_code) if isinstance(error, ResourceExhausted): - error_msg = f"Google GenAI quota exhausted: {error}" - await self._record_completion_failure(context, error_msg, None) - quota_error = self._build_request_error_from_context(context, error_msg) - quota_error.retry_after_seconds = self._extract_retry_after_seconds(error) + status_code = 429 + evidence = GoogleErrorEvidence( + status_code=status_code, + status=evidence.status, + message=evidence.message, + reason=evidence.reason, + quota_id=evidence.quota_id, + retry_delay_seconds=evidence.retry_delay_seconds or self._extract_retry_after_seconds(error), + ) + await self._record_completion_failure(context, evidence, status_code) + quota_error = self._build_request_error_from_context( + context, "Google GenAI quota exhausted.", status_code=status_code + ) + quota_error.retry_after_seconds = evidence.retry_delay_seconds return quota_error if isinstance(error, PermissionDenied): - error_msg = f"Google GenAI permission denied: {error}" - await self._record_completion_failure(context, error_msg, 403) - return self._build_request_error_from_context(context, error_msg) + status_code = 403 + await self._record_completion_failure(context, evidence, status_code) + return self._build_request_error_from_context(context, "Google GenAI request was denied.", status_code=status_code) if isinstance(error, InvalidArgument): - error_msg = f"Google GenAI invalid argument: {error}" - await self._record_completion_failure(context, error_msg, 400) - return self._build_request_error_from_context(context, error_msg) + status_code = 400 + await self._record_completion_failure(context, evidence, status_code) + return self._build_request_error_from_context( + context, "Google GenAI request is invalid.", status_code=status_code, error_type="invalid_request_error" + ) - error_msg = f"Google GenAI error: {error}" if context.proxy_url and self._is_proxy_connectivity_error(error): self._mark_proxy_health(context.proxy_url, success=False, error=str(error)) - await self._record_completion_failure( - context, - error_msg, - self._extract_status_code_from_exception(error), - ) - return self._build_request_error_from_context(context, error_msg) + await self._record_completion_failure(context, evidence, status_code) + return self._build_request_error_from_context(context, "Google GenAI request failed.", status_code=status_code) async def get_api_key_stats(self, include_unused_models: bool = False) -> Dict[str, Dict[str, Any]]: """Get current API key usage statistics (grouped by API key) diff --git a/smolrouter/mediator.py b/smolrouter/mediator.py index 03d704f..d03997f 100644 --- a/smolrouter/mediator.py +++ b/smolrouter/mediator.py @@ -14,7 +14,7 @@ from .interfaces import IModelStrategy, IAccessControl, ModelInfo, ClientContext from .caching import ModelAggregator, IModelCache from .providers import IModelProvider -from .google_genai_provider import GOOGLE_GENAI_IMAGE_ENDPOINT, GoogleGenAIProvider +from .google_genai_provider import GOOGLE_GENAI_IMAGE_ENDPOINT, GoogleGenAIProvider, GoogleGenAIRequestError from .dummy_provider import DummyProvider from .load_balancer import model_load_balancer from .request_metadata import RequestMetadata @@ -401,10 +401,11 @@ async def _handle_google_provider_request( self._attach_lb_instance(lb_instance, metadata), ) except Exception as e: - error_msg = str(e) - error_type = "api_error" - status_code = self._google_error_status_code(error_msg) - if status_code == 400: + is_provider_error = isinstance(e, GoogleGenAIRequestError) + error_msg = str(e) if is_provider_error else "Google GenAI request failed." + status_code = getattr(e, "status_code", None) or self._google_error_status_code(error_msg) + error_type = getattr(e, "error_type", "api_error") + if status_code == 400 and error_type == "api_error": error_type = "invalid_request_error" error_metadata = RequestMetadata( diff --git a/smolrouter/redis_backend.py b/smolrouter/redis_backend.py index 7601245..d8f5a4f 100644 --- a/smolrouter/redis_backend.py +++ b/smolrouter/redis_backend.py @@ -8,6 +8,7 @@ import logging import os import hashlib +import json from datetime import datetime, timedelta, timezone from typing import Dict, List, Optional, Any, Mapping, TypedDict, Unpack from zoneinfo import ZoneInfo @@ -30,6 +31,7 @@ REDIS_EMPTY_VALUES = ("", "None", None) REDIS_STATUS_NONE_VALUES = ("", "0", "None", None) _BODY_STORAGE_UPDATE_UNSET = object() +GOOGLE_INVALID_KEY_RECOVERY_AUDIT_LIMIT = 100 _COMPLETE_ONCE_LUA_SCRIPT = """ local inflight_key = KEYS[1] local completed_key = KEYS[2] @@ -332,7 +334,7 @@ def _prepare_quota_record_data(data: Dict[str, Any]) -> Dict[str, Any]: quota_data = dict(data) last_reset = quota_data.get("last_reset", "") - if not quota_data.get("last_reset_date"): + if last_reset: quota_data["last_reset_date"] = last_reset if not quota_data.get("api_key_hash") and quota_data.get("key_hash"): @@ -727,6 +729,7 @@ def record_failure(self): redis.call('HSET', quota_key, 'requests_today', 0) redis.call('HSET', quota_key, 'tokens_today', 0) redis.call('HSET', quota_key, 'last_reset', today) + redis.call('HSET', quota_key, 'last_reset_date', today) -- Clear quota exhaustion status on daily reset redis.call('HDEL', quota_key, 'quota_exhausted_at') redis.call('HDEL', quota_key, 'quota_cooldown_until') @@ -1147,6 +1150,10 @@ def google_rotary_counter_key(provider_id: str, model_name: str) -> str: def google_invalid_keys_key(provider_id: str) -> str: return f"google_invalid_keys:{provider_id}" + @staticmethod + def google_invalid_key_metadata_key(provider_id: str, api_key_hash: str) -> str: + return f"google_invalid_key_metadata:{provider_id}:{api_key_hash}" + @staticmethod async def get_or_create_quota( api_key: str, @@ -1221,6 +1228,7 @@ async def _increment_usage_with_pipeline_fallback( "requests_today": 0, "tokens_today": 0, "last_reset": today, + "last_reset_date": today, }, ) await client.hdel(quota_key, "quota_exhausted_at") @@ -1372,7 +1380,7 @@ async def _get_google_api_key_fallback_availability( return "invalid", None if quota.quota_cooldown_until and quota.quota_cooldown_until > now_dt: return "cooling_down", quota.quota_cooldown_until - if quota.last_reset_date == today and quota.quota_exhausted_at: + if getattr(quota, "last_reset", quota.last_reset_date) == today and quota.quota_exhausted_at: return "exhausted", None return "available", None @@ -1519,8 +1527,15 @@ async def get_provider_usage(provider_id: str) -> List["QuotaRecord"]: return quotas @staticmethod - async def mark_invalid(api_key_hash: str, provider_name: str) -> int: - """Mark all quota entries for a given key hash as invalid. + async def mark_invalid( + api_key_hash: str, + provider_name: str, + *, + reason: str = "operator", + status_code: Optional[int] = None, + request_id: Optional[str] = None, + ) -> int: + """Atomically mark one provider/key as invalid with bounded audit metadata. Args: api_key_hash: The hashed API key @@ -1531,23 +1546,87 @@ async def mark_invalid(api_key_hash: str, provider_name: str) -> int: """ client = get_redis() - # Get all quota keys for this key hash quota_keys = await client.smembers(f"quotas:by_key:{api_key_hash}") - await client.sadd(RedisApiKeyQuota.google_invalid_keys_key(provider_name), api_key_hash) + provider_quota_keys = [] + for quota_key in quota_keys: + if await client.hget(quota_key, "provider_id") == provider_name: + provider_quota_keys.append(quota_key) + + now = datetime.now(timezone.utc).isoformat() + metadata_key = RedisApiKeyQuota.google_invalid_key_metadata_key(provider_name, api_key_hash) + pipe = client.pipeline(transaction=True) + pipe.sadd(RedisApiKeyQuota.google_invalid_keys_key(provider_name), api_key_hash) + pipe.hsetnx(metadata_key, "first_reason", reason) + pipe.hsetnx(metadata_key, "first_status_code", str(status_code or "")) + pipe.hsetnx(metadata_key, "first_request_id", request_id or "") + pipe.hsetnx(metadata_key, "first_recorded_at", now) + pipe.hincrby(metadata_key, "occurrence_count", 1) + pipe.hset( + metadata_key, + mapping={ + "latest_reason": reason, + "latest_status_code": str(status_code or ""), + "latest_request_id": request_id or "", + "latest_recorded_at": now, + }, + ) + for quota_key in provider_quota_keys: + pipe.hset( + quota_key, + mapping={ + "invalid_key": "true", + "invalid_reason": reason, + "invalid_status_code": str(status_code or ""), + "invalidated_at": now, + "updated_at": now, + }, + ) + await pipe.execute() - count = 0 + logger.debug( + "Marked %s quota entries invalid for key_hash=%s provider=%s reason=%s", + len(provider_quota_keys), + api_key_hash, + provider_name, + reason, + ) + return len(provider_quota_keys) + + @staticmethod + async def recover_invalid(api_key_hash: str, provider_name: str, *, actor: str, reason: str) -> int: + """Atomically restore a key's eligibility while retaining quota state.""" + client = get_redis() + quota_keys = await client.smembers(f"quotas:by_key:{api_key_hash}") + provider_quota_keys = [] for quota_key in quota_keys: - # quota_key format: "quota:{provider_id}:{key_hash}:{model_name}" - # We need to match keys that have this provider_name in the right position - parts = quota_key.split(":") - if len(parts) >= 4 and parts[0] == "quota" and parts[1] == provider_name: - # This quota key matches both the key_hash (already filtered) and provider_name - await client.hset(quota_key, "invalid_key", "true") - await client.hset(quota_key, "updated_at", datetime.now(timezone.utc).isoformat()) - count += 1 - - logger.debug(f"Marked {count} quota entries as invalid for key_hash={api_key_hash}, provider={provider_name}") - return count + if await client.hget(quota_key, "provider_id") == provider_name: + provider_quota_keys.append(quota_key) + + metadata_key = RedisApiKeyQuota.google_invalid_key_metadata_key(provider_name, api_key_hash) + prior_metadata = await client.hgetall(metadata_key) + audit_key = f"google_invalid_key_recovery_audit:{provider_name}" + now = datetime.now(timezone.utc).isoformat() + pipe = client.pipeline(transaction=True) + pipe.srem(RedisApiKeyQuota.google_invalid_keys_key(provider_name), api_key_hash) + pipe.delete(metadata_key) + for quota_key in provider_quota_keys: + pipe.hdel(quota_key, "invalid_key", "invalid_reason", "invalid_status_code", "invalidated_at") + pipe.lpush( + audit_key, + json.dumps( + { + "actor": actor, + "reason": reason, + "api_key_hash": api_key_hash, + "recovered_at": now, + "prior_metadata": prior_metadata, + }, + sort_keys=True, + ), + ) + pipe.ltrim(audit_key, 0, GOOGLE_INVALID_KEY_RECOVERY_AUDIT_LIMIT - 1) + await pipe.execute() + return len(provider_quota_keys) @staticmethod async def mark_quota_exhausted( diff --git a/tests/unit/test_database.py b/tests/unit/test_database.py index d1ed2cd..edaebea 100644 --- a/tests/unit/test_database.py +++ b/tests/unit/test_database.py @@ -3,7 +3,7 @@ import smolrouter.database as database from unittest.mock import AsyncMock -from smolrouter.database import RequestLogEntry, estimate_tokens_from_request +from smolrouter.database import ApiKeyQuota, RequestLogEntry, estimate_tokens_from_request import pytest @@ -53,6 +53,21 @@ def fake_asyncio_run(coro): assert len(fallback_calls) == 1 +@pytest.mark.asyncio +async def test_api_key_quota_mark_invalid_by_hash_forwards_audit_metadata(monkeypatch): + mark_invalid = AsyncMock(return_value=2) + monkeypatch.setattr(database.RedisApiKeyQuota, "mark_invalid", mark_invalid) + + marked = await ApiKeyQuota.mark_invalid_by_hash( + "hashed-key", "google", reason="api_key_invalid", status_code=400, request_id="request-123" + ) + + assert marked == 2 + mark_invalid.assert_awaited_once_with( + "hashed-key", "google", reason="api_key_invalid", status_code=400, request_id="request-123" + ) + + @pytest.mark.asyncio async def test_record_exception_event_aggregates_signatures_and_summary(isolated_db): signature_payload = ValueError("failure 0x7f00 user@example.com") diff --git a/tests/unit/test_google_genai_helpers.py b/tests/unit/test_google_genai_helpers.py index 0673b63..374a5f9 100644 --- a/tests/unit/test_google_genai_helpers.py +++ b/tests/unit/test_google_genai_helpers.py @@ -16,6 +16,7 @@ import httpx import pytest +from google.genai import errors as google_genai_errors from smolrouter import google_genai_provider as ggp from smolrouter.google_genai_provider import ( @@ -924,14 +925,171 @@ def test_quota_cooldown_seconds_ambiguous_quota_is_transient(provider): def test_is_invalid_key_error(provider): - assert provider._is_invalid_key_error(status_code=403) is True - assert provider._is_invalid_key_error("403 upgrade your account to a paid plan", status_code=403) is False + assert provider._is_invalid_key_error(status_code=403) is False + assert provider._is_invalid_key_error("400 INVALID_ARGUMENT: Request contains an invalid argument.") is False assert provider._is_invalid_key_error("API key not valid") is True - assert provider._is_invalid_key_error("PERMISSION_DENIED") is True + assert provider._is_invalid_key_error("invalid API key") is True + assert provider._is_invalid_key_error("PERMISSION_DENIED") is False assert provider._is_invalid_key_error("transient network blip") is False assert provider._is_invalid_key_error(None) is False +def test_google_error_evidence_requires_explicit_key_specific_signal(provider): + generic_invalid_argument = { + "error": { + "code": 400, + "message": "Request contains an invalid argument.", + "status": "INVALID_ARGUMENT", + } + } + explicit_invalid_key = { + "error": { + "code": 400, + "message": "API key not valid.", + "status": "INVALID_ARGUMENT", + "details": [ + {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "reason": "API_KEY_INVALID"} + ], + } + } + + generic_evidence = provider._google_error_evidence_from_payload(generic_invalid_argument) + explicit_evidence = provider._google_error_evidence_from_payload(explicit_invalid_key) + + assert provider._invalid_key_reason(generic_evidence) is None + assert provider._invalid_key_reason(explicit_evidence) == "api_key_invalid" + + +def test_google_error_evidence_parses_detail_types_and_rejects_malformed_fields(provider): + payload = { + "error": { + "code": 429, + "message": "quota exhausted", + "status": "RESOURCE_EXHAUSTED", + "details": [ + None, + {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "reason": "RATE_LIMITED"}, + { + "@type": "type.googleapis.com/google.rpc.QuotaFailure", + "violations": [None, {"quotaId": "GenerateRequestsPerMinutePerProjectPerModel-FreeTier"}], + }, + {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "12.5s"}, + ], + } + } + malformed_payload = { + "error": { + "code": "not-a-status", + "message": None, + "status": None, + "details": [{"@type": "type.googleapis.com/google.rpc.QuotaFailure", "violations": "invalid"}], + } + } + + evidence = provider._google_error_evidence_from_payload(payload) + malformed = provider._google_error_evidence_from_payload(malformed_payload, fallback_status=502) + + assert evidence.reason == "RATE_LIMITED" + assert evidence.quota_id == "GenerateRequestsPerMinutePerProjectPerModel-FreeTier" + assert evidence.retry_delay_seconds == 12.5 + assert provider._is_per_minute_quota_error(evidence=evidence) is True + assert ( + provider._is_quota_exhausted_error( + evidence=ggp.GoogleErrorEvidence(status_code=None, status="RESOURCE_EXHAUSTED") + ) + is True + ) + assert malformed == ggp.GoogleErrorEvidence(status_code=502) + assert provider._google_error_evidence_from_payload(None, fallback_status=503) == ggp.GoogleErrorEvidence( + status_code=503 + ) + assert provider._google_error_evidence_from_payload({"error": []}, fallback_status=504) == ggp.GoogleErrorEvidence( + status_code=504 + ) + assert provider._google_quota_id( + {"@type": "type.googleapis.com/google.rpc.QuotaFailure", "violations": []} + ) == "" + + +def test_google_duration_parser_handles_numeric_invalid_and_non_string_values(provider): + assert provider._parse_google_duration(4) == 4.0 + assert provider._parse_google_duration(1.5) == 1.5 + assert provider._parse_google_duration(" 2.25s ") == 2.25 + assert provider._parse_google_duration("bad") is None + assert provider._parse_google_duration(None) is None + + +def test_google_sdk_error_evidence_preserves_explicit_daily_quota_id(provider): + error = google_genai_errors.ClientError( + 429, + { + "error": { + "code": 429, + "message": "You exceeded your current quota.", + "status": "RESOURCE_EXHAUSTED", + "details": [ + { + "@type": "type.googleapis.com/google.rpc.QuotaFailure", + "violations": [ + {"quotaId": "GenerateRequestsPerDayPerProjectPerModel-FreeTier"} + ], + } + ], + } + }, + ) + + evidence = provider._extract_google_error_evidence(error) + + assert evidence.status_code == 429 + assert evidence.quota_id == "GenerateRequestsPerDayPerProjectPerModel-FreeTier" + assert provider._is_per_day_quota_error(evidence=evidence) is True + + +def test_google_error_evidence_uses_response_and_legacy_exception_fields(provider): + response_payload = {"error": {"code": 418, "message": "response payload", "status": "TEAPOT"}} + response_error = RuntimeError("ignored") + response_error.details = {} + response_error.response = SimpleNamespace(status_code=503, json=lambda: response_payload) + + legacy_error = RuntimeError("ignored") + legacy_error.errors = [{"error": {"code": 401, "message": "legacy payload", "status": "UNAUTHENTICATED"}}] + + invalid_response_error = RuntimeError("ignored") + invalid_response_error.response = SimpleNamespace( + status_code=502, + json=lambda: (_ for _ in ()).throw(ValueError("invalid JSON")), + ) + invalid_response_error.errors = {"error": {"code": 500, "message": "fallback payload", "status": "INTERNAL"}} + + message_error = RuntimeError("ignored") + message_error.message = "top-level SDK message" + + assert provider._extract_google_error_evidence(response_error) == ggp.GoogleErrorEvidence( + status_code=418, status="TEAPOT", message="response payload" + ) + assert provider._extract_google_error_evidence(legacy_error) == ggp.GoogleErrorEvidence( + status_code=401, status="UNAUTHENTICATED", message="legacy payload" + ) + assert provider._extract_google_error_evidence(invalid_response_error) == ggp.GoogleErrorEvidence( + status_code=500, status="INTERNAL", message="fallback payload" + ) + assert provider._extract_google_error_evidence(message_error) == ggp.GoogleErrorEvidence( + status_code=None, message="top-level SDK message" + ) + + +def test_google_image_error_evidence_handles_invalid_json(provider): + class InvalidJsonResponse: + status_code = 502 + + @staticmethod + def json(): + raise ValueError("invalid JSON") + + assert provider._extract_image_error_evidence(InvalidJsonResponse()) == ggp.GoogleErrorEvidence(status_code=502) + + def test_extract_retry_delay(provider): assert provider._extract_retry_delay("Please retry in 20.5s") == 20.5 # pattern 1 assert provider._extract_retry_delay("retryDelay': '30s'") == 30.0 # pattern 2 @@ -941,6 +1099,10 @@ def test_extract_retry_delay(provider): def test_extract_status_code_from_exception(provider): + structured_error = RuntimeError("ignored") + structured_error.code = 418 + + assert provider._extract_status_code_from_exception(structured_error) == 418 assert provider._extract_status_code_from_exception(Exception("Got 403 permission")) == 403 assert provider._extract_status_code_from_exception(Exception("429 quota")) == 429 assert provider._extract_status_code_from_exception(Exception("401 unauthorized")) == 401 diff --git a/tests/unit/test_new_integrations.py b/tests/unit/test_new_integrations.py index f14573b..d21d352 100644 --- a/tests/unit/test_new_integrations.py +++ b/tests/unit/test_new_integrations.py @@ -18,6 +18,7 @@ from smolrouter.interfaces import ModelInfo, ProviderConfig, ProxyConfig, coerce_provider_proxy_settings from smolrouter.mediator import ModelMediator from smolrouter.google_genai_provider import ( + GoogleErrorEvidence, GoogleGenAICompletionContext, GoogleGenAIConfig, GoogleGenAIProvider, @@ -718,7 +719,7 @@ async def test_google_genai_generate_imagen_releases_rate_limiter_on_error(mock_ mock_client.return_value.post = AsyncMock(return_value=mock_response) mock_client.return_value.aclose = AsyncMock(return_value=None) - with pytest.raises(GoogleGenAIRequestError, match="paid plans"): + with pytest.raises(GoogleGenAIRequestError, match="Google GenAI request failed"): await provider.generate_image( { "model": "imagen-4.0-generate-001", @@ -769,6 +770,79 @@ async def test_google_genai_handle_completion_exception_branches(): ) +@pytest.mark.asyncio +async def test_google_genai_generic_invalid_argument_is_safe_and_non_terminal(): + provider = _make_google_provider() + context = GoogleGenAICompletionContext( + original_model="original", + observation_id="obs-400", + model_name="gemini-2.0-flash", + api_key="test-key", + ) + provider._record_completion_failure = AsyncMock() + error = InvalidArgument( + "Request contains an invalid argument.", + errors=[ + { + "error": { + "code": 400, + "message": "Request contains an invalid argument.", + "status": "INVALID_ARGUMENT", + } + } + ], + ) + + public_error = await provider._handle_completion_exception(context, error) + + assert public_error.status_code == 400 + assert public_error.error_type == "invalid_request_error" + assert str(public_error) == "Google GenAI request is invalid." + evidence = provider._record_completion_failure.await_args.args[1] + assert provider._invalid_key_reason(evidence) is None + assert evidence.message == "Request contains an invalid argument." + + +@pytest.mark.asyncio +async def test_google_genai_unstructured_rate_limit_preserves_status_only(): + provider = _make_google_provider() + context = GoogleGenAICompletionContext( + original_model="original", + observation_id="obs-unstructured-429", + model_name="gemini-2.0-flash", + api_key="test-key", + ) + provider._update_api_key_stats = AsyncMock() + + public_error = await provider._handle_completion_exception(context, RuntimeError("429 quota exhausted")) + + assert public_error.status_code == 429 + assert str(public_error) == "Google GenAI request failed." + provider._update_api_key_stats.assert_awaited_once() + assert provider._update_api_key_stats.await_args.kwargs["error"] == "google_upstream status=429" + assert provider._update_api_key_stats.await_args.kwargs["error_evidence"].status_code == 429 + + +@pytest.mark.asyncio +async def test_google_failure_recording_uses_bounded_evidence_fields(): + provider = _make_google_provider() + context = GoogleGenAICompletionContext(original_model="original", observation_id="obs-no-key", model_name="gemini") + provider._update_api_key_stats = AsyncMock() + evidence = GoogleErrorEvidence( + status_code=429, + status="RESOURCE_EXHAUSTED", + reason="RATE_LIMITED", + quota_id="GenerateRequestsPerMinutePerProjectPerModel-FreeTier", + ) + + assert provider._quota_error_record(evidence) == ( + "google_upstream status=429 google_status=RESOURCE_EXHAUSTED reason=RATE_LIMITED " + "quota_id=GenerateRequestsPerMinutePerProjectPerModel-FreeTier" + ) + await provider._record_completion_failure(context, evidence, 429) + provider._update_api_key_stats.assert_not_awaited() + + @pytest.mark.asyncio async def test_google_genai_proxy_health_and_transport_helpers(): provider = _make_google_provider(proxy_pool_enabled=True, proxy_pool=[ProxyConfig(https_proxy="http://127.0.0.1:8888")]) @@ -974,7 +1048,15 @@ async def test_google_genai_update_api_key_stats_covers_success_and_error_branch ) provider._get_quota_record = AsyncMock( - side_effect=[success_quota, invalid_quota, cooldown_quota, quota_exhausted, regular_quota, entitlement_quota] + side_effect=[ + success_quota, + invalid_quota, + cooldown_quota, + quota_exhausted, + regular_quota, + regular_quota, + entitlement_quota, + ] ) with patch("smolrouter.redis_backend.RedisApiKeyQuota.increment_usage", AsyncMock(return_value=None)), patch( @@ -989,11 +1071,12 @@ async def test_google_genai_update_api_key_stats_covers_success_and_error_branch "test-key", "gemini-2.0-flash", success=False, - error="permission denied 403", - status_code=403, + error="API key not valid", + status_code=400, + error_evidence=GoogleErrorEvidence(status_code=400, message="API key not valid"), ) assert api_key_backend.mark_invalid_by_hash.await_count == 1 - assert provider.config.api_keys == [] + assert provider.config.api_keys == ["test-key"] provider.config.api_keys = ["test-key"] await provider._update_api_key_stats( @@ -1031,6 +1114,17 @@ async def test_google_genai_update_api_key_stats_covers_success_and_error_branch assert regular_quota.error_count == 1 assert api_key_backend.mark_error.await_count == 1 + await provider._update_api_key_stats( + "test-key", + "gemini-2.0-flash", + success=False, + error="400 INVALID_ARGUMENT: Request contains an invalid argument.", + status_code=400, + ) + assert api_key_backend.mark_invalid_by_hash.await_count == 1 + assert api_key_backend.mark_error.await_count == 2 + assert provider.config.api_keys == ["test-key"] + provider.config.api_keys = ["test-key"] await provider._update_api_key_stats( "test-key", @@ -1041,7 +1135,7 @@ async def test_google_genai_update_api_key_stats_covers_success_and_error_branch ) assert entitlement_quota.error_count == 1 assert api_key_backend.mark_invalid_by_hash.await_count == 1 - assert api_key_backend.mark_error.await_count == 2 + assert api_key_backend.mark_error.await_count == 3 assert provider.config.api_keys == ["test-key"] diff --git a/tests/unit/test_redis_backend.py b/tests/unit/test_redis_backend.py index f3bd271..b8fc086 100644 --- a/tests/unit/test_redis_backend.py +++ b/tests/unit/test_redis_backend.py @@ -1,3 +1,4 @@ +import json from datetime import datetime, timedelta, timezone import pytest @@ -530,6 +531,71 @@ async def test_api_key_quota_round_trip_and_usage_markers(monkeypatch): assert provider_b_usage[0].invalid_key is False +@pytest.mark.asyncio +async def test_invalid_key_recovery_is_provider_scoped_and_preserves_quota_state(monkeypatch): + await redis_client.flushall() + monkeypatch.setattr(redis_backend_module, "_circuit_breaker", redis_backend_module.RedisCircuitBreaker()) + monkeypatch.setattr(RedisApiKeyQuota, "_script_initialized", False) + monkeypatch.setattr(RedisApiKeyQuota, "_script_sha", None) + monkeypatch.setattr(RedisApiKeyQuota, "_lua_disabled", True) + + await RedisApiKeyQuota.get_or_create_quota("sk-recover", "provider-a", "model-a") + await RedisApiKeyQuota.get_or_create_quota("sk-recover", "provider-a", "model-b") + await RedisApiKeyQuota.get_or_create_quota("sk-recover", "provider-b", "model-a") + await RedisApiKeyQuota.increment_usage("sk-recover", "provider-a", "model-a", request_count=2) + key_hash = RedisApiKeyQuota.hash_api_key("sk-recover") + + assert await RedisApiKeyQuota.mark_invalid( + key_hash, "provider-a", reason="api_key_invalid", status_code=400, request_id="request-1" + ) == 2 + assert await RedisApiKeyQuota.mark_invalid( + key_hash, "provider-a", reason="api_key_revoked", status_code=401, request_id="request-2" + ) == 2 + + metadata_key = RedisApiKeyQuota.google_invalid_key_metadata_key("provider-a", key_hash) + metadata = await redis_client.hgetall(metadata_key) + assert metadata["first_reason"] == "api_key_invalid" + assert metadata["first_status_code"] == "400" + assert metadata["first_request_id"] == "request-1" + assert metadata["occurrence_count"] == "2" + assert metadata["latest_reason"] == "api_key_revoked" + assert metadata["latest_status_code"] == "401" + assert metadata["latest_request_id"] == "request-2" + + assert await RedisApiKeyQuota.recover_invalid( + key_hash, "provider-a", actor="operator@example.test", reason="false positive" + ) == 2 + + provider_a_usage = await RedisApiKeyQuota.get_provider_usage("provider-a") + provider_b_usage = await RedisApiKeyQuota.get_provider_usage("provider-b") + assert {quota.model_name for quota in provider_a_usage if quota.invalid_key} == set() + assert next(quota for quota in provider_a_usage if quota.model_name == "model-a").requests_today == 2 + assert provider_b_usage[0].invalid_key is False + assert await redis_client.sismember(RedisApiKeyQuota.google_invalid_keys_key("provider-a"), key_hash) == 0 + assert await redis_client.exists(metadata_key) == 0 + + audit = json.loads(await redis_client.lindex("google_invalid_key_recovery_audit:provider-a", 0)) + assert audit["actor"] == "operator@example.test" + assert audit["reason"] == "false positive" + assert audit["prior_metadata"]["occurrence_count"] == "2" + + +@pytest.mark.asyncio +async def test_invalid_key_recovery_audit_is_bounded(monkeypatch): + await redis_client.flushall() + monkeypatch.setattr(redis_backend_module, "GOOGLE_INVALID_KEY_RECOVERY_AUDIT_LIMIT", 1) + await RedisApiKeyQuota.get_or_create_quota("sk-audit", "provider-a", "model-a") + key_hash = RedisApiKeyQuota.hash_api_key("sk-audit") + + for actor in ("first@example.test", "second@example.test"): + await RedisApiKeyQuota.mark_invalid(key_hash, "provider-a") + await RedisApiKeyQuota.recover_invalid(key_hash, "provider-a", actor=actor, reason="test") + + entries = await redis_client.lrange("google_invalid_key_recovery_audit:provider-a", 0, -1) + assert len(entries) == 1 + assert json.loads(entries[0])["actor"] == "second@example.test" + + @pytest.mark.asyncio async def test_mark_quota_cooldown_round_trips_and_clears_on_daily_reset(monkeypatch): await redis_client.flushall()