litellm - 💡(How to fix) Fix [Bug]: Organization budget does not reset after specified `budget_duration` expires [1 participants]

Check for existing issues

• I have searched the existing issues and checked that my issue is not a duplicate.

What happened?

When an organization is created with max_budget and budget_duration (e.g., "1d"), the organization's cumulative spend field in LiteLLM_OrganizationTable is never reset to zero after the budget period expires. This causes permanent blocking of the organization once the budget limit is exceeded, because the budget enforcement check (_organization_max_budget_check) compares the never-resetting spend against max_budget.

Budget reset works correctly for keys, users, teams, end-users, and team members -- but organizations are completely omitted from the reset cycle.

The root cause is that ResetBudgetJob.reset_budget() in litellm/proxy/common_utils/reset_budget_job.py processes all entity types except organizations. While the LiteLLM_BudgetTable entry linked to the organization has its budget_reset_at timestamp correctly rotated forward, the organization's own spend field is never zeroed out.

Root Cause Analysis

Architecture difference: organizations vs teams/users/keys

Organizations handle budgets differently from other entities:

Organizations delegate budget_duration, budget_reset_at, and max_budget to the linked LiteLLM_BudgetTable (schema lines 20-21, 14). The organization table itself has only spend and budget_id.

The missing reset: reset_budget_job.py:28-49

async def reset_budget(self):
    if self.prisma_client is not None:
        ### RESET KEY BUDGET ###
        await self.reset_budget_for_litellm_keys()          # line 40

        ### RESET USER BUDGET ###
        await self.reset_budget_for_litellm_users()         # line 43

        ## Reset Team Budget
        await self.reset_budget_for_litellm_teams()         # line 46

        ### RESET ENDUSER (Customer) BUDGET and corresponding Budget duration ###
        await self.reset_budget_for_litellm_budget_table()  # line 49

There is no reset_budget_for_litellm_organizations() call.

The budget table reset omits organizations: reset_budget_job.py:136-181

The reset_budget_for_litellm_budget_table() method handles entities linked via FK to LiteLLM_BudgetTable:

# Line 149-151: Find budgets that need reset
budgets_to_reset = await self.prisma_client.get_data(
    table_name="budget", query_type="find_all", reset_at=now
)

# Line 153-163: Update budget_reset_at timestamps on LiteLLM_BudgetTable
# (This DOES rotate the budget period correctly for the budget entry itself)

# Line 165-173: Reset END-USER spend (linked via budget_id FK)
endusers_to_reset = await self.prisma_client.get_data(
    table_name="enduser", query_type="find_all",
    budget_id_list=[budget.budget_id for budget in budgets_to_reset]
)

# Line 175-177: Reset TEAM MEMBER spend (linked via budget_id FK)
await self.reset_budget_for_litellm_team_members(budgets_to_reset=budgets_to_reset)

# Line 179-181: Reset KEY spend (linked via budget_id FK)
await self.reset_budget_for_keys_linked_to_budgets(budgets_to_reset=budgets_to_reset)

# MISSING: No code to reset ORGANIZATION spend (also linked via budget_id FK)

The budget table entry linked to the organization has its budget_reset_at correctly advanced to the next period (line 155-157). But the organization's spend field in LiteLLM_OrganizationTable is never touched.

Budget enforcement blocks permanently: auth_checks.py:3383-3467

async def _organization_max_budget_check(...):
    # Line 3437: Get max_budget from org's linked budget table
    org_max_budget = org_table.litellm_budget_table.max_budget

    # Line 3444: Compare against org's cumulative (never-reset) spend
    if org_table.spend >= org_max_budget:
        raise litellm.BudgetExceededError(
            message=f"Budget has been exceeded! Organization={org_id} "
                    f"Current cost: {org_table.spend}, Max budget: {org_max_budget}",
        )

Since org_table.spend only ever increases and is never reset, once the organization exceeds its budget, it remains blocked indefinitely regardless of budget_duration.

Spend accumulation works correctly

In litellm/proxy/db/db_spend_update_writer.py:605-634, organization spend IS tracked:

async def _update_org_db(self, response_cost, org_id, prisma_client):
    await self.spend_update_queue.add_update(
        update=SpendUpdateQueueItem(
            entity_type=Litellm_EntityType.ORGANIZATION,
            entity_id=org_id,
            response_cost=response_cost,
        )
    )

And committed to DB (lines 1337-1352):

batcher.litellm_organizationtable.update_many(
    where={"organization_id": org_id},
    data={"spend": {"increment": response_cost}},
)

Spend accumulates correctly. The only missing piece is the reset.

No spend counter cache for organizations

Unlike teams and keys which use spend_counter_cache with keys like spend:team:{team_id} and spend:key:{token}, organization spend is read directly from the database in _organization_max_budget_check(). This simplifies the fix (no cache invalidation needed).

