diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 4d44006369b..52dd142ea7e 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -547,6 +547,27 @@ async def _common_key_generation_helper( # noqa: PLR0915 }, ) + # Validate key duration against the team's max key duration + if ( + team_table is not None + and team_table.metadata is not None + and team_table.metadata.get("team_member_key_duration") + and data.duration is not None + ): + team_max_duration = team_table.metadata["team_member_key_duration"] + team_max_seconds = duration_in_seconds(duration=team_max_duration) + if data.duration == "-1": + user_key_duration: float = float("inf") + else: + user_key_duration = duration_in_seconds(duration=data.duration) + if user_key_duration > team_max_seconds: + raise HTTPException( + status_code=400, + detail={ + "error": f"Key duration exceeds team maximum. Requested: {data.duration}; Team maximum: {team_max_duration}" + }, + ) + # APPLY ENTERPRISE KEY MANAGEMENT PARAMS try: from litellm_enterprise.proxy.management_endpoints.key_management_endpoints import ( @@ -3372,6 +3393,55 @@ async def _insert_deprecated_key( "Failed to insert deprecated key for grace period: %s", deprecated_err, ) +async def _validate_regenerate_key_duration_against_team( + data: Optional[RegenerateKeyRequest], + key_in_db: LiteLLM_VerificationToken, + prisma_client: PrismaClient, + user_api_key_cache: DualCache, +) -> None: + """Raise HTTP 400 if the requested duration exceeds the team's max key duration.""" + if data is None or not data.duration: + return + + team_id = getattr(key_in_db, "team_id", None) + if not team_id: + return + + try: + team_table = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + check_db_only=True, + ) + except Exception: + return + + if ( + team_table is None + or team_table.metadata is None + or not team_table.metadata.get("team_member_key_duration") + ): + return + + team_max_duration = team_table.metadata["team_member_key_duration"] + team_max_seconds = duration_in_seconds(duration=team_max_duration) + + if data.duration == "-1": + user_seconds: float = float("inf") + else: + user_seconds = duration_in_seconds(duration=data.duration) + + if user_seconds > team_max_seconds: + raise HTTPException( + status_code=400, + detail={ + "error": f"Key duration exceeds team maximum. Requested: {data.duration}; Team maximum: {team_max_duration}" + }, + ) + + async def _execute_virtual_key_regeneration( *, prisma_client: PrismaClient, @@ -3387,6 +3457,13 @@ async def _execute_virtual_key_regeneration( """Generate new token, update DB, invalidate cache, and return response.""" from litellm.proxy.proxy_server import hash_token + await _validate_regenerate_key_duration_against_team( + data=data, + key_in_db=key_in_db, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + new_token = get_new_token(data=data) new_token_hash = hash_token(new_token) new_token_key_name = f"sk-...{new_token[-4:]}" 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 11c80351839..278dc522ab1 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 @@ -6436,7 +6436,7 @@ class TestValidateKeyAliasFormat: def test_validate_key_alias_format_invalid(self): from litellm.proxy.management_endpoints.key_management_endpoints import _validate_key_alias_format from litellm.proxy._types import ProxyException - + invalid_aliases = [ "", # empty " ", # whitespace @@ -6449,9 +6449,273 @@ class TestValidateKeyAliasFormat: " leading", "trailing ", ] - + for alias in invalid_aliases: with pytest.raises(ProxyException) as exc: _validate_key_alias_format(alias) assert str(exc.value.code) == "400" assert "Invalid key_alias format" in str(exc.value.message) + + +class TestCommonKeyGenerationHelperTeamDurationValidation: + """Tests for team duration validation inside _common_key_generation_helper.""" + + @pytest.mark.asyncio + async def test_duration_exceeds_team_max_raises(self): + """Raises HTTP 400 when the requested duration exceeds the team's max.""" + team = MagicMock(spec=LiteLLM_TeamTableCachedObj) + team.metadata = {"team_member_key_duration": "5d"} + + with pytest.raises(HTTPException) as exc_info: + await _common_key_generation_helper( + data=GenerateKeyRequest(user_id="user-1", duration="10d"), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1" + ), + litellm_changed_by=None, + team_table=team, + ) + assert exc_info.value.status_code == 400 + assert "Team maximum" in str(exc_info.value.detail) + + @pytest.mark.asyncio + async def test_never_expires_exceeds_team_max_raises(self): + """Duration '-1' (never expires) exceeds any finite team max.""" + team = MagicMock(spec=LiteLLM_TeamTableCachedObj) + team.metadata = {"team_member_key_duration": "30d"} + + with pytest.raises(HTTPException) as exc_info: + await _common_key_generation_helper( + data=GenerateKeyRequest(user_id="user-1", duration="-1"), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1" + ), + litellm_changed_by=None, + team_table=team, + ) + assert exc_info.value.status_code == 400 + assert "Team maximum" in str(exc_info.value.detail) + + @pytest.mark.asyncio + async def test_duration_within_team_max_passes(self): + """Duration within the team max proceeds past the validation block.""" + team = MagicMock(spec=LiteLLM_TeamTableCachedObj) + team.metadata = {"team_member_key_duration": "5d"} + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + new_callable=AsyncMock, + return_value={"key": "sk-test", "expires": None, "user_id": "user-1"}, + ), patch("litellm.proxy.proxy_server.prisma_client"), patch( + "litellm.proxy.proxy_server.llm_router" + ), patch( + "litellm.proxy.proxy_server.premium_user", False + ): + await _common_key_generation_helper( + data=GenerateKeyRequest(user_id="user-1", duration="3d"), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1" + ), + litellm_changed_by=None, + team_table=team, + ) + + @pytest.mark.asyncio + async def test_no_duration_skips_team_validation(self): + """No duration provided — team validation is skipped.""" + team = MagicMock(spec=LiteLLM_TeamTableCachedObj) + team.metadata = {"team_member_key_duration": "5d"} + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + new_callable=AsyncMock, + return_value={"key": "sk-test", "expires": None, "user_id": "user-1"}, + ), patch("litellm.proxy.proxy_server.prisma_client"), patch( + "litellm.proxy.proxy_server.llm_router" + ), patch( + "litellm.proxy.proxy_server.premium_user", False + ): + await _common_key_generation_helper( + data=GenerateKeyRequest(user_id="user-1", duration=None), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1" + ), + litellm_changed_by=None, + team_table=team, + ) + + @pytest.mark.asyncio + async def test_no_team_skips_team_validation(self): + """No team_table — team duration validation is skipped.""" + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + new_callable=AsyncMock, + return_value={"key": "sk-test", "expires": None, "user_id": "user-1"}, + ), patch("litellm.proxy.proxy_server.prisma_client"), patch( + "litellm.proxy.proxy_server.llm_router" + ), patch( + "litellm.proxy.proxy_server.premium_user", False + ): + await _common_key_generation_helper( + data=GenerateKeyRequest(user_id="user-1", duration="100d"), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1" + ), + litellm_changed_by=None, + team_table=None, + ) + + +class TestValidateRegenerateKeyDurationAgainstTeam: + """Tests for _validate_regenerate_key_duration_against_team.""" + + @pytest.mark.asyncio + async def test_no_data_skips_validation(self): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_regenerate_key_duration_against_team, + ) + + mock_key = MagicMock(spec=LiteLLM_VerificationToken) + mock_key.team_id = "team-123" + mock_prisma = AsyncMock() + mock_cache = MagicMock() + # Should not raise + await _validate_regenerate_key_duration_against_team( + data=None, + key_in_db=mock_key, + prisma_client=mock_prisma, + user_api_key_cache=mock_cache, + ) + + @pytest.mark.asyncio + async def test_no_duration_skips_validation(self): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_regenerate_key_duration_against_team, + ) + from litellm.proxy._types import RegenerateKeyRequest + + mock_key = MagicMock(spec=LiteLLM_VerificationToken) + mock_key.team_id = "team-123" + data = RegenerateKeyRequest(duration=None) + mock_prisma = AsyncMock() + mock_cache = MagicMock() + # Should not raise + await _validate_regenerate_key_duration_against_team( + data=data, + key_in_db=mock_key, + prisma_client=mock_prisma, + user_api_key_cache=mock_cache, + ) + + @pytest.mark.asyncio + async def test_no_team_id_skips_validation(self): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_regenerate_key_duration_against_team, + ) + from litellm.proxy._types import RegenerateKeyRequest + + mock_key = MagicMock(spec=LiteLLM_VerificationToken) + mock_key.team_id = None + data = RegenerateKeyRequest(duration="10d") + mock_prisma = AsyncMock() + mock_cache = MagicMock() + # Should not raise + await _validate_regenerate_key_duration_against_team( + data=data, + key_in_db=mock_key, + prisma_client=mock_prisma, + user_api_key_cache=mock_cache, + ) + + @pytest.mark.asyncio + async def test_duration_within_team_max_passes(self): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_regenerate_key_duration_against_team, + ) + from litellm.proxy._types import RegenerateKeyRequest, LiteLLM_TeamTableCachedObj + + mock_key = MagicMock(spec=LiteLLM_VerificationToken) + mock_key.team_id = "team-123" + data = RegenerateKeyRequest(duration="3d") + + mock_team = MagicMock(spec=LiteLLM_TeamTableCachedObj) + mock_team.metadata = {"team_member_key_duration": "5d"} + + mock_prisma = AsyncMock() + mock_cache = MagicMock() + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + new_callable=AsyncMock, + return_value=mock_team, + ): + # Should not raise + await _validate_regenerate_key_duration_against_team( + data=data, + key_in_db=mock_key, + prisma_client=mock_prisma, + user_api_key_cache=mock_cache, + ) + + @pytest.mark.asyncio + async def test_duration_exceeds_team_max_raises(self): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_regenerate_key_duration_against_team, + ) + from litellm.proxy._types import RegenerateKeyRequest, LiteLLM_TeamTableCachedObj + + mock_key = MagicMock(spec=LiteLLM_VerificationToken) + mock_key.team_id = "team-123" + data = RegenerateKeyRequest(duration="10d") + + mock_team = MagicMock(spec=LiteLLM_TeamTableCachedObj) + mock_team.metadata = {"team_member_key_duration": "5d"} + + mock_prisma = AsyncMock() + mock_cache = MagicMock() + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + new_callable=AsyncMock, + return_value=mock_team, + ): + with pytest.raises(HTTPException) as exc_info: + await _validate_regenerate_key_duration_against_team( + data=data, + key_in_db=mock_key, + prisma_client=mock_prisma, + user_api_key_cache=mock_cache, + ) + assert exc_info.value.status_code == 400 + assert "Team maximum" in str(exc_info.value.detail) + + @pytest.mark.asyncio + async def test_never_expires_duration_exceeds_team_max_raises(self): + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _validate_regenerate_key_duration_against_team, + ) + from litellm.proxy._types import RegenerateKeyRequest, LiteLLM_TeamTableCachedObj + + mock_key = MagicMock(spec=LiteLLM_VerificationToken) + mock_key.team_id = "team-123" + data = RegenerateKeyRequest(duration="-1") + + mock_team = MagicMock(spec=LiteLLM_TeamTableCachedObj) + mock_team.metadata = {"team_member_key_duration": "30d"} + + mock_prisma = AsyncMock() + mock_cache = MagicMock() + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + new_callable=AsyncMock, + return_value=mock_team, + ): + with pytest.raises(HTTPException) as exc_info: + await _validate_regenerate_key_duration_against_team( + data=data, + key_in_db=mock_key, + prisma_client=mock_prisma, + user_api_key_cache=mock_cache, + ) + assert exc_info.value.status_code == 400