diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 39f0af1948..03167c5a2d 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -68,6 +68,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( populate_request_with_path_params, ) from litellm.proxy.common_utils.realtime_utils import _realtime_request_body +from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.utils import ( PrismaClient, ProxyLogging, @@ -1977,6 +1978,16 @@ async def _run_centralized_common_checks( llm_router=llm_router, ) + # Merge x-litellm-tags into request_data BEFORE common_checks runs. + # _tag_max_budget_check inside common_checks only inspects request_data; + # without this pre-merge, header-supplied tags bypass tag-budget + # enforcement. + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( + request=request, + request_data=request_data, + user_api_key_dict=user_api_key_auth_obj, + ) + _ = await common_checks( request=request, request_body=request_data, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 00735e6be5..7cd099a729 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -25,6 +25,9 @@ from litellm.proxy._types import ( TeamCallbackMetadata, UserAPIKeyAuth, ) +from litellm.proxy.common_utils.callback_utils import ( + get_metadata_variable_name_from_kwargs, +) from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers # Cache special headers as a frozenset for O(1) lookup performance @@ -1187,6 +1190,71 @@ class LiteLLMProxyRequestSetup: return tags + @staticmethod + def apply_client_tag_policy_pre_auth( + request: Request, + request_data: dict, + user_api_key_dict: UserAPIKeyAuth, + ) -> None: + """ + Merge ``x-litellm-tags`` header tags into ``request_data`` BEFORE + auth budget gates run, so ``_tag_max_budget_check`` (which only + inspects ``request_data``) sees them. Without this, header-tagged + requests silently bypass per-tag budget enforcement. + + Why: ``add_litellm_data_to_request`` runs the equivalent merge + post-auth, after ``_tag_max_budget_check`` has already executed. + Header-supplied tags merged there are invisible to that check. + Running the merge here closes that gap; the post-auth merge in + ``add_litellm_data_to_request`` remains as defense-in-depth. + + How to apply: invoked from the auth chain just before + ``common_checks``. Mutates ``request_data`` in place; idempotent + when followed by ``add_litellm_data_to_request``. + """ + # No allow_client_tags opt-in: caller-supplied tags always flow + # into metadata.tags (see add_litellm_data_to_request). The pre-auth + # merge mirrors that so _tag_max_budget_check sees the same tags. + headers = _safe_get_request_headers(request=request) + raw_header_tags = headers.get("x-litellm-tags") + if not raw_header_tags: + return + + if isinstance(raw_header_tags, str): + header_tags: List[str] = [ + t.strip() for t in raw_header_tags.split(",") if t.strip() + ] + elif isinstance(raw_header_tags, list): + header_tags = [t for t in raw_header_tags if isinstance(t, str) and t] + else: + return + + if not header_tags: + return + + # Match the metadata key that get_tags_from_request_body will read + # from (litellm_metadata vs metadata) so the merged tags are visible + # to _tag_max_budget_check. + _metadata_variable_name = get_metadata_variable_name_from_kwargs(request_data) + metadata = request_data.get(_metadata_variable_name) + # metadata can arrive as a JSON string (multipart/form-data, extra_body). + # Parse it so existing tags survive the merge — overwriting the string + # with {} would let a caller bypass _tag_max_budget_check on an + # over-budget body tag by also sending a within-budget header tag. + if isinstance(metadata, str): + parsed = safe_json_loads(metadata) + metadata = parsed if isinstance(parsed, dict) else {} + request_data[_metadata_variable_name] = metadata + elif not isinstance(metadata, dict): + metadata = {} + request_data[_metadata_variable_name] = metadata + + existing_tags = metadata.get("tags") + metadata["tags"] = LiteLLMProxyRequestSetup._merge_tags( + request_tags=existing_tags if isinstance(existing_tags, list) else None, + tags_to_add=header_tags, + ) + async def add_litellm_data_to_request( # noqa: PLR0915 data: dict, diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 5412bf3a56..19204e8b8a 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -3850,6 +3850,287 @@ def test_get_guardrail_from_metadata_reads_litellm_metadata_when_no_metadata(): ], f"Expected guardrails from litellm_metadata fallback, got: {result}" +def _build_request_mock_with_headers(headers: dict) -> Request: + request_mock = MagicMock(spec=Request) + request_mock.url = MagicMock() + request_mock.url.path = "/v1/chat/completions" + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = headers + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + request_mock.state = MagicMock() + request_mock.state._cached_headers = None + return request_mock + + +class TestApplyClientTagPolicyPreAuth: + """Tests for ``LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth``. + + Regression coverage for the bug where ``x-litellm-tags`` header was + invisible to ``_tag_max_budget_check`` because the merge happened + post-auth in ``add_litellm_data_to_request``. + """ + + def test_merges_header_tags_into_metadata(self): + request_mock = _build_request_mock_with_headers( + {"x-litellm-tags": "tenant:acme,env:prod"} + ) + data = {"model": "gpt-3.5-turbo"} + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={}, + team_metadata={}, + ) + + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( + request=request_mock, + request_data=data, + user_api_key_dict=user_api_key_dict, + ) + + assert data["metadata"]["tags"] == ["tenant:acme", "env:prod"] + + def test_unions_header_tags_with_existing_metadata_tags(self): + request_mock = _build_request_mock_with_headers( + {"x-litellm-tags": "tenant:acme,env:prod"} + ) + data = { + "model": "gpt-3.5-turbo", + "metadata": {"tags": ["env:prod", "team:platform"]}, + } + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={}, + team_metadata={}, + ) + + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( + request=request_mock, + request_data=data, + user_api_key_dict=user_api_key_dict, + ) + + # Existing tags first, dedupe header tags + assert data["metadata"]["tags"] == ["env:prod", "team:platform", "tenant:acme"] + + def test_preserves_body_tags(self): + # Pre-auth must NOT touch body-supplied tags. _tag_max_budget_check + # (inside common_checks) enforces per-tag budgets on whatever tags + # it sees in request_data, including body tags. The helper only + # adds header tags to metadata.tags. + request_mock = _build_request_mock_with_headers( + {"x-litellm-tags": "tenant:acme"} + ) + data = { + "model": "gpt-3.5-turbo", + "tags": ["root-tag"], + "litellm_metadata": {"tags": ["litellm-meta-tag"]}, + } + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={}, + team_metadata={}, + ) + + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( + request=request_mock, + request_data=data, + user_api_key_dict=user_api_key_dict, + ) + + assert data["tags"] == ["root-tag"] + # litellm_metadata is the active metadata key (it's present), so + # header tags merge into it and union with existing tags there. + assert data["litellm_metadata"]["tags"] == [ + "litellm-meta-tag", + "tenant:acme", + ] + + def test_uses_litellm_metadata_when_present(self): + request_mock = _build_request_mock_with_headers( + {"x-litellm-tags": "tenant:acme"} + ) + data = { + "model": "gpt-3.5-turbo", + "litellm_metadata": {"foo": "bar"}, + } + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={}, + team_metadata={}, + ) + + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( + request=request_mock, + request_data=data, + user_api_key_dict=user_api_key_dict, + ) + + # get_metadata_variable_name_from_kwargs returns "litellm_metadata" + # when present, so header tags should land there to be visible to + # _tag_max_budget_check. + assert data["litellm_metadata"]["tags"] == ["tenant:acme"] + assert "tags" not in data.get("metadata", {}) + + def test_no_header_no_mutation(self): + request_mock = _build_request_mock_with_headers({}) + data = {"model": "gpt-3.5-turbo"} + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={}, + team_metadata={}, + ) + + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( + request=request_mock, + request_data=data, + user_api_key_dict=user_api_key_dict, + ) + + assert "metadata" not in data or "tags" not in data["metadata"] + + def test_string_metadata_tags_survive_header_merge(self): + # metadata can arrive as a JSON string (multipart/form-data, extra_body). + # The pre-auth merge must parse it so an over-budget body tag isn't + # silently dropped when a within-budget header tag is also present. + request_mock = _build_request_mock_with_headers({"x-litellm-tags": "free"}) + data = { + "model": "gpt-3.5-turbo", + "metadata": '{"tags": ["paid"]}', + } + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={}, + team_metadata={}, + ) + + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( + request=request_mock, + request_data=data, + user_api_key_dict=user_api_key_dict, + ) + + assert isinstance(data["metadata"], dict) + assert data["metadata"]["tags"] == ["paid", "free"] + + @pytest.mark.asyncio + async def test_string_metadata_does_not_bypass_tag_max_budget_check(self): + """Regression: string metadata containing an over-budget tag must not + be silently overwritten when an x-litellm-tags header is present.""" + from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_TagTable + from litellm.proxy.auth.auth_checks import _tag_max_budget_check + from litellm.proxy.utils import ProxyLogging + + request_mock = _build_request_mock_with_headers({"x-litellm-tags": "free"}) + data = { + "model": "gpt-3.5-turbo", + "metadata": '{"tags": ["paid"]}', + } + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={}, + team_metadata={}, + ) + + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( + request=request_mock, + request_data=data, + user_api_key_dict=user_api_key_dict, + ) + + paid_tag = LiteLLM_TagTable( + tag_name="paid", + spend=0.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), + ) + + async def mock_get_current_spend(counter_key, fallback_spend): + if counter_key == "spend:tag:paid": + return 0.50 + return fallback_spend + + with ( + patch( + "litellm.proxy.proxy_server.get_current_spend", + mock_get_current_spend, + ), + patch( + "litellm.proxy.auth.auth_checks.get_tag_objects_batch", + new_callable=AsyncMock, + return_value={"paid": paid_tag}, + ), + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _tag_max_budget_check( + request_body=data, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + valid_token=UserAPIKeyAuth(token="test-token"), + ) + assert exc_info.value.current_cost == 0.50 + assert exc_info.value.max_budget == 0.10 + + @pytest.mark.asyncio + async def test_header_tags_visible_to_tag_max_budget_check(self): + """End-to-end: helper + ``_tag_max_budget_check`` enforces budget on + header-supplied tags. Without the helper, this would silently pass.""" + from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_TagTable + from litellm.proxy.auth.auth_checks import _tag_max_budget_check + from litellm.proxy.utils import ProxyLogging + + request_mock = _build_request_mock_with_headers( + {"x-litellm-tags": "tenant:acme"} + ) + data = {"model": "gpt-3.5-turbo"} + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={}, + team_metadata={}, + ) + + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( + request=request_mock, + request_data=data, + user_api_key_dict=user_api_key_dict, + ) + + tag_object = LiteLLM_TagTable( + tag_name="tenant:acme", + spend=0.0, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.10), + ) + + async def mock_get_current_spend(counter_key, fallback_spend): + if counter_key == "spend:tag:tenant:acme": + return 0.50 + return fallback_spend + + with ( + patch( + "litellm.proxy.proxy_server.get_current_spend", + mock_get_current_spend, + ), + patch( + "litellm.proxy.auth.auth_checks.get_tag_objects_batch", + new_callable=AsyncMock, + return_value={"tenant:acme": tag_object}, + ), + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _tag_max_budget_check( + request_body=data, + prisma_client=MagicMock(), + user_api_key_cache=MagicMock(), + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + valid_token=UserAPIKeyAuth(token="test-token"), + ) + assert exc_info.value.current_cost == 0.50 + assert exc_info.value.max_budget == 0.10 + + # ============================================================================ # Tests for #27516: provider hint resolution from deployment when the # user-facing model name has no provider prefix. @@ -3864,9 +4145,7 @@ def test_resolve_provider_from_deployment_uses_litellm_params_model(): deployment.litellm_params.custom_llm_provider = None router.get_deployment_by_model_group_name.return_value = deployment - assert ( - _resolve_provider_from_deployment(router, "claude-sonnet-4.6") == "bedrock" - ) + assert _resolve_provider_from_deployment(router, "claude-sonnet-4.6") == "bedrock" def test_resolve_provider_from_deployment_prefers_custom_llm_provider(): @@ -3877,9 +4156,7 @@ def test_resolve_provider_from_deployment_prefers_custom_llm_provider(): deployment.litellm_params.custom_llm_provider = "bedrock" router.get_deployment_by_model_group_name.return_value = deployment - assert ( - _resolve_provider_from_deployment(router, "claude-sonnet-4.6") == "bedrock" - ) + assert _resolve_provider_from_deployment(router, "claude-sonnet-4.6") == "bedrock" def test_resolve_provider_from_deployment_no_match():