No existing tests for organization budget reset

The test file tests/test_litellm/proxy/common_utils/test_reset_budget_job.py covers keys, users, teams, end-users, and keys linked to budgets. There are zero tests for organization budget reset -- which is consistent with the feature being unimplemented.

Expected behavior

After budget_duration expires (e.g., after 24 hours for "1d"):

1. The organization's spend field in LiteLLM_OrganizationTable should be reset to 0.0
2. The organization's model_spend JSON field should be reset to {}
3. The linked LiteLLM_BudgetTable entry's budget_reset_at should advance to the next period (this already works)
4. The organization should be able to resume making requests within its budget

Actual behavior

1. The LiteLLM_BudgetTable.budget_reset_at is advanced correctly (the budget period rotates)
2. But LiteLLM_OrganizationTable.spend never resets -- it only increases
3. Once spend >= max_budget, the organization is permanently blocked with BudgetExceededError
4. The only workaround is manual database intervention: UPDATE "LiteLLM_OrganizationTable" SET spend = 0 WHERE organization_id = '...'

Error message

After the organization exceeds its budget and the duration expires without reset:

Budget has been exceeded! Organization=<org_id> Current cost: <accumulated_spend>, Max budget: <max_budget>

This error persists indefinitely even after the budget period should have reset.

Steps to Reproduce

1. Create an organization with budget settings:

   curl -X POST http://localhost:4000/organization/new \
     -H "Authorization: Bearer sk-..." \
     -H "Content-Type: application/json" \
     -d '{
       "organization_alias": "test-org",
       "max_budget": 10.0,
       "budget_duration": "1d"
     }'

2. Create API keys with organization_id set to the new organization
3. Generate traffic that consumes part of the budget (e.g., $5 of $10)
4. Verify LiteLLM_OrganizationTable.spend shows accumulated spend
5. Wait for budget_duration to expire (24h for "1d") or manually set the linked LiteLLM_BudgetTable.budget_reset_at to a past timestamp:

   UPDATE "LiteLLM_BudgetTable"
   SET budget_reset_at = NOW() - INTERVAL '1 hour'
   WHERE budget_id = '<org_budget_id>';

6. Wait for the reset job to run (runs every ~10 minutes, controlled by proxy_budget_rescheduler_min_time / proxy_budget_rescheduler_max_time)
7. Check LiteLLM_OrganizationTable.spend -- it remains unchanged
8. Check LiteLLM_BudgetTable.budget_reset_at -- this IS updated to the next period
9. If spend exceeded max_budget, requests continue to fail with BudgetExceededError

Suggested fix

Add organization spend reset to reset_budget_for_litellm_budget_table()

In litellm/proxy/common_utils/reset_budget_job.py, within reset_budget_for_litellm_budget_table(), after the existing team member and key reset calls (after line 181), add organization reset:

                await self.reset_budget_for_keys_linked_to_budgets(
                    budgets_to_reset=budgets_to_reset
                )

                # NEW: Reset organization spend for organizations linked to expiring budgets
                await self.reset_budget_for_litellm_organizations(
                    budgets_to_reset=budgets_to_reset
                )

Implement reset_budget_for_litellm_organizations()

Add a new method to ResetBudgetJob:

async def reset_budget_for_litellm_organizations(
    self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
):
    """
    Resets the spend for organizations linked to budget tiers that are being reset.

    Organizations store their budget config in a linked LiteLLM_BudgetTable entry
    (via budget_id FK) but accumulate spend on LiteLLM_OrganizationTable.spend.
    When the budget's duration expires, the org's spend must be zeroed out.
    """
    budget_ids = [
        budget.budget_id
        for budget in budgets_to_reset
        if budget.budget_id is not None
    ]
    if not budget_ids:
        return

    result = await self.prisma_client.db.litellm_organizationtable.update_many(
        where={
            "budget_id": {"in": budget_ids},
            "spend": {"gt": 0},  # Only reset orgs that have accumulated spend
        },
        data={
            "spend": 0,
            "model_spend": "{}",  # Also reset per-model spend tracking
        },
    )
    verbose_proxy_logger.debug(
        "Reset budget for %s organizations linked to expiring budgets", result
    )

Add tests

Add test cases in tests/test_litellm/proxy/common_utils/test_reset_budget_job.py mirroring the existing patterns for teams/users/keys:

1. Test organization spend resets when linked budget expires -- create org with budget_duration, accumulate spend, trigger reset, verify spend=0
2. Test multiple orgs sharing same budget -- both should reset when the shared budget expires
3. Test org spend does not reset when budget hasn't expired -- verify the reset only happens when budget_reset_at is in the past

Related code locations

Relevant log output

What part of LiteLLM is this about?

Proxy

What LiteLLM version are you on ?

v1.82.3

Twitter / LinkedIn details

No response