Fix early-return guard ignoring clear_budget_duration flag

When clear_budget_duration=True with all budget limits None, the guard
at the top of _upsert_budget_and_membership was short-circuiting to
disconnect the budget, bypassing the intended update.  Adding
`not clear_budget_duration` to the guard condition lets the function
reach the update path.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
RoyVivat 2026-04-06 16:01:18 -07:00
parent 9596d27e45
commit be7fd219be
No known key found for this signature in database
GPG key ID: 59743472EC86530E

View file

@ -376,7 +376,7 @@ async def _upsert_budget_and_membership(
If max_budget, tpm_limit, rpm_limit, and budget_duration are all None, the user's budget is removed from the team membership.
If any of these values exist, a budget is updated or created and linked to the team membership.
"""
if max_budget is None and tpm_limit is None and rpm_limit is None and budget_duration is None:
if max_budget is None and tpm_limit is None and rpm_limit is None and budget_duration is None and not clear_budget_duration:
# disconnect the budget since all limits are None
await tx.litellm_teammembership.update(
where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},