From cfda03ebe1229d5e0da1a48e9c23792172e76c52 Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Thu, 8 Jan 2026 14:55:36 -0300 Subject: [PATCH 1/3] fix(gemini): support snake_case for google_search tool parameters (#18451) * fix(gemini): support snake_case for google_search tool parameters Add snake_case aliases for Gemini tool names to match the pattern already used by other tools (url_context, google_maps, code_execution): - google_search -> googleSearch - google_search_retrieval -> googleSearchRetrieval - enterprise_web_search -> enterpriseWebSearch * test(gemini): add tests for snake_case google_search tool aliases * refactor(gemini): simplify get_tool_value calls formatting --- .../vertex_and_google_ai_studio_gemini.py | 27 ++++---- .../vertex_ai/gemini/test_transformation.py | 64 ++++++++++++++++++- 2 files changed, 77 insertions(+), 14 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index ba1788a217..3c93b1943e 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -480,20 +480,21 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): or tool_name == VertexToolName.CODE_EXECUTION.value ): # code_execution maintained for backwards compatibility code_execution = self.get_tool_value(tool, "codeExecution") - elif tool_name and tool_name == VertexToolName.GOOGLE_SEARCH.value: - googleSearch = self.get_tool_value( - tool, VertexToolName.GOOGLE_SEARCH.value - ) - elif ( - tool_name and tool_name == VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value + elif tool_name and ( + tool_name == VertexToolName.GOOGLE_SEARCH.value + or tool_name == "google_search" ): - googleSearchRetrieval = self.get_tool_value( - tool, VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value - ) - elif tool_name and tool_name == VertexToolName.ENTERPRISE_WEB_SEARCH.value: - enterpriseWebSearch = self.get_tool_value( - tool, VertexToolName.ENTERPRISE_WEB_SEARCH.value - ) + googleSearch = self.get_tool_value(tool, tool_name) + elif tool_name and ( + tool_name == VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value + or tool_name == "google_search_retrieval" + ): + googleSearchRetrieval = self.get_tool_value(tool, tool_name) + elif tool_name and ( + tool_name == VertexToolName.ENTERPRISE_WEB_SEARCH.value + or tool_name == "enterprise_web_search" + ): + enterpriseWebSearch = self.get_tool_value(tool, tool_name) elif tool_name and ( tool_name == VertexToolName.URL_CONTEXT.value or tool_name == "urlContext" diff --git a/tests/litellm/llms/vertex_ai/gemini/test_transformation.py b/tests/litellm/llms/vertex_ai/gemini/test_transformation.py index 6d005af28a..20f48b6f39 100644 --- a/tests/litellm/llms/vertex_ai/gemini/test_transformation.py +++ b/tests/litellm/llms/vertex_ai/gemini/test_transformation.py @@ -7,6 +7,7 @@ sys.path.insert( 0, os.path.abspath("../../../../..") ) # Adds the parent directory to the system path from litellm.llms.vertex_ai.gemini import transformation +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig from litellm.types.llms import openai from litellm.types import completion from litellm.types.llms.vertex_ai import RequestBody @@ -225,4 +226,65 @@ async def test__transform_request_body_image_config_with_image_size(): assert "generationConfig" in rb assert "imageConfig" in rb["generationConfig"] assert rb["generationConfig"]["imageConfig"]["aspectRatio"] == "16:9" - assert rb["generationConfig"]["imageConfig"]["imageSize"] == "4K" \ No newline at end of file + assert rb["generationConfig"]["imageConfig"]["imageSize"] == "4K" + + +def test_map_function_google_search_snake_case(): + """ + Test that google_search tool (snake_case) is properly mapped to googleSearch. + Fixes issue where tools=[{"google_search": {}}] was being stripped. + """ + config = VertexGeminiConfig() + optional_params = {} + + # Test snake_case google_search + tools = [{"google_search": {}}] + result = config._map_function(tools, optional_params) + + assert len(result) == 1 + assert "googleSearch" in result[0] + assert result[0]["googleSearch"] == {} + + +def test_map_function_google_search_camel_case(): + """ + Test that googleSearch tool (camelCase) still works. + """ + config = VertexGeminiConfig() + optional_params = {} + + # Test camelCase googleSearch + tools = [{"googleSearch": {}}] + result = config._map_function(tools, optional_params) + + assert len(result) == 1 + assert "googleSearch" in result[0] + assert result[0]["googleSearch"] == {} + + +def test_map_function_google_search_retrieval_snake_case(): + """ + Test that google_search_retrieval tool (snake_case) is properly mapped. + """ + config = VertexGeminiConfig() + optional_params = {} + + tools = [{"google_search_retrieval": {"dynamic_retrieval_config": {"mode": "MODE_DYNAMIC"}}}] + result = config._map_function(tools, optional_params) + + assert len(result) == 1 + assert "googleSearchRetrieval" in result[0] + + +def test_map_function_enterprise_web_search_snake_case(): + """ + Test that enterprise_web_search tool (snake_case) is properly mapped. + """ + config = VertexGeminiConfig() + optional_params = {} + + tools = [{"enterprise_web_search": {}}] + result = config._map_function(tools, optional_params) + + assert len(result) == 1 + assert "enterpriseWebSearch" in result[0] \ No newline at end of file From 3ebec39b740395efe6d97b175fc8119006a67607 Mon Sep 17 00:00:00 2001 From: Constantine Date: Thu, 8 Jan 2026 20:56:46 +0300 Subject: [PATCH 2/3] fix(proxy): use async anthropic client to prevent event loop blocking (#18435) Fixes #16716. Previously, synchronous Anthropic client was used for token counting, which blocked the event loop. This change switches to AsyncAnthropic and caches the client instance. --- litellm/proxy/utils.py | 10 ++++-- tests/test_litellm/test_utils_custom.py | 42 +++++++++++++++++++++++++ 2 files changed, 50 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/test_utils_custom.py diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d1a78534da..bd44cef954 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -131,6 +131,7 @@ else: unified_guardrail = UnifiedLLMGuardrails() +_anthropic_async_clients = {} def print_verbose(print_statement): """ @@ -4254,11 +4255,16 @@ async def count_tokens_with_anthropic_api( if anthropic_api_key and messages: # Call Anthropic API directly for more accurate token counting - client = anthropic.Anthropic(api_key=anthropic_api_key) + + # Use cached client if available to avoid socket exhaustion + if anthropic_api_key not in _anthropic_async_clients: + _anthropic_async_clients[anthropic_api_key] = anthropic.AsyncAnthropic(api_key=anthropic_api_key) + + client = _anthropic_async_clients[anthropic_api_key] # Call with explicit parameters to satisfy type checking # Type ignore for now since messages come from generic dict input - response = client.beta.messages.count_tokens( + response = await client.beta.messages.count_tokens( model=model_to_use, messages=messages, # type: ignore betas=["token-counting-2024-11-01"], diff --git a/tests/test_litellm/test_utils_custom.py b/tests/test_litellm/test_utils_custom.py new file mode 100644 index 0000000000..292da4132b --- /dev/null +++ b/tests/test_litellm/test_utils_custom.py @@ -0,0 +1,42 @@ +import pytest +from unittest.mock import MagicMock, patch, AsyncMock +from litellm.proxy.utils import count_tokens_with_anthropic_api, _anthropic_async_clients + +@pytest.mark.asyncio +async def test_count_tokens_caching(): + """ + Test that count_tokens_with_anthropic_api caches the client. + """ + # Clear cache + _anthropic_async_clients.clear() + + api_key = "sk-ant-test-key" + messages = [{"role": "user", "content": "hello"}] + model = "claude-3-opus-20240229" + + # Mock anthropic + with patch("anthropic.AsyncAnthropic") as mock_cls: + mock_client = MagicMock() + mock_cls.return_value = mock_client + + # Mock response + mock_response = MagicMock() + mock_response.input_tokens = 10 + + # Setup async return for count_tokens + mock_client.beta.messages.count_tokens = AsyncMock(return_value=mock_response) + + # First call + with patch.dict("os.environ", {"ANTHROPIC_API_KEY": api_key}): + await count_tokens_with_anthropic_api(model, messages) + + assert api_key in _anthropic_async_clients + assert _anthropic_async_clients[api_key] == mock_client + mock_cls.assert_called_once() # Should be called once + + # Second call + with patch.dict("os.environ", {"ANTHROPIC_API_KEY": api_key}): + await count_tokens_with_anthropic_api(model, messages) + + # Should still be called once (cached) + mock_cls.assert_called_once() From 516e4f8b9652cb6199b5e504de97835a51f07592 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Thu, 8 Jan 2026 23:53:36 +0530 Subject: [PATCH 3/3] fix: proactive RDS IAM token refresh to prevent 15-min connection failed (#18795) * fix: proactive RDS IAM token refresh to prevent 15-min connection failures (#16220) * fix: add noqa for PLR0915 in proxy_startup_event --- litellm/proxy/db/prisma_client.py | 309 +++++++++++++++--- litellm/proxy/proxy_server.py | 97 +++--- .../proxy/db/test_rds_iam_token_expiry.py | 275 ++++++++++++++++ 3 files changed, 602 insertions(+), 79 deletions(-) create mode 100644 tests/test_litellm/proxy/db/test_rds_iam_token_expiry.py diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 406ddceabf..c9c0cfe8f6 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -17,50 +17,141 @@ from litellm.secret_managers.main import str_to_bool class PrismaWrapper: + """ + Wrapper around Prisma client that handles RDS IAM token authentication. + + When iam_token_db_auth is enabled, this wrapper: + 1. Proactively refreshes IAM tokens before they expire (background task) + 2. Falls back to synchronous refresh if a token is found expired + 3. Uses proper locking to prevent race conditions during reconnection + + RDS IAM tokens are valid for 15 minutes. This wrapper refreshes them + 3 minutes before expiration to ensure uninterrupted database connectivity. + """ + + # Buffer time in seconds before token expiration to trigger refresh + # Refresh 3 minutes (180 seconds) before the token expires + TOKEN_REFRESH_BUFFER_SECONDS = 180 + + # Fallback refresh interval if token parsing fails (10 minutes) + FALLBACK_REFRESH_INTERVAL_SECONDS = 600 + def __init__(self, original_prisma: Any, iam_token_db_auth: bool): self._original_prisma = original_prisma self.iam_token_db_auth = iam_token_db_auth + # Background token refresh task management + self._token_refresh_task: Optional[asyncio.Task] = None + self._reconnection_lock = asyncio.Lock() + self._last_refresh_time: Optional[datetime] = None + + def _extract_token_from_db_url(self, db_url: Optional[str]) -> Optional[str]: + """ + Extract the token (password) from the DATABASE_URL. + + The token contains the AWS signature with X-Amz-Date and X-Amz-Expires parameters. + + Important: We must parse the URL while it's still encoded to preserve structure, + then decode the password portion. Otherwise the '?' in the token breaks URL parsing. + """ + if db_url is None: + return None + try: + # Parse URL while still encoded to preserve structure + parsed = urllib.parse.urlparse(db_url) + if parsed.password: + # Now decode just the password/token + return urllib.parse.unquote(parsed.password) + return None + except Exception: + return None + + def _parse_token_expiration(self, token: Optional[str]) -> Optional[datetime]: + """ + Parse the token to extract its expiration time. + + Returns the datetime when the token expires, or None if parsing fails. + """ + if token is None: + return None + + try: + # Token format: ...?X-Amz-Date=YYYYMMDDTHHMMSSZ&X-Amz-Expires=900&... + if "?" not in token: + return None + + query_string = token.split("?", 1)[1] + params = urllib.parse.parse_qs(query_string) + + expires_str = params.get("X-Amz-Expires", [None])[0] + date_str = params.get("X-Amz-Date", [None])[0] + + if not expires_str or not date_str: + return None + + token_created = datetime.strptime(date_str, "%Y%m%dT%H%M%SZ") + expires_in = int(expires_str) + + return token_created + timedelta(seconds=expires_in) + except Exception as e: + verbose_proxy_logger.debug(f"Failed to parse token expiration: {e}") + return None + + def _calculate_seconds_until_refresh(self) -> float: + """ + Calculate exactly how many seconds until we need to refresh the token. + + Uses precise timing: sleeps until (token_expiration - buffer_seconds). + For a 15-minute (900s) token with 180s buffer, this returns ~720s (12 min). + + Returns: + Number of seconds to sleep before the next refresh. + Returns 0 if token should be refreshed immediately. + Returns FALLBACK_REFRESH_INTERVAL_SECONDS if parsing fails. + """ + db_url = os.getenv("DATABASE_URL") + token = self._extract_token_from_db_url(db_url) + expiration_time = self._parse_token_expiration(token) + + if expiration_time is None: + # If we can't parse the token, use fallback interval + verbose_proxy_logger.debug( + f"Could not parse token expiration, using fallback interval of " + f"{self.FALLBACK_REFRESH_INTERVAL_SECONDS}s" + ) + return self.FALLBACK_REFRESH_INTERVAL_SECONDS + + # Calculate when we should refresh (expiration - buffer) + refresh_at = expiration_time - timedelta( + seconds=self.TOKEN_REFRESH_BUFFER_SECONDS + ) + + # How long until refresh time? + now = datetime.utcnow() + seconds_until_refresh = (refresh_at - now).total_seconds() + + # If already past refresh time, return 0 (refresh immediately) + return max(0, seconds_until_refresh) + def is_token_expired(self, token_url: Optional[str]) -> bool: + """Check if the token in the given URL is expired.""" if token_url is None: return True - # Decode the token URL to handle URL-encoded characters - decoded_url = urllib.parse.unquote(token_url) - # Parse the token URL - parsed_url = urllib.parse.urlparse(decoded_url) + token = self._extract_token_from_db_url(token_url) + expiration_time = self._parse_token_expiration(token) - # Parse the query parameters from the path component (if they exist there) - query_params = urllib.parse.parse_qs(parsed_url.query) + if expiration_time is None: + # If we can't parse the token, assume it's expired to trigger refresh + verbose_proxy_logger.debug( + "Could not parse token expiration, treating as expired" + ) + return True - # Get expiration time from the query parameters - expires = query_params.get("X-Amz-Expires", [None])[0] - if expires is None: - raise ValueError("X-Amz-Expires parameter is missing or invalid.") - - expires_int = int(expires) - - # Get the token's creation time from the X-Amz-Date parameter - token_time_str = query_params.get("X-Amz-Date", [""])[0] - if not token_time_str: - raise ValueError("X-Amz-Date parameter is missing or invalid.") - - # Ensure the token time string is parsed correctly - try: - token_time = datetime.strptime(token_time_str, "%Y%m%dT%H%M%SZ") - except ValueError as e: - raise ValueError(f"Invalid X-Amz-Date format: {e}") - - # Calculate the expiration time - expiration_time = token_time + timedelta(seconds=expires_int) - - # Current time in UTC - current_time = datetime.utcnow() - - # Check if the token is expired - return current_time > expiration_time + return datetime.utcnow() > expiration_time def get_rds_iam_token(self) -> Optional[str]: + """Generate a new RDS IAM token and update DATABASE_URL.""" if self.iam_token_db_auth: from litellm.proxy.auth.rds_iam_token import generate_iam_auth_token @@ -74,7 +165,6 @@ class PrismaWrapper: db_host=db_host, db_port=db_port, db_user=db_user ) - # print(f"token: {token}") _db_url = f"postgresql://{db_user}:{token}@{db_host}:{db_port}/{db_name}" if db_schema: _db_url += f"?schema={db_schema}" @@ -86,6 +176,7 @@ class PrismaWrapper: async def recreate_prisma_client( self, new_db_url: str, http_client: Optional[Any] = None ): + """Disconnect and reconnect the Prisma client with a new database URL.""" from prisma import Prisma # type: ignore try: @@ -100,21 +191,159 @@ class PrismaWrapper: await self._original_prisma.connect() + async def start_token_refresh_task(self) -> None: + """ + Start the background token refresh task. + + This task proactively refreshes RDS IAM tokens before they expire, + preventing connection failures. Should be called after the initial + Prisma client connection is established. + """ + if not self.iam_token_db_auth: + verbose_proxy_logger.debug( + "IAM token auth not enabled, skipping token refresh task" + ) + return + + if self._token_refresh_task is not None: + verbose_proxy_logger.debug("Token refresh task already running") + return + + self._token_refresh_task = asyncio.create_task(self._token_refresh_loop()) + verbose_proxy_logger.info( + "Started RDS IAM token proactive refresh background task" + ) + + async def stop_token_refresh_task(self) -> None: + """ + Stop the background token refresh task gracefully. + + Should be called during application shutdown to clean up resources. + """ + if self._token_refresh_task is None: + return + + self._token_refresh_task.cancel() + try: + await self._token_refresh_task + except asyncio.CancelledError: + pass + self._token_refresh_task = None + verbose_proxy_logger.info("Stopped RDS IAM token refresh background task") + + async def _token_refresh_loop(self) -> None: + """ + Background loop that proactively refreshes RDS IAM tokens before expiration. + + Uses precise timing: calculates the exact sleep duration until the token + needs to be refreshed (expiration - 3 minute buffer), then refreshes. + This is more efficient than polling, requiring only 1 wake-up per token cycle. + """ + verbose_proxy_logger.info( + f"RDS IAM token refresh loop started. " + f"Tokens will be refreshed {self.TOKEN_REFRESH_BUFFER_SECONDS}s before expiration." + ) + + while True: + try: + # Calculate exactly how long to sleep until next refresh + sleep_seconds = self._calculate_seconds_until_refresh() + + if sleep_seconds > 0: + verbose_proxy_logger.info( + f"RDS IAM token refresh scheduled in {sleep_seconds:.0f} seconds " + f"({sleep_seconds / 60:.1f} minutes)" + ) + await asyncio.sleep(sleep_seconds) + + # Refresh the token + verbose_proxy_logger.info("Proactively refreshing RDS IAM token...") + await self._safe_refresh_token() + + except asyncio.CancelledError: + verbose_proxy_logger.info("RDS IAM token refresh loop cancelled") + break + except Exception as e: + verbose_proxy_logger.error( + f"Error in RDS IAM token refresh loop: {e}. " + f"Retrying in {self.FALLBACK_REFRESH_INTERVAL_SECONDS}s..." + ) + # On error, wait before retrying to avoid tight error loops + try: + await asyncio.sleep(self.FALLBACK_REFRESH_INTERVAL_SECONDS) + except asyncio.CancelledError: + break + + async def _safe_refresh_token(self) -> None: + """ + Refresh the RDS IAM token with proper locking to prevent race conditions. + + Uses an asyncio lock to ensure only one refresh operation happens at a time, + preventing multiple concurrent reconnection attempts. + """ + async with self._reconnection_lock: + new_db_url = self.get_rds_iam_token() + if new_db_url: + await self.recreate_prisma_client(new_db_url) + self._last_refresh_time = datetime.utcnow() + verbose_proxy_logger.info( + "RDS IAM token refreshed successfully. New token valid for ~15 minutes." + ) + else: + verbose_proxy_logger.error( + "Failed to generate new RDS IAM token during proactive refresh" + ) + def __getattr__(self, name: str): + """ + Proxy attribute access to the underlying Prisma client. + + If IAM token auth is enabled and the token is expired, this method + provides a synchronous fallback to refresh the token. However, this + should rarely be needed since the background task proactively refreshes + tokens before they expire. + + FIXED: Now properly waits for reconnection to complete before returning, + instead of the previous fire-and-forget pattern that caused the bug. + """ original_attr = getattr(self._original_prisma, name) + if self.iam_token_db_auth: db_url = os.getenv("DATABASE_URL") - if self.is_token_expired(db_url): - db_url = self.get_rds_iam_token() - loop = asyncio.get_event_loop() - if db_url: + # Check if token is expired (should be rare if background task is running) + if self.is_token_expired(db_url): + verbose_proxy_logger.warning( + "RDS IAM token expired in __getattr__ - proactive refresh may have failed. " + "Triggering synchronous fallback refresh..." + ) + + new_db_url = self.get_rds_iam_token() + if new_db_url: + loop = asyncio.get_event_loop() + if loop.is_running(): - asyncio.run_coroutine_threadsafe( - self.recreate_prisma_client(db_url), loop + # FIXED: Actually wait for the reconnection to complete! + # The previous code used fire-and-forget which caused the bug. + future = asyncio.run_coroutine_threadsafe( + self.recreate_prisma_client(new_db_url), loop ) + try: + # Wait up to 30 seconds for reconnection + future.result(timeout=30) + verbose_proxy_logger.info( + "Synchronous token refresh completed successfully" + ) + except Exception as e: + verbose_proxy_logger.error( + f"Failed to refresh token synchronously: {e}" + ) + raise else: - asyncio.run(self.recreate_prisma_client(db_url)) + asyncio.run(self.recreate_prisma_client(new_db_url)) + + # Get the NEW attribute after reconnection + original_attr = getattr(self._original_prisma, name) else: raise ValueError("Failed to get RDS IAM token") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 06525e3913..a56a6379cb 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -533,9 +533,9 @@ except ImportError: server_root_path = os.getenv("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" ) @@ -658,7 +658,7 @@ async def _initialize_shared_aiohttp_session(): @asynccontextmanager -async def proxy_startup_event(app: FastAPI): +async def proxy_startup_event(app: FastAPI): # noqa: PLR0915 global prisma_client, master_key, use_background_health_checks, llm_router, llm_model_list, general_settings, proxy_budget_rescheduler_min_time, proxy_budget_rescheduler_max_time, litellm_proxy_admin_name, db_writer_client, store_model_in_db, premium_user, _license_check, proxy_batch_polling_interval, shared_aiohttp_session import json @@ -788,6 +788,17 @@ async def proxy_startup_event(app: FastAPI): except Exception as e: verbose_proxy_logger.error(f"Error closing shared aiohttp session: {e}") + # Shutdown event - stop RDS IAM token refresh background task + if ( + prisma_client is not None + and hasattr(prisma_client, "db") + and hasattr(prisma_client.db, "stop_token_refresh_task") + ): + try: + await prisma_client.db.stop_token_refresh_task() + except Exception as e: + verbose_proxy_logger.error(f"Error stopping token refresh task: {e}") + await proxy_shutdown_event() # type: ignore[reportGeneralTypeIssues] @@ -1083,9 +1094,7 @@ try: # In non-root Docker, we restructure in /var/lib/litellm/ui. try: _restructure_ui_html_files(ui_path) - verbose_proxy_logger.info( - f"Restructured UI directory: {ui_path}" - ) + verbose_proxy_logger.info(f"Restructured UI directory: {ui_path}") except PermissionError as e: verbose_proxy_logger.exception( f"Permission error while restructuring UI directory {ui_path}: {e}" @@ -1171,9 +1180,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 ) @@ -1181,9 +1190,9 @@ 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 polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache user_custom_auth = None @@ -1522,9 +1531,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 @@ -1876,7 +1885,6 @@ class ProxyConfig: "environment_variables" in config_to_save and config_to_save["environment_variables"] ): - # decrypt the environment_variables - in case a caller function has already encrypted the environment_variables decrypted_env_vars = self._decrypt_and_set_db_env_variables( environment_variables=config_to_save["environment_variables"], @@ -2794,21 +2802,21 @@ class ProxyConfig: verbose_proxy_logger.debug(f"_alerting_callbacks: {general_settings}") if _alerting_callbacks is None: return - + # Ensure proxy_logging_obj.alerting is set for all alerting types _alerting_value = general_settings.get("alerting", None) - verbose_proxy_logger.debug(f"_load_alerting_settings: Calling update_values with alerting={_alerting_value}") + verbose_proxy_logger.debug( + f"_load_alerting_settings: Calling update_values with alerting={_alerting_value}" + ) proxy_logging_obj.update_values( alerting=_alerting_value, alerting_threshold=general_settings.get("alerting_threshold", 600), alert_types=general_settings.get("alert_types", None), - alert_to_webhook_url=general_settings.get( - "alert_to_webhook_url", None - ), + alert_to_webhook_url=general_settings.get("alert_to_webhook_url", None), alerting_args=general_settings.get("alerting_args", None), redis_cache=redis_usage_cache, ) - + for _alert in _alerting_callbacks: if _alert == "slack": # [OLD] v0 implementation - already handled by update_values above @@ -3279,7 +3287,7 @@ class ProxyConfig: proxy_logging_obj: ProxyLogging """ _general_settings = config_data.get("general_settings", {}) - + if _general_settings is not None and "alerting" in _general_settings: if ( general_settings is not None @@ -3294,7 +3302,8 @@ class ProxyConfig: _merged_alerting = list(_yaml_alerting.union(_db_alerting)) # Preserve order: YAML values first, then DB values _merged_alerting = list(general_settings["alerting"]) + [ - item for item in _general_settings["alerting"] + item + for item in _general_settings["alerting"] if item not in general_settings["alerting"] ] verbose_proxy_logger.debug( @@ -3605,7 +3614,6 @@ class ProxyConfig: await self._init_vector_stores_in_db(prisma_client=prisma_client) if self._should_load_db_object(object_type="vector_store_indexes"): - await self._init_vector_store_indexes_in_db(prisma_client=prisma_client) if self._should_load_db_object(object_type="mcp"): @@ -3804,10 +3812,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) @@ -4134,9 +4142,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 @@ -4654,10 +4662,14 @@ class ProxyStartupEvent: replace_existing=True, misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) - verbose_proxy_logger.info("Responses cost check job scheduled successfully") + verbose_proxy_logger.info( + "Responses cost check job scheduled successfully" + ) except Exception as e: - verbose_proxy_logger.debug(f"Failed to setup responses cost checking: {e}") + verbose_proxy_logger.debug( + f"Failed to setup responses cost checking: {e}" + ) verbose_proxy_logger.debug( "Checking responses cost for LiteLLM Managed Files is an Enterprise Feature. Skipping..." ) @@ -4826,6 +4838,14 @@ class ProxyStartupEvent: await prisma_client.connect() + ## Start RDS IAM token refresh background task if enabled ## + # This proactively refreshes IAM tokens before they expire, + # preventing the 15-minute connection failure bug (#16220) + if hasattr(prisma_client, "db") and hasattr( + prisma_client.db, "start_token_refresh_task" + ): + await prisma_client.db.start_token_refresh_task() + ## Add necessary views to proxy ## asyncio.create_task( prisma_client.check_view_exists() @@ -5937,7 +5957,6 @@ async def realtime_websocket_endpoint( ), user_api_key_dict=Depends(user_api_key_auth_websocket), ): - await websocket.accept() # Only use explicit parameters, not all query params @@ -9525,9 +9544,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/tests/test_litellm/proxy/db/test_rds_iam_token_expiry.py b/tests/test_litellm/proxy/db/test_rds_iam_token_expiry.py new file mode 100644 index 0000000000..1492acb079 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_rds_iam_token_expiry.py @@ -0,0 +1,275 @@ +""" +Tests for the RDS IAM token proactive refresh implementation. + +Tests for GitHub Issue #16220: RDS IAM authentication connection failures after 15 minutes. + +The fix implements: +1. Proactive background token refresh (refreshes 3 min before expiration) +2. Precise sleep timing (1 wake-up per token cycle instead of polling) +3. Proper locking during reconnection +4. Fixed __getattr__ fallback that now waits for reconnection + +Run these tests: + poetry run pytest tests/test_litellm/proxy/db/test_rds_iam_token_expiry.py -v -s +""" + +import asyncio +import os +import urllib.parse +from datetime import datetime, timedelta +from unittest.mock import MagicMock, patch + +import pytest + + +class TestPrismaWrapperTokenRefresh: + """Tests for the PrismaWrapper RDS IAM token refresh implementation.""" + + @pytest.fixture + def setup_env(self): + """Setup environment variables for testing.""" + os.environ["DATABASE_HOST"] = "test-host.rds.amazonaws.com" + os.environ["DATABASE_PORT"] = "5432" + os.environ["DATABASE_USER"] = "test_user" + os.environ["DATABASE_NAME"] = "test_db" + os.environ["IAM_TOKEN_DB_AUTH"] = "True" + yield + # Cleanup + for key in [ + "DATABASE_HOST", + "DATABASE_PORT", + "DATABASE_USER", + "DATABASE_NAME", + "DATABASE_URL", + "IAM_TOKEN_DB_AUTH", + "DATABASE_SCHEMA", + ]: + os.environ.pop(key, None) + + def _generate_mock_token(self, expires_in_seconds: int = 900) -> str: + """Generate a mock IAM token with expiration info.""" + now = datetime.utcnow() + date_str = now.strftime("%Y%m%dT%H%M%SZ") + # Build the token like AWS does + token = f"mock-token?X-Amz-Date={date_str}&X-Amz-Expires={expires_in_seconds}&X-Amz-Signature=abc123" + return urllib.parse.quote(token, safe="") + + def _set_database_url_with_token(self, expires_in_seconds: int = 900): + """Set DATABASE_URL with a mock token.""" + token = self._generate_mock_token(expires_in_seconds) + os.environ[ + "DATABASE_URL" + ] = f"postgresql://test_user:{token}@test-host:5432/test_db" + + @pytest.mark.asyncio + async def test_is_token_expired_fresh(self, setup_env): + """Test that fresh token is not detected as expired.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + self._set_database_url_with_token(expires_in_seconds=900) + db_url = os.getenv("DATABASE_URL") + + assert wrapper.is_token_expired(db_url) is False + + @pytest.mark.asyncio + async def test_is_token_expired_old(self, setup_env): + """Test that old token is detected as expired.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + # Create an expired token + old_date = datetime.utcnow() - timedelta(seconds=901) + date_str = old_date.strftime("%Y%m%dT%H%M%SZ") + token = ( + f"mock-token?X-Amz-Date={date_str}&X-Amz-Expires=900&X-Amz-Signature=abc" + ) + encoded_token = urllib.parse.quote(token, safe="") + db_url = f"postgresql://test_user:{encoded_token}@test-host:5432/test_db" + + assert wrapper.is_token_expired(db_url) is True + + @pytest.mark.asyncio + async def test_start_stop_token_refresh_task(self, setup_env): + """Test that token refresh task starts and stops correctly.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + # Set a valid token + self._set_database_url_with_token(expires_in_seconds=900) + + # Start the task + await wrapper.start_token_refresh_task() + assert wrapper._token_refresh_task is not None + assert not wrapper._token_refresh_task.done() + + # Stop the task + await wrapper.stop_token_refresh_task() + assert wrapper._token_refresh_task is None + + @pytest.mark.asyncio + async def test_start_task_not_enabled(self, setup_env): + """Test that task doesn't start when IAM auth is not enabled.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + # IAM auth disabled + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=False) + + await wrapper.start_token_refresh_task() + assert wrapper._token_refresh_task is None + + @pytest.mark.asyncio + async def test_is_token_expired_null(self, setup_env): + """Test that None token is treated as expired.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + assert wrapper.is_token_expired(None) is True + + +class TestTokenExpirationParsing: + """Tests for token expiration parsing utilities.""" + + def test_parse_token_expiration_valid(self): + """Test parsing expiration from a valid token.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + # Create a token with known expiration + token = "mock-token?X-Amz-Date=20240101T120000Z&X-Amz-Expires=900&X-Amz-Signature=abc" + + expiration = wrapper._parse_token_expiration(token) + + assert expiration is not None + expected = datetime(2024, 1, 1, 12, 0, 0) + timedelta(seconds=900) + assert expiration == expected + + def test_parse_token_expiration_invalid(self): + """Test that invalid token returns None.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + # Invalid tokens + assert wrapper._parse_token_expiration(None) is None + assert wrapper._parse_token_expiration("no-query-params") is None + assert wrapper._parse_token_expiration("?missing=params") is None + + +class TestBackgroundRefreshLoop: + """Tests for the background refresh loop timing.""" + + @pytest.fixture + def setup_env(self): + """Setup environment variables for testing.""" + os.environ["DATABASE_HOST"] = "test-host.rds.amazonaws.com" + os.environ["DATABASE_PORT"] = "5432" + os.environ["DATABASE_USER"] = "test_user" + os.environ["DATABASE_NAME"] = "test_db" + yield + # Cleanup + for key in [ + "DATABASE_HOST", + "DATABASE_PORT", + "DATABASE_USER", + "DATABASE_NAME", + "DATABASE_URL", + ]: + os.environ.pop(key, None) + + @pytest.mark.asyncio + async def test_calculate_seconds_fallback_when_no_url(self, setup_env): + """Test that fallback is used when DATABASE_URL is not set.""" + from litellm.proxy.db.prisma_client import PrismaWrapper + + mock_prisma = MagicMock() + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + # Don't set DATABASE_URL + seconds = wrapper._calculate_seconds_until_refresh() + + # Should return fallback interval + assert seconds == wrapper.FALLBACK_REFRESH_INTERVAL_SECONDS + + +# ============================================================================ +# DEMONSTRATION SCRIPT +# ============================================================================ + + +async def demonstrate_fix(): + """ + Demonstrates the fix for the RDS IAM token expiration bug. + + Shows how the proactive refresh prevents the 15-minute connection failure. + """ + # Import the actual implementation + try: + from litellm.proxy.db.prisma_client import PrismaWrapper + except ImportError: + return + + # Setup mock environment + os.environ["DATABASE_HOST"] = "mock-rds.region.rds.amazonaws.com" + os.environ["DATABASE_PORT"] = "5432" + os.environ["DATABASE_USER"] = "iam_user" + os.environ["DATABASE_NAME"] = "litellm" + + # Create initial token (expires in 10 seconds for demo) + now = datetime.utcnow() + date_str = now.strftime("%Y%m%dT%H%M%SZ") + token = f"mock-token?X-Amz-Date={date_str}&X-Amz-Expires=10&X-Amz-Signature=abc123" + encoded_token = urllib.parse.quote(token, safe="") + os.environ[ + "DATABASE_URL" + ] = f"postgresql://iam_user:{encoded_token}@mock-rds:5432/litellm" + + # Create mock prisma client + mock_prisma = MagicMock() + + wrapper = PrismaWrapper(original_prisma=mock_prisma, iam_token_db_auth=True) + + # Override buffer for faster demo + wrapper.TOKEN_REFRESH_BUFFER_SECONDS = 3 + wrapper.FALLBACK_REFRESH_INTERVAL_SECONDS = 5 + _ = wrapper._calculate_seconds_until_refresh() # Verify calculation works + db_url = os.getenv("DATABASE_URL") + is_expired = wrapper.is_token_expired(db_url) + assert is_expired is False, "Fresh token should not be expired!" + + # Mock the _token_refresh_loop to prevent it from actually running + async def mock_loop(): + try: + await asyncio.sleep(1000) + except asyncio.CancelledError: + pass + + with patch.object(wrapper, "_token_refresh_loop", side_effect=mock_loop): + await wrapper.start_token_refresh_task() + await wrapper.stop_token_refresh_task() + + # Cleanup + for key in [ + "DATABASE_HOST", + "DATABASE_PORT", + "DATABASE_USER", + "DATABASE_NAME", + "DATABASE_URL", + ]: + os.environ.pop(key, None) + + +if __name__ == "__main__": + asyncio.run(demonstrate_fix())