diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 2eb69aad2de..e14c1091d3e 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -70,12 +70,18 @@ def _append_custom_keywords(base_keywords: list[str], custom_keywords: Optional[ return [*base_keywords, *deduped_custom.values()] -# Metadata keys that carry the parent request's budget reservation. These must not -# reach the classifier's internal acompletion call: the reservation belongs to the -# routed completion that the classifier is deciding on, not to the classifier call -# itself, and forwarding it would let the classifier's cost-tracking reconcile -# against a reservation it isn't responsible for. -_BUDGET_RESERVATION_METADATA_KEYS = frozenset({"user_api_key_budget_reservation", "user_api_key_auth"}) +# Metadata keys that carry only the parent request's budget reservation state. These +# must not reach internal sub-calls (classifier, embedding): the reservation belongs to +# the routed completion being decided on, not to the sub-call itself, and forwarding it +# would let the sub-call's cost callback finalize the reservation, causing the routed +# completion's callback to skip incrementing key/team budget counters. +# +# Note: user_api_key_auth itself is intentionally kept; it is required by +# _filter_deployments_by_model_access_groups to scope embedding/classifier model +# selection to the caller's authorized access groups. Only the budget reservation +# sub-field inside it is the problem -- stripping the whole object would allow +# an access-group-scoped caller to reach embedding deployments outside their group. +_BUDGET_RESERVATION_METADATA_KEYS = frozenset({"user_api_key_budget_reservation"}) def _classifier_call_metadata(metadata: Optional[dict[str, Any]]) -> Optional[dict[str, Any]]: diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 3d92700032b..5f408d1eae4 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -1254,15 +1254,19 @@ class TestLLMClassifier: "user_api_key": "sk-abc", "user_api_key_team_id": "team-1", "user_api_key_budget_reservation": {"reserved_cost": 1.0}, - "user_api_key_auth": {"budget_reservation": {"reserved_cost": 1.0}}, + "user_api_key_auth": {"models": ["gpt-4o"], "budget_reservation": {"reserved_cost": 1.0}}, } await llm_complexity_router.aclassify( "hi", request_kwargs={"litellm_metadata": request_metadata} ) call_kwargs = mock_router_instance.acompletion.call_args.kwargs + # user_api_key_budget_reservation is stripped (budget enforcement) but + # user_api_key_auth is kept so _filter_deployments_by_model_access_groups + # can scope the classifier's model selection to the caller's access groups. assert call_kwargs["metadata"] == { "user_api_key": "sk-abc", "user_api_key_team_id": "team-1", + "user_api_key_auth": {"models": ["gpt-4o"], "budget_reservation": {"reserved_cost": 1.0}}, } @pytest.mark.asyncio @@ -1638,7 +1642,7 @@ class TestSemanticKeywordTierRules: "user_api_key_hash": "hash-abc", "user_api_key_team_id": "team-1", "user_api_key_budget_reservation": {"reserved_cost": 1.0}, - "user_api_key_auth": {"budget_reservation": {"reserved_cost": 1.0}}, + "user_api_key_auth": {"models": ["voyage-3-5"], "budget_reservation": {"reserved_cost": 1.0}}, } await router.async_pre_routing_hook( model="test-model", @@ -1646,7 +1650,14 @@ class TestSemanticKeywordTierRules: messages=[{"role": "user", "content": "roll out my k8s cluster"}], ) assert fake_router.async_embedding_kwargs, "expected an embedding call for the prompt" - expected = {"user_api_key_hash": "hash-abc", "user_api_key_team_id": "team-1"} + # user_api_key_budget_reservation is stripped to prevent budget-bypass. + # user_api_key_auth is kept so _filter_deployments_by_model_access_groups + # scopes the embedding model selection to the caller's authorized groups. + expected = { + "user_api_key_hash": "hash-abc", + "user_api_key_team_id": "team-1", + "user_api_key_auth": {"models": ["voyage-3-5"], "budget_reservation": {"reserved_cost": 1.0}}, + } assert fake_router.async_embedding_kwargs[0]["metadata"] == expected assert fake_router.async_embedding_kwargs[0]["litellm_metadata"] == expected