From 36caeb013b7953f813fe097db30b09624973b35a Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 9 May 2026 18:08:36 -0700 Subject: [PATCH] fix: enforce tag budgets on x-litellm-tags header requests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The x-litellm-tags header was merged into request metadata only after the auth chain completed, so _tag_max_budget_check (which reads tags from the request body) silently failed open for header-tagged requests — spend accumulated past max_budget without any 400 budget_exceeded response. Move the client-tag policy (strip-or-merge gated on allow_client_tags) to run before common_checks so header tags are visible to budget enforcement. The post-auth strip+merge in add_litellm_data_to_request stays as defense-in-depth; the new pre-auth helper is idempotent with it. Co-Authored-By: Claude Opus 4.7 (1M context) --- litellm/proxy/auth/user_api_key_auth.py | 13 + litellm/proxy/litellm_pre_call_utils.py | 83 +++++++ .../proxy/test_litellm_pre_call_utils.py | 232 ++++++++++++++++++ 3 files changed, 328 insertions(+) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 4778549befc..995116e483c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1939,6 +1939,19 @@ async def _run_centralized_common_checks( llm_router=llm_router, ) + # Merge x-litellm-tags (or strip body tags when the key/team has not + # opted in via allow_client_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. + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + + 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 a63613c5836..43418428f83 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -24,6 +24,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 @@ -1177,6 +1180,86 @@ class LiteLLMProxyRequestSetup: return tags + @staticmethod + def apply_client_tag_policy_pre_auth( + request: Request, + request_data: dict, + user_api_key_dict: UserAPIKeyAuth, + ) -> None: + """ + Apply the client-tag policy BEFORE auth budget gates run, so + ``_tag_max_budget_check`` (which only inspects ``request_data``) + sees ``x-litellm-tags`` header tags. Without this, header-tagged + requests silently bypass per-tag budget enforcement. + + Mirrors the strip + merge that ``add_litellm_data_to_request`` + performs post-auth, gated on the same ``allow_client_tags`` flag. + + Why: ``add_litellm_data_to_request`` runs after the auth chain has + completed, so any header-supplied tags it merges in are invisible + to ``_tag_max_budget_check``. Running the merge here closes that + gap. The post-auth strip + merge 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``. + """ + _admin_allow_client_tags = False + for _admin_meta in ( + user_api_key_dict.metadata, + user_api_key_dict.team_metadata, + ): + if ( + isinstance(_admin_meta, dict) + and _admin_meta.get("allow_client_tags") is True + ): + _admin_allow_client_tags = True + break + + if not _admin_allow_client_tags: + # Strip any caller-supplied tags so the budget gate doesn't act + # on tags this key/team isn't authorized to set. Matches the + # post-auth strip in add_litellm_data_to_request. + for _meta_key in ("metadata", "litellm_metadata"): + _user_meta = request_data.get(_meta_key) + if isinstance(_user_meta, dict) and "tags" in _user_meta: + _user_meta.pop("tags", None) + if "tags" in request_data: + request_data.pop("tags", None) + return + + 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) + if 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 92611431a15..75dedb4b89b 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -4043,3 +4043,235 @@ def test_get_guardrail_from_metadata_reads_litellm_metadata_when_no_metadata(): assert result == [ "my-guardrail" ], 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_when_opted_in(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={"allow_client_tags": True}, + 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={"allow_client_tags": True}, + 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_strips_body_tags_when_not_opted_in(self): + request_mock = _build_request_mock_with_headers( + {"x-litellm-tags": "tenant:acme"} + ) + data = { + "model": "gpt-3.5-turbo", + "tags": ["root-tag"], + "metadata": {"tags": ["meta-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 "tags" not in data + assert "tags" not in data["metadata"] + assert "tags" not in data["litellm_metadata"] + + def test_does_not_merge_header_tags_when_not_opted_in(self): + # Even with the header set, no opt-in means the header is ignored + # and metadata.tags is not created from it. + 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, + ) + + assert "tags" not in data.get("metadata", {}) + + def test_team_metadata_opt_in_is_honored(self): + 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={"allow_client_tags": True}, + ) + + 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"] + + 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={"allow_client_tags": True}, + 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_when_opted_in(self): + request_mock = _build_request_mock_with_headers({}) + data = {"model": "gpt-3.5-turbo"} + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + metadata={"allow_client_tags": True}, + 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"] + + @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={"allow_client_tags": True}, + 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