From 4a4262aa0d0490ead1859fa34cf999d6dec53736 Mon Sep 17 00:00:00 2001
From: Abhimanyu Kapur <38531241+akapur99@users.noreply.github.com>
Date: Sat, 11 Jul 2026 11:57:32 -0700
Subject: [PATCH] fix(complexity_router): preserve user_api_key_auth in
 sub-call metadata

Removing user_api_key_auth entirely from classifier/embedding sub-call
metadata (as _BUDGET_RESERVATION_METADATA_KEYS previously did) prevented
_filter_deployments_by_model_access_groups from scoping those sub-calls to
the caller's authorized access groups. An access-group-scoped caller could
therefore reach embedding/classifier deployments outside their group.

Only strip user_api_key_budget_reservation, which is the actual budget-
reservation state that must not reach sub-calls. user_api_key_auth is now
kept so access-group filtering works correctly for both the embedding path
and the LLM classifier path.
---
 .../complexity_router/complexity_router.py     | 18 ++++++++++++------
 .../router_strategy/test_complexity_router.py  | 17 ++++++++++++++---
 2 files changed, 26 insertions(+), 9 deletions(-)

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