From 25f93bed918fbb35c44207b950632047593c40ed Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Fri, 3 Apr 2026 22:33:06 +0000 Subject: [PATCH 1/7] security: prevent API key leaks in error tracebacks, logs, and alerts Gemini API keys embedded in URLs as ?key= query parameters leak through httpx error tracebacks, which are then captured by traceback.format_exc() and forwarded to logging callbacks, Slack/Teams alerts, and HTTP client responses. Short-term: all httpx.HTTPStatusError handlers now raise MaskedHTTPStatusError(...) from None, which masks the URL and breaks exception chaining so the original error never appears in tracebacks. Long-term: moved all Gemini/Vertex URL constructions from ?key={api_key} to x-goog-api-key header (Google's documented auth method), so the key is never in the URL at all. WebSocket realtime is the only exception since WS clients cannot use custom headers. Additionally hardened all outbound credential paths: - WebSocket close reasons now pass through _redact_string() - Callback pipeline (failure_handler) redacts traceback_exception and error_str before forwarding to integrations (Langfuse, Datadog, etc.) - Slack/Teams alert messages redacted in send_llm_exception_alert, ProxyLogging.failure_handler, and post_call_failure_hook - HTTP error responses in proxy SSE and health endpoints redacted - Exception messages in exception_mapping_utils redacted - print_verbose() stdout output redacted when set_verbose=True - HTTPHandler.put() now has MaskedHTTPStatusError (was missing) --- litellm/_logging.py | 2 + .../exception_mapping_utils.py | 8 +- litellm/litellm_core_utils/litellm_logging.py | 10 +- litellm/llms/azure/realtime/handler.py | 4 +- litellm/llms/bedrock/realtime/handler.py | 4 +- litellm/llms/custom_httpx/http_handler.py | 125 +++++++++--------- litellm/llms/custom_httpx/llm_http_handler.py | 10 +- litellm/llms/gemini/common_utils.py | 5 +- litellm/llms/gemini/files/transformation.py | 6 +- .../gemini/interactions/transformation.py | 26 ++-- .../llms/gemini/realtime/transformation.py | 4 + .../gemini/vector_stores/transformation.py | 13 +- litellm/llms/openai/realtime/handler.py | 5 +- litellm/llms/vertex_ai/common_utils.py | 27 ++-- .../vertex_ai_context_caching.py | 14 +- litellm/llms/vertex_ai/vertex_llm_base.py | 3 +- litellm/main.py | 3 +- litellm/proxy/common_request_processing.py | 4 +- litellm/proxy/utils.py | 7 +- litellm/rag/ingestion/gemini_ingestion.py | 10 +- .../vertex_ai/test_gemini_batch_embeddings.py | 2 +- .../google_genai/test_google_genai_adapter.py | 118 +++++++---------- .../files/test_gemini_files_transformation.py | 26 ++-- 23 files changed, 217 insertions(+), 219 deletions(-) diff --git a/litellm/_logging.py b/litellm/_logging.py index 7824fcfa675..d072cc549d0 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -86,6 +86,8 @@ _SECRET_RE = _build_secret_patterns() def _redact_string(value: str) -> str: + if not _ENABLE_SECRET_REDACTION: + return value return _SECRET_RE.sub(_REDACTED, value) diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index bc54786420a..ef062ff47a3 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -6,7 +6,7 @@ from typing import Any, Optional import httpx import litellm -from litellm._logging import verbose_logger +from litellm._logging import _redact_string, verbose_logger from litellm.types.utils import LlmProviders from ..exceptions import ( @@ -2304,7 +2304,7 @@ def exception_type( # type: ignore # noqa: PLR0915 else: # if no status code then it is an APIConnectionError: https://github.com/openai/openai-python#handling-errors raise APIConnectionError( - message=f"{exception_provider} APIConnectionError - {message}\n{traceback.format_exc()}", + message=f"{exception_provider} APIConnectionError - {message}\n{_redact_string(traceback.format_exc())}", llm_provider="azure", model=model, litellm_debug_info=extra_information, @@ -2431,7 +2431,7 @@ def exception_type( # type: ignore # noqa: PLR0915 else: raise APIConnectionError( message="{}\n{}".format( - str(original_exception), traceback.format_exc() + str(original_exception), _redact_string(traceback.format_exc()) ), llm_provider=custom_llm_provider, model=model, @@ -2460,7 +2460,7 @@ def exception_type( # type: ignore # noqa: PLR0915 setattr(e, "litellm_response_headers", litellm_response_headers) raise e # it's already mapped raised_exc = APIConnectionError( - message="{}\n{}".format(original_exception, traceback.format_exc()), + message="{}\n{}".format(original_exception, _redact_string(traceback.format_exc())), llm_provider="", model="", ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e84c1e13a8b..455e651643a 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -36,7 +36,7 @@ from litellm import ( log_raw_request_response, turn_off_message_logging, ) -from litellm._logging import _is_debugging_on, verbose_logger +from litellm._logging import _is_debugging_on, _redact_string, verbose_logger from litellm._uuid import uuid from litellm.batches.batch_utils import _handle_completed_batch from litellm.caching.caching import DualCache, InMemoryCache @@ -2848,7 +2848,11 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["log_event_type"] = "failed_api_call" self.model_call_details["exception"] = exception - self.model_call_details["traceback_exception"] = traceback_exception + self.model_call_details["traceback_exception"] = ( + _redact_string(traceback_exception) + if isinstance(traceback_exception, str) + else traceback_exception + ) self.model_call_details["end_time"] = end_time self.model_call_details.setdefault("original_response", None) self.model_call_details["response_cost"] = 0 @@ -2871,7 +2875,7 @@ class Logging(LiteLLMLoggingBaseClass): end_time=end_time, logging_obj=self, status="failure", - error_str=str(exception), + error_str=_redact_string(str(exception)), original_exception=exception, standard_built_in_tools_params=self.standard_built_in_tools_params, ) diff --git a/litellm/llms/azure/realtime/handler.py b/litellm/llms/azure/realtime/handler.py index 6d00ecd51c9..1f3428f2ca5 100644 --- a/litellm/llms/azure/realtime/handler.py +++ b/litellm/llms/azure/realtime/handler.py @@ -6,7 +6,7 @@ This requires websockets, and is currently only supported on LiteLLM Proxy. from typing import Any, Optional, cast -from litellm._logging import verbose_proxy_logger +from litellm._logging import _redact_string, verbose_proxy_logger from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from ....litellm_core_utils.litellm_logging import Logging as LiteLLMLogging @@ -118,7 +118,7 @@ class AzureOpenAIRealtime(AzureChatCompletion): await realtime_streaming.bidirectional_forward() except websockets.exceptions.InvalidStatusCode as e: # type: ignore - await websocket.close(code=e.status_code, reason=str(e)) + await websocket.close(code=e.status_code, reason=_redact_string(str(e))) except Exception: verbose_proxy_logger.exception( "Error in AzureOpenAIRealtime.async_realtime" diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index cde9f3e6fce..8405ff500d7 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -8,7 +8,7 @@ import asyncio import json from typing import Any, Optional -from litellm._logging import verbose_proxy_logger +from litellm._logging import _redact_string, verbose_proxy_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from ..base_aws_llm import BaseAWSLLM @@ -152,7 +152,7 @@ class BedrockRealtime(BaseAWSLLM): f"Error in BedrockRealtime.async_realtime: {e}" ) try: - await websocket.close(code=1011, reason=f"Internal error: {str(e)}") + await websocket.close(code=1011, reason=_redact_string(f"Internal error: {str(e)}")) except Exception: pass raise diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 001547557d4..fdb05d1a91b 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -316,16 +316,65 @@ def mask_sensitive_info(error_message): return error_message +def _safe_get_response_text(response: httpx.Response) -> str: + """Safely read response text, falling back to empty string on decoding errors.""" + try: + return response.text + except Exception: + return "" + + +async def _safe_aread_response(response: httpx.Response) -> bytes: + """Safely read async response body, falling back to empty bytes on errors.""" + try: + return await response.aread() + except Exception: + return b"" + + +def _safe_read_response(response: httpx.Response) -> bytes: + """Safely read sync response body, falling back to empty bytes on errors.""" + try: + return response.read() + except Exception: + return b"" + + +def _raise_masked_sync_error(e: httpx.HTTPStatusError, stream: bool) -> None: + """Raise a MaskedHTTPStatusError for sync HTTP handlers.""" + if stream: + _body = mask_sensitive_info(_safe_read_response(e.response)) + raise MaskedHTTPStatusError(e, message=_body, text=_body) from None + _text = mask_sensitive_info(_safe_get_response_text(e.response)) + raise MaskedHTTPStatusError(e, message=_text, text=_text) from None + + +async def _raise_masked_async_error(e: httpx.HTTPStatusError, stream: bool) -> None: + """Raise a MaskedHTTPStatusError for async HTTP handlers.""" + if stream: + _body = await _safe_aread_response(e.response) + raise MaskedHTTPStatusError(e, message=_body, text=_body) from None + _text = mask_sensitive_info(_safe_get_response_text(e.response)) + raise MaskedHTTPStatusError(e, message=_text, text=_text) from None + + class MaskedHTTPStatusError(httpx.HTTPStatusError): def __init__( self, original_error, message: Optional[str] = None, text: Optional[str] = None ): # Create a new error with the masked URL masked_url = mask_sensitive_info(str(original_error.request.url)) - # Create a new error that looks like the original, but with a masked URL + # Mask the original exception message too (it contains the full URL) + masked_original_message = mask_sensitive_info(str(original_error)) + + # Safely access response content — decompression can fail (e.g. zlib error) + try: + response_content = original_error.response.content + except Exception: + response_content = b"" super().__init__( - message=original_error.message, + message=masked_original_message, request=httpx.Request( method=original_error.request.method, url=masked_url, @@ -334,12 +383,13 @@ class MaskedHTTPStatusError(httpx.HTTPStatusError): ), response=httpx.Response( status_code=original_error.response.status_code, - content=original_error.response.content, + content=response_content, headers=original_error.response.headers, ), ) self.message = message self.text = text + self.status_code = original_error.response.status_code class AsyncHTTPHandler: @@ -501,16 +551,7 @@ class AsyncHTTPHandler: headers=headers, ) except httpx.HTTPStatusError as e: - if stream is True: - setattr(e, "message", await e.response.aread()) - setattr(e, "text", await e.response.aread()) - else: - setattr(e, "message", mask_sensitive_info(e.response.text)) - setattr(e, "text", mask_sensitive_info(e.response.text)) - - setattr(e, "status_code", e.response.status_code) - - raise e + await _raise_masked_async_error(e, stream) except Exception as e: raise e @@ -571,12 +612,7 @@ class AsyncHTTPHandler: headers=headers, ) except httpx.HTTPStatusError as e: - setattr(e, "status_code", e.response.status_code) - if stream is True: - setattr(e, "message", await e.response.aread()) - else: - setattr(e, "message", e.response.text) - raise e + await _raise_masked_async_error(e, stream) except Exception as e: raise e @@ -637,12 +673,7 @@ class AsyncHTTPHandler: headers=headers, ) except httpx.HTTPStatusError as e: - setattr(e, "status_code", e.response.status_code) - if stream is True: - setattr(e, "message", await e.response.aread()) - else: - setattr(e, "message", e.response.text) - raise e + await _raise_masked_async_error(e, stream) except Exception as e: raise e @@ -690,12 +721,7 @@ class AsyncHTTPHandler: finally: await new_client.aclose() except httpx.HTTPStatusError as e: - setattr(e, "status_code", e.response.status_code) - if stream is True: - setattr(e, "message", await e.response.aread()) - else: - setattr(e, "message", e.response.text) - raise e + await _raise_masked_async_error(e, stream) except Exception as e: raise e @@ -1035,16 +1061,7 @@ class HTTPHandler: llm_provider="litellm-httpx-handler", ) except httpx.HTTPStatusError as e: - if stream is True: - setattr(e, "message", mask_sensitive_info(e.response.read())) - setattr(e, "text", mask_sensitive_info(e.response.read())) - else: - error_text = mask_sensitive_info(e.response.text) - setattr(e, "message", error_text) - setattr(e, "text", error_text) - - setattr(e, "status_code", e.response.status_code) - raise e + _raise_masked_sync_error(e, stream) except Exception as e: raise e @@ -1083,17 +1100,7 @@ class HTTPHandler: llm_provider="litellm-httpx-handler", ) except httpx.HTTPStatusError as e: - if stream is True: - setattr(e, "message", mask_sensitive_info(e.response.read())) - setattr(e, "text", mask_sensitive_info(e.response.read())) - else: - error_text = mask_sensitive_info(e.response.text) - setattr(e, "message", error_text) - setattr(e, "text", error_text) - - setattr(e, "status_code", e.response.status_code) - - raise e + _raise_masked_sync_error(e, stream) except Exception as e: raise e @@ -1130,6 +1137,8 @@ class HTTPHandler: model="default-model-name", llm_provider="litellm-httpx-handler", ) + except httpx.HTTPStatusError as e: + _raise_masked_sync_error(e, stream) except Exception as e: raise e @@ -1168,17 +1177,7 @@ class HTTPHandler: llm_provider="litellm-httpx-handler", ) except httpx.HTTPStatusError as e: - if stream is True: - setattr(e, "message", mask_sensitive_info(e.response.read())) - setattr(e, "text", mask_sensitive_info(e.response.read())) - else: - error_text = mask_sensitive_info(e.response.text) - setattr(e, "message", error_text) - setattr(e, "text", error_text) - - setattr(e, "status_code", e.response.status_code) - - raise e + _raise_masked_sync_error(e, stream) except Exception as e: raise e diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 7a8820a8785..d8ef7e74402 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -22,7 +22,7 @@ import litellm import litellm.litellm_core_utils import litellm.types import litellm.types.utils -from litellm._logging import verbose_logger +from litellm._logging import _redact_string, verbose_logger from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming @@ -4789,12 +4789,12 @@ class BaseLLMHTTPHandler: except websockets.exceptions.InvalidStatusCode as e: # type: ignore verbose_logger.exception(f"Error connecting to backend: {e}") - await websocket.close(code=e.status_code, reason=str(e)) + await websocket.close(code=e.status_code, reason=_redact_string(str(e))) except Exception as e: verbose_logger.exception(f"Error connecting to backend: {e}") try: await websocket.close( - code=1011, reason=f"Internal server error: {str(e)}" + code=1011, reason=_redact_string(f"Internal server error: {str(e)}") ) except RuntimeError as close_error: if "already completed" in str(close_error) or "websocket.close" in str( @@ -5076,12 +5076,12 @@ class BaseLLMHTTPHandler: except websockets.exceptions.InvalidStatusCode as e: # type: ignore verbose_logger.exception(f"Error connecting to responses WS backend: {e}") - await websocket.close(code=e.status_code, reason=str(e)) + await websocket.close(code=e.status_code, reason=_redact_string(str(e))) except Exception as e: verbose_logger.exception(f"Error in responses WS: {e}") try: await websocket.close( - code=1011, reason=f"Internal server error: {str(e)}" + code=1011, reason=_redact_string(f"Internal server error: {str(e)}") ) except RuntimeError as close_error: if "already completed" in str(close_error) or "websocket.close" in str( diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 87c107fab37..bc963d62b5f 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -28,7 +28,7 @@ class GeminiModelInfo(BaseLLMModelInfo): api_key: Optional[str] = None, api_base: Optional[str] = None, ) -> dict: - """Google AI Studio sends api key in query params""" + """Google AI Studio sends api key via x-goog-api-key header""" return headers @property @@ -75,7 +75,8 @@ class GeminiModelInfo(BaseLLMModelInfo): ) response = litellm.module_level_client.get( - url=f"{api_base}{endpoint}?key={api_key}", + url=f"{api_base}{endpoint}", + headers={"x-goog-api-key": api_key}, ) if response.status_code != 200: diff --git a/litellm/llms/gemini/files/transformation.py b/litellm/llms/gemini/files/transformation.py index a29ed66e63d..c30fba63263 100644 --- a/litellm/llms/gemini/files/transformation.py +++ b/litellm/llms/gemini/files/transformation.py @@ -86,7 +86,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): if not final_api_key: raise ValueError("api_key is required") - url = "{}/{}?key={}".format(api_base, endpoint, final_api_key) + url = "{}/{}".format(api_base, endpoint) return url def get_supported_openai_params( @@ -231,9 +231,9 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): ) api_base = api_base.rstrip("/") - url = f"{api_base}/v1beta/{file_part}?key={api_key}" + url = f"{api_base}/v1beta/{file_part}" - # Return empty params dict - API key is already in URL, no query params needed + # API key is passed via x-goog-api-key header (set in validate_environment) return url, {} def _normalize_gemini_file_id(self, file_id: str) -> str: diff --git a/litellm/llms/gemini/interactions/transformation.py b/litellm/llms/gemini/interactions/transformation.py index 772530342e1..c34da83cb8f 100644 --- a/litellm/llms/gemini/interactions/transformation.py +++ b/litellm/llms/gemini/interactions/transformation.py @@ -75,9 +75,13 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): model: str, litellm_params: Optional[GenericLiteLLMParams], ) -> dict: - """Google AI Studio uses API key in query params, not headers.""" + """Google AI Studio uses x-goog-api-key header for authentication.""" headers = headers or {} headers["Content-Type"] = "application/json" + if litellm_params: + api_key = GeminiModelInfo.get_api_key(litellm_params.get("api_key")) + if api_key: + headers["x-goog-api-key"] = api_key return headers def get_complete_url( @@ -98,11 +102,10 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): "Google API key is required. Set GOOGLE_API_KEY or GEMINI_API_KEY environment variable." ) - query_params = f"key={api_key}" if stream: - query_params += "&alt=sse" + return f"{api_base}/{self.api_version}/interactions?alt=sse" - return f"{api_base}/{self.api_version}/interactions?{query_params}" + return f"{api_base}/{self.api_version}/interactions" def transform_request( self, @@ -200,11 +203,10 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): ) -> Tuple[str, Dict]: """GET /{api_version}/interactions/{interaction_id}""" resolved_api_base = GeminiModelInfo.get_api_base(api_base) - api_key = GeminiModelInfo.get_api_key(litellm_params.api_key) - if not api_key: + if not GeminiModelInfo.get_api_key(litellm_params.api_key): raise ValueError("Google API key is required") return ( - f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}?key={api_key}", + f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}", {}, ) @@ -234,11 +236,10 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): ) -> Tuple[str, Dict]: """DELETE /{api_version}/interactions/{interaction_id}""" resolved_api_base = GeminiModelInfo.get_api_base(api_base) - api_key = GeminiModelInfo.get_api_key(litellm_params.api_key) - if not api_key: + if not GeminiModelInfo.get_api_key(litellm_params.api_key): raise ValueError("Google API key is required") return ( - f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}?key={api_key}", + f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}", {}, ) @@ -265,11 +266,10 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): ) -> Tuple[str, Dict]: """POST /{api_version}/interactions/{interaction_id}:cancel (if supported)""" resolved_api_base = GeminiModelInfo.get_api_base(api_base) - api_key = GeminiModelInfo.get_api_key(litellm_params.api_key) - if not api_key: + if not GeminiModelInfo.get_api_key(litellm_params.api_key): raise ValueError("Google API key is required") return ( - f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}:cancel?key={api_key}", + f"{resolved_api_base}/{self.api_version}/interactions/{interaction_id}:cancel", {}, ) diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index 2bb7bcd8b4f..4fac5aceb57 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -85,6 +85,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): raise ValueError("api_key is required for Gemini API calls") api_base = api_base.replace("https://", "wss://") api_base = api_base.replace("http://", "ws://") + # WebSocket connections do not support custom HTTP headers in all clients, + # so the API key must remain as a query parameter here. This is an accepted + # limitation; httpx is not used for WebSocket so MaskedHTTPStatusError + # already covers the main leak vector. return f"{api_base}/ws/google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent?key={api_key}" def map_model_turn_event( diff --git a/litellm/llms/gemini/vector_stores/transformation.py b/litellm/llms/gemini/vector_stores/transformation.py index 11fd77aecae..e6e8369643e 100644 --- a/litellm/llms/gemini/vector_stores/transformation.py +++ b/litellm/llms/gemini/vector_stores/transformation.py @@ -48,7 +48,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): def get_auth_credentials( self, litellm_params: dict ) -> BaseVectorStoreAuthCredentials: - """Gemini uses API key in query params, not headers.""" + """Gemini uses x-goog-api-key header for authentication.""" return {} def get_vector_store_endpoints_by_type(self) -> VectorStoreIndexEndpoints: @@ -79,6 +79,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): api_key = litellm_params.get("api_key") or get_api_key_from_env() if api_key: self._cached_api_key = api_key + headers["x-goog-api-key"] = api_key return headers @@ -133,13 +134,10 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): if model and model.startswith("gemini/"): model = model.replace("gemini/", "") - # Get API key - Gemini requires it as a query parameter api_key = litellm_params.get("api_key") or GeminiModelInfo.get_api_key() if not api_key: raise ValueError("GEMINI_API_KEY or GOOGLE_API_KEY is required") - - # Build the URL for generateContent with API key - url = f"{api_base}/models/{model}:generateContent?key={api_key}" + url = f"{api_base}/models/{model}:generateContent" # Build file_search tool configuration (using snake_case as per Gemini docs) file_search_config: Dict[str, Any] = { @@ -286,10 +284,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): """ url = f"{api_base}/fileSearchStores" - # Append API key as query parameter (required by Gemini) - api_key = self._cached_api_key or get_api_key_from_env() - if api_key: - url = f"{url}?key={api_key}" + # API key is passed via x-goog-api-key header (set in validate_environment) request_body: Dict[str, Any] = {} diff --git a/litellm/llms/openai/realtime/handler.py b/litellm/llms/openai/realtime/handler.py index c04857fc25f..5ca0a3186f7 100644 --- a/litellm/llms/openai/realtime/handler.py +++ b/litellm/llms/openai/realtime/handler.py @@ -6,6 +6,7 @@ This requires websockets, and is currently only supported on LiteLLM Proxy. from typing import Any, Optional, cast +from litellm._logging import _redact_string from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.types.realtime import RealtimeQueryParams @@ -148,11 +149,11 @@ class OpenAIRealtime(OpenAIChatCompletion): await realtime_streaming.bidirectional_forward() except websockets.exceptions.InvalidStatusCode as e: # type: ignore - await websocket.close(code=e.status_code, reason=str(e)) + await websocket.close(code=e.status_code, reason=_redact_string(str(e))) except Exception as e: try: await websocket.close( - code=1011, reason=f"Internal server error: {str(e)}" + code=1011, reason=_redact_string(f"Internal server error: {str(e)}") ) except RuntimeError as close_error: if "already completed" in str(close_error) or "websocket.close" in str( diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 5895a91f3aa..43e77f4fb75 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -337,8 +337,13 @@ def _get_gemini_url( mode: all_gemini_url_modes, model: str, stream: Optional[bool], - gemini_api_key: Optional[str], ) -> Tuple[str, str]: + """Build the Gemini API URL for the given mode. + + The API key is NOT included in the URL. Callers must pass it via the + ``x-goog-api-key`` header instead to avoid leaking credentials in + error tracebacks. + """ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) @@ -352,27 +357,27 @@ def _get_gemini_url( endpoint = "generateContent" if stream is True: endpoint = "streamGenerateContent" - url = "https://generativelanguage.googleapis.com/{}/{}:{}?key={}&alt=sse".format( - api_version, _gemini_model_name, endpoint, gemini_api_key + url = "https://generativelanguage.googleapis.com/{}/{}:{}?alt=sse".format( + api_version, _gemini_model_name, endpoint ) else: - url = "https://generativelanguage.googleapis.com/{}/{}:{}?key={}".format( - api_version, _gemini_model_name, endpoint, gemini_api_key + url = "https://generativelanguage.googleapis.com/{}/{}:{}".format( + api_version, _gemini_model_name, endpoint ) elif mode == "embedding": endpoint = "embedContent" - url = "https://generativelanguage.googleapis.com/v1beta/{}:{}?key={}".format( - _gemini_model_name, endpoint, gemini_api_key + url = "https://generativelanguage.googleapis.com/v1beta/{}:{}".format( + _gemini_model_name, endpoint ) elif mode == "batch_embedding": endpoint = "batchEmbedContents" - url = "https://generativelanguage.googleapis.com/v1beta/{}:{}?key={}".format( - _gemini_model_name, endpoint, gemini_api_key + url = "https://generativelanguage.googleapis.com/v1beta/{}:{}".format( + _gemini_model_name, endpoint ) elif mode == "count_tokens": endpoint = "countTokens" - url = "https://generativelanguage.googleapis.com/v1beta/{}:{}?key={}".format( - _gemini_model_name, endpoint, gemini_api_key + url = "https://generativelanguage.googleapis.com/v1beta/{}:{}".format( + _gemini_model_name, endpoint ) elif mode == "image_generation": raise ValueError( diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index b677cf3b1ec..0b872ddd8fd 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -62,10 +62,10 @@ class ContextCachingEndpoints(VertexBase): token, url """ if custom_llm_provider == "gemini": - auth_header = None + auth_header = {"x-goog-api-key": gemini_api_key} # type: ignore[assignment] endpoint = "cachedContents" - url = "https://generativelanguage.googleapis.com/v1beta/{}?key={}".format( - endpoint, gemini_api_key + url = "https://generativelanguage.googleapis.com/v1beta/{}".format( + endpoint ) elif custom_llm_provider == "vertex_ai": auth_header = vertex_auth_header @@ -353,7 +353,9 @@ class ContextCachingEndpoints(VertexBase): headers = { "Content-Type": "application/json", } - if token is not None: + if isinstance(token, dict): + headers.update(token) + elif token is not None: headers["Authorization"] = f"Bearer {token}" if extra_headers is not None: headers.update(extra_headers) @@ -501,7 +503,9 @@ class ContextCachingEndpoints(VertexBase): headers = { "Content-Type": "application/json", } - if token is not None: + if isinstance(token, dict): + headers.update(token) + elif token is not None: headers["Authorization"] = f"Bearer {token}" if extra_headers is not None: headers.update(extra_headers) diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 68d8f0d046d..430cc27adc8 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -473,9 +473,8 @@ class VertexBase: mode=mode, model=model, stream=stream, - gemini_api_key=gemini_api_key, ) - auth_header = None # this field is not used for gemin + auth_header = {"x-goog-api-key": gemini_api_key} # type: ignore[assignment] else: vertex_location = self.get_vertex_region( vertex_region=vertex_location, diff --git a/litellm/main.py b/litellm/main.py index ddd37b47536..22dffc0bbe8 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -40,6 +40,7 @@ from typing import ( get_args, ) +from litellm._logging import _redact_string from litellm._uuid import uuid if TYPE_CHECKING: @@ -7244,7 +7245,7 @@ async def ahealth_check( f"Mode {mode} not supported. See modes here: https://docs.litellm.ai/docs/proxy/health" ) except Exception as e: - stack_trace = traceback.format_exc() + stack_trace = _redact_string(traceback.format_exc()) if isinstance(stack_trace, str): stack_trace = stack_trace[:1000] diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 037f913ad07..c4717ad9cf3 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -22,7 +22,7 @@ from fastapi import HTTPException, Request, status from fastapi.responses import JSONResponse, Response, StreamingResponse import litellm -from litellm._logging import verbose_proxy_logger +from litellm._logging import _redact_string, verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import ( DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE, @@ -1785,7 +1785,7 @@ class ProxyBaseLLMRequestProcessing: if isinstance(e, HTTPException): raise e - error_traceback = traceback.format_exc() + error_traceback = _redact_string(traceback.format_exc()) error_msg = f"{str(e)}\n\n{error_traceback}" proxy_exception = ProxyException( message=getattr(e, "message", error_msg), diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a6f81986a6f..ac3bf6d498d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -77,7 +77,7 @@ from litellm import ( ModelResponseStream, Router, ) -from litellm._logging import verbose_proxy_logger +from litellm._logging import _redact_string, verbose_proxy_logger from litellm._service_logger import ServiceLogging, ServiceTypes from litellm.caching.caching import DualCache, RedisCache from litellm.caching.dual_cache import LimitedSizeOrderedDict @@ -155,7 +155,7 @@ def print_verbose(print_statement): verbose_proxy_logger.debug("{}\n{}".format(print_statement, traceback.format_exc())) if litellm.set_verbose: - print(f"LiteLLM Proxy: {print_statement}") # noqa + print(f"LiteLLM Proxy: {_redact_string(str(print_statement))}") # noqa def _get_email_logger_class(): @@ -1721,6 +1721,7 @@ class ProxyLogging: error_message = str(original_exception) if isinstance(traceback_str, str): error_message += traceback_str[:1000] + error_message = _redact_string(error_message) asyncio.create_task( self.alerting_handler( message=f"DB read/write call failed: {error_message}", @@ -1791,7 +1792,7 @@ class ProxyLogging: asyncio.create_task( self.alerting_handler( - message=f"LLM API call failed: `{exception_str}`", + message=_redact_string(f"LLM API call failed: `{exception_str}`"), level="High", alert_type=AlertType.llm_exceptions, request_data=request_data, diff --git a/litellm/rag/ingestion/gemini_ingestion.py b/litellm/rag/ingestion/gemini_ingestion.py index 96495b9f3ff..af6eb928e2c 100644 --- a/litellm/rag/ingestion/gemini_ingestion.py +++ b/litellm/rag/ingestion/gemini_ingestion.py @@ -143,7 +143,7 @@ class GeminiRAGIngestion(BaseRAGIngestion): Returns: Store name (format: fileSearchStores/xxxxxxx) """ - url = f"{base_url}/fileSearchStores?key={api_key}" + url = f"{base_url}/fileSearchStores" request_body = {"displayName": display_name} @@ -154,7 +154,10 @@ class GeminiRAGIngestion(BaseRAGIngestion): response = await client.post( url, json=request_body, - headers={"Content-Type": "application/json"}, + headers={ + "Content-Type": "application/json", + "x-goog-api-key": api_key, + }, ) if response.status_code != 200: @@ -228,7 +231,7 @@ class GeminiRAGIngestion(BaseRAGIngestion): # base_url is like: https://generativelanguage.googleapis.com/v1beta # We need: https://generativelanguage.googleapis.com/upload/v1beta/{store_id}:uploadToFileSearchStore api_base = base_url.replace("/v1beta", "") # Get base without version - url = f"{api_base}/upload/v1beta/{vector_store_id}:uploadToFileSearchStore?key={api_key}" + url = f"{api_base}/upload/v1beta/{vector_store_id}:uploadToFileSearchStore" # Build request body with chunking config and metadata if provided request_body: Dict[str, Any] = {"displayName": filename} @@ -263,6 +266,7 @@ class GeminiRAGIngestion(BaseRAGIngestion): "X-Goog-Upload-Header-Content-Length": str(file_size), "X-Goog-Upload-Header-Content-Type": content_type, "Content-Type": "application/json", + "x-goog-api-key": api_key, } verbose_logger.debug(f"Initiating resumable upload: {url}") diff --git a/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py b/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py index d814f8ec97f..8ca3a4d1492 100644 --- a/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py +++ b/tests/litellm/llms/vertex_ai/test_gemini_batch_embeddings.py @@ -310,7 +310,7 @@ def test_gemini_multimodal_embedding_e2e(): ) as mock_get_token: mock_get_token.return_value = ( {"x-goog-api-key": "test-key"}, - "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:embedContent?key=test-key" + "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:embedContent" ) mock_response = MagicMock() diff --git a/tests/test_litellm/google_genai/test_google_genai_adapter.py b/tests/test_litellm/google_genai/test_google_genai_adapter.py index 05b22098371..2af6867f092 100644 --- a/tests/test_litellm/google_genai/test_google_genai_adapter.py +++ b/tests/test_litellm/google_genai/test_google_genai_adapter.py @@ -1127,80 +1127,58 @@ async def test_google_generate_content_with_openai(): passed_fields = passed_fields - set(GenericLiteLLMParams.model_fields.keys()) # extra_headers is now explicitly passed through for providers that need custom headers assert passed_fields == set(["model", "messages", "extra_headers"]), f"Expected model, messages, and extra_headers to be passed through, got {passed_fields}" -@pytest.mark.asyncio -async def test_agenerate_content_x_goog_api_key_header(): +def test_validate_environment_sets_x_goog_api_key(): """ - Test that agenerate_content passes x-goog-api-key header correctly. - - This test verifies that when calling agenerate_content with a Google GenAI model, - the HTTP request includes the x-goog-api-key header with the correct API key value. - """ - import os - import unittest.mock + Test that VertexGeminiConfig.validate_environment correctly merges an + x-goog-api-key dict into the request headers. + + This is the mechanism by which Google AI Studio (Gemini) requests get + authenticated via header instead of a query-string ?key= parameter. + """ + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) - import httpx - test_api_key = "test-gemini-api-key-123" - - # Mock environment to ensure we use our test API key - with unittest.mock.patch.dict(os.environ, {"GEMINI_API_KEY": test_api_key}, clear=False): - # Mock the AsyncHTTPHandler's post method to capture headers - with unittest.mock.patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", new_callable=unittest.mock.AsyncMock) as mock_post: - # Mock a successful response - mock_response = unittest.mock.MagicMock() - mock_response.json.return_value = { - "candidates": [ - { - "content": { - "parts": [{"text": "Hello! How can I help you today?"}], - "role": "model" - }, - "finishReason": "STOP", - "index": 0 - } - ], - "usageMetadata": { - "promptTokenCount": 5, - "candidatesTokenCount": 10, - "totalTokenCount": 15 - } - } - mock_response.status_code = 200 - mock_response.headers = {} - mock_post.return_value = mock_response - - # Call agenerate_content with Google AI Studio model - try: - response = await agenerate_content( - model="gemini/gemini-1.5-flash", - contents=[ - {"role": "user", "parts": [{"text": "Hello, world!"}]} - ], - api_key=test_api_key - ) - except Exception: - # Ignore any response processing errors, we just want to check the headers - pass - - # Verify that AsyncHTTPHandler.post was called - mock_post.assert_called_once() - - # Get the arguments passed to the post call - call_args, call_kwargs = mock_post.call_args - - # Verify that headers contain x-goog-api-key - headers = call_kwargs.get("headers", {}) - assert "x-goog-api-key" in headers, f"x-goog-api-key header not found in headers: {list(headers.keys())}" - - # Verify the API key is set (could be our test key or from api_key parameter) - api_key_value = headers["x-goog-api-key"] - assert api_key_value == test_api_key, f"Expected x-goog-api-key to be {test_api_key}, got {api_key_value}" - - # Verify other expected headers - assert headers.get("Content-Type") == "application/json", f"Expected Content-Type application/json, got {headers.get('Content-Type')}" - print(f"✓ Test passed: x-goog-api-key header correctly set to {api_key_value}") - print(f"✓ All headers: {list(headers.keys())}") + # Simulate what _get_token_and_url returns for Gemini: a dict auth_header + auth_header_dict = {"x-goog-api-key": test_api_key} + + headers = VertexGeminiConfig().validate_environment( + api_key=auth_header_dict, + headers=None, + model="gemini-2.5-flash", + messages=[], + optional_params={}, + litellm_params={}, + ) + + assert "x-goog-api-key" in headers, f"x-goog-api-key not in headers: {headers}" + assert headers["x-goog-api-key"] == test_api_key + assert headers["Content-Type"] == "application/json" + + +def test_get_gemini_url_excludes_api_key(): + """ + Verify that _get_gemini_url never embeds the API key in the URL. + + API keys in URLs leak through httpx error tracebacks. The key must be + sent via the x-goog-api-key header instead. + """ + from litellm.llms.vertex_ai.common_utils import _get_gemini_url + + for mode in ("chat", "embedding", "batch_embedding", "count_tokens"): + url, _ = _get_gemini_url( + mode=mode, + model="gemini-2.5-flash", + stream=False, + ) + assert "key=" not in url, f"API key found in URL for mode={mode}: {url}" + + # Streaming chat should only have ?alt=sse + url, _ = _get_gemini_url(mode="chat", model="gemini-2.5-flash", stream=True) + assert "key=" not in url, f"API key found in streaming URL: {url}" + assert "alt=sse" in url, f"Missing alt=sse in streaming URL: {url}" def test_inline_data_base64_image_transformation(): diff --git a/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py b/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py index 6cc97cd95e6..21c036254e4 100644 --- a/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py +++ b/tests/test_litellm/llms/gemini/files/test_gemini_files_transformation.py @@ -37,12 +37,12 @@ class TestGoogleAIStudioFilesTransformation: litellm_params=litellm_params, ) - # Verify URL is constructed exactly as required: - # https://generativelanguage.googleapis.com/v1beta/files/{file_id}?key=API_KEY + # API key is passed via x-goog-api-key header, not in URL assert ( url - == "https://generativelanguage.googleapis.com/v1beta/files/test123?key=test-api-key" + == "https://generativelanguage.googleapis.com/v1beta/files/test123" ) + assert "key=" not in url # CRITICAL: params should be empty dict, not contain Content-Type or any other params # These would be incorrectly interpreted as query parameters @@ -64,12 +64,12 @@ class TestGoogleAIStudioFilesTransformation: litellm_params=litellm_params, ) - # Verify URL is constructed exactly as required: - # https://generativelanguage.googleapis.com/v1beta/files/{file_id}?key=API_KEY + # API key is passed via x-goog-api-key header, not in URL assert ( url - == "https://generativelanguage.googleapis.com/v1beta/files/test123?key=test-api-key" + == "https://generativelanguage.googleapis.com/v1beta/files/test123" ) + assert "key=" not in url # CRITICAL: params should be empty dict assert params == {}, f"Expected empty params dict, got: {params}" @@ -79,11 +79,10 @@ class TestGoogleAIStudioFilesTransformation: def test_transform_retrieve_file_request_with_raw_id_only(self): """ - Regression guard for the exact retrieval URL format. + Regression guard: API key must NOT appear in the URL. - If someone changes the method and stops producing: - https://generativelanguage.googleapis.com/v1beta/files/{file_id}?key=API_KEY - this test should fail. + The key is sent via x-goog-api-key header to prevent leaking + credentials in httpx error tracebacks. """ file_id = "cctqueckiggb" litellm_params = {"api_key": "test-api-key"} @@ -96,8 +95,9 @@ class TestGoogleAIStudioFilesTransformation: assert ( url - == "https://generativelanguage.googleapis.com/v1beta/files/cctqueckiggb?key=test-api-key" + == "https://generativelanguage.googleapis.com/v1beta/files/cctqueckiggb" ) + assert "key=" not in url assert params == {} @patch.dict("os.environ", {}, clear=True) @@ -285,10 +285,10 @@ class TestGoogleAIStudioFilesTransformation: litellm_params={}, ) - # Verify URL structure + # Verify URL structure - API key must NOT be in URL assert api_base in url assert "upload/v1beta/files" in url - assert f"key={api_key}" in url + assert "key=" not in url def test_transform_delete_file_request_with_full_uri(self): """Test delete file request transformation with full URI""" From 74f55b06715bf340775f1aec67bd27942bcf9478 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Fri, 3 Apr 2026 23:16:40 +0000 Subject: [PATCH 2/7] fix: apply mask_sensitive_info to async streaming error body MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses Greptile review — the async streaming path was missing mask_sensitive_info() on the response body, while the sync path had it. --- litellm/llms/custom_httpx/http_handler.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index fdb05d1a91b..859afcfe7e5 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -352,7 +352,7 @@ def _raise_masked_sync_error(e: httpx.HTTPStatusError, stream: bool) -> None: async def _raise_masked_async_error(e: httpx.HTTPStatusError, stream: bool) -> None: """Raise a MaskedHTTPStatusError for async HTTP handlers.""" if stream: - _body = await _safe_aread_response(e.response) + _body = mask_sensitive_info(await _safe_aread_response(e.response)) raise MaskedHTTPStatusError(e, message=_body, text=_body) from None _text = mask_sensitive_info(_safe_get_response_text(e.response)) raise MaskedHTTPStatusError(e, message=_text, text=_text) from None From b16d0b1d5ea5f3419f72a4ebac6d19659ccef1e9 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Sat, 4 Apr 2026 04:36:49 +0000 Subject: [PATCH 3/7] test: add coverage for credential leak prevention changes Add 50 tests across 3 files covering the new MaskedHTTPStatusError, safe response helpers, _redact_string in error paths, Gemini interactions x-goog-api-key header auth, and RAG ingestion header usage. Fix missing early-validation for Gemini API key in _get_token_and_url() which caused TypeError when key was None (headers got None value). Harmonize error messages between the two validation sites. Co-Authored-By: Claude Opus 4.6 (1M context) --- litellm/llms/vertex_ai/vertex_llm_base.py | 6 +- ...test_gemini_interactions_transformation.py | 148 +++++++++++ .../test_credential_leak_prevention.py | 234 +++++++++++++++++ .../test_redact_string_in_error_paths.py | 245 ++++++++++++++++++ 4 files changed, 632 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/interactions/test_gemini_interactions_transformation.py create mode 100644 tests/test_litellm/llms/custom_httpx/test_credential_leak_prevention.py create mode 100644 tests/test_litellm/test_redact_string_in_error_paths.py diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 430cc27adc8..46f4a807cd7 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -412,7 +412,7 @@ class VertexBase: url = "{}/models/{}:{}".format(api_base, model, endpoint) if gemini_api_key is None: raise ValueError( - "Missing gemini_api_key, please set `GEMINI_API_KEY`" + "Missing Gemini API key. Set the GEMINI_API_KEY or GOOGLE_API_KEY environment variable." ) if gemini_api_key is not None: auth_header = {"x-goog-api-key": gemini_api_key} # type: ignore[assignment] @@ -469,6 +469,10 @@ class VertexBase: """ version: Optional[Literal["v1beta1", "v1"]] = None if custom_llm_provider == "gemini": + if not gemini_api_key: + raise ValueError( + "Missing Gemini API key. Set the GEMINI_API_KEY or GOOGLE_API_KEY environment variable." + ) url, endpoint = _get_gemini_url( mode=mode, model=model, diff --git a/tests/test_litellm/interactions/test_gemini_interactions_transformation.py b/tests/test_litellm/interactions/test_gemini_interactions_transformation.py new file mode 100644 index 00000000000..465334d26fe --- /dev/null +++ b/tests/test_litellm/interactions/test_gemini_interactions_transformation.py @@ -0,0 +1,148 @@ +""" +Tests for Gemini Interactions API transformation. + +Covers credential leak prevention changes: +- validate_environment sets x-goog-api-key header +- get_complete_url excludes API key from URL +- get/delete/cancel interaction request URLs exclude API key +""" + +import os +import sys +from unittest.mock import patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.llms.gemini.interactions.transformation import ( + GoogleAIStudioInteractionsConfig, +) +from litellm.types.router import GenericLiteLLMParams + +_PATCH_GET_API_KEY = "litellm.llms.gemini.common_utils.GeminiModelInfo.get_api_key" + + +@pytest.fixture +def config(): + return GoogleAIStudioInteractionsConfig() + + +class TestValidateEnvironment: + def test_sets_x_goog_api_key_header(self, config): + litellm_params = GenericLiteLLMParams(api_key="test-api-key-123") + + headers = config.validate_environment( + headers={}, + model="gemini-2.5-flash", + litellm_params=litellm_params, + ) + + assert headers["x-goog-api-key"] == "test-api-key-123" + assert headers["Content-Type"] == "application/json" + + def test_no_api_key_skips_header(self, config): + litellm_params = GenericLiteLLMParams(api_key=None) + + with patch(_PATCH_GET_API_KEY, return_value=None): + headers = config.validate_environment( + headers={}, + model="gemini-2.5-flash", + litellm_params=litellm_params, + ) + + assert "x-goog-api-key" not in headers + assert headers["Content-Type"] == "application/json" + + def test_no_litellm_params_skips_header(self, config): + headers = config.validate_environment( + headers={}, + model="gemini-2.5-flash", + litellm_params=None, + ) + + assert "x-goog-api-key" not in headers + assert headers["Content-Type"] == "application/json" + + def test_preserves_existing_headers(self, config): + litellm_params = GenericLiteLLMParams(api_key="test-key") + + headers = config.validate_environment( + headers={"X-Custom": "value"}, + model="gemini-2.5-flash", + litellm_params=litellm_params, + ) + + assert headers["X-Custom"] == "value" + assert headers["x-goog-api-key"] == "test-key" + + +class TestGetCompleteUrl: + def test_url_excludes_api_key(self, config): + with patch(_PATCH_GET_API_KEY, return_value="secret-key"): + url = config.get_complete_url( + api_base=None, + model="gemini-2.5-flash", + litellm_params={"api_key": "secret-key"}, + ) + + assert "key=" not in url + assert "secret-key" not in url + assert url.endswith("/interactions") + + def test_stream_url_has_alt_sse_only(self, config): + with patch(_PATCH_GET_API_KEY, return_value="secret-key"): + url = config.get_complete_url( + api_base=None, + model="gemini-2.5-flash", + litellm_params={"api_key": "secret-key"}, + stream=True, + ) + + assert "key=" not in url + assert "secret-key" not in url + assert "alt=sse" in url + + def test_raises_without_api_key(self, config): + with patch(_PATCH_GET_API_KEY, return_value=None): + with pytest.raises(ValueError, match="Google API key is required"): + config.get_complete_url( + api_base=None, + model="gemini-2.5-flash", + litellm_params={"api_key": None}, + ) + + +class TestInteractionOperationUrls: + """Test that get/delete/cancel interaction URLs exclude API key.""" + + @pytest.mark.parametrize( + "method_name,interaction_id,expected_suffix", + [ + ("transform_get_interaction_request", "interaction-123", "interaction-123"), + ("transform_delete_interaction_request", "interaction-456", "interaction-456"), + ("transform_cancel_interaction_request", "interaction-789", "interaction-789:cancel"), + ], + ) + def test_url_excludes_key(self, config, method_name, interaction_id, expected_suffix): + with patch(_PATCH_GET_API_KEY, return_value="secret-key"): + url, params = getattr(config, method_name)( + interaction_id=interaction_id, + api_base="https://generativelanguage.googleapis.com", + litellm_params=GenericLiteLLMParams(api_key="secret-key"), + headers={}, + ) + + assert "key=" not in url + assert "secret-key" not in url + assert expected_suffix in url + + def test_get_interaction_raises_without_key(self, config): + with patch(_PATCH_GET_API_KEY, return_value=None): + with pytest.raises(ValueError, match="Google API key is required"): + config.transform_get_interaction_request( + interaction_id="interaction-123", + api_base="https://generativelanguage.googleapis.com", + litellm_params=GenericLiteLLMParams(api_key=None), + headers={}, + ) diff --git a/tests/test_litellm/llms/custom_httpx/test_credential_leak_prevention.py b/tests/test_litellm/llms/custom_httpx/test_credential_leak_prevention.py new file mode 100644 index 00000000000..543559422b0 --- /dev/null +++ b/tests/test_litellm/llms/custom_httpx/test_credential_leak_prevention.py @@ -0,0 +1,234 @@ +""" +Tests for credential leak prevention in HTTP handlers. + +Covers: +- MaskedHTTPStatusError construction and masking behavior +- _safe_get_response_text, _safe_aread_response, _safe_read_response helpers +- _raise_masked_sync_error and _raise_masked_async_error +""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + MaskedHTTPStatusError, + _raise_masked_async_error, + _raise_masked_sync_error, + _safe_aread_response, + _safe_get_response_text, + _safe_read_response, +) + + +def _make_httpx_status_error( + status_code: int = 400, + url: str = "https://example.com/v1/models?key=SECRET_KEY_123", + body: str = "Bad Request", +) -> httpx.HTTPStatusError: + """Create a real httpx.HTTPStatusError for testing.""" + request = httpx.Request("POST", url) + response = httpx.Response(status_code, request=request, content=body.encode()) + return httpx.HTTPStatusError( + message=f"Client error '{status_code}' for url '{url}'", + request=request, + response=response, + ) + + +class TestMaskedHTTPStatusError: + def test_masks_url_in_request(self): + orig = _make_httpx_status_error(url="https://api.example.com?key=MY_SECRET") + masked = MaskedHTTPStatusError(orig) + + assert "MY_SECRET" not in str(masked.request.url) + assert "[REDACTED_API_KEY]" in str(masked.request.url) + + def test_masks_original_message(self): + orig = _make_httpx_status_error(url="https://api.example.com?key=SUPER_SECRET") + masked = MaskedHTTPStatusError(orig) + + assert "SUPER_SECRET" not in str(masked) + assert "[REDACTED_API_KEY]" in str(masked) + + def test_preserves_status_code(self): + orig = _make_httpx_status_error(status_code=403) + masked = MaskedHTTPStatusError(orig) + + assert masked.status_code == 403 + assert masked.response.status_code == 403 + + def test_preserves_message_and_text_attrs(self): + orig = _make_httpx_status_error() + masked = MaskedHTTPStatusError(orig, message="custom msg", text="custom text") + + assert masked.message == "custom msg" + assert masked.text == "custom text" + + def test_handles_response_content_decompression_failure(self): + """If response.content raises (e.g. zlib error), should fall back to b''.""" + orig = _make_httpx_status_error() + + with patch.object( + type(orig.response), "content", + new_callable=lambda: property(lambda self: (_ for _ in ()).throw(Exception("zlib error"))), + ): + masked = MaskedHTTPStatusError(orig) + + assert masked.response.content == b"" + assert masked.status_code == 400 + + +class TestSafeResponseHelpers: + def test_safe_get_response_text_normal(self): + response = httpx.Response(200, content=b"hello world") + assert _safe_get_response_text(response) == "hello world" + + def test_safe_get_response_text_error(self): + response = MagicMock(spec=httpx.Response) + type(response).text = property(lambda self: (_ for _ in ()).throw(UnicodeDecodeError("utf-8", b"", 0, 1, "bad"))) + assert _safe_get_response_text(response) == "" + + def test_safe_read_response_normal(self): + response = httpx.Response(200, content=b"raw bytes") + result = _safe_read_response(response) + assert result == b"raw bytes" + + def test_safe_read_response_error(self): + response = MagicMock(spec=httpx.Response) + response.read.side_effect = Exception("read failure") + assert _safe_read_response(response) == b"" + + @pytest.mark.asyncio + async def test_safe_aread_response_normal(self): + response = MagicMock(spec=httpx.Response) + response.aread = AsyncMock(return_value=b"async bytes") + result = await _safe_aread_response(response) + assert result == b"async bytes" + + @pytest.mark.asyncio + async def test_safe_aread_response_error(self): + response = MagicMock(spec=httpx.Response) + response.aread = AsyncMock(side_effect=Exception("async read failure")) + result = await _safe_aread_response(response) + assert result == b"" + + +class TestRaiseMaskedError: + def test_sync_non_stream(self): + orig = _make_httpx_status_error( + url="https://api.example.com?key=LEAKED_KEY", body="error body" + ) + with pytest.raises(MaskedHTTPStatusError) as exc_info: + _raise_masked_sync_error(orig, stream=False) + + err = exc_info.value + assert "LEAKED_KEY" not in str(err.request.url) + assert err.status_code == 400 + assert err.text == "error body" + + def test_sync_stream(self): + orig = _make_httpx_status_error( + url="https://api.example.com?key=LEAKED_KEY", body="stream body" + ) + with pytest.raises(MaskedHTTPStatusError) as exc_info: + _raise_masked_sync_error(orig, stream=True) + + err = exc_info.value + assert "LEAKED_KEY" not in str(err.request.url) + assert err.message is not None + + def test_sync_breaks_exception_chain(self): + orig = _make_httpx_status_error() + with pytest.raises(MaskedHTTPStatusError) as exc_info: + _raise_masked_sync_error(orig, stream=False) + + assert exc_info.value.__cause__ is None + + @pytest.mark.asyncio + async def test_async_non_stream(self): + orig = _make_httpx_status_error( + url="https://api.example.com?key=LEAKED_KEY", body="async error" + ) + with pytest.raises(MaskedHTTPStatusError) as exc_info: + await _raise_masked_async_error(orig, stream=False) + + err = exc_info.value + assert "LEAKED_KEY" not in str(err.request.url) + assert err.status_code == 400 + assert err.text == "async error" + + @pytest.mark.asyncio + async def test_async_stream(self): + orig = _make_httpx_status_error( + url="https://api.example.com?key=LEAKED_KEY", body="async stream" + ) + with pytest.raises(MaskedHTTPStatusError) as exc_info: + await _raise_masked_async_error(orig, stream=True) + + err = exc_info.value + assert "LEAKED_KEY" not in str(err.request.url) + assert err.message is not None + + @pytest.mark.asyncio + async def test_async_breaks_chain(self): + orig = _make_httpx_status_error() + with pytest.raises(MaskedHTTPStatusError) as exc_info: + await _raise_masked_async_error(orig, stream=False) + + assert exc_info.value.__cause__ is None + + +class TestHTTPHandlerErrorPaths: + """Test that HTTP handler methods raise MaskedHTTPStatusError on HTTPStatusError.""" + + @pytest.fixture + def sync_handler(self): + handler = HTTPHandler() + yield handler + handler.close() + + @pytest.fixture + async def async_handler(self): + handler = AsyncHTTPHandler() + yield handler + await handler.close() + + @pytest.mark.parametrize("method", ["post", "put", "patch", "delete"]) + def test_sync_raises_masked_error(self, sync_handler, method): + with patch.object( + sync_handler.client, + "send", + side_effect=_make_httpx_status_error(url="https://api.test.com?key=SECRET"), + ): + with pytest.raises(MaskedHTTPStatusError) as exc_info: + kwargs = {"url": "https://api.test.com?key=SECRET"} + if method != "delete": + kwargs["data"] = {"test": 1} + getattr(sync_handler, method)(**kwargs) + + assert "SECRET" not in str(exc_info.value.request.url) + + @pytest.mark.parametrize("method", ["post", "put", "patch", "delete"]) + @pytest.mark.asyncio + async def test_async_raises_masked_error(self, async_handler, method): + with patch.object( + async_handler.client, + "send", + new_callable=AsyncMock, + side_effect=_make_httpx_status_error(url="https://api.test.com?key=SECRET"), + ): + with pytest.raises(MaskedHTTPStatusError) as exc_info: + kwargs = {"url": "https://api.test.com?key=SECRET"} + if method != "delete": + kwargs["data"] = {"test": 1} + await getattr(async_handler, method)(**kwargs) + + assert "SECRET" not in str(exc_info.value.request.url) diff --git a/tests/test_litellm/test_redact_string_in_error_paths.py b/tests/test_litellm/test_redact_string_in_error_paths.py new file mode 100644 index 00000000000..7d402ee68f1 --- /dev/null +++ b/tests/test_litellm/test_redact_string_in_error_paths.py @@ -0,0 +1,245 @@ +""" +Tests for _redact_string usage in error/logging paths. + +Covers actual execution of redaction in: +- WebSocket close reasons in realtime handlers (openai, azure, bedrock) +- Gemini RAG ingestion x-goog-api-key header usage +- Traceback redaction pattern used in proxy streaming +""" + +import os +import sys +import traceback +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm._logging import _ENABLE_SECRET_REDACTION, _redact_string + + +class TestRedactStringFunction: + def test_redacts_bearer_token(self): + text = "Authorization: Bearer sk-1234567890abcdefghij" + result = _redact_string(text) + assert "sk-1234567890abcdefghij" not in result + assert "REDACTED" in result + + def test_redacts_api_key_in_url(self): + text = "Error at https://example.com?api_key=my-secret-key-value-here" + result = _redact_string(text) + assert "my-secret-key-value-here" not in result + + def test_redacts_google_api_key(self): + text = "key=AIzaSyB1234567890abcdefghijklmnopqrstuvwx" + result = _redact_string(text) + assert "AIzaSyB1234567890abcdefghijklmnopqrstuvwx" not in result + + def test_passes_clean_text_through(self): + text = "This is a normal error message with no secrets" + assert _redact_string(text) == text + + @pytest.mark.skipif( + not _ENABLE_SECRET_REDACTION, reason="redaction disabled via env var" + ) + def test_redaction_enabled_by_default(self): + text = "Bearer sk-1234567890abcdefghij" + result = _redact_string(text) + assert "sk-1234567890abcdefghij" not in result + + +class TestOpenAIRealtimeRedaction: + """Test that OpenAI realtime handler redacts secrets in websocket close reasons.""" + + def _make_patches(self, handler): + """Shared patches for OpenAI realtime handler tests.""" + return ( + patch.object(handler, "_construct_url", return_value="wss://api.openai.com/v1/realtime?model=gpt-4"), + patch.object(handler, "_get_ssl_config", return_value=None), + patch.object(handler, "_get_additional_headers", return_value={}), + ) + + def _call_kwargs(self): + return dict( + model="gpt-4", + websocket=AsyncMock(), + logging_obj=MagicMock(), + api_base="https://api.openai.com/", + api_key="test-key", + ) + + @pytest.mark.asyncio + async def test_invalid_status_code_redacts_reason(self): + import websockets.exceptions + + from litellm.llms.openai.realtime.handler import OpenAIRealtime + + handler = OpenAIRealtime() + exc = websockets.exceptions.InvalidStatusCode(403, None) + exc.status_code = 403 + + kwargs = self._call_kwargs() + mock_ws = kwargs["websocket"] + p1, p2, p3 = self._make_patches(handler) + with p1, p2, p3, patch("websockets.connect", side_effect=exc): + await handler.async_realtime(**kwargs) + + mock_ws.close.assert_called_once() + assert mock_ws.close.call_args[1]["code"] == 403 + + @pytest.mark.asyncio + async def test_generic_exception_redacts_reason(self): + from litellm.llms.openai.realtime.handler import OpenAIRealtime + + handler = OpenAIRealtime() + secret_error = RuntimeError("Connection failed for api_key=sk-1234567890abcdefghij") + + kwargs = self._call_kwargs() + mock_ws = kwargs["websocket"] + p1, p2, p3 = self._make_patches(handler) + with p1, p2, p3, patch("websockets.connect", side_effect=secret_error): + await handler.async_realtime(**kwargs) + + mock_ws.close.assert_called_once() + assert mock_ws.close.call_args[1]["code"] == 1011 + assert "sk-1234567890abcdefghij" not in mock_ws.close.call_args[1]["reason"] + + +class TestAzureRealtimeRedaction: + """Test that Azure realtime handler redacts secrets in websocket close reasons.""" + + @pytest.mark.asyncio + async def test_invalid_status_code_redacts_reason(self): + import websockets.exceptions + + from litellm.llms.azure.realtime.handler import AzureOpenAIRealtime + + handler = AzureOpenAIRealtime() + mock_ws = AsyncMock() + exc = websockets.exceptions.InvalidStatusCode(403, None) + exc.status_code = 403 + + with patch.object(handler, "_construct_url", return_value="wss://test.openai.azure.com/openai/realtime"), \ + patch("websockets.connect", side_effect=exc): + await handler.async_realtime( + model="gpt-4", + websocket=mock_ws, + logging_obj=MagicMock(), + api_base="https://test.openai.azure.com/", + api_key="test-key", + api_version="2024-10-01-preview", + ) + + mock_ws.close.assert_called_once() + assert mock_ws.close.call_args[1]["code"] == 403 + + +class TestBedrockRealtimeRedaction: + """Test that _redact_string produces safe close reasons for Bedrock-style errors.""" + + def test_internal_error_message_redacted(self): + secret_error = RuntimeError("Failed with aws_secret_access_key=AKIAIOSFODNN7EXAMPLE123456") + reason = _redact_string(f"Internal error: {str(secret_error)}") + assert "AKIAIOSFODNN7EXAMPLE123456" not in reason + + +class TestLLMHTTPHandlerRealtimeRedaction: + """Test _redact_string on the exact patterns used in llm_http_handler WS close.""" + + def test_invalid_status_pattern(self): + error_msg = "InvalidStatusCode: 403 for wss://api.example.com?api_key=sk-leaked-key-here" + assert "sk-leaked-key-here" not in _redact_string(str(error_msg)) + + def test_internal_server_error_pattern(self): + error_msg = "Connection failed for api_key=sk-secret-key-12345678" + assert "sk-secret-key-12345678" not in _redact_string(f"Internal server error: {error_msg}") + + +class TestProxyStreamingDataGeneratorRedaction: + """Test _redact_string on traceback.format_exc() — the pattern at common_request_processing.py:1733.""" + + def test_redact_traceback_format_exc(self): + try: + raise RuntimeError( + "Failed connecting to api_key=sk-1234567890abcdefghij at https://api.example.com" + ) + except RuntimeError: + raw_tb = traceback.format_exc() + + redacted_tb = _redact_string(raw_tb) + + assert "sk-1234567890abcdefghij" not in redacted_tb + assert "Traceback" in redacted_tb + assert "RuntimeError" in redacted_tb + + +def _make_mock_ingest_options(): + mock = MagicMock() + mock.vector_store_config = {} + mock.ingest_name = "test" + mock.chunking_strategy = None + mock.embedding_model = None + mock.vector_db_type = "gemini" + return mock + + +class TestGeminiIngestionHeaders: + """Test that Gemini RAG ingestion uses x-goog-api-key header.""" + + @pytest.mark.asyncio + async def test_create_file_search_store_sends_header(self): + from litellm.rag.ingestion.gemini_ingestion import GeminiRAGIngestion + + ingestion = GeminiRAGIngestion(ingest_options=_make_mock_ingest_options()) + + mock_client = AsyncMock() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"name": "fileSearchStores/abc123"} + mock_client.post.return_value = mock_response + + with patch( + "litellm.rag.ingestion.gemini_ingestion.get_async_httpx_client", + return_value=mock_client, + ): + result = await ingestion._create_file_search_store( + api_key="test-gemini-key", + base_url="https://generativelanguage.googleapis.com/v1beta", + display_name="test-store", + ) + + assert result == "fileSearchStores/abc123" + call_kwargs = mock_client.post.call_args + assert call_kwargs[1]["headers"]["x-goog-api-key"] == "test-gemini-key" + assert "key=" not in call_kwargs[0][0] + + @pytest.mark.asyncio + async def test_initiate_resumable_upload_sends_header(self): + from litellm.rag.ingestion.gemini_ingestion import GeminiRAGIngestion + + ingestion = GeminiRAGIngestion(ingest_options=_make_mock_ingest_options()) + + mock_client = AsyncMock() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.headers = {"x-goog-upload-url": "https://upload.example.com/upload123"} + mock_client.post.return_value = mock_response + + with patch( + "litellm.rag.ingestion.gemini_ingestion.get_async_httpx_client", + return_value=mock_client, + ): + result = await ingestion._initiate_resumable_upload( + api_key="test-gemini-key", + base_url="https://generativelanguage.googleapis.com/v1beta", + vector_store_id="fileSearchStores/abc123", + filename="test.txt", + file_size=1024, + content_type="text/plain", + ) + + assert result == "https://upload.example.com/upload123" + call_kwargs = mock_client.post.call_args + assert call_kwargs[1]["headers"]["x-goog-api-key"] == "test-gemini-key" + assert "key=" not in call_kwargs[0][0] From abb8d8d4548e37a3dfadb02c30c0eee0983de601 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Sat, 4 Apr 2026 05:11:14 +0000 Subject: [PATCH 4/7] ci: trigger CI re-run for codecov Co-Authored-By: Claude Opus 4.6 (1M context) From f521e27371e957313cf75aa79968f1f7203a27be Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Tue, 14 Apr 2026 23:28:13 +0000 Subject: [PATCH 5/7] test(gemini): align API key expectations --- .../test_google_gemini_proxy_request.py | 37 ++++++++++--------- .../llms/vertex_ai/test_vertex.py | 14 ++++++- .../llms/vertex_ai/test_vertex_llm_base.py | 2 +- 3 files changed, 32 insertions(+), 21 deletions(-) diff --git a/tests/proxy_unit_tests/test_google_gemini_proxy_request.py b/tests/proxy_unit_tests/test_google_gemini_proxy_request.py index 90c2cac18d0..98eb731e871 100644 --- a/tests/proxy_unit_tests/test_google_gemini_proxy_request.py +++ b/tests/proxy_unit_tests/test_google_gemini_proxy_request.py @@ -320,29 +320,30 @@ async def test_generationconfig_to_config_mapping(sample_request_payload): for Google GenAI compatibility in the main functions. """ from litellm.google_genai.main import agenerate_content - + # Create a copy of the payload to avoid modifying the fixture test_data = sample_request_payload.copy() - - # Test that agenerate_content can handle generationConfig parameter - # This should not raise an error about parameter handling - try: - # This will fail due to missing API key, but should not fail due to parameter handling + + with patch( + "litellm.google_genai.main.base_llm_http_handler.generate_content_handler" + ) as mock_generate_content_handler: + mock_generate_content_handler.return_value = {"text": "mock response"} + await agenerate_content( model="gemini/gemini-2.5-flash", contents=test_data["contents"], - generationConfig=test_data["generationConfig"], # Pass as generationConfig - custom_llm_provider="gemini" + generationConfig=test_data["generationConfig"], + custom_llm_provider="gemini", ) - except Exception as e: - # Should not fail due to parameter handling issues - error_msg = str(e).lower() - if "generationconfig" in error_msg or "config" in error_msg or "parameter" in error_msg: - pytest.fail(f"Parameter handling failed: {e}") - # Other errors (like API key missing) are expected - print(f"✅ Parameter handling worked (API error expected): {type(e).__name__}") - - print("✅ generationConfig to config mapping test passed") + + mock_generate_content_handler.assert_called_once() + generate_content_config_dict = mock_generate_content_handler.call_args.kwargs[ + "generate_content_config_dict" + ] + assert generate_content_config_dict["temperature"] == 0 + assert generate_content_config_dict["topP"] == 1 + assert generate_content_config_dict["responseMimeType"] == "application/json" + assert "responseJsonSchema" in generate_content_config_dict @pytest.mark.asyncio @@ -405,7 +406,7 @@ async def test_gemini_custom_api_base_proxy_integration(): print(f"✅ Custom API base streaming URL test passed: {result_url_streaming}") # Test case 3: Error handling - missing API key - with pytest.raises(ValueError, match="Missing gemini_api_key"): + with pytest.raises(ValueError, match="Missing Gemini API key"): vertex_base._check_custom_proxy( api_base=custom_api_base, custom_llm_provider="gemini", diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex.py b/tests/test_litellm/llms/vertex_ai/test_vertex.py index 2bd6182a331..6facd0aab8e 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex.py @@ -93,6 +93,7 @@ def test_completion_pydantic_obj_2(): model="gemini/gemini-2.5-flash", messages=messages, response_format=EventsList, + api_key="test-api-key", client=client, ) # print(response) @@ -285,6 +286,7 @@ def test_function_calling_with_gemini(): }, }, ], + api_key="test-api-key", client=client, ) except Exception as e: @@ -372,7 +374,10 @@ def test_multiple_function_call(): with patch.object(client, "post", return_value=mock_response) as mock_post: r = litellm.completion( - messages=messages, model="gemini/gemini-1.5-flash-002", client=client + messages=messages, + model="gemini/gemini-1.5-flash-002", + api_key="test-api-key", + client=client, ) assert len(r.choices) > 0 @@ -478,7 +483,10 @@ def test_multiple_function_call_changed_text_pos(): with patch.object(client, "post", return_value=mock_response) as mock_post: resp = litellm.completion( - messages=messages, model="gemini/gemini-1.5-flash-002", client=client + messages=messages, + model="gemini/gemini-1.5-flash-002", + api_key="test-api-key", + client=client, ) assert len(resp.choices) > 0 mock_post.assert_called_once() @@ -599,6 +607,7 @@ def test_function_calling_with_gemini_multiple_results(): messages=messages, tools=tools, tool_choice="required", + api_key="test-api-key", client=client, ) print("Response\n", response) @@ -1182,6 +1191,7 @@ def test_logprobs(): {"role": "user", "content": "What's the weather like in San Francisco?"} ], logprobs=True, + api_key="test-api-key", client=client, ) print(resp) diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py index 78caf4b9778..2194cadf1b2 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py @@ -810,7 +810,7 @@ class TestVertexBase: if custom_llm_provider == "gemini" and api_base and gemini_api_key is None: # Test case 5: Should raise ValueError for Gemini without API key - with pytest.raises(ValueError, match="Missing gemini_api_key"): + with pytest.raises(ValueError, match="Missing Gemini API key"): vertex_base._check_custom_proxy( api_base=api_base, custom_llm_provider=custom_llm_provider, From b1bc3c166d21a5da801e8012de04644c828e6009 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Tue, 14 Apr 2026 23:37:13 +0000 Subject: [PATCH 6/7] test(prompts): isolate in-memory version tests --- .../proxy/prompts/test_prompt_endpoints.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/proxy/prompts/test_prompt_endpoints.py b/tests/test_litellm/proxy/prompts/test_prompt_endpoints.py index 6c2e5fa7667..d1a7c59aa3a 100644 --- a/tests/test_litellm/proxy/prompts/test_prompt_endpoints.py +++ b/tests/test_litellm/proxy/prompts/test_prompt_endpoints.py @@ -247,8 +247,10 @@ class TestPromptVersionsEndpoint: ), } - # Mock the IN_MEMORY_PROMPT_REGISTRY at the import location - with patch("litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY") as mock_registry: + # Force the in-memory path so this test is isolated from any leaked prisma mocks. + with patch("litellm.proxy.proxy_server.prisma_client", None), patch( + "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" + ) as mock_registry: mock_registry.IN_MEMORY_PROMPTS = mock_prompts # Test with base prompt ID @@ -293,7 +295,9 @@ class TestPromptVersionsEndpoint: user_role=LitellmUserRoles.PROXY_ADMIN ) - with patch("litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY") as mock_registry: + with patch("litellm.proxy.proxy_server.prisma_client", None), patch( + "litellm.proxy.prompts.prompt_registry.IN_MEMORY_PROMPT_REGISTRY" + ) as mock_registry: mock_registry.IN_MEMORY_PROMPTS = {} with pytest.raises(HTTPException) as exc_info: @@ -304,4 +308,3 @@ class TestPromptVersionsEndpoint: assert exc_info.value.status_code == 404 assert "No versions found" in exc_info.value.detail - From 2911d99d77ed9da61023489340d87eed403cece4 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Wed, 15 Apr 2026 00:28:39 +0000 Subject: [PATCH 7/7] test(gemini): stub API key for format param tests --- tests/test_litellm/test_main.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 40a3692ac69..f538bc4e2f0 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -213,6 +213,8 @@ async def test_url_with_format_param(model, sync_mode, monkeypatch): } ], } + if model.startswith("gemini/"): + args["api_key"] = "test-api-key" with patch.object(client, "post", new=MagicMock()) as mock_client: try: if sync_mode: