scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

fix(complexity_router): preserve user_api_key_auth in sub-call metadata

Browse source 
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.
This commit is contained in:
Abhimanyu Kapur 2026-07-11 11:57:32 -07:00
parent 4486d9fe0c
commit 4a4262aa0d
2 changed files with 26 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
litellm/router_strategy/complexity_router/complexity_router.py
View file

@ -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]]:

17
tests/test_litellm/router_strategy/test_complexity_router.py
View file

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