From 002eadb6b359581b0ae60c97bd6d558075fbf566 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 7 Nov 2026 01:42:35 +0000 Subject: [PATCH] fix(router): preserve routing tags after guardrail metadata enrichment Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_rust/core.py | 16 +- litellm/router_rust/steps/prerouting.py | 11 +- litellm/router_strategy/budget_limiter.py | 4 - litellm/router_strategy/tag_based_routing.py | 22 ++- .../test_unified_guardrail.py | 40 ++++- .../router_strategy/test_budget_limiter.py | 7 +- .../test_router_tag_routing.py | 163 +++++++++++++++--- 7 files changed, 219 insertions(+), 44 deletions(-) diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 3e222313f32..70a99853d8e 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -3248,15 +3248,17 @@ class RustCore( async def _answer_tags(self, step: Mapping[str, object], request: Mapping[str, object]) -> object: """``get_deployments_for_tag`` over the step's candidates (bridged: not ported to Rust yet).""" - from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs - from litellm.router_strategy.tag_based_routing import get_deployments_for_tag + from litellm.router_strategy.tag_based_routing import ( + get_deployments_for_tag, + get_tag_metadata_variable_name_from_kwargs, + ) kept: Final = await get_deployments_for_tag( ctx=RoutingContext(self), # pyright: ignore[reportArgumentType] # the filter reads registry() and settings() model=str(step.get("model")), healthy_deployments=self._records_of(step.get("candidates")), # pyright: ignore[reportArgumentType] # deployment dicts request_kwargs=dict(request), # pyright: ignore[reportArgumentType] # the request kwargs - metadata_variable_name=get_metadata_variable_name_from_kwargs(dict(request)), + metadata_variable_name=get_tag_metadata_variable_name_from_kwargs(request), request_deployments=self._records_of(step.get("pool")), # pyright: ignore[reportArgumentType] # deployment dicts ) return [found for record in kept if (found := picked_id(record)) is not None] @@ -3644,10 +3646,12 @@ class RustCore( """The cooling deployments in ``routable`` that tag routing would send the request to, each pick counted along with the picks it falls back to once that one is unavailable, judged without recording a tag match on the request.""" - from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs - from litellm.router_strategy.tag_based_routing import get_deployments_for_tag + from litellm.router_strategy.tag_based_routing import ( + get_deployments_for_tag, + get_tag_metadata_variable_name_from_kwargs, + ) - field: Final = get_metadata_variable_name_from_kwargs(dict(request)) + field: Final = get_tag_metadata_variable_name_from_kwargs(request) metadata: Final = request.get(field) if field in request and not isinstance(metadata, Mapping): return frozenset() diff --git a/litellm/router_rust/steps/prerouting.py b/litellm/router_rust/steps/prerouting.py index 47bd1e69046..4cb02334b6c 100644 --- a/litellm/router_rust/steps/prerouting.py +++ b/litellm/router_rust/steps/prerouting.py @@ -53,7 +53,11 @@ from litellm.router_strategy.complexity_router.context_compaction import ( is_native_compaction_call, reject_recursive_compactor, ) -from litellm.router_strategy.tag_based_routing import get_routing_tags_from_request_kwargs, is_valid_deployment_tag +from litellm.router_strategy.tag_based_routing import ( + get_routing_tags_from_request_kwargs, + get_tag_metadata_variable_name_from_kwargs, + is_valid_deployment_tag, +) from litellm.router_utils.fallback_event_handlers import record_pre_routing_selection from litellm.types.router import ConsumedRequestTagsStamp, Deployment, LiteLLM_Params, TaggedPreRoutingStrategy from litellm.types.utils import ( @@ -562,6 +566,11 @@ def _stamp_or_clear_metadata_key( if isinstance(bucket, dict): bucket.pop(key, None) return + if key == CONSUMED_REQUEST_TAGS_METADATA_KEY: + selected_metadata: Final = request_kwargs.get(get_tag_metadata_variable_name_from_kwargs(request_kwargs)) + if isinstance(selected_metadata, dict): + selected_metadata[key] = value + return _, metadata_bucket = get_or_create_metadata_bucket(request_kwargs) metadata_bucket[key] = value diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index c65820b9892..154140861c2 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -43,7 +43,6 @@ from litellm.litellm_core_utils.batch_utils import ( job_to_charge_once, ) from litellm.litellm_core_utils.core_helpers import ( - get_metadata_variable_name_from_kwargs, require_attributes, ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds_from @@ -262,7 +261,6 @@ class RouterBudgetLimiting(CustomLogger): potential_deployments=potential_deployments, request_tags=get_tags_from_request_kwargs( request_kwargs=request_kwargs, - metadata_variable_name=get_metadata_variable_name_from_kwargs(request_kwargs or {}), ), ) @@ -389,7 +387,6 @@ class RouterBudgetLimiting(CustomLogger): if self.tag_budget_config: _request_tags = get_tags_from_request_kwargs( request_kwargs=request_kwargs, - metadata_variable_name=get_metadata_variable_name_from_kwargs(request_kwargs or {}), ) for deployment in healthy_deployments: @@ -618,7 +615,6 @@ class RouterBudgetLimiting(CustomLogger): request_tags: Final = get_tags_from_request_kwargs( callback_fields, - metadata_variable_name=get_metadata_variable_name_from_kwargs(callback_fields), ) if len(request_tags) > 0: for _tag in request_tags: diff --git a/litellm/router_strategy/tag_based_routing.py b/litellm/router_strategy/tag_based_routing.py index f2aafa5b7df..10be1137648 100644 --- a/litellm/router_strategy/tag_based_routing.py +++ b/litellm/router_strategy/tag_based_routing.py @@ -772,7 +772,7 @@ def get_tags_from_request_kwargs( """ if request_kwargs is None: return [] - resolved_variable_name: Final = metadata_variable_name or get_metadata_variable_name_from_kwargs(request_kwargs) + resolved_variable_name: Final = metadata_variable_name or get_tag_metadata_variable_name_from_kwargs(request_kwargs) if resolved_variable_name in request_kwargs: return _tags_in_metadata(request_kwargs[resolved_variable_name]) if "litellm_params" in request_kwargs: @@ -787,8 +787,26 @@ def get_tags_from_request_kwargs( _get_tags_from_request_kwargs = get_tags_from_request_kwargs -def _request_metadata(request_kwargs: Mapping[str, object]) -> object: +def get_tag_metadata_variable_name_from_kwargs( + request_kwargs: Mapping[str, object], +) -> Literal["metadata", "litellm_metadata"]: metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(request_kwargs) + metadata: Final = request_kwargs.get("litellm_metadata") + legacy_metadata: Final = request_kwargs.get("metadata") + if ( + metadata_variable_name == "litellm_metadata" + and isinstance(metadata, Mapping) + and "tags" not in metadata + and "inherited_tags" not in metadata + and isinstance(legacy_metadata, Mapping) + and "inherited_tags" in legacy_metadata + ): + return "metadata" + return metadata_variable_name + + +def _request_metadata(request_kwargs: Mapping[str, object]) -> object: + metadata_variable_name: Final = get_tag_metadata_variable_name_from_kwargs(request_kwargs) if metadata_variable_name in request_kwargs: return request_kwargs[metadata_variable_name] litellm_params: Final = request_kwargs.get("litellm_params") diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 2229741fce5..b6990f270e4 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -48,6 +48,7 @@ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrai from litellm.proxy.mcp_server.guardrail_translation.handler import ( MCPGuardrailTranslationHandler, ) +from litellm.router_strategy.tag_based_routing import get_tags_from_request_kwargs from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.responses.main import build_web_search_call @@ -1315,9 +1316,9 @@ class TestStreamingTransform: [server_chunk, _stream_chunk("done", finish_reason="stop")], ) - delivered: Final = ( - [getattr(item.choices[0].delta, "provider_specific_fields", None) for item in out if item.choices] - ) + delivered: Final = [ + getattr(item.choices[0].delta, "provider_specific_fields", None) for item in out if item.choices + ] assert server_fields in delivered assert "".join(_delta_text(item) for item in out) == "DONE" @@ -3254,6 +3255,39 @@ class TestTranslationMappingsAreReadLive: ] +@pytest.mark.parametrize("hook", ["async_pre_call_hook", "async_moderation_hook"]) +@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) +@pytest.mark.asyncio +async def test_input_guardrails_preserve_routing_tags_and_existing_metadata( + monkeypatch: pytest.MonkeyPatch, hook: str, metadata_key: str +) -> None: + _patch_translation_mappings(monkeypatch, {CallTypes.acompletion: OpenAIChatCompletionsHandler}) + guardrail: Final = RecordingGuardrail() + metadata: Final = {"tags": ["region-b"], "inherited_tags": [], "user_api_key_user_id": "request-user"} + data: Final = { + "guardrail_to_apply": guardrail, + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], + metadata_key: metadata, + } + user_api_key_dict: Final = UserAPIKeyAuth(user_id="auth-user", user_email="auth@example.com") + + if hook == "async_pre_call_hook": + await UnifiedLLMGuardrails().async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=DualCache(), data=data, call_type=CallTypes.acompletion.value + ) + else: + await UnifiedLLMGuardrails().async_moderation_hook( + data=data, user_api_key_dict=user_api_key_dict, call_type=CallTypes.acompletion.value + ) + + assert len(guardrail.apply_calls) == 1 + assert guardrail.apply_calls[0]["inputs"]["texts"] == ["hello"] + assert get_tags_from_request_kwargs(data) == ["region-b"] + assert data[metadata_key] is metadata + assert metadata["user_api_key_user_id"] == "request-user" + + class AttachmentScanningGuardrail(RecordingGuardrail): """Records the raw request it is handed for an attachment scan.""" diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index d24d0598d87..7db85d221bf 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -123,9 +123,11 @@ async def test_budget_of_other_provider_is_untouched(disable_budget_sync): @pytest.mark.asyncio +@pytest.mark.parametrize("shadowed", [False, True]) async def test_tag_spend_is_incremented_for_request_metadata( disable_budget_sync: None, monkeypatch: pytest.MonkeyPatch, + shadowed: bool, ) -> None: monkeypatch.setattr(litellm, "tag_budget_config", {"release": {"budget_limit": 5.0, "time_period": "1h"}}) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) @@ -136,7 +138,8 @@ async def test_tag_spend_is_incremented_for_request_metadata( provider_in_payload=None, response_cost=0.73, ), - "metadata": {"tags": ["release"]}, + "metadata": {"tags": ["release"], "inherited_tags": []}, + **({"litellm_metadata": {"user_api_key_user_id": "auth-user"}} if shadowed else {}), } await _log_success(limiter, kwargs) @@ -503,7 +506,6 @@ async def test_batch_polls_charge_each_router_budget_once(monkeypatch, call_type assert tuple([await cache.async_get_cache(key) for key in keys]) == (1.0, 1.0, 1.0) - @pytest.mark.asyncio async def test_a_batch_retrieve_before_its_cost_is_final_charges_no_router_budget_even_for_its_guardrail(): from litellm.types.utils import LiteLLMBatch @@ -627,6 +629,7 @@ async def test_every_video_poll_pays_its_guardrail_and_only_the_finished_poll_pa keys: Final = ("provider_spend:openai:1d", "deployment_spend:deployment-1:1d") assert tuple([await cache.async_get_cache(key) for key in keys]) == pytest.approx((0.29, 0.29)) + @pytest.mark.asyncio async def test_batch_charge_is_shared_across_workers_and_requires_redis_confirmation(): from fakeredis import FakeServer, FakeRedis as FakeSyncRedis diff --git a/tests/unit/router_strategy/test_router_tag_routing.py b/tests/unit/router_strategy/test_router_tag_routing.py index e916e5b235d..5d5d8b4cb4e 100644 --- a/tests/unit/router_strategy/test_router_tag_routing.py +++ b/tests/unit/router_strategy/test_router_tag_routing.py @@ -11,7 +11,7 @@ import litellm from litellm._logging import verbose_logger from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.router_strategy.tag_based_routing import get_deployments_for_tag -from litellm.types.router import RouterErrors +from litellm.types.router import RouterErrors, RouterRateLimitError async def _routed_model_ids( @@ -269,9 +269,7 @@ async def test_error_from_tag_routing(): from litellm.types.router import RouterErrors - with pytest.raises( - Exception, match=RouterErrors.no_deployments_with_tag_routing.value - ): + with pytest.raises(Exception, match=RouterErrors.no_deployments_with_tag_routing.value): await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "Tell me a joke."}], @@ -452,6 +450,75 @@ def test_get_tags_from_request_kwargs_reads_no_tags_from_a_non_dict_shape(reques assert get_tags_from_request_kwargs(request_kwargs) == [] +@pytest.mark.parametrize( + "modern_metadata, expected", + [ + ({"user_api_key_user_id": "auth-user"}, ["region-b"]), + ({}, ["region-b"]), + ({"tags": ["region-a"]}, ["region-a"]), + ({"tags": []}, []), + ({"tags": None}, []), + ({"tags": "region-a"}, []), + ({"inherited_tags": [], "user_agent": "client"}, []), + (None, []), + ([], []), + ], +) +def test_auth_metadata_does_not_hide_legacy_tags(modern_metadata: object, expected: list[str]) -> None: + from litellm.router_strategy.tag_based_routing import get_tags_from_request_kwargs + + request_kwargs: Final = { + "metadata": {"tags": ["region-b"], "inherited_tags": []}, + "litellm_metadata": modern_metadata, + } + + assert get_tags_from_request_kwargs(request_kwargs) == expected + assert get_tags_from_request_kwargs(request_kwargs, metadata_variable_name="metadata") == ["region-b"] + assert get_tags_from_request_kwargs(request_kwargs, metadata_variable_name="litellm_metadata") == ( + expected if isinstance(modern_metadata, dict) and "tags" in modern_metadata else [] + ) + + +def test_auth_metadata_does_not_turn_baggage_tags_into_routing_tags() -> None: + from litellm.router_strategy.tag_based_routing import get_routing_tags_from_request_kwargs + + metadata: Final = _baggage_mapped_metadata(["region-b"], []) + request_kwargs: Final = {"metadata": metadata, "litellm_metadata": {"user_api_key_user_id": "auth-user"}} + + assert tuple(get_routing_tags_from_request_kwargs(request_kwargs)) == ("region-b",) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tags", [["region-b"], ["®ion-b"], ["!region-a"]]) +@pytest.mark.parametrize("inherited", [False, True]) +async def test_tag_filter_uses_policy_metadata_after_guardrail_enrichment(tags: list[str], inherited: bool) -> None: + router: Final = litellm.Router( + model_list=[ + { + "model_name": "chat", + "litellm_params": {"model": "openai/gpt-4o-mini", "tags": ["default", "region-a"], "order": 1}, + "model_info": {"id": "deployment-a"}, + }, + { + "model_name": "chat", + "litellm_params": {"model": "openai/gpt-4o-mini", "tags": ["default", "region-b"], "order": 2}, + "model_info": {"id": "deployment-b"}, + }, + ], + enable_tag_filtering=True, + ) + deployment: Final = await router.async_get_available_deployment( + model="chat", + request_kwargs={ + "metadata": {"tags": tags, "inherited_tags": tags if inherited else []}, + "litellm_metadata": {"user_api_key_user_id": "auth-user"}, + }, + messages=[{"role": "user", "content": "hello"}], + ) + + assert deployment["model_info"]["id"] == "deployment-b" + + def test_get_tags_from_request_kwargs_keeps_only_string_tags(): from litellm.router_strategy.tag_based_routing import get_tags_from_request_kwargs @@ -665,7 +732,7 @@ async def test_negation_all_excluded_raises(): enable_tag_filtering=True, ) - with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: + with pytest.raises(Exception, match="Not allowed to access model due to tags configuration\\.") as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -708,7 +775,7 @@ async def test_negation_ban_only_cannot_escape_default_pool(): # Sending only "!default" must NOT route to the paid deployment. # The base pool for ban-only is the default pool; banning the only # default deployment should raise rather than falling through to paid. - with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: + with pytest.raises(Exception, match="Not allowed to access model due to tags configuration\\.") as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -973,7 +1040,7 @@ async def test_negation_exhausts_entire_fallback_chain(): enable_tag_filtering=True, ) - with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: + with pytest.raises(Exception, match="Not allowed to access model due to tags configuration\\.") as exc_info: await router.acompletion( model="primary", messages=[{"role": "user", "content": "hi"}], @@ -1754,7 +1821,7 @@ async def test_required_and_unmatched_raises_by_default(): enable_tag_filtering=True, ) - with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: + with pytest.raises(Exception, match="Not allowed to access model due to tags configuration\\.") as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -1787,7 +1854,7 @@ async def test_required_and_combined_with_positive_unmatched_raises_by_default() enable_tag_filtering=True, ) - with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: + with pytest.raises(Exception, match="Not allowed to access model due to tags configuration\\.") as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -1856,9 +1923,7 @@ async def test_tag_fail_open_survives_deployment_narrowing( model=requested_model, request_kwargs={ "litellm_metadata": {"tags": list(tags), "inherited_tags": [], "user_api_key_team_id": "team-a"}, - **( - {"previous_response_id": _response_id_for_remaining_deployment()} if narrowing == "affinity" else {} - ), + **({"previous_response_id": _response_id_for_remaining_deployment()} if narrowing == "affinity" else {}), }, ) assert isinstance(candidates, (list, tuple)) @@ -2145,7 +2210,7 @@ async def test_negation_combined_with_positive_unmatched_raises_by_default(): enable_tag_filtering=True, ) - with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: + with pytest.raises(Exception, match="Not allowed to access model due to tags configuration\\.") as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -2303,7 +2368,7 @@ async def test_mixed_constraint_survivor_unmatched_by_positive_tag_raises_by_def enable_tag_filtering=True, ) - with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: + with pytest.raises(Exception, match="Not allowed to access model due to tags configuration\\.") as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -2396,7 +2461,7 @@ async def test_allow_fail_open_denied_when_request_includes_unknown_tag(): enable_tag_filtering=True, ) - with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: + with pytest.raises(Exception, match="Not allowed to access model due to tags configuration\\.") as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -2705,6 +2770,44 @@ def _cool_down(router: litellm.Router, *model_ids: str) -> None: ) +@pytest.mark.asyncio() +async def test_cooled_tag_probe_does_not_mutate_the_policy_bucket(): + router: Final = _quality_high_cost_low_router() + _cool_down(router, "quality-high-1", "quality-high-2", "cost-low-1", "cost-low-2") + metadata: Final = {"tags": ["quality:high"], "inherited_tags": []} + + with pytest.raises(RouterRateLimitError) as exc_info: + await router.async_get_available_deployment( + model="gpt-4", + request_kwargs={"metadata": metadata, "litellm_metadata": {"user_api_key_user_id": "auth-user"}}, + messages=[{"role": "user", "content": "hello"}], + ) + + assert metadata == {"tags": ["quality:high"], "inherited_tags": []} + assert exc_info.value.cooldown_deployments is not None + assert {deployment["model_id"] for deployment in exc_info.value.cooldown_deployments} == { + "quality-high-1", + "quality-high-2", + } + + +@pytest.mark.asyncio() +async def test_modern_route_keeps_header_regex_metadata_separate_from_provider_metadata(): + from tests.unit.router_strategy.test_router_tag_regex_routing import ALL_DEPLOYMENTS + + router: Final = litellm.Router(model_list=ALL_DEPLOYMENTS, enable_tag_filtering=True) + deployment: Final = await router.async_get_available_deployment( + model="claude-sonnet", + request_kwargs={ + "metadata": {"user_id": "provider-user", "tags": ["wrong"], "inherited_tags": ["wrong"]}, + "litellm_metadata": {"user_agent": "claude-code/1.2.3", "inherited_tags": []}, + }, + messages=[{"role": "user", "content": "hello"}], + ) + + assert deployment["model_info"]["id"] == "claude-code-deployment" + + @pytest.mark.asyncio() async def test_plain_tag_exhaustion_with_universal_default_tag_raises_by_default(): # Every deployment in the group is tagged "default" (a legitimate cross-cutting @@ -2715,7 +2818,7 @@ async def test_plain_tag_exhaustion_with_universal_default_tag_raises_by_default router = _quality_high_cost_low_router() _cool_down(router, "quality-high-1", "quality-high-2") - with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: + with pytest.raises(Exception, match="Not allowed to access model due to tags configuration\\.") as exc_info: await router.acompletion( model="gpt-4", messages=[{"role": "user", "content": "hi"}], @@ -2816,7 +2919,9 @@ def test_trusted_only_pool_discards_everything_when_inherited_sets_are_none(): deployments = ({"litellm_params": {"tags": ["region:us"]}},) # No origin info at all -> reproduce the pre-provenance unconditional # fall-open: the trusted-only pool ignores excluded_set/required_set entirely. - assert _trusted_only_pool(deployments, frozenset({"region:eu"}), frozenset({"region:apac"}), None, None) == deployments + assert ( + _trusted_only_pool(deployments, frozenset({"region:eu"}), frozenset({"region:apac"}), None, None) == deployments + ) def test_trusted_only_pool_keeps_constraint_backed_by_inherited_tags(): @@ -2825,9 +2930,7 @@ def test_trusted_only_pool_keeps_constraint_backed_by_inherited_tags(): eu = {"litellm_params": {"tags": ["region:eu"]}} us = {"litellm_params": {"tags": ["region:us"]}} # required_set={"region:eu"} IS in inherited_required_set -> protected, kept. - result = _trusted_only_pool( - (eu, us), frozenset(), frozenset({"region:eu"}), frozenset(), frozenset({"region:eu"}) - ) + result = _trusted_only_pool((eu, us), frozenset(), frozenset({"region:eu"}), frozenset(), frozenset({"region:eu"})) assert result == (eu,) @@ -2940,7 +3043,7 @@ async def test_allow_fail_open_raises_when_inherited_constraint_alone_is_unsatis # allow_fail_open unset. router = _eu_region_router() - with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: + with pytest.raises(Exception, match="Not allowed to access model due to tags configuration\\.") as exc_info: await router.acompletion( model="chat", messages=[{"role": "user", "content": "hi"}], @@ -3076,7 +3179,8 @@ async def test_router_selecting_tag_is_not_reapplied_to_the_routed_tier(): @pytest.mark.asyncio() -async def test_router_selecting_tag_is_consumed_on_litellm_metadata_shaped_requests(): +@pytest.mark.parametrize("shadowed", [False, True]) +async def test_router_selecting_tag_is_consumed_on_litellm_metadata_shaped_requests(shadowed: bool): # /v1/messages (and other litellm_metadata endpoints) store proxy metadata, # including x-litellm-tags header tags, under "litellm_metadata"; consumption # must read and stamp that same bucket instead of only "metadata". @@ -3084,7 +3188,14 @@ async def test_router_selecting_tag_is_consumed_on_litellm_metadata_shaped_reque deployment = await router.async_get_available_deployment( model="gpt4o", - request_kwargs={"litellm_metadata": {"tags": ["route"], "inherited_tags": []}}, + request_kwargs=( + { + "metadata": {"tags": ["route"], "inherited_tags": []}, + "litellm_metadata": {"user_api_key_user_id": "auth-user"}, + } + if shadowed + else {"litellm_metadata": {"tags": ["route"], "inherited_tags": []}} + ), messages=[{"role": "user", "content": "What is the capital of France?"}], ) @@ -3097,7 +3208,7 @@ async def test_tagged_request_direct_to_plain_group_still_rejected(): # tag filtering must reject exactly as before. router = _tagged_marker_router() - with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: + with pytest.raises(Exception, match="Not allowed to access model due to tags configuration\\.") as exc_info: await router.acompletion( model="gemini-flash", messages=[{"role": "user", "content": "hi"}], @@ -3118,7 +3229,7 @@ async def test_caller_forged_consumption_stamp_is_neutralized_by_the_hook(): # tag filtering runs. router = _tagged_marker_router() - with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: + with pytest.raises(Exception, match="Not allowed to access model due to tags configuration\\.") as exc_info: await router.acompletion( model="gemini-flash", messages=[{"role": "user", "content": "hi"}], @@ -3140,7 +3251,7 @@ async def test_inherited_constraint_still_applies_to_the_routed_tier(): # ®ion:eu comes from key/team policy (present in inherited_tags): # consuming the router-selecting "route" tag must not also discard the # inherited requirement, so a tier without the tag still raises... - with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info: + with pytest.raises(Exception, match="Not allowed to access model due to tags configuration\\.") as exc_info: await _tagged_marker_router().acompletion( model="gpt4o", messages=[{"role": "user", "content": "hi"}],