From c0c0c9a9ebc7443b1724c22e3b6b856aac1f750d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:05:10 -0700 Subject: [PATCH 1/6] fix(sdk): keep body and proxy headers on BadRequestError mapped from a litellm_proxy 400 The generic 400 branch of the OpenAI exception mapper dropped the wire body and no branch carried the response headers, so an application calling a LiteLLM proxy through a litellm_proxy/ model could not tell a guardrail block from any other failure without walking __cause__. BadRequestError now takes headers, filled for a litellm_proxy upstream, and the generic branch passes the body. The proxy edge treats the literal "None" type and param an older proxy sends as absent and stops forwarding an upstream proxy's date and server headers. --- litellm/constants.py | 6 +- litellm/exceptions.py | 4 +- .../exception_mapping_utils.py | 16 +++++ .../common_utils/openai_error_payload.py | 6 +- .../test_exception_mapping_utils.py | 68 +++++++++++++++++++ .../common_utils/test_openai_error_payload.py | 18 +++++ .../proxy/test_common_request_processing.py | 15 ++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 - 8 files changed, 129 insertions(+), 6 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 5751e6e46af..caf1aff6792 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1964,7 +1964,11 @@ BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset( } ) -UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS +ORIGIN_SERVER_HEADERS: Final[frozenset[str]] = frozenset({"date", "server"}) + +UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = ( + HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS | ORIGIN_SERVER_HEADERS +) # A retrieved response replays the usage of the call that created it, so pricing these # read/management routes like inference bills the same tokens twice. diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 3f22a4b2dcd..3d9e26e450b 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -10,7 +10,7 @@ ## LiteLLM versions of the OpenAI Exception Types import enum -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Any, Final import httpx @@ -226,6 +226,7 @@ class BadRequestError(openai.BadRequestError): max_retries: int | None = None, num_retries: int | None = None, body: dict | None = None, + headers: Mapping[str, str] | None = None, ): self.status_code = 400 self.message = f"litellm.BadRequestError: {message}" @@ -234,6 +235,7 @@ class BadRequestError(openai.BadRequestError): self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries + self.headers: dict[str, str] | None = {k: str(v) for k, v in headers.items()} if headers else None # Use response if it's a valid httpx.Response with a request, otherwise use minimal error response # Note: We check _request (not .request property) to avoid RuntimeError when _request is None if ( diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 82708d412c9..fd1ba666887 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -1,6 +1,7 @@ import json import re import traceback +from collections.abc import Mapping from typing import Any, Final, Protocol, cast import httpx @@ -254,6 +255,15 @@ class _ProviderHTTPException(Protocol): llm_provider: str +def _litellm_proxy_response_headers( + original_exception: _ProviderHTTPException, custom_llm_provider: str +) -> Mapping[str, str] | None: + if custom_llm_provider != "litellm_proxy": + return None + headers: Final = getattr(original_exception, "headers", None) + return headers if isinstance(headers, Mapping) else None + + def _map_openai_exception( *, model: str, @@ -264,6 +274,7 @@ def _map_openai_exception( exception_provider: str, extra_information: str, ) -> None: + upstream_headers: Final = _litellm_proxy_response_headers(original_exception, custom_llm_provider) # custom_llm_provider is openai, make it OpenAI message = get_error_message(error_obj=original_exception) if message is None: @@ -348,6 +359,7 @@ def _map_openai_exception( response=getattr(original_exception, "response", None), litellm_debug_info=extra_information, body=getattr(original_exception, "body", None), + headers=upstream_headers, ) elif "invalid_request_error" in error_str and "Incorrect API key provided" not in error_str: raise BadRequestError( @@ -357,6 +369,7 @@ def _map_openai_exception( response=getattr(original_exception, "response", None), litellm_debug_info=extra_information, body=getattr(original_exception, "body", None), + headers=upstream_headers, ) elif ( "Web server is returning an unknown error" in error_str @@ -404,6 +417,8 @@ def _map_openai_exception( model=model, response=getattr(original_exception, "response", None), litellm_debug_info=extra_information, + body=getattr(original_exception, "body", None), + headers=upstream_headers, ) elif original_exception.status_code == 401: raise AuthenticationError( @@ -436,6 +451,7 @@ def _map_openai_exception( response=getattr(original_exception, "response", None), litellm_debug_info=extra_information, body=getattr(original_exception, "body", None), + headers=upstream_headers, ) elif original_exception.status_code == 429: raise RateLimitError( diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py index 89f735ee8b6..90b3c998247 100644 --- a/litellm/proxy/common_utils/openai_error_payload.py +++ b/litellm/proxy/common_utils/openai_error_payload.py @@ -8,6 +8,8 @@ from typing import Final from fastapi import status +_STRINGIFIED_NONE: Final = "None" + _OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType( { status.HTTP_401_UNAUTHORIZED: "authentication_error", @@ -35,7 +37,7 @@ def openai_error_type(exc: object, status_code: int) -> str: """OpenAI types ``error.type`` as a required string, so an exception carrying none falls back to the type its status code stands for.""" carried: Final = attribute_of(exc, "type") - if isinstance(carried, str): + if isinstance(carried, str) and carried != _STRINGIFIED_NONE: return carried mapped: Final = _OPENAI_ERROR_TYPE_BY_STATUS.get(status_code) if mapped is not None: @@ -49,4 +51,4 @@ def openai_error_param(exc: object) -> str | None: """OpenAI types ``error.param`` as nullable, so an exception carrying none serializes as JSON ``null``.""" carried: Final = attribute_of(exc, "param") - return carried if isinstance(carried, str) else None + return carried if isinstance(carried, str) and carried != _STRINGIFIED_NONE else None diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index 42d3df76902..ea6ac17ad45 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -1409,3 +1409,71 @@ def test_bedrock_timeout_mapping_keeps_retry_after_readable(status_code): exception_headers = _get_response_headers(original_exception=exc_info.value) assert exception_headers is not None assert litellm.utils._get_retry_after_from_exception_header(response_headers=exception_headers) == 7 + + +_GUARDRAIL_BLOCK_ERROR = { + "message": "Content blocked: secret_project_codename pattern detected", + "param": "None", + "code": "400", + "provider_specific_fields": { + "error": "Content blocked: secret_project_codename pattern detected", + "pattern": "secret_project_codename", + "guardrail_name": "block-secret-project", + "guardrail_mode": "pre_call", + }, +} + + +def _openai_handler_error(error_type: str, headers: dict[str, str]) -> OpenAIError: + """What litellm/llms/openai/openai.py raises after the openai SDK rejects a 400: + the SDK's str() carries the wire body, and the handler copies headers and body over.""" + wire_error = {**_GUARDRAIL_BLOCK_ERROR, "type": error_type} + wire = httpx.Response( + status_code=400, + headers=headers, + json={"error": wire_error}, + request=httpx.Request("POST", "http://localhost:4000/v1/chat/completions"), + ) + return OpenAIError( + status_code=400, + message=f"Error code: 400 - {{'error': {wire_error}}}", + headers=wire.headers, + body=wire_error, + ) + + +@pytest.mark.parametrize("error_type", ["None", "invalid_request_error"]) +def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str): + """An SDK caller behind a proxy tells a guardrail block from any other 400 by the body's + provider_specific_fields and the proxy's x-litellm-* headers, so the mapped BadRequestError + must carry both whichever error.type the proxy version on the other end emits.""" + proxy_headers = {"x-litellm-call-id": "call-guardrail", "x-litellm-applied-guardrails": "block-secret-project"} + + with pytest.raises(litellm.BadRequestError) as exc_info: + exception_type( + model="claude-haiku-4-5", + original_exception=_openai_handler_error(error_type, proxy_headers), + custom_llm_provider="litellm_proxy", + completion_kwargs={}, + extra_kwargs={}, + ) + + assert exc_info.value.body["provider_specific_fields"]["guardrail_name"] == "block-secret-project" + assert exc_info.value.body["type"] == error_type + assert proxy_headers.items() <= exc_info.value.headers.items() + + +def test_openai_compatible_vendor_400_keeps_body_but_not_headers(): + """A vendor's own response headers stay on e.response the way every other mapped provider + error keeps them; only a LiteLLM proxy upstream puts headers on e.headers.""" + with pytest.raises(litellm.BadRequestError) as exc_info: + exception_type( + model="gpt-5.4-mini", + original_exception=_openai_handler_error("vendor_specific_error", {"openai-organization": "org-1"}), + custom_llm_provider="openai", + completion_kwargs={}, + extra_kwargs={}, + ) + + assert exc_info.value.body["type"] == "vendor_specific_error" + assert exc_info.value.headers is None diff --git a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py index 8b653ddfb71..db9fe3a2253 100644 --- a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py +++ b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py @@ -143,3 +143,21 @@ def test_a_status_carried_by_an_exception_drives_the_type_it_reports(): exc = HTTPException(status_code=403, detail="blocked by policy") assert openai_error_type(exc, error_status_code(exc, 400)) == "permission_error" + + +def test_the_stringified_none_an_older_upstream_proxy_sent_is_treated_as_absent(): + """A proxy fronting a proxy older than 1.102 receives {"type": "None", "param": "None"} on + the wire; the SDK now keeps that body on the mapped exception, and re-emitting the literal + is the exact bug this module exists to stop.""" + from litellm.exceptions import BadRequestError + + carried = BadRequestError( + message="Content blocked", + model="claude-haiku-4-5", + llm_provider="litellm_proxy", + body={"message": "Content blocked", "type": "None", "param": "None", "code": "400"}, + ) + + assert carried.type == "None" + assert openai_error_type(carried, 400) == "invalid_request_error" + assert openai_error_param(carried) is None diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 69e89d1c604..3cb58c44354 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -3985,6 +3985,21 @@ class TestHandleLLMApiExceptionFramingHeaders: assert proxy_exc.headers["x-custom-safe"] == "1" assert proxy_exc.headers["x-request-id"] == "abc-123" + async def test_strips_the_date_and_server_headers_of_an_upstream_litellm_proxy(self): + """A proxy fronting another LiteLLM proxy gets the upstream's date and server + on the mapped exception; forwarding them would duplicate the Date header + uvicorn adds to every response and leak the upstream server identity.""" + exc = litellm.BadRequestError( + message="Content blocked", + llm_provider="litellm_proxy", + model="claude-haiku-4-5", + headers={"date": "Sun, 13 Sep 2026 08:43:51 GMT", "server": "uvicorn", "x-request-id": "abc-123"}, + ) + proxy_exc = await self._invoke(exc) + assert "date" not in proxy_exc.headers + assert "server" not in proxy_exc.headers + assert proxy_exc.headers["x-request-id"] == "abc-123" + class TestAsyncStreamingDataGeneratorFastPath: """Fast/slow path branching in async_streaming_data_generator.""" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7eadaa6c991..839aa52fa84 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16781,7 +16781,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) @@ -16887,7 +16886,6 @@ export interface paths { * - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking. * - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x. - * - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests. * - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys) * - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. * - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys) From 825e4f17e949439d37f922bb02c0356f2bdd0dc5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 13 Sep 2026 02:23:34 -0700 Subject: [PATCH 2/6] fix(sdk): carry body and proxy headers on relayed litellm errors and content policy blocks too --- litellm/exceptions.py | 2 + .../exception_mapping_utils.py | 49 +++++++------- .../test_exception_mapping_utils.py | 65 +++++++++++++------ .../common_utils/test_openai_error_payload.py | 2 +- 4 files changed, 72 insertions(+), 46 deletions(-) diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 3d9e26e450b..fdc2cc1f169 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -626,6 +626,7 @@ class ContentPolicyViolationError(BadRequestError): litellm_debug_info: str | None = None, provider_specific_fields: dict | None = None, body: dict | None = None, + headers: Mapping[str, str] | None = None, ): self.status_code = 400 self.message = f"litellm.ContentPolicyViolationError: {message}" @@ -640,6 +641,7 @@ class ContentPolicyViolationError(BadRequestError): response=response, litellm_debug_info=self.litellm_debug_info, body=body, + headers=headers, ) # Call the base class constructor with the parameters it needs def __str__(self): diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index fd1ba666887..36b53a26c99 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -1,3 +1,4 @@ +import inspect import json import re import traceback @@ -203,11 +204,18 @@ def _get_response_headers(original_exception: Exception) -> httpx.Headers | None return _response_headers +def _accepted_init_kwargs(exception_class: type[Exception], candidates: Mapping[str, object]) -> Mapping[str, object]: + accepted: Final = inspect.signature(exception_class).parameters + return {name: value for name, value in candidates.items() if name in accepted} + + def extract_and_raise_litellm_exception( response: Any | None, error_str: str, model: str, custom_llm_provider: str, + body: object | None = None, + headers: Mapping[str, str] | None = None, ): """ Covers scenario where litellm sdk calling proxy. @@ -217,32 +225,19 @@ def extract_and_raise_litellm_exception( Relevant Issue: https://github.com/BerriAI/litellm/issues/7259 """ pattern: Final = r"litellm\.\w+Error" - - # Search for the exception in the error string match: Final = re.search(pattern, error_str) - - # Extract the exception if found - if match: - exception_name = match.group(0) - exception_name = exception_name.strip().replace("litellm.", "") - raised_exception_obj: Final = getattr(litellm, exception_name, None) - if raised_exception_obj: - # Try with response parameter first, fall back to without it - # Some exceptions (e.g., APIConnectionError) don't accept response param - try: - raise raised_exception_obj( - message=error_str, - llm_provider=custom_llm_provider, - model=model, - response=response, - ) - except TypeError: - # Exception doesn't accept response parameter - raise raised_exception_obj( - message=error_str, - llm_provider=custom_llm_provider, - model=model, - ) + if match is None: + return + exception_name: Final = match.group(0).removeprefix("litellm.") + raised_exception_obj: Final = getattr(litellm, exception_name, None) + if not raised_exception_obj: + return + raise raised_exception_obj( + message=error_str, + llm_provider=custom_llm_provider, + model=model, + **_accepted_init_kwargs(raised_exception_obj, {"response": response, "body": body, "headers": headers}), + ) class _ProviderHTTPException(Protocol): @@ -339,6 +334,8 @@ def _map_openai_exception( model=model, response=getattr(original_exception, "response", None), litellm_debug_info=extra_information, + body=getattr(original_exception, "body", None), + headers=upstream_headers, ) elif "invalid_encrypted_content" in error_str or "could not be verified" in error_str: helpful_message: Final = ( @@ -2443,6 +2440,8 @@ def exception_type( error_str=error_str, model=model, custom_llm_provider=custom_llm_provider, + body=getattr(original_exception, "body", None), + headers=_litellm_proxy_response_headers(mappable_exception, custom_llm_provider), ) if ( custom_llm_provider == "openai" diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index ea6ac17ad45..1ff2bbb9bdd 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -1424,35 +1424,37 @@ _GUARDRAIL_BLOCK_ERROR = { } -def _openai_handler_error(error_type: str, headers: dict[str, str]) -> OpenAIError: - """What litellm/llms/openai/openai.py raises after the openai SDK rejects a 400: +def _openai_handler_error( + error_type: str, + headers: dict[str, str], + status_code: int = 400, + message: str = _GUARDRAIL_BLOCK_ERROR["message"], +) -> OpenAIError: + """What litellm/llms/openai/openai.py raises after the openai SDK rejects a request: the SDK's str() carries the wire body, and the handler copies headers and body over.""" - wire_error = {**_GUARDRAIL_BLOCK_ERROR, "type": error_type} - wire = httpx.Response( - status_code=400, - headers=headers, - json={"error": wire_error}, - request=httpx.Request("POST", "http://localhost:4000/v1/chat/completions"), - ) + wire_error = {**_GUARDRAIL_BLOCK_ERROR, "type": error_type, "code": str(status_code), "message": message} return OpenAIError( - status_code=400, - message=f"Error code: 400 - {{'error': {wire_error}}}", - headers=wire.headers, + status_code=status_code, + message=f"Error code: {status_code} - {{'error': {wire_error}}}", + headers=httpx.Headers(headers), body=wire_error, ) -@pytest.mark.parametrize("error_type", ["None", "invalid_request_error"]) -def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str): - """An SDK caller behind a proxy tells a guardrail block from any other 400 by the body's - provider_specific_fields and the proxy's x-litellm-* headers, so the mapped BadRequestError - must carry both whichever error.type the proxy version on the other end emits.""" - proxy_headers = {"x-litellm-call-id": "call-guardrail", "x-litellm-applied-guardrails": "block-secret-project"} +_PROXY_HEADERS = {"x-litellm-call-id": "call-guardrail", "x-litellm-applied-guardrails": "block-secret-project"} + +@pytest.mark.parametrize( + ("error_type", "status_code"), [("None", 400), ("invalid_request_error", 400), ("None", 422)] +) +def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str, status_code: int): + """An SDK caller behind a proxy tells a guardrail block from any other 4xx by the body's + provider_specific_fields and the proxy's x-litellm-* headers, so the mapped BadRequestError + must carry both whichever error.type and status the proxy version on the other end emits.""" with pytest.raises(litellm.BadRequestError) as exc_info: exception_type( model="claude-haiku-4-5", - original_exception=_openai_handler_error(error_type, proxy_headers), + original_exception=_openai_handler_error(error_type, _PROXY_HEADERS, status_code=status_code), custom_llm_provider="litellm_proxy", completion_kwargs={}, extra_kwargs={}, @@ -1460,7 +1462,30 @@ def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str): assert exc_info.value.body["provider_specific_fields"]["guardrail_name"] == "block-secret-project" assert exc_info.value.body["type"] == error_type - assert proxy_headers.items() <= exc_info.value.headers.items() + assert exc_info.value.headers == _PROXY_HEADERS + + +@pytest.mark.parametrize( + "relayed_class", [litellm.BadRequestError, litellm.ContentPolicyViolationError] +) +def test_litellm_proxy_relayed_litellm_error_keeps_body_and_headers(relayed_class: type[litellm.BadRequestError]): + """A proxy relaying a provider's own litellm error names the class in the message, which + re-raises that class on the SDK side before the generic 400 mapping runs; it must carry the + body and the proxy headers the same way the generic mapping now does.""" + message = f"litellm.{relayed_class.__name__}: {_GUARDRAIL_BLOCK_ERROR['message']}" + + with pytest.raises(relayed_class) as exc_info: + exception_type( + model="claude-haiku-4-5", + original_exception=_openai_handler_error("None", _PROXY_HEADERS, message=message), + custom_llm_provider="litellm_proxy", + completion_kwargs={}, + extra_kwargs={}, + ) + + assert type(exc_info.value) is relayed_class + assert exc_info.value.body["provider_specific_fields"]["guardrail_name"] == "block-secret-project" + assert exc_info.value.headers == _PROXY_HEADERS def test_openai_compatible_vendor_400_keeps_body_but_not_headers(): diff --git a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py index db9fe3a2253..90f1da84a61 100644 --- a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py +++ b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py @@ -145,7 +145,7 @@ def test_a_status_carried_by_an_exception_drives_the_type_it_reports(): assert openai_error_type(exc, error_status_code(exc, 400)) == "permission_error" -def test_the_stringified_none_an_older_upstream_proxy_sent_is_treated_as_absent(): +def test_a_stringified_none_type_or_param_is_treated_as_absent(): """A proxy fronting a proxy older than 1.102 receives {"type": "None", "param": "None"} on the wire; the SDK now keeps that body on the mapped exception, and re-emitting the literal is the exact bug this module exists to stop.""" From f4f1e2eace561ee3b6ed15e916b0e8a20cffd73c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 21:46:22 -0700 Subject: [PATCH 3/6] refactor(sdk): move the None sentinel to constants and freeze the init kwargs filter --- litellm/constants.py | 2 ++ litellm/exceptions.py | 4 +++- litellm/litellm_core_utils/exception_mapping_utils.py | 7 +++++-- litellm/proxy/common_utils/openai_error_payload.py | 6 +++--- .../litellm_core_utils/test_exception_mapping_utils.py | 10 ---------- .../proxy/common_utils/test_openai_error_payload.py | 3 --- .../proxy/test_common_request_processing.py | 3 --- 7 files changed, 13 insertions(+), 22 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 1211a600747..cbf5efdbca0 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1980,6 +1980,8 @@ UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = ( HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS | ORIGIN_SERVER_HEADERS ) +STRINGIFIED_NONE: Final[str] = "None" + # A retrieved response replays the usage of the call that created it, so pricing these # read/management routes like inference bills the same tokens twice. NON_INFERENCE_CALL_TYPES: Final[frozenset[str]] = frozenset( diff --git a/litellm/exceptions.py b/litellm/exceptions.py index fdc2cc1f169..eb4b5f535ff 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -235,7 +235,9 @@ class BadRequestError(openai.BadRequestError): self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries - self.headers: dict[str, str] | None = {k: str(v) for k, v in headers.items()} if headers else None + self.headers = ( + {k: str(v) for k, v in headers.items()} if headers else None # mutable-ok: the proxy updates it in place + ) # Use response if it's a valid httpx.Response with a request, otherwise use minimal error response # Note: We check _request (not .request property) to avoid RuntimeError when _request is None if ( diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 36b53a26c99..b3dec655092 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -3,6 +3,7 @@ import json import re import traceback from collections.abc import Mapping +from types import MappingProxyType from typing import Any, Final, Protocol, cast import httpx @@ -206,7 +207,7 @@ def _get_response_headers(original_exception: Exception) -> httpx.Headers | None def _accepted_init_kwargs(exception_class: type[Exception], candidates: Mapping[str, object]) -> Mapping[str, object]: accepted: Final = inspect.signature(exception_class).parameters - return {name: value for name, value in candidates.items() if name in accepted} + return MappingProxyType({name: value for name, value in candidates.items() if name in accepted}) def extract_and_raise_litellm_exception( @@ -236,7 +237,9 @@ def extract_and_raise_litellm_exception( message=error_str, llm_provider=custom_llm_provider, model=model, - **_accepted_init_kwargs(raised_exception_obj, {"response": response, "body": body, "headers": headers}), + **_accepted_init_kwargs( + raised_exception_obj, MappingProxyType({"response": response, "body": body, "headers": headers}) + ), ) diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py index 90b3c998247..fe23ab2c4b6 100644 --- a/litellm/proxy/common_utils/openai_error_payload.py +++ b/litellm/proxy/common_utils/openai_error_payload.py @@ -8,7 +8,7 @@ from typing import Final from fastapi import status -_STRINGIFIED_NONE: Final = "None" +from litellm.constants import STRINGIFIED_NONE _OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType( { @@ -37,7 +37,7 @@ def openai_error_type(exc: object, status_code: int) -> str: """OpenAI types ``error.type`` as a required string, so an exception carrying none falls back to the type its status code stands for.""" carried: Final = attribute_of(exc, "type") - if isinstance(carried, str) and carried != _STRINGIFIED_NONE: + if isinstance(carried, str) and carried != STRINGIFIED_NONE: return carried mapped: Final = _OPENAI_ERROR_TYPE_BY_STATUS.get(status_code) if mapped is not None: @@ -51,4 +51,4 @@ def openai_error_param(exc: object) -> str | None: """OpenAI types ``error.param`` as nullable, so an exception carrying none serializes as JSON ``null``.""" carried: Final = attribute_of(exc, "param") - return carried if isinstance(carried, str) and carried != _STRINGIFIED_NONE else None + return carried if isinstance(carried, str) and carried != STRINGIFIED_NONE else None diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index 1ff2bbb9bdd..653c07d06ad 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -1430,8 +1430,6 @@ def _openai_handler_error( status_code: int = 400, message: str = _GUARDRAIL_BLOCK_ERROR["message"], ) -> OpenAIError: - """What litellm/llms/openai/openai.py raises after the openai SDK rejects a request: - the SDK's str() carries the wire body, and the handler copies headers and body over.""" wire_error = {**_GUARDRAIL_BLOCK_ERROR, "type": error_type, "code": str(status_code), "message": message} return OpenAIError( status_code=status_code, @@ -1448,9 +1446,6 @@ _PROXY_HEADERS = {"x-litellm-call-id": "call-guardrail", "x-litellm-applied-guar ("error_type", "status_code"), [("None", 400), ("invalid_request_error", 400), ("None", 422)] ) def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str, status_code: int): - """An SDK caller behind a proxy tells a guardrail block from any other 4xx by the body's - provider_specific_fields and the proxy's x-litellm-* headers, so the mapped BadRequestError - must carry both whichever error.type and status the proxy version on the other end emits.""" with pytest.raises(litellm.BadRequestError) as exc_info: exception_type( model="claude-haiku-4-5", @@ -1469,9 +1464,6 @@ def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str, s "relayed_class", [litellm.BadRequestError, litellm.ContentPolicyViolationError] ) def test_litellm_proxy_relayed_litellm_error_keeps_body_and_headers(relayed_class: type[litellm.BadRequestError]): - """A proxy relaying a provider's own litellm error names the class in the message, which - re-raises that class on the SDK side before the generic 400 mapping runs; it must carry the - body and the proxy headers the same way the generic mapping now does.""" message = f"litellm.{relayed_class.__name__}: {_GUARDRAIL_BLOCK_ERROR['message']}" with pytest.raises(relayed_class) as exc_info: @@ -1489,8 +1481,6 @@ def test_litellm_proxy_relayed_litellm_error_keeps_body_and_headers(relayed_clas def test_openai_compatible_vendor_400_keeps_body_but_not_headers(): - """A vendor's own response headers stay on e.response the way every other mapped provider - error keeps them; only a LiteLLM proxy upstream puts headers on e.headers.""" with pytest.raises(litellm.BadRequestError) as exc_info: exception_type( model="gpt-5.4-mini", diff --git a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py index 90f1da84a61..90850840ab4 100644 --- a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py +++ b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py @@ -146,9 +146,6 @@ def test_a_status_carried_by_an_exception_drives_the_type_it_reports(): def test_a_stringified_none_type_or_param_is_treated_as_absent(): - """A proxy fronting a proxy older than 1.102 receives {"type": "None", "param": "None"} on - the wire; the SDK now keeps that body on the mapped exception, and re-emitting the literal - is the exact bug this module exists to stop.""" from litellm.exceptions import BadRequestError carried = BadRequestError( diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 1f8fa762fbd..ef5741e472a 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -4064,9 +4064,6 @@ class TestHandleLLMApiExceptionFramingHeaders: assert proxy_exc.headers["x-request-id"] == "abc-123" async def test_strips_the_date_and_server_headers_of_an_upstream_litellm_proxy(self): - """A proxy fronting another LiteLLM proxy gets the upstream's date and server - on the mapped exception; forwarding them would duplicate the Date header - uvicorn adds to every response and leak the upstream server identity.""" exc = litellm.BadRequestError( message="Content blocked", llm_provider="litellm_proxy", From 6a635cbb64fdf7f567bb26058321447b92fb859e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 22:22:26 -0700 Subject: [PATCH 4/6] fix(sdk): carry a litellm_proxy error's headers on e.response, not e.headers A mapped litellm_proxy exception now attaches an httpx.Response that carries the proxy's response headers whenever the handler attached a header-less synthetic one, on every status branch and on the relay path. BadRequestError keeps its base-class contract: .headers stays the proxy-supplied channel, so the proxy edge keeps forwarding an upstream proxy's headers under the llm_provider- prefix and the date and server edge change is no longer needed. --- litellm/constants.py | 6 +- litellm/exceptions.py | 8 +- .../exception_mapping_utils.py | 64 ++- .../test_exception_mapping_utils.py | 121 ++--- .../proxy/test_common_request_processing.py | 455 ++++++------------ 5 files changed, 223 insertions(+), 431 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index cbf5efdbca0..09442d6151e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1974,11 +1974,7 @@ BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset( } ) -ORIGIN_SERVER_HEADERS: Final[frozenset[str]] = frozenset({"date", "server"}) - -UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = ( - HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS | ORIGIN_SERVER_HEADERS -) +UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS STRINGIFIED_NONE: Final[str] = "None" diff --git a/litellm/exceptions.py b/litellm/exceptions.py index eb4b5f535ff..3f22a4b2dcd 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -10,7 +10,7 @@ ## LiteLLM versions of the OpenAI Exception Types import enum -from collections.abc import Mapping, Sequence +from collections.abc import Sequence from typing import Any, Final import httpx @@ -226,7 +226,6 @@ class BadRequestError(openai.BadRequestError): max_retries: int | None = None, num_retries: int | None = None, body: dict | None = None, - headers: Mapping[str, str] | None = None, ): self.status_code = 400 self.message = f"litellm.BadRequestError: {message}" @@ -235,9 +234,6 @@ class BadRequestError(openai.BadRequestError): self.litellm_debug_info = litellm_debug_info self.max_retries = max_retries self.num_retries = num_retries - self.headers = ( - {k: str(v) for k, v in headers.items()} if headers else None # mutable-ok: the proxy updates it in place - ) # Use response if it's a valid httpx.Response with a request, otherwise use minimal error response # Note: We check _request (not .request property) to avoid RuntimeError when _request is None if ( @@ -628,7 +624,6 @@ class ContentPolicyViolationError(BadRequestError): litellm_debug_info: str | None = None, provider_specific_fields: dict | None = None, body: dict | None = None, - headers: Mapping[str, str] | None = None, ): self.status_code = 400 self.message = f"litellm.ContentPolicyViolationError: {message}" @@ -643,7 +638,6 @@ class ContentPolicyViolationError(BadRequestError): response=response, litellm_debug_info=self.litellm_debug_info, body=body, - headers=headers, ) # Call the base class constructor with the parameters it needs def __str__(self): diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index b3dec655092..e8406b87777 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -216,7 +216,6 @@ def extract_and_raise_litellm_exception( model: str, custom_llm_provider: str, body: object | None = None, - headers: Mapping[str, str] | None = None, ): """ Covers scenario where litellm sdk calling proxy. @@ -237,9 +236,7 @@ def extract_and_raise_litellm_exception( message=error_str, llm_provider=custom_llm_provider, model=model, - **_accepted_init_kwargs( - raised_exception_obj, MappingProxyType({"response": response, "body": body, "headers": headers}) - ), + **_accepted_init_kwargs(raised_exception_obj, MappingProxyType({"response": response, "body": body})), ) @@ -253,13 +250,20 @@ class _ProviderHTTPException(Protocol): llm_provider: str -def _litellm_proxy_response_headers( +def _litellm_proxy_response( original_exception: _ProviderHTTPException, custom_llm_provider: str -) -> Mapping[str, str] | None: - if custom_llm_provider != "litellm_proxy": - return None +) -> httpx.Response | None: + response: Final = getattr(original_exception, "response", None) + if custom_llm_provider != "litellm_proxy" or not isinstance(response, httpx.Response) or response.headers: + return response headers: Final = getattr(original_exception, "headers", None) - return headers if isinstance(headers, Mapping) else None + if not isinstance(headers, Mapping) or not headers: + return response + return httpx.Response( + status_code=response.status_code, + headers={str(k): str(v) for k, v in headers.items()}, + request=getattr(original_exception, "request", None), + ) def _map_openai_exception( @@ -272,7 +276,7 @@ def _map_openai_exception( exception_provider: str, extra_information: str, ) -> None: - upstream_headers: Final = _litellm_proxy_response_headers(original_exception, custom_llm_provider) + response: Final = _litellm_proxy_response(original_exception, custom_llm_provider) # custom_llm_provider is openai, make it OpenAI message = get_error_message(error_obj=original_exception) if message is None: @@ -301,14 +305,14 @@ def _map_openai_exception( message=f"RateLimitError: {exception_provider} - {message}", model=model, llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), + response=response, ) elif ExceptionCheckers.is_error_str_context_window_exceeded(error_str): raise ContextWindowExceededError( message=f"ContextWindowExceededError: {exception_provider} - {message}", llm_provider=custom_llm_provider, model=model, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, ) elif "invalid_request_error" in error_str and "model_not_found" in error_str: @@ -316,7 +320,7 @@ def _map_openai_exception( message=f"{exception_provider} - {message}", llm_provider=custom_llm_provider, model=model, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, ) elif "A timeout occurred" in error_str: @@ -335,10 +339,9 @@ def _map_openai_exception( message=f"ContentPolicyViolationError: {exception_provider} - {message}", llm_provider=custom_llm_provider, model=model, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, body=getattr(original_exception, "body", None), - headers=upstream_headers, ) elif "invalid_encrypted_content" in error_str or "could not be verified" in error_str: helpful_message: Final = ( @@ -356,20 +359,18 @@ def _map_openai_exception( message=helpful_message, llm_provider=custom_llm_provider, model=model, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, body=getattr(original_exception, "body", None), - headers=upstream_headers, ) elif "invalid_request_error" in error_str and "Incorrect API key provided" not in error_str: raise BadRequestError( message=f"{exception_provider} - {message}", llm_provider=custom_llm_provider, model=model, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, body=getattr(original_exception, "body", None), - headers=upstream_headers, ) elif ( "Web server is returning an unknown error" in error_str @@ -385,7 +386,7 @@ def _map_openai_exception( message=f"RateLimitError: {exception_provider} - {message}", model=model, llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, ) elif ( @@ -396,7 +397,7 @@ def _map_openai_exception( message=f"AuthenticationError: {exception_provider} - {message}", llm_provider=custom_llm_provider, model=model, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, ) elif "Mistral API raised a streaming error" in error_str: @@ -415,17 +416,16 @@ def _map_openai_exception( message=f"{exception_provider} - {message}", llm_provider=custom_llm_provider, model=model, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, body=getattr(original_exception, "body", None), - headers=upstream_headers, ) elif original_exception.status_code == 401: raise AuthenticationError( message=f"AuthenticationError: {exception_provider} - {message}", llm_provider=custom_llm_provider, model=model, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, ) elif original_exception.status_code == 404: @@ -433,7 +433,7 @@ def _map_openai_exception( message=f"NotFoundError: {exception_provider} - {message}", model=model, llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, ) elif original_exception.status_code == 408: @@ -448,17 +448,16 @@ def _map_openai_exception( message=f"{exception_provider} - {message}", model=model, llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, body=getattr(original_exception, "body", None), - headers=upstream_headers, ) elif original_exception.status_code == 429: raise RateLimitError( message=f"RateLimitError: {exception_provider} - {message}", model=model, llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, ) elif original_exception.status_code == 500: @@ -466,7 +465,7 @@ def _map_openai_exception( message=f"InternalServerError: {exception_provider} - {message}", model=model, llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, ) elif original_exception.status_code == 502: @@ -474,7 +473,7 @@ def _map_openai_exception( message=f"BadGatewayError: {exception_provider} - {message}", model=model, llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, ) elif original_exception.status_code == 503: @@ -482,7 +481,7 @@ def _map_openai_exception( message=f"ServiceUnavailableError: {exception_provider} - {message}", model=model, llm_provider=custom_llm_provider, - response=getattr(original_exception, "response", None), + response=response, litellm_debug_info=extra_information, ) elif original_exception.status_code == 504: # gateway timeout error @@ -2439,12 +2438,11 @@ def exception_type( custom_llm_provider == "litellm_proxy" ): # handle special case where calling litellm proxy + exception str contains error message extract_and_raise_litellm_exception( - response=getattr(original_exception, "response", None), + response=_litellm_proxy_response(mappable_exception, custom_llm_provider), error_str=error_str, model=model, custom_llm_provider=custom_llm_provider, body=getattr(original_exception, "body", None), - headers=_litellm_proxy_response_headers(mappable_exception, custom_llm_provider), ) if ( custom_llm_provider == "openai" diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index 653c07d06ad..a4869aac8fe 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -1,4 +1,3 @@ - import httpx import openai import pytest @@ -178,9 +177,7 @@ class TestExceptionCheckers: ] for error_str in error_strings: - result = ExceptionCheckers.is_azure_content_policy_violation_error( - error_str - ) + result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str) assert result is True, f"Should detect policy violation in: {error_str}" def test_is_azure_content_policy_violation_error_case_insensitive(self): @@ -194,12 +191,8 @@ class TestExceptionCheckers: ] for error_str in error_strings: - result = ExceptionCheckers.is_azure_content_policy_violation_error( - error_str - ) - assert ( - result is True - ), f"Should detect policy violation in uppercase: {error_str}" + result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str) + assert result is True, f"Should detect policy violation in uppercase: {error_str}" def test_is_azure_content_policy_violation_error_with_non_policy_errors(self): """Test that non-policy violation errors are not detected as policy violations""" @@ -216,12 +209,8 @@ class TestExceptionCheckers: ] for error_str in error_strings: - result = ExceptionCheckers.is_azure_content_policy_violation_error( - error_str - ) - assert ( - result is False - ), f"Should NOT detect policy violation in: {error_str}" + result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str) + assert result is False, f"Should NOT detect policy violation in: {error_str}" def test_is_azure_content_policy_violation_error_with_partial_matches(self): """Test that partial keyword matches work correctly""" @@ -234,9 +223,7 @@ class TestExceptionCheckers: ] for error_str in positive_cases: - result = ExceptionCheckers.is_azure_content_policy_violation_error( - error_str - ) + result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str) assert result is True, f"Should detect policy violation in: {error_str}" # These should not match even though they contain similar words @@ -248,12 +235,8 @@ class TestExceptionCheckers: ] for error_str in negative_cases: - result = ExceptionCheckers.is_azure_content_policy_violation_error( - error_str - ) - assert ( - result is False - ), f"Should NOT detect policy violation in: {error_str}" + result = ExceptionCheckers.is_azure_content_policy_violation_error(error_str) + assert result is False, f"Should NOT detect policy violation in: {error_str}" gemini_context_window_test_cases = [ @@ -271,12 +254,8 @@ gemini_context_window_test_cases = [ ] -@pytest.mark.parametrize( - "error_message, should_raise_context_window", gemini_context_window_test_cases -) -def test_gemini_context_window_error_mapping( - error_message, should_raise_context_window -): +@pytest.mark.parametrize("error_message, should_raise_context_window", gemini_context_window_test_cases) +def test_gemini_context_window_error_mapping(error_message, should_raise_context_window): """ Tests that the exception_type function correctly maps Gemini's context window exceeded errors to litellm.ContextWindowExceededError. @@ -421,9 +400,7 @@ vertex_rate_limit_test_cases = [ ] -@pytest.mark.parametrize( - "error_message, should_raise_rate_limit", vertex_rate_limit_test_cases -) +@pytest.mark.parametrize("error_message, should_raise_rate_limit", vertex_rate_limit_test_cases) def test_vertex_ai_rate_limit_error_mapping(error_message, should_raise_rate_limit): """ Tests that the exception_type function correctly maps Vertex AI's @@ -458,10 +435,7 @@ class TestGetBodyErrorCode: """Unit tests for _get_body_error_code helper.""" def test_parses_int_code(self): - body = ( - '{"error":{"message":"high demand","type":"upstream_error",' - '"param":"","code":429}}' - ) + body = '{"error":{"message":"high demand","type":"upstream_error","param":"","code":429}}' assert _get_body_error_code(body) == 429 def test_parses_string_code(self): @@ -498,8 +472,7 @@ gemini_body_code_429_test_cases = [ ), ( 503, - '{"error":{"message":"upstream unavailable","type":"upstream_error",' - '"param":"","code":429}}', + '{"error":{"message":"upstream unavailable","type":"upstream_error","param":"","code":429}}', litellm.RateLimitError, "HTTP 503 envelope with body code:429 -> RateLimitError", ), @@ -769,9 +742,7 @@ class _UpstreamHTTPError(Exception): self.message = "upstream failure" self.status_code = status_code self.request = httpx.Request("POST", "https://api.example.com/v1/chat/completions") - self.response = httpx.Response( - status_code=status_code, request=self.request, text="upstream failure" - ) + self.response = httpx.Response(status_code=status_code, request=self.request, text="upstream failure") UPSTREAM_STATUS_CODES = (400, 401, 403, 404, 408, 422, 429, 500, 503) @@ -892,15 +863,13 @@ PROVIDERS_WITHOUT_A_HANDLER = tuple( MINIMAX_401_BODY = ( '{"type":"error","error":{"type":"authorized_error","message":"login fail: Please carry the API secret key ' - "in the 'Authorization' field of the request header (1004)\",\"http_code\":\"401\"}," + 'in the \'Authorization\' field of the request header (1004)","http_code":"401"},' '"request_id":"06ddc9ba97ee6340e38f10e09787f547"}' ) def _expected_for(provider: str, status_code: int) -> tuple[type[Exception], int]: - return DEVIATIONS_FROM_THE_OPENAI_SHAPE.get(provider, {}).get( - status_code, OPENAI_SHAPED[status_code] - ) + return DEVIATIONS_FROM_THE_OPENAI_SHAPE.get(provider, {}).get(status_code, OPENAI_SHAPED[status_code]) @pytest.fixture @@ -910,9 +879,7 @@ def quiet_exception_mapping(monkeypatch): @pytest.mark.parametrize("status_code", UPSTREAM_STATUS_CODES) @pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER) -def test_an_upstream_status_maps_to_one_exception_per_provider( - provider, status_code, quiet_exception_mapping -): +def test_an_upstream_status_maps_to_one_exception_per_provider(provider, status_code, quiet_exception_mapping): expected_class, expected_status = _expected_for(provider, status_code) with pytest.raises(openai.APIError) as raised: @@ -928,9 +895,7 @@ def test_an_upstream_status_maps_to_one_exception_per_provider( @pytest.mark.parametrize("status_code", UPSTREAM_STATUS_CODES) @pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER) -def test_a_mapped_exception_keeps_the_provider_and_model_it_came_from( - provider, status_code, quiet_exception_mapping -): +def test_a_mapped_exception_keeps_the_provider_and_model_it_came_from(provider, status_code, quiet_exception_mapping): with pytest.raises(openai.APIError) as raised: exception_type( model="test-model", @@ -943,12 +908,8 @@ def test_a_mapped_exception_keeps_the_provider_and_model_it_came_from( @pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER) -def test_an_already_mapped_litellm_exception_passes_through_untouched( - provider, quiet_exception_mapping -): - already_mapped = litellm.RateLimitError( - message="already mapped", llm_provider=provider, model="test-model" - ) +def test_an_already_mapped_litellm_exception_passes_through_untouched(provider, quiet_exception_mapping): + already_mapped = litellm.RateLimitError(message="already mapped", llm_provider=provider, model="test-model") returned = exception_type( model="test-model", @@ -961,9 +922,7 @@ def test_an_already_mapped_litellm_exception_passes_through_untouched( @pytest.mark.parametrize("status_code", UPSTREAM_STATUS_CODES) @pytest.mark.parametrize("provider", PROVIDERS_WITHOUT_A_HANDLER) -def test_a_provider_without_a_handler_maps_by_the_upstream_status( - provider, status_code, quiet_exception_mapping -): +def test_a_provider_without_a_handler_maps_by_the_upstream_status(provider, status_code, quiet_exception_mapping): expected_class, expected_status = STATUS_KEYED[status_code] with pytest.raises(openai.APIError) as raised: @@ -1015,9 +974,7 @@ def test_an_unmapped_exception_with_no_model_or_provider_is_a_connection_error(q assert "boom" in raised.value.message -def _raise_and_map( - model: str | None, original_exception: Exception, custom_llm_provider: str | None -) -> None: +def _raise_and_map(model: str | None, original_exception: Exception, custom_llm_provider: str | None) -> None: """Calls exception_type() from inside the except block, as litellm/main.py does, so traceback.format_exc() has a real stack.""" try: @@ -1058,9 +1015,7 @@ def test_an_unmapped_exception_with_no_model_or_provider_message_keeps_traceback CONTEXT_WINDOW_MESSAGE = "This model's maximum context length is 4096 tokens." -CONTENT_POLICY_MESSAGE = ( - '{"error": {"type": "invalid_request_error", "code": "content_policy_violation"}}' -) +CONTENT_POLICY_MESSAGE = '{"error": {"type": "invalid_request_error", "code": "content_policy_violation"}}' TIMEOUT_MESSAGE = "Request timed out." PROVIDERS_THAT_RECOGNISE_A_FULL_CONTEXT_WINDOW = ( @@ -1103,15 +1058,11 @@ class _UpstreamErrorWithMessage(_UpstreamHTTPError): super().__init__(status_code=status_code) self.args = (message,) self.message = message - self.response = httpx.Response( - status_code=status_code, request=self.request, text=message - ) + self.response = httpx.Response(status_code=status_code, request=self.request, text=message) @pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER) -def test_a_full_context_window_reaches_the_caller_as_the_router_needs_it( - provider, quiet_exception_mapping -): +def test_a_full_context_window_reaches_the_caller_as_the_router_needs_it(provider, quiet_exception_mapping): if provider in PROVIDERS_THAT_RECOGNISE_A_FULL_CONTEXT_WINDOW: expected_class, expected_status = litellm.ContextWindowExceededError, 400 else: @@ -1129,9 +1080,7 @@ def test_a_full_context_window_reaches_the_caller_as_the_router_needs_it( @pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER) -def test_a_content_policy_block_reaches_the_caller_as_the_router_needs_it( - provider, quiet_exception_mapping -): +def test_a_content_policy_block_reaches_the_caller_as_the_router_needs_it(provider, quiet_exception_mapping): if provider in PROVIDERS_THAT_RECOGNISE_A_CONTENT_POLICY_BLOCK: expected_class, expected_status = litellm.ContentPolicyViolationError, 400 else: @@ -1149,9 +1098,7 @@ def test_a_content_policy_block_reaches_the_caller_as_the_router_needs_it( @pytest.mark.parametrize("provider", PROVIDERS_WITH_A_HANDLER) -def test_a_timed_out_request_is_a_timeout_for_every_provider( - provider, quiet_exception_mapping -): +def test_a_timed_out_request_is_a_timeout_for_every_provider(provider, quiet_exception_mapping): with pytest.raises(litellm.Timeout) as raised: exception_type( model="test-model", @@ -1442,9 +1389,7 @@ def _openai_handler_error( _PROXY_HEADERS = {"x-litellm-call-id": "call-guardrail", "x-litellm-applied-guardrails": "block-secret-project"} -@pytest.mark.parametrize( - ("error_type", "status_code"), [("None", 400), ("invalid_request_error", 400), ("None", 422)] -) +@pytest.mark.parametrize(("error_type", "status_code"), [("None", 400), ("invalid_request_error", 400), ("None", 422)]) def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str, status_code: int): with pytest.raises(litellm.BadRequestError) as exc_info: exception_type( @@ -1457,12 +1402,10 @@ def test_litellm_proxy_guardrail_block_keeps_body_and_headers(error_type: str, s assert exc_info.value.body["provider_specific_fields"]["guardrail_name"] == "block-secret-project" assert exc_info.value.body["type"] == error_type - assert exc_info.value.headers == _PROXY_HEADERS + assert dict(exc_info.value.response.headers) == _PROXY_HEADERS -@pytest.mark.parametrize( - "relayed_class", [litellm.BadRequestError, litellm.ContentPolicyViolationError] -) +@pytest.mark.parametrize("relayed_class", [litellm.BadRequestError, litellm.ContentPolicyViolationError]) def test_litellm_proxy_relayed_litellm_error_keeps_body_and_headers(relayed_class: type[litellm.BadRequestError]): message = f"litellm.{relayed_class.__name__}: {_GUARDRAIL_BLOCK_ERROR['message']}" @@ -1477,7 +1420,7 @@ def test_litellm_proxy_relayed_litellm_error_keeps_body_and_headers(relayed_clas assert type(exc_info.value) is relayed_class assert exc_info.value.body["provider_specific_fields"]["guardrail_name"] == "block-secret-project" - assert exc_info.value.headers == _PROXY_HEADERS + assert dict(exc_info.value.response.headers) == _PROXY_HEADERS def test_openai_compatible_vendor_400_keeps_body_but_not_headers(): @@ -1491,4 +1434,4 @@ def test_openai_compatible_vendor_400_keeps_body_but_not_headers(): ) assert exc_info.value.body["type"] == "vendor_specific_error" - assert exc_info.value.headers is None + assert not exc_info.value.response.headers diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index ef5741e472a..a33b491fce2 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -127,16 +127,12 @@ class TestProxyBaseLLMRequestProcessing: assert json.loads(result.body) == guardrailed_body @pytest.mark.asyncio - async def test_handle_non_streaming_allm_passthrough_route_forwards_upstream_headers( - self, monkeypatch - ): + async def test_handle_non_streaming_allm_passthrough_route_forwards_upstream_headers(self, monkeypatch): """The guardrail JSON path must forward upstream response headers (e.g. x-amzn-requestid) alongside the x-litellm-* headers, matching the non-guardrail passthrough path, while dropping length headers that no longer match the rewritten body.""" - processing_obj = ProxyBaseLLMRequestProcessing( - data={"custom_llm_provider": "bedrock"} - ) + processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"}) monkeypatch.setattr( processing_obj, "_has_post_call_guardrails_for_passthrough", @@ -176,14 +172,10 @@ class TestProxyBaseLLMRequestProcessing: assert result.headers["content-length"] == str(len(result.body)) @pytest.mark.asyncio - async def test_handle_event_stream_allm_passthrough_route_forwards_upstream_headers( - self, monkeypatch - ): + async def test_handle_event_stream_allm_passthrough_route_forwards_upstream_headers(self, monkeypatch): """The guardrail event-stream branch must also forward upstream response headers alongside the x-litellm-* headers.""" - processing_obj = ProxyBaseLLMRequestProcessing( - data={"custom_llm_provider": "bedrock"} - ) + processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"}) monkeypatch.setattr( processing_obj, "_has_post_call_guardrails_for_passthrough", @@ -225,15 +217,11 @@ class TestProxyBaseLLMRequestProcessing: assert result.headers["x-litellm-call-id"] == "test-call-id" @pytest.mark.asyncio - async def test_handle_non_streaming_allm_passthrough_route_applies_response_headers_hook( - self, monkeypatch - ): + async def test_handle_non_streaming_allm_passthrough_route_applies_response_headers_hook(self, monkeypatch): """Guardrailed non-streaming passthrough responses must include headers injected by post_call_response_headers_hook, matching the headers a non-guardrailed passthrough response would carry.""" - processing_obj = ProxyBaseLLMRequestProcessing( - data={"custom_llm_provider": "bedrock"} - ) + processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"}) monkeypatch.setattr( processing_obj, "_has_post_call_guardrails_for_passthrough", @@ -252,9 +240,7 @@ class TestProxyBaseLLMRequestProcessing: return kwargs["response"] proxy_logging_obj.post_call_success_hook = fake_post_call_success_hook - proxy_logging_obj.post_call_response_headers_hook = AsyncMock( - return_value={"x-litellm-custom": "from-hook"} - ) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={"x-litellm-custom": "from-hook"}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=upstream, @@ -378,9 +364,7 @@ class TestProxyBaseLLMRequestProcessing: json.dumps(persisted_body) @pytest.mark.asyncio - async def test_common_processing_pre_call_logic_arms_auto_router_compression_before_guardrails( - self, monkeypatch - ): + async def test_common_processing_pre_call_logic_arms_auto_router_compression_before_guardrails(self, monkeypatch): """arm_pre_call must run before pre_call_hook: an auto router's own compression policy has to be in `data["metadata"]` (naming the model-side guardrail so it runs even if it isn't default_on) by the time guardrails see the request.""" @@ -2191,16 +2175,10 @@ class TestCommonRequestProcessingHelpers: def _stringified_none_paths(node: object, path: str = "error") -> tuple[str, ...]: if isinstance(node, dict): - return tuple( - found - for key, value in node.items() - for found in _stringified_none_paths(value, f"{path}.{key}") - ) + return tuple(found for key, value in node.items() for found in _stringified_none_paths(value, f"{path}.{key}")) if isinstance(node, (list, tuple)): return tuple( - found - for index, value in enumerate(node) - for found in _stringified_none_paths(value, f"{path}[{index}]") + found for index, value in enumerate(node) for found in _stringified_none_paths(value, f"{path}[{index}]") ) return (path,) if node == "None" else () @@ -2947,9 +2925,7 @@ class TestStreamingOverheadHeader: user_api_key_dict=mock_user_api_key_dict, call_id="test-call-id", hidden_params={}, - litellm_logging_obj=self._timing_logging_obj( - {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} - ), + litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}), ) assert headers["x-litellm-response-duration-ms"] == "500.0" @@ -2968,9 +2944,7 @@ class TestStreamingOverheadHeader: user_api_key_dict=mock_user_api_key_dict, call_id="test-call-id", hidden_params={}, - litellm_logging_obj=self._timing_logging_obj( - {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} - ), + litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}), read_timing_from_logging_obj=False, ) @@ -2991,9 +2965,7 @@ class TestStreamingOverheadHeader: user_api_key_dict=mock_user_api_key_dict, call_id="test-call-id", hidden_params={"_response_ms": 300.0}, - litellm_logging_obj=self._timing_logging_obj( - {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} - ), + litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}), ) assert headers["x-litellm-response-duration-ms"] == "300.0" @@ -3034,9 +3006,7 @@ class TestStreamingOverheadHeader: user_api_key_dict=mock_user_api_key_dict, call_id="test-call-id", hidden_params={"_response_ms": 300.0, "litellm_overhead_time_ms": 7.5}, - litellm_logging_obj=self._timing_logging_obj( - {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} - ), + litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}), ) assert headers["x-litellm-response-duration-ms"] == "300.0" @@ -3488,9 +3458,7 @@ class TestStreamCloseOnDisconnect: finally: closed.set() - response = _UpstreamClosingStreamingResponse( - body(), media_type="text/event-stream" - ) + response = _UpstreamClosingStreamingResponse(body(), media_type="text/event-stream") async def receive(): await asyncio.Event().wait() @@ -3521,9 +3489,7 @@ class TestStreamCloseOnDisconnect: finally: closed.set() - response = _UpstreamClosingStreamingResponse( - body(), media_type="text/event-stream" - ) + response = _UpstreamClosingStreamingResponse(body(), media_type="text/event-stream") async def receive(): await disconnected.wait() @@ -3594,9 +3560,7 @@ class TestStreamCloseOnDisconnect: finally: inner_closed.set() - response = await create_response( - generator=wrapped(), media_type="text/event-stream", headers={} - ) + response = await create_response(generator=wrapped(), media_type="text/event-stream", headers={}) async def receive(): await asyncio.Event().wait() @@ -3828,9 +3792,7 @@ class TestStreamCloseOnDisconnect: with pytest.raises(_ClientDisconnectedBeforeFirstChunk): await asyncio.wait_for( - _buffer_first_chunk_honoring_disconnect( - AcloseRaises(), request=self._request_that_disconnects() - ), + _buffer_first_chunk_honoring_disconnect(AcloseRaises(), request=self._request_that_disconnects()), timeout=5, ) @@ -3846,9 +3808,7 @@ class TestStreamCloseOnDisconnect: with pytest.raises(_ClientDisconnectedBeforeFirstChunk): await asyncio.wait_for( - _buffer_first_chunk_honoring_disconnect( - blocking_gen(), request=self._request_that_disconnects() - ), + _buffer_first_chunk_honoring_disconnect(blocking_gen(), request=self._request_that_disconnects()), timeout=5, ) assert closed.is_set() @@ -3864,9 +3824,7 @@ class TestHandleLLMApiExceptionRetryAfter: user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") proxy_logging_obj = MagicMock() proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) - proxy_logging_obj.post_call_response_headers_hook = AsyncMock( - return_value=callback_headers or {} - ) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=callback_headers or {}) try: await processor._handle_llm_api_exception( @@ -3918,9 +3876,7 @@ class TestHandleLLMApiExceptionRetryAfter: enable_pre_call_checks=False, cooldown_list=[], ) - proxy_exc = await self._invoke( - exc, callback_headers={"retry-after": "", "x-custom": "1"} - ) + proxy_exc = await self._invoke(exc, callback_headers={"retry-after": "", "x-custom": "1"}) assert proxy_exc.headers["retry-after"] == "43" assert proxy_exc.headers["x-custom"] == "1" @@ -4063,18 +4019,6 @@ class TestHandleLLMApiExceptionFramingHeaders: assert proxy_exc.headers["x-custom-safe"] == "1" assert proxy_exc.headers["x-request-id"] == "abc-123" - async def test_strips_the_date_and_server_headers_of_an_upstream_litellm_proxy(self): - exc = litellm.BadRequestError( - message="Content blocked", - llm_provider="litellm_proxy", - model="claude-haiku-4-5", - headers={"date": "Sun, 13 Sep 2026 08:43:51 GMT", "server": "uvicorn", "x-request-id": "abc-123"}, - ) - proxy_exc = await self._invoke(exc) - assert "date" not in proxy_exc.headers - assert "server" not in proxy_exc.headers - assert proxy_exc.headers["x-request-id"] == "abc-123" - class TestAsyncStreamingDataGeneratorFastPath: """Fast/slow path branching in async_streaming_data_generator.""" @@ -4161,9 +4105,7 @@ class TestDisconnectGatherCleanup: return Request(scope={"type": "http", "headers": []}, receive=receive) @pytest.mark.asyncio - async def test_base_process_llm_request_raises_499_on_client_disconnect( - self, monkeypatch - ): + async def test_base_process_llm_request_raises_499_on_client_disconnect(self, monkeypatch): """With cancel_on_disconnect enabled, base_process_llm_request returns 499.""" import asyncio @@ -4192,9 +4134,7 @@ class TestDisconnectGatherCleanup: "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)), ) - monkeypatch.setattr( - processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False) - ) + monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)) with pytest.raises(HTTPException) as exc_info: await processing_obj.base_process_llm_request( @@ -4212,9 +4152,7 @@ class TestDisconnectGatherCleanup: assert "disconnected" in exc_info.value.detail.lower() @pytest.mark.asyncio - async def test_base_process_llm_request_reraises_cancelled_error_without_client_disconnect( - self, monkeypatch - ): + async def test_base_process_llm_request_reraises_cancelled_error_without_client_disconnect(self, monkeypatch): import asyncio import litellm.proxy.common_request_processing as cpr @@ -4239,9 +4177,7 @@ class TestDisconnectGatherCleanup: "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)), ) - monkeypatch.setattr( - processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False) - ) + monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)) monkeypatch.setattr( cpr, "route_request", @@ -4302,9 +4238,7 @@ class TestDisconnectGatherCleanup: "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)), ) - monkeypatch.setattr( - processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False) - ) + monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)) with pytest.raises(HTTPException): await processing_obj.base_process_llm_request( @@ -4355,9 +4289,7 @@ class TestDisconnectGatherCleanup: assert task.done() @pytest.mark.asyncio - async def test_base_process_llm_request_preserves_llm_error_after_gather( - self, monkeypatch - ): + async def test_base_process_llm_request_preserves_llm_error_after_gather(self, monkeypatch): import litellm.proxy.common_request_processing as cpr from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -4386,9 +4318,7 @@ class TestDisconnectGatherCleanup: "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)), ) - monkeypatch.setattr( - processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False) - ) + monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)) mock_request = MagicMock(spec=Request) mock_request.is_disconnected = AsyncMock(return_value=False) @@ -4425,19 +4355,13 @@ class TestStreamingClientDisconnectLogging: "litellm_params": {"metadata": {}}, } - recorded = await _record_streaming_client_disconnect_if_needed( - mock_request, request_data - ) + recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data) assert recorded is True assert request_data["metadata"]["client_disconnected"] is True + assert request_data["metadata"]["error_information"]["error_code"] == "499" assert ( - request_data["metadata"]["error_information"]["error_code"] == "499" - ) - assert ( - mock_logging_obj.model_call_details["litellm_params"]["metadata"][ - "error_information" - ]["error_code"] + mock_logging_obj.model_call_details["litellm_params"]["metadata"]["error_information"]["error_code"] == "499" ) @@ -4451,9 +4375,7 @@ class TestStreamingClientDisconnectLogging: mock_request.is_disconnected = AsyncMock(return_value=False) request_data = {"metadata": {}} - recorded = await _record_streaming_client_disconnect_if_needed( - mock_request, request_data - ) + recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data) assert recorded is False assert "client_disconnected" not in request_data["metadata"] @@ -4478,22 +4400,12 @@ class TestStreamingClientDisconnectLogging: "litellm_params": {"metadata": {}}, } - recorded = await _record_streaming_client_disconnect_if_needed( - mock_request, request_data - ) + recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data) assert recorded is True assert request_data["metadata"]["client_disconnected"] is True - assert ( - mock_logging_obj.model_call_details["litellm_params"]["metadata"][ - "client_disconnected" - ] - is True - ) - assert ( - mock_logging_obj.model_call_details["metadata"]["client_disconnected"] - is True - ) + assert mock_logging_obj.model_call_details["litellm_params"]["metadata"]["client_disconnected"] is True + assert mock_logging_obj.model_call_details["metadata"]["client_disconnected"] is True @pytest.mark.asyncio async def test_record_streaming_client_disconnect_handles_none_request_data_metadata(self): @@ -4509,15 +4421,11 @@ class TestStreamingClientDisconnectLogging: "litellm_params": {"metadata": None}, } - recorded = await _record_streaming_client_disconnect_if_needed( - mock_request, request_data - ) + recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data) assert recorded is True assert request_data["metadata"]["client_disconnected"] is True - assert ( - request_data["litellm_params"]["metadata"]["client_disconnected"] is True - ) + assert request_data["litellm_params"]["metadata"]["client_disconnected"] is True @pytest.mark.asyncio async def test_apply_client_disconnect_metadata_none_returns_early(self): @@ -4528,9 +4436,7 @@ class TestStreamingClientDisconnectLogging: _apply_client_disconnect_metadata(None) @pytest.mark.asyncio - async def test_finalize_streaming_generator_cleanup_fires_deferred_logging( - self, monkeypatch - ): + async def test_finalize_streaming_generator_cleanup_fires_deferred_logging(self, monkeypatch): from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ) @@ -4562,9 +4468,7 @@ class TestStreamingClientDisconnectLogging: assert request_data["metadata"]["error_information"]["error_code"] == "499" @pytest.mark.asyncio - async def test_finalize_streaming_generator_cleanup_skips_disconnect_after_completion( - self, monkeypatch - ): + async def test_finalize_streaming_generator_cleanup_skips_disconnect_after_completion(self, monkeypatch): from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ) @@ -4594,9 +4498,7 @@ class TestStreamingClientDisconnectLogging: assert "client_disconnected" not in request_data["metadata"] @pytest.mark.asyncio - async def test_async_streaming_data_generator_records_499_on_early_aclose( - self, monkeypatch - ): + async def test_async_streaming_data_generator_records_499_on_early_aclose(self, monkeypatch): from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ) @@ -4611,9 +4513,7 @@ class TestStreamingClientDisconnectLogging: yield {"choices": [{"delta": {"content": " there"}}]} mock_proxy_logging = MagicMock(spec=ProxyLogging) - mock_proxy_logging.async_post_call_streaming_iterator_hook = ( - mock_streaming_iterator - ) + mock_proxy_logging.async_post_call_streaming_iterator_hook = mock_streaming_iterator ProxyLogging._callback_capabilities_cache.clear() mock_request = MagicMock(spec=Request) @@ -4624,9 +4524,7 @@ class TestStreamingClientDisconnectLogging: "model": "gemini-2.0-flash", "metadata": {}, "litellm_params": {"metadata": {}}, - "litellm_logging_obj": MagicMock( - model_call_details={"metadata": {}, "litellm_params": {}} - ), + "litellm_logging_obj": MagicMock(model_call_details={"metadata": {}, "litellm_params": {}}), } gen = ProxyBaseLLMRequestProcessing.async_streaming_data_generator( @@ -4645,6 +4543,8 @@ class TestStreamingClientDisconnectLogging: assert request_data["metadata"]["error_information"]["error_code"] == "499" ProxyLogging._callback_capabilities_cache.clear() + + class TestCancelOnDisconnect: """ Coverage for the opt-in `general_settings.cancel_on_disconnect` flag: @@ -4671,23 +4571,17 @@ class TestCancelOnDisconnect: llm_call = asyncio.get_running_loop().create_future() disconnect_event = asyncio.Event() - await _cancel_llm_call_on_client_disconnect( - request, llm_call, disconnect_event - ) + await _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event) assert llm_call.cancelled() assert disconnect_event.is_set() async def test_monitor_is_noop_while_client_stays_connected(self): - request = self._request( - [{"type": "http.request", "body": b"", "more_body": False}] - ) + request = self._request([{"type": "http.request", "body": b"", "more_body": False}]) llm_call = asyncio.get_running_loop().create_future() disconnect_event = asyncio.Event() - monitor = asyncio.create_task( - _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event) - ) + monitor = asyncio.create_task(_cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event)) await asyncio.sleep(0.01) assert not monitor.done() @@ -4706,9 +4600,7 @@ class TestCancelOnDisconnect: llm_call = asyncio.get_running_loop().create_future() disconnect_event = asyncio.Event() - await _cancel_llm_call_on_client_disconnect( - request, llm_call, disconnect_event - ) + await _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event) assert not llm_call.cancelled() assert not disconnect_event.is_set() @@ -4723,9 +4615,7 @@ class TestCancelOnDisconnect: with pytest.raises(asyncio.CancelledError): await _await_llm_call_cancelling_on_disconnect(request, llm_call) - async def _drive_base_process_llm_request( - self, monkeypatch, general_settings: dict, llm_call, request: Request - ): + async def _drive_base_process_llm_request(self, monkeypatch, general_settings: dict, llm_call, request: Request): from litellm.proxy._types import UserAPIKeyAuth logging_obj = MagicMock() @@ -4734,9 +4624,7 @@ class TestCancelOnDisconnect: logging_obj._on_deferred_stream_complete = None logging_obj.cost_breakdown = None - processor = ProxyBaseLLMRequestProcessing( - data={"model": "fake-model", "litellm_logging_obj": logging_obj} - ) + processor = ProxyBaseLLMRequestProcessing(data={"model": "fake-model", "litellm_logging_obj": logging_obj}) proxy_logging_obj = MagicMock(spec=ProxyLogging) proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) @@ -4744,9 +4632,7 @@ class TestCancelOnDisconnect: proxy_logging_obj.post_call_success_hook = AsyncMock( side_effect=lambda data, user_api_key_dict, response: response ) - proxy_logging_obj.post_call_response_headers_hook = AsyncMock( - return_value=None - ) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=None) async def fake_route_request(**kwargs): return llm_call() @@ -4825,9 +4711,7 @@ class TestCancelOnDisconnect: with pytest.raises(ProxyException) as exc_info: await processor._handle_llm_api_exception( - e=HTTPException( - status_code=499, detail="Client disconnected the request" - ), + e=HTTPException(status_code=499, detail="Client disconnected the request"), user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), proxy_logging_obj=proxy_logging_obj, ) @@ -4893,7 +4777,9 @@ class TestAllmPassthroughRoutePostCallGuardrails: proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock()) monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", capture_hook) - with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True): + with patch.object( + ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True + ): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=httpx_response, @@ -4943,7 +4829,9 @@ class TestAllmPassthroughRoutePostCallGuardrails: proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock()) monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", non_dict_hook) - with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True): + with patch.object( + ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True + ): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=httpx_response, @@ -4981,7 +4869,9 @@ class TestAllmPassthroughRoutePostCallGuardrails: hook_spy = AsyncMock() monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", hook_spy) - with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True): + with patch.object( + ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True + ): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=httpx_response, @@ -5022,7 +4912,9 @@ class TestAllmPassthroughRoutePostCallGuardrails: hook_spy = AsyncMock() monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", hook_spy) - with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=False): + with patch.object( + ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=False + ): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=httpx_response, @@ -5134,7 +5026,9 @@ class TestEventStreamAllmPassthroughRoute: "content-length": "99", } - with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True): + with patch.object( + ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True + ): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=mock_response, @@ -5165,9 +5059,7 @@ class TestAllmPassthroughStreamingProviderGate: de-anonymized. """ - def _build_processing_obj( - self, custom_llm_provider: str, endpoint: str = "" - ) -> ProxyBaseLLMRequestProcessing: + def _build_processing_obj(self, custom_llm_provider: str, endpoint: str = "") -> ProxyBaseLLMRequestProcessing: logging_obj = MagicMock() logging_obj.litellm_call_id = "call-123" logging_obj.cost_breakdown = None @@ -5254,14 +5146,17 @@ class TestAllmPassthroughStreamingProviderGate: processing_obj = self._build_processing_obj("anthropic") chunks = [b"chunk-1", b"chunk-2"] - with patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=True, + with ( + patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), + patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=True, + ), ): result = await self._run(processing_obj, monkeypatch, chunks) @@ -5270,27 +5165,27 @@ class TestAllmPassthroughStreamingProviderGate: assert streamed == chunks @pytest.mark.asyncio - async def test_bedrock_converse_stream_is_buffered_through_handler( - self, monkeypatch - ): - processing_obj = self._build_processing_obj( - "bedrock", "model/us.amazon.nova-lite-v1:0/converse-stream" - ) + async def test_bedrock_converse_stream_is_buffered_through_handler(self, monkeypatch): + processing_obj = self._build_processing_obj("bedrock", "model/us.amazon.nova-lite-v1:0/converse-stream") chunks = [b"raw-1", b"raw-2"] - with patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=True, - ), patch( - "litellm.llms.bedrock.passthrough.guardrail_translation.handler." - "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", - new=AsyncMock(return_value=b"modified-body"), - ) as mock_handler: + with ( + patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), + patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=True, + ), + patch( + "litellm.llms.bedrock.passthrough.guardrail_translation.handler." + "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", + new=AsyncMock(return_value=b"modified-body"), + ) as mock_handler, + ): result = await self._run(processing_obj, monkeypatch, chunks) assert isinstance(result, Response) @@ -5306,19 +5201,23 @@ class TestAllmPassthroughStreamingProviderGate: ) chunks = [b"raw-1", b"raw-2"] - with patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=True, - ), patch( - "litellm.llms.bedrock.passthrough.guardrail_translation.handler." - "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", - new=AsyncMock(return_value=b"modified-body"), - ) as mock_handler: + with ( + patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), + patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=True, + ), + patch( + "litellm.llms.bedrock.passthrough.guardrail_translation.handler." + "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", + new=AsyncMock(return_value=b"modified-body"), + ) as mock_handler, + ): result = await self._run(processing_obj, monkeypatch, chunks) assert isinstance(result, StreamingResponse) @@ -5340,14 +5239,17 @@ class TestAllmPassthroughStreamingProviderGate: ) chunks = [b"raw-1", b"raw-2"] - with patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=False, + with ( + patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), + patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=False, + ), ): result = await self._run(processing_obj, monkeypatch, chunks) @@ -5366,14 +5268,17 @@ class TestAllmPassthroughStreamingProviderGate: processing_obj = self._build_processing_obj("anthropic") chunks = [b"chunk-1", b"chunk-2"] - with patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=False, + with ( + patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), + patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=False, + ), ): result = await self._run(processing_obj, monkeypatch, chunks) @@ -5821,9 +5726,7 @@ class TestCostHeadersForCallsPricedAtZero: fastapi_response = Response() processing_obj = ProxyBaseLLMRequestProcessing(data={"litellm_logging_obj": logging_obj}) - with patch.object( - ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails", return_value=False - ): + with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails", return_value=False): await processing_obj.base_process_llm_request( request=MagicMock(spec=Request, headers={}), fastapi_response=fastapi_response, @@ -5894,9 +5797,7 @@ class TestCostHeadersForCallsPricedAtZero: assert breakdown.tool_usage_cost == 0.0 def test_cost_breakdown_stays_empty_for_an_inference_call(self): - breakdown = _get_cost_breakdown_from_logging_obj( - litellm_logging_obj=self._logging_obj(call_type="acompletion") - ) + breakdown = _get_cost_breakdown_from_logging_obj(litellm_logging_obj=self._logging_obj(call_type="acompletion")) assert breakdown == CostBreakdownHeaderValues() @@ -5923,7 +5824,6 @@ class TestCostHeadersForCallsPricedAtZero: class TestPreCallWithFallbacksOnLocalRateLimit: - @pytest.mark.asyncio async def test_fallback_triggered_on_local_rate_limit(self): from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError @@ -6075,9 +5975,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo"]}] user_api_key_dict = MagicMock() - user_api_key_dict.router_settings = { - "fallbacks": [{"gpt-4": ["claude-3-haiku"]}] - } + user_api_key_dict.router_settings = {"fallbacks": [{"gpt-4": ["claude-3-haiku"]}]} with patch.object( processor, @@ -6108,9 +6006,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing - processor = ProxyBaseLLMRequestProcessing( - data={"model": "gpt-4", "disable_fallbacks": True} - ) + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4", "disable_fallbacks": True}) async def mock_pre_call_logic(**kwargs): raise ProxyRateLimitError( @@ -6236,9 +6132,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: # Real per-key per-model TPM limiter + a key carrying the customer's # `model_tpm_limit` metadata (only the primary is capped). - limiter = _PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) + limiter = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) user_api_key_dict = UserAPIKeyAuth( api_key="sk-lit3890", metadata={"model_tpm_limit": {primary_model: 100}}, @@ -6246,10 +6140,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: # Pre-seed the primary's per-model token counter at the cap so the very # next request trips it. The counter key uses the *hashed* api_key. - counter_key = ( - f"{user_api_key_dict.api_key}::{primary_model}" - f"::{precise_minute}::request_count" - ) + counter_key = f"{user_api_key_dict.api_key}::{primary_model}::{precise_minute}::request_count" await limiter.internal_usage_cache.async_set_cache( key=counter_key, value={"current_requests": 0, "current_tpm": 100, "current_rpm": 0}, @@ -6280,9 +6171,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: mock_router = MagicMock() mock_router.fallbacks = [{primary_model: [fallback_model]}] - with patch( - "litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock - ): + with patch("litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock): with patch.object( processor, "common_processing_pre_call_logic", @@ -6312,9 +6201,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: # Sanity-check the premise: the limiter genuinely raises a # ProxyRateLimitError for the capped primary under the frozen clock. - with patch( - "litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock - ): + with patch("litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock): with pytest.raises(ProxyRateLimitError): await limiter.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -6675,16 +6562,12 @@ class TestStreamingClientDisconnectBilling: prompt_tokens=1000, completion_tokens=10, total_tokens=1010, - prompt_tokens_details=PromptTokensDetailsWrapper( - cached_tokens=500 - ), + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=500), ), ) ) - event = await self._bill_and_collect_success_event( - append_openai_style_cached_usage_chunk - ) + event = await self._bill_and_collect_success_event(append_openai_style_cached_usage_chunk) usage = event["response_obj"].usage assert getattr(usage, "cache_read_input_tokens", None) == 500 @@ -7454,9 +7337,7 @@ class TestInjectCostIntoUsageDict: logging_obj.model_call_details["custom_llm_provider"] = "anthropic" assert logging_obj.cost_breakdown is None - model_response = ModelResponse( - usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224) - ) + model_response = ModelResponse(usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224)) cost = ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, logging_obj) assert cost is not None and cost > 0 @@ -7485,9 +7366,7 @@ class TestInjectCostIntoUsageDict: ) existing = logging_obj.cost_breakdown - model_response = ModelResponse( - usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224) - ) + model_response = ModelResponse(usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224)) ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, logging_obj) assert logging_obj.cost_breakdown is existing @@ -7782,9 +7661,7 @@ def test_ttft_keepalive_interval_only_arms_for_a_streaming_request(request_data, @pytest.mark.asyncio @pytest.mark.parametrize("stream_requested, expect_ping", [(True, True), (False, False)]) -async def test_base_process_llm_request_pings_while_the_upstream_call_is_still_running( - stream_requested, expect_ping -): +async def test_base_process_llm_request_pings_while_the_upstream_call_is_still_running(stream_requested, expect_ping): """The wiring, not the helper: every route funnels through this method, and the whole time-to-first-token is spent inside the call it wraps.""" @@ -7930,9 +7807,7 @@ async def test_a_late_failure_is_reported_to_the_failure_hook(): async def record(exc): audited.append(exc) - response = await open_sse_before_first_byte( - slow_failure(), ping_interval_seconds=0.05, on_late_failure=record - ) + response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=record) collected = await _drain(response) assert [type(exc).__name__ for exc in audited] == ["HTTPException"] @@ -7949,9 +7824,7 @@ async def test_a_failing_audit_hook_never_costs_the_client_its_error_frame(): async def broken_hook(exc): raise RuntimeError("the audit backend is down") - response = await open_sse_before_first_byte( - slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook - ) + response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook) collected = await _drain(response) error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip()) @@ -8005,9 +7878,7 @@ async def test_base_process_llm_request_audits_a_failure_that_lands_after_its_ke [(0, False), (None, True)], ids=["operator-hard-disabled-this-deployment", "deployment-says-nothing"], ) -async def test_base_process_llm_request_honours_a_deployment_hard_disable( - deployment_keepalive, expect_ping -): +async def test_base_process_llm_request_honours_a_deployment_hard_disable(deployment_keepalive, expect_ping): """`keepalive_seconds: 0` is documented as a disable a request cannot lift. The funnel has to hand its router to the gate for that to hold before the upstream has answered, since no deployment has served the request yet.""" @@ -8053,9 +7924,7 @@ async def test_a_hook_returning_a_replacement_decides_what_the_client_sees(): async def sanitize(exc): return HTTPException(status_code=502, detail="upstream unavailable") - response = await open_sse_before_first_byte( - slow_failure(), ping_interval_seconds=0.05, on_late_failure=sanitize - ) + response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=sanitize) collected = await _drain(response) error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip()) @@ -8094,9 +7963,7 @@ async def test_a_hook_that_returns_nothing_leaves_the_real_error_intact(): async def audit_only(exc): return None - response = await open_sse_before_first_byte( - slow_failure(), ping_interval_seconds=0.05, on_late_failure=audit_only - ) + response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=audit_only) collected = await _drain(response) error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip()) @@ -8113,9 +7980,7 @@ async def test_a_broken_hook_does_not_replace_the_real_error_with_its_own_bug(): async def broken_hook(exc): raise RuntimeError("the audit backend is down") - response = await open_sse_before_first_byte( - slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook - ) + response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook) collected = await _drain(response) error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip()) @@ -8328,9 +8193,7 @@ class TestStreamingResponseHeadersFollowFallback: proxy_logging_obj.post_call_success_hook = AsyncMock( side_effect=lambda data, user_api_key_dict, response: response ) - proxy_logging_obj.post_call_response_headers_hook = AsyncMock( - return_value={"x-callback-header": "kept"} - ) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={"x-callback-header": "kept"}) async def fake_route_request(**kwargs): async def call(): @@ -8338,9 +8201,7 @@ class TestStreamingResponseHeadersFollowFallback: return call() - monkeypatch.setattr( - litellm.proxy.common_request_processing, "route_request", fake_route_request - ) + monkeypatch.setattr(litellm.proxy.common_request_processing, "route_request", fake_route_request) result = await processor.base_process_llm_request( request=Request(scope={"type": "http", "headers": []}), From 62b2b36ce90e6054a77b5c648c4013bce09c271f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 22:39:19 -0700 Subject: [PATCH 5/6] test(proxy): drop the reformat-only diff of the request processing tests The proxy edge test file no longer carries any test of this change, and the remaining diff was the scoped format gate reflowing the whole file to the 120 limit, so it goes back to the merge base bytes --- .../proxy/test_common_request_processing.py | 443 +++++++++++------- 1 file changed, 285 insertions(+), 158 deletions(-) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index a33b491fce2..812fd8ed47d 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -127,12 +127,16 @@ class TestProxyBaseLLMRequestProcessing: assert json.loads(result.body) == guardrailed_body @pytest.mark.asyncio - async def test_handle_non_streaming_allm_passthrough_route_forwards_upstream_headers(self, monkeypatch): + async def test_handle_non_streaming_allm_passthrough_route_forwards_upstream_headers( + self, monkeypatch + ): """The guardrail JSON path must forward upstream response headers (e.g. x-amzn-requestid) alongside the x-litellm-* headers, matching the non-guardrail passthrough path, while dropping length headers that no longer match the rewritten body.""" - processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"}) + processing_obj = ProxyBaseLLMRequestProcessing( + data={"custom_llm_provider": "bedrock"} + ) monkeypatch.setattr( processing_obj, "_has_post_call_guardrails_for_passthrough", @@ -172,10 +176,14 @@ class TestProxyBaseLLMRequestProcessing: assert result.headers["content-length"] == str(len(result.body)) @pytest.mark.asyncio - async def test_handle_event_stream_allm_passthrough_route_forwards_upstream_headers(self, monkeypatch): + async def test_handle_event_stream_allm_passthrough_route_forwards_upstream_headers( + self, monkeypatch + ): """The guardrail event-stream branch must also forward upstream response headers alongside the x-litellm-* headers.""" - processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"}) + processing_obj = ProxyBaseLLMRequestProcessing( + data={"custom_llm_provider": "bedrock"} + ) monkeypatch.setattr( processing_obj, "_has_post_call_guardrails_for_passthrough", @@ -217,11 +225,15 @@ class TestProxyBaseLLMRequestProcessing: assert result.headers["x-litellm-call-id"] == "test-call-id" @pytest.mark.asyncio - async def test_handle_non_streaming_allm_passthrough_route_applies_response_headers_hook(self, monkeypatch): + async def test_handle_non_streaming_allm_passthrough_route_applies_response_headers_hook( + self, monkeypatch + ): """Guardrailed non-streaming passthrough responses must include headers injected by post_call_response_headers_hook, matching the headers a non-guardrailed passthrough response would carry.""" - processing_obj = ProxyBaseLLMRequestProcessing(data={"custom_llm_provider": "bedrock"}) + processing_obj = ProxyBaseLLMRequestProcessing( + data={"custom_llm_provider": "bedrock"} + ) monkeypatch.setattr( processing_obj, "_has_post_call_guardrails_for_passthrough", @@ -240,7 +252,9 @@ class TestProxyBaseLLMRequestProcessing: return kwargs["response"] proxy_logging_obj.post_call_success_hook = fake_post_call_success_hook - proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={"x-litellm-custom": "from-hook"}) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value={"x-litellm-custom": "from-hook"} + ) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=upstream, @@ -364,7 +378,9 @@ class TestProxyBaseLLMRequestProcessing: json.dumps(persisted_body) @pytest.mark.asyncio - async def test_common_processing_pre_call_logic_arms_auto_router_compression_before_guardrails(self, monkeypatch): + async def test_common_processing_pre_call_logic_arms_auto_router_compression_before_guardrails( + self, monkeypatch + ): """arm_pre_call must run before pre_call_hook: an auto router's own compression policy has to be in `data["metadata"]` (naming the model-side guardrail so it runs even if it isn't default_on) by the time guardrails see the request.""" @@ -2175,10 +2191,16 @@ class TestCommonRequestProcessingHelpers: def _stringified_none_paths(node: object, path: str = "error") -> tuple[str, ...]: if isinstance(node, dict): - return tuple(found for key, value in node.items() for found in _stringified_none_paths(value, f"{path}.{key}")) + return tuple( + found + for key, value in node.items() + for found in _stringified_none_paths(value, f"{path}.{key}") + ) if isinstance(node, (list, tuple)): return tuple( - found for index, value in enumerate(node) for found in _stringified_none_paths(value, f"{path}[{index}]") + found + for index, value in enumerate(node) + for found in _stringified_none_paths(value, f"{path}[{index}]") ) return (path,) if node == "None" else () @@ -2925,7 +2947,9 @@ class TestStreamingOverheadHeader: user_api_key_dict=mock_user_api_key_dict, call_id="test-call-id", hidden_params={}, - litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}), + litellm_logging_obj=self._timing_logging_obj( + {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} + ), ) assert headers["x-litellm-response-duration-ms"] == "500.0" @@ -2944,7 +2968,9 @@ class TestStreamingOverheadHeader: user_api_key_dict=mock_user_api_key_dict, call_id="test-call-id", hidden_params={}, - litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}), + litellm_logging_obj=self._timing_logging_obj( + {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} + ), read_timing_from_logging_obj=False, ) @@ -2965,7 +2991,9 @@ class TestStreamingOverheadHeader: user_api_key_dict=mock_user_api_key_dict, call_id="test-call-id", hidden_params={"_response_ms": 300.0}, - litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}), + litellm_logging_obj=self._timing_logging_obj( + {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} + ), ) assert headers["x-litellm-response-duration-ms"] == "300.0" @@ -3006,7 +3034,9 @@ class TestStreamingOverheadHeader: user_api_key_dict=mock_user_api_key_dict, call_id="test-call-id", hidden_params={"_response_ms": 300.0, "litellm_overhead_time_ms": 7.5}, - litellm_logging_obj=self._timing_logging_obj({"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5}), + litellm_logging_obj=self._timing_logging_obj( + {"_response_ms": 500.0, "litellm_overhead_time_ms": 42.5} + ), ) assert headers["x-litellm-response-duration-ms"] == "300.0" @@ -3458,7 +3488,9 @@ class TestStreamCloseOnDisconnect: finally: closed.set() - response = _UpstreamClosingStreamingResponse(body(), media_type="text/event-stream") + response = _UpstreamClosingStreamingResponse( + body(), media_type="text/event-stream" + ) async def receive(): await asyncio.Event().wait() @@ -3489,7 +3521,9 @@ class TestStreamCloseOnDisconnect: finally: closed.set() - response = _UpstreamClosingStreamingResponse(body(), media_type="text/event-stream") + response = _UpstreamClosingStreamingResponse( + body(), media_type="text/event-stream" + ) async def receive(): await disconnected.wait() @@ -3560,7 +3594,9 @@ class TestStreamCloseOnDisconnect: finally: inner_closed.set() - response = await create_response(generator=wrapped(), media_type="text/event-stream", headers={}) + response = await create_response( + generator=wrapped(), media_type="text/event-stream", headers={} + ) async def receive(): await asyncio.Event().wait() @@ -3792,7 +3828,9 @@ class TestStreamCloseOnDisconnect: with pytest.raises(_ClientDisconnectedBeforeFirstChunk): await asyncio.wait_for( - _buffer_first_chunk_honoring_disconnect(AcloseRaises(), request=self._request_that_disconnects()), + _buffer_first_chunk_honoring_disconnect( + AcloseRaises(), request=self._request_that_disconnects() + ), timeout=5, ) @@ -3808,7 +3846,9 @@ class TestStreamCloseOnDisconnect: with pytest.raises(_ClientDisconnectedBeforeFirstChunk): await asyncio.wait_for( - _buffer_first_chunk_honoring_disconnect(blocking_gen(), request=self._request_that_disconnects()), + _buffer_first_chunk_honoring_disconnect( + blocking_gen(), request=self._request_that_disconnects() + ), timeout=5, ) assert closed.is_set() @@ -3824,7 +3864,9 @@ class TestHandleLLMApiExceptionRetryAfter: user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") proxy_logging_obj = MagicMock() proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) - proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=callback_headers or {}) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value=callback_headers or {} + ) try: await processor._handle_llm_api_exception( @@ -3876,7 +3918,9 @@ class TestHandleLLMApiExceptionRetryAfter: enable_pre_call_checks=False, cooldown_list=[], ) - proxy_exc = await self._invoke(exc, callback_headers={"retry-after": "", "x-custom": "1"}) + proxy_exc = await self._invoke( + exc, callback_headers={"retry-after": "", "x-custom": "1"} + ) assert proxy_exc.headers["retry-after"] == "43" assert proxy_exc.headers["x-custom"] == "1" @@ -4105,7 +4149,9 @@ class TestDisconnectGatherCleanup: return Request(scope={"type": "http", "headers": []}, receive=receive) @pytest.mark.asyncio - async def test_base_process_llm_request_raises_499_on_client_disconnect(self, monkeypatch): + async def test_base_process_llm_request_raises_499_on_client_disconnect( + self, monkeypatch + ): """With cancel_on_disconnect enabled, base_process_llm_request returns 499.""" import asyncio @@ -4134,7 +4180,9 @@ class TestDisconnectGatherCleanup: "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)), ) - monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)) + monkeypatch.setattr( + processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False) + ) with pytest.raises(HTTPException) as exc_info: await processing_obj.base_process_llm_request( @@ -4152,7 +4200,9 @@ class TestDisconnectGatherCleanup: assert "disconnected" in exc_info.value.detail.lower() @pytest.mark.asyncio - async def test_base_process_llm_request_reraises_cancelled_error_without_client_disconnect(self, monkeypatch): + async def test_base_process_llm_request_reraises_cancelled_error_without_client_disconnect( + self, monkeypatch + ): import asyncio import litellm.proxy.common_request_processing as cpr @@ -4177,7 +4227,9 @@ class TestDisconnectGatherCleanup: "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)), ) - monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)) + monkeypatch.setattr( + processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False) + ) monkeypatch.setattr( cpr, "route_request", @@ -4238,7 +4290,9 @@ class TestDisconnectGatherCleanup: "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)), ) - monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)) + monkeypatch.setattr( + processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False) + ) with pytest.raises(HTTPException): await processing_obj.base_process_llm_request( @@ -4289,7 +4343,9 @@ class TestDisconnectGatherCleanup: assert task.done() @pytest.mark.asyncio - async def test_base_process_llm_request_preserves_llm_error_after_gather(self, monkeypatch): + async def test_base_process_llm_request_preserves_llm_error_after_gather( + self, monkeypatch + ): import litellm.proxy.common_request_processing as cpr from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -4318,7 +4374,9 @@ class TestDisconnectGatherCleanup: "common_processing_pre_call_logic", AsyncMock(return_value=({"model": "gemini-2.0-flash"}, mock_logging_obj)), ) - monkeypatch.setattr(processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False)) + monkeypatch.setattr( + processing_obj, "_has_post_call_guardrails", MagicMock(return_value=False) + ) mock_request = MagicMock(spec=Request) mock_request.is_disconnected = AsyncMock(return_value=False) @@ -4355,13 +4413,19 @@ class TestStreamingClientDisconnectLogging: "litellm_params": {"metadata": {}}, } - recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data) + recorded = await _record_streaming_client_disconnect_if_needed( + mock_request, request_data + ) assert recorded is True assert request_data["metadata"]["client_disconnected"] is True - assert request_data["metadata"]["error_information"]["error_code"] == "499" assert ( - mock_logging_obj.model_call_details["litellm_params"]["metadata"]["error_information"]["error_code"] + request_data["metadata"]["error_information"]["error_code"] == "499" + ) + assert ( + mock_logging_obj.model_call_details["litellm_params"]["metadata"][ + "error_information" + ]["error_code"] == "499" ) @@ -4375,7 +4439,9 @@ class TestStreamingClientDisconnectLogging: mock_request.is_disconnected = AsyncMock(return_value=False) request_data = {"metadata": {}} - recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data) + recorded = await _record_streaming_client_disconnect_if_needed( + mock_request, request_data + ) assert recorded is False assert "client_disconnected" not in request_data["metadata"] @@ -4400,12 +4466,22 @@ class TestStreamingClientDisconnectLogging: "litellm_params": {"metadata": {}}, } - recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data) + recorded = await _record_streaming_client_disconnect_if_needed( + mock_request, request_data + ) assert recorded is True assert request_data["metadata"]["client_disconnected"] is True - assert mock_logging_obj.model_call_details["litellm_params"]["metadata"]["client_disconnected"] is True - assert mock_logging_obj.model_call_details["metadata"]["client_disconnected"] is True + assert ( + mock_logging_obj.model_call_details["litellm_params"]["metadata"][ + "client_disconnected" + ] + is True + ) + assert ( + mock_logging_obj.model_call_details["metadata"]["client_disconnected"] + is True + ) @pytest.mark.asyncio async def test_record_streaming_client_disconnect_handles_none_request_data_metadata(self): @@ -4421,11 +4497,15 @@ class TestStreamingClientDisconnectLogging: "litellm_params": {"metadata": None}, } - recorded = await _record_streaming_client_disconnect_if_needed(mock_request, request_data) + recorded = await _record_streaming_client_disconnect_if_needed( + mock_request, request_data + ) assert recorded is True assert request_data["metadata"]["client_disconnected"] is True - assert request_data["litellm_params"]["metadata"]["client_disconnected"] is True + assert ( + request_data["litellm_params"]["metadata"]["client_disconnected"] is True + ) @pytest.mark.asyncio async def test_apply_client_disconnect_metadata_none_returns_early(self): @@ -4436,7 +4516,9 @@ class TestStreamingClientDisconnectLogging: _apply_client_disconnect_metadata(None) @pytest.mark.asyncio - async def test_finalize_streaming_generator_cleanup_fires_deferred_logging(self, monkeypatch): + async def test_finalize_streaming_generator_cleanup_fires_deferred_logging( + self, monkeypatch + ): from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ) @@ -4468,7 +4550,9 @@ class TestStreamingClientDisconnectLogging: assert request_data["metadata"]["error_information"]["error_code"] == "499" @pytest.mark.asyncio - async def test_finalize_streaming_generator_cleanup_skips_disconnect_after_completion(self, monkeypatch): + async def test_finalize_streaming_generator_cleanup_skips_disconnect_after_completion( + self, monkeypatch + ): from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ) @@ -4498,7 +4582,9 @@ class TestStreamingClientDisconnectLogging: assert "client_disconnected" not in request_data["metadata"] @pytest.mark.asyncio - async def test_async_streaming_data_generator_records_499_on_early_aclose(self, monkeypatch): + async def test_async_streaming_data_generator_records_499_on_early_aclose( + self, monkeypatch + ): from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, ) @@ -4513,7 +4599,9 @@ class TestStreamingClientDisconnectLogging: yield {"choices": [{"delta": {"content": " there"}}]} mock_proxy_logging = MagicMock(spec=ProxyLogging) - mock_proxy_logging.async_post_call_streaming_iterator_hook = mock_streaming_iterator + mock_proxy_logging.async_post_call_streaming_iterator_hook = ( + mock_streaming_iterator + ) ProxyLogging._callback_capabilities_cache.clear() mock_request = MagicMock(spec=Request) @@ -4524,7 +4612,9 @@ class TestStreamingClientDisconnectLogging: "model": "gemini-2.0-flash", "metadata": {}, "litellm_params": {"metadata": {}}, - "litellm_logging_obj": MagicMock(model_call_details={"metadata": {}, "litellm_params": {}}), + "litellm_logging_obj": MagicMock( + model_call_details={"metadata": {}, "litellm_params": {}} + ), } gen = ProxyBaseLLMRequestProcessing.async_streaming_data_generator( @@ -4543,8 +4633,6 @@ class TestStreamingClientDisconnectLogging: assert request_data["metadata"]["error_information"]["error_code"] == "499" ProxyLogging._callback_capabilities_cache.clear() - - class TestCancelOnDisconnect: """ Coverage for the opt-in `general_settings.cancel_on_disconnect` flag: @@ -4571,17 +4659,23 @@ class TestCancelOnDisconnect: llm_call = asyncio.get_running_loop().create_future() disconnect_event = asyncio.Event() - await _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event) + await _cancel_llm_call_on_client_disconnect( + request, llm_call, disconnect_event + ) assert llm_call.cancelled() assert disconnect_event.is_set() async def test_monitor_is_noop_while_client_stays_connected(self): - request = self._request([{"type": "http.request", "body": b"", "more_body": False}]) + request = self._request( + [{"type": "http.request", "body": b"", "more_body": False}] + ) llm_call = asyncio.get_running_loop().create_future() disconnect_event = asyncio.Event() - monitor = asyncio.create_task(_cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event)) + monitor = asyncio.create_task( + _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event) + ) await asyncio.sleep(0.01) assert not monitor.done() @@ -4600,7 +4694,9 @@ class TestCancelOnDisconnect: llm_call = asyncio.get_running_loop().create_future() disconnect_event = asyncio.Event() - await _cancel_llm_call_on_client_disconnect(request, llm_call, disconnect_event) + await _cancel_llm_call_on_client_disconnect( + request, llm_call, disconnect_event + ) assert not llm_call.cancelled() assert not disconnect_event.is_set() @@ -4615,7 +4711,9 @@ class TestCancelOnDisconnect: with pytest.raises(asyncio.CancelledError): await _await_llm_call_cancelling_on_disconnect(request, llm_call) - async def _drive_base_process_llm_request(self, monkeypatch, general_settings: dict, llm_call, request: Request): + async def _drive_base_process_llm_request( + self, monkeypatch, general_settings: dict, llm_call, request: Request + ): from litellm.proxy._types import UserAPIKeyAuth logging_obj = MagicMock() @@ -4624,7 +4722,9 @@ class TestCancelOnDisconnect: logging_obj._on_deferred_stream_complete = None logging_obj.cost_breakdown = None - processor = ProxyBaseLLMRequestProcessing(data={"model": "fake-model", "litellm_logging_obj": logging_obj}) + processor = ProxyBaseLLMRequestProcessing( + data={"model": "fake-model", "litellm_logging_obj": logging_obj} + ) proxy_logging_obj = MagicMock(spec=ProxyLogging) proxy_logging_obj.during_call_hook = AsyncMock(return_value=None) @@ -4632,7 +4732,9 @@ class TestCancelOnDisconnect: proxy_logging_obj.post_call_success_hook = AsyncMock( side_effect=lambda data, user_api_key_dict, response: response ) - proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=None) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value=None + ) async def fake_route_request(**kwargs): return llm_call() @@ -4711,7 +4813,9 @@ class TestCancelOnDisconnect: with pytest.raises(ProxyException) as exc_info: await processor._handle_llm_api_exception( - e=HTTPException(status_code=499, detail="Client disconnected the request"), + e=HTTPException( + status_code=499, detail="Client disconnected the request" + ), user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), proxy_logging_obj=proxy_logging_obj, ) @@ -4777,9 +4881,7 @@ class TestAllmPassthroughRoutePostCallGuardrails: proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock()) monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", capture_hook) - with patch.object( - ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True - ): + with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=httpx_response, @@ -4829,9 +4931,7 @@ class TestAllmPassthroughRoutePostCallGuardrails: proxy_logging_obj = ProxyLogging(user_api_key_cache=MagicMock()) monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", non_dict_hook) - with patch.object( - ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True - ): + with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=httpx_response, @@ -4869,9 +4969,7 @@ class TestAllmPassthroughRoutePostCallGuardrails: hook_spy = AsyncMock() monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", hook_spy) - with patch.object( - ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True - ): + with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=httpx_response, @@ -4912,9 +5010,7 @@ class TestAllmPassthroughRoutePostCallGuardrails: hook_spy = AsyncMock() monkeypatch.setattr(proxy_logging_obj, "post_call_success_hook", hook_spy) - with patch.object( - ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=False - ): + with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=False): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=httpx_response, @@ -5026,9 +5122,7 @@ class TestEventStreamAllmPassthroughRoute: "content-length": "99", } - with patch.object( - ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True - ): + with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails_for_passthrough", return_value=True): processing_obj = ProxyBaseLLMRequestProcessing(data={}) result = await processing_obj._handle_non_streaming_allm_passthrough_route( response=mock_response, @@ -5059,7 +5153,9 @@ class TestAllmPassthroughStreamingProviderGate: de-anonymized. """ - def _build_processing_obj(self, custom_llm_provider: str, endpoint: str = "") -> ProxyBaseLLMRequestProcessing: + def _build_processing_obj( + self, custom_llm_provider: str, endpoint: str = "" + ) -> ProxyBaseLLMRequestProcessing: logging_obj = MagicMock() logging_obj.litellm_call_id = "call-123" logging_obj.cost_breakdown = None @@ -5146,17 +5242,14 @@ class TestAllmPassthroughStreamingProviderGate: processing_obj = self._build_processing_obj("anthropic") chunks = [b"chunk-1", b"chunk-2"] - with ( - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=True, - ), + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=True, ): result = await self._run(processing_obj, monkeypatch, chunks) @@ -5165,27 +5258,27 @@ class TestAllmPassthroughStreamingProviderGate: assert streamed == chunks @pytest.mark.asyncio - async def test_bedrock_converse_stream_is_buffered_through_handler(self, monkeypatch): - processing_obj = self._build_processing_obj("bedrock", "model/us.amazon.nova-lite-v1:0/converse-stream") + async def test_bedrock_converse_stream_is_buffered_through_handler( + self, monkeypatch + ): + processing_obj = self._build_processing_obj( + "bedrock", "model/us.amazon.nova-lite-v1:0/converse-stream" + ) chunks = [b"raw-1", b"raw-2"] - with ( - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=True, - ), - patch( - "litellm.llms.bedrock.passthrough.guardrail_translation.handler." - "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", - new=AsyncMock(return_value=b"modified-body"), - ) as mock_handler, - ): + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=True, + ), patch( + "litellm.llms.bedrock.passthrough.guardrail_translation.handler." + "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", + new=AsyncMock(return_value=b"modified-body"), + ) as mock_handler: result = await self._run(processing_obj, monkeypatch, chunks) assert isinstance(result, Response) @@ -5201,23 +5294,19 @@ class TestAllmPassthroughStreamingProviderGate: ) chunks = [b"raw-1", b"raw-2"] - with ( - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=True, - ), - patch( - "litellm.llms.bedrock.passthrough.guardrail_translation.handler." - "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", - new=AsyncMock(return_value=b"modified-body"), - ) as mock_handler, - ): + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=True, + ), patch( + "litellm.llms.bedrock.passthrough.guardrail_translation.handler." + "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", + new=AsyncMock(return_value=b"modified-body"), + ) as mock_handler: result = await self._run(processing_obj, monkeypatch, chunks) assert isinstance(result, StreamingResponse) @@ -5239,17 +5328,14 @@ class TestAllmPassthroughStreamingProviderGate: ) chunks = [b"raw-1", b"raw-2"] - with ( - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=False, - ), + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=False, ): result = await self._run(processing_obj, monkeypatch, chunks) @@ -5268,17 +5354,14 @@ class TestAllmPassthroughStreamingProviderGate: processing_obj = self._build_processing_obj("anthropic") chunks = [b"chunk-1", b"chunk-2"] - with ( - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails", - return_value=False, - ), - patch.object( - ProxyBaseLLMRequestProcessing, - "_has_post_call_guardrails_for_passthrough", - return_value=False, - ), + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=False, ): result = await self._run(processing_obj, monkeypatch, chunks) @@ -5726,7 +5809,9 @@ class TestCostHeadersForCallsPricedAtZero: fastapi_response = Response() processing_obj = ProxyBaseLLMRequestProcessing(data={"litellm_logging_obj": logging_obj}) - with patch.object(ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails", return_value=False): + with patch.object( + ProxyBaseLLMRequestProcessing, "_has_post_call_guardrails", return_value=False + ): await processing_obj.base_process_llm_request( request=MagicMock(spec=Request, headers={}), fastapi_response=fastapi_response, @@ -5797,7 +5882,9 @@ class TestCostHeadersForCallsPricedAtZero: assert breakdown.tool_usage_cost == 0.0 def test_cost_breakdown_stays_empty_for_an_inference_call(self): - breakdown = _get_cost_breakdown_from_logging_obj(litellm_logging_obj=self._logging_obj(call_type="acompletion")) + breakdown = _get_cost_breakdown_from_logging_obj( + litellm_logging_obj=self._logging_obj(call_type="acompletion") + ) assert breakdown == CostBreakdownHeaderValues() @@ -5824,6 +5911,7 @@ class TestCostHeadersForCallsPricedAtZero: class TestPreCallWithFallbacksOnLocalRateLimit: + @pytest.mark.asyncio async def test_fallback_triggered_on_local_rate_limit(self): from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError @@ -5975,7 +6063,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo"]}] user_api_key_dict = MagicMock() - user_api_key_dict.router_settings = {"fallbacks": [{"gpt-4": ["claude-3-haiku"]}]} + user_api_key_dict.router_settings = { + "fallbacks": [{"gpt-4": ["claude-3-haiku"]}] + } with patch.object( processor, @@ -6006,7 +6096,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing - processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4", "disable_fallbacks": True}) + processor = ProxyBaseLLMRequestProcessing( + data={"model": "gpt-4", "disable_fallbacks": True} + ) async def mock_pre_call_logic(**kwargs): raise ProxyRateLimitError( @@ -6132,7 +6224,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: # Real per-key per-model TPM limiter + a key carrying the customer's # `model_tpm_limit` metadata (only the primary is capped). - limiter = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache())) + limiter = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) user_api_key_dict = UserAPIKeyAuth( api_key="sk-lit3890", metadata={"model_tpm_limit": {primary_model: 100}}, @@ -6140,7 +6234,10 @@ class TestPreCallWithFallbacksOnLocalRateLimit: # Pre-seed the primary's per-model token counter at the cap so the very # next request trips it. The counter key uses the *hashed* api_key. - counter_key = f"{user_api_key_dict.api_key}::{primary_model}::{precise_minute}::request_count" + counter_key = ( + f"{user_api_key_dict.api_key}::{primary_model}" + f"::{precise_minute}::request_count" + ) await limiter.internal_usage_cache.async_set_cache( key=counter_key, value={"current_requests": 0, "current_tpm": 100, "current_rpm": 0}, @@ -6171,7 +6268,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: mock_router = MagicMock() mock_router.fallbacks = [{primary_model: [fallback_model]}] - with patch("litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock): + with patch( + "litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock + ): with patch.object( processor, "common_processing_pre_call_logic", @@ -6201,7 +6300,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: # Sanity-check the premise: the limiter genuinely raises a # ProxyRateLimitError for the capped primary under the frozen clock. - with patch("litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock): + with patch( + "litellm.proxy.hooks.parallel_request_limiter.datetime", _FrozenClock + ): with pytest.raises(ProxyRateLimitError): await limiter.async_pre_call_hook( user_api_key_dict=user_api_key_dict, @@ -6562,12 +6663,16 @@ class TestStreamingClientDisconnectBilling: prompt_tokens=1000, completion_tokens=10, total_tokens=1010, - prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=500), + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=500 + ), ), ) ) - event = await self._bill_and_collect_success_event(append_openai_style_cached_usage_chunk) + event = await self._bill_and_collect_success_event( + append_openai_style_cached_usage_chunk + ) usage = event["response_obj"].usage assert getattr(usage, "cache_read_input_tokens", None) == 500 @@ -7337,7 +7442,9 @@ class TestInjectCostIntoUsageDict: logging_obj.model_call_details["custom_llm_provider"] = "anthropic" assert logging_obj.cost_breakdown is None - model_response = ModelResponse(usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224)) + model_response = ModelResponse( + usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224) + ) cost = ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, logging_obj) assert cost is not None and cost > 0 @@ -7366,7 +7473,9 @@ class TestInjectCostIntoUsageDict: ) existing = logging_obj.cost_breakdown - model_response = ModelResponse(usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224)) + model_response = ModelResponse( + usage=Usage(prompt_tokens=3216, completion_tokens=8, total_tokens=3224) + ) ProxyBaseLLMRequestProcessing._logging_obj_cost_or_none(model_response, logging_obj) assert logging_obj.cost_breakdown is existing @@ -7661,7 +7770,9 @@ def test_ttft_keepalive_interval_only_arms_for_a_streaming_request(request_data, @pytest.mark.asyncio @pytest.mark.parametrize("stream_requested, expect_ping", [(True, True), (False, False)]) -async def test_base_process_llm_request_pings_while_the_upstream_call_is_still_running(stream_requested, expect_ping): +async def test_base_process_llm_request_pings_while_the_upstream_call_is_still_running( + stream_requested, expect_ping +): """The wiring, not the helper: every route funnels through this method, and the whole time-to-first-token is spent inside the call it wraps.""" @@ -7807,7 +7918,9 @@ async def test_a_late_failure_is_reported_to_the_failure_hook(): async def record(exc): audited.append(exc) - response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=record) + response = await open_sse_before_first_byte( + slow_failure(), ping_interval_seconds=0.05, on_late_failure=record + ) collected = await _drain(response) assert [type(exc).__name__ for exc in audited] == ["HTTPException"] @@ -7824,7 +7937,9 @@ async def test_a_failing_audit_hook_never_costs_the_client_its_error_frame(): async def broken_hook(exc): raise RuntimeError("the audit backend is down") - response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook) + response = await open_sse_before_first_byte( + slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook + ) collected = await _drain(response) error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip()) @@ -7878,7 +7993,9 @@ async def test_base_process_llm_request_audits_a_failure_that_lands_after_its_ke [(0, False), (None, True)], ids=["operator-hard-disabled-this-deployment", "deployment-says-nothing"], ) -async def test_base_process_llm_request_honours_a_deployment_hard_disable(deployment_keepalive, expect_ping): +async def test_base_process_llm_request_honours_a_deployment_hard_disable( + deployment_keepalive, expect_ping +): """`keepalive_seconds: 0` is documented as a disable a request cannot lift. The funnel has to hand its router to the gate for that to hold before the upstream has answered, since no deployment has served the request yet.""" @@ -7924,7 +8041,9 @@ async def test_a_hook_returning_a_replacement_decides_what_the_client_sees(): async def sanitize(exc): return HTTPException(status_code=502, detail="upstream unavailable") - response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=sanitize) + response = await open_sse_before_first_byte( + slow_failure(), ping_interval_seconds=0.05, on_late_failure=sanitize + ) collected = await _drain(response) error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip()) @@ -7963,7 +8082,9 @@ async def test_a_hook_that_returns_nothing_leaves_the_real_error_intact(): async def audit_only(exc): return None - response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=audit_only) + response = await open_sse_before_first_byte( + slow_failure(), ping_interval_seconds=0.05, on_late_failure=audit_only + ) collected = await _drain(response) error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip()) @@ -7980,7 +8101,9 @@ async def test_a_broken_hook_does_not_replace_the_real_error_with_its_own_bug(): async def broken_hook(exc): raise RuntimeError("the audit backend is down") - response = await open_sse_before_first_byte(slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook) + response = await open_sse_before_first_byte( + slow_failure(), ping_interval_seconds=0.05, on_late_failure=broken_hook + ) collected = await _drain(response) error_frame = json.loads(collected[-2].decode().removeprefix("data: ").strip()) @@ -8193,7 +8316,9 @@ class TestStreamingResponseHeadersFollowFallback: proxy_logging_obj.post_call_success_hook = AsyncMock( side_effect=lambda data, user_api_key_dict, response: response ) - proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={"x-callback-header": "kept"}) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock( + return_value={"x-callback-header": "kept"} + ) async def fake_route_request(**kwargs): async def call(): @@ -8201,7 +8326,9 @@ class TestStreamingResponseHeadersFollowFallback: return call() - monkeypatch.setattr(litellm.proxy.common_request_processing, "route_request", fake_route_request) + monkeypatch.setattr( + litellm.proxy.common_request_processing, "route_request", fake_route_request + ) result = await processor.base_process_llm_request( request=Request(scope={"type": "http", "headers": []}), From 9fb94ea761f38feb350e5225e6b3467e3d641405 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 14 Sep 2026 22:59:40 -0700 Subject: [PATCH 6/6] fix(exceptions): keep repeated litellm_proxy response headers on the rebuilt response httpx.Headers.items() comma-joins repeated header names, so the rebuilt response iterates multi_items() and keeps every value, matching what the raw openai client exposes on e.response.headers --- .../exception_mapping_utils.py | 3 ++- .../test_exception_mapping_utils.py | 17 ++++++++++++++++- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index e8406b87777..70675966dfc 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -259,9 +259,10 @@ def _litellm_proxy_response( headers: Final = getattr(original_exception, "headers", None) if not isinstance(headers, Mapping) or not headers: return response + pairs: Final = headers.multi_items() if isinstance(headers, httpx.Headers) else headers.items() return httpx.Response( status_code=response.status_code, - headers={str(k): str(v) for k, v in headers.items()}, + headers=[(str(k), str(v)) for k, v in pairs], request=getattr(original_exception, "request", None), ) diff --git a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py index a4869aac8fe..acc6248bf3e 100644 --- a/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_exception_mapping_utils.py @@ -1373,7 +1373,7 @@ _GUARDRAIL_BLOCK_ERROR = { def _openai_handler_error( error_type: str, - headers: dict[str, str], + headers: dict[str, str] | list[tuple[str, str]], status_code: int = 400, message: str = _GUARDRAIL_BLOCK_ERROR["message"], ) -> OpenAIError: @@ -1435,3 +1435,18 @@ def test_openai_compatible_vendor_400_keeps_body_but_not_headers(): assert exc_info.value.body["type"] == "vendor_specific_error" assert not exc_info.value.response.headers + + +def test_litellm_proxy_repeated_response_header_keeps_each_value(): + repeated = [("x-litellm-call-id", "call-guardrail"), ("set-cookie", "a=1"), ("set-cookie", "b=2")] + + with pytest.raises(litellm.BadRequestError) as exc_info: + exception_type( + model="claude-haiku-4-5", + original_exception=_openai_handler_error("None", repeated), + custom_llm_provider="litellm_proxy", + completion_kwargs={}, + extra_kwargs={}, + ) + + assert exc_info.value.response.headers.multi_items() == repeated