diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 6ddf2cfeb20..05d532c8e21 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -699,6 +699,11 @@ async def common_checks( if valid_token is not None: from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route( + request_data=request_body, + route=route, + ) + LiteLLMProxyRequestSetup.apply_key_tags_pre_auth( request_data=request_body, user_api_key_dict=valid_token, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 4a2df18b93b..e439f6a5998 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2396,6 +2396,17 @@ async def _run_centralized_common_checks( llm_router=llm_router, ) + # Pin the metadata variable name (litellm_metadata vs metadata) before + # any tag merge runs. Without this, header tags from + # apply_client_tag_policy_pre_auth would land in `metadata` while the + # later seed in common_checks pushes key tags and the + # _tag_max_budget_check read into `litellm_metadata`, hiding header + # tags from per-tag budget enforcement on LITELLM_METADATA_ROUTES. + LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route( + request_data=request_data, + route=route, + ) + # 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 diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 177fced5cd4..c0cdf84dfb6 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -108,6 +108,7 @@ def parse_cache_control(cache_control): LITELLM_METADATA_ROUTES = ( "batches", + "bedrock", "/v1/messages", "responses", "files", @@ -1237,6 +1238,27 @@ def add_request_tag_to_metadata( return tags + @staticmethod + def pre_seed_litellm_metadata_for_route( + request_data: dict, + route: str, + ) -> None: + """Pre-seed ``litellm_metadata`` for routes that track tags there. + + Routes in ``LITELLM_METADATA_ROUTES`` (e.g. Bedrock, ``/v1/messages``, + responses, batches, files) store request-scoped tag metadata in + ``litellm_metadata`` rather than the provider-facing ``metadata`` + field. ``get_metadata_variable_name_from_kwargs`` picks the target + based on whether ``litellm_metadata`` is present, so it must be + seeded BEFORE any tag merge runs; otherwise header tags from + ``apply_client_tag_policy_pre_auth`` land in ``metadata`` while + key tags from ``apply_key_tags_pre_auth`` and the read in + ``_tag_max_budget_check`` resolve to ``litellm_metadata``, leaving + header tags invisible to per-tag budget enforcement. + """ + if any(metadata_route in route for metadata_route in LITELLM_METADATA_ROUTES): + request_data.setdefault("litellm_metadata", {}) + @staticmethod def apply_key_tags_pre_auth( request_data: dict, @@ -1468,8 +1490,7 @@ async def add_litellm_data_to_request( _metadata_variable_name=_metadata_variable_name, ) - # Add headers to metadata for guardrails to access (fixes #17477) - # Guardrails use metadata["headers"] to access request headers (e.g., User-Agent) + # Expose request headers under the metadata field for guardrails (fixes #17477) if _metadata_variable_name in data and isinstance( data[_metadata_variable_name], dict ): diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 52634cc25fe..cf8448c0383 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -1750,6 +1750,60 @@ async def test_reject_clientside_metadata_tags_allows_key_tags_without_client_ta assert request_body["metadata"]["tags"] == ["engineering"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route", + [ + "/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke", + "/v1/messages", + ], +) +async def test_common_checks_metadata_route_keeps_key_tags_out_of_provider_metadata( + route, +): + """GH#30629: on routes that track tags in litellm_metadata (bedrock, /v1/messages, + responses, ...) key-level tags must land in litellm_metadata, never in the + provider-facing metadata field (Bedrock rejects non-user_id metadata with HTTP 400). + The auth-time pre-seed keys off LITELLM_METADATA_ROUTES, so hardcoding a single route + or dropping the pre-seed makes apply_key_tags_pre_auth fall back to metadata; this + guards that regression. + """ + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + request_body = {"messages": [{"role": "user", "content": "test"}]} + + mock_request = MagicMock(spec=Request) + valid_token = UserAPIKeyAuth( + token="test-token", + metadata={"tags": ["engineering"]}, + ) + + with patch( + "litellm.proxy.auth.auth_checks.get_tag_objects_batch", + new_callable=AsyncMock, + return_value={}, + ): + result = await common_checks( + request_body=request_body, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route=route, + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=valid_token, + request=mock_request, + ) + + assert result is True + assert request_body["litellm_metadata"]["tags"] == ["engineering"] + assert "metadata" not in request_body + + @pytest.mark.asyncio async def test_virtual_key_soft_budget_check_with_user_obj(): """Test _virtual_key_soft_budget_check includes user_email when user_obj is provided""" diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 27c1b04fbd9..7219ab58799 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -2689,6 +2689,67 @@ async def test_centralized_common_checks_runs_for_standard_auth(): setattr(_proxy_server_mod, k, v) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "route", + [ + "/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke", + "/v1/messages", + ], +) +async def test_centralized_common_checks_routes_header_tags_to_litellm_metadata(route): + """GH#30629: on LITELLM_METADATA_ROUTES the tag-budget read resolves to + litellm_metadata, so the litellm_metadata pre-seed must run before + apply_client_tag_policy_pre_auth merges x-litellm-tags. Otherwise header tags + land in metadata and silently escape _tag_max_budget_check. This guards the + pre-seed call site in _run_centralized_common_checks; dropping it routes header + tags back into metadata. + """ + import litellm.proxy.proxy_server as _proxy_server_mod + from fastapi import Request + from starlette.datastructures import URL + + token = UserAPIKeyAuth(api_key="sk-test", user_id="u1") + request = Request( + scope={ + "type": "http", + "method": "POST", + "headers": [(b"x-litellm-tags", b"tenant:acme")], + "query_string": b"", + } + ) + request._url = URL(url=route) + request_data: dict = {"model": "us.anthropic.claude-sonnet-4-6"} + + attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch( + "litellm.proxy.auth.user_api_key_auth.common_checks", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.auth.user_api_key_auth._reserve_budget_after_common_checks", + new_callable=AsyncMock, + ), + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data=request_data, + route=route, + ) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + assert request_data["litellm_metadata"]["tags"] == ["tenant:acme"] + assert "metadata" not in request_data + + @pytest.mark.asyncio async def test_centralized_common_checks_skipped_for_custom_auth_without_flag(): """Existing RPS guarantee: custom-auth deployments without diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 362f4986c62..379e219b29b 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -2682,15 +2682,19 @@ async def test_add_litellm_data_to_request_adds_headers_to_metadata(): version="1.0", ) - # Verify headers are added to metadata for guardrails - assert "metadata" in result, "metadata should be present in result" - assert "headers" in result["metadata"], "headers should be present in metadata" + # Verify headers are added to litellm_metadata for guardrails. + # Bedrock passthrough uses litellm_metadata to prevent key-level + # tags from leaking into the provider payload (GH#30629). + assert "litellm_metadata" in result, "litellm_metadata should be present in result" + assert ( + "headers" in result["litellm_metadata"] + ), "headers should be present in litellm_metadata" assert isinstance( - result["metadata"]["headers"], dict + result["litellm_metadata"]["headers"], dict ), "headers should be a dictionary" # Verify specific headers are accessible (important for guardrails) - headers = result["metadata"]["headers"] + headers = result["litellm_metadata"]["headers"] assert ( "user-agent" in headers or "User-Agent" in headers ), "User-Agent header should be accessible in metadata" 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 ec262d75ab8..74f681bb97e 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -96,6 +96,20 @@ def test_returns_metadata_for_embeddings(self): request = self._make_request("/v1/embeddings") assert _get_metadata_variable_name(request) == "metadata" + def test_returns_litellm_metadata_for_bedrock_invoke(self): + # GH#30629: bedrock passthrough must use litellm_metadata + # to prevent key-level tags from leaking into provider body + request = self._make_request( + "/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke" + ) + assert _get_metadata_variable_name(request) == "litellm_metadata" + + def test_returns_litellm_metadata_for_bedrock_converse(self): + request = self._make_request( + "/bedrock/model/us.anthropic.claude-sonnet-4-6/converse" + ) + assert _get_metadata_variable_name(request) == "litellm_metadata" + def test_get_enforced_params_for_service_account_settings(): """ @@ -4256,6 +4270,99 @@ async def mock_get_current_spend( assert exc_info.value.current_cost == 0.50 assert exc_info.value.max_budget == 0.10 + @pytest.mark.asyncio + @pytest.mark.parametrize( + "route", + [ + "/bedrock/model/us.anthropic.claude-sonnet-4-6/invoke", + "/v1/messages", + ], + ) + async def test_header_tags_visible_to_tag_max_budget_check_on_metadata_route( + self, route + ): + """Regression: on LITELLM_METADATA_ROUTES (bedrock, /v1/messages, ...), + common_checks pre-seeds ``litellm_metadata`` and writes key tags there + before ``_tag_max_budget_check`` reads from the same key. The auth wrapper + calls ``apply_client_tag_policy_pre_auth`` first, so without an earlier + pre-seed header tags land in ``metadata`` and the budget check (now + resolving to ``litellm_metadata``) silently ignores them. This test mirrors + the actual auth-time call order and verifies that an over-budget + header-supplied tag still trips ``_tag_max_budget_check``. + """ + from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_TagTable + from litellm.proxy.auth.auth_checks import common_checks + from litellm.proxy.utils import ProxyLogging + + request_mock = _build_request_mock_with_headers( + {"x-litellm-tags": "tenant:acme"} + ) + data = {"model": "us.anthropic.claude-sonnet-4-6"} + valid_token = UserAPIKeyAuth( + token="test-token", + api_key="hashed-key", + metadata={}, + team_metadata={}, + ) + + LiteLLMProxyRequestSetup.pre_seed_litellm_metadata_for_route( + request_data=data, + route=route, + ) + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( + request=request_mock, + request_data=data, + user_api_key_dict=valid_token, + ) + + 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, max_budget=None, **kwargs + ): + if counter_key == "spend:tag:tenant:acme": + return 0.50 + return fallback_spend + + with ( + patch( + "litellm.proxy.proxy_server.prisma_client", + MagicMock(), + ), + 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 common_checks( + request_body=data, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route=route, + llm_router=None, + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + valid_token=valid_token, + request=request_mock, + ) + assert exc_info.value.current_cost == 0.50 + assert exc_info.value.max_budget == 0.10 + + assert "metadata" not in data + assert data["litellm_metadata"]["tags"] == ["tenant:acme"] + class TestApplyKeyTagsPreAuth: def test_merges_key_tags_into_metadata(self):