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/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/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 14f198e0f12..b1a109169c2 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -532,6 +532,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 1e87dcaef1c..747d18c3457 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1238,7 +1238,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 [] @@ -1295,16 +1297,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( @@ -1314,11 +1318,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) @@ -1368,7 +1387,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, ), ) @@ -1390,12 +1411,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 03278633928..4bdedda956d 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -144,12 +144,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, ) @@ -1023,6 +1025,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: @@ -1440,6 +1443,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: @@ -1568,6 +1572,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 @@ -2148,6 +2153,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) @@ -2450,6 +2456,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: @@ -2591,6 +2598,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/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/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a67d8d934bf..b5ab14fcd32 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -883,7 +883,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/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/pyproject.toml b/pyproject.toml index 8dedca241ad..3e0b99a2205 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.87.0" +version = "1.87.1" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.14" @@ -253,7 +253,7 @@ source-exclude = [ profile = "black" [tool.commitizen] -version = "1.87.0" +version = "1.87.1" version_files = [ "pyproject.toml:^version", ] 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. 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"} diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 68e1636d380..f324eeeac4c 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -382,6 +382,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( 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; 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 38bf2d2c915..c3d61943288 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 @@ -3032,6 +3032,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 # ============================================ diff --git a/uv.lock b/uv.lock index fe3e0e037cf..9c110bfc5a9 100644 --- a/uv.lock +++ b/uv.lock @@ -3269,7 +3269,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.87.0" +version = "1.87.1" source = { editable = "." } dependencies = [ { name = "aiohttp" },