[Feature] Key reset_spend endpoint

[yuneng-jiang]

Relevant issues

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

🆕 New Feature
✅ Test

Changes

Adds a new /key/{key_id}/reset_spend endpoint that allows proxy admins and team admins to reset a key's spend to a specific value. The endpoint validates that the reset value is non-negative, does not exceed the current spend, and respects budget limits. Includes authorization checks to ensure only proxy admins or team admins can reset spend for keys in their teams.

Screenshots

[vercel]

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

Project	Deployment	Actions	Updated (UTC)
litellm	Ready	Preview, Comment	Feb 2, 2026 11:54pm

[greptile-apps]

Greptile Overview

Greptile Summary

This PR adds a new /key/{key_id}/reset_spend endpoint that allows proxy admins and team admins to adjust a key's spend value downward, useful for correcting tracking errors or refunding charges.

Key changes:

• Added ResetSpendRequest model and KEY_RESET_SPEND route enum
• Implemented reset_key_spend_fn endpoint with validation ensuring reset value is non-negative, doesn't exceed current spend, and respects budget limits
• Added _check_proxy_or_team_admin_for_key helper for authorization
• Properly invalidates cache after spend update
• Comprehensive test coverage including validation, authorization, and edge cases

Implementation notes:
The endpoint only allows resetting spend to a value <= current spend (reducing spend, not increasing). It validates against both key-level and budget-level max_budget constraints. The implementation follows existing patterns in the codebase for key management endpoints.

Confidence Score: 4/5

• This PR is safe to merge with minimal risk
• The implementation is solid with proper validation, authorization, and cache invalidation. Comprehensive tests cover all scenarios. Minor deduction because the PR doesn't clarify whether budget tables track aggregate spend across keys (though schema suggests they don't).
• No files require special attention

Important Files Changed

Filename	Overview
litellm/proxy/_types.py	Added KEY_RESET_SPEND route enum and ResetSpendRequest model for the new endpoint. Changes are straightforward and follow existing patterns.
litellm/proxy/management_endpoints/key_management_endpoints.py	Implements reset_key_spend_fn endpoint with validation and authorization helpers. Properly validates input, checks permissions, updates database, and invalidates cache. Minor concern: doesn't update budget table's aggregate spend if shared budgets track total spend.
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py	Comprehensive test coverage for reset_key_spend_fn including validation, authorization, success cases, edge cases, and error handling. Tests are well-structured and thorough.

Sequence Diagram

sequenceDiagram
    participant Client
    participant Endpoint as reset_key_spend_fn
    participant Validator as _validate_reset_spend_value
    participant AuthCheck as _check_proxy_or_team_admin_for_key
    participant DB as Prisma Client
    participant Cache as user_api_key_cache

    Client->>Endpoint: POST /key/{key}/reset_spend<br/>{reset_to: float}
    Endpoint->>Endpoint: Check prisma_client exists
    Endpoint->>Endpoint: Hash key if starts with "sk"
    Endpoint->>DB: find_unique(token, include budget_table)
    DB-->>Endpoint: Return key_in_db or None
    alt Key not found
        Endpoint-->>Client: 404 Key not found
    end
    Endpoint->>Validator: _validate_reset_spend_value(reset_to, key_in_db)
    Validator->>Validator: Check type, non-negative, <= current_spend
    Validator->>Validator: Check against max_budget (key & budget table)
    alt Validation fails
        Validator-->>Client: 400 Validation error
    end
    Validator-->>Endpoint: Return validated reset_to
    Endpoint->>AuthCheck: _check_proxy_or_team_admin_for_key
    AuthCheck->>AuthCheck: Check if PROXY_ADMIN
    alt Not PROXY_ADMIN
        AuthCheck->>DB: get_team_object(team_id)
        DB-->>AuthCheck: Return team_table
        AuthCheck->>AuthCheck: Check if user is team admin
    end
    alt Not authorized
        AuthCheck-->>Client: 403 Forbidden
    end
    AuthCheck-->>Endpoint: Authorization passed
    Endpoint->>DB: update(token, spend=reset_to)
    DB-->>Endpoint: Return updated_key
    Endpoint->>Cache: _delete_cache_key_object(hashed_token)
    Cache-->>Endpoint: Cache invalidated
    Endpoint-->>Client: 200 {spend, previous_spend, max_budget, etc.}

Loading

[greptile-apps]

_{3 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}