From 7f3f8fae2dd6a5470a1f61f325f7a0ca3f009de4 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:33:50 +0000 Subject: [PATCH 01/20] feat(proxy): temporary budget increase for team members Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../migration.sql | 5 + .../litellm_proxy_extras/schema.prisma | 2 + litellm/models/budget.py | 2 + litellm/proxy/_types.py | 17 +++ litellm/proxy/auth/auth_checks.py | 24 +++- .../management_endpoints/team_endpoints.py | 4 + litellm/proxy/schema.prisma | 2 + schema.prisma | 2 + .../proxy/auth/test_auth_checks.py | 118 ++++++++++++++++++ .../test_team_endpoints.py | 23 ++++ 10 files changed, 198 insertions(+), 1 deletion(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260917000000_add_budget_temp_increase/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260917000000_add_budget_temp_increase/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260917000000_add_budget_temp_increase/migration.sql new file mode 100644 index 00000000000..a1c431274a3 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260917000000_add_budget_temp_increase/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable +ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN IF NOT EXISTS "temp_budget_increase" DOUBLE PRECISION; + +-- AlterTable +ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN IF NOT EXISTS "temp_budget_expiry" TIMESTAMP(3); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 139fb031671..547491e7dc6 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -22,6 +22,8 @@ model LiteLLM_BudgetTable { budget_duration String? budget_reset_at DateTime? allowed_models String[] @default([]) // per-member model scope; empty = inherit team models + temp_budget_increase Float? + temp_budget_expiry DateTime? created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") diff --git a/litellm/models/budget.py b/litellm/models/budget.py index 125ce739d6a..ddc694743c4 100644 --- a/litellm/models/budget.py +++ b/litellm/models/budget.py @@ -30,6 +30,8 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase): model_max_budget: dict | None = None budget_duration: str | None = None allowed_models: list[str] | None = None # per-member model scope; empty = inherit team models + temp_budget_increase: float | None = None + temp_budget_expiry: datetime | None = None model_config = ConfigDict(protected_namespaces=()) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 680c63393e8..df891bea5e0 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4397,6 +4397,21 @@ class TeamMemberUpdateRequest(TeamMemberDeleteRequest): default=None, description="List of models this team member can access. Pass an empty list to remove per-member model restrictions.", ) + temp_budget_increase: float | None = Field( + default=None, + description="Temporary additive budget increase for this team member, active until temp_budget_expiry", + ) + temp_budget_expiry: datetime | None = Field( + default=None, + description="UTC expiry for temp_budget_increase", + ) + + @model_validator(mode="after") + def validate_temp_budget(self) -> "TeamMemberUpdateRequest": + if self.temp_budget_increase is not None or self.temp_budget_expiry is not None: + if self.temp_budget_increase is None or self.temp_budget_expiry is None: + raise ValueError("temp_budget_increase and temp_budget_expiry must be set together") + return self class TeamMemberUpdateResponse(MemberUpdateResponse): @@ -4406,6 +4421,8 @@ class TeamMemberUpdateResponse(MemberUpdateResponse): rpm_limit: int | None = None budget_duration: str | None = None allowed_models: list[str] | None = None + temp_budget_increase: float | None = None + temp_budget_expiry: datetime | None = None class TeamModelAddRequest(BaseModel): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3dd2e2d8eb2..fcb35fc41a6 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -14,6 +14,7 @@ import math import re import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast @@ -5295,6 +5296,24 @@ async def _virtual_key_max_budget_alert_check( ) +def _effective_team_member_budget(budget: LiteLLM_BudgetTable, now: datetime) -> float | None: + """Per-member cap including an unexpired temp_budget_increase. Naive + temp_budget_expiry values are treated as UTC (same convention as + _get_temp_budget_increase for keys).""" + if budget.max_budget is None: + return None + if budget.temp_budget_increase is None or budget.temp_budget_expiry is None: + return budget.max_budget + expiry: Final = ( + budget.temp_budget_expiry.replace(tzinfo=timezone.utc) + if budget.temp_budget_expiry.tzinfo is None + else budget.temp_budget_expiry + ) + if expiry <= now: + return budget.max_budget + return budget.max_budget + budget.temp_budget_increase + + async def _check_team_member_budget( team_object: LiteLLM_TeamTable | None, user_object: LiteLLM_UserTable | None, @@ -5330,7 +5349,10 @@ async def _check_team_member_budget( and loaded_membership.litellm_budget_table is not None and loaded_membership.litellm_budget_table.max_budget is not None ): - team_member_budget = loaded_membership.litellm_budget_table.max_budget + team_member_budget = _effective_team_member_budget( + loaded_membership.litellm_budget_table, + now=get_utc_datetime(), + ) else: default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") if isinstance(default_budget_id, str): diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index d16fc0fb40c..0eb0f59e09c 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -3692,6 +3692,8 @@ _MEMBER_BUDGET_PATCH_FIELDS: Final = { "rpm_limit": "rpm_limit", "budget_duration": "budget_duration", "allowed_models": "allowed_models", + "temp_budget_increase": "temp_budget_increase", + "temp_budget_expiry": "temp_budget_expiry", } @@ -3862,6 +3864,8 @@ async def team_member_update( rpm_limit=data.rpm_limit, budget_duration=data.budget_duration, allowed_models=data.allowed_models, + temp_budget_increase=data.temp_budget_increase, + temp_budget_expiry=data.temp_budget_expiry, ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 139fb031671..547491e7dc6 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -22,6 +22,8 @@ model LiteLLM_BudgetTable { budget_duration String? budget_reset_at DateTime? allowed_models String[] @default([]) // per-member model scope; empty = inherit team models + temp_budget_increase Float? + temp_budget_expiry DateTime? created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") diff --git a/schema.prisma b/schema.prisma index 139fb031671..547491e7dc6 100644 --- a/schema.prisma +++ b/schema.prisma @@ -22,6 +22,8 @@ model LiteLLM_BudgetTable { budget_duration String? budget_reset_at DateTime? allowed_models String[] @default([]) // per-member model scope; empty = inherit team models + temp_budget_increase Float? + temp_budget_expiry DateTime? created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index f480e096081..fbaa8c371f2 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -8461,3 +8461,121 @@ def test_route_skips_budget_checks_marks_only_spend_free_routes() -> None: def test_request_skips_budget_checks_extends_route_rule_with_zero_cost_models() -> None: assert request_skips_budget_checks(route="/v1/models", model=None, llm_router=None) is True assert request_skips_budget_checks(route="/v1/chat/completions", model=None, llm_router=None) is False + + +def test_effective_team_member_budget_applies_unexpired_increase() -> None: + from litellm.proxy.auth.auth_checks import _effective_team_member_budget + + budget: Final = LiteLLM_BudgetTable( + max_budget=100.0, + temp_budget_increase=50.0, + temp_budget_expiry=datetime(2100, 1, 1), + ) + assert _effective_team_member_budget(budget, now=datetime(2026, 1, 1, tzinfo=timezone.utc)) == 150.0 + + +def test_effective_team_member_budget_ignores_expired_increase() -> None: + from litellm.proxy.auth.auth_checks import _effective_team_member_budget + + budget: Final = LiteLLM_BudgetTable( + max_budget=100.0, + temp_budget_increase=50.0, + temp_budget_expiry=datetime(2020, 1, 1, tzinfo=timezone.utc), + ) + assert _effective_team_member_budget(budget, now=datetime(2026, 1, 1, tzinfo=timezone.utc)) == 100.0 + + +def test_effective_team_member_budget_without_increase() -> None: + from litellm.proxy.auth.auth_checks import _effective_team_member_budget + + now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + assert _effective_team_member_budget(LiteLLM_BudgetTable(max_budget=100.0), now=now) == 100.0 + assert _effective_team_member_budget(LiteLLM_BudgetTable(max_budget=None), now=now) is None + + +@pytest.mark.asyncio +async def test_team_member_budget_check_temp_budget_increase_extends_cap(): + """Spend above max_budget but below max_budget + active temp increase + must not raise; once the increase expires the same spend must raise.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy._types import LiteLLM_TeamMembership + from litellm.proxy.utils import ProxyLogging + + team_object = LiteLLM_TeamTable(team_id="test-team", metadata={}) + user_object = LiteLLM_UserTable(user_id="test-user") + valid_token = UserAPIKeyAuth( + token="test-token", + user_id="test-user", + team_id="test-team", + ) + + team_membership = LiteLLM_TeamMembership( + user_id="test-user", + team_id="test-team", + spend=0.0, + budget_id="budget-1", + litellm_budget_table=LiteLLM_BudgetTable( + max_budget=100.0, + temp_budget_increase=100.0, + temp_budget_expiry=datetime.now(timezone.utc) + timedelta(hours=1), + ), + ) + + proxy_logging_obj = ProxyLogging(user_api_key_cache=None) + + prisma_client = MagicMock() + prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None) + + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + if counter_key == "spend:team_member:test-user:test-team": + return 150.0 + return fallback_spend + + # $150 spend is over the $100 cap but under the $200 temp-extended cap. + with ( + patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), + patch( + "litellm.proxy.auth.auth_checks.get_team_membership", + new_callable=AsyncMock, + return_value=team_membership, + ), + ): + await _check_team_member_budget( + team_object=team_object, + user_object=user_object, + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=DualCache(), + proxy_logging_obj=proxy_logging_obj, + ) + + expired_membership = LiteLLM_TeamMembership( + user_id="test-user", + team_id="test-team", + spend=0.0, + budget_id="budget-1", + litellm_budget_table=LiteLLM_BudgetTable( + max_budget=100.0, + temp_budget_increase=100.0, + temp_budget_expiry=datetime.now(timezone.utc) - timedelta(hours=1), + ), + ) + with ( + patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), + patch( + "litellm.proxy.auth.auth_checks.get_team_membership", + new_callable=AsyncMock, + return_value=expired_membership, + ), + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_budget( + team_object=team_object, + user_object=user_object, + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=DualCache(), + proxy_logging_obj=proxy_logging_obj, + ) + assert exc_info.value.current_cost == 150.0 + assert exc_info.value.max_budget == 100.0 diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index a89bc9a8a3e..0b9597da057 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -15422,3 +15422,26 @@ async def test_team_info_reports_what_the_caller_may_edit(caller, org_admin, ena ) assert response["team_info"].caller_edit_access.model_dump(mode="json") == expected + + +def test_build_member_budget_patch_maps_temp_budget_fields() -> None: + from litellm.proxy.management_endpoints.team_endpoints import _build_member_budget_patch + + expiry: Final = datetime(2030, 1, 1, tzinfo=timezone.utc) + request: Final = TeamMemberUpdateRequest( + team_id="team-1", + user_id="user-1", + temp_budget_increase=50.0, + temp_budget_expiry=expiry, + ) + assert _build_member_budget_patch(request) == { + "temp_budget_increase": 50.0, + "temp_budget_expiry": expiry, + } + + +def test_team_member_update_request_temp_budget_fields_must_be_set_together() -> None: + with pytest.raises(ValidationError, match="temp_budget_increase and temp_budget_expiry must be set together"): + TeamMemberUpdateRequest(team_id="team-1", user_id="user-1", temp_budget_increase=50.0) + with pytest.raises(ValidationError, match="temp_budget_increase and temp_budget_expiry must be set together"): + TeamMemberUpdateRequest(team_id="team-1", user_id="user-1", temp_budget_expiry="2030-01-01T00:00:00Z") From 25e7253fdafac66c608336336cb73a26e4b054ef Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:34:49 +0000 Subject: [PATCH 02/20] refactor(proxy): drop comments from team member temp budget helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 3 --- tests/test_litellm/proxy/auth/test_auth_checks.py | 1 - 2 files changed, 4 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index fcb35fc41a6..a33cda43758 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5297,9 +5297,6 @@ async def _virtual_key_max_budget_alert_check( def _effective_team_member_budget(budget: LiteLLM_BudgetTable, now: datetime) -> float | None: - """Per-member cap including an unexpired temp_budget_increase. Naive - temp_budget_expiry values are treated as UTC (same convention as - _get_temp_budget_increase for keys).""" if budget.max_budget is None: return None if budget.temp_budget_increase is None or budget.temp_budget_expiry is None: diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index fbaa8c371f2..a773a75eb1e 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -8531,7 +8531,6 @@ async def test_team_member_budget_check_temp_budget_increase_extends_cap(): return 150.0 return fallback_spend - # $150 spend is over the $100 cap but under the $200 temp-extended cap. with ( patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), patch( From e43f19fc7cf3e4ec382d467564c682e086cd3211 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:42:46 +0000 Subject: [PATCH 03/20] docs(proxy): document temp budget fields on organization endpoints Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/management_endpoints/organization_endpoints.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index c6a76a920f6..685037a0d5b 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -376,6 +376,8 @@ async def new_organization( - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias) - object_permission: Optional[LiteLLM_ObjectPermissionBase] - organization-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"]}. IF null or {} then no object permission. - allowed_models: Optional[List[str]] - List of models the organization is allowed to access. If not set, defaults to the models field. + - temp_budget_increase: *Optional[float]* - Temporary additive budget increase for the org, active until temp_budget_expiry. + - temp_budget_expiry: *Optional[str]* - UTC expiry for temp_budget_increase. Case 1: Create new org **without** a budget_id ```bash From 32d1dd0cde1f17dbadf81bfd477c4919e20dfc00 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:48:33 +0000 Subject: [PATCH 04/20] fix(proxy): apply temp budget increase at member spend admission and reservation checks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/user_api_key_auth.py | 6 +++++- litellm/proxy/spend_tracking/budget_reservation.py | 7 ++++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 4cbd4213463..c2dc9b16735 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -46,6 +46,7 @@ from litellm.proxy.auth.auth_checks import ( _can_object_call_model, _check_end_user_budget, _delete_cache_key_object, + _effective_team_member_budget, _get_user_role, _is_model_cost_zero, _is_user_proxy_admin, @@ -2248,7 +2249,10 @@ async def _user_api_key_auth_builder( ) if team_member_info is not None and team_member_info.litellm_budget_table is not None: - team_member_budget: Final = team_member_info.litellm_budget_table.max_budget + team_member_budget: Final = _effective_team_member_budget( + team_member_info.litellm_budget_table, + now=datetime.now(timezone.utc), + ) if team_member_budget is not None and team_member_budget > 0: # Read from cross-pod counter (Redis-first) if available from litellm.proxy.proxy_server import get_current_spend diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 373f2d0fe36..1c6f20e515b 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -690,7 +690,12 @@ async def _get_team_member_budget_counter( team_member_budget: float | None = None if team_membership is not None and team_membership.litellm_budget_table is not None: - team_member_budget = team_membership.litellm_budget_table.max_budget + from litellm.proxy.auth.auth_checks import _effective_team_member_budget + + team_member_budget = _effective_team_member_budget( + team_membership.litellm_budget_table, + now=datetime.now(timezone.utc), + ) else: default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") if isinstance(default_budget_id, str): From 7c1eb197bf9a5ace99a74de74f4a648276addb54 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:52:20 +0000 Subject: [PATCH 05/20] chore(ui): regenerate dashboard API types for team member temp budget fields Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 872875cc535..6878306f7da 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -10701,6 +10701,8 @@ export interface paths { * - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias) * - object_permission: Optional[LiteLLM_ObjectPermissionBase] - organization-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"]}. IF null or {} then no object permission. * - allowed_models: Optional[List[str]] - List of models the organization is allowed to access. If not set, defaults to the models field. + * - temp_budget_increase: *Optional[float]* - Temporary additive budget increase for the org, active until temp_budget_expiry. + * - temp_budget_expiry: *Optional[str]* - UTC expiry for temp_budget_increase. * Case 1: Create new org **without** a budget_id * * ```bash @@ -29379,6 +29381,10 @@ export interface components { rpm_limit?: number | null; /** Soft Budget */ soft_budget?: number | null; + /** Temp Budget Expiry */ + temp_budget_expiry?: string | null; + /** Temp Budget Increase */ + temp_budget_increase?: number | null; /** Tpd Limit */ tpd_limit?: number | null; /** Tpm Limit */ @@ -29414,6 +29420,10 @@ export interface components { rpm_limit?: number | null; /** Soft Budget */ soft_budget?: number | null; + /** Temp Budget Expiry */ + temp_budget_expiry?: string | null; + /** Temp Budget Increase */ + temp_budget_increase?: number | null; /** Tpd Limit */ tpd_limit?: number | null; /** Tpm Limit */ @@ -33179,6 +33189,10 @@ export interface components { rpm_limit?: number | null; /** Soft Budget */ soft_budget?: number | null; + /** Temp Budget Expiry */ + temp_budget_expiry?: string | null; + /** Temp Budget Increase */ + temp_budget_increase?: number | null; /** Tpd Limit */ tpd_limit?: number | null; /** Tpm Limit */ @@ -33302,6 +33316,10 @@ export interface components { tags?: string[] | null; /** Team Id */ team_id: string; + /** Temp Budget Expiry */ + temp_budget_expiry?: string | null; + /** Temp Budget Increase */ + temp_budget_increase?: number | null; /** Tpd Limit */ tpd_limit?: number | null; /** Tpm Limit */ @@ -38013,6 +38031,16 @@ export interface components { rpm_limit?: number | null; /** Team Id */ team_id: string; + /** + * Temp Budget Expiry + * @description UTC expiry for temp_budget_increase + */ + temp_budget_expiry?: string | null; + /** + * Temp Budget Increase + * @description Temporary additive budget increase for this team member, active until temp_budget_expiry + */ + temp_budget_increase?: number | null; /** * Tpm Limit * @description Tokens per minute limit for this team member @@ -38035,6 +38063,10 @@ export interface components { rpm_limit?: number | null; /** Team Id */ team_id: string; + /** Temp Budget Expiry */ + temp_budget_expiry?: string | null; + /** Temp Budget Increase */ + temp_budget_increase?: number | null; /** Tpm Limit */ tpm_limit?: number | null; /** User Email */ @@ -39201,6 +39233,10 @@ export interface components { tags?: string[] | null; /** Team Id */ team_id?: string | null; + /** Temp Budget Expiry */ + temp_budget_expiry?: string | null; + /** Temp Budget Increase */ + temp_budget_increase?: number | null; /** Tpd Limit */ tpd_limit?: number | null; /** Tpm Limit */ From b94cd21707d3262ce388c5e09436c144ae14f58c Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:57:42 +0000 Subject: [PATCH 06/20] test(proxy): suppress TQ008 on member temp budget patches Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/auth/test_auth_checks.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index a773a75eb1e..518dad8c48f 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -8532,8 +8532,8 @@ async def test_team_member_budget_check_temp_budget_increase_extends_cap(): return fallback_spend with ( - patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), - patch( + patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), # test-quality-ok: [TQ008] no seam on the cross-pod spend counter + patch( # test-quality-ok: [TQ008] isolates the check from the DB fetch "litellm.proxy.auth.auth_checks.get_team_membership", new_callable=AsyncMock, return_value=team_membership, @@ -8560,8 +8560,8 @@ async def test_team_member_budget_check_temp_budget_increase_extends_cap(): ), ) with ( - patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), - patch( + patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), # test-quality-ok: [TQ008] no seam on the cross-pod spend counter + patch( # test-quality-ok: [TQ008] isolates the check from the DB fetch "litellm.proxy.auth.auth_checks.get_team_membership", new_callable=AsyncMock, return_value=expired_membership, From f972fddafcb5c0da1966ab82583a9bab333bc8ab Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 18:10:00 +0000 Subject: [PATCH 07/20] test(proxy): include temp budget fields in customer budget table fixture Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/management_endpoints/test_customer_endpoints.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py index 1510d8f671d..77e52f30bb7 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py @@ -806,6 +806,8 @@ _EXPECTED_CUSTOMER = { "model_max_budget": None, "budget_duration": "30d", "allowed_models": [], + "temp_budget_increase": None, + "temp_budget_expiry": None, "budget_reset_at": "2024-02-01T00:00:00", "created_at": "2024-01-01T00:00:00", }, From 701c8809222bf11ac1e98e7303829c24f656075f Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 18:56:30 +0000 Subject: [PATCH 08/20] feat(ui): temporary budget increase controls for team members Adds temp_budget_increase and temp_budget_expiry to the team member edit form with pair validation, seeds stored values into edit mode, sends both through /team/member_update, and adds cached-key auth and reservation regression tests for active and expired increases Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../organization_endpoints.py | 4 +- .../proxy/auth/test_user_api_key_auth.py | 106 ++++++++++++++++++ .../spend_tracking/test_budget_reservation.py | 56 ++++++++- .../src/components/networking.tsx | 8 ++ .../team/EditMembership.integration.test.tsx | 65 +++++++++++ .../src/components/team/EditMembership.tsx | 22 +++- .../src/components/team/TeamInfo.tsx | 31 +++++ .../components/team/TeamMemberTab.test.tsx | 31 ++++- .../src/components/team/TeamMemberTab.tsx | 26 +++-- .../components/team/memberFormValues.test.ts | 79 ++++++++++++- .../src/components/team/memberFormValues.ts | 23 +++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +- 12 files changed, 432 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 685037a0d5b..cc4f8ad5dad 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -376,8 +376,8 @@ async def new_organization( - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias) - object_permission: Optional[LiteLLM_ObjectPermissionBase] - organization-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"]}. IF null or {} then no object permission. - allowed_models: Optional[List[str]] - List of models the organization is allowed to access. If not set, defaults to the models field. - - temp_budget_increase: *Optional[float]* - Temporary additive budget increase for the org, active until temp_budget_expiry. - - temp_budget_expiry: *Optional[str]* - UTC expiry for temp_budget_increase. + - temp_budget_increase: *Optional[float]* - Stored on the org budget row but only enforced for team member budgets today. + - temp_budget_expiry: *Optional[str]* - Stored on the org budget row but only enforced for team member budgets today. Case 1: Create new org **without** a budget_id ```bash diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index ba3e98ee718..a78984ecb2f 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -7573,6 +7573,112 @@ async def test_cached_key_team_member_budget_blocks_at_exact_cap(team_member_spe assert f"TeamMember={user_id}:{team_id}" in exc_info.value.message +@pytest.mark.asyncio +@pytest.mark.parametrize( + "expiry_offset, expect_blocked", + [ + (timedelta(days=1), False), + (timedelta(days=-1), True), + ], +) +async def test_cached_key_team_member_budget_honours_temp_increase(expiry_offset, expect_blocked): + """A member over their permanent cap is admitted while a temp_budget_increase is unexpired + and blocked again once it expires, on the cached-key auth path.""" + from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj + from litellm.proxy.common_utils.user_api_key_cache import team_membership_auth_cache_key + from litellm.proxy.utils import hash_token + + api_key = "sk-team-member-temp-budget" + hashed_token = hash_token(api_key) + team_id = "team-temp-budget" + user_id = "user-temp-budget" + team_member_spend = 2.5 + + user_api_key_cache = DualCache() + await _cache_key_object( + hashed_token=hashed_token, + user_api_key_obj=UserAPIKeyAuth( + token=hashed_token, + team_id=team_id, + user_id=user_id, + team_member_spend=team_member_spend, + ), + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=None, + ) + await user_api_key_cache.async_set_cache( + key=f"team_id:{team_id}", + value=LiteLLM_TeamTableCachedObj(team_id=team_id), + ) + await user_api_key_cache.async_set_cache( + key=user_id, + value=LiteLLM_UserTable(user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER), + ) + await user_api_key_cache.async_set_cache( + key=team_membership_auth_cache_key(team_id=team_id, user_id=user_id), + value=LiteLLM_TeamMembership( + user_id=user_id, + team_id=team_id, + spend=team_member_spend, + budget_id="budget-temp", + litellm_budget_table=LiteLLM_BudgetTable( + max_budget=2.0, + temp_budget_increase=1.0, + temp_budget_expiry=datetime.now(timezone.utc) + expiry_offset, + ), + ), + ) + + mock_request = MagicMock() + mock_request.url.path = "/v1/messages" + mock_request.method = "POST" + mock_request.headers = {"authorization": f"Bearer {api_key}"} + mock_request.query_params = {} + mock_request.state = SimpleNamespace() + + proxy_logging_obj = MagicMock() + proxy_logging_obj.budget_alerts = AsyncMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + + async def _auth(): + return await _user_api_key_auth_builder( + request=mock_request, + api_key=f"Bearer {api_key}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={"model": "claude-sonnet-5", "messages": [{"role": "user", "content": "hi"}]}, + ) + + with ( + patch( # test-quality-ok: the builder reads proxy settings from module globals, no injection seam + "litellm.proxy.proxy_server.general_settings", {"disable_budget_reservation": True} + ), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), # test-quality-ok: module-global proxy state + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: module-global proxy state + patch( # test-quality-ok: seed the cached key, team and membership without a DB + "litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache + ), + patch( # test-quality-ok: module-global proxy state + "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj + ), + patch( # test-quality-ok: the live counter needs Redis or a DB; pin the spend the check compares + "litellm.proxy.proxy_server.get_current_spend", + new=AsyncMock(return_value=team_member_spend), + ), + ): + if not expect_blocked: + result = await _auth() + assert result.team_member_spend == team_member_spend + return + with pytest.raises(ProxyException) as exc_info: + await _auth() + + assert exc_info.value.type == ProxyErrorTypes.budget_exceeded + assert "Max budget: 2.0" in exc_info.value.message + + async def _proxy_exception_for_key( api_key: str, general_settings: dict[str, bool], diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index 3e0acf917aa..05cad26ce10 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -2,6 +2,7 @@ from __future__ import annotations import json import math +from datetime import datetime, timedelta, timezone from types import MappingProxyType from typing import Final @@ -9,10 +10,20 @@ import pytest import litellm from litellm.caching import DualCache +from litellm.models.budget import LiteLLM_BudgetTable from litellm.proxy import proxy_server -from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy._types import ( + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_UserTable, + UserAPIKeyAuth, +) +from litellm.proxy.common_utils.user_api_key_cache import ( + UserApiKeyCache, + team_membership_reservation_cache_key, +) from litellm.proxy.spend_tracking.budget_reservation import ( + _get_team_member_budget_counter, count_request_input_tokens, estimate_request_max_cost, reserve_budget_for_request, @@ -445,3 +456,44 @@ async def test_models_without_a_rust_tokenizer_stay_in_python( assert factory.calls == [] assert dict(counts) == dict(python_counts) assert counts[model] not in RUST_INPUT_TOKENS_BY_TOKENIZER.values() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "expiry_offset, expected_max_budget", + [ + (timedelta(days=1), 3.0), + (timedelta(days=-1), 2.0), + ], +) +async def test_team_member_reservation_counter_honours_temp_budget_increase( + expiry_offset: timedelta, expected_max_budget: float +) -> None: + user_id: Final = "member-temp" + team_id: Final = "team-temp" + cache: Final = UserApiKeyCache() + await cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), + value=LiteLLM_TeamMembership( + user_id=user_id, + team_id=team_id, + spend=0.5, + budget_id="budget-temp", + litellm_budget_table=LiteLLM_BudgetTable( + max_budget=2.0, + temp_budget_increase=1.0, + temp_budget_expiry=datetime.now(timezone.utc) + expiry_offset, + ), + ), + ) + + counter: Final = await _get_team_member_budget_counter( + valid_token=UserAPIKeyAuth(token="hashed", user_id=user_id, team_id=team_id), + team_object=LiteLLM_TeamTable(team_id=team_id), + user_object=LiteLLM_UserTable(user_id=user_id), + user_api_key_cache=cache, + ) + + assert counter is not None + assert counter.max_budget == expected_max_budget + assert counter.fallback_spend == 0.5 diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index e77c8ba7e41..feaf334ade7 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -2952,6 +2952,8 @@ export interface Member { rpm_limit?: number | null; budget_duration?: string | null; allowed_models?: string[] | null; + temp_budget_increase?: number | null; + temp_budget_expiry?: string | null; } export const teamMemberAddCall = async (accessToken: string, teamId: string, formValues: Member) => { @@ -3086,6 +3088,12 @@ export const teamMemberUpdateCall = async ( if (formValues.allowed_models !== undefined) { requestBody.allowed_models = formValues.allowed_models; } + if ("temp_budget_increase" in formValues) { + requestBody.temp_budget_increase = orNull(formValues.temp_budget_increase); + } + if ("temp_budget_expiry" in formValues) { + requestBody.temp_budget_expiry = orNull(formValues.temp_budget_expiry); + } const response = await fetch(url, { method: "POST", diff --git a/ui/litellm-dashboard/src/components/team/EditMembership.integration.test.tsx b/ui/litellm-dashboard/src/components/team/EditMembership.integration.test.tsx index a82f475512c..d542b19f718 100644 --- a/ui/litellm-dashboard/src/components/team/EditMembership.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/team/EditMembership.integration.test.tsx @@ -28,6 +28,27 @@ const additionalFields = [ const teamMemberConfig = { title: "Edit Member", showEmail: true, showUserId: true, roleOptions, additionalFields }; +const tempBudgetConfig = { + ...teamMemberConfig, + additionalFields: [ + ...additionalFields, + { name: "temp_budget_increase", label: "Temporary Budget Increase (USD)", type: "numerical" as const, step: 0.01 }, + { name: "temp_budget_expiry", label: "Temporary Budget Expiry (UTC)", type: "utc-datetime" as const }, + ], +}; + +const TEMP_BUDGET_PAIR_MESSAGE = "Set both a temporary budget increase and its expiry, or neither"; + +const cappedMember = { user_id: "u1", user_email: "a@b.com", role: "user", max_budget_in_team: 10 }; + +const tempBudgetMember = { + user_id: "u1", + user_email: "a@b.com", + role: "user", + temp_budget_increase: 25, + temp_budget_expiry: "2030-01-02T03:04:00Z", +}; + const orgMemberConfig = { title: "Edit Member", showEmail: true, showUserId: true, roleOptions }; type Member = Record; @@ -242,6 +263,50 @@ describe("EditMembership submit payload", () => { await waitFor(() => expect(onSubmit).not.toHaveBeenCalled()); }); + it("submits a typed temporary increase with its expiry as a UTC ISO timestamp", async () => { + renderEdit(tempBudgetConfig, cappedMember); + + fireEvent.change(screen.getByLabelText("Temporary Budget Increase (USD)"), { target: { value: "25" } }); + fireEvent.change(screen.getByLabelText("Temporary Budget Expiry (UTC)"), { target: { value: "2030-01-02T03:04" } }); + + save(); + + await waitFor(() => expect(onSubmit).toHaveBeenCalledOnce()); + expect(submitted().max_budget_in_team).toBe(10); + expect(submitted().temp_budget_increase).toBe("25"); + expect(submitted().temp_budget_expiry).toBe("2030-01-02T03:04:00.000Z"); + }); + + it("seeds a stored temporary budget into the controls and clears both to null when the operator blanks them", async () => { + renderEdit(tempBudgetConfig, tempBudgetMember); + + expect(screen.getByLabelText("Temporary Budget Increase (USD)")).toHaveValue(25); + expect(screen.getByLabelText("Temporary Budget Expiry (UTC)")).toHaveValue("2030-01-02T03:04"); + + fireEvent.change(screen.getByLabelText("Temporary Budget Increase (USD)"), { target: { value: "" } }); + fireEvent.change(screen.getByLabelText("Temporary Budget Expiry (UTC)"), { target: { value: "" } }); + + save(); + + await waitFor(() => expect(onSubmit).toHaveBeenCalledOnce()); + expect(submitted().temp_budget_increase).toBeNull(); + expect(submitted().temp_budget_expiry).toBeNull(); + }); + + it.each([ + ["Temporary Budget Increase (USD)", "25"], + ["Temporary Budget Expiry (UTC)", "2030-01-02T03:04"], + ])("blocks submission when only %s is set", async (label, value) => { + renderEdit(tempBudgetConfig, { user_id: "u1", user_email: "a@b.com", role: "user" }); + + fireEvent.change(screen.getByLabelText(label), { target: { value } }); + + save(); + + expect(await screen.findByText(TEMP_BUDGET_PAIR_MESSAGE)).toBeInTheDocument(); + expect(onSubmit).not.toHaveBeenCalled(); + }); + it("clears the fields once the submit handler resolves", async () => { renderEdit(orgMemberConfig, { user_id: "u1", user_email: "a@b.com", role: "user" }); diff --git a/ui/litellm-dashboard/src/components/team/EditMembership.tsx b/ui/litellm-dashboard/src/components/team/EditMembership.tsx index 909b5d56c97..5f342571964 100644 --- a/ui/litellm-dashboard/src/components/team/EditMembership.tsx +++ b/ui/litellm-dashboard/src/components/team/EditMembership.tsx @@ -1,4 +1,6 @@ import React, { useEffect, useMemo, useState } from "react"; +import dayjs from "dayjs"; +import utc from "dayjs/plugin/utc"; import { z } from "zod/v4"; import NumericalInput from "../shared/numerical_input"; import BudgetDurationDropdown from "../common_components/budget_duration_dropdown"; @@ -9,17 +11,22 @@ import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; +import { UtcDateTimeInput } from "@/components/shared/form/UtcDateTimeInput"; import { useZodForm } from "@/lib/forms/useZodForm"; import { buildMemberFormData, buildMemberFormValues, emptyMemberFormValues, + TEMP_BUDGET_PAIR_MESSAGE, + tempBudgetPairError, type MemberAdditionalField, type MemberFieldsConfig, type MemberFormValues, } from "./memberFormValues"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +dayjs.extend(utc); + interface BaseMember { user_email?: string; user_id?: string; @@ -53,7 +60,10 @@ const buildMemberSchema = (config: ModalConfig): z.ZodType [field.name, memberFieldSchema])), }; - return z.object(shape); + return z.object(shape).superRefine((values, ctx) => { + const path = tempBudgetPairError(values); + if (path !== null) ctx.addIssue({ code: "custom", path: [path], message: TEMP_BUDGET_PAIR_MESSAGE }); + }); }; const MemberModal = ({ @@ -160,6 +170,16 @@ const MemberModal = ({ onChange={(next) => onChange(mode === "add" ? next ?? undefined : next)} /> ); + case "utc-datetime": + return ( + onChange(next === null ? null : next.toISOString())} + /> + ); default: return null; } diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index df7b06661c2..719198c03cb 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -264,6 +264,8 @@ export interface TeamMembership { budget_duration: string | null; budget_reset_at: string | null; allowed_models?: string[] | null; + temp_budget_increase?: number | null; + temp_budget_expiry?: string | null; }; } @@ -799,6 +801,8 @@ const TeamInfoView: React.FC = ({ rpm_limit: values.rpm_limit, budget_duration: values.budget_duration, allowed_models: values.allowed_models, + temp_budget_increase: values.temp_budget_increase, + temp_budget_expiry: values.temp_budget_expiry, }; toast.dismiss(); // Remove all existing toasts @@ -2306,6 +2310,33 @@ const TeamInfoView: React.FC = ({ ), type: "budget-duration" as const, }, + { + name: "temp_budget_increase", + label: ( + + Temporary Budget Increase (USD){" "} + + + + + ), + type: "numerical" as const, + step: 0.01, + min: 0, + placeholder: "Extra budget for this member until the expiry", + }, + { + name: "temp_budget_expiry", + label: ( + + Temporary Budget Expiry (UTC){" "} + + + + + ), + type: "utc-datetime" as const, + }, { name: "tpm_limit", label: ( diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx index 2c119bb5848..760074d5dc9 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx @@ -3,7 +3,7 @@ import userEvent from "@testing-library/user-event"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders } from "../../../tests/test-utils"; import { TeamData } from "./TeamInfo"; -import TeamMembersComponent from "./TeamMemberTab"; +import TeamMembersComponent, { seedMemberBudgetFields } from "./TeamMemberTab"; vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ useUISettings: vi.fn(), @@ -380,6 +380,35 @@ describe("TeamMembersComponent", () => { expect(mockSetSelectedEditMember).toHaveBeenCalledWith(expect.objectContaining(zeroLimitsMember)); }); + it("seeds the edit payload with the stored temporary budget increase and expiry, keeping a 0 increase as 0", () => { + const budget = { + ...createMockTeamData().team_memberships[0].litellm_budget_table, + temp_budget_increase: 0, + temp_budget_expiry: "2030-01-02T03:04:00Z", + }; + + const seeded = { + user_id: "user1@test.com", + role: "member", + max_budget_in_team: 1000, + tpm_limit: 10000, + rpm_limit: 100, + budget_duration: null, + allowed_models: [], + temp_budget_increase: 0, + temp_budget_expiry: "2030-01-02T03:04:00Z", + }; + expect(seedMemberBudgetFields({ user_id: "user1@test.com", role: "member" }, budget)).toStrictEqual(seeded); + }); + + it("seeds null temporary budget fields for a member without a budget row", () => { + expect(seedMemberBudgetFields({ user_id: "user2@test.com", role: "admin" }, undefined)).toMatchObject({ + max_budget_in_team: null, + temp_budget_increase: null, + temp_budget_expiry: null, + }); + }); + it("should call setIsAddMemberModalVisible when Add Member button is clicked", async () => { const user = userEvent.setup(); diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx index 3780770315c..16f3d12d71c 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx @@ -8,7 +8,21 @@ import { formatNumberWithCommas } from "@/utils/dataUtils"; import { isProxyAdminRole, isUserTeamAdminForSingleTeam } from "@/utils/roles"; import { CircleHelp } from "lucide-react"; import type { ComponentProps } from "react"; -import { TeamData } from "./TeamInfo"; +import { TeamData, TeamMembership } from "./TeamInfo"; + +export const seedMemberBudgetFields = ( + record: Member, + budget: TeamMembership["litellm_budget_table"] | undefined, +): Member => ({ + ...record, + max_budget_in_team: budget?.max_budget ?? null, + tpm_limit: budget?.tpm_limit ?? null, + rpm_limit: budget?.rpm_limit ?? null, + budget_duration: budget?.budget_duration || null, + allowed_models: budget?.allowed_models || [], + temp_budget_increase: budget?.temp_budget_increase ?? null, + temp_budget_expiry: budget?.temp_budget_expiry ?? null, +}); interface TeamMemberTabProps { teamData: TeamData; @@ -192,15 +206,7 @@ export default function TeamMemberTab({ canEdit={canEditTeam} onEdit={(record) => { const membership = teamData.team_memberships.find((tm) => tm.user_id === record.user_id); - const enhancedMember = { - ...record, - max_budget_in_team: membership?.litellm_budget_table?.max_budget ?? null, - tpm_limit: membership?.litellm_budget_table?.tpm_limit ?? null, - rpm_limit: membership?.litellm_budget_table?.rpm_limit ?? null, - budget_duration: membership?.litellm_budget_table?.budget_duration || null, - allowed_models: membership?.litellm_budget_table?.allowed_models || [], - }; - setSelectedEditMember(enhancedMember); + setSelectedEditMember(seedMemberBudgetFields(record, membership?.litellm_budget_table)); setIsEditMemberModalVisible(true); }} onDelete={handleMemberDelete} diff --git a/ui/litellm-dashboard/src/components/team/memberFormValues.test.ts b/ui/litellm-dashboard/src/components/team/memberFormValues.test.ts index 7a1fcd85a80..fb6602e9f86 100644 --- a/ui/litellm-dashboard/src/components/team/memberFormValues.test.ts +++ b/ui/litellm-dashboard/src/components/team/memberFormValues.test.ts @@ -4,6 +4,7 @@ import { buildMemberFormValues, emptyMemberFormValues, memberFieldNames, + tempBudgetPairError, type MemberFieldsConfig, } from "./memberFormValues"; @@ -25,6 +26,16 @@ const teamConfig: MemberFieldsConfig = { ], }; +const tempBudgetConfig: MemberFieldsConfig = { + roleOptions, + showUserId: true, + additionalFields: [ + { name: "max_budget_in_team", label: "Budget", type: "numerical" }, + { name: "temp_budget_increase", label: "Temp Increase", type: "numerical" }, + { name: "temp_budget_expiry", label: "Temp Expiry", type: "utc-datetime" }, + ], +}; + const orgConfig: MemberFieldsConfig = { roleOptions, showEmail: true, showUserId: true }; describe("memberFieldNames", () => { @@ -117,6 +128,30 @@ describe("buildMemberFormValues", () => { ).toStrictEqual(unlimitedMember); }); + it("seeds a stored temporary budget increase and its expiry, keeping a 0 increase as 0", () => { + const tempBudgetMember = { + user_id: "u1", + role: "user", + max_budget_in_team: 10, + temp_budget_increase: 0, + temp_budget_expiry: "2030-01-01T00:00:00Z", + }; + expect(buildMemberFormValues("edit", tempBudgetMember, tempBudgetConfig)).toStrictEqual(tempBudgetMember); + }); + + it("collapses a missing temporary budget increase and expiry to null", () => { + const noTempBudget = { + user_id: "u1", + role: "user", + max_budget_in_team: null, + temp_budget_increase: null, + temp_budget_expiry: null, + }; + expect(buildMemberFormValues("edit", { user_id: "u1", role: "user" }, tempBudgetConfig)).toStrictEqual( + noTempBudget, + ); + }); + it("falls back to the configured default role when the member has none", () => { expect(buildMemberFormValues("edit", { user_id: "u1", role: "" }, { ...orgConfig, defaultRole: "user" }).role).toBe( "user", @@ -157,6 +192,17 @@ describe("emptyMemberFormValues", () => { }); }); + it("clears a utc-datetime field to null", () => { + const cleared = { + user_id: "", + role: "", + max_budget_in_team: null, + temp_budget_increase: null, + temp_budget_expiry: null, + }; + expect(emptyMemberFormValues(tempBudgetConfig)).toStrictEqual(cleared); + }); + it("clears numeric, duration and multi-select fields to values their controls accept", () => { expect( emptyMemberFormValues({ @@ -199,9 +245,12 @@ describe("buildMemberFormData", () => { }); }); - it.each(["max_budget_in_team", "tpm_limit", "rpm_limit"])("turns a blank %s into null", (key) => { - expect(buildMemberFormData({ [key]: " " })[key]).toBeNull(); - }); + it.each(["max_budget_in_team", "tpm_limit", "rpm_limit", "temp_budget_increase"])( + "turns a blank %s into null", + (key) => { + expect(buildMemberFormData({ [key]: " " })[key]).toBeNull(); + }, + ); it.each(["user_email", "user_id", "budget_duration"])("leaves a blank %s as an empty string", (key) => { expect(buildMemberFormData({ [key]: " " })[key]).toBe(""); @@ -226,3 +275,27 @@ describe("buildMemberFormData", () => { ]); }); }); + +describe("tempBudgetPairError", () => { + it.each([ + [{ temp_budget_increase: 50, temp_budget_expiry: "2030-01-01T00:00:00.000Z" }], + [{ temp_budget_increase: "0", temp_budget_expiry: "2030-01-01T00:00:00.000Z" }], + [{ temp_budget_increase: 0, temp_budget_expiry: "2030-01-01T00:00:00.000Z" }], + [{ temp_budget_increase: null, temp_budget_expiry: null }], + [{ temp_budget_increase: "", temp_budget_expiry: null }], + [{}], + ])("accepts %j", (values) => { + expect(tempBudgetPairError(values)).toBeNull(); + }); + + it("points at the missing increase when only the expiry is set", () => { + expect(tempBudgetPairError({ temp_budget_increase: "", temp_budget_expiry: "2030-01-01T00:00:00.000Z" })).toBe( + "temp_budget_increase", + ); + }); + + it("points at the missing expiry when only the increase is set", () => { + expect(tempBudgetPairError({ temp_budget_increase: 25, temp_budget_expiry: null })).toBe("temp_budget_expiry"); + expect(tempBudgetPairError({ temp_budget_increase: 0, temp_budget_expiry: undefined })).toBe("temp_budget_expiry"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/team/memberFormValues.ts b/ui/litellm-dashboard/src/components/team/memberFormValues.ts index 51b8fac71c1..77cf2dd3e43 100644 --- a/ui/litellm-dashboard/src/components/team/memberFormValues.ts +++ b/ui/litellm-dashboard/src/components/team/memberFormValues.ts @@ -2,7 +2,7 @@ export type MemberFieldValue = string | number | null | undefined | string[]; export type MemberFormValues = Record; -export type MemberFieldType = "input" | "select" | "numerical" | "multi-select" | "budget-duration"; +export type MemberFieldType = "input" | "select" | "numerical" | "multi-select" | "budget-duration" | "utc-datetime"; export interface MemberAdditionalField { name: string; @@ -22,7 +22,23 @@ export interface MemberFieldsConfig { additionalFields?: Array; } -const NULLABLE_NUMERIC_FIELDS: ReadonlySet = new Set(["max_budget_in_team", "tpm_limit", "rpm_limit"]); +const NULLABLE_NUMERIC_FIELDS: ReadonlySet = new Set([ + "max_budget_in_team", + "tpm_limit", + "rpm_limit", + "temp_budget_increase", +]); + +export const TEMP_BUDGET_PAIR_MESSAGE = "Set both a temporary budget increase and its expiry, or neither"; + +const isUnset = (value: MemberFieldValue): boolean => value === null || value === undefined || value === ""; + +export const tempBudgetPairError = (values: MemberFormValues): "temp_budget_increase" | "temp_budget_expiry" | null => { + const increaseUnset = isUnset(values.temp_budget_increase); + const expiryUnset = isUnset(values.temp_budget_expiry); + if (increaseUnset === expiryUnset) return null; + return increaseUnset ? "temp_budget_increase" : "temp_budget_expiry"; +}; export const memberFieldNames = (config: MemberFieldsConfig): string[] => [ ...(config.showEmail ? ["user_email"] : []), @@ -48,6 +64,8 @@ export const buildMemberFormValues = ( rpm_limit: initialData.rpm_limit ?? null, budget_duration: initialData.budget_duration || null, allowed_models: initialData.allowed_models || [], + temp_budget_increase: initialData.temp_budget_increase ?? null, + temp_budget_expiry: initialData.temp_budget_expiry || null, }; return pickFieldNames(config, seeded); @@ -62,6 +80,7 @@ const emptyValueForType = (type: MemberFieldType | undefined): MemberFieldValue return []; case "numerical": case "budget-duration": + case "utc-datetime": return null; default: return ""; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6878306f7da..27ffb905237 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -10701,8 +10701,8 @@ export interface paths { * - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias) * - object_permission: Optional[LiteLLM_ObjectPermissionBase] - organization-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"]}. IF null or {} then no object permission. * - allowed_models: Optional[List[str]] - List of models the organization is allowed to access. If not set, defaults to the models field. - * - temp_budget_increase: *Optional[float]* - Temporary additive budget increase for the org, active until temp_budget_expiry. - * - temp_budget_expiry: *Optional[str]* - UTC expiry for temp_budget_increase. + * - temp_budget_increase: *Optional[float]* - Stored on the org budget row but only enforced for team member budgets today. + * - temp_budget_expiry: *Optional[str]* - Stored on the org budget row but only enforced for team member budgets today. * Case 1: Create new org **without** a budget_id * * ```bash From b9d0008d970d4301bf897f1131d0cd7a78f90ac7 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 19:50:50 +0000 Subject: [PATCH 09/20] fix(proxy): keep temp budget fields out of organization metadata on create Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../organization_endpoints.py | 5 ++- .../test_organization_endpoints.py | 43 +++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index cc4f8ad5dad..b0b0ed379c3 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -291,6 +291,9 @@ async def _verify_org_access( _STR_OBJECT_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) _BUDGET_SETTABLE_FIELDS: Final = frozenset(LiteLLM_BudgetTable.model_fields.keys()) - {"budget_id"} _ORG_COLUMN_FIELDS: Final = frozenset({"organization_alias", "models"}) +_ORG_METADATA_FIELDS: Final = tuple( + field for field in LiteLLM_ManagementEndpoint_MetadataFields if field not in _BUDGET_SETTABLE_FIELDS +) def build_budget_write_data(budget_updates: Mapping[str, object], updated_by: str) -> Mapping[str, object]: @@ -514,7 +517,7 @@ async def new_organization( organization_payload["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name organization_row: Final = LiteLLM_OrganizationTable.model_validate(organization_payload) - for field in LiteLLM_ManagementEndpoint_MetadataFields: + for field in _ORG_METADATA_FIELDS: if getattr(data, field, None) is not None: _set_object_metadata_field( object_data=organization_row, diff --git a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py index 47ee5dc1dd2..c13d93ab862 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py @@ -1346,6 +1346,49 @@ async def test_new_organization_rejects_shared_alias_tool_permission_key(): prisma_client.db.litellm_objectpermissiontable.create.assert_not_called() +@pytest.mark.asyncio +async def test_new_organization_temp_budget_fields_go_to_budget_row_not_metadata(monkeypatch): + """temp_budget_increase/expiry are budget columns and also key-metadata field names, so + /organization/new must write them to the budget row and keep the datetime out of the org + metadata JSON (a datetime there broke JSON serialization and 500'd the request).""" + from datetime import datetime, timezone + + from litellm.proxy._types import LitellmUserRoles, NewOrganizationRequest, UserAPIKeyAuth + from litellm.proxy.management_endpoints.organization_endpoints import new_organization + from litellm.proxy.utils import PrismaClient + + expiry = datetime(2099, 1, 1, tzinfo=timezone.utc) + prisma_client = MagicMock() + prisma_client.jsonify_object = MagicMock(side_effect=lambda data: PrismaClient.jsonify_object(prisma_client, data)) + prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_budgettable.create = AsyncMock(return_value=MagicMock(budget_id="budget-1")) + prisma_client.db.litellm_organizationtable.create = AsyncMock(return_value={"organization_id": "org-1"}) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True, raising=False) + + response = await new_organization( + data=NewOrganizationRequest( + organization_alias="org", + max_budget=10, + temp_budget_increase=5, + temp_budget_expiry=expiry, + ), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert response == {"organization_id": "org-1"} + budget_write = prisma_client.db.litellm_budgettable.create.await_args.kwargs["data"] + assert (budget_write["max_budget"], budget_write["temp_budget_increase"], budget_write["temp_budget_expiry"]) == ( + 10, + 5, + expiry, + ) + org_write = prisma_client.db.litellm_organizationtable.create.await_args.kwargs["data"] + assert org_write["budget_id"] == "budget-1" + assert json.loads(org_write.get("metadata", "{}")) == {} + + def test_v2_update_organization_is_in_openapi_schema(): """PATCH /v2/organization/{organization_id} is documented in the generated OpenAPI spec.""" from fastapi import FastAPI From d9ddc4b9010f2a5d71cd27fd26461274ca7e4848 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 19:50:50 +0000 Subject: [PATCH 10/20] test(proxy): assert temp budget increase stops at the exact expiry instant Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/auth/test_auth_checks.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 518dad8c48f..81bf1471a00 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -8477,12 +8477,10 @@ def test_effective_team_member_budget_applies_unexpired_increase() -> None: def test_effective_team_member_budget_ignores_expired_increase() -> None: from litellm.proxy.auth.auth_checks import _effective_team_member_budget - budget: Final = LiteLLM_BudgetTable( - max_budget=100.0, - temp_budget_increase=50.0, - temp_budget_expiry=datetime(2020, 1, 1, tzinfo=timezone.utc), - ) + expiry: Final = datetime(2020, 1, 1, tzinfo=timezone.utc) + budget: Final = LiteLLM_BudgetTable(max_budget=100.0, temp_budget_increase=50.0, temp_budget_expiry=expiry) assert _effective_team_member_budget(budget, now=datetime(2026, 1, 1, tzinfo=timezone.utc)) == 100.0 + assert _effective_team_member_budget(budget, now=expiry) == 100.0 def test_effective_team_member_budget_without_increase() -> None: From e93245131261ad4c4a7b05feba8d1325925accc7 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 19:56:35 +0000 Subject: [PATCH 11/20] refactor(proxy): move effective member budget onto the budget model and reject negative temp increases Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/models/budget.py | 17 ++++++++++- litellm/proxy/_types.py | 1 + litellm/proxy/auth/auth_checks.py | 21 +------------ litellm/proxy/auth/user_api_key_auth.py | 4 +-- .../spend_tracking/budget_reservation.py | 7 +---- tests/test_litellm/models/test_models.py | 23 +++++++++++++- .../proxy/auth/test_auth_checks.py | 30 ------------------- .../test_team_endpoints.py | 7 +++++ 8 files changed, 49 insertions(+), 61 deletions(-) diff --git a/litellm/models/budget.py b/litellm/models/budget.py index ddc694743c4..2bef5d279d6 100644 --- a/litellm/models/budget.py +++ b/litellm/models/budget.py @@ -5,7 +5,8 @@ Canonical definition for ``litellm_budgettable``. Re-exported from ``litellm.proxy._types`` for backwards compatibility. """ -from datetime import datetime +from datetime import datetime, timezone +from typing import Final from pydantic import ConfigDict @@ -35,6 +36,20 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase): model_config = ConfigDict(protected_namespaces=()) + def effective_max_budget(self, now: datetime) -> float | None: + if self.max_budget is None: + return None + if self.temp_budget_increase is None or self.temp_budget_expiry is None: + return self.max_budget + expiry: Final = ( + self.temp_budget_expiry.replace(tzinfo=timezone.utc) + if self.temp_budget_expiry.tzinfo is None + else self.temp_budget_expiry + ) + if expiry <= now: + return self.max_budget + return self.max_budget + self.temp_budget_increase + class LiteLLM_BudgetTableFull(LiteLLM_BudgetTable): """LiteLLM_BudgetTable + server-managed fields returned on API responses.""" diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index df891bea5e0..7728c7ee34f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4399,6 +4399,7 @@ class TeamMemberUpdateRequest(TeamMemberDeleteRequest): ) temp_budget_increase: float | None = Field( default=None, + ge=0, description="Temporary additive budget increase for this team member, active until temp_budget_expiry", ) temp_budget_expiry: datetime | None = Field( diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index a33cda43758..a36ef548e2f 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -14,7 +14,6 @@ import math import re import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence -from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast @@ -5296,21 +5295,6 @@ async def _virtual_key_max_budget_alert_check( ) -def _effective_team_member_budget(budget: LiteLLM_BudgetTable, now: datetime) -> float | None: - if budget.max_budget is None: - return None - if budget.temp_budget_increase is None or budget.temp_budget_expiry is None: - return budget.max_budget - expiry: Final = ( - budget.temp_budget_expiry.replace(tzinfo=timezone.utc) - if budget.temp_budget_expiry.tzinfo is None - else budget.temp_budget_expiry - ) - if expiry <= now: - return budget.max_budget - return budget.max_budget + budget.temp_budget_increase - - async def _check_team_member_budget( team_object: LiteLLM_TeamTable | None, user_object: LiteLLM_UserTable | None, @@ -5346,10 +5330,7 @@ async def _check_team_member_budget( and loaded_membership.litellm_budget_table is not None and loaded_membership.litellm_budget_table.max_budget is not None ): - team_member_budget = _effective_team_member_budget( - loaded_membership.litellm_budget_table, - now=get_utc_datetime(), - ) + team_member_budget = loaded_membership.litellm_budget_table.effective_max_budget(now=get_utc_datetime()) else: default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") if isinstance(default_budget_id, str): diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index c2dc9b16735..cffb745fff2 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -46,7 +46,6 @@ from litellm.proxy.auth.auth_checks import ( _can_object_call_model, _check_end_user_budget, _delete_cache_key_object, - _effective_team_member_budget, _get_user_role, _is_model_cost_zero, _is_user_proxy_admin, @@ -2249,8 +2248,7 @@ async def _user_api_key_auth_builder( ) if team_member_info is not None and team_member_info.litellm_budget_table is not None: - team_member_budget: Final = _effective_team_member_budget( - team_member_info.litellm_budget_table, + team_member_budget: Final = team_member_info.litellm_budget_table.effective_max_budget( now=datetime.now(timezone.utc), ) if team_member_budget is not None and team_member_budget > 0: diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 1c6f20e515b..4028fcaf2ae 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -690,12 +690,7 @@ async def _get_team_member_budget_counter( team_member_budget: float | None = None if team_membership is not None and team_membership.litellm_budget_table is not None: - from litellm.proxy.auth.auth_checks import _effective_team_member_budget - - team_member_budget = _effective_team_member_budget( - team_membership.litellm_budget_table, - now=datetime.now(timezone.utc), - ) + team_member_budget = team_membership.litellm_budget_table.effective_max_budget(now=datetime.now(timezone.utc)) else: default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") if isinstance(default_budget_id, str): diff --git a/tests/test_litellm/models/test_models.py b/tests/test_litellm/models/test_models.py index 9b803c14062..46648efdfff 100644 --- a/tests/test_litellm/models/test_models.py +++ b/tests/test_litellm/models/test_models.py @@ -2,7 +2,7 @@ Tests for backend domain models. """ -from datetime import datetime +from datetime import datetime, timezone import pytest from pydantic import BaseModel, TypeAdapter @@ -71,6 +71,27 @@ class TestBudget: assert budget.max_budget is None assert budget.allowed_models is None + def test_effective_max_budget_applies_unexpired_increase(self): + budget = LiteLLM_BudgetTable( + max_budget=100.0, + temp_budget_increase=50.0, + temp_budget_expiry=datetime(2100, 1, 1), + ) + assert budget.effective_max_budget(now=datetime(2026, 1, 1, tzinfo=timezone.utc)) == 150.0 + + def test_effective_max_budget_ignores_expired_increase(self): + budget = LiteLLM_BudgetTable( + max_budget=100.0, + temp_budget_increase=50.0, + temp_budget_expiry=datetime(2020, 1, 1, tzinfo=timezone.utc), + ) + assert budget.effective_max_budget(now=datetime(2026, 1, 1, tzinfo=timezone.utc)) == 100.0 + + def test_effective_max_budget_without_increase(self): + now = datetime(2026, 1, 1, tzinfo=timezone.utc) + assert LiteLLM_BudgetTable(max_budget=100.0).effective_max_budget(now=now) == 100.0 + assert LiteLLM_BudgetTable(max_budget=None, temp_budget_increase=50.0).effective_max_budget(now=now) is None + class TestCredentials: def test_credentials_creation(self): diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 518dad8c48f..3048d4a6a02 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -8463,36 +8463,6 @@ def test_request_skips_budget_checks_extends_route_rule_with_zero_cost_models() assert request_skips_budget_checks(route="/v1/chat/completions", model=None, llm_router=None) is False -def test_effective_team_member_budget_applies_unexpired_increase() -> None: - from litellm.proxy.auth.auth_checks import _effective_team_member_budget - - budget: Final = LiteLLM_BudgetTable( - max_budget=100.0, - temp_budget_increase=50.0, - temp_budget_expiry=datetime(2100, 1, 1), - ) - assert _effective_team_member_budget(budget, now=datetime(2026, 1, 1, tzinfo=timezone.utc)) == 150.0 - - -def test_effective_team_member_budget_ignores_expired_increase() -> None: - from litellm.proxy.auth.auth_checks import _effective_team_member_budget - - budget: Final = LiteLLM_BudgetTable( - max_budget=100.0, - temp_budget_increase=50.0, - temp_budget_expiry=datetime(2020, 1, 1, tzinfo=timezone.utc), - ) - assert _effective_team_member_budget(budget, now=datetime(2026, 1, 1, tzinfo=timezone.utc)) == 100.0 - - -def test_effective_team_member_budget_without_increase() -> None: - from litellm.proxy.auth.auth_checks import _effective_team_member_budget - - now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) - assert _effective_team_member_budget(LiteLLM_BudgetTable(max_budget=100.0), now=now) == 100.0 - assert _effective_team_member_budget(LiteLLM_BudgetTable(max_budget=None), now=now) is None - - @pytest.mark.asyncio async def test_team_member_budget_check_temp_budget_increase_extends_cap(): """Spend above max_budget but below max_budget + active temp increase diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 0b9597da057..b458b85602d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -15445,3 +15445,10 @@ def test_team_member_update_request_temp_budget_fields_must_be_set_together() -> TeamMemberUpdateRequest(team_id="team-1", user_id="user-1", temp_budget_increase=50.0) with pytest.raises(ValidationError, match="temp_budget_increase and temp_budget_expiry must be set together"): TeamMemberUpdateRequest(team_id="team-1", user_id="user-1", temp_budget_expiry="2030-01-01T00:00:00Z") + + +def test_team_member_update_request_rejects_negative_temp_budget_increase() -> None: + with pytest.raises(ValidationError, match="greater than or equal to 0"): + TeamMemberUpdateRequest( + team_id="team-1", user_id="user-1", temp_budget_increase=-1.0, temp_budget_expiry="2030-01-01T00:00:00Z" + ) From 67c522fe737eea4fc7e72bfb154b60caddebdf0f Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 20:07:33 +0000 Subject: [PATCH 12/20] fix(proxy): reject non-finite temp budget increases on team member update Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 1 + .../proxy/management_endpoints/test_team_endpoints.py | 10 +++++++--- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 7728c7ee34f..d9529f7d575 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4400,6 +4400,7 @@ class TeamMemberUpdateRequest(TeamMemberDeleteRequest): temp_budget_increase: float | None = Field( default=None, ge=0, + allow_inf_nan=False, description="Temporary additive budget increase for this team member, active until temp_budget_expiry", ) temp_budget_expiry: datetime | None = Field( diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index b458b85602d..a7748220e44 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -15447,8 +15447,12 @@ def test_team_member_update_request_temp_budget_fields_must_be_set_together() -> TeamMemberUpdateRequest(team_id="team-1", user_id="user-1", temp_budget_expiry="2030-01-01T00:00:00Z") -def test_team_member_update_request_rejects_negative_temp_budget_increase() -> None: - with pytest.raises(ValidationError, match="greater than or equal to 0"): +@pytest.mark.parametrize( + ("increase", "message"), + [(-1.0, "greater than or equal to 0"), (float("inf"), "finite number")], +) +def test_team_member_update_request_rejects_unusable_temp_budget_increase(increase: float, message: str) -> None: + with pytest.raises(ValidationError, match=message): TeamMemberUpdateRequest( - team_id="team-1", user_id="user-1", temp_budget_increase=-1.0, temp_budget_expiry="2030-01-01T00:00:00Z" + team_id="team-1", user_id="user-1", temp_budget_increase=increase, temp_budget_expiry="2030-01-01T00:00:00Z" ) From 97c50acc1944af6260265cdf21ef8966bbb81b07 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 20:30:29 +0000 Subject: [PATCH 13/20] fix(proxy): persist a temp budget pair for members without a private budget row Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/common_utils.py | 2 ++ .../test_upsert_budget_membership.py | 24 +++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 973311608ed..702f2ed4ced 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -487,6 +487,8 @@ _TEAM_MEMBER_BUDGET_LIMIT_FIELDS: Final = ( "model_max_budget", "budget_duration", "allowed_models", + "temp_budget_increase", + "temp_budget_expiry", ) diff --git a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py index e9b4f11e891..eba05362bbf 100644 --- a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py +++ b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py @@ -219,6 +219,30 @@ async def test_create_seeds_reset_at_and_links(mock_tx, fake_user): ) +# TEST: with no existing budget, a patch carrying only the temporary increase +# pair must still create a budget row and link it, so the fields the 200 +# response echoes are actually stored. +@pytest.mark.asyncio +async def test_create_from_temp_budget_pair_only(mock_tx, fake_user): + expiry = datetime(2100, 1, 1, tzinfo=timezone.utc) + await _upsert_budget_and_membership( + mock_tx, + team_id="team-new", + user_id="user-new", + existing_budget_id=None, + user_api_key_dict=fake_user, + budget_patch={"temp_budget_increase": 5.0, "temp_budget_expiry": expiry}, + ) + + mock_tx.litellm_budgettable.create.assert_awaited_once() + data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"] + assert data["temp_budget_increase"] == 5.0 + assert data["temp_budget_expiry"] == expiry + assert "max_budget" not in data + mock_tx.litellm_teammembership.upsert.assert_awaited_once() + mock_tx.litellm_teammembership.update.assert_not_called() + + # TEST: clone-on-write when the membership still points at the team's shared # default budget. Editing this member must fork a private budget instead of # mutating the shared row, and cloning a duration must seed a fresh reset time. From e73b8d49dab0a59d31340ee46976c8b0ce2902f5 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 20:49:58 +0000 Subject: [PATCH 14/20] fix(proxy): seed a new member budget row from the team default cap Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/common_utils.py | 6 ++-- .../test_upsert_budget_membership.py | 28 +++++++++++++++++++ 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 702f2ed4ced..116a9e6ff42 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -570,8 +570,10 @@ async def _upsert_budget_and_membership( "updated_by": user_api_key_dict.user_id or "", } - if is_shared_default: - default_budget_row: Final = await tx.litellm_budgettable.find_unique(where={"budget_id": existing_budget_id}) + if team_default_budget_id is not None: + default_budget_row: Final = await tx.litellm_budgettable.find_unique( + where={"budget_id": team_default_budget_id} + ) if default_budget_row is not None: default_budget_dict: Final = default_budget_row.model_dump() for field in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS: diff --git a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py index eba05362bbf..3b7bd054874 100644 --- a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py +++ b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py @@ -243,6 +243,34 @@ async def test_create_from_temp_budget_pair_only(mock_tx, fake_user): mock_tx.litellm_teammembership.update.assert_not_called() +# TEST: a member with no budget row who falls back to the team default at +# enforcement time must keep that default cap on the new private row, or the +# temporary increase has nothing to add to. +@pytest.mark.asyncio +async def test_create_from_temp_pair_keeps_team_default_cap(mock_tx, fake_user): + expiry = datetime(2100, 1, 1, tzinfo=timezone.utc) + mock_tx.litellm_budgettable.find_unique = AsyncMock( + return_value=budget_row(budget_id="team-default-budget-1", max_budget=0.4, allowed_models=[]) + ) + await _upsert_budget_and_membership( + mock_tx, + team_id="team-default", + user_id="user-unlinked", + existing_budget_id=None, + user_api_key_dict=fake_user, + budget_patch={"temp_budget_increase": 1.0, "temp_budget_expiry": expiry}, + team_default_budget_id="team-default-budget-1", + ) + + mock_tx.litellm_budgettable.find_unique.assert_awaited_once_with(where={"budget_id": "team-default-budget-1"}) + data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"] + assert data["max_budget"] == 0.4 + assert data["temp_budget_increase"] == 1.0 + assert data["temp_budget_expiry"] == expiry + assert "allowed_models" not in data + mock_tx.litellm_teammembership.upsert.assert_awaited_once() + + # TEST: clone-on-write when the membership still points at the team's shared # default budget. Editing this member must fork a private budget instead of # mutating the shared row, and cloning a duration must seed a fresh reset time. From 664688f3372f8e20270ca90eeadcd744f05ef214 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 21:00:31 +0000 Subject: [PATCH 15/20] fix(proxy): do not carry a zero team default cap onto a new member budget row Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/common_utils.py | 2 ++ .../test_upsert_budget_membership.py | 29 +++++++++++++++---- 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 116a9e6ff42..546078cce09 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -578,6 +578,8 @@ async def _upsert_budget_and_membership( default_budget_dict: Final = default_budget_row.model_dump() for field in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS: value = default_budget_dict.get(field) + if field == "max_budget" and value == 0 and not is_shared_default: + continue if _is_set_budget_value(value): create_data[field] = value diff --git a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py index 3b7bd054874..fbd6ef51f05 100644 --- a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py +++ b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py @@ -219,9 +219,6 @@ async def test_create_seeds_reset_at_and_links(mock_tx, fake_user): ) -# TEST: with no existing budget, a patch carrying only the temporary increase -# pair must still create a budget row and link it, so the fields the 200 -# response echoes are actually stored. @pytest.mark.asyncio async def test_create_from_temp_budget_pair_only(mock_tx, fake_user): expiry = datetime(2100, 1, 1, tzinfo=timezone.utc) @@ -243,9 +240,6 @@ async def test_create_from_temp_budget_pair_only(mock_tx, fake_user): mock_tx.litellm_teammembership.update.assert_not_called() -# TEST: a member with no budget row who falls back to the team default at -# enforcement time must keep that default cap on the new private row, or the -# temporary increase has nothing to add to. @pytest.mark.asyncio async def test_create_from_temp_pair_keeps_team_default_cap(mock_tx, fake_user): expiry = datetime(2100, 1, 1, tzinfo=timezone.utc) @@ -271,6 +265,29 @@ async def test_create_from_temp_pair_keeps_team_default_cap(mock_tx, fake_user): mock_tx.litellm_teammembership.upsert.assert_awaited_once() +@pytest.mark.asyncio +async def test_create_from_temp_pair_skips_zero_team_default_cap(mock_tx, fake_user): + expiry = datetime(2100, 1, 1, tzinfo=timezone.utc) + mock_tx.litellm_budgettable.find_unique = AsyncMock( + return_value=budget_row(budget_id="team-default-budget-1", max_budget=0, rpm_limit=10) + ) + await _upsert_budget_and_membership( + mock_tx, + team_id="team-default", + user_id="user-unlinked", + existing_budget_id=None, + user_api_key_dict=fake_user, + budget_patch={"temp_budget_increase": 1.0, "temp_budget_expiry": expiry}, + team_default_budget_id="team-default-budget-1", + ) + + data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"] + assert "max_budget" not in data + assert data["rpm_limit"] == 10 + assert data["temp_budget_increase"] == 1.0 + mock_tx.litellm_teammembership.upsert.assert_awaited_once() + + # TEST: clone-on-write when the membership still points at the team's shared # default budget. Editing this member must fork a private budget instead of # mutating the shared row, and cloning a duration must seed a fresh reset time. From 7d65d9d774056e256ab3a12bac8d52aaced5e190 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 23:17:13 +0000 Subject: [PATCH 16/20] fix(proxy): seed member budget forks from the row the membership points at Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/management_endpoints/common_utils.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 6f88b05a8ae..c65c344992d 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -591,16 +591,12 @@ async def _upsert_budget_and_membership( "updated_by": user_api_key_dict.user_id or "", } - if team_default_budget_id is not None: - default_budget_row: Final = await tx.litellm_budgettable.find_unique( - where={"budget_id": team_default_budget_id} - ) + if is_shared_default: + default_budget_row: Final = await tx.litellm_budgettable.find_unique(where={"budget_id": existing_budget_id}) if default_budget_row is not None: default_budget_dict: Final = default_budget_row.model_dump() for field in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS: value = default_budget_dict.get(field) - if field == "max_budget" and value == 0 and not is_shared_default: - continue if _is_set_budget_value(value): create_data[field] = value From 44f90b5ba73c4b20b62aac9c8b3e7f1f1b4bb426 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 17 Sep 2026 23:30:43 +0000 Subject: [PATCH 17/20] fix(proxy): seed member budget creates from the right source row Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/common_utils.py | 20 ++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index c65c344992d..33bfa1f8b44 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -591,12 +591,22 @@ async def _upsert_budget_and_membership( "updated_by": user_api_key_dict.user_id or "", } - if is_shared_default: - default_budget_row: Final = await tx.litellm_budgettable.find_unique(where={"budget_id": existing_budget_id}) - if default_budget_row is not None: - default_budget_dict: Final = default_budget_row.model_dump() + seed_row_id: Final = ( + existing_budget_id + if is_shared_default + else team_default_budget_id + if team_default_budget_id is not None + and ("temp_budget_increase" in write_data or "temp_budget_expiry" in write_data) + else None + ) + if seed_row_id is not None: + seed_row: Final = await tx.litellm_budgettable.find_unique(where={"budget_id": seed_row_id}) + if seed_row is not None: + seed_dict: Final = seed_row.model_dump() for field in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS: - value = default_budget_dict.get(field) + value = seed_dict.get(field) + if field == "max_budget" and value == 0 and not is_shared_default: + continue if _is_set_budget_value(value): create_data[field] = value From 3926e74e9df086f1c77bb1745da672de25d60934 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 01:28:49 +0000 Subject: [PATCH 18/20] test(proxy): lock plain member updates to the merge-patch contract without a team default snapshot Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_upsert_budget_membership.py | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py index e17a1a287bf..cabf5e24933 100644 --- a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py +++ b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py @@ -280,6 +280,29 @@ async def test_create_from_temp_pair_skips_zero_team_default_cap(mock_tx, fake_u mock_tx.litellm_teammembership.upsert.assert_awaited_once() +@pytest.mark.asyncio +async def test_create_from_plain_patch_does_not_snapshot_team_default(mock_tx, fake_user): + mock_tx.litellm_budgettable.find_unique = AsyncMock( + return_value=budget_row(budget_id="team-default-budget-1", max_budget=0.4, rpm_limit=10) + ) + await _upsert_budget_and_membership( + mock_tx, + team_id="team-default", + user_id="user-unlinked", + existing_budget_id=None, + user_api_key_dict=fake_user, + budget_patch={"tpm_limit": 500}, + team_default_budget_id="team-default-budget-1", + ) + + mock_tx.litellm_budgettable.find_unique.assert_not_awaited() + data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"] + assert data["tpm_limit"] == 500 + assert "max_budget" not in data + assert "rpm_limit" not in data + mock_tx.litellm_teammembership.upsert.assert_awaited_once() + + # TEST: clone-on-write when the membership still points at the team's shared # default budget. Editing this member must fork a private budget instead of # mutating the shared row, and cloning a duration must seed a fresh reset time. From 507219e26ff490a348e26340a8e3704df1a71034 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 02:09:09 +0000 Subject: [PATCH 19/20] fix(proxy): resolve temporary budget grants against the live team default instead of a snapshot A temporary-only member update no longer clones the team default budget into the private row. The row stores just the temp pair and auth, spend admission and reservation add the active increase to the current shared default, so a later lowering of the default reaches members with an active grant Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/models/budget.py | 15 +-- litellm/proxy/auth/auth_checks.py | 14 +-- .../management_endpoints/common_utils.py | 28 +++--- .../spend_tracking/budget_reservation.py | 12 ++- tests/test_litellm/models/test_models.py | 9 ++ .../proxy/auth/test_auth_checks.py | 71 ++++++++++++++ .../test_upsert_budget_membership.py | 98 ++++++++++++++----- .../spend_tracking/test_budget_reservation.py | 49 ++++++++++ 8 files changed, 236 insertions(+), 60 deletions(-) diff --git a/litellm/models/budget.py b/litellm/models/budget.py index 2bef5d279d6..61123810fd1 100644 --- a/litellm/models/budget.py +++ b/litellm/models/budget.py @@ -36,19 +36,20 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase): model_config = ConfigDict(protected_namespaces=()) - def effective_max_budget(self, now: datetime) -> float | None: - if self.max_budget is None: - return None + def active_temp_budget_increase(self, now: datetime) -> float: if self.temp_budget_increase is None or self.temp_budget_expiry is None: - return self.max_budget + return 0.0 expiry: Final = ( self.temp_budget_expiry.replace(tzinfo=timezone.utc) if self.temp_budget_expiry.tzinfo is None else self.temp_budget_expiry ) - if expiry <= now: - return self.max_budget - return self.max_budget + self.temp_budget_increase + return 0.0 if expiry <= now else self.temp_budget_increase + + def effective_max_budget(self, now: datetime) -> float | None: + if self.max_budget is None: + return None + return self.max_budget + self.active_temp_budget_increase(now) class LiteLLM_BudgetTableFull(LiteLLM_BudgetTable): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index a36ef548e2f..4247bf49e71 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5325,12 +5325,10 @@ async def _check_team_member_budget( # Per-member override wins; otherwise fall back to the team-level # default configured via team.metadata["team_member_budget_id"]. team_member_budget: float | None = None - if ( - loaded_membership is not None - and loaded_membership.litellm_budget_table is not None - and loaded_membership.litellm_budget_table.max_budget is not None - ): - team_member_budget = loaded_membership.litellm_budget_table.effective_max_budget(now=get_utc_datetime()) + member_budget_row: Final = loaded_membership.litellm_budget_table if loaded_membership is not None else None + now: Final = get_utc_datetime() + if member_budget_row is not None and member_budget_row.max_budget is not None: + team_member_budget = member_budget_row.effective_max_budget(now=now) else: default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") if isinstance(default_budget_id, str): @@ -5346,7 +5344,9 @@ async def _check_team_member_budget( and default_budget.max_budget is not None and default_budget.max_budget > 0 ): - team_member_budget = default_budget.max_budget + team_member_budget = default_budget.max_budget + ( + member_budget_row.active_temp_budget_increase(now=now) if member_budget_row is not None else 0.0 + ) if team_member_budget is not None: team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0 diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index a5937ea8c7c..78e3ac7bd66 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -480,6 +480,8 @@ _TEAM_MEMBER_BUDGET_LIMIT_FIELDS: Final = ( "temp_budget_expiry", ) +_TEMP_BUDGET_FIELDS: Final = frozenset({"temp_budget_increase", "temp_budget_expiry"}) + MEMBER_BUDGET_PATCH_FIELDS: Final = MappingProxyType( { @@ -552,6 +554,8 @@ async def _upsert_budget_and_membership( ``shared_budget_ids`` extends that protection to any other row more than one membership points at, which a caller patching several members at once has already counted; a row listed there is cloned rather than written in place. + A patch that only touches the temporary budget pair never copies permanent + limits into a new row, so the member keeps inheriting the live team default. """ if not budget_patch: return @@ -566,6 +570,7 @@ async def _upsert_budget_and_membership( is_shared_default: Final = existing_budget_id is not None and ( existing_budget_id == team_default_budget_id or existing_budget_id in (shared_budget_ids or frozenset()) ) + temp_only: Final = frozenset(write_data) <= _TEMP_BUDGET_FIELDS async def _disconnect(): await tx.litellm_teammembership.update( @@ -586,30 +591,19 @@ async def _upsert_budget_and_membership( ) return - seeds_temp_budget: Final = "temp_budget_increase" in write_data or "temp_budget_expiry" in write_data - source_row_id: Final = ( - existing_budget_id - if is_shared_default - else team_default_budget_id - if team_default_budget_id is not None and seeds_temp_budget - else None - ) source_row: Final = ( - await tx.litellm_budgettable.find_unique(where={"budget_id": source_row_id}) - if source_row_id is not None + await tx.litellm_budgettable.find_unique(where={"budget_id": existing_budget_id}) + if is_shared_default and not temp_only else None ) source: Final[Mapping[str, Any]] = source_row.model_dump() if source_row is not None else MappingProxyType({}) - def _seeds(field: str) -> bool: - if field == "max_budget" and source.get(field) == 0 and not is_shared_default: - return False - return _is_set_budget_value(source.get(field)) - create_data: Final[dict[str, Any]] = { # mutable-ok: Prisma create payloads are dict-shaped "created_by": user_api_key_dict.user_id or "", "updated_by": user_api_key_dict.user_id or "", - **MappingProxyType({f: source[f] for f in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS if _seeds(f)}), + **MappingProxyType( + {f: source[f] for f in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS if _is_set_budget_value(source.get(f))} + ), **write_data, } @@ -621,7 +615,7 @@ async def _upsert_budget_and_membership( create_data.pop("budget_reset_at", None) if not _has_meaningful_budget_limit(create_data): - if existing_budget_id is not None: + if existing_budget_id is not None and not temp_only: await _disconnect() return diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 4028fcaf2ae..528f63064c6 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -688,16 +688,22 @@ async def _get_team_member_budget_counter( elif isinstance(cached_team_membership, dict): team_membership = LiteLLM_TeamMembership(**cached_team_membership) + member_budget_row: Final = team_membership.litellm_budget_table if team_membership is not None else None + now: Final = datetime.now(timezone.utc) team_member_budget: float | None = None - if team_membership is not None and team_membership.litellm_budget_table is not None: - team_member_budget = team_membership.litellm_budget_table.effective_max_budget(now=datetime.now(timezone.utc)) + if member_budget_row is not None and member_budget_row.max_budget is not None: + team_member_budget = member_budget_row.effective_max_budget(now=now) else: default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id") if isinstance(default_budget_id, str): default_budget: Final = await user_api_key_cache.async_get_cache( key=f"team_member_default_budget:{default_budget_id}", ) - team_member_budget = _to_float(_get_value(default_budget, "max_budget")) + default_cap: Final = _to_float(_get_value(default_budget, "max_budget")) + if default_cap is not None and default_cap > 0: + team_member_budget = default_cap + ( + member_budget_row.active_temp_budget_increase(now=now) if member_budget_row is not None else 0.0 + ) if team_member_budget is None or team_member_budget <= 0: return None diff --git a/tests/test_litellm/models/test_models.py b/tests/test_litellm/models/test_models.py index 4845dd8a7df..7774f6b543d 100644 --- a/tests/test_litellm/models/test_models.py +++ b/tests/test_litellm/models/test_models.py @@ -90,6 +90,15 @@ class TestBudget: assert LiteLLM_BudgetTable(max_budget=100.0).effective_max_budget(now=now) == 100.0 assert LiteLLM_BudgetTable(max_budget=None, temp_budget_increase=50.0).effective_max_budget(now=now) is None + def test_active_temp_budget_increase_is_independent_of_max_budget(self): + now = datetime(2026, 1, 1, tzinfo=timezone.utc) + bare = LiteLLM_BudgetTable(max_budget=None, temp_budget_increase=50.0, temp_budget_expiry=datetime(2100, 1, 1)) + assert bare.active_temp_budget_increase(now=now) == 50.0 + assert bare.effective_max_budget(now=now) is None + expired = LiteLLM_BudgetTable(max_budget=None, temp_budget_increase=50.0, temp_budget_expiry=now) + assert expired.active_temp_budget_increase(now=now) == 0.0 + assert LiteLLM_BudgetTable(max_budget=None).active_temp_budget_increase(now=now) == 0.0 + class TestCredentials: def test_credentials_creation(self): diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 3048d4a6a02..348973f2161 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -8548,3 +8548,74 @@ async def test_team_member_budget_check_temp_budget_increase_extends_cap(): ) assert exc_info.value.current_cost == 150.0 assert exc_info.value.max_budget == 100.0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "default_cap, expiry_offset, spend, expected_cap", + [ + (0.4, timedelta(hours=1), 1.0, None), + (0.4, timedelta(hours=-1), 1.0, 0.4), + (0.0, timedelta(hours=1), 1.0, None), + ], +) +async def test_team_member_budget_check_adds_temp_increase_to_live_team_default( + default_cap: float, expiry_offset: timedelta, spend: float, expected_cap: float | None +): + """A member row that carries only the temporary pair inherits the team default + cap live: the increase is added to it while active, the default alone applies + once it expires, and a zero default stays uncapped.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy._types import LiteLLM_TeamMembership + from litellm.proxy.utils import ProxyLogging + + cache = DualCache() + await cache.async_set_cache( + key="team_member_default_budget:default-budget-1", + value=LiteLLM_BudgetTable(budget_id="default-budget-1", max_budget=default_cap), + ) + team_object = LiteLLM_TeamTable(team_id="test-team", metadata={"team_member_budget_id": "default-budget-1"}) + valid_token = UserAPIKeyAuth(token="test-token", user_id="test-user", team_id="test-team") + team_membership = LiteLLM_TeamMembership( + user_id="test-user", + team_id="test-team", + spend=spend, + budget_id="budget-1", + litellm_budget_table=LiteLLM_BudgetTable( + max_budget=None, + temp_budget_increase=1.0, + temp_budget_expiry=datetime.now(timezone.utc) + expiry_offset, + ), + ) + + async def mock_get_current_spend(counter_key, fallback_spend, max_budget=None, **kwargs): + return fallback_spend + + with ( + patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend), # test-quality-ok: [TQ008] no seam on the cross-pod spend counter + patch( # test-quality-ok: [TQ008] isolates the check from the DB fetch + "litellm.proxy.auth.auth_checks.get_team_membership", + new_callable=AsyncMock, + return_value=team_membership, + ), + ): + if expected_cap is None: + await _check_team_member_budget( + team_object=team_object, + user_object=LiteLLM_UserTable(user_id="test-user"), + valid_token=valid_token, + prisma_client=MagicMock(), + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + ) + return + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_budget( + team_object=team_object, + user_object=LiteLLM_UserTable(user_id="test-user"), + valid_token=valid_token, + prisma_client=MagicMock(), + user_api_key_cache=cache, + proxy_logging_obj=ProxyLogging(user_api_key_cache=None), + ) + assert exc_info.value.max_budget == expected_cap diff --git a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py index cabf5e24933..08768f1636f 100644 --- a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py +++ b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py @@ -233,10 +233,10 @@ async def test_create_from_temp_budget_pair_only(mock_tx, fake_user): @pytest.mark.asyncio -async def test_create_from_temp_pair_keeps_team_default_cap(mock_tx, fake_user): +async def test_create_from_temp_pair_never_snapshots_team_default(mock_tx, fake_user): expiry = datetime(2100, 1, 1, tzinfo=timezone.utc) mock_tx.litellm_budgettable.find_unique = AsyncMock( - return_value=budget_row(budget_id="team-default-budget-1", max_budget=0.4, allowed_models=[]) + return_value=budget_row(budget_id="team-default-budget-1", max_budget=0.4, rpm_limit=10) ) await _upsert_budget_and_membership( mock_tx, @@ -248,34 +248,80 @@ async def test_create_from_temp_pair_keeps_team_default_cap(mock_tx, fake_user): team_default_budget_id="team-default-budget-1", ) + mock_tx.litellm_budgettable.find_unique.assert_not_awaited() + data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"] + assert "max_budget" not in data + assert "rpm_limit" not in data + assert data["temp_budget_increase"] == 1.0 + assert data["temp_budget_expiry"] == expiry + mock_tx.litellm_teammembership.upsert.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_temp_pair_on_shared_default_member_creates_bare_row(mock_tx, fake_user): + expiry = datetime(2100, 1, 1, tzinfo=timezone.utc) + mock_tx.litellm_budgettable.find_unique = AsyncMock( + return_value=budget_row(budget_id="team-default-budget-1", max_budget=0.4, rpm_limit=10) + ) + await _upsert_budget_and_membership( + mock_tx, + team_id="team-default", + user_id="user-on-default", + existing_budget_id="team-default-budget-1", + user_api_key_dict=fake_user, + budget_patch={"temp_budget_increase": 1.0, "temp_budget_expiry": expiry}, + team_default_budget_id="team-default-budget-1", + ) + + mock_tx.litellm_budgettable.find_unique.assert_not_awaited() + mock_tx.litellm_budgettable.update.assert_not_called() + data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"] + assert "max_budget" not in data + assert "rpm_limit" not in data + assert data["temp_budget_increase"] == 1.0 + assert data["temp_budget_expiry"] == expiry + mock_tx.litellm_teammembership.upsert.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_clearing_temp_pair_on_shared_default_member_is_noop(mock_tx, fake_user): + await _upsert_budget_and_membership( + mock_tx, + team_id="team-default", + user_id="user-on-default", + existing_budget_id="team-default-budget-1", + user_api_key_dict=fake_user, + budget_patch={"temp_budget_increase": None, "temp_budget_expiry": None}, + team_default_budget_id="team-default-budget-1", + ) + + mock_tx.litellm_budgettable.create.assert_not_called() + mock_tx.litellm_budgettable.update.assert_not_called() + mock_tx.litellm_teammembership.update.assert_not_called() + mock_tx.litellm_teammembership.upsert.assert_not_called() + + +@pytest.mark.asyncio +async def test_temp_pair_with_permanent_field_still_clones_shared_default(mock_tx, fake_user): + expiry = datetime(2100, 1, 1, tzinfo=timezone.utc) + mock_tx.litellm_budgettable.find_unique = AsyncMock( + return_value=budget_row(budget_id="team-default-budget-1", max_budget=0.4, rpm_limit=10) + ) + await _upsert_budget_and_membership( + mock_tx, + team_id="team-default", + user_id="user-on-default", + existing_budget_id="team-default-budget-1", + user_api_key_dict=fake_user, + budget_patch={"temp_budget_increase": 1.0, "temp_budget_expiry": expiry, "tpm_limit": 500}, + team_default_budget_id="team-default-budget-1", + ) + mock_tx.litellm_budgettable.find_unique.assert_awaited_once_with(where={"budget_id": "team-default-budget-1"}) data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"] assert data["max_budget"] == 0.4 - assert data["temp_budget_increase"] == 1.0 - assert data["temp_budget_expiry"] == expiry - assert "allowed_models" not in data - mock_tx.litellm_teammembership.upsert.assert_awaited_once() - - -@pytest.mark.asyncio -async def test_create_from_temp_pair_skips_zero_team_default_cap(mock_tx, fake_user): - expiry = datetime(2100, 1, 1, tzinfo=timezone.utc) - mock_tx.litellm_budgettable.find_unique = AsyncMock( - return_value=budget_row(budget_id="team-default-budget-1", max_budget=0, rpm_limit=10) - ) - await _upsert_budget_and_membership( - mock_tx, - team_id="team-default", - user_id="user-unlinked", - existing_budget_id=None, - user_api_key_dict=fake_user, - budget_patch={"temp_budget_increase": 1.0, "temp_budget_expiry": expiry}, - team_default_budget_id="team-default-budget-1", - ) - - data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"] - assert "max_budget" not in data assert data["rpm_limit"] == 10 + assert data["tpm_limit"] == 500 assert data["temp_budget_increase"] == 1.0 mock_tx.litellm_teammembership.upsert.assert_awaited_once() diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index 05cad26ce10..0e0025f5194 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -497,3 +497,52 @@ async def test_team_member_reservation_counter_honours_temp_budget_increase( assert counter is not None assert counter.max_budget == expected_max_budget assert counter.fallback_spend == 0.5 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "default_cap, expiry_offset, expected_max_budget", + [ + (2.0, timedelta(days=1), 3.0), + (2.0, timedelta(days=-1), 2.0), + (0.0, timedelta(days=1), None), + ], +) +async def test_team_member_reservation_counter_adds_temp_increase_to_live_team_default( + default_cap: float, expiry_offset: timedelta, expected_max_budget: float | None +) -> None: + user_id: Final = "member-bare" + team_id: Final = "team-bare" + cache: Final = UserApiKeyCache() + await cache.async_set_cache( + key="team_member_default_budget:default-bare", + value=LiteLLM_BudgetTable(budget_id="default-bare", max_budget=default_cap), + ) + await cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), + value=LiteLLM_TeamMembership( + user_id=user_id, + team_id=team_id, + spend=0.5, + budget_id="budget-bare", + litellm_budget_table=LiteLLM_BudgetTable( + max_budget=None, + temp_budget_increase=1.0, + temp_budget_expiry=datetime.now(timezone.utc) + expiry_offset, + ), + ), + ) + + counter: Final = await _get_team_member_budget_counter( + valid_token=UserAPIKeyAuth(token="hashed", user_id=user_id, team_id=team_id), + team_object=LiteLLM_TeamTable(team_id=team_id, metadata={"team_member_budget_id": "default-bare"}), + user_object=LiteLLM_UserTable(user_id=user_id), + user_api_key_cache=cache, + ) + + if expected_max_budget is None: + assert counter is None + return + assert counter is not None + assert counter.max_budget == expected_max_budget + assert counter.fallback_spend == 0.5 From 504b464417f17bc3f5b05cddc758c6f5da040022 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 02:12:12 +0000 Subject: [PATCH 20/20] test(proxy): assert the persisted member budget row as a whole instead of field absence Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_upsert_budget_membership.py | 28 +++++++------------ 1 file changed, 10 insertions(+), 18 deletions(-) diff --git a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py index 08768f1636f..7a75ec395f1 100644 --- a/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py +++ b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py @@ -57,6 +57,12 @@ def assert_future_reset_time(value): assert value > datetime.now(timezone.utc) +def stored_budget_row(mock_tx): + """The budget row the create call persists, minus the audit columns.""" + data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"] + return {k: v for k, v in data.items() if k not in ("created_by", "updated_by")} + + # TEST: an empty patch (caller sent no budget fields) leaves everything alone. # This is the merge-patch contract: absent != clear. Updating only a member's # role must not silently wipe their budget. @@ -224,10 +230,7 @@ async def test_create_from_temp_budget_pair_only(mock_tx, fake_user): ) mock_tx.litellm_budgettable.create.assert_awaited_once() - data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"] - assert data["temp_budget_increase"] == 5.0 - assert data["temp_budget_expiry"] == expiry - assert "max_budget" not in data + assert stored_budget_row(mock_tx) == {"temp_budget_increase": 5.0, "temp_budget_expiry": expiry} mock_tx.litellm_teammembership.upsert.assert_awaited_once() mock_tx.litellm_teammembership.update.assert_not_called() @@ -249,11 +252,7 @@ async def test_create_from_temp_pair_never_snapshots_team_default(mock_tx, fake_ ) mock_tx.litellm_budgettable.find_unique.assert_not_awaited() - data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"] - assert "max_budget" not in data - assert "rpm_limit" not in data - assert data["temp_budget_increase"] == 1.0 - assert data["temp_budget_expiry"] == expiry + assert stored_budget_row(mock_tx) == {"temp_budget_increase": 1.0, "temp_budget_expiry": expiry} mock_tx.litellm_teammembership.upsert.assert_awaited_once() @@ -275,11 +274,7 @@ async def test_temp_pair_on_shared_default_member_creates_bare_row(mock_tx, fake mock_tx.litellm_budgettable.find_unique.assert_not_awaited() mock_tx.litellm_budgettable.update.assert_not_called() - data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"] - assert "max_budget" not in data - assert "rpm_limit" not in data - assert data["temp_budget_increase"] == 1.0 - assert data["temp_budget_expiry"] == expiry + assert stored_budget_row(mock_tx) == {"temp_budget_increase": 1.0, "temp_budget_expiry": expiry} mock_tx.litellm_teammembership.upsert.assert_awaited_once() @@ -342,10 +337,7 @@ async def test_create_from_plain_patch_does_not_snapshot_team_default(mock_tx, f ) mock_tx.litellm_budgettable.find_unique.assert_not_awaited() - data = mock_tx.litellm_budgettable.create.await_args.kwargs["data"] - assert data["tpm_limit"] == 500 - assert "max_budget" not in data - assert "rpm_limit" not in data + assert stored_budget_row(mock_tx) == {"tpm_limit": 500} mock_tx.litellm_teammembership.upsert.assert_awaited_once()