From 1fe434dbc33180041a7d3cdc4256043adb1925c7 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 25 May 2026 21:08:52 -0700 Subject: [PATCH 1/8] fix(azure): preserve AD token refresh in v1 OpenAI client path (#28627) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(azure): preserve AD token refresh in v1 OpenAI client path The /openai/v1/ code path (api_version in {"v1", "latest", "preview"}) constructs a plain OpenAI/AsyncOpenAI client, but only forwarded `api_key` from `azure_client_params`. When `enable_azure_ad_token_refresh` is set (or any AD-only auth), `api_key` is None and the client constructor raised "The api_key client option must be set...", breaking every Azure call with a v1 api_version. The OpenAI SDK (>=2.20.0) accepts a callable for `api_key` and re-invokes it on every request via `_refresh_api_key`, so we now forward `azure_ad_token_provider` directly — preserving the per-request token refresh behavior of the regular AzureOpenAI client and avoiding the expiry hole that resolving the token once at client-creation time would introduce. Static `azure_ad_token` strings fall through to `api_key`. For the async path we wrap the sync provider returned by azure-identity in an async function since AsyncOpenAI expects `Callable[[], Awaitable[str]]`. Fixes #27945 https://claude.ai/code/session_01UnzrDSFUUgp5T2wRoPMxq5 * fix(azure): offload sync token provider to thread in v1 async wrapper * fix(azure): include AD credential identity in v1 client cache key --------- Co-authored-by: Claude (cherry picked from commit 96a2e8b16dcad1d6f1175731048b26d4d7b25ad4) --- litellm/llms/azure/common_utils.py | 46 ++- .../llms/azure/test_azure_common_utils.py | 330 ++++++++++++++++++ 2 files changed, 374 insertions(+), 2 deletions(-) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 4fc1ae960b8..e1ac1858912 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -1,3 +1,5 @@ +import asyncio +import hashlib import json import os from typing import Any, Callable, Dict, Literal, NamedTuple, Optional, Union, cast @@ -449,6 +451,25 @@ class BaseAzureLLM(BaseOpenAILLM): ] = None client_initialization_params: dict = locals() client_initialization_params["is_async"] = _is_async + _lp = litellm_params or {} + _ad_provider = _lp.get("azure_ad_token_provider") + _ad_token = _lp.get("azure_ad_token") + _client_secret = _lp.get("client_secret") + _azure_password = _lp.get("azure_password") + client_initialization_params["azure_ad_token"] = ( + hashlib.sha256(_ad_token.encode()).hexdigest() + if isinstance(_ad_token, str) + else None + ) + client_initialization_params["azure_ad_token_provider"] = ( + f"provider_id={id(_ad_provider) if callable(_ad_provider) else None}" + f"|tenant_id={_lp.get('tenant_id')}" + f"|client_id={_lp.get('client_id')}" + f"|client_secret={hashlib.sha256(_client_secret.encode()).hexdigest() if isinstance(_client_secret, str) else None}" + f"|azure_username={_lp.get('azure_username')}" + f"|azure_password={hashlib.sha256(_azure_password.encode()).hexdigest() if isinstance(_azure_password, str) else None}" + f"|azure_scope={_lp.get('azure_scope')}" + ) if client is None: cached_client = self.get_cached_openai_client( client_initialization_params=client_initialization_params, @@ -474,8 +495,29 @@ class BaseAzureLLM(BaseOpenAILLM): if self._is_azure_v1_api_version(api_version): # Extract only params that OpenAI client accepts # Always use /openai/v1/ regardless of whether user passed "v1", "latest", or "preview" - v1_params = { - "api_key": azure_client_params.get("api_key"), + # The OpenAI client accepts a callable for `api_key` and re-invokes it + # on every request (via `_refresh_api_key`), so passing + # `azure_ad_token_provider` directly preserves Azure AD token refresh + # behavior that the regular AzureOpenAI client provides. + v1_api_key: Optional[Union[str, Callable[[], Any]]] = ( + azure_client_params.get("api_key") + or azure_client_params.get("azure_ad_token_provider") + or azure_client_params.get("azure_ad_token") + ) + if _is_async is True and callable(v1_api_key): + # AsyncOpenAI expects an async provider; wrap the sync provider + # returned by azure-identity. Offload to a thread so a token + # refresh (blocking HTTP call to AAD on cache miss) does not + # stall the event loop. + _sync_provider = v1_api_key + + async def _async_v1_api_key() -> str: + return await asyncio.to_thread(_sync_provider) + + v1_api_key = _async_v1_api_key + + v1_params: Dict[str, Any] = { + "api_key": v1_api_key, "base_url": f"{api_base}/openai/v1/", } if "timeout" in azure_client_params: diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index 3fa794375e7..413241adf37 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -1646,6 +1646,336 @@ def test_azure_v1_api_uses_openai_client(api_version): ), f"base_url should contain /openai/v1/, got {async_client.base_url}" +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_api_with_azure_ad_token_provider(api_version): + """ + The v1 OpenAI client path must forward `azure_ad_token_provider` so Azure AD + auth works for `api_version` in {"v1", "latest", "preview"}. + + Regression: https://github.com/BerriAI/litellm/issues/27945 — before the fix + the v1 branch only forwarded `api_key`, so AD-only configs raised + "The api_key client option must be set" on every request. + + The OpenAI SDK accepts a callable for `api_key` and re-invokes it on every + request, so passing the provider directly preserves token refresh. + """ + from openai import AsyncOpenAI, OpenAI + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + token_value = "mock-azure-ad-token-from-provider" + + def token_provider(): + return token_value + + init_return = { + "api_key": None, + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": None, + "azure_ad_token_provider": token_provider, + } + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = init_return + + client = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + _is_async=False, + ) + + assert isinstance(client, OpenAI) + # The SDK stores callables as `_api_key_provider` and refreshes + # `self.api_key` before each request. + assert client._api_key_provider is token_provider + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = init_return + + async_client = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + _is_async=True, + ) + + assert isinstance(async_client, AsyncOpenAI) + # Async client requires an async provider; we wrap the sync provider + # so the SDK can `await` it. + assert async_client._api_key_provider is not None + assert async_client._api_key_provider is not token_provider + + +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_api_async_token_provider_resolves_to_current_token(api_version): + """ + The async wrapper must call the underlying sync provider on each invocation + (not cache its first return value), so token rotation is honored. + """ + import asyncio + + from openai import AsyncOpenAI + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + tokens = iter(["token-1", "token-2", "token-3"]) + + def rotating_provider(): + return next(tokens) + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = { + "api_key": None, + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": None, + "azure_ad_token_provider": rotating_provider, + } + + async_client = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + _is_async=True, + ) + + assert isinstance(async_client, AsyncOpenAI) + loop = asyncio.new_event_loop() + try: + first = loop.run_until_complete(async_client._api_key_provider()) + second = loop.run_until_complete(async_client._api_key_provider()) + finally: + loop.close() + + assert first == "token-1" + assert second == "token-2" + + +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_api_with_static_azure_ad_token(api_version): + """ + When only `azure_ad_token` (a static string) is set, the v1 client should + receive it as `api_key`. + """ + from openai import OpenAI + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + token_value = "static-azure-ad-token" + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = { + "api_key": None, + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": token_value, + "azure_ad_token_provider": None, + } + + client = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + _is_async=False, + ) + + assert isinstance(client, OpenAI) + assert client.api_key == token_value + + +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_api_key_wins_over_ad_token(api_version): + """ + Explicit `api_key` takes precedence over `azure_ad_token_provider` / + `azure_ad_token`, matching the priority documented in + `initialize_azure_sdk_client`. + """ + from openai import OpenAI + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = { + "api_key": "explicit-key", + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": "should-be-ignored", + "azure_ad_token_provider": lambda: "also-ignored", + } + + client = base_llm.get_azure_openai_client( + api_key="explicit-key", + api_base=api_base, + api_version=api_version, + _is_async=False, + ) + + assert isinstance(client, OpenAI) + assert client.api_key == "explicit-key" + assert client._api_key_provider is None + + +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_client_cache_separates_distinct_ad_providers(api_version): + """ + Two configs sharing api_base/api_version but with different AD token + providers must not share a cached OpenAI client, otherwise requests for + one config would be sent with another config's AD credentials. + """ + from openai import AsyncOpenAI + + litellm.in_memory_llm_clients_cache._cache = {} + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + + def provider_a(): + return "token-a" + + def provider_b(): + return "token-b" + + def _init_for(provider): + return { + "api_key": None, + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": None, + "azure_ad_token_provider": provider, + } + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = _init_for(provider_a) + client_a = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + litellm_params={"azure_ad_token_provider": provider_a}, + _is_async=True, + ) + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = _init_for(provider_b) + client_b = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + litellm_params={"azure_ad_token_provider": provider_b}, + _is_async=True, + ) + + assert isinstance(client_a, AsyncOpenAI) + assert isinstance(client_b, AsyncOpenAI) + assert client_a is not client_b + + +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_client_cache_separates_distinct_entra_credentials(api_version): + """ + Configs that synthesize an AD provider from tenant_id/client_id/client_secret + must not share a cached client when those inputs differ. + """ + from openai import AsyncOpenAI + + litellm.in_memory_llm_clients_cache._cache = {} + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + + def synth_provider(): + return "synthesized-token" + + def _init_synth(): + return { + "api_key": None, + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": None, + "azure_ad_token_provider": synth_provider, + } + + common = { + "api_key": None, + "api_base": api_base, + "api_version": api_version, + "_is_async": True, + } + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = _init_synth() + client_a = base_llm.get_azure_openai_client( + litellm_params={ + "tenant_id": "tenant-a", + "client_id": "client-a", + "client_secret": "secret-a", + }, + **common, + ) + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = _init_synth() + client_b = base_llm.get_azure_openai_client( + litellm_params={ + "tenant_id": "tenant-b", + "client_id": "client-b", + "client_secret": "secret-b", + }, + **common, + ) + + assert isinstance(client_a, AsyncOpenAI) + assert isinstance(client_b, AsyncOpenAI) + assert client_a is not client_b + + +@pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) +def test_azure_v1_client_cache_reuses_for_identical_ad_config(api_version): + """ + Identical AD configs should still share a cached client (regression guard + so the cache-key change doesn't accidentally disable caching). + """ + from openai import AsyncOpenAI + + litellm.in_memory_llm_clients_cache._cache = {} + + base_llm = BaseAzureLLM() + api_base = "https://test.openai.azure.com" + + def provider(): + return "tok" + + init_return = { + "api_key": None, + "azure_endpoint": api_base, + "api_version": api_version, + "azure_ad_token": None, + "azure_ad_token_provider": provider, + } + + with patch.object(base_llm, "initialize_azure_sdk_client") as mock_init: + mock_init.return_value = init_return + client_a = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + litellm_params={"azure_ad_token_provider": provider}, + _is_async=True, + ) + client_b = base_llm.get_azure_openai_client( + api_key=None, + api_base=api_base, + api_version=api_version, + litellm_params={"azure_ad_token_provider": provider}, + _is_async=True, + ) + + assert isinstance(client_a, AsyncOpenAI) + assert client_a is client_b + + def test_azure_traditional_api_uses_azure_openai_client(): """ Test that traditional Azure API versions still use AzureOpenAI client. From f18377d3a2f3769dda8631961b6fa72f5172c82b Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Sat, 30 May 2026 08:28:04 +0530 Subject: [PATCH 2/8] fix(proxy): map stripped batch body.model to proxy alias for auth (#29264) * fix(proxy): map stripped batch body.model to proxy alias for auth replace_model_in_jsonl rewrites JSONL body.model to the provider id before upload; batch file access checks must resolve that id back to model_name so keys granted the proxy alias are not rejected with 403. Co-authored-by: Cursor * fix(proxy): surface resolved proxy alias in batch file 403 detail --------- Co-authored-by: Cursor Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> (cherry picked from commit 70d2748d802738df804568b5f85533c9d0afa4ea) --- litellm/proxy/hooks/batch_rate_limiter.py | 10 ++++- .../proxy/hooks/test_batch_file_validation.py | 42 ++++++++++++++++++- 2 files changed, 49 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index f740d5dd40c..1c14e7d751f 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -308,9 +308,15 @@ class _PROXY_BatchRateLimiter(CustomLogger): llm_model_list = llm_router.model_list if llm_router is not None else None for model in models: + # body.model may be the provider id after replace_model_in_jsonl; map to proxy model_name for auth. + model_to_check = model + if llm_router is not None: + proxy_model_name = llm_router.resolve_model_name_from_model_id(model) + if proxy_model_name is not None: + model_to_check = proxy_model_name try: await can_key_call_model( - model=model, + model=model_to_check, llm_model_list=llm_model_list, valid_token=user_api_key_dict, llm_router=llm_router, @@ -326,7 +332,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): detail={ "error": ( "Batch input file references a model the caller is " - f"not authorized to use: model={model}, reason={str(e)}" + f"not authorized to use: model={model_to_check}, reason={str(e)}" ) }, ) diff --git a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py index 7f1006543bb..f047d625479 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py +++ b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py @@ -14,7 +14,6 @@ from fastapi import HTTPException from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth - # --------------------------------------------------------------------------- # Token counter — covers all three batch payload shapes # --------------------------------------------------------------------------- @@ -260,6 +259,47 @@ async def test_pre_call_allows_authorized_model_in_batch_file(): ) +@pytest.mark.asyncio +async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias(): + """After replace_model_in_jsonl, body.model is the provider id (e.g. gpt-5.5). + Auth must check the proxy model_name the key was granted, not the stripped id.""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + proxy_alias = "openai/openai/gpt-5.5-batch" + file_dict = [ + {"body": {"model": "gpt-5.5", "messages": [{"role": "user", "content": "x"}]}} + ] + user = UserAPIKeyAuth( + api_key="sk-ok", + user_id="alice", + models=[proxy_alias], + user_role=LitellmUserRoles.INTERNAL_USER.value, + ) + mock_router = MagicMock() + mock_router.model_list = [] + mock_router.resolve_model_name_from_model_id.return_value = proxy_alias + can_key_call_model = AsyncMock(return_value=True) + + with ( + patch( + "litellm.proxy.auth.auth_checks.can_key_call_model", + new=can_key_call_model, + ), + patch("litellm.proxy.proxy_server.llm_router", mock_router), + ): + await rate_limiter._enforce_batch_file_model_access( + user_api_key_dict=user, + file_content_as_dict=file_dict, + ) + + can_key_call_model.assert_awaited_once() + assert can_key_call_model.await_args.kwargs["model"] == proxy_alias + + @pytest.mark.asyncio async def test_pre_call_skips_check_when_no_models_present(): """Files without any `body.model` (corrupt or empty) must not 500; From b281a9c4fabdd8958be329277de05ba57c4e09ea Mon Sep 17 00:00:00 2001 From: Shivam Rawat Date: Tue, 2 Jun 2026 19:31:36 -0700 Subject: [PATCH 3/8] fix(proxy): resolve managed video model ids for auth (#29545) * fix(proxy): resolve managed video model ids for auth Co-authored-by: Cursor * test(proxy): cover character_id router model resolution Co-authored-by: Cursor --------- Co-authored-by: Cursor (cherry picked from commit d45e9e4d5605ec92d4f798f62b30486ef5d413ec) --- litellm/proxy/auth/auth_checks.py | 1 + litellm/proxy/auth/auth_utils.py | 35 ++++++++++-- litellm/proxy/auth/user_api_key_auth.py | 8 +++ .../spend_tracking/budget_reservation.py | 4 +- .../proxy/auth/test_auth_utils.py | 56 +++++++++++++++++++ 5 files changed, 96 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 0b30999aa21..d4fd76f3841 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -494,6 +494,7 @@ async def common_checks( # noqa: PLR0915 route=route, request_headers=_safe_get_request_headers(request=request), request_query_params=_safe_get_request_query_params(request=request), + llm_router=llm_router, ) # 1. If team is blocked diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 28c58c9af6e..ee6be24ea59 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1186,7 +1186,9 @@ def _route_uses_model_routing_sources(route: str) -> bool: def _extract_models_from_managed_resource_id( - resource_id: Any, resource_id_field: Optional[str] = None + resource_id: Any, + resource_id_field: Optional[str] = None, + llm_router: Optional[Router] = None, ) -> List[str]: if not isinstance(resource_id, str) or not resource_id: return [] @@ -1243,16 +1245,18 @@ def _extract_models_from_managed_resource_id( ) if resource_id_field == "video_id": + model_id = decode_video_id_with_provider(resource_id).get("model_id") _append_model_candidates( candidates=candidates, - value=decode_video_id_with_provider(resource_id).get("model_id"), + value=_resolve_model_id_with_router(model_id, llm_router), ) else: + model_id = decode_character_id_with_provider(resource_id).get( + "model_id" + ) _append_model_candidates( candidates=candidates, - value=decode_character_id_with_provider(resource_id).get( - "model_id" - ), + value=_resolve_model_id_with_router(model_id, llm_router), ) except Exception as e: verbose_proxy_logger.debug( @@ -1262,11 +1266,26 @@ def _extract_models_from_managed_resource_id( return _dedupe_model_candidates(candidates) +def _resolve_model_id_with_router( + model_id: Optional[str], llm_router: Optional[Router] +) -> Optional[str]: + if model_id is None or llm_router is None: + return model_id + try: + return llm_router.resolve_model_name_from_model_id(model_id) or model_id + except Exception as e: + verbose_proxy_logger.debug( + "Unable to resolve model_id from managed resource ID: %s", str(e) + ) + return model_id + + def _extract_model_candidates_from_request( request_data: dict, route: str, request_headers: Optional[Mapping[str, Any]] = None, request_query_params: Optional[Mapping[str, Any]] = None, + llm_router: Optional[Router] = None, ) -> List[str]: candidates: List[str] = [] uses_model_routing_sources = _route_uses_model_routing_sources(route=route) @@ -1316,7 +1335,9 @@ def _extract_model_candidates_from_request( _append_model_candidates( candidates, _extract_models_from_managed_resource_id( - request_data.get(field), resource_id_field=field + request_data.get(field), + resource_id_field=field, + llm_router=llm_router, ), ) @@ -1338,12 +1359,14 @@ def get_model_from_request( route: str, request_headers: Optional[Mapping[str, Any]] = None, request_query_params: Optional[Mapping[str, Any]] = None, + llm_router: Optional[Router] = None, ) -> Optional[Union[str, List[str]]]: candidates = _extract_model_candidates_from_request( request_data=request_data, route=route, request_headers=request_headers, request_query_params=request_query_params, + llm_router=llm_router, ) model = _format_model_candidates(candidates) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 03167c5a2dc..8aa1ff051df 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -142,12 +142,14 @@ def _get_model_from_request_context( request_data: dict, route: str, request: Optional[Request], + llm_router: Optional[Any] = None, ) -> Optional[Union[str, List[str]]]: return get_model_from_request( request_data=request_data, route=route, request_headers=_safe_get_request_headers(request=request), request_query_params=_safe_get_request_query_params(request=request), + llm_router=llm_router, ) @@ -970,6 +972,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 request_data=request_data, route=route, request=request, + llm_router=llm_router, ) skip_budget_checks = False if model is not None and llm_router is not None: @@ -1379,6 +1382,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 request_data=request_data, route=route, request=request, + llm_router=llm_router, ) skip_budget_checks = False if model is not None and llm_router is not None: @@ -1507,6 +1511,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 request_data=request_data, route=route, request=request, + llm_router=llm_router, ) current_models = _get_model_names_for_budget_checks( model=current_model @@ -2073,6 +2078,7 @@ def _should_skip_budget_checks( request_data=request_data, route=route, request=request, + llm_router=llm_router, ) if model is not None and llm_router is not None: return _is_model_cost_zero(model=model, llm_router=llm_router) @@ -2347,6 +2353,7 @@ async def _enforce_key_and_fallback_model_access( request_data=request_data, route=route, request=request, + llm_router=llm_router, ) if model is not None: @@ -2488,6 +2495,7 @@ async def _run_post_custom_auth_checks( request_data=request_data, route=route, request=request, + llm_router=llm_router, ) current_models = _get_model_names_for_budget_checks(model=current_model) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 200a17e2368..eb8af3b073e 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -72,7 +72,7 @@ async def reserve_budget_for_request( return None if route in {"/models", "/v1/models", "/utils/token_counter"}: return None - if get_model_from_request(request_body, route) is None: + if get_model_from_request(request_body, route, llm_router=llm_router) is None: return None counters = await _get_budget_counters( @@ -797,7 +797,7 @@ def estimate_request_max_cost( route: str, llm_router: Optional[Router], ) -> Optional[float]: - model = get_model_from_request(request_body, route) + model = get_model_from_request(request_body, route, llm_router=llm_router) if model is None: return None diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 70e8812c99a..957d697432e 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -381,6 +381,62 @@ def test_get_model_from_request_extracts_video_id_model(): ) +def test_get_model_from_request_resolves_video_id_model_with_router(): + from litellm.types.videos.utils import encode_video_id_with_provider + + provider_video_id = ( + "projects/test-project/locations/us-central1/publishers/google/models/" + "veo-3.1-generate-001/operations/operation-id" + ) + video_id = encode_video_id_with_provider( + video_id=provider_video_id, + provider="vertex_ai", + model_id="veo-3.1-generate-001", + ) + llm_router = MagicMock() + llm_router.resolve_model_name_from_model_id.return_value = ( + "gcp/google/veo-3.1-generate-001" + ) + + assert ( + get_model_from_request( + request_data={"video_id": video_id}, + route="/v1/videos/{video_id}", + llm_router=llm_router, + ) + == "gcp/google/veo-3.1-generate-001" + ) + llm_router.resolve_model_name_from_model_id.assert_called_once_with( + "veo-3.1-generate-001" + ) + + +def test_get_model_from_request_resolves_character_id_model_with_router(): + from litellm.types.videos.utils import encode_character_id_with_provider + + character_id = encode_character_id_with_provider( + character_id="character-provider-id", + provider="vertex_ai", + model_id="veo-3.1-generate-001", + ) + llm_router = MagicMock() + llm_router.resolve_model_name_from_model_id.return_value = ( + "gcp/google/veo-3.1-generate-001" + ) + + assert ( + get_model_from_request( + request_data={"character_id": character_id}, + route="/v1/videos/characters/{character_id}", + llm_router=llm_router, + ) + == "gcp/google/veo-3.1-generate-001" + ) + llm_router.resolve_model_name_from_model_id.assert_called_once_with( + "veo-3.1-generate-001" + ) + + def test_get_model_from_request_only_runs_media_decoders_for_matching_fields(): with ( patch( From 97d0a99c7445e2f463820c2ec9053ae615f76643 Mon Sep 17 00:00:00 2001 From: milan-berri Date: Wed, 3 Jun 2026 19:55:45 +0300 Subject: [PATCH 4/8] fix(key_generate): allow team members to create keys on org-scoped teams (#29310) * fix(key_generate): allow team members to create keys on org-scoped teams When a virtual key is created for a team, enterprise logic inherits the team's organization_id onto the key (add_team_organization_id). Since the VERIA-55 org-IDOR fix, /key/generate then required the caller to be an explicit LiteLLM_OrganizationMembership member of that org, returning 403 "Caller is not a member of organization_id=". Admins normally only add users to teams (not orgs), so self-serve key creation regressed for any user on an org-scoped team (regression since v1.84.0-rc.1). Skip the org-membership check when organization_id was inherited from the key's team (organization_id == team_table.organization_id). Team-level authorization already gates this path, so team membership is sufficient. The membership check still runs when a caller assigns an organization_id that did not come from the key's team, preserving the IDOR protection. Adds regression tests covering both the team-inherited (allowed) and foreign-org (still blocked) cases. Co-authored-by: Cursor * test(key_generate): cover mismatched team org IDOR path on generate Add test_generate_key_foreign_org_with_mismatched_team_still_enforces_membership for the case where a team is present but request organization_id differs from team_table.organization_id. Enterprise inheritance is no-op'd in the test so the guard is exercised directly; membership validation must still run. Addresses Greptile review on #29310. Co-authored-by: Cursor --------- Co-authored-by: Cursor (cherry picked from commit b11833c737fe329512828c086090ffeca8a53082) --- .../key_management_endpoints.py | 7 +- .../test_key_management_endpoints.py | 243 ++++++++++++++++++ 2 files changed, 249 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 636d70e4327..8102ca79906 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -863,7 +863,12 @@ async def _common_key_generation_helper( # noqa: PLR0915 user_api_key_dict.user_role is not None and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value ) - if not _is_proxy_admin: + _org_inherited_from_team = ( + team_table is not None + and team_table.organization_id is not None + and data.organization_id == team_table.organization_id + ) + if not _is_proxy_admin and not _org_inherited_from_team: await _validate_caller_can_assign_key_org( user_api_key_dict=user_api_key_dict, organization_id=data.organization_id, diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 333630d8b5e..1818ffd6232 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -2999,6 +2999,249 @@ async def test_generate_key_with_object_permission(): assert "object_permission" not in key_data +@pytest.mark.asyncio +async def test_generate_key_team_member_inherits_org_skips_membership_check(): + """Regression: a team member creating a key for an org-scoped team must not + be blocked by the org-membership check. + + When ``organization_id`` is inherited from the key's team (via + ``apply_enterprise_key_management_params`` -> ``add_team_organization_id``), + the caller already passed team-level authorization. Requiring an explicit + ``LiteLLM_OrganizationMembership`` row on top of that broke the normal admin + workflow (admins only add users to teams). This asserts the org-membership + check is skipped when the org id came from the caller's team. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles + from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _common_key_generation_helper, + ) + + org_id = "org-from-team" + + # Team belongs to an org; caller is a team member but NOT an explicit member + # of that organization (the regression scenario). + mock_team_table = MagicMock() + mock_team_table.organization_id = org_id + mock_team_table.metadata = None + + mock_validate_org = AsyncMock() + mock_generate_key = AsyncMock( + return_value={ + "key": "sk-test-key", + "expires": None, + "user_id": "alice", + "team_id": "team-1", + } + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_mcp_servers_against_team", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_search_tools_against_team", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._validate_caller_can_assign_key_org", + mock_validate_org, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_org_object", + new_callable=AsyncMock, + return_value=MagicMock(litellm_budget_table=None), + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._check_org_key_limits", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + mock_generate_key, + ), + ): + result = await _common_key_generation_helper( + data=GenerateKeyRequest( + user_id="alice", + team_id="team-1", + organization_id=org_id, + ), + user_api_key_dict=UserAPIKeyAuth( + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ), + litellm_changed_by=None, + team_table=mock_team_table, + ) + + # Key creation proceeded for the team member ... + mock_generate_key.assert_awaited_once() + assert result is not None + # ... and the org-membership check was bypassed because organization_id was + # inherited from the caller's team. + mock_validate_org.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_generate_key_foreign_org_without_team_still_enforces_membership(): + """VERIA-55: a caller assigning a key to an organization that was NOT + inherited from a team must still pass the org-membership check. + + This guards the IDOR fix: ``team_table is None`` (or an org id that does not + match the team) means the org id did not come from team context, so the + explicit membership validation must run. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles + from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _common_key_generation_helper, + ) + + foreign_org_id = "someone-elses-org" + + mock_validate_org = AsyncMock() + mock_generate_key = AsyncMock( + return_value={ + "key": "sk-test-key", + "expires": None, + "user_id": "alice", + "team_id": None, + } + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._validate_caller_can_assign_key_org", + mock_validate_org, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_org_object", + new_callable=AsyncMock, + return_value=MagicMock(litellm_budget_table=None), + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._check_org_key_limits", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + mock_generate_key, + ), + ): + await _common_key_generation_helper( + data=GenerateKeyRequest( + user_id="alice", + organization_id=foreign_org_id, + ), + user_api_key_dict=UserAPIKeyAuth( + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ), + litellm_changed_by=None, + team_table=None, + ) + + # No team context -> the org-membership check must still run. + mock_validate_org.assert_awaited_once() + assert mock_validate_org.call_args.kwargs["organization_id"] == foreign_org_id + + +@pytest.mark.asyncio +async def test_generate_key_foreign_org_with_mismatched_team_still_enforces_membership(): + """VERIA-55: when a team is present but its organization_id differs from the + organization_id on the key request, the org-membership check must still run.""" + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.proxy._types import GenerateKeyRequest, LitellmUserRoles + from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _common_key_generation_helper, + ) + + team_org_id = "other-org" + foreign_org_id = "someone-elses-org" + + mock_team_table = MagicMock() + mock_team_table.organization_id = team_org_id + mock_team_table.metadata = None + + mock_validate_org = AsyncMock() + mock_generate_key = AsyncMock( + return_value={ + "key": "sk-test-key", + "expires": None, + "user_id": "alice", + "team_id": "team-1", + } + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_mcp_servers_against_team", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_search_tools_against_team", + new_callable=AsyncMock, + ), + patch( + "litellm_enterprise.proxy.management_endpoints.key_management_endpoints.apply_enterprise_key_management_params", + side_effect=lambda data, team_table: data, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._validate_caller_can_assign_key_org", + mock_validate_org, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_org_object", + new_callable=AsyncMock, + return_value=MagicMock(litellm_budget_table=None), + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints._check_org_key_limits", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + mock_generate_key, + ), + ): + await _common_key_generation_helper( + data=GenerateKeyRequest( + user_id="alice", + team_id="team-1", + organization_id=foreign_org_id, + ), + user_api_key_dict=UserAPIKeyAuth( + user_id="alice", + user_role=LitellmUserRoles.INTERNAL_USER.value, + ), + litellm_changed_by=None, + team_table=mock_team_table, + ) + + mock_validate_org.assert_awaited_once() + assert mock_validate_org.call_args.kwargs["organization_id"] == foreign_org_id + + # ============================================ # Organization Key Limit Tests # ============================================ From 6a57332f20ebbb1e48ed28188427f30e8e9c115a Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 3 Jun 2026 11:34:04 -0700 Subject: [PATCH 5/8] fix(vertex): strip output_config.effort for Vertex Claude models that reject it (Haiku 4.5) (#29585) * fix(vertex): strip output_config.effort for models that reject it Haiku 4.5 on Vertex AI does not support output_config.effort and 400s with "output_config.effort: Extra inputs are not permitted". PR #27074 emptied VERTEX_UNSUPPORTED_OUTPUT_CONFIG_KEYS so effort would forward for Opus/Sonnet 4.6+, but that made the strip unconditional across every Vertex Anthropic model, including ones that don't support it. Claude Code injects effort into its default Messages payload, so `claude --model claude-haiku-4.5` started failing. Make the sanitizer model-aware: drop output_config.effort for models that don't advertise output_config support (or any reasoning effort level) while forwarding it for those that do. The fix covers both the chat-completion and Messages pass-through transformation paths since they share the helper. * chore(vertex): log at debug when dropping unsupported output_config.effort Operators pointing an unregistered Vertex Claude alias that does support effort would otherwise see it stripped with no signal. Debug level keeps it out of normal logs since Claude Code sends effort on every request. (cherry picked from commit cc55662e5fdc6af4f118a1f3ff885068b75450d9) --- .../transformation.py | 2 +- .../anthropic/output_params_utils.py | 51 ++++++++++++++----- .../anthropic/transformation.py | 2 +- ...artner_models_anthropic_messages_config.py | 34 +++++++++++++ ...partner_models_anthropic_transformation.py | 46 ++++++++++++++--- 5 files changed, 112 insertions(+), 23 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py index 4be4c2d5e78..1e92754857b 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py @@ -159,6 +159,6 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert "model", None ) # do not pass model in request body to vertex ai - sanitize_vertex_anthropic_output_params(anthropic_messages_request) + sanitize_vertex_anthropic_output_params(anthropic_messages_request, model) return anthropic_messages_request diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/output_params_utils.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/output_params_utils.py index a33ad677789..280cc1c888a 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/output_params_utils.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/output_params_utils.py @@ -10,23 +10,38 @@ import; extracting the helper into a leaf module resolves the warning and keeps the parent module's import surface narrow. """ -# Keys inside ``output_config`` that Vertex AI Claude does not accept. -# Add an entry only when a 400 "Extra inputs are not permitted" is -# reproducible against the live Vertex endpoint. +# Keys inside ``output_config`` that Vertex AI Claude rejects regardless of +# the target model. Add an entry only when a 400 "Extra inputs are not +# permitted" is reproducible against the live Vertex endpoint for every model. VERTEX_UNSUPPORTED_OUTPUT_CONFIG_KEYS: frozenset = frozenset() -def sanitize_vertex_anthropic_output_params(data: dict) -> None: +def _model_accepts_output_config_effort(model: str) -> bool: + """Whether ``model`` accepts ``output_config.effort`` on Vertex. + + Opus/Sonnet 4.6+ advertise ``supports_output_config`` (or a reasoning + effort level) and accept it; Haiku 4.5 advertises neither and 400s on + ``output_config.effort: Extra inputs are not permitted``. Imported lazily + so this stays a leaf module (see module docstring). + """ + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + return AnthropicConfig._model_supports_effort_param(model) + + +def sanitize_vertex_anthropic_output_params(data: dict, model: str) -> None: """ Strip Vertex-unsupported keys from ``output_config`` / ``output_format`` in-place; forward whatever remains. Behavior: - * ``output_config`` containing only unsupported keys (e.g. ``effort`` - alone) is removed entirely so the request body has no empty dict. - * ``output_config`` containing a mix of supported + unsupported keys - has the unsupported subset filtered out and the rest forwarded. - * ``output_config`` that is supported in full passes through unchanged. + * ``output_config.effort`` is dropped for models that don't accept it + (e.g. Haiku 4.5) and forwarded for those that do (Opus/Sonnet 4.6+). + Clients like Claude Code inject it into every Messages payload, so the + gate has to live here rather than rely on the caller. + * Keys in ``VERTEX_UNSUPPORTED_OUTPUT_CONFIG_KEYS`` are always filtered. + * ``output_config`` left empty after filtering is removed so the request + body has no empty dict. * ``output_format`` is forwarded as-is (Vertex AI Claude accepts it). * Non-dict values for ``output_config`` are dropped to avoid sending malformed payloads downstream. @@ -37,11 +52,19 @@ def sanitize_vertex_anthropic_output_params(data: dict) -> None: if not isinstance(output_config, dict): data.pop("output_config", None) return - sanitized = { - k: v - for k, v in output_config.items() - if k not in VERTEX_UNSUPPORTED_OUTPUT_CONFIG_KEYS - } + + drop_keys = set(VERTEX_UNSUPPORTED_OUTPUT_CONFIG_KEYS) + if "effort" in output_config and not _model_accepts_output_config_effort(model): + from litellm._logging import verbose_logger + + verbose_logger.debug( + "Dropping unsupported output_config.effort for vertex_ai model=%s " + "(no supports_output_config in the model map)", + model, + ) + drop_keys.add("effort") + + sanitized = {k: v for k, v in output_config.items() if k not in drop_keys} if sanitized: data["output_config"] = sanitized else: diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index 4627d9f6df3..c852909d475 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -106,7 +106,7 @@ class VertexAIAnthropicConfig(AnthropicConfig): data.pop("model", None) # vertex anthropic doesn't accept 'model' parameter - sanitize_vertex_anthropic_output_params(data) + sanitize_vertex_anthropic_output_params(data, model) tools = optional_params.get("tools") tool_search_used = self.is_tool_search_used(tools) diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index b8cd65d3c99..6f4bb4e59c2 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -313,6 +313,40 @@ def test_transform_anthropic_messages_request_removes_scope_from_cache_control() assert result["messages"][0]["content"][0]["cache_control"]["type"] == "ephemeral" +def test_messages_request_strips_effort_for_haiku_45(): + """Regression: Claude Code (``claude --model claude-haiku-4.5``) sends + ``output_config.effort`` in its default Messages payload. Haiku 4.5 on + Vertex rejects it with 400 ``output_config.effort: Extra inputs are not + permitted``, so the pass-through must strip it for Haiku while keeping it + for Opus/Sonnet 4.6+.""" + config = VertexAIPartnerModelsAnthropicMessagesConfig() + messages = [{"role": "user", "content": "Hello"}] + + haiku_result = config.transform_anthropic_messages_request( + model="claude-haiku-4-5@20251001", + messages=messages, + anthropic_messages_optional_request_params={ + "max_tokens": 1024, + "output_config": {"effort": "high"}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert "output_config" not in haiku_result + + opus_result = config.transform_anthropic_messages_request( + model="claude-opus-4-6", + messages=messages, + anthropic_messages_optional_request_params={ + "max_tokens": 1024, + "output_config": {"effort": "high"}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + assert opus_result["output_config"] == {"effort": "high"} + + def test_provider_config_manager_reuses_vertex_anthropic_messages_config_instance(): """ Regression test: repeated provider config lookups for the same Vertex Claude model diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py index d89d09a4e63..ac2368130d8 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py @@ -675,28 +675,60 @@ def test_sanitize_vertex_anthropic_output_params_unit(): sanitize_vertex_anthropic_output_params, ) + supported = "claude-opus-4-6" + # No-op when output_config absent. data: dict = {"max_tokens": 8} - sanitize_vertex_anthropic_output_params(data) + sanitize_vertex_anthropic_output_params(data, supported) assert data == {"max_tokens": 8} - # Effort-only → preserved (Vertex 4.6/4.7 accept it on rawPredict). + # Effort-only on a supporting model → preserved (Vertex 4.6/4.7 accept it). data = {"output_config": {"effort": "high"}} - sanitize_vertex_anthropic_output_params(data) + sanitize_vertex_anthropic_output_params(data, supported) assert data["output_config"] == {"effort": "high"} # Format-only → preserved unchanged. fmt = {"format": {"type": "json_schema", "schema": {"type": "object"}}} data = {"output_config": dict(fmt)} - sanitize_vertex_anthropic_output_params(data) + sanitize_vertex_anthropic_output_params(data, supported) assert data["output_config"] == fmt - # Mixed → both effort and format kept (no current Vertex-unsupported keys). + # Mixed on a supporting model → both effort and format kept. data = {"output_config": {"format": fmt["format"], "effort": "high"}} - sanitize_vertex_anthropic_output_params(data) + sanitize_vertex_anthropic_output_params(data, supported) assert data["output_config"] == {"format": fmt["format"], "effort": "high"} # Non-dict → dropped defensively. data = {"output_config": "garbage"} - sanitize_vertex_anthropic_output_params(data) + sanitize_vertex_anthropic_output_params(data, supported) assert "output_config" not in data + + +def test_sanitize_strips_effort_for_haiku_45(): + """Regression: Haiku 4.5 on Vertex does not support ``output_config.effort`` + and 400s with ``Extra inputs are not permitted``. Claude Code injects + ``effort`` into every Messages payload, so the helper must strip it for + models that don't advertise output_config support while leaving it intact + for Opus/Sonnet 4.6+.""" + from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.output_params_utils import ( + sanitize_vertex_anthropic_output_params, + ) + + haiku = "claude-haiku-4-5@20251001" + + # Effort-only → output_config removed entirely (no empty dict on the wire). + data: dict = {"output_config": {"effort": "high"}, "max_tokens": 8} + sanitize_vertex_anthropic_output_params(data, haiku) + assert "output_config" not in data + assert data["max_tokens"] == 8 + + # Mixed → effort stripped, format preserved. + fmt = {"type": "json_schema", "schema": {"type": "object"}} + data = {"output_config": {"effort": "high", "format": fmt}} + sanitize_vertex_anthropic_output_params(data, haiku) + assert data["output_config"] == {"format": fmt} + + # Same payload on a supporting model keeps effort untouched. + data = {"output_config": {"effort": "high"}} + sanitize_vertex_anthropic_output_params(data, "vertex_ai/claude-opus-4-6") + assert data["output_config"] == {"effort": "high"} From cd0e7914127049297ad63f6adbe7a928bcb272ec Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 3 Jun 2026 12:13:02 -0700 Subject: [PATCH 6/8] fix: passthrough endpoints duplicate logs (#29598) * fix duplicate cost callbacks for anthropic streaming pass-through Two bugs caused _PROXY_track_cost_callback to see stream=True + complete_streaming_response=None on every streaming pass-through request, making the dedup guard in dispatch_success_handlers permanently inactive: 1. pass_through_endpoints.py created the Logging object with stream=False for all requests. _is_assembled_stream_success short-circuits on self.stream is not True, so has_dispatched_final_stream_success was never set and any second dispatch went through unchecked. Fix: set logging_obj.stream = True after stream detection. 2. _create_anthropic_response_logging_payload set complete_streaming_response inside the try block after litellm.completion_cost(), so a pricing error caused an early return without setting it on model_call_details. Fix: set complete_streaming_response before the try block. Co-Authored-By: Claude Sonnet 4.6 * fix stream * add stream to logging obj * test(pass_through): give mock logging object a real model_call_details dict The anthropic passthrough logging payload now records the assembled response on model_call_details before cost calculation, which requires model_call_details to support item assignment. In production it is always a dict; the existing unit test stubbed the logging object with a bare Mock whose attribute is not subscriptable, so the new assignment raised TypeError. Use a real dict to match the production logging object. * test(pass_through): cover streaming logging-obj stream flag The streaming branch of pass_through_request that marks the logging object as streaming (logging_obj.stream and model_call_details["stream"]) had no unit coverage, so the patch coverage gate flagged it. Add a regression test that drives a streaming pass-through request through pass_through_request and asserts the logging object is flagged as a stream before dispatch. * test(pass_through): cover SSE-response stream flag fallback branch The auto-detected streaming branch of pass_through_request (when a request that was not flagged as streaming returns a text/event-stream response) sets logging_obj.stream and model_call_details["stream"] but had no unit coverage, so the codecov patch gate failed at 60%. Drive a non-streaming pass-through request whose upstream response is SSE through pass_through_request and assert the logging object is flagged as a stream before dispatch. * fix(pass_through): gate complete_streaming_response on stream flag perform_redaction only scrubs complete_streaming_response when model_call_details["stream"] is True. Setting it unconditionally for non-streaming Anthropic pass-through responses left the assembled response unredacted in model_call_details, which is handed to logging callbacks as kwargs when message logging is disabled. Only record it for actual streaming responses so redaction always applies. --------- Co-authored-by: mubashir1osmani Co-authored-by: Claude Sonnet 4.6 (cherry picked from commit 2bbdbfa5c348e198eb21461731c784e12897f01f) --- .../anthropic_passthrough_logging_handler.py | 7 + .../pass_through_endpoints.py | 6 + .../test_unit_test_anthropic_pass_through.py | 1 + ...t_anthropic_passthrough_logging_handler.py | 332 ++++++++++++++++++ .../test_pass_through_endpoints.py | 125 +++++++ 5 files changed, 471 insertions(+) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index c42faa59cf0..82260eb2ae5 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -114,6 +114,13 @@ class AnthropicPassthroughLoggingHandler: handles streaming and non-streaming responses """ + # Only record complete_streaming_response for actual streaming responses. + # perform_redaction scrubs this field only when stream is True, so setting + # it on a non-streaming response would bypass message redaction. + if logging_obj.model_call_details.get("stream") is True: + logging_obj.model_call_details["complete_streaming_response"] = ( + litellm_model_response + ) try: # Get custom_llm_provider from logging object if available (e.g., azure_ai for Azure Anthropic) custom_llm_provider = logging_obj.model_call_details.get( diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 26f89e6b315..9c196b9cbd3 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -871,6 +871,9 @@ async def pass_through_request( # noqa: PLR0915 ) if stream: + logging_obj.stream = True + logging_obj.model_call_details["stream"] = True + if is_multipart: response = ( await HttpPassThroughEndpointHelpers.make_multipart_http_request( @@ -931,6 +934,9 @@ async def pass_through_request( # noqa: PLR0915 verbose_proxy_logger.debug("response.headers= %s", response.headers) if _is_streaming_response(response) is True: + logging_obj.stream = True + logging_obj.model_call_details["stream"] = True + try: response.raise_for_status() except httpx.HTTPStatusError as e: diff --git a/tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py b/tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py index 97a1f2eecc7..1513ee26473 100644 --- a/tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py +++ b/tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py @@ -318,6 +318,7 @@ def test_handle_logging_anthropic_collected_chunks(all_chunks): from litellm.types.utils import ModelResponse litellm_logging_obj = Mock() + litellm_logging_obj.model_call_details = {} pass_through_logging_obj = Mock() sent_args = { diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index c16c42decc0..7593ea2f6db 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -684,3 +684,335 @@ class TestAnthropicBatchPassthroughCostTracking: mock_proxy_logging_obj.get_proxy_hook.assert_called_once_with( "managed_files" ) + + +class TestStreamFalseDeduplication: + """ + Regression tests for the duplicate-callback bug where a streaming pass-through + request had stream=False hardcoded on its Logging object. + + Before the fix: + - logging_obj.stream was always False for pass-through requests + - _is_assembled_stream_success() checked `self.stream is not True` and returned + False immediately, so has_dispatched_final_stream_success was never set + - Any second dispatch_success_handlers call went through unchecked + + After the fix: + - pass_through_endpoints.py sets logging_obj.stream = True after detecting stream + - _create_anthropic_response_logging_payload sets complete_streaming_response on + model_call_details so callbacks see the correct assembled response state + - _is_assembled_stream_success returns True, dedup guard fires on first dispatch + """ + + @staticmethod + def _sse(event, data): + return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() + + @staticmethod + def _make_logging_obj(stream: bool = False) -> LiteLLMLoggingObj: + logging_obj = LiteLLMLoggingObj( + model="claude-3-5-sonnet-20241022", + messages=[{"role": "user", "content": "hello"}], + stream=stream, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="test-call-id", + function_id="1245", + ) + return logging_obj + + @staticmethod + def _build_chunks(): + frames = [ + TestStreamFalseDeduplication._sse( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_abc", + "type": "message", + "role": "assistant", + "model": "claude-3-5-sonnet-20241022", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 0}, + }, + }, + ), + TestStreamFalseDeduplication._sse( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + ), + TestStreamFalseDeduplication._sse( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hello"}, + }, + ), + TestStreamFalseDeduplication._sse( + "content_block_stop", {"type": "content_block_stop", "index": 0} + ), + TestStreamFalseDeduplication._sse( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 5}, + }, + ), + TestStreamFalseDeduplication._sse("message_stop", {"type": "message_stop"}), + ] + from litellm.proxy.pass_through_endpoints.streaming_handler import ( + PassThroughStreamingHandler, + ) + + return PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(frames) + + def test_complete_streaming_response_set_on_model_call_details(self): + """ + After the fix, _create_anthropic_response_logging_payload must set + complete_streaming_response on logging_obj.model_call_details so that + callbacks like _PROXY_track_cost_callback see the assembled response + instead of None. + + Before the fix: model_call_details had no complete_streaming_response key. + The log showed: "kwargs stream: True + complete streaming response: None" + """ + from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + EndpointType, + ) + + # pass_through_request sets the stream flag before the streaming handler + # reconstructs the response; mirror that here. + logging_obj = self._make_logging_obj(stream=True) + logging_obj.model_call_details["stream"] = True + all_chunks = list(self._build_chunks()) + + result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/anthropic/v1/messages", + request_body={"model": "claude-3-5-sonnet-20241022", "stream": True}, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + all_chunks=all_chunks, + end_time=datetime.now(), + ) + + # The assembled response must be stored on model_call_details so callbacks + # can identify this as a completed streaming call, not an in-progress one. + assert ( + logging_obj.model_call_details.get("complete_streaming_response") + is not None + ), "complete_streaming_response must be set on model_call_details after assembly" + + # The returned result must match what was stored + assert result["result"] is logging_obj.model_call_details.get( + "complete_streaming_response" + ) + + def test_dedup_guard_fires_when_stream_true_on_logging_obj(self): + """ + When logging_obj.stream is True (set by pass_through_endpoints.py after + detecting a streaming request), dispatch_success_handlers must set + has_dispatched_final_stream_success=True on the first call so that any + second call is a no-op. + + This is the _is_assembled_stream_success gate: with stream=False it + always returned False and the guard was permanently disabled. + """ + from litellm.types.passthrough_endpoints.pass_through_endpoints import ( + EndpointType, + ) + from litellm.types.utils import ModelResponse + + # Simulate what pass_through_endpoints.py now does after stream detection + logging_obj = self._make_logging_obj(stream=False) + logging_obj.stream = True # fix applied + logging_obj.model_call_details["stream"] = True + + # Simulate what _create_anthropic_response_logging_payload now does + mock_response = ModelResponse(model="claude-3-5-sonnet-20241022") + logging_obj.model_call_details["complete_streaming_response"] = mock_response + + assert logging_obj._is_assembled_stream_success(result=mock_response) is True + + # First dispatch sets the flag + assert not logging_obj.model_call_details.get( + "has_dispatched_final_stream_success" + ) + logging_obj.model_call_details["has_dispatched_final_stream_success"] = True + + # Second dispatch would be blocked — simulate the guard check + would_skip = bool( + logging_obj._is_assembled_stream_success(result=mock_response) + and logging_obj.model_call_details.get( + "has_dispatched_final_stream_success" + ) + ) + assert would_skip is True, ( + "Dedup guard must block a second dispatch_success_handlers call for the " + "same assembled streaming response" + ) + + def test_sse_fallback_path_sets_stream_true_for_dedup(self): + """ + When a nominally non-streaming request receives an SSE response + (_is_streaming_response returns True), the fallback branch in + pass_through_endpoints.py must set logging_obj.stream = True so the + dedup guard activates. + + Before the fix the fallback path never set stream=True, so + _is_assembled_stream_success always returned False and duplicate + callback dispatches were never blocked. + """ + from litellm.types.utils import ModelResponse + + # logging_obj starts with stream=False, as created before the request + logging_obj = self._make_logging_obj(stream=False) + assert logging_obj._is_assembled_stream_success(result=MagicMock()) is False + + # Simulate what the SSE fallback branch in pass_through_endpoints.py now does + logging_obj.stream = True + logging_obj.model_call_details["stream"] = True + + mock_response = ModelResponse(model="claude-3-5-sonnet-20241022") + logging_obj.model_call_details["complete_streaming_response"] = mock_response + + # With stream=True the dedup guard must be active + assert logging_obj._is_assembled_stream_success(result=mock_response) is True + + logging_obj.model_call_details["has_dispatched_final_stream_success"] = True + + would_skip = bool( + logging_obj._is_assembled_stream_success(result=mock_response) + and logging_obj.model_call_details.get( + "has_dispatched_final_stream_success" + ) + ) + assert would_skip is True + + def test_stream_false_logging_obj_bypasses_dedup_guard(self): + """ + Demonstrates the pre-fix state: with stream=False on the logging object, + _is_assembled_stream_success always returns False regardless of whether + complete_streaming_response is set. This means the dedup guard can never + fire, so duplicate dispatches go through unchecked. + + This test documents the old broken behavior so the fix is clearly justified. + """ + from litellm.types.utils import ModelResponse + + logging_obj = self._make_logging_obj(stream=False) + mock_response = ModelResponse(model="claude-3-5-sonnet-20241022") + logging_obj.model_call_details["complete_streaming_response"] = mock_response + + # With stream=False, _is_assembled_stream_success returns False even though + # complete_streaming_response is present — the guard is permanently disabled. + assert logging_obj._is_assembled_stream_success(result=mock_response) is False + + +class TestNonStreamingResponseRedaction: + """ + Regression tests ensuring _create_anthropic_response_logging_payload only sets + complete_streaming_response for streaming responses. perform_redaction scrubs + that field exclusively when model_call_details["stream"] is True, so storing it + on a non-streaming response would deliver the unredacted response to logging + callbacks when message logging is disabled. + """ + + @staticmethod + def _make_logging_obj(stream: bool) -> LiteLLMLoggingObj: + logging_obj = LiteLLMLoggingObj( + model="claude-3-5-sonnet-20241022", + messages=[{"role": "user", "content": "hello"}], + stream=stream, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="test-call-id", + function_id="1245", + ) + # pass_through_request mirrors the stream flag onto model_call_details, + # which is the key perform_redaction inspects. + logging_obj.model_call_details["stream"] = stream + return logging_obj + + def test_non_streaming_does_not_set_complete_streaming_response(self): + from litellm.types.utils import ModelResponse + + logging_obj = self._make_logging_obj(stream=False) + response = ModelResponse(model="claude-3-5-sonnet-20241022") + + AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=response, + model="claude-3-5-sonnet-20241022", + kwargs={}, + start_time=datetime.now(), + end_time=datetime.now(), + logging_obj=logging_obj, + ) + + assert ( + "complete_streaming_response" not in logging_obj.model_call_details + ), "non-streaming responses must not populate complete_streaming_response" + + def test_streaming_sets_complete_streaming_response(self): + from litellm.types.utils import ModelResponse + + logging_obj = self._make_logging_obj(stream=True) + response = ModelResponse(model="claude-3-5-sonnet-20241022") + + AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=response, + model="claude-3-5-sonnet-20241022", + kwargs={}, + start_time=datetime.now(), + end_time=datetime.now(), + logging_obj=logging_obj, + ) + + assert ( + logging_obj.model_call_details.get("complete_streaming_response") + is response + ) + + def test_non_streaming_response_is_redacted_when_message_logging_off(self): + from litellm.litellm_core_utils.redact_messages import ( + redact_message_input_output_from_logging, + ) + from litellm.types.utils import Choices, Message, ModelResponse + + logging_obj = self._make_logging_obj(stream=False) + response = ModelResponse( + model="claude-3-5-sonnet-20241022", + choices=[Choices(message=Message(role="assistant", content="secret"))], + ) + + AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=response, + model="claude-3-5-sonnet-20241022", + kwargs={}, + start_time=datetime.now(), + end_time=datetime.now(), + logging_obj=logging_obj, + ) + + logging_obj.model_call_details["litellm_params"] = { + "metadata": {"headers": {"x-litellm-enable-message-redaction": True}} + } + + redacted = redact_message_input_output_from_logging( + model_call_details=logging_obj.model_call_details, + result=response, + ) + + leaked = logging_obj.model_call_details.get("complete_streaming_response") + assert leaked is None + assert redacted.choices[0].message.content == "redacted-by-litellm" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 97a21136198..89b58973079 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -989,6 +989,131 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs(): assert metadata["user_api_key_user_id"] == "test-user-id" +@pytest.mark.asyncio +async def test_pass_through_request_streaming_marks_logging_obj_as_stream(): + """ + Regression: a streaming pass-through request must flag its logging object as + streaming (logging_obj.stream and model_call_details["stream"]) before the + response is dispatched, so cost/success callbacks treat it as a stream and the + streaming dedup guard fires instead of double-logging. + """ + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.PassThroughStreamingHandler.chunk_processor" + ) as mock_chunk_processor: + mock_proxy_logging.pre_call_hook = AsyncMock( + return_value={"model": "claude-3", "stream": True} + ) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + + upstream_response = MagicMock() + upstream_response.status_code = 200 + upstream_response.headers = {} + upstream_response.raise_for_status = MagicMock() + + async_client = MagicMock() + async_client.build_request = MagicMock(return_value=MagicMock()) + async_client.send = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + + async def _empty_chunks(*args, **kwargs): + return + yield # pragma: no cover + + mock_chunk_processor.return_value = _empty_chunks() + + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = "http://test-proxy.com/v1/messages" + mock_request.body = AsyncMock( + return_value=b'{"model": "claude-3", "stream": true}' + ) + mock_request.headers = Headers({}) + mock_request.query_params = QueryParams({}) + + await pass_through_request( + request=mock_request, + target="http://target-api.com/v1/messages", + custom_headers={}, + user_api_key_dict=MagicMock(), + stream=True, + ) + + async_client.send.assert_awaited_once() + assert async_client.send.call_args.kwargs["stream"] is True + + mock_chunk_processor.assert_called_once() + logging_obj = mock_chunk_processor.call_args.kwargs[ + "litellm_logging_obj" + ] + assert logging_obj.stream is True + assert logging_obj.model_call_details["stream"] is True + + +@pytest.mark.asyncio +async def test_pass_through_request_sse_response_marks_logging_obj_as_stream(): + """ + Regression: a request that is not flagged as streaming up front but whose + upstream response comes back as an SSE stream (content-type text/event-stream) + must still flag its logging object as streaming before dispatch. Otherwise the + cost/success callbacks treat the assembled stream as a non-stream and the dedup + guard never fires, double-logging the request. + """ + with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client" + ) as mock_get_client: + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.PassThroughStreamingHandler.chunk_processor" + ) as mock_chunk_processor: + mock_proxy_logging.pre_call_hook = AsyncMock( + return_value={"model": "claude-3"} + ) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + + upstream_response = MagicMock() + upstream_response.status_code = 200 + upstream_response.headers = {"content-type": "text/event-stream"} + upstream_response.raise_for_status = MagicMock() + + async_client = MagicMock() + async_client.request = AsyncMock(return_value=upstream_response) + mock_get_client.return_value = MagicMock(client=async_client) + + async def _empty_chunks(*args, **kwargs): + return + yield # pragma: no cover + + mock_chunk_processor.return_value = _empty_chunks() + + mock_request = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.url = "http://test-proxy.com/v1/messages" + mock_request.body = AsyncMock(return_value=b'{"model": "claude-3"}') + mock_request.headers = Headers({}) + mock_request.query_params = QueryParams({}) + + await pass_through_request( + request=mock_request, + target="http://target-api.com/v1/messages", + custom_headers={}, + user_api_key_dict=MagicMock(), + stream=False, + ) + + async_client.request.assert_awaited_once() + + mock_chunk_processor.assert_called_once() + logging_obj = mock_chunk_processor.call_args.kwargs[ + "litellm_logging_obj" + ] + assert logging_obj.stream is True + assert logging_obj.model_call_details["stream"] is True + + @pytest.mark.asyncio async def test_create_pass_through_endpoint(): """ From c38c2bf7fe4ea12cd8ce27da997c13460c2eb27b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 3 Jun 2026 22:53:32 +0000 Subject: [PATCH 7/8] =?UTF-8?q?bump:=20version=201.85.3=20=E2=86=92=201.85?= =?UTF-8?q?.4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 4a11854cc68..0563d0e6093 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.85.3" +version = "1.85.4" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.14" @@ -250,7 +250,7 @@ source-exclude = [ profile = "black" [tool.commitizen] -version = "1.85.3" +version = "1.85.4" version_files = [ "pyproject.toml:^version", ] From aaa30164ef25433bd5e2ca6c06a23bef1f2b1243 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 3 Jun 2026 22:53:32 +0000 Subject: [PATCH 8/8] chore: update uv.lock for 1.85.4 --- uv.lock | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/uv.lock b/uv.lock index 0f347fcd1e4..97d007a3843 100644 --- a/uv.lock +++ b/uv.lock @@ -3189,7 +3189,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.85.3" +version = "1.85.4" source = { editable = "." } dependencies = [ { name = "aiohttp" },