From b6ae59c616fb213b829ce803c7f896988407fe3e Mon Sep 17 00:00:00 2001
From: rahul dhanawade <rdhanawade56@gmail.com>
Date: Fri, 6 Feb 2026 12:59:35 +0530
Subject: [PATCH] Fix #20532: Add organization budget reset to ResetBudgetJob

- Reset org spend when their linked budget expires
- Query expired budgets, then reset orgs by budget_id match
- Run org reset before budget_table reset (order matters)
- Follows existing pattern for keys, users, teams, endusers
---
 .../proxy/common_utils/reset_budget_job.py    | 80 +++++++++++++++++++
 1 file changed, 80 insertions(+)

diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py
index fb600cee26b..d0cd96461ba 100644
--- a/litellm/proxy/common_utils/reset_budget_job.py
+++ b/litellm/proxy/common_utils/reset_budget_job.py
@@ -8,6 +8,7 @@ from litellm._logging import verbose_proxy_logger
 from litellm.proxy._types import (
     LiteLLM_BudgetTableFull,
     LiteLLM_EndUserTable,
+    LiteLLM_OrganizationTable,
     LiteLLM_TeamTable,
     LiteLLM_UserTable,
     LiteLLM_VerificationToken,
@@ -48,6 +49,12 @@ class ResetBudgetJob:
             ### RESET ENDUSER (Customer) BUDGET and corresponding Budget duration ###
             await self.reset_budget_for_litellm_budget_table()
 
+            ##Reset The organization Budget
+            await self.reset_budget_for_litellm_organizations() 
+
+
+            
+
     async def reset_budget_for_litellm_team_members(
         self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
     ):
@@ -481,6 +488,79 @@ class ResetBudgetJob:
             )
             verbose_proxy_logger.exception("Failed to reset budget for teams: %s", e)
 
+
+    async def reset_budget_for_litellm_organizations(self):
+        """
+        Resets spend for LiteLLM Organizations whose linked budget has expired. 
+
+        NOTE: Must run AFTER reset_budget_for_litellm_budget_table() to capture
+        first-time budgets that had budget_reset_at computed.
+        """
+        now = datetime.now(timezone.utc)
+        start_time = time.time()
+        budgets_to_reset: Optional[List[LiteLLM_BudgetTableFull]] = None
+
+        try:
+            # Find budgets whose budget_reset_at is due
+            budgets_to_reset = await self.prisma_client.get_data(
+                table_name="budget", query_type="find_all", reset_at=now
+            )
+
+            budget_ids: List[str] = []
+            if budgets_to_reset:
+                
+                from typing import cast
+                budget_ids = [
+                    cast(str, b.budget_id)
+                    for b in budgets_to_reset
+                    if b.budget_id is not None
+                ]
+
+            if not budget_ids:
+                return
+
+            #Reset org spend for orgs linked to these budgets
+            await self.prisma_client.db.litellm_organizationtable.update_many(
+                where={"budget_id": {"in": budget_ids}},
+                data={
+                    "spend": 0.0,
+                    "model_spend": {},
+                },
+            )
+
+            end_time = time.time()
+            asyncio.create_task(
+                self.proxy_logging_obj.service_logging_obj.async_service_success_hook(
+                    service=ServiceTypes.RESET_BUDGET_JOB,
+                    duration=end_time - start_time,
+                    call_type="reset_budget_organizations",
+                    start_time=start_time,
+                    end_time=end_time,
+                    event_metadata={
+                        "num_budgets_found": len(budgets_to_reset) if budgets_to_reset else 0,
+                        "budget_ids_reset": json.dumps(budget_ids),
+                    },
+                )
+            )
+        except Exception as e:
+            end_time = time.time()
+            asyncio.create_task(
+                self.proxy_logging_obj.service_logging_obj.async_service_failure_hook(
+                    service=ServiceTypes.RESET_BUDGET_JOB,
+                    duration=end_time - start_time,
+                    error=e,
+                    call_type="reset_budget_organizations",
+                    start_time=start_time,
+                    end_time=end_time,
+                    event_metadata={
+                        "num_budgets_found": len(budgets_to_reset) if budgets_to_reset else 0,
+                    },
+                )
+            )
+            verbose_proxy_logger.exception(
+                "Failed to reset budget for organizations: %s", e
+            )
+
     @staticmethod
     async def _reset_budget_common(
         item: Union[LiteLLM_TeamTable, LiteLLM_UserTable, LiteLLM_VerificationToken],