""" LLM Client with retry logic, exponential backoff, and safety guards. R6 Implementation: Network retry logic with exponential backoff. Integrates with: - outbound.py (SSRF protection, payload sanitization) - session_manager.py (shared session, rate/concurrency limits) IMPORTANT SAFETY NOTES: - POST requests may cause duplicate charges if retried incorrectly - Streaming requests should NOT be retried after connection established - All requests use allow_redirects=False and treat any 3xx as failure (SSRF protection) """ import asyncio import aiohttp import uuid import time import random import logging from dataclasses import dataclass, field from typing import Optional, Tuple, Dict, Any try: from .services.outbound_policy import is_destination_policy_error except ImportError as import_error: from import_compat import ensure_absolute_import_fallback_allowed ensure_absolute_import_fallback_allowed(import_error) from services.outbound_policy import is_destination_policy_error logger = logging.getLogger(__name__) # Stable failure kinds for callers that must not reflect raw upstream/network text. FAILURE_DESTINATION_NOT_ALLOWED = "destination_not_allowed" FAILURE_HTTP_STATUS = "http_status" FAILURE_REDIRECT = "redirect" FAILURE_RATE_LIMITED = "rate_limited" FAILURE_TIMEOUT = "timeout" FAILURE_NETWORK = "network" @dataclass class RetryConfig: """Configuration for retry behavior.""" # Retry limits max_retries: int = 2 # Conservative default # Timing base_delay: float = 1.0 # Initial delay (seconds) max_delay: float = 30.0 # Cap on delay exponential_base: float = 2.0 # Backoff multiplier # Timeouts request_timeout_seconds: float = 60.0 # Per-attempt timeout total_timeout_seconds: float = 180.0 # Entire operation timeout # Retryable conditions retryable_network_errors: bool = True # ClientError (timeout, refused) retry_on_429: bool = True # Rate limit (with Retry-After) retry_on_5xx: bool = False # Disabled by default (risk of dupe) max_5xx_retries: int = 1 # If enabled, limit to 1 retryable_5xx_codes: Tuple[int, ...] = (502, 503, 504) # Exclude 500 # Idempotency (OpenAI standard) add_idempotency_key: bool = True idempotency_header: str = "Idempotency-Key" @dataclass class RetryResult: """Result of a retry-aware request.""" response: Optional[aiohttp.ClientResponse] success: bool attempts: int total_time: float error: Optional[str] = None idempotency_key: Optional[str] = None failure_kind: Optional[str] = None status: Optional[int] = None def _safe_release(response: aiohttp.ClientResponse) -> None: """ Safely release response resources. Note: response.close() is synchronous in aiohttp. This prevents connection pool leaks when retrying. """ try: response.close() except Exception: pass def _get_retry_after( response: aiohttp.ClientResponse, config: RetryConfig, attempt: int ) -> float: """Parse Retry-After header or calculate backoff.""" retry_after = response.headers.get("Retry-After") if retry_after: try: return min(float(retry_after), config.max_delay) except ValueError: pass return _calculate_backoff(attempt, config) def _calculate_backoff(attempt: int, config: RetryConfig) -> float: """Calculate exponential backoff with jitter.""" delay = config.base_delay * (config.exponential_base ** (attempt - 1)) jitter = random.uniform(0, delay * 0.1) # 10% jitter return min(delay + jitter, config.max_delay) async def llm_request_with_retry( session: aiohttp.ClientSession, method: str, url: str, *, json: Optional[Dict[str, Any]] = None, headers: Optional[Dict[str, str]] = None, config: Optional[RetryConfig] = None, is_streaming: bool = False, ) -> RetryResult: """ Make HTTP request with safe retry logic. Args: session: aiohttp ClientSession to use method: HTTP method (GET, POST, etc.) url: Target URL json: JSON payload (optional) headers: HTTP headers (optional) config: Retry configuration (optional, uses defaults) is_streaming: If True, marks this as a streaming request Returns: RetryResult with response, success status, and metadata IMPORTANT: - For streaming: retries are DISABLED after connection established - Caller must use `async with result.response:` for proper cleanup - All requests use allow_redirects=False (SSRF protection) Example: result = await llm_request_with_retry(session, "POST", url, json=payload) if result.success: async with result.response as resp: data = await resp.json() """ config = config or RetryConfig() headers = dict(headers) if headers else {} # Add idempotency key for POST requests idempotency_key: Optional[str] = None if config.add_idempotency_key and method.upper() == "POST": idempotency_key = str(uuid.uuid4()) headers[config.idempotency_header] = idempotency_key logger.debug(f"Added idempotency key: {idempotency_key}") start_time = time.monotonic() attempts = 0 last_error: Optional[str] = None failure_kind: Optional[str] = None response: Optional[aiohttp.ClientResponse] = None while attempts <= config.max_retries: # Check total timeout budget elapsed = time.monotonic() - start_time if elapsed >= config.total_timeout_seconds: logger.warning(f"Total timeout exceeded ({config.total_timeout_seconds}s)") return RetryResult( response=None, success=False, attempts=attempts, total_time=elapsed, error=f"Total timeout exceeded ({config.total_timeout_seconds}s)", idempotency_key=idempotency_key, failure_kind=FAILURE_TIMEOUT, ) attempts += 1 remaining_time = config.total_timeout_seconds - elapsed per_request_timeout = min(config.request_timeout_seconds, remaining_time) try: timeout = aiohttp.ClientTimeout(total=per_request_timeout) response = await session.request( method, url, json=json, headers=headers, allow_redirects=False, # SSRF protection (required) timeout=timeout, ) # === Redirects are failures (3xx) === # CRITICAL: never treat a redirect as success; a 3xx from an approved destination must not # be followed or handed to callers as a usable response (SSRF redirect pivot). if 300 <= response.status < 400: _safe_release(response) logger.warning(f"Redirect response {response.status} rejected (redirects are not followed)") return RetryResult( response=None, success=False, attempts=attempts, total_time=time.monotonic() - start_time, error=f"HTTP {response.status} redirect not followed", idempotency_key=idempotency_key, failure_kind=FAILURE_REDIRECT, status=response.status, ) # === Success path (2xx) === if response.status < 300: logger.debug(f"Request succeeded on attempt {attempts}") return RetryResult( response=response, success=True, attempts=attempts, total_time=time.monotonic() - start_time, idempotency_key=idempotency_key, ) # === 429 Rate Limit === if response.status == 429 and config.retry_on_429: delay = _get_retry_after(response, config, attempts) _safe_release(response) if attempts <= config.max_retries: logger.warning( f"429 Rate Limited, waiting {delay:.1f}s " f"(attempt {attempts}/{config.max_retries + 1})" ) await asyncio.sleep(delay) continue else: return RetryResult( response=None, success=False, attempts=attempts, total_time=time.monotonic() - start_time, error="Rate limit exceeded after max retries", idempotency_key=idempotency_key, failure_kind=FAILURE_RATE_LIMITED, status=429, ) # === 5xx Server Error === if response.status in config.retryable_5xx_codes and config.retry_on_5xx: if attempts <= config.max_5xx_retries: delay = _calculate_backoff(attempts, config) _safe_release(response) logger.warning( f"{response.status} Server Error, retry in {delay:.1f}s " f"(attempt {attempts}/{config.max_5xx_retries + 1})" ) await asyncio.sleep(delay) continue # === Non-retryable status (4xx except 429, 500, or exhausted retries) === logger.warning(f"Non-retryable status {response.status}") return RetryResult( response=response, success=False, attempts=attempts, total_time=time.monotonic() - start_time, error=f"HTTP {response.status}", idempotency_key=idempotency_key, failure_kind=FAILURE_HTTP_STATUS, status=response.status, ) except asyncio.TimeoutError: last_error = f"Request timeout ({per_request_timeout:.1f}s)" failure_kind = FAILURE_TIMEOUT logger.warning(f"Timeout on attempt {attempts}") if config.retryable_network_errors and attempts <= config.max_retries: delay = _calculate_backoff(attempts, config) logger.info(f"Retrying in {delay:.1f}s") await asyncio.sleep(delay) continue break except aiohttp.ClientError as e: if is_destination_policy_error(e): # Policy refusals are final; never retry towards a forbidden address. last_error = "Destination refused by the outbound policy" failure_kind = FAILURE_DESTINATION_NOT_ALLOWED logger.warning("Outbound destination refused at connect time") break last_error = str(e) failure_kind = FAILURE_NETWORK logger.warning(f"Network error on attempt {attempts}: {e}") if config.retryable_network_errors and attempts <= config.max_retries: delay = _calculate_backoff(attempts, config) logger.info(f"Retrying in {delay:.1f}s") await asyncio.sleep(delay) continue break # Exhausted all retries logger.error(f"Request failed after {attempts} attempts: {last_error}") return RetryResult( response=response, success=False, attempts=attempts, total_time=time.monotonic() - start_time, error=last_error, idempotency_key=idempotency_key, failure_kind=failure_kind, ) # Convenience function for non-streaming POST requests async def llm_post_with_retry( session: aiohttp.ClientSession, url: str, *, json: Optional[Dict[str, Any]] = None, headers: Optional[Dict[str, str]] = None, config: Optional[RetryConfig] = None, ) -> RetryResult: """ Convenience wrapper for POST requests with retry. See llm_request_with_retry for full documentation. """ return await llm_request_with_retry( session, "POST", url, json=json, headers=headers, config=config, is_streaming=False )