mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-29 01:42:19 +00:00
fix(auth): enforce and alert on team member budgets only in common_checks
The builder re-checked the team member budget inline before common_checks ran the same check, so one request that crossed a team_member_max_budget_alert_emails threshold dispatched two alerts. Drop the inline check; common_checks is the single authorization point and already covers per-member rows, the team default member budget, zero-cost skips and the cross-pod spend counter. Its 422 message now uses the TeamMember=user:team form the builder and budget reservation already returned.
This commit is contained in:
ryan-crabbe-berri
parent
703e754b46
commit
92ecef29b5
3 changed files with 59 additions and 115 deletions
|
|
@ -5646,12 +5646,16 @@ async def _check_team_member_budget(
|
|||
)
|
||||
|
||||
if team_member_spend >= team_member_budget:
|
||||
entity_id: Final = f"{valid_token.user_id}:{team_object.team_id}"
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=team_member_spend,
|
||||
max_budget=team_member_budget,
|
||||
message=f"Budget has been exceeded! User={valid_token.user_id} in Team={team_object.team_id} Current cost: {team_member_spend}, Max budget: {team_member_budget}",
|
||||
message=(
|
||||
f"Budget has been exceeded! TeamMember={entity_id} "
|
||||
f"Current cost: {team_member_spend}, Max budget: {team_member_budget}"
|
||||
),
|
||||
entity_type=Litellm_EntityType.TEAM_MEMBER.value,
|
||||
entity_id=f"{valid_token.user_id}:{team_object.team_id}",
|
||||
entity_id=entity_id,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -50,7 +50,6 @@ from litellm.proxy.auth.auth_checks import (
|
|||
_get_user_role,
|
||||
_is_model_cost_zero,
|
||||
_is_user_proxy_admin,
|
||||
_team_member_max_budget_alert_check,
|
||||
_virtual_key_max_budget_alert_check,
|
||||
_virtual_key_max_budget_check,
|
||||
_virtual_key_soft_budget_check,
|
||||
|
|
@ -119,7 +118,6 @@ from litellm.proxy.common_utils.model_listing_utils import claude_code_requested
|
|||
from litellm.proxy.common_utils.realtime_utils import _realtime_request_body
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
UserApiKeyCache,
|
||||
team_membership_auth_cache_key,
|
||||
)
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
||||
|
|
@ -134,7 +132,6 @@ from litellm.proxy.utils import (
|
|||
ProxyLogging,
|
||||
normalize_route_for_root_path,
|
||||
)
|
||||
from litellm.repositories.table_repositories import TeamMembershipRepository
|
||||
from litellm.router_utils.common_utils import resolve_model_group_alias
|
||||
from litellm.secret_managers.main import get_secret_bool
|
||||
from litellm.types.services import ServiceTypes
|
||||
|
|
@ -2239,74 +2236,6 @@ async def _user_api_key_auth_builder(
|
|||
if skip_budget_checks:
|
||||
verbose_proxy_logger.info("Skipping all budget checks for zero-cost model: %s", model)
|
||||
|
||||
# Check 3. Check if user is in their team budget
|
||||
if not skip_budget_checks and valid_token.team_member_spend is not None:
|
||||
_user_id: Final = valid_token.user_id
|
||||
_team_id: Final = valid_token.team_id
|
||||
if prisma_client is not None and _user_id is not None and _team_id is not None:
|
||||
_cache_key: Final = team_membership_auth_cache_key(team_id=_team_id, user_id=_user_id)
|
||||
|
||||
team_member_info = await user_api_key_cache.async_get_cache(
|
||||
key=_cache_key,
|
||||
model_type=LiteLLM_TeamMembership,
|
||||
)
|
||||
if team_member_info is None:
|
||||
# read from DB
|
||||
_db_member: Final = await TeamMembershipRepository(prisma_client).table.find_first(
|
||||
where={
|
||||
"user_id": _user_id,
|
||||
"team_id": _team_id,
|
||||
},
|
||||
include={"litellm_budget_table": True},
|
||||
)
|
||||
if _db_member is not None:
|
||||
team_member_info = LiteLLM_TeamMembership(**_db_member.model_dump())
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=_cache_key,
|
||||
value=team_member_info,
|
||||
model_type=LiteLLM_TeamMembership,
|
||||
ttl=5,
|
||||
)
|
||||
|
||||
if team_member_info is not None and team_member_info.litellm_budget_table is not None:
|
||||
team_member_budget: Final = team_member_info.litellm_budget_table.effective_max_budget(
|
||||
now=datetime.now(timezone.utc),
|
||||
)
|
||||
if team_member_budget is not None and team_member_budget > 0:
|
||||
# Read from cross-pod counter (Redis-first) if available
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
||||
team_member_spend = valid_token.team_member_spend
|
||||
if valid_token.user_id is not None and valid_token.team_id is not None:
|
||||
team_member_spend = await get_current_spend(
|
||||
counter_key=f"spend:team_member:{valid_token.user_id}:{valid_token.team_id}",
|
||||
fallback_spend=team_member_spend,
|
||||
max_budget=team_member_budget,
|
||||
)
|
||||
_team_member_max_budget_alert_check(
|
||||
team_id=_team_id,
|
||||
team_alias=valid_token.team_alias,
|
||||
team_metadata=valid_token.team_metadata,
|
||||
organization_id=valid_token.org_id,
|
||||
user_id=_user_id,
|
||||
user_email=user_obj.user_email if user_obj is not None else None,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
spend=team_member_spend,
|
||||
max_budget=team_member_budget,
|
||||
)
|
||||
if team_member_spend >= team_member_budget:
|
||||
_entity_id: Final = f"{valid_token.user_id}:{valid_token.team_id}"
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=team_member_spend,
|
||||
max_budget=team_member_budget,
|
||||
message=(
|
||||
f"Budget has been exceeded! TeamMember={_entity_id} "
|
||||
f"Current cost: {team_member_spend}, Max budget: {team_member_budget}"
|
||||
),
|
||||
entity_type=Litellm_EntityType.TEAM_MEMBER.value,
|
||||
entity_id=_entity_id,
|
||||
)
|
||||
|
||||
# Check 3. If token is expired
|
||||
if valid_token.expires is not None:
|
||||
current_time = datetime.now(timezone.utc)
|
||||
|
|
|
|||
|
|
@ -7855,6 +7855,30 @@ async def test_temp_budget_increase_applied_for_cached_key():
|
|||
assert cached_after.max_budget == 2.0
|
||||
|
||||
|
||||
async def _authenticate_and_authorize(mock_request, api_key):
|
||||
"""Builder then the single common_checks gate, the same sequence user_api_key_auth runs."""
|
||||
from litellm.proxy.auth.user_api_key_auth import _authorize_authenticated_request
|
||||
|
||||
request_data = {"model": "claude-sonnet-5", "messages": [{"role": "user", "content": "hi"}]}
|
||||
auth_obj = await _user_api_key_auth_builder(
|
||||
request=mock_request,
|
||||
api_key=f"Bearer {api_key}",
|
||||
azure_api_key_header="",
|
||||
anthropic_api_key_header=None,
|
||||
google_ai_studio_api_key_header=None,
|
||||
azure_apim_header=None,
|
||||
request_data=request_data,
|
||||
)
|
||||
recovered = await _authorize_authenticated_request(
|
||||
user_api_key_auth_obj=auth_obj,
|
||||
request=mock_request,
|
||||
request_data=request_data,
|
||||
route="/v1/messages",
|
||||
api_key=f"Bearer {api_key}",
|
||||
)
|
||||
return recovered or auth_obj
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"team_member_spend, expect_blocked",
|
||||
|
|
@ -7868,7 +7892,7 @@ async def test_cached_key_team_member_budget_blocks_at_exact_cap(team_member_spe
|
|||
"""A team member counter sitting exactly at the cap (where a resized reservation
|
||||
lands it) must be rejected by the cached-key auth path like every other budget check."""
|
||||
from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj
|
||||
from litellm.proxy.common_utils.user_api_key_cache import team_membership_auth_cache_key
|
||||
from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key
|
||||
from litellm.proxy.utils import hash_token
|
||||
|
||||
api_key = "sk-team-member-exact-cap"
|
||||
|
|
@ -7898,7 +7922,7 @@ async def test_cached_key_team_member_budget_blocks_at_exact_cap(team_member_spe
|
|||
value=LiteLLM_UserTable(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER),
|
||||
)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=team_membership_auth_cache_key(team_id=team_id, user_id=user_id),
|
||||
key=team_membership_reservation_cache_key(team_id=team_id, user_id=user_id),
|
||||
value=LiteLLM_TeamMembership(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
|
|
@ -7920,15 +7944,7 @@ async def test_cached_key_team_member_budget_blocks_at_exact_cap(team_member_spe
|
|||
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
|
||||
|
||||
async def _auth():
|
||||
return await _user_api_key_auth_builder(
|
||||
request=mock_request,
|
||||
api_key=f"Bearer {api_key}",
|
||||
azure_api_key_header="",
|
||||
anthropic_api_key_header=None,
|
||||
google_ai_studio_api_key_header=None,
|
||||
azure_apim_header=None,
|
||||
request_data={"model": "claude-sonnet-5", "messages": [{"role": "user", "content": "hi"}]},
|
||||
)
|
||||
return await _authenticate_and_authorize(mock_request, api_key)
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: the builder reads proxy settings from module globals, no injection seam
|
||||
|
|
@ -7973,7 +7989,10 @@ async def test_cached_key_team_member_budget_emails_configured_thresholds(
|
|||
"""The team's team_member_max_budget_alert_emails thresholds fire from the cached-key auth path,
|
||||
including on the request that trips the hard cap, and stay silent below the lowest threshold."""
|
||||
from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj
|
||||
from litellm.proxy.common_utils.user_api_key_cache import team_membership_auth_cache_key
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
team_membership_auth_cache_key,
|
||||
team_membership_reservation_cache_key,
|
||||
)
|
||||
from litellm.proxy.utils import hash_token
|
||||
|
||||
api_key = "sk-team-member-alert-thresholds"
|
||||
|
|
@ -7998,7 +8017,11 @@ async def test_cached_key_team_member_budget_emails_configured_thresholds(
|
|||
)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=f"team_id:{team_id}",
|
||||
value=LiteLLM_TeamTableCachedObj(team_id=team_id),
|
||||
value=LiteLLM_TeamTableCachedObj(
|
||||
team_id=team_id,
|
||||
team_alias="platform",
|
||||
metadata={"team_member_max_budget_alert_emails": alert_emails},
|
||||
),
|
||||
)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=user_id,
|
||||
|
|
@ -8006,16 +8029,20 @@ async def test_cached_key_team_member_budget_emails_configured_thresholds(
|
|||
user_id=user_id, user_email="member@example.com", user_role=LitellmUserRoles.INTERNAL_USER
|
||||
),
|
||||
)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=team_membership_auth_cache_key(team_id=team_id, user_id=user_id),
|
||||
value=LiteLLM_TeamMembership(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
spend=team_member_spend,
|
||||
budget_id="budget-alert-thresholds",
|
||||
litellm_budget_table=LiteLLM_BudgetTable(max_budget=2.4),
|
||||
),
|
||||
membership = LiteLLM_TeamMembership(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
spend=team_member_spend,
|
||||
budget_id="budget-alert-thresholds",
|
||||
litellm_budget_table=LiteLLM_BudgetTable(max_budget=2.4),
|
||||
)
|
||||
# A live proxy holds the row under both keys, so any second team-member check in the
|
||||
# auth flow would find it too and send a duplicate alert.
|
||||
for membership_cache_key in (
|
||||
team_membership_reservation_cache_key(team_id=team_id, user_id=user_id),
|
||||
team_membership_auth_cache_key(team_id=team_id, user_id=user_id),
|
||||
):
|
||||
await user_api_key_cache.async_set_cache(key=membership_cache_key, value=membership)
|
||||
|
||||
mock_request = MagicMock()
|
||||
mock_request.url.path = "/v1/messages"
|
||||
|
|
@ -8030,15 +8057,7 @@ async def test_cached_key_team_member_budget_emails_configured_thresholds(
|
|||
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock(return_value=None)
|
||||
|
||||
async def _auth():
|
||||
return await _user_api_key_auth_builder(
|
||||
request=mock_request,
|
||||
api_key=f"Bearer {api_key}",
|
||||
azure_api_key_header="",
|
||||
anthropic_api_key_header=None,
|
||||
google_ai_studio_api_key_header=None,
|
||||
azure_apim_header=None,
|
||||
request_data={"model": "claude-sonnet-5", "messages": [{"role": "user", "content": "hi"}]},
|
||||
)
|
||||
return await _authenticate_and_authorize(mock_request, api_key)
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: the builder reads proxy settings from module globals, no injection seam
|
||||
|
|
@ -8089,7 +8108,7 @@ async def test_cached_key_team_member_budget_honours_temp_increase(expiry_offset
|
|||
"""A member over their permanent cap is admitted while a temp_budget_increase is unexpired
|
||||
and blocked again once it expires, on the cached-key auth path."""
|
||||
from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj
|
||||
from litellm.proxy.common_utils.user_api_key_cache import team_membership_auth_cache_key
|
||||
from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key
|
||||
from litellm.proxy.utils import hash_token
|
||||
|
||||
api_key = "sk-team-member-temp-budget"
|
||||
|
|
@ -8119,7 +8138,7 @@ async def test_cached_key_team_member_budget_honours_temp_increase(expiry_offset
|
|||
value=LiteLLM_UserTable(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER),
|
||||
)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=team_membership_auth_cache_key(team_id=team_id, user_id=user_id),
|
||||
key=team_membership_reservation_cache_key(team_id=team_id, user_id=user_id),
|
||||
value=LiteLLM_TeamMembership(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
|
|
@ -8145,15 +8164,7 @@ async def test_cached_key_team_member_budget_honours_temp_increase(expiry_offset
|
|||
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
|
||||
|
||||
async def _auth():
|
||||
return await _user_api_key_auth_builder(
|
||||
request=mock_request,
|
||||
api_key=f"Bearer {api_key}",
|
||||
azure_api_key_header="",
|
||||
anthropic_api_key_header=None,
|
||||
google_ai_studio_api_key_header=None,
|
||||
azure_apim_header=None,
|
||||
request_data={"model": "claude-sonnet-5", "messages": [{"role": "user", "content": "hi"}]},
|
||||
)
|
||||
return await _authenticate_and_authorize(mock_request, api_key)
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: the builder reads proxy settings from module globals, no injection seam
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue