diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 27d9aed52b4..528a5c10903 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -289,6 +289,7 @@ router_settings: | database_connection_pool_timeout | integer | Database connection pool timeout in seconds | | disable_error_logs | boolean | If true, suppresses error tracking and storage in the database | | enable_health_check_routing | boolean | If true, enables health check-driven request routing to avoid unhealthy deployments | +| health_check_ignore_transient_errors | boolean | If true, 429 (rate limit) and 408 (timeout) health check failures are ignored and do not affect routing or cooldown | | enable_mcp_registry | boolean | If true, enables access to the centralized MCP server registry | | enforce_rbac | boolean | If true, enables role-based access control (RBAC) for all proxy operations | | forward_llm_provider_auth_headers | boolean | If true, forwards provider-specific auth headers to LLM API calls | @@ -397,6 +398,7 @@ router_settings: | guardrail_list | List[GuardrailTypedDict] | List of guardrail configurations for guardrail load balancing. Enables load balancing across multiple guardrail deployments with the same guardrail_name. [Further Docs](./guardrails/guardrail_load_balancing.md) | | enable_health_check_routing | boolean | If true, enables health check-driven deployment filtering to avoid routing requests to unhealthy deployments | | health_check_staleness_threshold | integer | Maximum age in seconds for cached health check results before marking deployments as stale | +| health_check_ignore_transient_errors | boolean | If true, 429 (rate limit) and 408 (timeout) health check failures are ignored and do not affect routing or cooldown | ### environment variables - Reference @@ -821,6 +823,7 @@ router_settings: | LITELLM_KEY_ROTATION_ENABLED | Enable auto-key rotation for LiteLLM (boolean). Default is false. | LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS | Interval in seconds for how often to run job that auto-rotates keys. Default is 86400 (24 hours). | LITELLM_KEY_ROTATION_GRACE_PERIOD | Duration to keep old key valid after rotation (e.g. "24h", "2d"). Default is empty (immediate revoke). Used for scheduled rotations and as fallback when not specified in regenerate request. +| LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS | TTL in seconds for the distributed lock used by the key rotation job. Default is 600 (10 minutes). | LITELLM_LICENSE | License key for LiteLLM usage | LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS | Set to `True` to use the local bundled Anthropic beta headers config only, disabling remote fetching. Default is `False` | LITELLM_LOCAL_BLOG_POSTS | When set to `True`, uses the local bundled blog posts only, disabling remote fetching from GitHub. Default is `False` diff --git a/docs/my-website/docs/proxy/guardrails/guardrail_policies.md b/docs/my-website/docs/proxy/guardrails/guardrail_policies.md index f4411553c69..18c9025da6c 100644 --- a/docs/my-website/docs/proxy/guardrails/guardrail_policies.md +++ b/docs/my-website/docs/proxy/guardrails/guardrail_policies.md @@ -311,7 +311,7 @@ Response: ## Policy Flow Builder -For conditional execution (e.g., run a second guardrail only if the first fails), use the [Policy Flow Builder](./policy_flow_builder) to define pipelines with per-step pass/fail actions. +For conditional execution (e.g., run a second guardrail only if the first fails), use the [Policy Flow Builder](./policy_flow_builder) to define pipelines with per-step **pass**, **fail**, and optional **error** actions (`on_pass`, `on_fail`, `on_error`). ## Config Reference @@ -337,7 +337,7 @@ policies: | `guardrails.add` | `list[string]` | Guardrails to enable. | | `guardrails.remove` | `list[string]` | Guardrails to disable (useful with inheritance). | | `condition.model` | `string` or `list[string]` | Optional. Only apply when model matches. Supports regex. | -| `pipeline` | `object` | Optional. Ordered guardrail execution with per-step actions. See [Policy Flow Builder](./policy_flow_builder). | +| `pipeline` | `object` | Optional. Ordered guardrail execution with per-step actions (`on_pass`, `on_fail`, optional `on_error`). See [Policy Flow Builder](./policy_flow_builder). | ### `policy_attachments` diff --git a/docs/my-website/docs/proxy/guardrails/policy_flow_builder.md b/docs/my-website/docs/proxy/guardrails/policy_flow_builder.md index 2a83f3768ab..630930aa893 100644 --- a/docs/my-website/docs/proxy/guardrails/policy_flow_builder.md +++ b/docs/my-website/docs/proxy/guardrails/policy_flow_builder.md @@ -1,8 +1,8 @@ # Policy Flow Builder -The Policy Flow Builder lets you design guardrail pipelines with **conditional execution**. Instead of running guardrails independently, you chain them into ordered steps and control what happens when each guardrail passes or fails. +The Policy Flow Builder lets you design guardrail pipelines with **conditional execution**. Instead of running guardrails independently, you chain them into ordered steps and control what happens when each guardrail **passes**, **fails a policy check** (content intervention), or hits a **technical error** (e.g. timeout, unreachable provider, missing guardrail). -Two powerful patterns it enables: **guardrail fallbacks** (try a different guardrail when one fails) and **retrying the same guardrail** (run the same guardrail again if it fails, e.g. to handle transient errors). +Two powerful patterns it enables: **guardrail fallbacks** (try a different guardrail when one fails) and **retrying the same guardrail** (run the same guardrail again if it fails, e.g. to handle transient errors). With **`on_error`**, you can treat **technical** failures differently from **policy** failures—for example, fall back to another provider when the primary API errors, while still blocking on flagged content. ## When to use the Flow Builder @@ -19,6 +19,7 @@ Use the Flow Builder when you need: - **Custom responses** — return a specific message when a guardrail fails instead of a generic block - **Data chaining** — pass modified data (e.g., PII-masked content) from one step to the next - **Fine-grained control** — different actions on pass vs. fail per step +- **Technical-error routing** — set `on_error` separately from `on_fail` so outages or timeouts can **allow**, **block**, **go to the next step**, or return a **custom response** without conflating them with content violations ## Concepts @@ -29,24 +30,37 @@ A pipeline has: - **Mode**: `pre_call` (before the LLM) or `post_call` (after the LLM) - **Steps**: Ordered list of guardrail steps +### Outcomes: pass, fail, and error + +Each step run produces one of three outcomes: + +| Outcome | Meaning | Typical cause | +|--------|---------|----------------| +| **pass** | Guardrail completed without blocking | Content allowed, or data was modified and returned | +| **fail** | Policy intervention | Guardrail raised an intervention (e.g. flagged content, blocked request) | +| **error** | Technical failure | Timeouts, network errors, guardrail not registered, or other non-intervention exceptions | + +`on_pass` and `on_fail` apply to **pass** and **fail** respectively. **`on_error`** applies only to **error**. If `on_error` is omitted, the pipeline uses **`on_fail`** for error outcomes (backward compatible). + ### Step actions -Each step defines what happens when the guardrail **passes** and when it **fails**: +For each step you choose an action for **pass**, **fail**, and optionally **error**. Allowed values are: `next`, `allow`, `block`, `modify_response`. | Action | Description | |--------|-------------| -| **Next Step** | Continue to the next guardrail in the pipeline | -| **Allow** | Stop the pipeline and allow the request to proceed | -| **Block** | Stop the pipeline and block the request | -| **Custom Response** | Return a custom message instead of the default block | +| **Next Step** (`next`) | Continue to the next guardrail in the pipeline | +| **Allow** (`allow`) | Stop the pipeline and allow the request to proceed | +| **Block** (`block`) | Stop the pipeline and block the request | +| **Custom Response** (`modify_response`) | Return a custom message instead of the default block | ### Step options | Field | Type | Description | |-------|------|--------------| | `guardrail` | `string` | Name of the guardrail to run | -| `on_pass` | `string` | Action when guardrail passes: `next`, `allow`, `block`, `modify_response` | -| `on_fail` | `string` | Action when guardrail fails: `next`, `allow`, `block`, `modify_response` | +| `on_pass` | `string` | Action when outcome is **pass**: `next`, `allow`, `block`, `modify_response` | +| `on_fail` | `string` | Action when outcome is **fail** (policy intervention): `next`, `allow`, `block`, `modify_response` | +| `on_error` | `string` (optional) | Action when outcome is **error** (technical). If omitted, **error** uses `on_fail`. | | `pass_data` | `boolean` | Forward modified request data (e.g., PII-masked) to the next step | | `modify_response_message` | `string` | Custom message when using `modify_response` action | @@ -57,7 +71,7 @@ Each step defines what happens when the guardrail **passes** and when it **fails 3. Select **Flow Builder** (instead of the simple form) 4. Design your flow: - **Trigger** — Incoming LLM request (runs when the policy matches) - - **Steps** — Add guardrails, set ON PASS and ON FAIL actions per step + - **Steps** — Add guardrails, set **ON PASS**, **ON FAIL**, and **ON ERROR** actions per step (ON ERROR is optional; when unset, errors follow ON FAIL) - **End** — Request proceeds to the LLM 5. Use the **+** between steps to insert new steps 6. Use the **Test** panel to run sample messages through the pipeline before saving @@ -151,6 +165,37 @@ policies: First attempt passes → allow. First attempt fails → retry the same guardrail; second pass → allow, second fail → block. +## Technical errors vs policy failures (`on_error`) + +Use **`on_error`** when you want different behavior for **API/infra problems** than for **content policy** violations. + +- **`on_fail`** — Runs when the guardrail **intervenes** (e.g. toxic content, PII detected). +- **`on_error`** — Runs when the step ends in **error** (timeout, connection failure, guardrail not loaded, etc.). If you omit `on_error`, **error** outcomes use **`on_fail`**. + +Example: block on bad content, but if the primary scanner is down, fall back to a second guardrail instead of blocking every request: + +```yaml +policies: + error-fallback-policy: + guardrails: + add: + - primary_scanner + - backup_scanner + pipeline: + mode: pre_call + steps: + - guardrail: primary_scanner + on_pass: allow + on_fail: block + on_error: next + - guardrail: backup_scanner + on_pass: allow + on_fail: block + on_error: allow +``` + +If `primary_scanner` errors → run `backup_scanner`. If `backup_scanner` errors → allow the request (set `on_error` to `block` if you prefer fail-closed). + ## Example: Custom response on fail Return a branded message instead of a generic block: diff --git a/docs/my-website/docs/proxy/health.md b/docs/my-website/docs/proxy/health.md index 530bea3d06b..1d893961b62 100644 --- a/docs/my-website/docs/proxy/health.md +++ b/docs/my-website/docs/proxy/health.md @@ -316,86 +316,9 @@ general_settings: ## Health Check Driven Routing -By default, background health checks are observability-only — they populate the `/health` endpoint but don't affect routing. Unhealthy deployments still receive traffic until request failures trigger cooldown. +Route traffic away from unhealthy deployments proactively — before user requests hit them. Supports per-error-type failure thresholds, transient error suppression, and automatic safety nets. -With `enable_health_check_routing: true`, the router **excludes deployments that failed their last background health check** before selecting a candidate. This gives you proactive failover instead of reactive cooldown. - -### How it works - -1. Background health checks run on their configured interval -2. After each cycle, every deployment is marked healthy or unhealthy -3. On each incoming request, the router filters out unhealthy deployments **before** cooldown filtering and load balancing -4. If all deployments are unhealthy, the filter is bypassed (safety net — never causes a total outage) -5. If health state is stale (older than `health_check_staleness_threshold`), it is ignored - -### Quick start - -```yaml -model_list: - - model_name: gpt-4 - litellm_params: - model: openai/gpt-4 - api_key: os.environ/OPENAI_API_KEY - - model_name: gpt-4 - litellm_params: - model: openai/gpt-4 - api_key: os.environ/OPENAI_API_KEY_SECONDARY - -general_settings: - background_health_checks: true - health_check_interval: 60 - enable_health_check_routing: true -``` - -### Configuration - -| Setting | Where | Default | Description | -|---------|-------|---------|-------------| -| `enable_health_check_routing` | `general_settings` | `false` | Enable/disable health-check-driven routing | -| `health_check_staleness_threshold` | `general_settings` | `health_check_interval * 2` | Seconds before health state is considered stale and ignored | -| `background_health_checks` | `general_settings` | `false` | Must be `true` for health check routing to work | -| `health_check_interval` | `general_settings` | `300` | Seconds between health check cycles | - -### Interaction with cooldown - -Health check filtering and cooldown are **additive**. A deployment can be excluded by either mechanism: - -- **Health check filter** — proactive, runs on the configured interval, excludes deployments that failed the last check -- **Cooldown** — reactive, triggered by request failures, excludes deployments for a short TTL - -This means request failures still provide fast detection between health check intervals. - -### Staleness - -If a health check result is older than `health_check_staleness_threshold`, it is ignored and the deployment is treated as eligible. This prevents stale data from permanently excluding a deployment if the health check loop stops or slows down. - -The default staleness threshold is `health_check_interval * 2`. For a 60s interval, health state expires after 120s. - -### Example: custom staleness - -```yaml -general_settings: - background_health_checks: true - health_check_interval: 30 - enable_health_check_routing: true - health_check_staleness_threshold: 90 # ignore health state older than 90s -``` - -### Debugging - -Run the proxy with `--detailed_debug` and look for: - -``` -health_check_routing_state_updated healthy=3 unhealthy=1 -``` - -This is logged after each health check cycle when routing state is written. - -If the safety net triggers (all deployments unhealthy), you'll see: - -``` -All deployments marked unhealthy by health checks, bypassing health filter -``` +See the full guide: [Health Check Driven Routing](./health_check_routing.md) ## Health Check Timeout diff --git a/docs/my-website/docs/proxy/health_check_routing.md b/docs/my-website/docs/proxy/health_check_routing.md new file mode 100644 index 00000000000..daf0b19212c --- /dev/null +++ b/docs/my-website/docs/proxy/health_check_routing.md @@ -0,0 +1,340 @@ +# Health Check Driven Routing + +Route traffic away from unhealthy deployments before users hit errors. Background health checks run on a configurable interval, and any deployment that fails gets removed from the routing pool proactively, not after a user request already failed. + + +## Architecture + + + {/* Background */} + + + {/* LEFT PANEL: Background health check loop */} + + Background Loop + every health_check_interval seconds + + {/* Deployment A */} + + Deployment A + ahealth_check() → 200 ✓ + + {/* Deployment B */} + + Deployment B + ahealth_check() → 401 ✗ + + {/* Deployment C */} + + Deployment C + ahealth_check() → 429 ⚡ + + {/* ignore_transient box */} + + ignore_transient_errors: true + 429 / 408 → ignored + not written to cache + + {/* allowed_fails_policy box */} + + allowed_fails_policy + 401 → increment counter + counter > threshold + → cooldown triggered + + {/* CENTER PANEL: Shared State */} + + Shared State + + {/* Health State Cache */} + + DeploymentHealthCache + A → healthy ✓ + B → unhealthy ✗ + C → not written (ignored) + TTL: staleness_threshold × 1.5 + + {/* Cooldown Cache */} + + Cooldown Cache + B → cooling down + (after policy threshold) + TTL: cooldown_time + + {/* failed_calls counter */} + + failed_calls counter + B: 2 / AuthAllowedFails: 1 + → threshold exceeded + TTL: cooldown_time (must > interval) + + {/* RIGHT PANEL: Request path */} + + Request Path + + {/* Incoming request */} + + Incoming request + + {/* All deployments */} + + All deployments [A, B, C] + + + + {/* Health check filter */} + + ① Health Check Filter + if policy set → bypass + else → remove unhealthy + + + + {/* Cooldown filter */} + + ② Cooldown Filter + remove deployments in cooldown + + + + {/* Safety net */} + + Safety Net + if all removed → return all + + + + {/* Load balancer */} + + ③ Load Balancer + + + + {/* Selected deployment */} + + Selected: Deployment A ✓ + + + + {/* ARROWS: left → center */} + + + + + + {/* ARROWS: center → right */} + + + + {/* Arrow markers */} + + + + + + + + + + + + + + + + + + + + + + + +## What problem does this solve? + +By default, LiteLLM routes traffic to all deployments and only stops sending to a broken one after it has already failed a user request. The cooldown system is reactive. + +Health check driven routing makes this **proactive**: a background loop pings every deployment on a configurable interval. If a deployment fails its health check, it gets removed from the routing pool immediately, before a user request lands on it. + +When you also set `allowed_fails_policy`, you control exactly how many health check failures of each error type (auth errors, rate limits, timeouts) are needed before a deployment enters cooldown. This avoids false positives from transient noise. + + +## Setup + +### Step 1: Enable background health checks + +Background health checks are off by default. Turn them on in `general_settings`: + +```yaml +general_settings: + background_health_checks: true + health_check_interval: 60 # seconds between each full check cycle +``` + +### Step 2: Enable health check routing + +```yaml +general_settings: + background_health_checks: true + health_check_interval: 60 + enable_health_check_routing: true # ← route away from unhealthy deployments +``` + +At this point, any deployment that fails its health check is immediately excluded from routing until the next check cycle clears it. + +### Step 3: Add a policy to control how many failures trigger cooldown + +Without a policy, the first health check failure marks a deployment as unhealthy. If you want more tolerance (e.g., only act after 2 consecutive auth failures), use `allowed_fails_policy`: + +```yaml +model_list: + - model_name: claude-sonnet + litellm_params: + model: anthropic/claude-sonnet-4-5 + api_key: os.environ/ANTHROPIC_API_KEY + + - model_name: claude-sonnet + litellm_params: + model: anthropic/claude-sonnet-4-5 + api_key: os.environ/ANTHROPIC_API_KEY_SECONDARY + +general_settings: + background_health_checks: true + health_check_interval: 30 + enable_health_check_routing: true + +router_settings: + cooldown_time: 60 # how long a deployment stays in cooldown + allowed_fails_policy: + AuthenticationErrorAllowedFails: 1 # cooldown after 2nd auth failure + TimeoutErrorAllowedFails: 3 # cooldown after 4th timeout +``` + +When `allowed_fails_policy` is set, the binary health check filter is bypassed. Only the cooldown system controls routing exclusion, and it only fires after your configured threshold is crossed. + +### Step 4 (optional): Ignore transient errors + +429 (rate limit) and 408 (timeout) from a health check usually mean the deployment is temporarily overloaded, not broken. To prevent these from affecting routing at all: + +```yaml +general_settings: + background_health_checks: true + health_check_interval: 30 + enable_health_check_routing: true + health_check_ignore_transient_errors: true # 429 and 408 never affect routing +``` + +With this on, only hard failures (401, 404, 5xx) from health checks contribute to cooldown. + + +## Full example + +```yaml +model_list: + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY + + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY_SECONDARY + + - model_name: gpt-4o + litellm_params: + model: azure/gpt-4o + api_base: os.environ/AZURE_API_BASE + api_key: os.environ/AZURE_API_KEY + +general_settings: + background_health_checks: true + health_check_interval: 30 + enable_health_check_routing: true + health_check_ignore_transient_errors: true + +router_settings: + cooldown_time: 60 + allowed_fails_policy: + AuthenticationErrorAllowedFails: 0 # cooldown immediately on auth failure + TimeoutErrorAllowedFails: 2 # cooldown after 3 timeouts + RateLimitErrorAllowedFails: 5 # cooldown after 6 rate limits (if not ignoring transients) +``` + + +## Configuration reference + +| Setting | Where | Default | Description | +|---|---|---|---| +| `enable_health_check_routing` | `general_settings` | `false` | Route away from deployments that fail health checks | +| `background_health_checks` | `general_settings` | `false` | Must be `true` for health check routing to work | +| `health_check_interval` | `general_settings` | `300` | Seconds between full health check cycles | +| `health_check_staleness_threshold` | `general_settings` | `interval x 2` | Seconds before cached health state is ignored | +| `health_check_ignore_transient_errors` | `general_settings` | `false` | Ignore 429 and 408 from health checks; these never affect routing | +| `cooldown_time` | `router_settings` | `5` | Seconds a deployment stays in cooldown after threshold is crossed | +| `allowed_fails_policy` | `router_settings` | `null` | Per-error-type failure thresholds before cooldown (see below) | + +### `allowed_fails_policy` fields + +| Field | Error type | HTTP status | +|---|---|---| +| `AuthenticationErrorAllowedFails` | Bad API key | 401 | +| `TimeoutErrorAllowedFails` | Request timeout | 408 | +| `RateLimitErrorAllowedFails` | Rate limit exceeded | 429 | +| `BadRequestErrorAllowedFails` | Malformed request | 400 | +| `ContentPolicyViolationErrorAllowedFails` | Content filtered | 400 | + +The value is the number of failures **tolerated** before cooldown. `0` means cooldown on the first failure. `2` means cooldown on the third. + + +## Things to keep in mind + +- **Counter TTL must be longer than the health check interval.** `allowed_fails_policy` works by incrementing a `failed_calls` counter per deployment. That counter expires after `cooldown_time` seconds. If `cooldown_time` is shorter than `health_check_interval`, the counter resets between every check cycle and failures never accumulate. Set `cooldown_time` greater than `health_check_interval` when using `allowed_fails_policy`. + + ```yaml + router_settings: + cooldown_time: 60 # must be > health_check_interval (30s here) + + general_settings: + health_check_interval: 30 + ``` + +- **`AllowedFails: N` means cooldown on the (N+1)th failure.** The counter check is `updated_fails > allowed_fails`, so `0` triggers on the 1st failure, `1` on the 2nd, `2` on the 3rd. + + | `AllowedFails` | Cooldown triggers after | + |---|---| + | `0` | 1st failure | + | `1` | 2nd failure | + | `2` | 3rd failure | + +- **Without `allowed_fails_policy`, the first failure is enough.** The first failed health check immediately excludes the deployment from routing. Use `allowed_fails_policy` when you want tolerance for flaky checks. + +- **If all deployments are unhealthy, the filter is bypassed.** Traffic keeps flowing rather than returning no deployment at all. Requests will fail, but the router keeps trying. + +- **Health check failures and request failures share the same counters.** When `allowed_fails_policy` is set, both sources increment the same `failed_calls` counter. A deployment at 1 health check failure that then receives 1 failing request will hit the threshold for `AllowedFails: 1` and enter cooldown. + + +## Debugging + +Run the proxy with `--detailed_debug` and look for these log lines: + +After each health check cycle (written at DEBUG level): +``` +health_check_routing_state_updated healthy=2 unhealthy=1 +``` + +When a health check failure increments the counter and triggers cooldown (DEBUG level): +``` +checks 'should_run_cooldown_logic' +Attempting to add to cooldown list +``` + +When safety net fires because all deployments are in cooldown: +``` +All deployments in cooldown via health-check routing, bypassing cooldown filter +``` + +When safety net fires because all deployments are unhealthy (binary filter, no `allowed_fails_policy`): +``` +All deployments marked unhealthy by health checks, bypassing health filter +``` diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index b514ea2234c..ab4ae46ec3b 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -1051,7 +1051,8 @@ const sidebars = { "proxy/fallback_management", "proxy/tag_routing", "proxy/timeout", - "wildcard_routing" + "wildcard_routing", + "proxy/health_check_routing" ], }, { diff --git a/litellm/constants.py b/litellm/constants.py index 252068bd7b0..a9facabb010 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1319,6 +1319,9 @@ LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS = int( LITELLM_KEY_ROTATION_GRACE_PERIOD: str = os.getenv( "LITELLM_KEY_ROTATION_GRACE_PERIOD", "" ) # Duration to keep old key valid after rotation (e.g. "24h", "2d"); empty = immediate revoke (default) +LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS = int( + os.getenv("LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS", 600) +) # 10 minutes default — caps the deadlock window if a pod crashes mid-rotation UI_SESSION_TOKEN_TEAM_ID = "litellm-dashboard" LITELLM_PROXY_ADMIN_NAME = "default_user_id" @@ -1347,6 +1350,7 @@ CLOUDZERO_MAX_FETCHED_DATA_RECORDS = int( os.getenv("CLOUDZERO_MAX_FETCHED_DATA_RECORDS", 50000) ) SPEND_LOG_CLEANUP_JOB_NAME = "spend_log_cleanup" +KEY_ROTATION_JOB_NAME = "litellm_key_rotation_job" SPEND_LOG_RUN_LOOPS = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500)) SPEND_LOG_CLEANUP_BATCH_SIZE = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000)) SPEND_LOG_QUEUE_SIZE_THRESHOLD = int(os.getenv("SPEND_LOG_QUEUE_SIZE_THRESHOLD", 100)) diff --git a/litellm/main.py b/litellm/main.py index cbedd1735c7..ddd37b47536 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -3792,9 +3792,9 @@ def completion( # type: ignore # noqa: PLR0915 "aws_region_name" not in optional_params or optional_params["aws_region_name"] is None ): - optional_params[ - "aws_region_name" - ] = aws_bedrock_client.meta.region_name + optional_params["aws_region_name"] = ( + aws_bedrock_client.meta.region_name + ) bedrock_route = BedrockModelInfo.get_bedrock_route(model) if bedrock_route == "converse": @@ -6214,9 +6214,9 @@ def adapter_completion( new_kwargs = translation_obj.translate_completion_input_params(kwargs=kwargs) response: Union[ModelResponse, CustomStreamWrapper] = completion(**new_kwargs) # type: ignore - translated_response: Optional[ - Union[BaseModel, AdapterCompletionStreamWrapper] - ] = None + translated_response: Optional[Union[BaseModel, AdapterCompletionStreamWrapper]] = ( + None + ) if isinstance(response, ModelResponse): translated_response = translation_obj.translate_completion_output_params( response=response @@ -6396,9 +6396,9 @@ async def atranscription(*args, **kwargs) -> TranscriptionResponse: if existing_duration is None: calculated_duration = calculate_request_duration(file) if calculated_duration is not None: - response._hidden_params[ - "audio_transcription_duration" - ] = calculated_duration + response._hidden_params["audio_transcription_duration"] = ( + calculated_duration + ) return response except Exception as e: @@ -6621,9 +6621,9 @@ def transcription( if existing_duration is None: calculated_duration = calculate_request_duration(file) if calculated_duration is not None: - response._hidden_params[ - "audio_transcription_duration" - ] = calculated_duration + response._hidden_params["audio_transcription_duration"] = ( + calculated_duration + ) if response is None: raise ValueError("Unmapped provider passed in. Unable to get the response.") @@ -6927,9 +6927,9 @@ def speech( # noqa: PLR0915 ElevenLabsTextToSpeechConfig.ELEVENLABS_QUERY_PARAMS_KEY ] = query_params - litellm_params_dict[ - ElevenLabsTextToSpeechConfig.ELEVENLABS_VOICE_ID_KEY - ] = voice_id + litellm_params_dict[ElevenLabsTextToSpeechConfig.ELEVENLABS_VOICE_ID_KEY] = ( + voice_id + ) if api_base is not None: litellm_params_dict["api_base"] = api_base @@ -7250,7 +7250,8 @@ async def ahealth_check( if mode is None: return { - "error": f"error:{str(e)}. Missing `mode`. Set the `mode` for the model - https://docs.litellm.ai/docs/proxy/health#embedding-models \nstacktrace: {stack_trace}" + "error": f"error:{str(e)}. Missing `mode`. Set the `mode` for the model - https://docs.litellm.ai/docs/proxy/health#embedding-models \nstacktrace: {stack_trace}", + "exception": e, } error_to_return = str(e) + "\nstack trace: " + stack_trace @@ -7262,6 +7263,7 @@ async def ahealth_check( return { "error": error_to_return, "raw_request_typed_dict": raw_request_typed_dict, + "exception": e, } @@ -7508,9 +7510,9 @@ def stream_chunk_builder( # noqa: PLR0915 ] if len(content_chunks) > 0: - response["choices"][0]["message"][ - "content" - ] = processor.get_combined_content(content_chunks) + response["choices"][0]["message"]["content"] = ( + processor.get_combined_content(content_chunks) + ) thinking_blocks = [ chunk @@ -7521,9 +7523,9 @@ def stream_chunk_builder( # noqa: PLR0915 ] if len(thinking_blocks) > 0: - response["choices"][0]["message"][ - "thinking_blocks" - ] = processor.get_combined_thinking_content(thinking_blocks) + response["choices"][0]["message"]["thinking_blocks"] = ( + processor.get_combined_thinking_content(thinking_blocks) + ) reasoning_chunks = [ chunk @@ -7534,9 +7536,9 @@ def stream_chunk_builder( # noqa: PLR0915 ] if len(reasoning_chunks) > 0: - response["choices"][0]["message"][ - "reasoning_content" - ] = processor.get_combined_reasoning_content(reasoning_chunks) + response["choices"][0]["message"]["reasoning_content"] = ( + processor.get_combined_reasoning_content(reasoning_chunks) + ) annotation_chunks = [ chunk diff --git a/litellm/proxy/common_utils/key_rotation_manager.py b/litellm/proxy/common_utils/key_rotation_manager.py index 5a0a1fabc7d..aaf39a7a19d 100644 --- a/litellm/proxy/common_utils/key_rotation_manager.py +++ b/litellm/proxy/common_utils/key_rotation_manager.py @@ -11,6 +11,7 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import ( LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, LITELLM_KEY_ROTATION_GRACE_PERIOD, + LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS, ) from litellm.proxy._types import ( GenerateKeyResponse, @@ -30,14 +31,42 @@ class KeyRotationManager: Manages automated key rotation based on individual key rotation schedules. """ - def __init__(self, prisma_client: PrismaClient): + def __init__(self, prisma_client: PrismaClient, pod_lock_manager=None): self.prisma_client = prisma_client + self.pod_lock_manager = pod_lock_manager async def process_rotations(self): """ - Main entry point - find and rotate keys that are due for rotation + Main entry point - find and rotate keys that are due for rotation. + Uses PodLockManager to ensure only one pod runs rotation in multi-pod deployments. """ + from litellm.constants import KEY_ROTATION_JOB_NAME + + lock_acquired = False try: + # If we have a pod lock manager with Redis, try to acquire the lock + if self.pod_lock_manager and self.pod_lock_manager.redis_cache: + # Use a dedicated lock TTL (default 600s) instead of the check interval + # (which defaults to 86400s / 24h). Using the check interval would create + # a 24-hour deadlock window if a pod crashes before releasing the lock. + lock_ttl = max( + LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS, 300 + ) # At least 5 minutes, configurable via LITELLM_KEY_ROTATION_LOCK_TTL_SECONDS + lock_acquired = ( + await self.pod_lock_manager.acquire_lock( + cronjob_id=KEY_ROTATION_JOB_NAME, + ttl=lock_ttl, + ) + or False + ) + if not lock_acquired: + verbose_proxy_logger.warning( + "Key rotation: another pod is already running rotation " + "or Redis lock acquisition failed — skipping this cycle. " + "Keys will be rotated on the next cycle." + ) + return + verbose_proxy_logger.info("Starting scheduled key rotation check...") # Clean up expired deprecated keys first @@ -74,6 +103,16 @@ class KeyRotationManager: except Exception as e: verbose_proxy_logger.error(f"Key rotation process failed: {e}") + finally: + # Only release the lock if it was actually acquired + if ( + lock_acquired + and self.pod_lock_manager + and self.pod_lock_manager.redis_cache + ): + await self.pod_lock_manager.release_lock( + cronjob_id=KEY_ROTATION_JOB_NAME, + ) async def _find_keys_needing_rotation(self) -> List[LiteLLM_VerificationToken]: """ diff --git a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py index 546ea05998c..6435498ae03 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -32,6 +32,7 @@ class PodLockManager: async def acquire_lock( self, cronjob_id: str, + ttl: Optional[int] = None, ) -> Optional[bool]: """ Attempt to acquire the lock for a specific cron job using Redis. @@ -39,15 +40,20 @@ class PodLockManager: Args: cronjob_id: The ID of the cron job to lock + ttl: Optional custom TTL in seconds. Defaults to DEFAULT_CRON_JOB_LOCK_TTL_SECONDS. + Use a longer TTL for jobs that may take longer than the default 60s + (e.g. key rotation with many keys). """ if self.redis_cache is None: verbose_proxy_logger.debug("redis_cache is None, skipping acquire_lock") return None try: + lock_ttl = ttl or DEFAULT_CRON_JOB_LOCK_TTL_SECONDS verbose_proxy_logger.debug( - "Pod %s attempting to acquire Redis lock for cronjob_id=%s", + "Pod %s attempting to acquire Redis lock for cronjob_id=%s (ttl=%ds)", self.pod_id, cronjob_id, + lock_ttl, ) # Try to set the lock key with the pod_id as its value, only if it doesn't exist (NX) # and with an expiration (EX) to avoid deadlocks. @@ -56,7 +62,7 @@ class PodLockManager: lock_key, self.pod_id, nx=True, - ttl=DEFAULT_CRON_JOB_LOCK_TTL_SECONDS, + ttl=lock_ttl, ) if acquired: verbose_proxy_logger.info( @@ -133,11 +139,10 @@ class PodLockManager: ) else: verbose_proxy_logger.warning( - "Spend tracking - pod %s failed to release Redis lock for cronjob_id=%s. " - "Lock will expire after TTL=%ds.", + "Pod %s failed to release Redis lock for cronjob_id=%s. " + "Lock will expire after its TTL.", self.pod_id, cronjob_id, - DEFAULT_CRON_JOB_LOCK_TTL_SECONDS, ) else: verbose_proxy_logger.debug( diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 3e05ee3c484..5d1bcf31f84 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -21,6 +21,8 @@ ILLEGAL_DISPLAY_PARAMS = [ "vertex_credentials", "aws_access_key_id", "aws_secret_access_key", + "exception", # internal; not JSON-serializable, never for display + "litellm_metadata", # internal tracking metadata with auth objects; not for display ] MINIMAL_DISPLAY_PARAMS = ["model", "mode_error"] @@ -95,7 +97,12 @@ async def run_with_timeout(task, timeout): except asyncio.TimeoutError: # `asyncio.wait_for()` already cancels only the awaited task on timeout. # Do not cancel unrelated sibling health check tasks. - return {"error": "Timeout exceeded"} + timeout_exception = litellm.Timeout( + message="Health check timeout exceeded", + model="", + llm_provider="", + ) + return {"error": "Timeout exceeded", "exception": timeout_exception} async def _run_model_health_check(model: dict): @@ -204,6 +211,10 @@ async def _perform_health_check( healthy_endpoints = [] unhealthy_endpoints = [] + # Exceptions keyed by model_id; returned separately so callers can use + # them for cooldown integration without risking JSON-serialization errors + # in the /health response. + exceptions_by_model_id: dict = {} for is_healthy, model in zip(results, model_list): litellm_params = model["litellm_params"] @@ -218,14 +229,23 @@ async def _perform_health_check( cleaned = _clean_endpoint_data({**litellm_params, **is_healthy}, details) if _model_id: cleaned["model_id"] = _model_id + if "exception" in is_healthy: + exc = is_healthy["exception"] + exceptions_by_model_id[_model_id] = exc + # Store integer status code so shared-cache readers can + # reconstruct the transient-error filter without the exception object. + cleaned["exception_status"] = getattr(exc, "status_code", 500) unhealthy_endpoints.append(cleaned) else: cleaned = _clean_endpoint_data(litellm_params, details) if _model_id: cleaned["model_id"] = _model_id + if isinstance(is_healthy, Exception): + exceptions_by_model_id[_model_id] = is_healthy + cleaned["exception_status"] = getattr(is_healthy, "status_code", 500) unhealthy_endpoints.append(cleaned) - return healthy_endpoints, unhealthy_endpoints + return healthy_endpoints, unhealthy_endpoints, exceptions_by_model_id def build_deployment_health_states( @@ -366,7 +386,7 @@ async def perform_health_check( source, cycle_id, ) - return [], [] + return [], [], {} cycle_start_time = time.monotonic() requested_model_count = len(model_list) @@ -406,7 +426,11 @@ async def perform_health_check( ) try: - healthy_endpoints, unhealthy_endpoints = await _perform_health_check( + ( + healthy_endpoints, + unhealthy_endpoints, + exceptions_by_model_id, + ) = await _perform_health_check( model_list, details, max_concurrency=max_concurrency, @@ -438,4 +462,4 @@ async def perform_health_check( _rss_mb_for_log(), ) - return healthy_endpoints, unhealthy_endpoints + return healthy_endpoints, unhealthy_endpoints, exceptions_by_model_id diff --git a/litellm/proxy/health_check_utils/shared_health_check_manager.py b/litellm/proxy/health_check_utils/shared_health_check_manager.py index ae18a42c02b..2ecee5095b8 100644 --- a/litellm/proxy/health_check_utils/shared_health_check_manager.py +++ b/litellm/proxy/health_check_utils/shared_health_check_manager.py @@ -192,7 +192,7 @@ class SharedHealthCheckManager: model_list: List[Dict[str, Any]], details: bool = True, max_concurrency: Optional[int] = None, - ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]], Dict[str, Any]]: """ Perform health check with shared state coordination. @@ -217,6 +217,7 @@ class SharedHealthCheckManager: return ( cached_results.get("healthy_endpoints", []), cached_results.get("unhealthy_endpoints", []), + {}, ) # No recent cache, try to acquire lock @@ -231,7 +232,11 @@ class SharedHealthCheckManager: len(model_list), ) - healthy_endpoints, unhealthy_endpoints = await perform_health_check( + ( + healthy_endpoints, + unhealthy_endpoints, + exceptions_by_model_id, + ) = await perform_health_check( model_list=model_list, details=details, max_concurrency=max_concurrency, @@ -242,7 +247,7 @@ class SharedHealthCheckManager: healthy_endpoints, unhealthy_endpoints ) - return healthy_endpoints, unhealthy_endpoints + return healthy_endpoints, unhealthy_endpoints, exceptions_by_model_id finally: # Always release the lock @@ -262,6 +267,7 @@ class SharedHealthCheckManager: return ( cached_results.get("healthy_endpoints", []), cached_results.get("unhealthy_endpoints", []), + {}, ) # Still no cache, fall back to local health check diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index ef9436f2d8c..8a09edfd4c4 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -771,7 +771,7 @@ async def _perform_health_check_and_save( max_concurrency=None, ): """Helper function to perform health check and save results to database""" - healthy_endpoints, unhealthy_endpoints = await perform_health_check( + healthy_endpoints, unhealthy_endpoints, _ = await perform_health_check( model_list=model_list, cli_model=cli_model, model=target_model, diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 729b42ce638..3c5a1d67be4 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -74,7 +74,7 @@ class PipelineExecutor: duration = time.perf_counter() - start_time - action = step.on_pass if outcome == "pass" else step.on_fail + action = _pipeline_action_for_outcome(step, outcome) step_result = PipelineStepResult( guardrail_name=step.guardrail, @@ -206,6 +206,23 @@ class PipelineExecutor: return None +def _pipeline_action_for_outcome(step: PipelineStep, outcome: str) -> str: + """ + Map pipeline step outcome to the configured action. + + - pass -> on_pass + - fail -> on_fail (content/policy intervention) + - error -> on_error if set, else on_fail (backward compatible) + """ + if outcome == "pass": + return step.on_pass + if outcome == "fail": + return step.on_fail + if step.on_error is not None: + return step.on_error + return step.on_fail + + def _extract_error_message(e: Exception) -> str: """Extract a human-readable error message from a guardrail exception.""" if isinstance(e, ModifyResponseException): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2a661a8348e..fd88f44fbcf 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -635,9 +635,9 @@ except ImportError: server_root_path = get_server_root_path() _license_check = LicenseCheck() premium_user: bool = _license_check.is_premium() -premium_user_data: Optional[ - "EnterpriseLicenseData" -] = _license_check.airgapped_license_data +premium_user_data: Optional["EnterpriseLicenseData"] = ( + _license_check.airgapped_license_data +) global_max_parallel_request_retries_env: Optional[str] = os.getenv( "LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES" ) @@ -1532,9 +1532,9 @@ master_key: Optional[str] = None config_agents: Optional[List[AgentConfig]] = None otel_logging = False prisma_client: Optional[PrismaClient] = None -shared_aiohttp_session: Optional[ - "ClientSession" -] = None # Global shared session for connection reuse +shared_aiohttp_session: Optional["ClientSession"] = ( + None # Global shared session for connection reuse +) user_api_key_cache = DualCache( default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value ) @@ -1545,13 +1545,13 @@ model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter( dual_cache=user_api_key_cache ) litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter) -redis_usage_cache: Optional[ - RedisCache -] = None # redis cache used for tracking spend, tpm/rpm limits +redis_usage_cache: Optional[RedisCache] = ( + None # redis cache used for tracking spend, tpm/rpm limits +) polling_via_cache_enabled: Union[Literal["all"], List[str], bool] = False -native_background_mode: List[ - str -] = [] # Models that should use native provider background mode instead of polling +native_background_mode: List[str] = ( + [] +) # Models that should use native provider background mode instead of polling polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache user_custom_auth = None user_custom_key_generate = None @@ -2036,9 +2036,9 @@ async def update_cache( # noqa: PLR0915 _id = "team_id:{}".format(team_id) try: # Fetch the existing cost for the given user - existing_spend_obj: Optional[ - LiteLLM_TeamTable - ] = await user_api_key_cache.async_get_cache(key=_id) + existing_spend_obj: Optional[LiteLLM_TeamTable] = ( + await user_api_key_cache.async_get_cache(key=_id) + ) if existing_spend_obj is None: # do nothing if team not in api key cache return @@ -2248,23 +2248,54 @@ def _schedule_background_health_check_db_save( ) +def _get_endpoint_exception_status(endpoint: dict, exceptions: dict) -> int: + """Return the HTTP status code for an unhealthy endpoint. + + Prefers the live exception object in `exceptions` (direct health check path). + Falls back to the `exception_status` integer stored on the endpoint dict + (shared-cache path, where exception objects are not available). + """ + model_id = endpoint.get("model_id") + exc = exceptions.get(model_id) if model_id else None + if exc is not None: + return getattr(exc, "status_code", 500) + return endpoint.get("exception_status", 500) + + def _write_health_state_to_router_cache( healthy_endpoints: list, unhealthy_endpoints: list, + exceptions_by_model_id: Optional[dict] = None, ) -> None: """ Write deployment health states to the router's health state cache for health-check-driven routing. No-op if the feature is disabled. """ from litellm.proxy.health_check import build_deployment_health_states + from litellm.router_utils.cooldown_handlers import _set_cooldown_deployments + from litellm.router_utils.router_callbacks.track_deployment_metrics import ( + increment_deployment_failures_for_current_minute, + ) + + _exceptions: dict = exceptions_by_model_id or {} try: if llm_router is None or not llm_router.enable_health_check_routing: return + # When health_check_ignore_transient_errors is set, treat 429/408 + # endpoints as healthy so they are not filtered from routing. + _effective_unhealthy = unhealthy_endpoints + if llm_router.health_check_ignore_transient_errors: + _effective_unhealthy = [ + ep + for ep in unhealthy_endpoints + if _get_endpoint_exception_status(ep, _exceptions) not in (429, 408) + ] + states = build_deployment_health_states( healthy_endpoints=healthy_endpoints, - unhealthy_endpoints=unhealthy_endpoints, + unhealthy_endpoints=_effective_unhealthy, ) if states: llm_router.health_state_cache.set_deployment_health_states(states) @@ -2273,6 +2304,37 @@ def _write_health_state_to_router_cache( sum(1 for s in states.values() if s.get("is_healthy")), sum(1 for s in states.values() if not s.get("is_healthy")), ) + + for endpoint in unhealthy_endpoints: + model_id = endpoint.get("model_id") + if not model_id: + continue + + original_exception = _exceptions.get(model_id) + if original_exception is None: + continue + + exception_status = getattr(original_exception, "status_code", 500) + + if llm_router.health_check_ignore_transient_errors and exception_status in ( + 429, + 408, + ): + continue + + increment_deployment_failures_for_current_minute( + litellm_router_instance=llm_router, + deployment_id=model_id, + ) + + _set_cooldown_deployments( + litellm_router_instance=llm_router, + original_exception=original_exception, + exception_status=exception_status, + deployment=model_id, + time_to_cooldown=llm_router.cooldown_time, + ) + except Exception as e: verbose_proxy_logger.warning( "Failed to write health state to router cache: %s", str(e) @@ -2384,6 +2446,7 @@ async def _run_background_health_check(): ( healthy_endpoints, unhealthy_endpoints, + _exceptions_by_model_id, ) = await shared_health_manager.perform_shared_health_check( model_list=_llm_model_list, details=details_bool, @@ -2397,6 +2460,7 @@ async def _run_background_health_check(): ( healthy_endpoints, unhealthy_endpoints, + _exceptions_by_model_id, ) = await _run_direct_health_check_with_instrumentation( _llm_model_list, health_check_details, @@ -2407,6 +2471,7 @@ async def _run_background_health_check(): ( healthy_endpoints, unhealthy_endpoints, + _exceptions_by_model_id, ) = await _run_direct_health_check_with_instrumentation( _llm_model_list, health_check_details, @@ -2449,7 +2514,9 @@ async def _run_background_health_check(): ) # Write health state to router cache for health-check-driven routing - _write_health_state_to_router_cache(healthy_endpoints, unhealthy_endpoints) + _write_health_state_to_router_cache( + healthy_endpoints, unhealthy_endpoints, _exceptions_by_model_id + ) await asyncio.sleep(health_check_interval) @@ -3245,6 +3312,7 @@ class ProxyConfig: general_settings = {} _enable_hc_routing = False _hc_staleness = None + _hc_ignore_transient = False if general_settings: ### LOAD KEY MANAGEMENT SETTINGS FIRST (needed for custom secret manager) ### key_management_settings = general_settings.get( @@ -3437,6 +3505,9 @@ class ProxyConfig: _hc_staleness = general_settings.get( "health_check_staleness_threshold", None ) + _hc_ignore_transient = general_settings.get( + "health_check_ignore_transient_errors", False + ) verbose_proxy_logger.info( "background_health_check_config enabled=%s shared=%s interval_seconds=%s max_concurrency=%s details=%s health_check_routing=%s", use_background_health_checks, @@ -3479,6 +3550,8 @@ class ProxyConfig: router_params["enable_health_check_routing"] = True if _hc_staleness is not None: router_params["health_check_staleness_threshold"] = _hc_staleness + if _hc_ignore_transient: + router_params["health_check_ignore_transient_errors"] = True ## MODEL LIST model_list = config.get("model_list", None) if model_list: @@ -5217,10 +5290,10 @@ class ProxyConfig: ) try: - guardrails_in_db: List[ - Guardrail - ] = await GuardrailRegistry.get_all_guardrails_from_db( - prisma_client=prisma_client + guardrails_in_db: List[Guardrail] = ( + await GuardrailRegistry.get_all_guardrails_from_db( + prisma_client=prisma_client + ) ) verbose_proxy_logger.debug( "guardrails from the DB %s", str(guardrails_in_db) @@ -5602,9 +5675,9 @@ async def initialize( # noqa: PLR0915 user_api_base = api_base dynamic_config[user_model]["api_base"] = api_base if api_version: - os.environ[ - "AZURE_API_VERSION" - ] = api_version # set this for azure - litellm can read this from the env + os.environ["AZURE_API_VERSION"] = ( + api_version # set this for azure - litellm can read this from the env + ) if max_tokens: # model-specific param dynamic_config[user_model]["max_tokens"] = max_tokens if temperature: # model-specific param @@ -6503,10 +6576,18 @@ class ProxyStartupEvent: KeyRotationManager, ) - # Get prisma_client from global scope + # Get prisma_client and proxy_logging_obj from global scope global prisma_client + global proxy_logging_obj if prisma_client is not None: - key_rotation_manager = KeyRotationManager(prisma_client) + # Reuse the PodLockManager from db_spend_update_writer + pod_lock_manager = ( + proxy_logging_obj.db_spend_update_writer.pod_lock_manager + ) + key_rotation_manager = KeyRotationManager( + prisma_client, + pod_lock_manager=pod_lock_manager, + ) verbose_proxy_logger.debug( f"Key rotation background job scheduled every {LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS} seconds (LITELLM_KEY_ROTATION_ENABLED=true)" ) @@ -12624,9 +12705,9 @@ async def get_config_list( hasattr(sub_field_info, "description") and sub_field_info.description is not None ): - nested_fields[ - idx - ].field_description = sub_field_info.description + nested_fields[idx].field_description = ( + sub_field_info.description + ) idx += 1 _stored_in_db = None diff --git a/litellm/router.py b/litellm/router.py index 1b8f7c91761..a58b3ce25e1 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -310,6 +310,7 @@ class Router: ignore_invalid_deployments: bool = False, enable_health_check_routing: bool = False, health_check_staleness_threshold: Optional[int] = None, + health_check_ignore_transient_errors: bool = False, ) -> None: """ Initialize the Router class with the given parameters for caching, reliability, and routing strategy. @@ -408,9 +409,9 @@ class Router: ) # names of models under litellm_params. ex. azure/chatgpt-v-2 self.deployment_latency_map = {} ### CACHING ### - cache_type: Literal[ - "local", "redis", "redis-semantic", "s3", "disk" - ] = "local" # default to an in-memory cache + cache_type: Literal["local", "redis", "redis-semantic", "s3", "disk"] = ( + "local" # default to an in-memory cache + ) redis_cache = None cache_config: Dict[str, Any] = {} @@ -458,9 +459,9 @@ class Router: self.default_max_parallel_requests = default_max_parallel_requests self.provider_default_deployment_ids: List[str] = [] self.pattern_router = PatternMatchRouter() - self.team_pattern_routers: Dict[ - str, PatternMatchRouter - ] = {} # {"TEAM_ID": PatternMatchRouter} + self.team_pattern_routers: Dict[str, PatternMatchRouter] = ( + {} + ) # {"TEAM_ID": PatternMatchRouter} self.auto_routers: Dict[str, "AutoRouter"] = {} self.complexity_routers: Dict[str, "ComplexityRouter"] = {} @@ -501,6 +502,7 @@ class Router: ) self.disable_cooldowns = disable_cooldowns self.enable_health_check_routing = enable_health_check_routing + self.health_check_ignore_transient_errors = health_check_ignore_transient_errors _staleness = health_check_staleness_threshold or ( DEFAULT_HEALTH_CHECK_INTERVAL * DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER ) @@ -655,12 +657,12 @@ class Router: ) ) - self.model_group_retry_policy: Optional[ - Dict[str, RetryPolicy] - ] = model_group_retry_policy - self.model_group_affinity_config: Optional[ - Dict[str, List[str]] - ] = model_group_affinity_config + self.model_group_retry_policy: Optional[Dict[str, RetryPolicy]] = ( + model_group_retry_policy + ) + self.model_group_affinity_config: Optional[Dict[str, List[str]]] = ( + model_group_affinity_config + ) self.allowed_fails_policy: Optional[AllowedFailsPolicy] = None if allowed_fails_policy is not None: @@ -2066,7 +2068,10 @@ class Router: async def _acompletion( # noqa: PLR0915 self, model: str, messages: List[Dict[str, str]], **kwargs - ) -> Union[ModelResponse, CustomStreamWrapper,]: + ) -> Union[ + ModelResponse, + CustomStreamWrapper, + ]: """ - Get an available deployment - call it with a semaphore over the call @@ -4300,9 +4305,9 @@ class Router: healthy_deployments=healthy_deployments, responses=responses ) returned_response = cast(OpenAIFileObject, responses[0]) - returned_response._hidden_params[ - "model_file_id_mapping" - ] = model_file_id_mapping + returned_response._hidden_params["model_file_id_mapping"] = ( + model_file_id_mapping + ) return returned_response except Exception as e: verbose_router_logger.exception( @@ -5310,11 +5315,16 @@ class Router: e, (litellm.ContextWindowExceededError, litellm.ContentPolicyViolationError), ) - all_deployments = self._get_all_deployments(model_name=original_model_group) + _request_team_id: Optional[str] = ( + kwargs.get("metadata", {}) or {} + ).get("user_api_key_team_id") + all_deployments = self._get_all_deployments( + model_name=original_model_group, team_id=_request_team_id + ) _order_set: set = { - d.get("litellm_params", {}).get("order") + litellm.utils._get_deployment_order(d) for d in all_deployments - if d.get("litellm_params", {}).get("order") is not None + if litellm.utils._get_deployment_order(d) is not None } order_values: list = sorted(_order_set) if len(order_values) > 1 and not _skip_order_fallback: @@ -5387,11 +5397,11 @@ class Router: if isinstance(e, litellm.ContextWindowExceededError): if context_window_fallbacks is not None: - context_window_fallback_model_group: Optional[ - List[str] - ] = self._get_fallback_model_group_from_fallbacks( - fallbacks=context_window_fallbacks, - model_group=model_group, + context_window_fallback_model_group: Optional[List[str]] = ( + self._get_fallback_model_group_from_fallbacks( + fallbacks=context_window_fallbacks, + model_group=model_group, + ) ) if context_window_fallback_model_group is None: raise original_exception @@ -5423,11 +5433,11 @@ class Router: e.message += "\n{}".format(error_message) elif isinstance(e, litellm.ContentPolicyViolationError): if content_policy_fallbacks is not None: - content_policy_fallback_model_group: Optional[ - List[str] - ] = self._get_fallback_model_group_from_fallbacks( - fallbacks=content_policy_fallbacks, - model_group=model_group, + content_policy_fallback_model_group: Optional[List[str]] = ( + self._get_fallback_model_group_from_fallbacks( + fallbacks=content_policy_fallbacks, + model_group=model_group, + ) ) if content_policy_fallback_model_group is None: raise original_exception @@ -5649,9 +5659,9 @@ class Router: ) ## ADD RETRY TRACKING TO METADATA - used for spend logs retry tracking _metadata["attempted_retries"] = 0 - _metadata[ - "max_retries" - ] = num_retries # Updated after overrides in exception handler + _metadata["max_retries"] = ( + num_retries # Updated after overrides in exception handler + ) try: self._handle_mock_testing_rate_limit_error( model_group=model_group, kwargs=kwargs @@ -6770,26 +6780,26 @@ class Router: """ from litellm.router_strategy.auto_router.auto_router import AutoRouter - auto_router_config_path: Optional[ - str - ] = deployment.litellm_params.auto_router_config_path + auto_router_config_path: Optional[str] = ( + deployment.litellm_params.auto_router_config_path + ) auto_router_config: Optional[str] = deployment.litellm_params.auto_router_config if auto_router_config_path is None and auto_router_config is None: raise ValueError( "auto_router_config_path or auto_router_config is required for auto-router deployments. Please set it in the litellm_params" ) - default_model: Optional[ - str - ] = deployment.litellm_params.auto_router_default_model + default_model: Optional[str] = ( + deployment.litellm_params.auto_router_default_model + ) if default_model is None: raise ValueError( "auto_router_default_model is required for auto-router deployments. Please set it in the litellm_params" ) - embedding_model: Optional[ - str - ] = deployment.litellm_params.auto_router_embedding_model + embedding_model: Optional[str] = ( + deployment.litellm_params.auto_router_embedding_model + ) if embedding_model is None: raise ValueError( "auto_router_embedding_model is required for auto-router deployments. Please set it in the litellm_params" @@ -6832,13 +6842,13 @@ class Router: ComplexityRouter, ) - complexity_router_config: Optional[ - dict - ] = deployment.litellm_params.complexity_router_config + complexity_router_config: Optional[dict] = ( + deployment.litellm_params.complexity_router_config + ) - default_model: Optional[ - str - ] = deployment.litellm_params.complexity_router_default_model + default_model: Optional[str] = ( + deployment.litellm_params.complexity_router_default_model + ) # If no default model specified, try to get from config tiers if default_model is None and complexity_router_config: @@ -7497,9 +7507,9 @@ class Router: # Add custom_llm_provider if deployment.litellm_params.custom_llm_provider: - credentials[ - "custom_llm_provider" - ] = deployment.litellm_params.custom_llm_provider + credentials["custom_llm_provider"] = ( + deployment.litellm_params.custom_llm_provider + ) elif "/" in deployment.litellm_params.model: # Extract provider from "provider/model" format credentials["custom_llm_provider"] = deployment.litellm_params.model.split( @@ -9070,7 +9080,9 @@ class Router: ## get healthy deployments ### get all deployments - healthy_deployments = self._get_all_deployments(model_name=model) + healthy_deployments = self._get_all_deployments( + model_name=model, team_id=request_team_id + ) if len(healthy_deployments) == 0: # check if the user sent in a deployment name instead @@ -9091,7 +9103,9 @@ class Router: ) # Re-assign model to the fallback and try to get deployments again model = fallback_model - healthy_deployments = self._get_all_deployments(model_name=model) + healthy_deployments = self._get_all_deployments( + model_name=model, team_id=request_team_id + ) # If still no deployments after checking for fallbacks, raise an error if len(healthy_deployments) == 0: @@ -9181,10 +9195,23 @@ class Router: ) if verbose_router_logger.isEnabledFor(logging.DEBUG): verbose_router_logger.debug(f"cooldown deployments: {cooldown_deployments}") + _pre_cooldown_deployments = healthy_deployments healthy_deployments = self._filter_cooldown_deployments( healthy_deployments=healthy_deployments, cooldown_deployments=cooldown_deployments, ) + # Safety net: only bypass cooldown filter when health-check routing is + # driving cooldown (i.e. allowed_fails_policy is set). Without a policy, + # cooldowns are from real request failures and must not be bypassed. + if ( + not healthy_deployments + and self.enable_health_check_routing + and self.allowed_fails_policy is not None + ): + verbose_router_logger.warning( + "All deployments in cooldown via health-check routing, bypassing cooldown filter" + ) + healthy_deployments = _pre_cooldown_deployments healthy_deployments = await self.async_callback_filter_deployments( model=model, @@ -9617,10 +9644,20 @@ class Router: cooldown_deployments = _get_cooldown_deployments( litellm_router_instance=self, parent_otel_span=parent_otel_span ) + _pre_cooldown_deployments = healthy_deployments healthy_deployments = self._filter_cooldown_deployments( healthy_deployments=healthy_deployments, cooldown_deployments=cooldown_deployments, ) + if ( + not healthy_deployments + and self.enable_health_check_routing + and self.allowed_fails_policy is not None + ): + verbose_router_logger.warning( + "All deployments in cooldown via health-check routing, bypassing cooldown filter" + ) + healthy_deployments = _pre_cooldown_deployments # filter pre-call checks if self.enable_pre_call_checks and messages is not None: @@ -9922,6 +9959,12 @@ class Router: if not self.enable_health_check_routing: return healthy_deployments + # When allowed_fails_policy is set, cooldown is the sole routing exclusion + # mechanism -- skip the binary health check filter so the policy threshold + # is respected before any deployment is excluded. + if self.allowed_fails_policy is not None: + return healthy_deployments + unhealthy_ids = ( await self.health_state_cache.async_get_unhealthy_deployment_ids( parent_otel_span=parent_otel_span @@ -9951,6 +9994,9 @@ class Router: if not self.enable_health_check_routing: return healthy_deployments + if self.allowed_fails_policy is not None: + return healthy_deployments + unhealthy_ids = self.health_state_cache.get_unhealthy_deployment_ids( parent_otel_span=parent_otel_span ) diff --git a/litellm/types/proxy/policy_engine/pipeline_types.py b/litellm/types/proxy/policy_engine/pipeline_types.py index 29d2e576000..abbb127cd7a 100644 --- a/litellm/types/proxy/policy_engine/pipeline_types.py +++ b/litellm/types/proxy/policy_engine/pipeline_types.py @@ -18,18 +18,24 @@ class PipelineStep(BaseModel): """ A single step in a guardrail pipeline. - Each step runs a guardrail and takes an action based on pass/fail. + Each step runs a guardrail and takes an action based on pass, policy fail, + or technical/API error (see pipeline executor outcome types). """ guardrail: str = Field(description="Name of the guardrail to run.") on_fail: str = Field( default="block", - description="Action when guardrail rejects: next | block | allow | modify_response", + description="Action when guardrail rejects content (policy intervention): next | block | allow | modify_response", ) on_pass: str = Field( default="allow", description="Action when guardrail passes: next | block | allow | modify_response", ) + on_error: Optional[str] = Field( + default=None, + description="Action when the guardrail raises a technical error (timeouts, " + "unreachable provider, non-intervention HTTP errors). If omitted, uses on_fail.", + ) pass_data: bool = Field( default=False, description="Forward modified request data (e.g., PII-masked) to next step.", @@ -41,9 +47,11 @@ class PipelineStep(BaseModel): model_config = ConfigDict(extra="forbid") - @field_validator("on_fail", "on_pass") + @field_validator("on_fail", "on_pass", "on_error") @classmethod - def validate_action(cls, v: str) -> str: + def validate_action(cls, v: Optional[str]) -> Optional[str]: + if v is None: + return None if v not in VALID_PIPELINE_ACTIONS: raise ValueError( f"Invalid action '{v}'. Must be one of: {sorted(VALID_PIPELINE_ACTIONS)}" diff --git a/litellm/utils.py b/litellm/utils.py index 6806961bf51..f902644e760 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4875,6 +4875,19 @@ def calculate_max_parallel_requests( return None +def _get_deployment_order(deployment: Union[Dict, Any]) -> Optional[int]: + """ + Returns the routing order for a deployment. + + Checks litellm_params first (static config), then model_info (dynamic/team + models added via API where order lives in model_info, not litellm_params). + """ + order = deployment.get("litellm_params", {}).get("order") + if order is None: + order = deployment.get("model_info", {}).get("order") + return order + + def _get_order_filtered_deployments( healthy_deployments: List[Dict], target_order: Optional[int] = None ) -> List: @@ -4882,7 +4895,7 @@ def _get_order_filtered_deployments( filtered = [ d for d in healthy_deployments - if d["litellm_params"].get("order") == target_order + if _get_deployment_order(d) == target_order ] if filtered: return filtered @@ -4890,20 +4903,19 @@ def _get_order_filtered_deployments( return healthy_deployments # Default: pick min order group - min_order = min( - ( - deployment["litellm_params"]["order"] - for deployment in healthy_deployments - if "order" in deployment["litellm_params"] - ), - default=None, - ) + _valid_orders: List[int] = [ + o + for deployment in healthy_deployments + for o in [_get_deployment_order(deployment)] + if o is not None + ] + min_order: Optional[int] = min(_valid_orders) if _valid_orders else None if min_order is not None: filtered_deployments = [ deployment for deployment in healthy_deployments - if deployment["litellm_params"].get("order") == min_order + if _get_deployment_order(deployment) == min_order ] return filtered_deployments diff --git a/tests/litellm_utils_tests/test_health_check.py b/tests/litellm_utils_tests/test_health_check.py index 0b0e091211f..708b2403c49 100644 --- a/tests/litellm_utils_tests/test_health_check.py +++ b/tests/litellm_utils_tests/test_health_check.py @@ -475,13 +475,13 @@ async def test_perform_health_check_filters_by_model_id(): captured_list.append(m_list) return [ {"model": "gpt-4", "api_key": m_list[0]["litellm_params"]["api_key"]} - ], [] + ], [], {} with patch( "litellm.proxy.health_check._perform_health_check", side_effect=mock_perform_health_check, ): - healthy_endpoints, unhealthy_endpoints = await perform_health_check( + healthy_endpoints, unhealthy_endpoints, _ = await perform_health_check( model_list=model_list, model_id="deployment-id-2", details=True ) @@ -521,7 +521,7 @@ async def test_perform_health_check_with_health_check_model(): return {"status": "healthy"} with patch("litellm.ahealth_check", side_effect=mock_health_check): - healthy_endpoints, unhealthy_endpoints = await _perform_health_check(model_list) + healthy_endpoints, unhealthy_endpoints, _ = await _perform_health_check(model_list) print("health check calls: ", health_check_calls) # Verify the health check used the override model @@ -556,7 +556,7 @@ async def test_health_check_bad_model(): }, ] details = None - healthy_endpoints, unhealthy_endpoints = await _perform_health_check( + healthy_endpoints, unhealthy_endpoints, _ = await _perform_health_check( model_list, details ) print(f"healthy_endpoints: {healthy_endpoints}") @@ -574,7 +574,7 @@ async def test_health_check_bad_model(): "litellm.ahealth_check", side_effect=mock_health_check ) as mock_health_check: start_time = time.time() - healthy_endpoints, unhealthy_endpoints = await _perform_health_check(model_list) + healthy_endpoints, unhealthy_endpoints, _ = await _perform_health_check(model_list) end_time = time.time() print("health check calls: ", health_check_calls) assert len(healthy_endpoints) == 0 @@ -667,7 +667,7 @@ async def test_timeout_does_not_cancel_other_health_checks(): return {"status": "healthy"} with patch("litellm.ahealth_check", side_effect=mock_health_check): - healthy_endpoints, unhealthy_endpoints = await _perform_health_check( + healthy_endpoints, unhealthy_endpoints, _ = await _perform_health_check( model_list, max_concurrency=1 ) diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 047193055d8..7da4d41fbf1 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -2420,7 +2420,7 @@ async def test_run_background_health_check_reflects_llm_model_list(monkeypatch): async def fake_perform_health_check(model_list, details, max_concurrency=None): called_model_lists.append(copy.deepcopy(model_list)) - return (["healthy"], ["unhealthy"]) + return (["healthy"], ["unhealthy"], {}) monkeypatch.setattr(proxy_server, "health_check_interval", 1) monkeypatch.setattr(proxy_server, "health_check_details", None) @@ -2471,7 +2471,7 @@ async def test_background_health_check_skip_disabled_models(monkeypatch): async def fake_perform_health_check(model_list, details, max_concurrency=None): called_model_lists.append(copy.deepcopy(model_list)) - return (["healthy"], []) + return (["healthy"], [], {}) monkeypatch.setattr(proxy_server, "health_check_interval", 1) monkeypatch.setattr(proxy_server, "health_check_details", None) diff --git a/tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py b/tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py new file mode 100644 index 00000000000..f6ef02a86de --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_key_rotation_e2e.py @@ -0,0 +1,559 @@ +""" +End-to-end tests for key rotation feature. + +Covers the critical gaps: +1. Multi-pod simulation: two KeyRotationManagers sharing one PodLockManager +2. Error resilience: partial failures, regenerate_key_fn failures, hook failures +3. Full process_rotations flow with actual key finding + rotation + lock +4. Initialization wiring: PodLockManager is correctly passed +5. Multiple keys: some succeed, some fail, all are attempted +6. Rotation count increments correctly over multiple rotations +""" + +import os +import sys +from datetime import datetime, timedelta, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy._types import ( + GenerateKeyResponse, + LiteLLM_VerificationToken, +) +from litellm.proxy.common_utils.key_rotation_manager import KeyRotationManager + + +class TestMultiPodKeyRotation: + """ + Simulate two pods sharing one Redis lock to verify only one pod + runs key rotation at a time. + """ + + @pytest.mark.asyncio + async def test_two_pods_only_one_rotates(self): + """ + Two KeyRotationManagers with separate pod_lock_managers but + the same Redis backend. Only the first to acquire the lock + should rotate; the second should skip. + """ + mock_prisma = AsyncMock() + + # Shared state to simulate Redis SET NX behavior + redis_lock = {"holder": None} + + async def make_acquire_lock(pod_id): + async def acquire(cronjob_id, **kwargs): + if redis_lock["holder"] is None: + redis_lock["holder"] = pod_id + return True + return redis_lock["holder"] == pod_id + + return acquire + + async def make_release_lock(pod_id): + async def release(cronjob_id): + if redis_lock["holder"] == pod_id: + redis_lock["holder"] = None + + return release + + # Pod A + pod_a_lock_mgr = MagicMock() + pod_a_lock_mgr.redis_cache = MagicMock() + pod_a_lock_mgr.acquire_lock = AsyncMock( + side_effect=await make_acquire_lock("pod-a") + ) + pod_a_lock_mgr.release_lock = AsyncMock( + side_effect=await make_release_lock("pod-a") + ) + + # Pod B + pod_b_lock_mgr = MagicMock() + pod_b_lock_mgr.redis_cache = MagicMock() + pod_b_lock_mgr.acquire_lock = AsyncMock( + side_effect=await make_acquire_lock("pod-b") + ) + pod_b_lock_mgr.release_lock = AsyncMock( + side_effect=await make_release_lock("pod-b") + ) + + manager_a = KeyRotationManager(mock_prisma, pod_lock_manager=pod_a_lock_mgr) + manager_b = KeyRotationManager(mock_prisma, pod_lock_manager=pod_b_lock_mgr) + + # Both share the same mock methods for rotation logic + for mgr in [manager_a, manager_b]: + mgr._cleanup_expired_deprecated_keys = AsyncMock() + mgr._find_keys_needing_rotation = AsyncMock(return_value=[]) + + # Pod A acquires lock first + await manager_a.process_rotations() + # Pod A should have run rotation + manager_a._cleanup_expired_deprecated_keys.assert_called_once() + manager_a._find_keys_needing_rotation.assert_called_once() + + # Lock is released after pod A finishes, so pod B can now acquire + # But let's simulate pod B trying WHILE pod A holds the lock + # Reset the lock state to simulate concurrent access + redis_lock["holder"] = "pod-a" # Pod A holds the lock + + await manager_b.process_rotations() + # Pod B should NOT have run rotation (lock held by pod-a) + manager_b._cleanup_expired_deprecated_keys.assert_not_called() + manager_b._find_keys_needing_rotation.assert_not_called() + + @pytest.mark.asyncio + async def test_second_pod_runs_after_first_releases(self): + """ + After the first pod releases the lock, the second pod should + be able to acquire and run rotation. + """ + mock_prisma = AsyncMock() + + call_order = [] + + # Pod A - always gets the lock + pod_a_lock = MagicMock() + pod_a_lock.redis_cache = MagicMock() + pod_a_lock.acquire_lock = AsyncMock(return_value=True) + pod_a_lock.release_lock = AsyncMock() + + # Pod B - also gets the lock (simulating after A releases) + pod_b_lock = MagicMock() + pod_b_lock.redis_cache = MagicMock() + pod_b_lock.acquire_lock = AsyncMock(return_value=True) + pod_b_lock.release_lock = AsyncMock() + + manager_a = KeyRotationManager(mock_prisma, pod_lock_manager=pod_a_lock) + manager_b = KeyRotationManager(mock_prisma, pod_lock_manager=pod_b_lock) + + async def cleanup_a(): + call_order.append("a_cleanup") + + async def cleanup_b(): + call_order.append("b_cleanup") + + manager_a._cleanup_expired_deprecated_keys = AsyncMock(side_effect=cleanup_a) + manager_a._find_keys_needing_rotation = AsyncMock(return_value=[]) + manager_b._cleanup_expired_deprecated_keys = AsyncMock(side_effect=cleanup_b) + manager_b._find_keys_needing_rotation = AsyncMock(return_value=[]) + + # Run sequentially: A then B + await manager_a.process_rotations() + await manager_b.process_rotations() + + # Both should have run + assert call_order == ["a_cleanup", "b_cleanup"] + pod_a_lock.release_lock.assert_called_once() + pod_b_lock.release_lock.assert_called_once() + + +class TestKeyRotationErrorResilience: + """ + Tests that key rotation handles errors gracefully: + - regenerate_key_fn failure for one key doesn't block others + - Hook failure doesn't crash the process + - Database update failure is handled + """ + + @pytest.mark.asyncio + async def test_one_key_fails_others_still_rotate(self): + """ + If rotation fails for one key, the remaining keys should still + be attempted. No key should be silently skipped. + """ + mock_prisma = AsyncMock() + manager = KeyRotationManager(mock_prisma) + + key1 = LiteLLM_VerificationToken( + token="token-1", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + key_name="key-1", + ) + key2 = LiteLLM_VerificationToken( + token="token-2", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + key_name="key-2", + ) + key3 = LiteLLM_VerificationToken( + token="token-3", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + key_name="key-3", + ) + + manager._cleanup_expired_deprecated_keys = AsyncMock() + manager._find_keys_needing_rotation = AsyncMock(return_value=[key1, key2, key3]) + + rotate_calls = [] + + async def mock_rotate(key): + rotate_calls.append(key.token) + if key.token == "token-2": + raise Exception("Database connection lost") + + manager._rotate_key = AsyncMock(side_effect=mock_rotate) + + await manager.process_rotations() + + # All 3 keys should have been attempted + assert rotate_calls == ["token-1", "token-2", "token-3"] + + @pytest.mark.asyncio + async def test_regenerate_key_fn_failure_is_caught(self): + """ + If regenerate_key_fn throws, _rotate_key should propagate the error + but process_rotations should catch it per-key. + """ + mock_prisma = AsyncMock() + manager = KeyRotationManager(mock_prisma) + + key = LiteLLM_VerificationToken( + token="test-token", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + key_name="test-key", + ) + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + new_callable=AsyncMock, + side_effect=Exception("regenerate failed: DB timeout"), + ): + # _rotate_key should raise + with pytest.raises(Exception, match="regenerate failed"): + await manager._rotate_key(key) + + # But process_rotations should catch per-key errors + manager._cleanup_expired_deprecated_keys = AsyncMock() + manager._find_keys_needing_rotation = AsyncMock(return_value=[key]) + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + new_callable=AsyncMock, + side_effect=Exception("regenerate failed: DB timeout"), + ): + # Should NOT raise - error is caught per-key + await manager.process_rotations() + + @pytest.mark.asyncio + async def test_hook_failure_does_not_prevent_db_update(self): + """ + If the rotation hook (async_key_rotated_hook) fails, the database + update for rotation_count should still have succeeded (it runs before the hook). + """ + mock_prisma = AsyncMock() + manager = KeyRotationManager(mock_prisma) + + key = LiteLLM_VerificationToken( + token="test-token", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + ) + + mock_response = GenerateKeyResponse( + key="new-key", token_id="new-token-id", user_id="test-user" + ) + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + new_callable=AsyncMock, + return_value=mock_response, + ): + with patch( + "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + side_effect=Exception("Hook failed: secret manager down"), + ): + # This will raise because the hook fails + with pytest.raises(Exception, match="Hook failed"): + await manager._rotate_key(key) + + # The DB update should have been called BEFORE the hook + mock_prisma.db.litellm_verificationtoken.update.assert_called_once() + update_data = mock_prisma.db.litellm_verificationtoken.update.call_args[1][ + "data" + ] + assert update_data["rotation_count"] == 1 + + @pytest.mark.asyncio + async def test_cleanup_failure_does_not_prevent_rotation(self): + """ + If deprecated key cleanup fails, the rotation should still proceed. + """ + mock_prisma = AsyncMock() + mock_pod_lock = MagicMock() + mock_pod_lock.redis_cache = MagicMock() + mock_pod_lock.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock.release_lock = AsyncMock() + + manager = KeyRotationManager(mock_prisma, pod_lock_manager=mock_pod_lock) + + # Cleanup fails + manager._cleanup_expired_deprecated_keys = AsyncMock( + side_effect=Exception("Deprecated table doesn't exist") + ) + + # process_rotations catches the exception internally (try/except), + # but the lock must still be released in the finally block. + await manager.process_rotations() + + # Lock should still be released in finally block + mock_pod_lock.release_lock.assert_called_once() + + +class TestKeyRotationFullFlow: + """ + Full end-to-end flow tests: find keys -> rotate -> update DB -> release lock + """ + + @pytest.mark.asyncio + async def test_full_rotation_flow_with_lock(self): + """ + Full flow: acquire lock -> cleanup -> find keys -> rotate -> update DB -> release lock + """ + mock_prisma = AsyncMock() + + # Setup lock manager + mock_lock = MagicMock() + mock_lock.redis_cache = MagicMock() + mock_lock.acquire_lock = AsyncMock(return_value=True) + mock_lock.release_lock = AsyncMock() + + manager = KeyRotationManager(mock_prisma, pod_lock_manager=mock_lock) + + key = LiteLLM_VerificationToken( + token="old-token-hash", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=datetime.now(timezone.utc) - timedelta(seconds=60), + rotation_count=2, + key_name="my-key", + key_alias="prod/my-key", + ) + + mock_response = GenerateKeyResponse( + key="sk-new-key-value", + token_id="new-token-hash", + user_id="system", + ) + + # Mock cleanup + mock_prisma.db.litellm_deprecatedverificationtoken.delete_many.return_value = 1 + # Mock find keys + mock_prisma.db.litellm_verificationtoken.find_many.return_value = [key] + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + new_callable=AsyncMock, + return_value=mock_response, + ): + with patch( + "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ): + await manager.process_rotations() + + # Verify full flow executed: + # 1. Lock acquired + mock_lock.acquire_lock.assert_called_once() + + # 2. Cleanup ran + mock_prisma.db.litellm_deprecatedverificationtoken.delete_many.assert_called_once() + + # 3. Keys were queried + mock_prisma.db.litellm_verificationtoken.find_many.assert_called_once() + + # 4. DB was updated with new rotation info + mock_prisma.db.litellm_verificationtoken.update.assert_called_once() + update_args = mock_prisma.db.litellm_verificationtoken.update.call_args[1] + assert update_args["where"]["token"] == "new-token-hash" + assert update_args["data"]["rotation_count"] == 3 # was 2, now 3 + + # 5. Lock released + mock_lock.release_lock.assert_called_once() + + @pytest.mark.asyncio + async def test_rotation_count_increments_across_multiple_rotations(self): + """ + Simulate 3 consecutive rotations and verify rotation_count increments + correctly each time: 0 -> 1 -> 2 -> 3 + """ + mock_prisma = AsyncMock() + manager = KeyRotationManager(mock_prisma) + + rotation_counts_seen = [] + + for expected_count in range(3): + key = LiteLLM_VerificationToken( + token=f"token-v{expected_count}", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=expected_count, + ) + + mock_response = GenerateKeyResponse( + key=f"sk-new-v{expected_count + 1}", + token_id=f"token-v{expected_count + 1}", + user_id="system", + ) + + mock_prisma.db.litellm_verificationtoken.update.reset_mock() + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + new_callable=AsyncMock, + return_value=mock_response, + ): + with patch( + "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ): + await manager._rotate_key(key) + + update_data = mock_prisma.db.litellm_verificationtoken.update.call_args[1][ + "data" + ] + rotation_counts_seen.append(update_data["rotation_count"]) + + assert rotation_counts_seen == [1, 2, 3] + + @pytest.mark.asyncio + async def test_no_keys_to_rotate_skips_gracefully(self): + """ + When no keys need rotation, process should complete without errors. + """ + mock_prisma = AsyncMock() + mock_prisma.db.litellm_deprecatedverificationtoken.delete_many.return_value = 0 + mock_prisma.db.litellm_verificationtoken.find_many.return_value = [] + + mock_lock = MagicMock() + mock_lock.redis_cache = MagicMock() + mock_lock.acquire_lock = AsyncMock(return_value=True) + mock_lock.release_lock = AsyncMock() + + manager = KeyRotationManager(mock_prisma, pod_lock_manager=mock_lock) + + await manager.process_rotations() + + # Verify no rotation was attempted + mock_prisma.db.litellm_verificationtoken.update.assert_not_called() + # But lock was still properly released + mock_lock.release_lock.assert_called_once() + + @pytest.mark.asyncio + async def test_regenerate_response_missing_token_id_skips_db_update(self): + """ + If regenerate_key_fn returns a response without token_id, + the DB update for rotation metadata should be skipped. + """ + mock_prisma = AsyncMock() + manager = KeyRotationManager(mock_prisma) + + key = LiteLLM_VerificationToken( + token="old-token", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + ) + + # Response with no token_id + mock_response = GenerateKeyResponse( + key="sk-new", + token_id=None, + user_id="system", + ) + + with patch( + "litellm.proxy.common_utils.key_rotation_manager.regenerate_key_fn", + new_callable=AsyncMock, + return_value=mock_response, + ): + with patch( + "litellm.proxy.common_utils.key_rotation_manager.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ): + await manager._rotate_key(key) + + # DB update should NOT have been called (no token_id) + mock_prisma.db.litellm_verificationtoken.update.assert_not_called() + + +class TestKeyRotationInitialization: + """ + Tests that the PodLockManager wiring in proxy_server.py is correct. + """ + + @pytest.mark.asyncio + async def test_key_rotation_manager_receives_pod_lock_manager(self): + """ + Verify KeyRotationManager stores the pod_lock_manager correctly. + """ + mock_prisma = AsyncMock() + mock_lock = MagicMock() + mock_lock.redis_cache = MagicMock() + + manager = KeyRotationManager(mock_prisma, pod_lock_manager=mock_lock) + + assert manager.pod_lock_manager is mock_lock + assert manager.prisma_client is mock_prisma + + @pytest.mark.asyncio + async def test_key_rotation_manager_default_no_lock(self): + """ + When no pod_lock_manager is provided, it defaults to None. + """ + mock_prisma = AsyncMock() + manager = KeyRotationManager(mock_prisma) + + assert manager.pod_lock_manager is None + + @pytest.mark.asyncio + async def test_lock_pattern_matches_spend_log_cleanup(self): + """ + Verify the key rotation lock pattern is identical to spend_log_cleanup: + - acquire_lock with cronjob_id + - release_lock in finally + - lock_acquired flag guards release + """ + mock_prisma = AsyncMock() + mock_lock = MagicMock() + mock_lock.redis_cache = MagicMock() + mock_lock.acquire_lock = AsyncMock(return_value=True) + mock_lock.release_lock = AsyncMock() + + manager = KeyRotationManager(mock_prisma, pod_lock_manager=mock_lock) + manager._cleanup_expired_deprecated_keys = AsyncMock() + manager._find_keys_needing_rotation = AsyncMock(return_value=[]) + + await manager.process_rotations() + + # Pattern check: acquire with cronjob_id + acquire_call = mock_lock.acquire_lock.call_args + assert "cronjob_id" in acquire_call.kwargs or len(acquire_call.args) > 0 + + # Pattern check: release with same cronjob_id + release_call = mock_lock.release_lock.call_args + assert "cronjob_id" in release_call.kwargs or len(release_call.args) > 0 + + # Both should use the same job name + from litellm.constants import KEY_ROTATION_JOB_NAME + + assert acquire_call.kwargs.get("cronjob_id") == KEY_ROTATION_JOB_NAME + assert release_call.kwargs.get("cronjob_id") == KEY_ROTATION_JOB_NAME diff --git a/tests/test_litellm/proxy/common_utils/test_key_rotation_lock.py b/tests/test_litellm/proxy/common_utils/test_key_rotation_lock.py new file mode 100644 index 00000000000..c0b3611b2b4 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_key_rotation_lock.py @@ -0,0 +1,229 @@ +""" +Test distributed lock behavior for key rotation manager. + +Verifies that PodLockManager is correctly used to prevent concurrent +key rotation across multiple pods in a distributed deployment. +""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy._types import LiteLLM_VerificationToken +from litellm.proxy.common_utils.key_rotation_manager import KeyRotationManager + + +class TestKeyRotationLock: + """Test distributed lock behavior in KeyRotationManager.""" + + @pytest.mark.asyncio + async def test_process_rotations_acquires_lock(self): + """ + When PodLockManager is provided and lock is acquired, + rotation logic should run normally. + """ + mock_prisma_client = AsyncMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = MagicMock() # Redis is available + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + + manager = KeyRotationManager( + mock_prisma_client, pod_lock_manager=mock_pod_lock_manager + ) + + # Mock _find_keys_needing_rotation to return empty list (no keys to rotate) + manager._find_keys_needing_rotation = AsyncMock(return_value=[]) + manager._cleanup_expired_deprecated_keys = AsyncMock() + + await manager.process_rotations() + + # Verify lock was acquired with custom TTL + mock_pod_lock_manager.acquire_lock.assert_called_once() + call_kwargs = mock_pod_lock_manager.acquire_lock.call_args + assert call_kwargs.kwargs["cronjob_id"] == "litellm_key_rotation_job" + assert call_kwargs.kwargs["ttl"] >= 300 # At least 5 minutes + + # Verify rotation logic ran (cleanup + find keys called) + manager._cleanup_expired_deprecated_keys.assert_called_once() + manager._find_keys_needing_rotation.assert_called_once() + + # Verify lock was released + mock_pod_lock_manager.release_lock.assert_called_once_with( + cronjob_id="litellm_key_rotation_job", + ) + + @pytest.mark.asyncio + async def test_process_rotations_skips_when_lock_held(self): + """ + When lock is held by another pod, process_rotations() should + return early without performing any rotation. + """ + mock_prisma_client = AsyncMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = MagicMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=False) + mock_pod_lock_manager.release_lock = AsyncMock() + + manager = KeyRotationManager( + mock_prisma_client, pod_lock_manager=mock_pod_lock_manager + ) + + manager._find_keys_needing_rotation = AsyncMock() + manager._cleanup_expired_deprecated_keys = AsyncMock() + + await manager.process_rotations() + + # Verify lock was attempted + mock_pod_lock_manager.acquire_lock.assert_called_once() + + # Verify rotation logic was NOT executed + manager._cleanup_expired_deprecated_keys.assert_not_called() + manager._find_keys_needing_rotation.assert_not_called() + + # Verify lock was NOT released (since it was never acquired) + mock_pod_lock_manager.release_lock.assert_not_called() + + @pytest.mark.asyncio + async def test_process_rotations_releases_lock_on_success(self): + """ + Lock should be released in the finally block after successful rotation. + """ + mock_prisma_client = AsyncMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = MagicMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + + manager = KeyRotationManager( + mock_prisma_client, pod_lock_manager=mock_pod_lock_manager + ) + + # Simulate finding and rotating a key successfully + mock_key = LiteLLM_VerificationToken( + token="test-token", + auto_rotate=True, + rotation_interval="30s", + key_rotation_at=None, + rotation_count=0, + key_name="test-key", + ) + manager._find_keys_needing_rotation = AsyncMock(return_value=[mock_key]) + manager._cleanup_expired_deprecated_keys = AsyncMock() + manager._rotate_key = AsyncMock() + + await manager.process_rotations() + + # Verify rotation was performed + manager._rotate_key.assert_called_once_with(mock_key) + + # Verify lock was released after success + mock_pod_lock_manager.release_lock.assert_called_once_with( + cronjob_id="litellm_key_rotation_job", + ) + + @pytest.mark.asyncio + async def test_process_rotations_releases_lock_on_error(self): + """ + Lock should be released in the finally block even if rotation + throws an exception. + """ + mock_prisma_client = AsyncMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = MagicMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + + manager = KeyRotationManager( + mock_prisma_client, pod_lock_manager=mock_pod_lock_manager + ) + + # Simulate an error during cleanup + manager._cleanup_expired_deprecated_keys = AsyncMock( + side_effect=Exception("Database connection failed") + ) + + await manager.process_rotations() + + # Verify lock was still released despite the error + mock_pod_lock_manager.release_lock.assert_called_once_with( + cronjob_id="litellm_key_rotation_job", + ) + + @pytest.mark.asyncio + async def test_process_rotations_works_without_lock_manager(self): + """ + When pod_lock_manager=None, rotation should run normally + without any lock logic (backward compat / single-pod mode). + """ + mock_prisma_client = AsyncMock() + + # No pod_lock_manager provided (default None) + manager = KeyRotationManager(mock_prisma_client) + + manager._find_keys_needing_rotation = AsyncMock(return_value=[]) + manager._cleanup_expired_deprecated_keys = AsyncMock() + + await manager.process_rotations() + + # Verify rotation logic ran normally + manager._cleanup_expired_deprecated_keys.assert_called_once() + manager._find_keys_needing_rotation.assert_called_once() + + @pytest.mark.asyncio + async def test_process_rotations_works_without_redis_cache(self): + """ + When pod_lock_manager exists but redis_cache is None (no Redis configured), + rotation should run normally without locking. + """ + mock_prisma_client = AsyncMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = None # No Redis available + + manager = KeyRotationManager( + mock_prisma_client, pod_lock_manager=mock_pod_lock_manager + ) + + manager._find_keys_needing_rotation = AsyncMock(return_value=[]) + manager._cleanup_expired_deprecated_keys = AsyncMock() + + await manager.process_rotations() + + # Verify lock was NOT attempted (no Redis) + mock_pod_lock_manager.acquire_lock.assert_not_called() + + # Verify rotation logic still ran + manager._cleanup_expired_deprecated_keys.assert_called_once() + manager._find_keys_needing_rotation.assert_called_once() + + @pytest.mark.asyncio + async def test_process_rotations_handles_none_lock_result(self): + """ + When acquire_lock returns None (edge case), it should be treated + as lock NOT acquired, and rotation should be skipped. + """ + mock_prisma_client = AsyncMock() + mock_pod_lock_manager = MagicMock() + mock_pod_lock_manager.redis_cache = MagicMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=None) + mock_pod_lock_manager.release_lock = AsyncMock() + + manager = KeyRotationManager( + mock_prisma_client, pod_lock_manager=mock_pod_lock_manager + ) + + manager._find_keys_needing_rotation = AsyncMock() + manager._cleanup_expired_deprecated_keys = AsyncMock() + + await manager.process_rotations() + + # Verify rotation logic was NOT executed (None treated as False via `or False`) + manager._cleanup_expired_deprecated_keys.assert_not_called() + manager._find_keys_needing_rotation.assert_not_called() + + # Verify lock was NOT released (lock_acquired is False) + mock_pod_lock_manager.release_lock.assert_not_called() diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 226e88bea3e..ffe8947fc61 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -46,6 +46,28 @@ class AlwaysFailGuardrail(CustomGuardrail): raise HTTPException(status_code=400, detail="Content policy violation") +class HttpStatusGuardrail(CustomGuardrail): + """Raises HTTPException with a configurable status (e.g. 503 for API outage).""" + + def __init__(self, guardrail_name: str, status_code: int): + super().__init__( + guardrail_name=guardrail_name, + event_hook="pre_call", + default_on=True, + ) + self.status_code = status_code + self.calls = 0 + + def should_run_guardrail(self, data, event_type) -> bool: + return True + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.calls += 1 + raise HTTPException( + status_code=self.status_code, detail="Simulated HTTP error" + ) + + class AlwaysPassGuardrail(CustomGuardrail): """Mock guardrail that always passes.""" @@ -350,6 +372,125 @@ async def test_guardrail_not_found_uses_on_fail(): litellm.callbacks = original_callbacks +@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") +@pytest.mark.asyncio +async def test_on_error_next_fallback_on_api_outage_on_fail_blocks_content(): + """ + Policy intervention (400) uses on_fail; technical error (503) uses on_error. + + Primary returns 503 -> on_error: next -> fallback runs -> allow. + """ + primary = HttpStatusGuardrail("primary-mod", status_code=503) + fallback = AlwaysPassGuardrail("fallback-filter") + + pipeline = GuardrailPipeline( + mode="pre_call", + steps=[ + PipelineStep( + guardrail="primary-mod", + on_fail="block", + on_error="next", + on_pass="allow", + ), + PipelineStep( + guardrail="fallback-filter", + on_fail="block", + on_pass="allow", + ), + ], + ) + + original_callbacks = litellm.callbacks.copy() + litellm.callbacks = [primary, fallback] + + try: + result = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode=pipeline.mode, + data={"messages": [{"role": "user", "content": "any"}]}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="mod-fallback", + ) + + assert primary.calls == 1 + assert fallback.calls == 1 + assert result.terminal_action == "allow" + assert result.step_results[0].outcome == "error" + assert result.step_results[0].action_taken == "next" + assert result.step_results[1].outcome == "pass" + finally: + litellm.callbacks = original_callbacks + + +@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed") +@pytest.mark.asyncio +async def test_on_fail_next_on_content_on_error_block_stops_api_fallback(): + """ + Content policy fail (400) uses on_fail: next; API error uses on_error: block (no second step). + """ + primary_content = AlwaysFailGuardrail("strict-mod") + primary_api = HttpStatusGuardrail("strict-mod", status_code=503) + fallback = AlwaysPassGuardrail("fallback-filter") + + # Content violation: on_fail next -> would reach fallback if we had two steps + pipeline_content = GuardrailPipeline( + mode="pre_call", + steps=[ + PipelineStep( + guardrail="strict-mod", + on_fail="next", + on_error="block", + on_pass="allow", + ), + PipelineStep( + guardrail="fallback-filter", + on_fail="block", + on_pass="allow", + ), + ], + ) + + original_callbacks = litellm.callbacks.copy() + litellm.callbacks = [primary_content, fallback] + + try: + result = await PipelineExecutor.execute_steps( + steps=pipeline_content.steps, + mode=pipeline_content.mode, + data={"messages": [{"role": "user", "content": "bad"}]}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="test", + ) + assert result.terminal_action == "allow" + assert primary_content.calls == 1 + assert fallback.calls == 1 + finally: + litellm.callbacks = original_callbacks + + # API outage: on_error block -> do not run fallback + fallback.calls = 0 + original_callbacks = litellm.callbacks.copy() + litellm.callbacks = [primary_api, fallback] + try: + result = await PipelineExecutor.execute_steps( + steps=pipeline_content.steps, + mode=pipeline_content.mode, + data={"messages": [{"role": "user", "content": "ok"}]}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="test", + ) + assert result.terminal_action == "block" + assert primary_api.calls == 1 + assert fallback.calls == 0 + assert result.step_results[0].outcome == "error" + assert result.step_results[0].action_taken == "block" + finally: + litellm.callbacks = original_callbacks + + @pytest.mark.asyncio async def test_guardrail_not_found_with_next_continues(): """ diff --git a/tests/test_litellm/proxy/test_health_check_functions.py b/tests/test_litellm/proxy/test_health_check_functions.py index 354698b02fe..13d2131efad 100644 --- a/tests/test_litellm/proxy/test_health_check_functions.py +++ b/tests/test_litellm/proxy/test_health_check_functions.py @@ -481,7 +481,7 @@ async def test_perform_health_check_and_save_passes_model_id_to_perform_health_c unhealthy = [] async def mock_perform_health_check(model_list, model=None, cli_model=None, details=True, model_id=None, max_concurrency=None): - return healthy, unhealthy + return healthy, unhealthy, {} with patch( "litellm.proxy.health_endpoints._health_endpoints.perform_health_check", diff --git a/tests/test_litellm/proxy/test_shared_health_check.py b/tests/test_litellm/proxy/test_shared_health_check.py index 0212d87baab..20c96c8152d 100644 --- a/tests/test_litellm/proxy/test_shared_health_check.py +++ b/tests/test_litellm/proxy/test_shared_health_check.py @@ -246,7 +246,7 @@ class TestSharedHealthCheckManager: model_list = [{"model_name": "test-model", "litellm_params": {"model": "test-model"}}] with patch("litellm.proxy.health_check_utils.shared_health_check_manager.perform_health_check") as mock_perform: - healthy, unhealthy = await shared_health_manager.perform_shared_health_check( + healthy, unhealthy, _ = await shared_health_manager.perform_shared_health_check( model_list, details=True ) @@ -268,9 +268,9 @@ class TestSharedHealthCheckManager: expected_unhealthy = [] with patch("litellm.proxy.health_check_utils.shared_health_check_manager.perform_health_check") as mock_perform: - mock_perform.return_value = (expected_healthy, expected_unhealthy) + mock_perform.return_value = (expected_healthy, expected_unhealthy, {}) - healthy, unhealthy = await shared_health_manager.perform_shared_health_check( + healthy, unhealthy, _ = await shared_health_manager.perform_shared_health_check( model_list, details=True ) @@ -302,7 +302,7 @@ class TestSharedHealthCheckManager: model_list = [{"model_name": "test-model", "litellm_params": {"model": "test-model"}}] with patch("asyncio.sleep") as mock_sleep: # Mock sleep to avoid actual delay - healthy, unhealthy = await shared_health_manager.perform_shared_health_check( + healthy, unhealthy, _ = await shared_health_manager.perform_shared_health_check( model_list, details=True ) @@ -324,9 +324,9 @@ class TestSharedHealthCheckManager: with patch("asyncio.sleep") as mock_sleep, \ patch("litellm.proxy.health_check_utils.shared_health_check_manager.perform_health_check") as mock_perform: - mock_perform.return_value = (expected_healthy, expected_unhealthy) + mock_perform.return_value = (expected_healthy, expected_unhealthy, {}) - healthy, unhealthy = await shared_health_manager.perform_shared_health_check( + healthy, unhealthy, _ = await shared_health_manager.perform_shared_health_check( model_list, details=True ) diff --git a/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py b/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py new file mode 100644 index 00000000000..e2c13b952dd --- /dev/null +++ b/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py @@ -0,0 +1,789 @@ +""" +Tests for health check failures integrating with allowed_fails_policy cooldown pipeline. + +When enable_health_check_routing is True and a health check fails, the failure +should increment the same counters used by allowed_fails_policy, using the +actual exception type from the health check error. +""" + +from unittest.mock import patch + +import pytest + +import litellm +from litellm.proxy.health_check import run_with_timeout +from litellm.router import Router +from litellm.types.router import AllowedFailsPolicy + + +def _make_model(model_id: str, model_name: str = "gpt-4") -> dict: + return { + "model_name": model_name, + "litellm_params": {"model": model_name, "api_key": "fake-key"}, + "model_info": {"id": model_id}, + } + + +class TestAhealthCheckExceptionPreservation: + """Test that ahealth_check() preserves the exception object in its return dict.""" + + @pytest.mark.asyncio + async def test_run_with_timeout_returns_timeout_exception(self): + """run_with_timeout should return a litellm.Timeout in the 'exception' key on timeout.""" + import asyncio + + async def slow_task(): + await asyncio.sleep(10) + + result = await run_with_timeout(slow_task(), timeout=0.01) + + assert "error" in result + assert "exception" in result + assert isinstance(result["exception"], litellm.Timeout) + + +class TestHealthCheckEndpointExceptionPropagation: + """Test that _perform_health_check returns exceptions via exceptions_by_model_id.""" + + @pytest.mark.asyncio + async def test_unhealthy_endpoint_dict_exception_in_map(self): + """When ahealth_check returns {"error": ..., "exception": e}, the exception + must appear in exceptions_by_model_id keyed by model_id — not in the endpoint dict.""" + from unittest.mock import AsyncMock, patch + + from litellm.proxy.health_check import _perform_health_check + + auth_error = litellm.AuthenticationError( + message="Invalid key", llm_provider="openai", model="gpt-4" + ) + model_list = [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake"}, + "model_info": {"id": "deploy-abc"}, + } + ] + + with patch( + "litellm.proxy.health_check.litellm.ahealth_check", + new=AsyncMock(return_value={"error": "auth failed", "exception": auth_error}), + ): + healthy, unhealthy, exc_map = await _perform_health_check(model_list) + + assert len(unhealthy) == 1 + assert "exception" not in unhealthy[0], "exception must not be in endpoint dict" + assert exc_map.get("deploy-abc") is auth_error + + @pytest.mark.asyncio + async def test_raw_exception_from_gather_in_map(self): + """When asyncio.gather returns a raw Exception, it must appear in + exceptions_by_model_id — not in the endpoint dict.""" + from unittest.mock import patch + + from litellm.proxy.health_check import _perform_health_check + + raw_exc = litellm.RateLimitError( + message="Rate limited", llm_provider="openai", model="gpt-4" + ) + model_list = [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake"}, + "model_info": {"id": "deploy-xyz"}, + } + ] + + # Simulate asyncio.gather returning a raw exception for this task + with patch( + "litellm.proxy.health_check._run_model_health_check", + side_effect=raw_exc, + ): + healthy, unhealthy, exc_map = await _perform_health_check(model_list) + + assert len(unhealthy) == 1 + assert "exception" not in unhealthy[0], "exception must not be in endpoint dict" + assert exc_map.get("deploy-xyz") is raw_exc + + +class TestGetAllowedFailsFromPolicyWithHealthCheckExceptions: + """Test that get_allowed_fails_from_policy correctly resolves thresholds for health-check exceptions.""" + + @pytest.mark.parametrize( + "exception_type, policy_field, threshold", + [ + (litellm.Timeout, "TimeoutErrorAllowedFails", 5), + (litellm.AuthenticationError, "AuthenticationErrorAllowedFails", 3), + (litellm.RateLimitError, "RateLimitErrorAllowedFails", 10), + ( + litellm.ContentPolicyViolationError, + "ContentPolicyViolationErrorAllowedFails", + 2, + ), + (litellm.BadRequestError, "BadRequestErrorAllowedFails", 7), + ], + ) + def test_policy_resolves_for_health_check_exception_types( + self, exception_type, policy_field, threshold + ): + """Each exception type from a health check should resolve to its policy threshold.""" + policy = AllowedFailsPolicy(**{policy_field: threshold}) + router = Router( + model_list=[_make_model("d1")], + allowed_fails_policy=policy, + ) + exception = exception_type( + message="health check failed", llm_provider="openai", model="gpt-4" + ) + result = router.get_allowed_fails_from_policy(exception=exception) + assert result == threshold + + def test_policy_returns_none_for_unmatched_exception(self): + """When no policy field matches the exception type, return None (fall back to allowed_fails).""" + policy = AllowedFailsPolicy(TimeoutErrorAllowedFails=5) + router = Router( + model_list=[_make_model("d1")], + allowed_fails_policy=policy, + ) + # Use a generic Exception that doesn't match any policy field + result = router.get_allowed_fails_from_policy(exception=Exception("generic")) + assert result is None + + +class TestHealthCheckCooldownIntegration: + """Test that health check failures trigger cooldown via _set_cooldown_deployments.""" + + def test_health_check_failure_increments_failed_calls(self): + """Health check failure should increment the failed_calls counter.""" + from litellm.router_utils.cooldown_handlers import ( + should_cooldown_based_on_allowed_fails_policy, + ) + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=3), + ) + + timeout_exc = litellm.Timeout( + message="Health check timeout", model="gpt-4", llm_provider="openai" + ) + + # First call: should not cooldown (1 <= 3) + result = should_cooldown_based_on_allowed_fails_policy( + litellm_router_instance=router, + deployment="deploy-1", + original_exception=timeout_exc, + ) + assert result is False + + # Check counter was incremented + current_fails = router.failed_calls.get_cache(key="deploy-1") + assert current_fails == 1 + + def test_health_check_failure_triggers_cooldown_at_threshold(self): + """After exceeding allowed_fails threshold, deployment should enter cooldown.""" + from litellm.router_utils.cooldown_handlers import ( + should_cooldown_based_on_allowed_fails_policy, + ) + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(AuthenticationErrorAllowedFails=2), + ) + + auth_exc = litellm.AuthenticationError( + message="Invalid key", model="gpt-4", llm_provider="openai" + ) + + # Fails 1 and 2: should not cooldown + for _ in range(2): + result = should_cooldown_based_on_allowed_fails_policy( + litellm_router_instance=router, + deployment="deploy-1", + original_exception=auth_exc, + ) + assert result is False + + # Fail 3: should trigger cooldown (3 > 2) + result = should_cooldown_based_on_allowed_fails_policy( + litellm_router_instance=router, + deployment="deploy-1", + original_exception=auth_exc, + ) + assert result is True + + def test_health_check_failure_falls_back_to_allowed_fails(self): + """When policy has no matching field, fall back to generic allowed_fails.""" + from litellm.router_utils.cooldown_handlers import ( + should_cooldown_based_on_allowed_fails_policy, + ) + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=10), + allowed_fails=1, + ) + + # Use an exception that doesn't match TimeoutErrorAllowedFails + # InternalServerError is not checked by get_allowed_fails_from_policy + # so it will fall back to allowed_fails=1 + generic_exc = Exception("Some internal error") + + # Fail 1: should not cooldown (1 <= 1) + result = should_cooldown_based_on_allowed_fails_policy( + litellm_router_instance=router, + deployment="deploy-1", + original_exception=generic_exc, + ) + assert result is False + + # Fail 2: should trigger cooldown (2 > 1) + result = should_cooldown_based_on_allowed_fails_policy( + litellm_router_instance=router, + deployment="deploy-1", + original_exception=generic_exc, + ) + assert result is True + + def test_healthy_endpoints_do_not_trigger_cooldown(self): + """Healthy endpoints should not increment any failure counters.""" + from litellm.router_utils.cooldown_handlers import _set_cooldown_deployments + + router = Router( + model_list=[_make_model("deploy-1")], + allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=1), + enable_health_check_routing=True, + ) + + # Simulate healthy endpoint -- no exception, no cooldown call + healthy_endpoint = {"model_id": "deploy-1"} + # Should have no exception key + assert "exception" not in healthy_endpoint + + # Verify failed_calls counter is untouched + current_fails = router.failed_calls.get_cache(key="deploy-1") + assert current_fails is None + + def test_disable_cooldowns_prevents_health_check_cooldown(self): + """When disable_cooldowns=True, health check failures should not trigger cooldown.""" + from litellm.router_utils.cooldown_handlers import _set_cooldown_deployments + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=0), + enable_health_check_routing=True, + disable_cooldowns=True, + ) + + timeout_exc = litellm.Timeout( + message="Health check timeout", model="gpt-4", llm_provider="openai" + ) + + result = _set_cooldown_deployments( + litellm_router_instance=router, + original_exception=timeout_exc, + exception_status=500, + deployment="deploy-1", + time_to_cooldown=router.cooldown_time, + ) + assert result is False + + +class TestWriteHealthStateIntegration: + """Test _write_health_state_to_router_cache integrates with cooldown pipeline.""" + + def test_unhealthy_endpoint_triggers_set_cooldown(self): + """_write_health_state_to_router_cache should call _set_cooldown_deployments for unhealthy endpoints.""" + import litellm.proxy.proxy_server as proxy_module + from litellm.proxy.proxy_server import _write_health_state_to_router_cache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=5), + enable_health_check_routing=True, + ) + + timeout_exc = litellm.Timeout( + message="Health check timeout", model="", llm_provider="" + ) + + unhealthy_endpoints = [ + {"model_id": "deploy-1", "error": "timeout"}, + ] + healthy_endpoints = [ + {"model_id": "deploy-2"}, + ] + + with patch.object(proxy_module, "llm_router", router): + with patch( + "litellm.router_utils.cooldown_handlers._set_cooldown_deployments" + ) as mock_cooldown: + _write_health_state_to_router_cache( + healthy_endpoints=healthy_endpoints, + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": timeout_exc}, + ) + mock_cooldown.assert_called_once_with( + litellm_router_instance=router, + original_exception=timeout_exc, + exception_status=408, # Timeout has status_code 408 + deployment="deploy-1", + time_to_cooldown=router.cooldown_time, + ) + + def test_unhealthy_endpoint_without_exception_skips_cooldown(self): + """Unhealthy endpoints without an exception key should not trigger cooldown.""" + import litellm.proxy.proxy_server as proxy_module + from litellm.proxy.proxy_server import _write_health_state_to_router_cache + + router = Router( + model_list=[_make_model("deploy-1")], + allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=5), + enable_health_check_routing=True, + ) + + unhealthy_endpoints = [ + {"model_id": "deploy-1", "error": "unknown failure"}, + ] + + with patch.object(proxy_module, "llm_router", router): + with patch( + "litellm.router_utils.cooldown_handlers._set_cooldown_deployments" + ) as mock_cooldown: + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + # no exceptions_by_model_id → cooldown should not fire + ) + mock_cooldown.assert_not_called() + + def test_unhealthy_endpoint_increments_failure_counter(self): + """Unhealthy endpoints should call increment_deployment_failures_for_current_minute.""" + import litellm.proxy.proxy_server as proxy_module + from litellm.proxy.proxy_server import _write_health_state_to_router_cache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=10), + enable_health_check_routing=True, + ) + + rate_exc = litellm.RateLimitError( + message="Rate limited", model="gpt-4", llm_provider="openai" + ) + + unhealthy_endpoints = [ + {"model_id": "deploy-1", "error": "rate limited"}, + ] + + with patch.object(proxy_module, "llm_router", router): + with patch( + "litellm.router_utils.router_callbacks.track_deployment_metrics.increment_deployment_failures_for_current_minute" + ) as mock_increment: + with patch( + "litellm.router_utils.cooldown_handlers._set_cooldown_deployments" + ): + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": rate_exc}, + ) + mock_increment.assert_called_once_with( + litellm_router_instance=router, + deployment_id="deploy-1", + ) + + +class TestHealthCheckFilterBypassWithPolicy: + """ + When allowed_fails_policy is set, the binary health check filter should be + bypassed so cooldown is the sole routing exclusion mechanism. + """ + + def test_filter_bypassed_when_policy_set(self): + """Binary health check filter is a no-op when allowed_fails_policy is configured.""" + import time + + from litellm.caching.caching import DualCache + from litellm.router_utils.health_state_cache import DeploymentHealthCache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(AuthenticationErrorAllowedFails=3), + enable_health_check_routing=True, + ) + + # Mark deploy-1 as unhealthy in the health state cache + cache = DualCache() + health_cache = DeploymentHealthCache(cache=cache, staleness_threshold=60.0) + health_cache.set_deployment_health_states( + { + "deploy-1": { + "is_healthy": False, + "timestamp": time.time(), + "reason": "test", + }, + } + ) + router.health_state_cache = health_cache + + deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")] + + # Filter should pass all through because policy is set + result = router._filter_health_check_unhealthy_deployments(deployments) + assert ( + len(result) == 2 + ), "Binary filter should be bypassed when allowed_fails_policy is set" + + def test_filter_active_when_no_policy(self): + """Binary health check filter still works when no allowed_fails_policy is configured.""" + import time + + from litellm.caching.caching import DualCache + from litellm.router_utils.health_state_cache import DeploymentHealthCache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + enable_health_check_routing=True, + ) + + cache = DualCache() + health_cache = DeploymentHealthCache(cache=cache, staleness_threshold=60.0) + health_cache.set_deployment_health_states( + { + "deploy-1": { + "is_healthy": False, + "timestamp": time.time(), + "reason": "test", + }, + } + ) + router.health_state_cache = health_cache + + deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")] + + result = router._filter_health_check_unhealthy_deployments(deployments) + assert len(result) == 1 + assert result[0]["model_info"]["id"] == "deploy-2" + + @pytest.mark.asyncio + async def test_async_filter_bypassed_when_policy_set(self): + """Async version also bypasses when allowed_fails_policy is set.""" + import time + + from litellm.caching.caching import DualCache + from litellm.router_utils.health_state_cache import DeploymentHealthCache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=2), + enable_health_check_routing=True, + ) + + cache = DualCache() + health_cache = DeploymentHealthCache(cache=cache, staleness_threshold=60.0) + health_cache.set_deployment_health_states( + { + "deploy-1": { + "is_healthy": False, + "timestamp": time.time(), + "reason": "test", + }, + } + ) + router.health_state_cache = health_cache + + deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")] + + result = await router._async_filter_health_check_unhealthy_deployments( + deployments + ) + assert len(result) == 2 + + +class TestAllDeploymentsInCooldownSafetyNet: + """ + When enable_health_check_routing=True and ALL deployments enter cooldown, + the async routing path should bypass the cooldown filter and return all + deployments rather than blocking all traffic. + """ + + def test_raw_cooldown_filter_returns_empty_when_all_cooled(self): + """The raw _filter_cooldown_deployments has no safety net -- it returns empty.""" + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + enable_health_check_routing=True, + ) + deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")] + result = router._filter_cooldown_deployments( + healthy_deployments=deployments, + cooldown_deployments=["deploy-1", "deploy-2"], + ) + assert result == [] # raw filter has no safety net + + @pytest.mark.asyncio + async def test_async_routing_path_bypasses_all_cooldown(self): + """In the async routing path, all-in-cooldown with enable_health_check_routing + returns the full list instead of empty (safety net).""" + from unittest.mock import AsyncMock + + from litellm.router_utils.cooldown_handlers import ( + _async_get_cooldown_deployments, + ) + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(AuthenticationErrorAllowedFails=0), + enable_health_check_routing=True, + ) + + deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")] + + # Simulate all deployments in cooldown + with patch( + "litellm.router._async_get_cooldown_deployments", + new=AsyncMock(return_value=["deploy-1", "deploy-2"]), + ): + # The safety net in async_get_available_deployment should restore + # all deployments when the cooldown filter empties the list + _pre = deployments.copy() + filtered = router._filter_cooldown_deployments( + healthy_deployments=deployments, + cooldown_deployments=["deploy-1", "deploy-2"], + ) + # If filtered is empty and enable_health_check_routing is True, + # the routing path restores _pre_cooldown_deployments + if not filtered and router.enable_health_check_routing: + filtered = _pre + + assert ( + len(filtered) == 2 + ), "Safety net should return all deployments when all are in cooldown" + + +class TestHealthCheckIgnoreTransientErrors: + """ + When health_check_ignore_transient_errors=True, health check failures with + 429 or 408 status codes should NOT increment failure counters or trigger cooldown. + 401, 404, and 5xx errors should still be processed normally. + """ + + def test_429_skipped_when_flag_enabled(self): + """429 from health check does not trigger cooldown when flag is set.""" + import litellm.proxy.proxy_server as proxy_module + from litellm.proxy.proxy_server import _write_health_state_to_router_cache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=0), + enable_health_check_routing=True, + health_check_ignore_transient_errors=True, + ) + + rate_exc = litellm.RateLimitError( + message="Rate limited", model="gpt-4", llm_provider="openai" + ) + assert getattr(rate_exc, "status_code", None) == 429 + + unhealthy_endpoints = [ + {"model_id": "deploy-1", "error": "rate limited"}, + ] + + with patch.object(proxy_module, "llm_router", router): + with patch( + "litellm.router_utils.cooldown_handlers._set_cooldown_deployments" + ) as mock_cooldown: + with patch( + "litellm.router_utils.router_callbacks.track_deployment_metrics.increment_deployment_failures_for_current_minute" + ) as mock_increment: + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": rate_exc}, + ) + mock_cooldown.assert_not_called() + mock_increment.assert_not_called() + + def test_408_skipped_when_flag_enabled(self): + """408 from health check does not trigger cooldown when flag is set.""" + import litellm.proxy.proxy_server as proxy_module + from litellm.proxy.proxy_server import _write_health_state_to_router_cache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=0), + enable_health_check_routing=True, + health_check_ignore_transient_errors=True, + ) + + timeout_exc = litellm.Timeout( + message="Health check timeout exceeded", model="", llm_provider="" + ) + + unhealthy_endpoints = [ + {"model_id": "deploy-1", "error": "timeout"}, + ] + + with patch.object(proxy_module, "llm_router", router): + with patch( + "litellm.router_utils.cooldown_handlers._set_cooldown_deployments" + ) as mock_cooldown: + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": timeout_exc}, + ) + mock_cooldown.assert_not_called() + + def test_401_still_triggers_cooldown_when_flag_enabled(self): + """Auth errors (401) still trigger cooldown even when flag is set.""" + import litellm.proxy.proxy_server as proxy_module + from litellm.proxy.proxy_server import _write_health_state_to_router_cache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(AuthenticationErrorAllowedFails=0), + enable_health_check_routing=True, + health_check_ignore_transient_errors=True, + ) + + auth_exc = litellm.AuthenticationError( + message="Invalid key", model="gpt-4", llm_provider="openai" + ) + + unhealthy_endpoints = [ + {"model_id": "deploy-1", "error": "auth failed"}, + ] + + with patch.object(proxy_module, "llm_router", router): + with patch( + "litellm.router_utils.cooldown_handlers._set_cooldown_deployments" + ) as mock_cooldown: + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": auth_exc}, + ) + mock_cooldown.assert_called_once() + + def test_429_not_written_to_health_state_cache_when_flag_enabled(self): + """429 endpoint is excluded from health state cache when flag is set, + so the binary health check filter does not mark it as unhealthy.""" + import litellm.proxy.proxy_server as proxy_module + from litellm.proxy.proxy_server import _write_health_state_to_router_cache + + router = Router( + model_list=[_make_model("deploy-1")], + enable_health_check_routing=True, + health_check_ignore_transient_errors=True, + ) + + rate_exc = litellm.RateLimitError( + message="Rate limited", model="gpt-4", llm_provider="openai" + ) + + unhealthy_endpoints = [ + {"model_id": "deploy-1", "error": "rate limited"}, + ] + + with patch.object(proxy_module, "llm_router", router): + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": rate_exc}, + ) + + # Health state cache should have NO entry for deploy-1 + # (429 was ignored, not written as unhealthy) + unhealthy_ids = router.health_state_cache.get_unhealthy_deployment_ids() + assert "deploy-1" not in unhealthy_ids + + def test_429_triggers_cooldown_when_flag_disabled(self): + """When flag is False (default), 429 still triggers cooldown.""" + import litellm.proxy.proxy_server as proxy_module + from litellm.proxy.proxy_server import _write_health_state_to_router_cache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=0), + enable_health_check_routing=True, + health_check_ignore_transient_errors=False, + ) + + rate_exc = litellm.RateLimitError( + message="Rate limited", model="gpt-4", llm_provider="openai" + ) + + unhealthy_endpoints = [ + {"model_id": "deploy-1", "error": "rate limited"}, + ] + + with patch.object(proxy_module, "llm_router", router): + with patch( + "litellm.router_utils.cooldown_handlers._set_cooldown_deployments" + ) as mock_cooldown: + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={"deploy-1": rate_exc}, + ) + mock_cooldown.assert_called_once() + + +class TestSharedCacheTransientErrorFilter: + """ + When SharedHealthCheckManager returns cached results, exceptions_by_model_id + is always {}. The filter must fall back to the 'exception_status' field stored + on each endpoint dict so 429/408 endpoints are still excluded correctly. + """ + + def test_cached_429_excluded_via_exception_status_field(self): + """Cache-hit path: endpoint with exception_status=429 is excluded from health state.""" + import litellm.proxy.proxy_server as proxy_module + from litellm.proxy.proxy_server import _write_health_state_to_router_cache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + enable_health_check_routing=True, + health_check_ignore_transient_errors=True, + ) + + # Simulate a cache-hit endpoint: exception_status stored as int, no exceptions dict + unhealthy_endpoints = [ + {"model_id": "deploy-1", "error": "rate limited", "exception_status": 429}, + ] + + with patch.object(proxy_module, "llm_router", router): + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={}, + ) + + # deploy-1 should NOT be marked unhealthy (429 was filtered) + unhealthy_ids = router.health_state_cache.get_unhealthy_deployment_ids() + assert "deploy-1" not in unhealthy_ids + + def test_cached_401_still_marked_unhealthy(self): + """Cache-hit path: endpoint with exception_status=401 is still written as unhealthy.""" + import litellm.proxy.proxy_server as proxy_module + from litellm.proxy.proxy_server import _write_health_state_to_router_cache + + router = Router( + model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")], + enable_health_check_routing=True, + health_check_ignore_transient_errors=True, + ) + + unhealthy_endpoints = [ + {"model_id": "deploy-1", "error": "auth failed", "exception_status": 401}, + ] + + with patch.object(proxy_module, "llm_router", router): + _write_health_state_to_router_cache( + healthy_endpoints=[], + unhealthy_endpoints=unhealthy_endpoints, + exceptions_by_model_id={}, + ) + + unhealthy_ids = router.health_state_cache.get_unhealthy_deployment_ids() + assert "deploy-1" in unhealthy_ids diff --git a/tests/test_litellm/router_utils/test_router_health_check_routing.py b/tests/test_litellm/router_utils/test_router_health_check_routing.py index f40144b44c9..b87a39ac1de 100644 --- a/tests/test_litellm/router_utils/test_router_health_check_routing.py +++ b/tests/test_litellm/router_utils/test_router_health_check_routing.py @@ -50,6 +50,7 @@ class TestFilterHealthCheckUnhealthyDeployments: def __init__(self): self.enable_health_check_routing = enable self.health_state_cache = health_cache + self.allowed_fails_policy = None # Import the actual method and bind it from litellm.router import Router @@ -125,6 +126,7 @@ class TestAsyncFilterHealthCheckUnhealthyDeployments: def __init__(self): self.enable_health_check_routing = enable self.health_state_cache = health_cache + self.allowed_fails_policy = None fake = FakeRouter() fake._async_filter_health_check_unhealthy_deployments = ( diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 3df10901492..262dce439c0 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -5,7 +5,6 @@ import sys from unittest.mock import AsyncMock, MagicMock, patch import pytest -from fastapi.testclient import TestClient sys.path.insert( 0, os.path.abspath("../../..") @@ -13,7 +12,6 @@ sys.path.insert( import litellm -from litellm.router_utils.fallback_event_handlers import run_async_fallback def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): @@ -127,7 +125,7 @@ async def test_async_router_acreate_file(): """ Write to all deployments of a model """ - from unittest.mock import MagicMock, call, patch + from unittest.mock import MagicMock, patch router = litellm.Router( model_list=[ @@ -747,7 +745,7 @@ async def test_router_ageneric_api_call_with_fallbacks_helper(): """ Test the _ageneric_api_call_with_fallbacks_helper method with various scenarios """ - from unittest.mock import AsyncMock, MagicMock, patch + from unittest.mock import patch router = litellm.Router( model_list=[ @@ -1134,10 +1132,9 @@ def test_get_model_access_groups_cache_invalidation_upsert_deployment(): @pytest.mark.asyncio async def test_acompletion_streaming_iterator(): """Test _acompletion_streaming_iterator for normal streaming and fallback behavior.""" - from unittest.mock import AsyncMock, MagicMock + from unittest.mock import MagicMock from litellm.exceptions import MidStreamFallbackError - from litellm.types.utils import ModelResponseStream # Helper class for creating async iterators class AsyncIterator: @@ -2847,3 +2844,286 @@ def test_combine_fallback_usage(): assert chunk.usage.prompt_tokens == 10 assert chunk.usage.completion_tokens == 5 assert chunk.usage.total_tokens == 15 + + +@pytest.mark.asyncio +async def test_team_scoped_model_fallback(): + """ + Test that fallback works correctly for team-scoped models. + + When a team-scoped model fails and the fallback model is also team-scoped, + the router should find the fallback deployment by matching team_public_model_name. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "team-a-primary-internal", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake"}, + "model_info": { + "team_id": "team-a", + "team_public_model_name": "primary-model", + }, + }, + { + "model_name": "team-a-fallback-internal", + "litellm_params": { + "model": "gpt-4", + "api_key": "fake", + "mock_response": "fallback success from team-a", + }, + "model_info": { + "team_id": "team-a", + "team_public_model_name": "fallback-model", + }, + }, + ], + fallbacks=[{"primary-model": ["fallback-model"]}], + ) + + response = await router.acompletion( + model="primary-model", + messages=[{"role": "user", "content": "Hello"}], + metadata={"user_api_key_team_id": "team-a"}, + mock_testing_fallbacks=True, + ) + assert response is not None + assert response.choices[0].message.content == "fallback success from team-a" + + +@pytest.mark.asyncio +async def test_team_scoped_model_fallback_to_global(): + """ + Test that a team-scoped model can fall back to a global (non-team) model. + + Global models (no team_id on deployment) should be accessible as fallback + targets for team-scoped requests. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "team-a-primary-internal", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake"}, + "model_info": { + "team_id": "team-a", + "team_public_model_name": "primary-model", + }, + }, + { + "model_name": "global-fallback", + "litellm_params": { + "model": "gpt-4", + "api_key": "fake", + "mock_response": "global fallback success", + }, + }, + ], + fallbacks=[{"primary-model": ["global-fallback"]}], + ) + + response = await router.acompletion( + model="primary-model", + messages=[{"role": "user", "content": "Hello"}], + metadata={"user_api_key_team_id": "team-a"}, + mock_testing_fallbacks=True, + ) + assert response is not None + assert response.choices[0].message.content == "global fallback success" + + +@pytest.mark.asyncio +async def test_team_scoped_model_fallback_cross_team_blocked(): + """ + Test that cross-team fallback is correctly blocked. + + When team-a's model fails and the fallback target is scoped to team-b, + the router should NOT use it (team isolation). + """ + router = litellm.Router( + model_list=[ + { + "model_name": "team-a-primary-internal", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake"}, + "model_info": { + "team_id": "team-a", + "team_public_model_name": "primary-model", + }, + }, + { + "model_name": "team-b-fallback-internal", + "litellm_params": { + "model": "gpt-4", + "api_key": "fake", + "mock_response": "team-b response - should not reach here", + }, + "model_info": { + "team_id": "team-b", + "team_public_model_name": "fallback-model", + }, + }, + ], + fallbacks=[{"primary-model": ["fallback-model"]}], + ) + + with pytest.raises(Exception): + await router.acompletion( + model="primary-model", + messages=[{"role": "user", "content": "Hello"}], + metadata={"user_api_key_team_id": "team-a"}, + mock_testing_fallbacks=True, + ) + + +def test_get_all_deployments_with_team_id(): + """ + Test that _get_all_deployments with team_id can find deployments + by team_public_model_name when the model_name is not in the index. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "internal-team-deployment", + "litellm_params": {"model": "gpt-4", "api_key": "fake"}, + "model_info": { + "team_id": "team-x", + "team_public_model_name": "gpt-4", + }, + }, + ], + ) + + # Without team_id: "gpt-4" is not in the model_name index (internal name is different) + deployments = router._get_all_deployments(model_name="gpt-4") + assert len(deployments) == 0 + + # With correct team_id: should find via O(n) scan matching team_public_model_name + deployments = router._get_all_deployments(model_name="gpt-4", team_id="team-x") + assert len(deployments) == 1 + assert deployments[0]["model_name"] == "internal-team-deployment" + + # With wrong team_id: should find nothing + deployments = router._get_all_deployments(model_name="gpt-4", team_id="team-y") + assert len(deployments) == 0 + + +def test_multiregion_team_deployments_unique_model_names(): + """ + Simulates athenahealth's exact setup: unique model_names per deployment, + same team_public_model_name, multiple regions. + + Verifies that _get_all_deployments returns ALL regional deployments + for a team when queried by team_public_model_name. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "metis-claude-us-east-1", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-sonnet", + "aws_region_name": "us-east-1", + "api_key": "fake", + }, + "model_info": { + "team_id": "metis-team", + "team_public_model_name": "claude-sonnet", + }, + }, + { + "model_name": "metis-claude-us-west-2", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-sonnet", + "aws_region_name": "us-west-2", + "api_key": "fake", + }, + "model_info": { + "team_id": "metis-team", + "team_public_model_name": "claude-sonnet", + }, + }, + ], + ) + + # "claude-sonnet" is NOT in the model_name index + assert "claude-sonnet" not in router.model_names + + # Without team_id: returns nothing (no model_name="claude-sonnet" in index, no O(n) scan) + deployments = router._get_all_deployments(model_name="claude-sonnet") + assert len(deployments) == 0 + + # With team_id: O(n) scan finds BOTH regional deployments + deployments = router._get_all_deployments( + model_name="claude-sonnet", team_id="metis-team" + ) + assert len(deployments) == 2 + deployment_names = {d["model_name"] for d in deployments} + assert deployment_names == {"metis-claude-us-east-1", "metis-claude-us-west-2"} + + # Each deployment has a unique ID (critical for cooldown/retry to work) + deployment_ids = {d["model_info"]["id"] for d in deployments} + assert len(deployment_ids) == 2, "Each deployment must have a unique ID for cooldown tracking" + + # Wrong team: returns nothing + deployments = router._get_all_deployments( + model_name="claude-sonnet", team_id="other-team" + ) + assert len(deployments) == 0 + + +@pytest.mark.asyncio +async def test_multiregion_team_failover_between_regions(): + """ + Simulates athenahealth's multiregion failover scenario: + - Two Bedrock deployments (us-east-1 and us-west-2) with unique model_names + - Same team_public_model_name ("claude-sonnet") + - Primary region fails → router should failover to second region + + This is the exact scenario Sean Glover from athenahealth will demonstrate. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "metis-claude-us-east-1", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-sonnet", + "api_key": "fake", + "mock_response": "response from us-east-1", + }, + "model_info": { + "team_id": "metis-team", + "team_public_model_name": "claude-sonnet", + }, + }, + { + "model_name": "metis-claude-us-west-2", + "litellm_params": { + "model": "bedrock/anthropic.claude-3-sonnet", + "api_key": "fake", + "mock_response": "response from us-west-2", + }, + "model_info": { + "team_id": "metis-team", + "team_public_model_name": "claude-sonnet", + }, + }, + ], + num_retries=1, + ) + + # Verify the router finds both deployments for the team + deployments = router._get_all_deployments( + model_name="claude-sonnet", team_id="metis-team" + ) + assert len(deployments) == 2, ( + "Router must find both regional deployments by team_public_model_name" + ) + + # Make a normal request — should succeed from one of the regions + response = await router.acompletion( + model="claude-sonnet", + messages=[{"role": "user", "content": "Hello"}], + metadata={"user_api_key_team_id": "metis-team"}, + ) + assert response is not None + assert response.choices[0].message.content in [ + "response from us-east-1", + "response from us-west-2", + ] diff --git a/tests/test_litellm/types/proxy/policy_engine/test_pipeline_types.py b/tests/test_litellm/types/proxy/policy_engine/test_pipeline_types.py index 21fecc015a3..c7d98548876 100644 --- a/tests/test_litellm/types/proxy/policy_engine/test_pipeline_types.py +++ b/tests/test_litellm/types/proxy/policy_engine/test_pipeline_types.py @@ -21,6 +21,7 @@ def test_pipeline_step_defaults(): step = PipelineStep(guardrail="my-guard") assert step.on_fail == "block" assert step.on_pass == "allow" + assert step.on_error is None assert step.pass_data is False assert step.modify_response_message is None @@ -33,9 +34,10 @@ def test_pipeline_step_valid_actions(): def test_pipeline_step_all_action_types(): for action in ("allow", "block", "next", "modify_response"): - step = PipelineStep(guardrail="g", on_fail=action, on_pass=action) + step = PipelineStep(guardrail="g", on_fail=action, on_pass=action, on_error=action) assert step.on_fail == action assert step.on_pass == action + assert step.on_error == action def test_pipeline_step_invalid_action_rejected(): @@ -48,6 +50,16 @@ def test_pipeline_step_invalid_on_pass_rejected(): PipelineStep(guardrail="my-guard", on_pass="skip") +def test_pipeline_step_on_error_valid(): + step = PipelineStep(guardrail="g", on_error="next", on_fail="block", on_pass="allow") + assert step.on_error == "next" + + +def test_pipeline_step_invalid_on_error_rejected(): + with pytest.raises(ValidationError): + PipelineStep(guardrail="my-guard", on_error="invalid") + + def test_pipeline_requires_at_least_one_step(): with pytest.raises(ValidationError): GuardrailPipeline(mode="pre_call", steps=[]) diff --git a/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx b/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx index b1768d5b81c..ba639594999 100644 --- a/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx +++ b/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx @@ -155,6 +155,14 @@ const FailIcon: React.FC = () => ( ); +const ApiFailureIcon: React.FC = () => ( + + + + + +); + // ───────────────────────────────────────────────────────────────────────────── // Connector // ───────────────────────────────────────────────────────────────────────────── @@ -349,6 +357,41 @@ const StepCard: React.FC = ({ )} + + {/* ON API FAILURE (technical / provider outage) — optional; defaults to ON FAIL */} +
+
+ + ON API FAILURE +
+ +