Fix/organization max budget not enforced (#17334)

* test: add failing tests for organization budget enforcement bug

Add comprehensive tests exposing that organization-level budgets are
retrieved but never enforced during request authentication. Tests verify:

1. Basic org budget exceeded scenario (team under budget, org over)
2. Multiple teams collectively exceeding org budget
3. Organization budget fields exist but are never checked
4. Inconsistency between team budget enforcement (works) and org (doesn't)

Tests intentionally fail to document the bug. Will be fixed in next commit.

Related to organization_max_budget not being enforced in auth_checks.py

* fix: enforce organization budget in auth checks

Add organization budget enforcement to common_checks() in auth_checks.py.
Previously, organization_max_budget was retrieved from DB but never checked,
allowing teams to collectively exceed their organization's budget limit.

Changes:
- Add _organization_max_budget_check() function following team budget pattern
- Call org budget check after team budget check in common_checks()
- Add "organization_budget" to budget_alerts type literals
- Update tests to verify org budget is enforced

Budget hierarchy is now properly enforced:
  Organization Budget (hard ceiling)
    └─ Team Budget (sub-allocation)
        └─ Team Member Budget (per-user within team)
            └─ Key Budget (per-key)

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

* fix: add organization_id to budget alerts, fix enum comparison and linting of newly added code

- Add organization_id field to CallInfo class for better alert context
- Include organization_id in budget alerts (token, soft, team, org)
- Fix event_group enum comparison (was comparing enum to string)
- Add OrganizationBudgetAlert class for organization budget alerting
- Add organization_budget to test parameterizations
- Apply Black formatting to slack_alerting.py

---------

Co-authored-by: Claude <noreply@anthropic.com>

litellm/integrations/SlackAlerting/budget_alert_types.py

@@ -50,6 +50,14 @@ class TeamBudgetAlert(BaseBudgetAlertType):
         return user_info.team_id or "default_id"
 
 
+class OrganizationBudgetAlert(BaseBudgetAlertType):
+    def get_event_message(self) -> str:
+        return "Organization Budget: "
+
+    def get_id(self, user_info: CallInfo) -> str:
+        return user_info.organization_id or "default_id"
+
+
 class TokenBudgetAlert(BaseBudgetAlertType):
     def get_event_message(self) -> str:
         return "Key Budget: "
@@ -72,6 +80,7 @@ def get_budget_alert_type(
         "soft_budget",
         "user_budget",
         "team_budget",
+        "organization_budget",
         "proxy_budget",
         "projected_limit_exceeded",
     ],
@@ -83,6 +92,7 @@ def get_budget_alert_type(
         "soft_budget": SoftBudgetAlert(),
         "user_budget": UserBudgetAlert(),
         "team_budget": TeamBudgetAlert(),
+        "organization_budget": OrganizationBudgetAlert(),
         "token_budget": TokenBudgetAlert(),
         "projected_limit_exceeded": ProjectedLimitExceededAlert(),
     }

litellm/integrations/SlackAlerting/slack_alerting.py

@@ -134,19 +134,25 @@ class SlackAlerting(CustomBatchLogger):
         if llm_router is not None:
             self.llm_router = llm_router
 
-    def _prepare_outage_value_for_cache(self, outage_value: Union[dict, ProviderRegionOutageModel, OutageModel]) -> dict:
+    def _prepare_outage_value_for_cache(
+        self, outage_value: Union[dict, ProviderRegionOutageModel, OutageModel]
+    ) -> dict:
         """
         Helper method to prepare outage value for Redis caching.
         Converts set objects to lists for JSON serialization.
         """
         # Convert to dict for processing
         cache_value = dict(outage_value)
-        
-        if "deployment_ids" in cache_value and isinstance(cache_value["deployment_ids"], set):
+
+        if "deployment_ids" in cache_value and isinstance(
+            cache_value["deployment_ids"], set
+        ):
             cache_value["deployment_ids"] = list(cache_value["deployment_ids"])
         return cache_value
 
-    def _restore_outage_value_from_cache(self, outage_value: Optional[dict]) -> Optional[dict]:
+    def _restore_outage_value_from_cache(
+        self, outage_value: Optional[dict]
+    ) -> Optional[dict]:
         """
         Helper method to restore outage value after retrieving from cache.
         Converts list objects back to sets for proper handling.
@@ -528,6 +534,7 @@ class SlackAlerting(CustomBatchLogger):
             "soft_budget",
             "user_budget",
             "team_budget",
+            "organization_budget",
             "proxy_budget",
             "projected_limit_exceeded",
         ],
@@ -1338,7 +1345,7 @@ Model Info:
             subject=email_event["subject"],
             html=email_event["html"],
         )
-        if webhook_event.event_group == "team":
+        if webhook_event.event_group == Litellm_EntityType.TEAM:
             from litellm.integrations.email_alerting import send_team_budget_alert
 
             await send_team_budget_alert(webhook_event=webhook_event)
@@ -1399,7 +1406,7 @@ Model Info:
         current_time = datetime.now().strftime("%H:%M:%S")
         _proxy_base_url = os.getenv("PROXY_BASE_URL", None)
         # Use .name if it's an enum, otherwise use as is
-        alert_type_name = getattr(alert_type, 'name', alert_type)
+        alert_type_name = getattr(alert_type, "name", alert_type)
         alert_type_formatted = f"Alert type: `{alert_type_name}`"
         if alert_type == "daily_reports" or alert_type == "new_model_added":
             formatted_message = alert_type_formatted + message

litellm/proxy/_types.py

@@ -2444,6 +2444,7 @@ class CallInfo(LiteLLMPydanticObjectBase):
     user_id: Optional[str] = None
     team_id: Optional[str] = None
     team_alias: Optional[str] = None
+    organization_id: Optional[str] = None
     user_email: Optional[str] = None
     key_alias: Optional[str] = None
     projected_exceeded_date: Optional[str] = None

litellm/proxy/auth/auth_checks.py

@@ -143,6 +143,14 @@ async def common_checks(
         valid_token=valid_token,
     )
 
+    # 3.1. If organization is in budget
+    await _organization_max_budget_check(
+        valid_token=valid_token,
+        prisma_client=prisma_client,
+        user_api_key_cache=user_api_key_cache,
+        proxy_logging_obj=proxy_logging_obj,
+    )
+
     await _tag_max_budget_check(
         request_body=request_body,
         prisma_client=prisma_client,
@@ -1893,6 +1901,7 @@ async def _virtual_key_max_budget_check(
             max_budget=valid_token.max_budget,
             user_id=valid_token.user_id,
             team_id=valid_token.team_id,
+            organization_id=valid_token.org_id,
             user_email=user_email,
             key_alias=valid_token.key_alias,
             event_group=Litellm_EntityType.KEY,
@@ -1939,6 +1948,7 @@ async def _virtual_key_soft_budget_check(
             user_id=valid_token.user_id,
             team_id=valid_token.team_id,
             team_alias=valid_token.team_alias,
+            organization_id=valid_token.org_id,
             user_email=None,
             key_alias=valid_token.key_alias,
             event_group=Litellm_EntityType.KEY,
@@ -1977,6 +1987,7 @@ async def _team_max_budget_check(
                 user_id=valid_token.user_id,
                 team_id=valid_token.team_id,
                 team_alias=valid_token.team_alias,
+                organization_id=valid_token.org_id,
                 event_group=Litellm_EntityType.TEAM,
             )
             asyncio.create_task(
@@ -1993,6 +2004,65 @@ async def _team_max_budget_check(
         )
 
 
+async def _organization_max_budget_check(
+    valid_token: Optional[UserAPIKeyAuth],
+    prisma_client: Optional[PrismaClient],
+    user_api_key_cache: DualCache,
+    proxy_logging_obj: ProxyLogging,
+):
+    """
+    Check if the organization is over its max budget.
+
+    Raises:
+        BudgetExceededError if the organization is over its max budget.
+        Triggers a budget alert if the organization is over its max budget.
+    """
+    # Only check if token has organization info and organization_max_budget is set
+    if (
+        valid_token is None
+        or valid_token.org_id is None
+        or valid_token.organization_max_budget is None
+        or valid_token.organization_max_budget <= 0
+    ):
+        return
+
+    # Get organization object to check current spend
+    if prisma_client is not None:
+        org_table = await get_org_object(
+            org_id=valid_token.org_id,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+        )
+
+        if (
+            org_table is not None
+            and org_table.spend >= valid_token.organization_max_budget
+        ):
+            # Trigger budget alert
+            call_info = CallInfo(
+                token=valid_token.token,
+                spend=org_table.spend,
+                max_budget=valid_token.organization_max_budget,
+                user_id=valid_token.user_id,
+                team_id=valid_token.team_id,
+                team_alias=valid_token.team_alias,
+                organization_id=valid_token.org_id,
+                event_group=Litellm_EntityType.ORGANIZATION,
+            )
+            asyncio.create_task(
+                proxy_logging_obj.budget_alerts(
+                    type="organization_budget",
+                    user_info=call_info,
+                )
+            )
+
+            raise litellm.BudgetExceededError(
+                current_cost=org_table.spend,
+                max_budget=valid_token.organization_max_budget,
+                message=f"Budget has been exceeded! Organization={valid_token.org_id} Current cost: {org_table.spend}, Max budget: {valid_token.organization_max_budget}",
+            )
+
+
 async def _tag_max_budget_check(
     request_body: dict,
     prisma_client: Optional[PrismaClient],

litellm/proxy/utils.py

@@ -1072,6 +1072,7 @@ class ProxyLogging:
             "user_budget",
             "soft_budget",
             "team_budget",
+            "organization_budget",
             "proxy_budget",
             "projected_limit_exceeded",
         ],

tests/logging_callback_tests/test_alerting.py

@@ -477,6 +477,7 @@ async def test_send_daily_reports_all_zero_or_none():
         "token_budget",
         "user_budget",
         "team_budget",
+        "organization_budget",
         "proxy_budget",
         "projected_limit_exceeded",
     ],
@@ -514,6 +515,7 @@ async def test_send_token_budget_crossed_alerts(alerting_type):
         "token_budget",
         "user_budget",
         "team_budget",
+        "organization_budget",
         "proxy_budget",
         "projected_limit_exceeded",
     ],

tests/test_litellm/proxy/auth/test_organization_budget_enforcement.py

@@ -0,0 +1,344 @@
+"""
+Tests for organization budget enforcement.
+
+These tests verify that organization-level budgets are properly enforced during
+request authentication. When an organization's spend exceeds its max_budget,
+requests should fail with BudgetExceededError.
+
+This prevents teams within an organization from collectively exceeding the
+organization's budget limit.
+"""
+
+import asyncio
+import os
+import sys
+from typing import Optional
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+sys.path.insert(0, os.path.abspath("../../../"))
+
+import litellm
+from litellm.proxy._types import (
+    LiteLLM_BudgetTable,
+    LiteLLM_OrganizationTable,
+    LiteLLM_TeamTable,
+    UserAPIKeyAuth,
+)
+from litellm.proxy.auth.auth_checks import common_checks
+from litellm.proxy.utils import ProxyLogging
+
+
+@pytest.mark.asyncio
+async def test_organization_budget_exceeded_blocks_request():
+    """
+    Bug: Organization budget is retrieved but NEVER enforced.
+
+    When organization spend >= organization_max_budget, requests should fail
+    with BudgetExceededError. Currently this passes because no check exists.
+    """
+    org_id = "test-org-budget-exceeded"
+
+    # Organization with max_budget of 100, but spend is 150
+    org_object = LiteLLM_OrganizationTable(
+        organization_id=org_id,
+        budget_id="org-budget-1",
+        spend=150.0,  # Over budget!
+        models=["gpt-4"],
+        created_by="test",
+        updated_by="test",
+        litellm_budget_table=LiteLLM_BudgetTable(
+            max_budget=100.0,  # Budget is 100
+        ),
+    )
+
+    # Team within the organization (team itself is under budget)
+    team_object = LiteLLM_TeamTable(
+        team_id="test-team-1",
+        organization_id=org_id,
+        max_budget=50.0,  # Team budget is 50
+        spend=10.0,       # Team spend is only 10 - under budget
+        models=["gpt-4"],
+    )
+
+    # Valid token with organization info
+    valid_token = UserAPIKeyAuth(
+        token="sk-test-123",
+        team_id="test-team-1",
+        org_id=org_id,
+        organization_max_budget=100.0,  # This is set but never checked!
+    )
+
+    mock_request = MagicMock()
+    mock_request.url.path = "/v1/chat/completions"
+
+    mock_proxy_logging = MagicMock(spec=ProxyLogging)
+    mock_proxy_logging.budget_alerts = AsyncMock()
+
+    with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
+        with patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache:
+            with patch("litellm.proxy.auth.auth_checks.get_org_object", new_callable=AsyncMock) as mock_get_org:
+                mock_get_org.return_value = org_object
+
+                # BUG: This should raise BudgetExceededError but currently passes
+                with pytest.raises(litellm.BudgetExceededError) as exc_info:
+                    await common_checks(
+                        request_body={"model": "gpt-4"},
+                        team_object=team_object,
+                        user_object=None,
+                        end_user_object=None,
+                        global_proxy_spend=None,
+                        general_settings={},
+                        route="/v1/chat/completions",
+                        llm_router=None,
+                        proxy_logging_obj=mock_proxy_logging,
+                        valid_token=valid_token,
+                        request=mock_request,
+                    )
+
+                assert "Organization" in str(exc_info.value.message)
+                assert exc_info.value.current_cost == 150.0
+                assert exc_info.value.max_budget == 100.0
+
+
+@pytest.mark.asyncio
+async def test_multiple_teams_exceed_organization_budget():
+    """
+    Test that organization budget is enforced even when individual teams are under budget.
+
+    Scenario:
+    - Organization max_budget = $5000, spend = $5000 (at limit)
+    - Team A spend = $1500 (under team budget of $2000)
+    - Request via Team A should FAIL because org is at budget limit
+
+    Expected: Request fails with BudgetExceededError
+    """
+    org_id = "multi-team-org"
+
+    # Organization at budget limit
+    org_object = LiteLLM_OrganizationTable(
+        organization_id=org_id,
+        budget_id="org-budget-2",
+        spend=5000.0,  # At $5000 limit
+        models=["gpt-4"],
+        created_by="test",
+        updated_by="test",
+        litellm_budget_table=LiteLLM_BudgetTable(
+            max_budget=5000.0,  # Org budget is $5000
+        ),
+    )
+
+    # Team A - under its own budget, but org is almost at limit
+    team_a = LiteLLM_TeamTable(
+        team_id="team-a",
+        organization_id=org_id,
+        max_budget=2000.0,
+        spend=1500.0,  # Team A has spent $1500 of its $2000 budget
+        models=["gpt-4"],
+    )
+
+    valid_token = UserAPIKeyAuth(
+        token="sk-team-a-key",
+        team_id="team-a",
+        org_id=org_id,
+        organization_max_budget=5000.0,  # Set but never enforced
+    )
+
+    mock_request = MagicMock()
+    mock_request.url.path = "/v1/chat/completions"
+
+    mock_proxy_logging = MagicMock(spec=ProxyLogging)
+    mock_proxy_logging.budget_alerts = AsyncMock()
+
+    with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
+        with patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache:
+            with patch("litellm.proxy.auth.auth_checks.get_org_object", new_callable=AsyncMock) as mock_get_org:
+                mock_get_org.return_value = org_object
+
+                # Org is at budget limit, should raise BudgetExceededError
+                with pytest.raises(litellm.BudgetExceededError) as exc_info:
+                    await common_checks(
+                        request_body={"model": "gpt-4"},
+                        team_object=team_a,
+                        user_object=None,
+                        end_user_object=None,
+                        global_proxy_spend=None,
+                        general_settings={},
+                        route="/v1/chat/completions",
+                        llm_router=None,
+                        proxy_logging_obj=mock_proxy_logging,
+                        valid_token=valid_token,
+                        request=mock_request,
+                    )
+
+                # Verify the error message mentions organization
+                assert "Organization" in str(exc_info.value.message)
+                assert exc_info.value.current_cost == 5000.0
+                assert exc_info.value.max_budget == 5000.0
+
+
+@pytest.mark.asyncio
+async def test_organization_budget_fields_are_checked():
+    """
+    Verify that organization_max_budget is populated in UserAPIKeyAuth
+    and BudgetExceededError is raised when organization is over budget.
+    """
+    # Token has org budget info
+    valid_token = UserAPIKeyAuth(
+        token="sk-test",
+        team_id="test-team",
+        org_id="test-org",
+        organization_max_budget=100.0,  # Budget is $100
+    )
+
+    # Verify the field exists and is set
+    assert valid_token.organization_max_budget == 100.0
+    assert valid_token.org_id == "test-org"
+
+    team_object = LiteLLM_TeamTable(
+        team_id="test-team",
+        organization_id="test-org",
+        max_budget=None,
+        spend=0.0,
+        models=["gpt-4"],
+    )
+
+    mock_request = MagicMock()
+    mock_request.url.path = "/v1/chat/completions"
+
+    mock_proxy_logging = MagicMock(spec=ProxyLogging)
+    mock_proxy_logging.budget_alerts = AsyncMock()
+
+    # Organization is over budget
+    org_over_budget = LiteLLM_OrganizationTable(
+        organization_id="test-org",
+        budget_id="budget-1",
+        spend=150.0,  # Over $100 budget
+        models=["gpt-4"],
+        created_by="test",
+        updated_by="test",
+        litellm_budget_table=LiteLLM_BudgetTable(max_budget=100.0),
+    )
+
+    with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
+        with patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache:
+            with patch("litellm.proxy.auth.auth_checks.get_org_object", new_callable=AsyncMock) as mock_get_org:
+                mock_get_org.return_value = org_over_budget
+
+                # Organization is over budget, should raise BudgetExceededError
+                with pytest.raises(litellm.BudgetExceededError) as exc_info:
+                    await common_checks(
+                        request_body={"model": "gpt-4"},
+                        team_object=team_object,
+                        user_object=None,
+                        end_user_object=None,
+                        global_proxy_spend=None,
+                        general_settings={},
+                        route="/v1/chat/completions",
+                        llm_router=None,
+                        proxy_logging_obj=mock_proxy_logging,
+                        valid_token=valid_token,
+                        request=mock_request,
+                    )
+
+                assert exc_info.value.current_cost == 150.0
+                assert exc_info.value.max_budget == 100.0
+
+
+@pytest.mark.asyncio
+async def test_both_team_and_org_budget_enforced():
+    """
+    Verify that both team budget and organization budget are enforced consistently.
+
+    This test verifies:
+    1. Team over budget raises BudgetExceededError
+    2. Organization over budget also raises BudgetExceededError
+    """
+    mock_request = MagicMock()
+    mock_request.url.path = "/v1/chat/completions"
+
+    mock_proxy_logging = MagicMock(spec=ProxyLogging)
+    mock_proxy_logging.budget_alerts = AsyncMock()
+
+    # Scenario A: Team over budget - should raise BudgetExceededError
+    team_over_budget = LiteLLM_TeamTable(
+        team_id="team-over",
+        max_budget=100.0,
+        spend=150.0,  # Over budget
+        models=["gpt-4"],
+    )
+
+    valid_token_team = UserAPIKeyAuth(
+        token="sk-team-test",
+        team_id="team-over",
+    )
+
+    with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
+        with patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache:
+            with pytest.raises(litellm.BudgetExceededError) as exc_info:
+                await common_checks(
+                    request_body={"model": "gpt-4"},
+                    team_object=team_over_budget,
+                    user_object=None,
+                    end_user_object=None,
+                    global_proxy_spend=None,
+                    general_settings={},
+                    route="/v1/chat/completions",
+                    llm_router=None,
+                    proxy_logging_obj=mock_proxy_logging,
+                    valid_token=valid_token_team,
+                    request=mock_request,
+                )
+            assert "Team" in str(exc_info.value.message)
+
+    # Scenario B: Org over budget - should also raise BudgetExceededError
+    org_over_budget = LiteLLM_OrganizationTable(
+        organization_id="org-over",
+        budget_id="budget-1",
+        spend=150.0,  # Over $100 budget
+        models=["gpt-4"],
+        created_by="test",
+        updated_by="test",
+        litellm_budget_table=LiteLLM_BudgetTable(max_budget=100.0),
+    )
+
+    team_under_budget = LiteLLM_TeamTable(
+        team_id="team-under",
+        organization_id="org-over",
+        max_budget=50.0,
+        spend=10.0,  # Team is fine
+        models=["gpt-4"],
+    )
+
+    valid_token_org = UserAPIKeyAuth(
+        token="sk-org-test",
+        team_id="team-under",
+        org_id="org-over",
+        organization_max_budget=100.0,
+    )
+
+    with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
+        with patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache:
+            with patch("litellm.proxy.auth.auth_checks.get_org_object", new_callable=AsyncMock) as mock_get_org:
+                mock_get_org.return_value = org_over_budget
+
+                # Organization is over budget, should raise BudgetExceededError
+                with pytest.raises(litellm.BudgetExceededError) as exc_info:
+                    await common_checks(
+                        request_body={"model": "gpt-4"},
+                        team_object=team_under_budget,
+                        user_object=None,
+                        end_user_object=None,
+                        global_proxy_spend=None,
+                        general_settings={},
+                        route="/v1/chat/completions",
+                        llm_router=None,
+                        proxy_logging_obj=mock_proxy_logging,
+                        valid_token=valid_token_org,
+                        request=mock_request,
+                    )
+
+                assert "Organization" in str(exc_info.value.message)
+                assert exc_info.value.current_cost == 150.0
+                assert exc_info.value.max_budget == 100.0