From cdeca638b8b43e91c2c8ddd2ee4a99042910441d Mon Sep 17 00:00:00 2001 From: Atharva Jaiswal <92455570+AtharvaJaiswal005@users.noreply.github.com> Date: Tue, 24 Feb 2026 10:15:53 +0530 Subject: [PATCH] fix: team assignment fails for keys with special model names (#21919) * auth_with_role_name add region_name arg for cross-account sts * update tests to include case with aws_region_name for _auth_with_aws_role * Only pass region_name to STS client when aws_region_name is set * Add optional aws_sts_endpoint to _auth_with_aws_role * Parametrize ambient-credentials test for no opts, region_name, and aws_sts_endpoint * consistently passing region and endpoint args into explicit credentials irsa * fix env var leakage * fix: bedrock openai-compatible imported-model should also have model arn encoded * fix: team assignment fails for keys with special model names (#21880) --------- Co-authored-by: An Tang Co-authored-by: Sameer Kankute --- .../key_management_endpoints.py | 3 + .../test_key_management_endpoints.py | 100 ++++++++++++++++++ 2 files changed, 103 insertions(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a230c2e9336..6ddedfcda3c 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2165,8 +2165,11 @@ async def validate_key_team_change( - The person initiating the change must be either Proxy Admin or Team Admin """ # Check if the team has access to the key's models + _special_model_names = {e.value for e in SpecialModelNames} if len(key.models) > 0: for model in key.models: + if model in _special_model_names: + continue await can_team_access_model( model=model, team_object=team, diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index ffb4e955426..44d379a3882 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -1468,6 +1468,106 @@ async def test_validate_key_team_change_with_member_permissions(): ) +@pytest.mark.asyncio +async def test_validate_key_team_change_with_special_model_names(): + """ + Test that validate_key_team_change skips validation for special model names + like 'all-team-models', 'all-proxy-models', and 'no-default-models'. + + These are placeholder values, not real model names, so they should not be + checked against the team's actual model list. + + Related: https://github.com/BerriAI/litellm/issues/21880 + """ + mock_key = MagicMock() + mock_key.user_id = "test-user-123" + mock_key.models = ["all-team-models"] + mock_key.tpm_limit = None + mock_key.rpm_limit = None + + mock_team = MagicMock() + mock_team.team_id = "test-team-456" + mock_team.members_with_roles = [] + mock_team.tpm_limit = None + mock_team.rpm_limit = None + + mock_change_initiator = MagicMock() + mock_change_initiator.user_id = "test-user-123" + mock_change_initiator.user_role = LitellmUserRoles.PROXY_ADMIN.value + + mock_router = MagicMock() + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.can_team_access_model", + new_callable=AsyncMock, + ) as mock_can_access: + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._get_user_in_team" + ) as mock_get_user: + mock_get_user.return_value = MagicMock() + + # Should NOT raise - "all-team-models" should be skipped + await validate_key_team_change( + key=mock_key, + team=mock_team, + change_initiated_by=mock_change_initiator, + llm_router=mock_router, + ) + + # can_team_access_model should NOT have been called + mock_can_access.assert_not_called() + + +@pytest.mark.asyncio +async def test_validate_key_team_change_special_models_mixed_with_real(): + """ + Test that when a key has both special and real model names, + only real model names are validated against the team. + + Related: https://github.com/BerriAI/litellm/issues/21880 + """ + mock_key = MagicMock() + mock_key.user_id = "test-user-123" + mock_key.models = ["all-team-models", "gpt-4"] + mock_key.tpm_limit = None + mock_key.rpm_limit = None + + mock_team = MagicMock() + mock_team.team_id = "test-team-456" + mock_team.members_with_roles = [] + mock_team.tpm_limit = None + mock_team.rpm_limit = None + + mock_change_initiator = MagicMock() + mock_change_initiator.user_id = "test-user-123" + mock_change_initiator.user_role = LitellmUserRoles.PROXY_ADMIN.value + + mock_router = MagicMock() + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.can_team_access_model", + new_callable=AsyncMock, + ) as mock_can_access: + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._get_user_in_team" + ) as mock_get_user: + mock_get_user.return_value = MagicMock() + + await validate_key_team_change( + key=mock_key, + team=mock_team, + change_initiated_by=mock_change_initiator, + llm_router=mock_router, + ) + + # Should only be called for "gpt-4", not "all-team-models" + mock_can_access.assert_called_once_with( + model="gpt-4", + team_object=mock_team, + llm_router=mock_router, + ) + + def test_key_rotation_fields_helper(): """ Test the key data update logic for rotation fields.