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