ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #27573 from BerriAI/litellm_tag_budget_header_enforcement

Browse Source 
fix: enforce tag budgets on x-litellm-tags header requests
This commit is contained in:
yuneng-jiang 2026-05-12 18:34:19 -07:00 committed by GitHub
commit 924e8b243f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 362 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
litellm/proxy/auth/user_api_key_auth.py
View File

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

68
litellm/proxy/litellm_pre_call_utils.py
View File

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

289
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View File

@ -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():