From 691505cdae1abbad3f296ca44fcd8a644624d8d3 Mon Sep 17 00:00:00 2001
From: Shivam Rawat <161387515+shivamrawat1@users.noreply.github.com>
Date: Fri, 9 Jan 2026 11:14:40 -0800
Subject: [PATCH] added fix for org level budget enforcement (#18813)

---
 litellm/proxy/auth/auth_checks.py | 108 ++++++++++++++++++------------
 1 file changed, 67 insertions(+), 41 deletions(-)

diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py
index de4973ecc6..26778ece60 100644
--- a/litellm/proxy/auth/auth_checks.py
+++ b/litellm/proxy/auth/auth_checks.py
@@ -147,6 +147,7 @@ async def common_checks(
     # 3.1. If organization is in budget
     await _organization_max_budget_check(
         valid_token=valid_token,
+        team_object=team_object,
         prisma_client=prisma_client,
         user_api_key_cache=user_api_key_cache,
         proxy_logging_obj=proxy_logging_obj,
@@ -2310,61 +2311,86 @@ async def _team_max_budget_check(
 
 async def _organization_max_budget_check(
     valid_token: Optional[UserAPIKeyAuth],
+    team_object: Optional[LiteLLM_TeamTable],
     prisma_client: Optional[PrismaClient],
     user_api_key_cache: DualCache,
     proxy_logging_obj: ProxyLogging,
 ):
     """
     Check if the organization is over its max budget.
+    
+    This function checks the organization budget using:
+    1. First, tries to use valid_token.org_id (if key has organization_id set)
+    2. Falls back to team_object.organization_id (if key doesn't have org_id but team does)
+    
+    This ensures organization budget checks work even when keys don't have organization_id
+    set directly, as long as their team belongs to an organization.
 
     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
-    ):
+    if valid_token is None or prisma_client is None:
         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,
+    # Determine organization_id: first try from token, then fallback to team
+    org_id: Optional[str] = None
+    if valid_token.org_id is not None:
+        org_id = valid_token.org_id
+    elif team_object is not None and team_object.organization_id is not None:
+        org_id = team_object.organization_id
+    
+    # If no organization_id found, skip the check
+    if org_id is None:
+        return
+
+    # Get organization object with budget table to check current spend and max budget
+    try:
+        org_table = await prisma_client.db.litellm_organizationtable.find_unique(
+            where={"organization_id": org_id},
+            include={"litellm_budget_table": True},
+        )
+    except Exception:
+        # If organization lookup fails, skip the check
+        return
+
+    if org_table is None:
+        return
+
+    # Get max_budget from organization's budget table
+    org_max_budget: Optional[float] = None
+    if org_table.litellm_budget_table is not None:
+        org_max_budget = org_table.litellm_budget_table.max_budget
+
+    # Only check if organization has a valid max_budget set
+    if org_max_budget is None or org_max_budget <= 0:
+        return
+
+    # Check if organization spend exceeds max budget
+    if org_table.spend >= org_max_budget:
+        # Trigger budget alert
+        call_info = CallInfo(
+            token=valid_token.token,
+            spend=org_table.spend,
+            max_budget=org_max_budget,
+            user_id=valid_token.user_id,
+            team_id=valid_token.team_id,
+            team_alias=valid_token.team_alias,
+            organization_id=org_id,
+            event_group=Litellm_EntityType.ORGANIZATION,
+        )
+        asyncio.create_task(
+            proxy_logging_obj.budget_alerts(
+                type="organization_budget",
+                user_info=call_info,
+            )
         )
 
-        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}",
-            )
+        raise litellm.BudgetExceededError(
+            current_cost=org_table.spend,
+            max_budget=org_max_budget,
+            message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_table.spend}, Max budget: {org_max_budget}",
+        )
 
 
 async def _tag_max_budget_check(
@@ -2601,4 +2627,4 @@ def _can_object_call_vector_stores(
                 code=status.HTTP_401_UNAUTHORIZED,
             )
 
-    return True
+    return True
\ No newline at end of file