[Fix] Keys with budget_id don't inherit budget_duration, causing budget resets to never trigger

[shivamrawat1]

Relevant issues

When creating a virtual key with budget_id pointing to a pre-defined budget tier (e.g., one with budget_duration: "7d"), the key's spend is never reset. The only workaround is manually setting budget_duration on each key individually.

Root Cause: /key/generate stores budget_id as a foreign key but never copies budget_duration from the linked budget tier to the key's own budget_duration / budget_reset_at fields. The reset job queries keys by budget_reset_at < now — keys with budget_reset_at = NULL are never matched, so their spend is never reset.

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have Added testing in the tests/litellm/ directory, Adding at least 1 test is a hard requirement - see details
• My PR passes all unit tests on make test-unit
• My PR's scope is as isolated as possible, it only solves 1 specific problem
• [ ]

CI (LiteLLM team)

CI status guideline:

• 50-55 passing tests: main is stable with minor issues.
• 45-49 passing tests: acceptable but needs attention
• <= 40 passing tests: unstable; be careful with your merges and assess the risk.

• Branch creation CI run
  Link:

• CI run for the last commit
  Link:

• Merge / cherry-pick CI run
  Links:

Type

🐛 Bug Fix
✅ Test

Changes

• Key creation (key_management_endpoints.py): When budget_id is provided without an explicit budget_duration, look up the linked budget tier and inherit its budget_duration onto the key.

• Reset job fallback (reset_budget_job.py): When the budget table reset fires, also reset spend on keys linked via budget_id that have no budget_duration set — covering pre-existing keys created before this fix.

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
litellm	Ready	Preview, Comment	Feb 8, 2026 2:18am

[greptile-apps]

Greptile Overview

Greptile Summary

This PR fixes a critical bug where keys created with budget_id pointing to a budget tier never had their spend reset. The root cause was that these keys didn't inherit budget_duration from their linked budget tier, leaving their budget_reset_at field NULL, which caused the reset job to skip them.

Key changes:

• Key creation: When a key is created with budget_id but without explicit budget_duration, the system now looks up the linked budget tier and inherits its budget_duration onto the key
• Reset job fallback: Added reset_budget_for_keys_linked_to_budgets() method to handle pre-existing keys (created before this fix) that have budget_id but no budget_duration
• Avoids double-reset: Keys with their own budget_duration are excluded from the new fallback reset since they're already handled by reset_budget_for_litellm_keys()
• Comprehensive tests: Added tests for inheritance behavior, explicit override, and the new reset job logic

The implementation correctly uses elif to avoid unnecessary database lookups when budget_duration is explicitly provided, and the new database query happens in the key generation path (management endpoint), not the critical API request path.

Confidence Score: 5/5

• This PR is safe to merge with no identified risks
• The fix is well-designed with proper conditional logic (using elif to avoid unnecessary DB queries), comprehensive test coverage for both forward and backward compatibility, and clear documentation. The database query happens in the admin/management path (not the critical request path), and the reset job includes proper guards against double-resetting.
• No files require special attention

Important Files Changed

Filename	Overview
litellm/proxy/management_endpoints/key_management_endpoints.py	Inherits budget_duration from linked budget tier when not explicitly set on key creation
litellm/proxy/common_utils/reset_budget_job.py	Adds fallback reset for keys linked to budget tiers that lack their own budget_duration

Sequence Diagram

sequenceDiagram
    participant Client
    participant KeyEndpoint as Key Management Endpoint
    participant BudgetTable as Budget Table (DB)
    participant KeyTable as Key Table (DB)
    participant ResetJob as Reset Budget Job

    Note over Client,ResetJob: Key Creation Flow (New Behavior)
    
    Client->>KeyEndpoint: POST /key/generate with budget_id, no budget_duration
    KeyEndpoint->>KeyEndpoint: Check if budget_duration in request
    alt budget_duration NOT provided
        KeyEndpoint->>BudgetTable: find_unique by budget_id
        BudgetTable-->>KeyEndpoint: Return budget tier with duration
        KeyEndpoint->>KeyEndpoint: Inherit budget_duration from tier
    else budget_duration provided
        KeyEndpoint->>KeyEndpoint: Use explicit budget_duration
    end
    KeyEndpoint->>KeyTable: Create key with budget_duration
    KeyEndpoint-->>Client: Return created key

    Note over Client,ResetJob: Budget Reset Flow (Dual Path)
    
    ResetJob->>ResetJob: Scheduled job runs
    
    par Keys with own budget_duration
        ResetJob->>KeyTable: reset_budget_for_litellm_keys()<br/>WHERE budget_reset_at < now
        KeyTable-->>ResetJob: Reset keys with own schedule
    and Budget tier resets
        ResetJob->>BudgetTable: Find expired budget tiers
        BudgetTable-->>ResetJob: Return budgets_to_reset
        ResetJob->>KeyTable: reset_budget_for_keys_linked_to_budgets()<br/>WHERE budget_id IN list AND budget_duration IS NULL
        KeyTable-->>ResetJob: Reset pre-existing keys without duration
    end
    
    ResetJob->>ResetJob: Complete reset cycle

Loading

[greptile-apps]

_{2 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[krrishdholakia]

@shivamrawat1 instead of storing the budget from the budget table onto the key (which would complicate future updates) - can we have the reset logic check for the budget_table.budget_duration?

this would allow for any updates to the budget_table's budget duration to just work with any keys associated with it

[krrishdholakia]

@shin-bot-litellm can you monitor this PR and alert me when it's ready for review (set a heartbeat for 1 check / day until it's merged, please bump @shivamrawat1 for the fix if it's not done)