scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

fix: enforce tag budgets on x-litellm-tags header requests

Browse source 
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) <noreply@anthropic.com>
This commit is contained in:
shivam 2026-05-09 18:08:36 -07:00
parent 02edaef50c
commit 36caeb013b
No known key found for this signature in database
3 changed files with 328 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
litellm/proxy/auth/user_api_key_auth.py
View file

@ -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,

83
litellm/proxy/litellm_pre_call_utils.py
View file

@ -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,

232
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View file

@ -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