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 e335ebabd8b..87df171d653 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 @@ -2854,7 +2854,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 @@ -2877,7 +2881,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..859afcfe7e5 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 = 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 + + 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..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,13 +469,16 @@ 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, 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/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/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/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/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""" 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, 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 - 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: 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]