From 08f2e924f3c642d3716cc11bdead7809cec9f505 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 5 Jun 2026 09:13:30 -0700 Subject: [PATCH] feat(proxy): auth_v2 data-plane honors model access groups (v1 parity) The inference gate now allows a model when the key lists an access-group name the model belongs to, mirroring v1 model_in_access_group. The group lookup is resolved from the router in the entry point and injected into can_call_model, which stays a pure predicate. Closes the access-group parity gap flagged when the data plane moved off casbin; name, wildcard, and sentinel matching are unchanged. --- litellm/proxy/auth/v2/data_plane.py | 21 +++++++++++++---- litellm/proxy/auth/v2/entry.py | 23 +++++++++++++++++-- .../proxy/auth/v2/test_data_plane.py | 23 +++++++++++++++++++ 3 files changed, 60 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/auth/v2/data_plane.py b/litellm/proxy/auth/v2/data_plane.py index 0ee599126f0..ae9aaa014b0 100644 --- a/litellm/proxy/auth/v2/data_plane.py +++ b/litellm/proxy/auth/v2/data_plane.py @@ -1,5 +1,5 @@ import re -from typing import List, Optional +from typing import Iterable, List, Optional # Sentinels that mean "any model" in the existing key/team model lists. _UNRESTRICTED_SENTINELS = {"*", "all-proxy-models", "all-team-models"} @@ -19,19 +19,30 @@ def _matches_pattern(requested_model: str, pattern: str) -> bool: return bool(re.match("^" + pattern.replace("*", ".*") + "$", requested_model)) -def can_call_model(allowed_models: Optional[List[str]], requested_model: str) -> bool: +def can_call_model( + allowed_models: Optional[List[str]], + requested_model: str, + model_access_groups: Optional[Iterable[str]] = None, +) -> bool: """Decide whether a principal with ``allowed_models`` may call ``requested_model``. Data-plane access is a direct membership/pattern predicate, not a policy engine: it runs on the inference hot path where a casbin evaluation would be pure overhead for what is a list check. Empty list or a sentinel means unrestricted; an exact name matches; a wildcard pattern (e.g. ``bedrock/*``) - matches using v1's pattern semantics. Access-group expansion is not yet - honored here (tracked as a parity follow-up). + matches using v1's pattern semantics. + + ``model_access_groups`` are the access-group names ``requested_model`` belongs + to (from the router); if the key lists any of them the call is allowed, + mirroring v1 ``model_in_access_group``. Injected rather than resolved here so + this stays a pure predicate. """ models = list(allowed_models or []) if _is_unrestricted(models): return True if requested_model in models: return True - return any(_matches_pattern(requested_model, model) for model in models) + if any(_matches_pattern(requested_model, model) for model in models): + return True + groups = model_access_groups or () + return any(model in groups for model in models) diff --git a/litellm/proxy/auth/v2/entry.py b/litellm/proxy/auth/v2/entry.py index 5663ce9e1f7..2bc700ac1a7 100644 --- a/litellm/proxy/auth/v2/entry.py +++ b/litellm/proxy/auth/v2/entry.py @@ -17,6 +17,22 @@ async def _anonymous_identity(api_key: Optional[str]) -> Any: return UserAPIKeyAuth(api_key=api_key) +def _model_access_groups(requested_model: str) -> Any: + """Access-group names the requested model belongs to, via the router. + + Returns an empty tuple when no router is configured so the caller degrades to + plain name/pattern matching. + """ + from litellm.proxy.proxy_server import llm_router + + if llm_router is None: + return () + try: + return llm_router.get_model_access_groups(model_name=requested_model) + except Exception: + return () + + async def _best_effort_identity(api_key: Optional[str], ctx: AuthContext) -> Any: """On loud-open routes, use the real identity if a usable key is present, otherwise fall back to an anonymous principal. Never fails the request.""" @@ -80,14 +96,17 @@ async def user_api_key_auth_v2( return identity if is_inference_route(route): - # Data plane: plain allowed-model predicate over the principal's key. + # Data plane: plain allowed-model predicate over the principal's key, + # with access-group expansion resolved from the router. identity = await authenticate(token, ctx) request_data = await _read_request_body(request=request) requested_model = ( request_data.get("model") if isinstance(request_data, dict) else None ) if requested_model and not can_call_model( - getattr(identity, "models", None), requested_model + getattr(identity, "models", None), + requested_model, + _model_access_groups(requested_model), ): raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, diff --git a/tests/test_litellm/proxy/auth/v2/test_data_plane.py b/tests/test_litellm/proxy/auth/v2/test_data_plane.py index 7ebd28dcc49..8837bd54db7 100644 --- a/tests/test_litellm/proxy/auth/v2/test_data_plane.py +++ b/tests/test_litellm/proxy/auth/v2/test_data_plane.py @@ -49,3 +49,26 @@ def test_partial_wildcard_within_provider(): def test_exact_name_without_wildcard_does_not_pattern_match(): # No '*' -> exact membership only, never a substring/regex match. assert can_call_model(["gpt-4o"], "gpt-4o-mini") is False + + +def test_access_group_membership_grants_access(): + # Mirrors v1 model_in_access_group: the requested model belongs to "beta", and + # the key lists the group name, so the call is allowed. + groups = {"beta": ["o1", "o1-mini"]} + assert can_call_model(["beta"], "o1", model_access_groups=groups) is True + + +def test_unlisted_access_group_is_denied(): + groups = {"beta": ["o1"]} + # Model is in group "beta" but the key only lists group "gamma". + assert can_call_model(["gamma"], "o1", model_access_groups=groups) is False + + +def test_access_groups_accept_any_iterable_of_names(): + assert can_call_model(["beta"], "o1", model_access_groups={"beta", "alpha"}) is True + + +def test_no_access_groups_falls_back_to_name_and_pattern(): + # Without groups, only name/pattern matching applies (existing behavior). + assert can_call_model(["beta"], "o1") is False + assert can_call_model(["o1"], "o1", model_access_groups=None) is True