From 437a179612c939ee79a09b569be2ac94dce05233 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 21 Apr 2026 15:18:23 +0530 Subject: [PATCH] fix(router): constrain same-name deployment routing by access groups Filter router candidate deployments by caller-authorized model access groups when access is granted via group membership, preventing cross-group load balancing for shared public model names. Made-with: Cursor --- litellm/router.py | 77 ++++++++++++++++++++ tests/test_litellm/test_router.py | 115 ++++++++++++++++++++++++++++++ 2 files changed, 192 insertions(+) diff --git a/litellm/router.py b/litellm/router.py index f6976109bc7..2e436973e7f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -9241,6 +9241,12 @@ class Router: healthy_deployments = self._get_all_deployments( model_name=model, team_id=request_team_id ) + healthy_deployments = self._filter_deployments_by_model_access_groups( + model=model, + healthy_deployments=healthy_deployments, + request_kwargs=request_kwargs, + request_team_id=request_team_id, + ) if len(healthy_deployments) == 0: # check if the user sent in a deployment name instead @@ -9264,6 +9270,14 @@ class Router: healthy_deployments = self._get_all_deployments( model_name=model, team_id=request_team_id ) + healthy_deployments = ( + self._filter_deployments_by_model_access_groups( + model=model, + healthy_deployments=healthy_deployments, + request_kwargs=request_kwargs, + request_team_id=request_team_id, + ) + ) # If still no deployments after checking for fallbacks, raise an error if len(healthy_deployments) == 0: @@ -9289,6 +9303,68 @@ class Router: return model, healthy_deployments + def _filter_deployments_by_model_access_groups( + self, + model: str, + healthy_deployments: List, + request_kwargs: Optional[Dict], + request_team_id: Optional[str], + ) -> List: + """ + Restrict candidate deployments to caller-authorized model access groups. + + This is only applied when: + - request metadata includes `user_api_key_auth`, and + - caller permissions for this model are access-group-only + (no explicit model, wildcard, or all-proxy grants). + """ + if not healthy_deployments or request_kwargs is None: + return healthy_deployments + + metadata = request_kwargs.get("metadata") or {} + litellm_metadata = request_kwargs.get("litellm_metadata") or {} + user_api_key_auth = metadata.get("user_api_key_auth") or litellm_metadata.get( + "user_api_key_auth" + ) + if user_api_key_auth is None: + return healthy_deployments + + object_models = set(getattr(user_api_key_auth, "models", []) or []) + object_team_models = set(getattr(user_api_key_auth, "team_models", []) or []) + allowed_models = object_models | object_team_models + if not allowed_models: + return healthy_deployments + + # If caller has direct model/wildcard/all-proxy access, do not constrain + # deployment choice by access group. + if ( + model in allowed_models + or "*" in allowed_models + or "all-proxy-models" in allowed_models + ): + return healthy_deployments + + access_groups_for_model = self.get_model_access_groups( + model_name=model, team_id=request_team_id + ) + if len(access_groups_for_model) == 0: + return healthy_deployments + + allowed_access_groups = set(access_groups_for_model.keys()) & allowed_models + if not allowed_access_groups: + return healthy_deployments + + filtered_deployments = [] + for deployment in healthy_deployments: + deployment_model_info = deployment.get("model_info") or {} + deployment_access_groups = set( + deployment_model_info.get("access_groups", []) or [] + ) + if deployment_access_groups & allowed_access_groups: + filtered_deployments.append(deployment) + + return filtered_deployments + async def async_get_healthy_deployments( self, model: str, @@ -9796,6 +9872,7 @@ class Router: messages=messages, input=input, specific_deployment=specific_deployment, + request_kwargs=request_kwargs, ) if isinstance(healthy_deployments, dict): diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 2ae54f55103..af8633361ac 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -3204,3 +3204,118 @@ async def test_multiregion_team_failover_between_regions(): "response from us-east-1", "response from us-west-2", ] + + +def test_access_group_scoped_key_filters_deployments_with_same_public_model(): + """ + If a key can access a model only via access group membership, + router candidate deployments for that public model should be constrained + to deployments in the allowed access group. + """ + from litellm.proxy._types import UserAPIKeyAuth + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-5", + "litellm_params": { + "model": "openai/gpt-5.1", + "api_key": "key1", + "mock_response": "response-via-AG1", + }, + "model_info": {"access_groups": ["AG1"]}, + }, + { + "model_name": "gpt-5", + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "key2", + "mock_response": "response-via-AG2", + }, + "model_info": {"access_groups": ["AG2"]}, + }, + ] + ) + + scoped_key = UserAPIKeyAuth( + api_key="hashed-key", + team_id="team2", + models=["AG2"], + team_models=["AG2"], + ) + + _model, deployments = router._common_checks_available_deployment( + model="gpt-5", + request_kwargs={ + "metadata": { + "user_api_key_team_id": "team2", + "user_api_key_auth": scoped_key, + } + }, + ) + + assert len(deployments) == 1 + assert deployments[0].get("model_info", {}).get("access_groups") == ["AG2"] + + seen = set() + for _ in range(20): + response = router.completion( + model="gpt-5", + messages=[{"role": "user", "content": "hello"}], + metadata={"user_api_key_team_id": "team2", "user_api_key_auth": scoped_key}, + ) + seen.add(response.choices[0].message.content) + + assert seen == {"response-via-AG2"} + + +def test_explicit_model_access_does_not_force_access_group_filtering(): + """ + If a key has explicit model access in addition to access group entries, + do not force access-group-only filtering for deployment selection. + """ + from litellm.proxy._types import UserAPIKeyAuth + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-5", + "litellm_params": { + "model": "openai/gpt-5.1", + "api_key": "key1", + "mock_response": "response-via-AG1", + }, + "model_info": {"access_groups": ["AG1"]}, + }, + { + "model_name": "gpt-5", + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "key2", + "mock_response": "response-via-AG2", + }, + "model_info": {"access_groups": ["AG2"]}, + }, + ] + ) + + explicit_key = UserAPIKeyAuth( + api_key="hashed-key", + team_id="team2", + models=["AG2", "gpt-5"], + team_models=["AG2", "gpt-5"], + ) + + _model, deployments = router._common_checks_available_deployment( + model="gpt-5", + request_kwargs={ + "metadata": { + "user_api_key_team_id": "team2", + "user_api_key_auth": explicit_key, + } + }, + ) + + deployment_groups = [d.get("model_info", {}).get("access_groups") for d in deployments] + assert ["AG1"] in deployment_groups + assert ["AG2"] in deployment_groups