Keep budget_reset_at off the user-settable budget allowlist

LiteLLM_BudgetTable is documented as "user-controllable params" and its
model_fields.keys() is used as the allowlist for extracting budget fields
from incoming API request bodies (management_helpers/utils.py:88,
organization_endpoints.py:112/255/537/549, project_endpoints.py:197/245/632,
customer_endpoints.py:598). Request models like NewOrganizationRequest
inherit from LiteLLM_BudgetTable, so anything on the base class becomes
user-settable — a caller could set budget_reset_at far in the future and
evade budget cycling.

Move budget_reset_at from the base class to LiteLLM_BudgetTableFull so it
appears on API responses without becoming writable, and type
LiteLLM_TeamMembership.litellm_budget_table as Union[Full, Base] so
Pydantic picks Full when the data has server-managed fields (/team/info
reads Prisma rows that include budget_reset_at and created_at) and Base
when callers construct with only user-settable fields (existing auth
tests and caches).
This commit is contained in:
Ryan Crabbe 2026-04-21 15:38:58 -07:00
parent e5f3e15969
commit 1a0ac9634c
No known key found for this signature in database

View file

@ -1997,7 +1997,12 @@ class TeamRequest(LiteLLMPydanticObjectBase):
class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase):
"""Represents user-controllable params for a LiteLLM_BudgetTable record"""
"""Represents user-controllable params for a LiteLLM_BudgetTable record.
Budget-write paths use `model_fields.keys()` on this class as an allowlist
for user input. Keep server-managed fields (e.g. `budget_reset_at`) on
`LiteLLM_BudgetTableFull` so they aren't user-settable.
"""
budget_id: Optional[str] = None
soft_budget: Optional[float] = None
@ -2007,7 +2012,6 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase):
rpm_limit: Optional[int] = None
model_max_budget: Optional[dict] = None
budget_duration: Optional[str] = None
budget_reset_at: Optional[datetime] = None
allowed_models: Optional[List[str]] = (
None # per-member model scope; empty = inherit team models
)
@ -2016,8 +2020,9 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase):
class LiteLLM_BudgetTableFull(LiteLLM_BudgetTable):
"""Represents all params for a LiteLLM_BudgetTable record"""
"""LiteLLM_BudgetTable + server-managed fields returned on API responses."""
budget_reset_at: Optional[datetime] = None
created_at: datetime
@ -3696,7 +3701,12 @@ class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase):
budget_id: Optional[str] = None
spend: Optional[float] = 0.0
total_spend: Optional[float] = 0.0
litellm_budget_table: Optional[LiteLLM_BudgetTable]
# Union so Pydantic picks Full when data has server-managed fields
# (/team/info) and Base when callers/tests construct with only
# user-settable fields.
litellm_budget_table: Optional[
Union[LiteLLM_BudgetTableFull, LiteLLM_BudgetTable]
]
def safe_get_team_member_rpm_limit(self) -> Optional[int]:
if self.litellm_budget_table is not None: