diff --git a/litellm/main.py b/litellm/main.py index eace9c630ba..8f14b1865f7 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": @@ -6198,9 +6198,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 @@ -6380,9 +6380,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: @@ -6605,9 +6605,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.") @@ -6911,9 +6911,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 @@ -7234,7 +7234,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 @@ -7246,6 +7247,7 @@ async def ahealth_check( return { "error": error_to_return, "raw_request_typed_dict": raw_request_typed_dict, + "exception": e, } @@ -7492,9 +7494,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 @@ -7505,9 +7507,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 @@ -7518,9 +7520,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/health_check.py b/litellm/proxy/health_check.py index bad9a754518..70621c5e979 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -95,7 +95,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): @@ -218,11 +223,15 @@ 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: + cleaned["exception"] = is_healthy["exception"] 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): + cleaned["exception"] = is_healthy unhealthy_endpoints.append(cleaned) return healthy_endpoints, unhealthy_endpoints diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3e0a5272522..55214d5726b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -375,9 +375,7 @@ from litellm.proxy.management_endpoints.fallback_management_endpoints import ( from litellm.proxy.management_endpoints.internal_user_endpoints import ( router as internal_user_router, ) -from litellm.proxy.management_endpoints.internal_user_endpoints import ( - user_update, -) +from litellm.proxy.management_endpoints.internal_user_endpoints import user_update from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import ( router as jwt_key_mapping_router, ) @@ -446,9 +444,7 @@ from litellm.proxy.openai_evals_endpoints.endpoints import router as evals_route from litellm.proxy.openai_files_endpoints.files_endpoints import ( router as openai_files_router, ) -from litellm.proxy.openai_files_endpoints.files_endpoints import ( - set_files_config, -) +from litellm.proxy.openai_files_endpoints.files_endpoints import set_files_config from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( passthrough_endpoint_router, ) @@ -503,7 +499,9 @@ from litellm.proxy.utils import ( get_error_message_str, get_server_root_path, handle_exception_on_proxy, + hash_password, hash_token, + migrate_passwords_to_scrypt_async, model_dump_with_preserved_fields, update_spend, ) @@ -550,9 +548,7 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( LiteLLM_UpperboundKeyGenerateParams, ) from litellm.types.realtime import RealtimeQueryParams -from litellm.types.router import ( - DeploymentTypedDict, -) +from litellm.types.router import DeploymentTypedDict from litellm.types.router import ModelInfo as RouterModelInfo from litellm.types.router import ( RouterGeneralSettings, @@ -639,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" ) @@ -870,6 +866,17 @@ async def proxy_startup_event(app: FastAPI): # noqa: PLR0915 user_api_key_cache=user_api_key_cache, ) + if prisma_client is not None: + + async def _run_pw_migration(): + try: + result = await migrate_passwords_to_scrypt_async(prisma_client) + verbose_proxy_logger.info(f"Password migration: {result}") + except Exception as e: + verbose_proxy_logger.warning(f"Password migration skipped: {e}") + + asyncio.create_task(_run_pw_migration()) + ProxyStartupEvent._initialize_startup_logging( llm_router=llm_router, proxy_logging_obj=proxy_logging_obj, @@ -1524,23 +1531,26 @@ 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 ) +spend_counter_cache = DualCache( + default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value +) 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 @@ -1694,6 +1704,130 @@ def cost_tracking(): ) +async def get_current_spend(counter_key: str, fallback_spend: float) -> float: + """ + Read current spend from the cross-pod spend counter. + + Reads Redis FIRST (authoritative cross-pod value), not DualCache's + async_get_cache which returns in-memory first. This is critical: + DualCache.async_get_cache returns stale per-pod values because each + pod's in-memory cache is only updated by that pod's own increments. + + Fallback chain: + 1. Redis counter (cross-pod, authoritative) + 2. In-memory counter (single-instance or Redis failure) + 3. Cached object's .spend from DB (cold start, no counter yet) + """ + # 1. Try Redis first (cross-pod authoritative) + if spend_counter_cache.redis_cache is not None: + try: + val = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) + if val is not None: + return float(val) + except Exception as e: + verbose_proxy_logger.debug( + "get_current_spend: Redis read failed for %s, falling back to in-memory: %s", + counter_key, + e, + ) + + # 2. Fall back to in-memory counter (single-instance or Redis failure) + val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + if val is not None: + return float(val) + + # 3. Final fallback: cached object's spend from DB + return fallback_spend + + +async def increment_spend_counters( + token: Optional[str], + team_id: Optional[str], + user_id: Optional[str], + response_cost: Optional[float], +): + """ + Atomically increment spend counters for budget enforcement. + + Uses spend_counter_cache (DualCache with Redis backend when available) + so counters are shared across all pods. Budget check functions read + from these counters via get_current_spend() (Redis-first). + + Awaited (not create_task) in the cost callback, so the counter is + updated before the next request's auth check runs. + """ + if response_cost is None or response_cost == 0: + return + + if token is not None: + # token arrives pre-hashed from metadata["user_api_key"] (auth flow + # hashes raw "sk-..." keys before they reach the callback). The + # startswith("sk-") check is a safety net matching update_cache — + # if a raw key somehow arrives, hash it; otherwise use as-is to + # avoid double-hashing (budget checks read valid_token.token which + # is single-hashed). + hashed_token = ( + hash_token(token=token) + if isinstance(token, str) and token.startswith("sk-") + else token + ) + await _init_and_increment_spend_counter( + counter_key=f"spend:key:{hashed_token}", + source_cache_key=hashed_token, + increment=response_cost, + ) + + if team_id is not None: + await _init_and_increment_spend_counter( + counter_key=f"spend:team:{team_id}", + source_cache_key=f"team_id:{team_id}", + increment=response_cost, + ) + + if user_id is not None and team_id is not None: + await _init_and_increment_spend_counter( + counter_key=f"spend:team_member:{user_id}:{team_id}", + source_cache_key=f"team_membership:{user_id}:{team_id}", + increment=response_cost, + ) + + +async def _init_and_increment_spend_counter( + counter_key: str, + source_cache_key: str, + increment: float, +): + """ + Initialize counter from cached object's DB-loaded spend if not yet set, + then atomically increment in both in-memory and Redis. + + On first access per pod: + 1. Check spend_counter_cache (in-memory -> Redis via DualCache for init check) + 2. If not found anywhere, read base spend from user_api_key_cache (DB-loaded object) + 3. Seed counter via async_increment_cache (not async_set_cache) to avoid a + check-then-set race: if two pods cold-start simultaneously, both may see + the counter as absent and seed it. Using increment instead of set means + the worst case is over-counting (conservative — blocks slightly early) + rather than under-counting (would allow overspend). + 4. Increment atomically (both in-memory + Redis) + """ + current = await spend_counter_cache.async_get_cache(key=counter_key) + if current is None: + source = await user_api_key_cache.async_get_cache(key=source_cache_key) + base_spend = 0.0 + if source is not None: + if isinstance(source, dict): + base_spend = source.get("spend", 0.0) or 0.0 + else: + base_spend = getattr(source, "spend", 0.0) or 0.0 + if base_spend > 0: + await spend_counter_cache.async_increment_cache( + key=counter_key, value=base_spend + ) + + await spend_counter_cache.async_increment_cache(key=counter_key, value=increment) + + async def update_cache( # noqa: PLR0915 token: Optional[str], user_id: Optional[str], @@ -1900,9 +2034,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 @@ -2121,6 +2255,10 @@ def _write_health_state_to_router_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, + ) try: if llm_router is None or not llm_router.enable_health_check_routing: @@ -2137,6 +2275,31 @@ 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 = endpoint.get("exception") + if original_exception is None: + continue + + exception_status = getattr(original_exception, "status_code", 500) + + 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) @@ -2559,6 +2722,7 @@ class ProxyConfig: ): ## INIT PROXY REDIS USAGE CLIENT ## redis_usage_cache = litellm.cache.cache + spend_counter_cache.redis_cache = redis_usage_cache # Note: PKCE verifier storage uses redis_usage_cache directly (not # user_api_key_cache) to avoid routing all API-key lookups through Redis. @@ -2726,12 +2890,36 @@ class ProxyConfig: return search_tools_parsed if search_tools_parsed else None + # Environment variable keys that must not be overridden via config because + # they can alter process execution, library loading, or network routing. + _BLOCKED_ENV_KEYS: Set[str] = { + "PATH", + "LD_PRELOAD", + "LD_LIBRARY_PATH", + "DYLD_LIBRARY_PATH", + "DYLD_INSERT_LIBRARIES", + "PYTHONPATH", + "PYTHONSTARTUP", + "PYTHONHOME", + "HOME", + "USER", + "SHELL", + "LOGNAME", + "NO_PROXY", + "no_proxy", + } + def _load_environment_variables(self, config: dict): ## ENVIRONMENT VARIABLES global premium_user environment_variables = config.get("environment_variables", None) if environment_variables: for key, value in environment_variables.items(): + if key in self._BLOCKED_ENV_KEYS: + verbose_proxy_logger.warning( + "Skipping blocked environment variable key: %s", key + ) + continue ######################################################### # handles this scenario: # ```yaml @@ -5032,10 +5220,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) @@ -5417,9 +5605,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 @@ -5547,6 +5735,11 @@ def _restamp_streaming_chunk_model( if _is_azure_model_router_request(requested_model_from_client): return chunk, model_mismatch_logged + # For fastest_response batch completions, preserve the winning model's name + # instead of stamping the comma-separated list the client sent. + if request_data.get("fastest_response", False): + return chunk, model_mismatch_logged + downstream_model = ( chunk.get("model") if isinstance(chunk, dict) else getattr(chunk, "model", None) ) @@ -5756,9 +5949,9 @@ class ProxyStartupEvent: """ from litellm.secret_managers.main import str_to_bool - _use_redis_transaction_buffer: Optional[ - Union[bool, str] - ] = general_settings.get("use_redis_transaction_buffer", False) + _use_redis_transaction_buffer: Optional[Union[bool, str]] = ( + general_settings.get("use_redis_transaction_buffer", False) + ) if isinstance(_use_redis_transaction_buffer, str): _use_redis_transaction_buffer = str_to_bool(_use_redis_transaction_buffer) @@ -9041,7 +9234,7 @@ def _add_team_models_to_all_models( for team_object in team_db_objects_typed: if ( - len(team_object.models) == 0 # empty list = all model access + not team_object.models # None or empty list = all model access or SpecialModelNames.all_proxy_models.value in team_object.models ): model_list = llm_router.get_model_list() @@ -9075,6 +9268,70 @@ def _add_team_models_to_all_models( return team_models +async def _add_access_group_models_to_team_models( + team_db_objects_typed: List[LiteLLM_TeamTable], + llm_router: Router, + prisma_client: PrismaClient, + team_models: Dict[str, Set[str]], +) -> Dict[str, Set[str]]: + """ + Resolve models reachable via team access groups and merge them into team_models. + + Batch-fetches all distinct access groups in a single DB query, then resolves + each eligible team's access group models via the pre-fetched map. + + This ensures models associated with a team only through access groups + (not directly in team.models) are included in the UI model listing. + """ + # First pass: identify eligible teams and collect all distinct access group IDs + eligible_teams: List[LiteLLM_TeamTable] = [] + all_access_group_ids: Set[str] = set() + + for team_object in team_db_objects_typed: + if not team_object.access_group_ids: + continue + + # Skip teams with empty models list — they already have access to everything + # (handled by _add_team_models_to_all_models) + if ( + not team_object.models + or SpecialModelNames.all_proxy_models.value in team_object.models + ): + continue + + eligible_teams.append(team_object) + all_access_group_ids.update(team_object.access_group_ids) + + if not eligible_teams: + return team_models + + # Single batch fetch for all access groups + access_group_rows = await prisma_client.db.litellm_accessgrouptable.find_many( + where={"access_group_id": {"in": list(all_access_group_ids)}} + ) + ag_model_map: Dict[str, List[str]] = { + row.access_group_id: row.access_model_names or [] for row in access_group_rows + } + + # Second pass: resolve deployments for each eligible team + for team_object in eligible_teams: + model_names: Set[str] = set() + for ag_id in team_object.access_group_ids or []: + model_names.update(ag_model_map.get(ag_id, [])) + + for model_name in model_names: + deployments = llm_router.get_model_list( + model_name=model_name, team_id=team_object.team_id + ) + if deployments is not None: + for deployment in deployments: + model_id = deployment.get("model_info", {}).get("id", None) + if model_id is not None: + team_models.setdefault(model_id, set()).add(team_object.team_id) + + return team_models + + async def get_all_team_models( user_teams: Union[List[str], Literal["*"]], prisma_client: PrismaClient, @@ -9111,6 +9368,14 @@ async def get_all_team_models( llm_router=llm_router, ) + # Also resolve models reachable via team access groups + team_models = await _add_access_group_models_to_team_models( + team_db_objects_typed=team_db_objects_typed, + llm_router=llm_router, + prisma_client=prisma_client, + team_models=team_models, + ) + # convert set to list returned_team_models: Dict[str, List[str]] = {} for model_id, team_ids in team_models.items(): @@ -9618,7 +9883,7 @@ async def _filter_models_by_team_id( team_accessible_model_ids: Set[str] = set() if ( - len(team_object.models) == 0 # empty list = all model access + not team_object.models # empty list = all model access or SpecialModelNames.all_proxy_models.value in team_object.models ): # Team has access to all models @@ -11504,9 +11769,9 @@ async def claim_onboarding_link(data: InvitationClaim): }, ) ### UPDATE USER OBJECT ### - hash_password = hash_token(token=data.password) + hashed_pw = hash_password(data.password) user_obj = await prisma_client.db.litellm_usertable.update( - where={"user_id": invite_obj.user_id}, data={"password": hash_password} + where={"user_id": invite_obj.user_id}, data={"password": hashed_pw} ) if user_obj is None: @@ -11526,6 +11791,8 @@ async def claim_onboarding_link(data: InvitationClaim): }, ) + if user_obj and hasattr(user_obj, "__dict__"): + user_obj.__dict__.pop("password", None) return user_obj @@ -11964,7 +12231,10 @@ async def invitation_delete( dependencies=[Depends(user_api_key_auth)], include_in_schema=False, ) -async def update_config(config_info: ConfigYAML): # noqa: PLR0915 +async def update_config( # noqa: PLR0915 + config_info: ConfigYAML, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): """ For Admin UI - allows admin to update config via UI @@ -11972,6 +12242,10 @@ async def update_config(config_info: ConfigYAML): # noqa: PLR0915 """ global llm_router, llm_model_list, general_settings, proxy_config, proxy_logging_obj, master_key, prisma_client try: + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, detail="Only proxy admins can update config" + ) import base64 """ @@ -12353,9 +12627,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 7c4073746ed..52833815248 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -170,11 +170,7 @@ from litellm.types.utils import ( ) from litellm.types.utils import ModelInfo from litellm.types.utils import ModelInfo as ModelMapInfo -from litellm.types.utils import ( - ModelResponseStream, - StandardLoggingPayload, - Usage, -) +from litellm.types.utils import ModelResponseStream, StandardLoggingPayload, Usage from litellm.utils import ( CustomStreamWrapper, EmbeddingResponse, @@ -408,9 +404,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 +454,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"] = {} @@ -655,12 +651,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 +2062,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 +4299,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( @@ -5387,11 +5386,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 +5422,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 +5648,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 +6769,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 +6831,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 +7496,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( @@ -9175,12 +9174,16 @@ class Router: cooldown_deployments = await _async_get_cooldown_deployments( litellm_router_instance=self, parent_otel_span=parent_otel_span ) - 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, ) + if not healthy_deployments and self.enable_health_check_routing: + 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, @@ -9613,10 +9616,16 @@ 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: + 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: @@ -9918,6 +9927,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 @@ -9947,6 +9962,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/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..7aca9b79653 --- /dev/null +++ b/tests/test_litellm/router_utils/test_health_check_allowed_fails_integration.py @@ -0,0 +1,532 @@ +""" +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 propagates exception objects through to unhealthy_endpoints.""" + + def test_unhealthy_endpoint_with_exception_dict(self): + """When health check returns {"error": ..., "exception": e}, exception should be in the endpoint.""" + from litellm.proxy.health_check import _clean_endpoint_data + + auth_error = litellm.AuthenticationError( + message="Invalid key", llm_provider="openai", model="gpt-4" + ) + + # Simulate what _perform_health_check does for an unhealthy dict result + is_healthy = {"error": "auth failed", "exception": auth_error} + litellm_params = {"model": "gpt-4", "api_key": "fake"} + cleaned = _clean_endpoint_data({**litellm_params, **is_healthy}, details=True) + # Exception should be preserved after cleaning + if "exception" in is_healthy: + cleaned["exception"] = is_healthy["exception"] + + assert cleaned["exception"] is auth_error + + def test_unhealthy_endpoint_raw_exception(self): + """When gather returns a raw Exception, it should be stored in the endpoint dict.""" + raw_exc = litellm.RateLimitError( + message="Rate limited", llm_provider="openai", model="gpt-4" + ) + + # Simulate the else branch in _perform_health_check + from litellm.proxy.health_check import _clean_endpoint_data + + litellm_params = {"model": "gpt-4"} + cleaned = _clean_endpoint_data(litellm_params, details=True) + if isinstance(raw_exc, Exception): + cleaned["exception"] = raw_exc + + assert cleaned["exception"] 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", "exception": timeout_exc}, + ] + 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, + ) + 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"}, # no "exception" key + ] + + 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, + ) + 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", "exception": rate_exc}, + ] + + 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, + ) + 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"