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 1894518e51d..cc7793d4bb2 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..61123810fd1 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 @@ -30,9 +31,26 @@ 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=()) + def active_temp_budget_increase(self, now: datetime) -> float: + if self.temp_budget_increase is None or self.temp_budget_expiry is None: + 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 + ) + 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): """LiteLLM_BudgetTable + server-managed fields returned on API responses.""" diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e3ca88147d7..d2a7daf99ca 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4412,6 +4412,23 @@ 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, + 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( + 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): @@ -4421,6 +4438,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 84aa5d16639..cdada970956 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5395,12 +5395,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.max_budget + 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): @@ -5416,7 +5414,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/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index bf24ed19113..7315219dd47 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2250,7 +2250,9 @@ 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 = 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: # Read from cross-pod counter (Redis-first) if available from litellm.proxy.proxy_server import get_current_spend diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 29f24d2465f..78e3ac7bd66 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -476,8 +476,12 @@ _TEAM_MEMBER_BUDGET_LIMIT_FIELDS: Final = ( "model_max_budget", "budget_duration", "allowed_models", + "temp_budget_increase", + "temp_budget_expiry", ) +_TEMP_BUDGET_FIELDS: Final = frozenset({"temp_budget_increase", "temp_budget_expiry"}) + MEMBER_BUDGET_PATCH_FIELDS: Final = MappingProxyType( { @@ -486,6 +490,8 @@ MEMBER_BUDGET_PATCH_FIELDS: Final = MappingProxyType( "rpm_limit": "rpm_limit", "budget_duration": "budget_duration", "allowed_models": "allowed_models", + "temp_budget_increase": "temp_budget_increase", + "temp_budget_expiry": "temp_budget_expiry", } ) @@ -548,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 @@ -562,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( @@ -583,7 +592,9 @@ async def _upsert_budget_and_membership( return source_row: Final = ( - await tx.litellm_budgettable.find_unique(where={"budget_id": existing_budget_id}) if is_shared_default else 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({}) @@ -604,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/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 29c31cc9c18..24bbd2b4b1f 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -306,6 +306,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]: @@ -391,6 +394,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]* - 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 @@ -527,7 +532,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/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index b37d7970e50..4957b3bd925 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -3848,6 +3848,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 1894518e51d..cc7793d4bb2 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/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 373f2d0fe36..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.max_budget + 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/schema.prisma b/schema.prisma index 1894518e51d..cc7793d4bb2 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/models/test_models.py b/tests/test_litellm/models/test_models.py index 9b803c14062..7774f6b543d 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,34 @@ 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): + expiry = datetime(2020, 1, 1, tzinfo=timezone.utc) + budget = LiteLLM_BudgetTable(max_budget=100.0, temp_budget_increase=50.0, temp_budget_expiry=expiry) + assert budget.effective_max_budget(now=datetime(2026, 1, 1, tzinfo=timezone.utc)) == 100.0 + assert budget.effective_max_budget(now=expiry) == 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 + + 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 51a53d5d41e..79acf37eeff 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -8686,3 +8686,161 @@ 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 + + +@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 + + 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, + ), + ): + 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), # 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, + ), + ): + 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 + + +@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/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 674eec9738c..e6975edd4cb 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 @@ -7754,6 +7754,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/common_utils/test_upsert_budget_membership.py b/tests/test_litellm/proxy/common_utils/test_upsert_budget_membership.py index 2cc0d9f74f5..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. @@ -211,6 +217,130 @@ async def test_create_seeds_reset_at_and_links(mock_tx, fake_user): ) +@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() + 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() + + +@pytest.mark.asyncio +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, 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", + ) + + mock_tx.litellm_budgettable.find_unique.assert_not_awaited() + assert stored_budget_row(mock_tx) == {"temp_budget_increase": 1.0, "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() + assert stored_budget_row(mock_tx) == {"temp_budget_increase": 1.0, "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["rpm_limit"] == 10 + assert data["tpm_limit"] == 500 + assert data["temp_budget_increase"] == 1.0 + 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() + assert stored_budget_row(mock_tx) == {"tpm_limit": 500} + 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. 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", }, 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 47acc091672..3c6afa86c45 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py @@ -1352,6 +1352,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 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 ec7c989e342..4f5c5066367 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -15571,3 +15571,37 @@ 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_member_budget_patch_maps_temp_budget_fields() -> None: + from litellm.proxy.management_endpoints.common_utils import 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 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") + + +@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=increase, temp_budget_expiry="2030-01-01T00:00:00Z" + ) 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..0e0025f5194 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,93 @@ 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 + + +@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 diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 6ce30148f1c..8a4fd586afc 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -2821,6 +2821,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) => { @@ -2955,6 +2957,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 434893c806c..6eb42d7f54f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -10743,6 +10743,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]* - 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 @@ -29617,6 +29619,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 */ @@ -29652,6 +29658,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 */ @@ -33465,6 +33475,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 */ @@ -33588,6 +33602,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 */ @@ -38371,6 +38389,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 @@ -38393,6 +38421,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 */ @@ -39561,6 +39593,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 */