diff --git a/litellm/proxy/db/create_views.py b/litellm/proxy/db/create_views.py index 141ce92f172..5ea9cba8018 100644 --- a/litellm/proxy/db/create_views.py +++ b/litellm/proxy/db/create_views.py @@ -1,15 +1,50 @@ -from typing import Any, Final +from typing import Any, Final, Protocol from litellm import verbose_logger _db = Any + +class SupportsExecuteRaw(Protocol): + """The one database operation create_view_tolerating_race needs. + + Narrower than the `_db = Any` the rest of this module still uses, so the + helper's contract is checkable at its call sites without retyping every + function here. + """ + + async def execute_raw(self, query: str, *args: object) -> int: ... + + # Markers that indicate a view/relation does not yet exist in the database. # Keeping these in one place avoids repeating the check across all view blocks # and prevents overly broad matches (e.g. bare 'undefined' would also match # 'undefined function' or 'column undefined_col referenced in query'). _VIEW_NOT_FOUND_MARKERS: Final = ("does not exist", "no such table", "undefined table") +# Markers for the inverse condition: another replica created the view between +# our existence probe and our CREATE. +_VIEW_ALREADY_EXISTS_MARKERS: Final = ("already exists", "duplicate object", "duplicate table") + + +async def create_view_tolerating_race(db: SupportsExecuteRaw, view_name: str, ddl: str) -> None: + """ + Create a view, treating "a concurrent creator won" as success. + + Every replica booting against the same fresh database observes the view as + absent and issues the CREATE; Postgres fails all but one with a + duplicate-object error. The desired end state is still reached, so losing + that race is success. Without this, the loser's exception propagates out of + a detached startup task and the remaining views are never created. + """ + try: + await db.execute_raw(ddl) + verbose_logger.debug("%s Created!", view_name) + except Exception as e: + if not any(marker in str(e).lower() for marker in _VIEW_ALREADY_EXISTS_MARKERS): + raise + verbose_logger.debug("%s already created by a concurrent replica", view_name) + async def create_missing_views(db: _db): """ @@ -34,7 +69,10 @@ async def create_missing_views(db: _db): if not any(marker in error_msg for marker in _VIEW_NOT_FOUND_MARKERS): raise # If an error occurs, the view does not exist, so create it - await db.execute_raw(""" + await create_view_tolerating_race( + db, + "LiteLLM_VerificationTokenView", + """ CREATE VIEW "LiteLLM_VerificationTokenView" AS SELECT v.*, @@ -46,9 +84,8 @@ async def create_missing_views(db: _db): FROM "LiteLLM_VerificationToken" v LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id LEFT JOIN "LiteLLM_ProjectTable" p ON v.project_id = p.project_id; - """) - - verbose_logger.debug("LiteLLM_VerificationTokenView Created!") + """, + ) try: await db.query_raw("""SELECT 1 FROM "MonthlyGlobalSpend" LIMIT 1""") @@ -69,9 +106,7 @@ async def create_missing_views(db: _db): GROUP BY DATE("startTime"); """ - await db.execute_raw(query=sql_query) - - verbose_logger.debug("MonthlyGlobalSpend Created!") + await create_view_tolerating_race(db, "MonthlyGlobalSpend", sql_query) try: await db.query_raw("""SELECT 1 FROM "Last30dKeysBySpend" LIMIT 1""") @@ -100,9 +135,7 @@ async def create_missing_views(db: _db): ORDER BY total_spend DESC; """ - await db.execute_raw(query=sql_query) - - verbose_logger.debug("Last30dKeysBySpend Created!") + await create_view_tolerating_race(db, "Last30dKeysBySpend", sql_query) try: await db.query_raw("""SELECT 1 FROM "Last30dModelsBySpend" LIMIT 1""") @@ -126,9 +159,7 @@ async def create_missing_views(db: _db): ORDER BY total_spend DESC; """ - await db.execute_raw(query=sql_query) - - verbose_logger.debug("Last30dModelsBySpend Created!") + await create_view_tolerating_race(db, "Last30dModelsBySpend", sql_query) try: await db.query_raw("""SELECT 1 FROM "MonthlyGlobalSpendPerKey" LIMIT 1""") verbose_logger.debug("MonthlyGlobalSpendPerKey Exists!") @@ -150,9 +181,7 @@ async def create_missing_views(db: _db): DATE("startTime"), api_key; """ - await db.execute_raw(query=sql_query) - - verbose_logger.debug("MonthlyGlobalSpendPerKey Created!") + await create_view_tolerating_race(db, "MonthlyGlobalSpendPerKey", sql_query) try: await db.query_raw("""SELECT 1 FROM "MonthlyGlobalSpendPerUserPerKey" LIMIT 1""") verbose_logger.debug("MonthlyGlobalSpendPerUserPerKey Exists!") @@ -176,9 +205,7 @@ async def create_missing_views(db: _db): "user", api_key; """ - await db.execute_raw(query=sql_query) - - verbose_logger.debug("MonthlyGlobalSpendPerUserPerKey Created!") + await create_view_tolerating_race(db, "MonthlyGlobalSpendPerUserPerKey", sql_query) try: await db.query_raw("""SELECT 1 FROM "DailyTagSpend" LIMIT 1""") @@ -197,9 +224,7 @@ async def create_missing_views(db: _db): FROM "LiteLLM_SpendLogs" s GROUP BY individual_request_tag, DATE(s."startTime"); """ - await db.execute_raw(query=sql_query) - - verbose_logger.debug("DailyTagSpend Created!") + await create_view_tolerating_race(db, "DailyTagSpend", sql_query) try: await db.query_raw("""SELECT 1 FROM "Last30dTopEndUsersSpend" LIMIT 1""") @@ -218,9 +243,7 @@ async def create_missing_views(db: _db): ORDER BY total_spend DESC LIMIT 100; """ - await db.execute_raw(query=sql_query) - - verbose_logger.debug("Last30dTopEndUsersSpend Created!") + await create_view_tolerating_race(db, "Last30dTopEndUsersSpend", sql_query) async def should_create_missing_views(db: _db) -> bool: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2a385c4c42a..116dea464ff 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -888,17 +888,24 @@ async def _common_key_generation_helper( if litellm.default_key_generate_params is not None: for elem in data: key, value = elem - if value is None and key in [ - "max_budget", - "user_id", - "team_id", - "max_parallel_requests", - "tpm_limit", - "rpm_limit", - "budget_duration", - "duration", - ]: - setattr(data, key, litellm.default_key_generate_params.get(key, None)) + if ( + value is None + and (key != "budget_duration" or key not in data.model_fields_set) + and key + in [ + "max_budget", + "user_id", + "team_id", + "max_parallel_requests", + "tpm_limit", + "rpm_limit", + "budget_duration", + "duration", + ] + ): + default_value = litellm.default_key_generate_params.get(key) + if default_value is not None: + setattr(data, key, default_value) elif key == "models" and value == []: setattr(data, key, litellm.default_key_generate_params.get(key, [])) elif key == "metadata" and value == {}: diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 60d3d650d00..ae1bb52278a 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1313,8 +1313,9 @@ async def new_team( if isinstance(default_organization_id, str): data.organization_id = default_organization_id - # Apply defaults from litellm.default_team_params for any fields - # not explicitly provided in the request. + # Apply defaults from litellm.default_team_params to null fields. + # budget_duration alone distinguishes explicit null (a deliberate + # never-resetting budget, which the default must not override) from omitted. for field in ( "max_budget", "budget_duration", @@ -1322,7 +1323,9 @@ async def new_team( "rpm_limit", "team_member_permissions", ): - if getattr(data, field, None) is None: + if getattr(data, field, None) is None and ( + field != "budget_duration" or field not in data.model_fields_set + ): default_value = _get_default_team_param(field) if default_value is not None: setattr(data, field, default_value) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index b3feb5bd8d6..5584dae9e15 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -666,7 +666,7 @@ async def get_internal_user_settings(): ) async def get_default_team_settings(): """ - Get all SSO settings from the litellm_settings configuration. + Get the default team parameters (litellm_settings.default_team_params). Returns a structured object with values and descriptions for UI display. """ from litellm.proxy.proxy_server import proxy_config @@ -894,8 +894,9 @@ async def update_default_team_settings( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ - Update the default team parameters for SSO users. - These settings will be applied to new teams created from SSO. + Update the default team parameters (litellm_settings.default_team_params). + Applied to every new team for fields not explicitly provided in the create request; + `models` only applies to teams automatically created via SSO Groups. """ if settings.organization_id is not None: await _validate_default_organization_exists(settings.organization_id) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8a1fae42789..ef0376ade84 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -106,6 +106,7 @@ from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_c from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.create_views import ( create_missing_views, + create_view_tolerating_race, should_create_missing_views, ) from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter @@ -3273,7 +3274,10 @@ class PrismaClient: ## check if required view exists ## if ret[0]["view_names"] and required_view not in ret[0]["view_names"]: await self.health_check() # make sure we can connect to db - await self.db.execute_raw(""" + await create_view_tolerating_race( + self.db, + "LiteLLM_VerificationTokenView", + """ CREATE VIEW "LiteLLM_VerificationTokenView" AS SELECT v.*, @@ -3283,9 +3287,8 @@ class PrismaClient: t.rpm_limit AS team_rpm_limit FROM "LiteLLM_VerificationToken" v LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id; - """) - - verbose_proxy_logger.info("LiteLLM_VerificationTokenView Created in DB!") + """, + ) else: should_create_views: Final = await should_create_missing_views(db=self.db) if should_create_views: diff --git a/litellm/types/proxy/management_endpoints/ui_sso.py b/litellm/types/proxy/management_endpoints/ui_sso.py index 0585057c22e..b4691b9b08c 100644 --- a/litellm/types/proxy/management_endpoints/ui_sso.py +++ b/litellm/types/proxy/management_endpoints/ui_sso.py @@ -202,28 +202,29 @@ class SSOConfig(LiteLLMPydanticObjectBase): class DefaultTeamSSOParams(LiteLLMPydanticObjectBase): """ - Default parameters to apply when a new team is automatically created by LiteLLM via SSO Groups + Default parameters applied to every /team/new call for fields not explicitly provided in the request. + `models` is the exception: it only applies to teams automatically created by LiteLLM via SSO Groups. """ models: list[str] = Field( default=[], - description="Default list of models that new automatically created teams can access", + description="Default list of models for teams automatically created via SSO Groups", ) max_budget: float | None = Field( default=None, - description="Default maximum budget (in USD) for new automatically created teams", + description="Default maximum budget (in USD) for new teams, when not explicitly provided", ) budget_duration: str | None = Field( default=None, - description="Default budget duration for new automatically created teams (e.g. 'daily', 'weekly', 'monthly')", + description="Default budget duration for new teams, when not explicitly provided (e.g. '24h', '7d', '30d')", ) tpm_limit: int | None = Field( default=None, - description="Default tpm limit for new automatically created teams", + description="Default tpm limit for new teams, when not explicitly provided", ) rpm_limit: int | None = Field( default=None, - description="Default rpm limit for new automatically created teams", + description="Default rpm limit for new teams, when not explicitly provided", ) team_member_permissions: list[KeyManagementRoutes] | None = Field( default=None, diff --git a/tests/test_litellm/proxy/db/test_create_views.py b/tests/test_litellm/proxy/db/test_create_views.py index c0c09d0137b..ecc6d70123e 100644 --- a/tests/test_litellm/proxy/db/test_create_views.py +++ b/tests/test_litellm/proxy/db/test_create_views.py @@ -189,3 +189,66 @@ async def test_create_views_creates_view_on_undefined_table_error(): await create_missing_views(mock_db) mock_db.execute_raw.assert_called_once() + + +# Every view create_missing_views is responsible for. Hard-coded rather than +# derived from the module, so adding a view without guarding it fails here. +EXPECTED_VIEW_COUNT = 8 + + +@pytest.mark.asyncio +async def test_create_views_tolerates_a_concurrent_creator_on_every_view(): + """A replica that loses the CREATE race must attempt every view regardless. + + Regression: two proxy pods booting on a fresh DB both see every view as + absent and both issue the CREATE, and Postgres fails the loser with a + duplicate-object error on whichever views the winner got to first. Any + creation site still calling execute_raw unguarded re-raises that error and + aborts the rest of the function. + + Every CREATE loses here, which is what pins the guard to all of them: an + earlier version of this fix converted only the first and the last site and + still died on MonthlyGlobalSpend against a real Postgres. Counting the + attempts is the assertion, because a partial fix simply stops early. + """ + from litellm.proxy.db.create_views import create_missing_views + + mock_db = MagicMock() + mock_db.query_raw = AsyncMock(side_effect=Exception("relation does not exist")) + mock_db.execute_raw = AsyncMock( + side_effect=Exception('relation "some_view" already exists') + ) + + await create_missing_views(mock_db) + + assert mock_db.execute_raw.await_count == EXPECTED_VIEW_COUNT, ( + f"every view must still be attempted when the replica loses every race; " + f"got {mock_db.execute_raw.await_count} of {EXPECTED_VIEW_COUNT}, so a " + f"creation site is still unguarded and aborted the rest" + ) + + +@pytest.mark.asyncio +async def test_create_views_reraises_genuine_ddl_error(): + """An already-exists guard must not swallow real DDL failures.""" + from litellm.proxy.db.create_views import create_missing_views + + mock_db = MagicMock() + mock_db.query_raw = AsyncMock(side_effect=Exception("relation does not exist")) + mock_db.execute_raw = AsyncMock(side_effect=Exception("syntax error at or near")) + + with pytest.raises(Exception, match="syntax error"): + await create_missing_views(mock_db) + + +@pytest.mark.asyncio +async def test_create_view_tolerating_race_swallows_only_already_exists(): + from litellm.proxy.db.create_views import create_view_tolerating_race + + mock_db = MagicMock() + mock_db.execute_raw = AsyncMock(side_effect=Exception("duplicate object")) + await create_view_tolerating_race(mock_db, "SomeView", "CREATE VIEW ...") + + mock_db.execute_raw = AsyncMock(side_effect=Exception("permission denied")) + with pytest.raises(Exception, match="permission denied"): + await create_view_tolerating_race(mock_db, "SomeView", "CREATE VIEW ...") diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 939607dd139..0ebdc07d282 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -15600,3 +15600,106 @@ async def test_unblock_key_stamps_settings_updated_at(monkeypatch): sent = mock_prisma_client.db.litellm_verificationtoken.update.call_args.kwargs["data"] assert sent["blocked"] is False assert before <= sent["settings_updated_at"] <= after + + +def _wire_key_generation_prisma(monkeypatch): + created_key = MagicMock(token="hashed_token_123", litellm_budget_table=None, object_permission=None) + + mock_prisma_client = AsyncMock() + mock_prisma_client.insert_data = AsyncMock(return_value=created_key) + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.litellm_verificationtoken = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(return_value=created_key) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + return mock_prisma_client.insert_data + + +async def _generate_key_and_get_persisted_row(data: GenerateKeyRequest, mock_insert_data): + await _common_key_generation_helper( + data=data, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + litellm_changed_by=None, + team_table=None, + ) + key_call = next(c for c in mock_insert_data.call_args_list if c.kwargs["table_name"] == "key") + return key_call.kwargs["data"] + + +@pytest.mark.asyncio +async def test_key_generate_explicit_null_budget_duration_beats_default_key_generate_params(monkeypatch): + """An explicit `"budget_duration": null` asks for a budget that never resets. + + Gating on the value alone made that indistinguishable from omitting the field, + so the configured default overrode the opt-out and budget_reset_at got stamped. + """ + monkeypatch.setattr(litellm, "default_key_generate_params", {"budget_duration": "30d"}) + mock_insert_data = _wire_key_generation_prisma(monkeypatch) + + key_row = await _generate_key_and_get_persisted_row(GenerateKeyRequest(budget_duration=None), mock_insert_data) + + assert key_row["budget_duration"] is None + assert key_row["budget_reset_at"] is None + + +@pytest.mark.asyncio +async def test_key_generate_omitted_budget_duration_still_takes_default_key_generate_params(monkeypatch): + """Omitting the field keeps applying the default, the behavior the explicit-null fix must not break.""" + monkeypatch.setattr(litellm, "default_key_generate_params", {"budget_duration": "30d"}) + mock_insert_data = _wire_key_generation_prisma(monkeypatch) + + key_row = await _generate_key_and_get_persisted_row(GenerateKeyRequest(), mock_insert_data) + + assert key_row["budget_duration"] == "30d" + assert key_row["budget_reset_at"] is not None + + +@pytest.mark.asyncio +async def test_key_generate_explicit_null_budget_duration_cannot_bypass_upperbound(monkeypatch): + """upperbound_key_generate_params is an admin ceiling: an explicit null must not mint an uncapped key, + otherwise any key creator could bypass configured limits (duration, budgets, rate limits).""" + from litellm.types.proxy.management_endpoints.ui_sso import ( + LiteLLM_UpperboundKeyGenerateParams, + ) + + monkeypatch.setattr(litellm, "default_key_generate_params", None) + monkeypatch.setattr( + litellm, + "upperbound_key_generate_params", + LiteLLM_UpperboundKeyGenerateParams(budget_duration="30d"), + ) + mock_insert_data = _wire_key_generation_prisma(monkeypatch) + + key_row = await _generate_key_and_get_persisted_row(GenerateKeyRequest(budget_duration=None), mock_insert_data) + + assert key_row["budget_duration"] == "30d" + assert key_row["budget_reset_at"] is not None + + +@pytest.mark.asyncio +async def test_key_generate_omitted_budget_duration_still_filled_by_upperbound(monkeypatch): + """The upperbound's long-standing fill-on-omitted behavior stays untouched.""" + from litellm.types.proxy.management_endpoints.ui_sso import ( + LiteLLM_UpperboundKeyGenerateParams, + ) + + monkeypatch.setattr(litellm, "default_key_generate_params", None) + monkeypatch.setattr( + litellm, + "upperbound_key_generate_params", + LiteLLM_UpperboundKeyGenerateParams(budget_duration="30d"), + ) + mock_insert_data = _wire_key_generation_prisma(monkeypatch) + + key_row = await _generate_key_and_get_persisted_row(GenerateKeyRequest(), mock_insert_data) + + assert key_row["budget_duration"] == "30d" + assert key_row["budget_reset_at"] is not None 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 073f1ba782e..5190df7521e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -11094,3 +11094,104 @@ async def test_new_team_output_token_estimate_rejected_for_non_admin(): assert str(exc.value.code) == "403" assert "on a team" in str(exc.value.message) + + +def _wire_new_team_prisma(mock_db_client): + mock_db_client.jsonify_team_object = lambda db_data: db_data + mock_db_client.get_data = AsyncMock(return_value=None) + mock_db_client.db = MagicMock() + + created_team = MagicMock(team_id="team-defaults") + created_team.model_dump.return_value = {"team_id": "team-defaults"} + + mock_db_client.db.litellm_teamtable = MagicMock() + mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_db_client.db.litellm_teamtable.create = AsyncMock(return_value=created_team) + mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=created_team) + mock_db_client.db.litellm_usertable = MagicMock() + mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) + + return mock_db_client.db.litellm_teamtable.create + + +@pytest.mark.asyncio +async def test_new_team_explicit_null_budget_duration_beats_configured_default( + mock_db_client, mock_admin_auth, monkeypatch +): + """An explicit `"budget_duration": null` asks for a lifetime budget that never resets. + + Gating on the value alone made that indistinguishable from omitting the field, + so the default overrode the opt-out and budget_reset_at got stamped. + """ + from fastapi import Request + + import litellm + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + monkeypatch.setattr(litellm, "default_team_settings", None) + monkeypatch.setattr(litellm, "default_team_params", {"budget_duration": "30d"}) + mock_team_create = _wire_new_team_prisma(mock_db_client) + + await new_team( + data=NewTeamRequest(team_alias="lifetime-budget-team", budget_duration=None), + http_request=MagicMock(spec=Request), + user_api_key_dict=mock_admin_auth, + ) + + team_data = mock_team_create.call_args.kwargs["data"] + assert team_data.get("budget_duration") is None + assert team_data.get("budget_reset_at") is None + + +@pytest.mark.asyncio +async def test_new_team_omitted_budget_duration_still_takes_configured_default( + mock_db_client, mock_admin_auth, monkeypatch +): + """Omitting the field keeps applying the default, the behavior the explicit-null fix must not break.""" + from fastapi import Request + + import litellm + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + monkeypatch.setattr(litellm, "default_team_settings", None) + monkeypatch.setattr(litellm, "default_team_params", {"budget_duration": "30d"}) + mock_team_create = _wire_new_team_prisma(mock_db_client) + + await new_team( + data=NewTeamRequest(team_alias="default-budget-team"), + http_request=MagicMock(spec=Request), + user_api_key_dict=mock_admin_auth, + ) + + team_data = mock_team_create.call_args.kwargs["data"] + assert team_data.get("budget_duration") == "30d" + assert team_data.get("budget_reset_at") is not None + + +@pytest.mark.asyncio +async def test_new_team_explicit_null_max_budget_still_takes_configured_default( + mock_db_client, mock_admin_auth, monkeypatch +): + """The explicit-null opt-out is budget_duration-only: nulling limit fields + (max_budget, tpm/rpm) must not skip configured defaults, or any team creator + could mint uncapped teams (veria finding on PR #36699).""" + from fastapi import Request + + import litellm + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + monkeypatch.setattr(litellm, "default_team_settings", None) + monkeypatch.setattr(litellm, "default_team_params", {"max_budget": 100.0}) + mock_team_create = _wire_new_team_prisma(mock_db_client) + + await new_team( + data=NewTeamRequest(team_alias="unlimited-budget-team", max_budget=None), + http_request=MagicMock(spec=Request), + user_api_key_dict=mock_admin_auth, + ) + + team_data = mock_team_create.call_args.kwargs["data"] + assert team_data.get("max_budget") == 100.0 diff --git a/ui/litellm-dashboard/src/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx index c5c77c0c878..79d994140e9 100644 --- a/ui/litellm-dashboard/src/components/Teams.test.tsx +++ b/ui/litellm-dashboard/src/components/Teams.test.tsx @@ -1,12 +1,19 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { act, fireEvent, render, screen, waitFor, within } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; import { NuqsTestingAdapter, OnUrlUpdateFunction } from "nuqs/adapters/testing"; import React from "react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { useTeamMetadataSchema } from "@/app/(dashboard)/hooks/teams/useTeamMetadataSchema"; import NotificationsManager from "./molecules/notifications_manager"; import { fetchAvailableModelsForTeamOrKey } from "./key_team_helpers/fetch_available_models_team_key"; -import { fetchMCPAccessGroups, getGuardrailsList, getPoliciesList, teamCreateCall } from "./networking"; +import { + fetchMCPAccessGroups, + getDefaultTeamSettings, + getGuardrailsList, + getPoliciesList, + teamCreateCall, +} from "./networking"; import Teams from "./Teams"; const can = vi.fn(); @@ -34,6 +41,7 @@ vi.mock("./networking", () => ({ v2TeamListCall: vi.fn(), getGuardrailsList: vi.fn().mockResolvedValue({ guardrails: [] }), getPoliciesList: vi.fn().mockResolvedValue({ policies: [] }), + getDefaultTeamSettings: vi.fn().mockResolvedValue({ values: {} }), })); // Teams invalidates teamsTableKeys on mutations; the selected team is passed up from the table. @@ -649,6 +657,105 @@ describe("Teams - access_group_ids in team create", () => { }); }); +describe("Teams - Reset Budget in team create", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockTeamInfoView.mockClear(); + vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4"]); + vi.mocked(fetchMCPAccessGroups).mockResolvedValue([]); + vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [] }); + vi.mocked(getDefaultTeamSettings).mockResolvedValue({ values: { budget_duration: "30d" } }); + vi.mocked(teamCreateCall).mockResolvedValue({ + team_id: "new-team-1", + team_alias: "Test Team", + models: ["gpt-4"], + organization_id: null, + keys: [], + members_with_roles: [], + spend: 0, + }); + mockUseOrganizations.mockReturnValue({ data: null }); + }); + + const openCreateModal = async () => { + renderWithQueryClient(); + + const createButton = screen.getAllByRole("button", { name: /create team/i })[0]; + act(() => { + fireEvent.click(createButton); + }); + + await waitFor(() => { + expect(screen.getByLabelText(/team name/i)).toBeInTheDocument(); + }); + }; + + const resetBudgetField = () => screen.getByText("Reset Budget").closest(".ant-form-item") as HTMLElement; + + const submitCreateModal = async () => { + fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "Test Team" } }); + + const createTeamSubmitButtons = screen.getAllByRole("button", { name: /create team/i }); + fireEvent.click(createTeamSubmitButtons[createTeamSubmitButtons.length - 1]); + + await waitFor(() => { + expect(teamCreateCall).toHaveBeenCalled(); + }); + + return vi.mocked(teamCreateCall).mock.calls[0][1]; + }; + + it("should send an explicit null budget_duration when Never resets is selected", async () => { + await openCreateModal(); + + await userEvent.click(within(resetBudgetField()).getByRole("combobox")); + await userEvent.click(await screen.findByText("Never resets")); + + const payload = await submitCreateModal(); + + expect(payload.budget_duration).toBeNull(); + expect(JSON.stringify(payload)).toContain('"budget_duration":null'); + }); + + it("should omit budget_duration entirely when Reset Budget is left untouched", async () => { + await openCreateModal(); + + const payload = await submitCreateModal(); + + expect(payload.budget_duration).toBeUndefined(); + expect(JSON.stringify(payload)).not.toContain("budget_duration"); + }); + + it("should send the picked duration when one is selected", async () => { + await openCreateModal(); + + await userEvent.click(within(resetBudgetField()).getByRole("combobox")); + await userEvent.click(await screen.findByText("weekly")); + + const payload = await submitCreateModal(); + + expect(payload.budget_duration).toBe("7d"); + }); + + it("should show the configured server default as the Reset Budget placeholder", async () => { + await openCreateModal(); + + await waitFor(() => { + expect(within(resetBudgetField()).getByText("Default: monthly (30d)")).toBeInTheDocument(); + }); + }); + + it("should fall back to the n/a placeholder when the default settings fetch fails", async () => { + vi.mocked(getDefaultTeamSettings).mockRejectedValue(new Error("Unauthorized")); + + await openCreateModal(); + + await waitFor(() => { + expect(within(resetBudgetField()).getByText("n/a")).toBeInTheDocument(); + }); + }); +}); + describe("Teams - metadata key-value pairs in team create", () => { beforeEach(() => { vi.clearAllMocks(); diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index 7c269555ec0..becbe0e2b48 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -9,7 +9,7 @@ import { Accordion, AccordionBody, AccordionHeader, TextInput } from "@tremor/re import { Button, Form, Input, Layout, Modal, Select, Switch, Tabs, theme, Tooltip, Typography } from "antd"; import { Plus, Users } from "lucide-react"; import React, { useEffect, useState } from "react"; -import { useQueryClient } from "@tanstack/react-query"; +import { useQuery, useQueryClient } from "@tanstack/react-query"; import { PageHeader } from "@/components/shared/PageHeader"; import { Button as UIButton } from "@/components/ui/button"; import { teamsTableKeys } from "@/app/(dashboard)/hooks/teams/useTeams"; @@ -29,7 +29,11 @@ import MCPServerSelector from "./mcp_server_management/MCPServerSelector"; import MCPToolPermissions from "./mcp_server_management/MCPToolPermissions"; import NotificationsManager from "./molecules/notifications_manager"; import { extractProxyErrorMessage } from "@/lib/http/client"; -import { Organization, getGuardrailsList, getPoliciesList, teamDeleteCall } from "./networking"; +import BudgetDurationDropdown, { + getBudgetDurationLabel, + NEVER_RESETS_BUDGET_DURATION, +} from "./common_components/budget_duration_dropdown"; +import { Organization, getDefaultTeamSettings, getGuardrailsList, getPoliciesList, teamDeleteCall } from "./networking"; import NumericalInput from "./shared/numerical_input"; import VectorStoreSelector from "./vector_store_management/VectorStoreSelector"; import SearchToolSelector from "./search_tools/SearchToolSelector"; @@ -116,6 +120,18 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser const [routerSettings, setRouterSettings] = useState(null); const [routerSettingsKey, setRouterSettingsKey] = useState(0); + const { data: defaultTeamSettings } = useQuery({ + queryKey: ["defaultTeamSettings"], + queryFn: () => getDefaultTeamSettings(accessToken as string), + enabled: isTeamModalVisible && accessToken != null, + retry: false, + staleTime: 60_000, + }); + const defaultBudgetDuration: string | undefined = defaultTeamSettings?.values?.budget_duration ?? undefined; + const budgetDurationPlaceholder = defaultBudgetDuration + ? `Default: ${getBudgetDurationLabel(defaultBudgetDuration)} (${defaultBudgetDuration})` + : "n/a"; + useEffect(() => { form.setFieldValue("models", []); }, [currentOrgForCreateTeam, userModels]); @@ -249,6 +265,10 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser formValues.organization_id = organizationId.trim(); } + if (formValues.budget_duration === NEVER_RESETS_BUDGET_DURATION) { + formValues.budget_duration = null; + } + NotificationsManager.info("Creating Team"); const metadataObject = { @@ -645,11 +665,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser - + diff --git a/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx b/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx index 847a6ca1949..4db36f27553 100644 --- a/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx @@ -3,12 +3,15 @@ import { Select } from "antd"; const { Option } = Select; +export const NEVER_RESETS_BUDGET_DURATION = "none"; + interface BudgetDurationDropdownProps { value?: string | null; onChange?: (value: string | undefined) => void; className?: string; style?: React.CSSProperties; placeholder?: string; + showNeverResets?: boolean; } const BudgetDurationDropdown: React.FC = ({ @@ -17,6 +20,7 @@ const BudgetDurationDropdown: React.FC = ({ className = "", style = {}, placeholder = "n/a", + showNeverResets = false, }) => { return ( onChange?.(event.target.value)}> + + {showNeverResets ? : null} + + + ), +})); vi.mock("../common_components/check_openapi_schema", () => ({ default: () => null })); vi.mock("../common_components/KeyLifecycleSettings", () => ({ default: () => null })); vi.mock("../common_components/ModelAliasManager", () => ({ default: () => null })); @@ -820,4 +829,58 @@ describe("CreateKey", () => { expect(screen.queryByPlaceholderText(PROMPTS_PLACEHOLDER)).not.toBeInTheDocument(); }); }); + + describe("budget reset", () => { + const openModal = async () => { + renderWithProviders(); + + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create new key/i })); + }); + + await waitFor(() => { + expect(screen.getByTestId("budget-duration-dropdown")).toBeInTheDocument(); + }); + }; + + const submit = async () => { + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create key/i })); + }); + + await waitFor(() => { + expect(mockKeyCreateCall).toHaveBeenCalled(); + }); + + return mockKeyCreateCall.mock.calls[0][2]; + }; + + it("should send an explicit null budget_duration when 'Never resets' is selected", async () => { + await openModal(); + + expect(screen.getByRole("option", { name: "Never resets" })).toBeInTheDocument(); + + act(() => { + fireEvent.change(screen.getByTestId("budget-duration-dropdown"), { target: { value: "none" } }); + formMock.setFieldValue("key_alias", "Never Resets Key"); + }); + + const formValues = await submit(); + + expect("budget_duration" in formValues).toBe(true); + expect(formValues.budget_duration).toBeNull(); + }); + + it("should omit budget_duration entirely when the reset dropdown is untouched", async () => { + await openModal(); + + act(() => { + formMock.setFieldValue("key_alias", "Inherits Default Key"); + }); + + const formValues = await submit(); + + expect("budget_duration" in formValues).toBe(false); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 8951a471f84..c747cfe36fe 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -18,7 +18,7 @@ import { rolesWithWriteAccess } from "../../utils/roles"; import AgentSelector from "../agent_management/AgentSelector"; import { mapDisplayToInternalNames } from "../callback_info_helpers"; import AccessGroupSelector from "../common_components/AccessGroupSelector"; -import BudgetDurationDropdown from "../common_components/budget_duration_dropdown"; +import BudgetDurationDropdown, { NEVER_RESETS_BUDGET_DURATION } from "../common_components/budget_duration_dropdown"; import SchemaFormFields from "../common_components/check_openapi_schema"; import KeyLifecycleSettings from "../common_components/KeyLifecycleSettings"; import ModelAliasManager from "../common_components/ModelAliasManager"; @@ -542,6 +542,10 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp formValues.budget_fallbacks = budgetFallbacks; } + if (formValues.budget_duration === NEVER_RESETS_BUDGET_DURATION) { + formValues.budget_duration = null; + } + let response; if (keyOwner === "service_account") { response = await keyCreateServiceAccountCall(accessToken, formValues); @@ -1064,6 +1068,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp help={`Team Reset Budget: ${team?.budget_duration !== null && team?.budget_duration !== undefined ? team?.budget_duration : "None"}`} > form.setFieldValue("budget_duration", value)} /> diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index d6df58ec0f8..df7523dfe4e 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -959,6 +959,83 @@ describe("TeamInfoView", () => { ); }); }); + + const openSettingsEditorForTeam = async ( + user: ReturnType, + teamOverrides: Record, + ) => { + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData(teamOverrides)); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0); + }); + + await user.click(screen.getByRole("tab", { name: "Settings" })); + await waitFor(() => { + expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /edit settings/i })); + await waitFor(() => { + expect(screen.getByLabelText("Team Name")).toBeInTheDocument(); + }); + + return screen.getByText("Reset Budget").closest(".ant-form-item") as HTMLElement; + }; + + it("should send an explicit null budget_duration when a stored Reset Budget is cleared", async () => { + const user = userEvent.setup({ delay: null }); + const resetBudgetItem = await openSettingsEditorForTeam(user, { budget_duration: "30d" }); + + const clearIcon = resetBudgetItem.querySelector(".ant-select-clear"); + expect(clearIcon).not.toBeNull(); + fireEvent.mouseDown(clearIcon as Element); + + await waitFor(() => { + expect(within(resetBudgetItem).getByText("Never resets")).toBeInTheDocument(); + }); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalled(); + }); + + const updateArg = vi.mocked(networking.teamUpdateCall).mock.calls[0][1]; + expect(updateArg.budget_duration).toBeNull(); + expect(JSON.stringify(updateArg)).toContain('"budget_duration":null'); + }); + + it("should keep a stored Reset Budget when the form is saved untouched", async () => { + const user = userEvent.setup({ delay: null }); + await openSettingsEditorForTeam(user, { budget_duration: "30d" }); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalled(); + }); + + expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].budget_duration).toBe("30d"); + }); + + it("should send the newly picked budget_duration when one is selected", async () => { + const user = userEvent.setup({ delay: null }); + const resetBudgetItem = await openSettingsEditorForTeam(user, { budget_duration: null }); + + await user.click(within(resetBudgetItem).getByRole("combobox")); + await user.click(await screen.findByText("weekly")); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalled(); + }); + + expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].budget_duration).toBe("7d"); + }); }); describe("metadata key-value editing", () => { diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index eae11d481e3..a23cf77e00e 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -36,6 +36,7 @@ import { CheckIcon, CopyIcon } from "lucide-react"; import React, { useEffect, useMemo, useState } from "react"; import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils"; import AccessGroupSelector from "../common_components/AccessGroupSelector"; +import BudgetDurationDropdown from "../common_components/budget_duration_dropdown"; import { computeTeamModelBadges, normalizeTeamModelSelection, @@ -526,7 +527,7 @@ const TeamInfoView: React.FC = ({ model_rpm_limit: modelRpmLimit, max_budget: values.max_budget, soft_budget: sanitizeNumeric(values.soft_budget), - budget_duration: values.budget_duration, + budget_duration: values.budget_duration ?? null, metadata: { ...parsedMetadata, ...passthroughRoutesMetadata, @@ -1152,11 +1153,7 @@ const TeamInfoView: React.FC = ({ - + diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a2a4c01c25f..b6dbe0dfc6d 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -4420,7 +4420,7 @@ export interface paths { }; /** * Get Default Team Settings - * @description Get all SSO settings from the litellm_settings configuration. + * @description Get the default team parameters (litellm_settings.default_team_params). * Returns a structured object with values and descriptions for UI display. */ get: operations["get_default_team_settings_get_default_team_settings_get"]; @@ -14810,8 +14810,9 @@ export interface paths { head?: never; /** * Update Default Team Settings - * @description Update the default team parameters for SSO users. - * These settings will be applied to new teams created from SSO. + * @description Update the default team parameters (litellm_settings.default_team_params). + * Applied to every new team for fields not explicitly provided in the create request; + * `models` only applies to teams automatically created via SSO Groups. */ patch: operations["update_default_team_settings_update_default_team_settings_patch"]; trace?: never; @@ -24431,22 +24432,23 @@ export interface components { }; /** * DefaultTeamSSOParams - * @description Default parameters to apply when a new team is automatically created by LiteLLM via SSO Groups + * @description Default parameters applied to every /team/new call for fields not explicitly provided in the request. + * `models` is the exception: it only applies to teams automatically created by LiteLLM via SSO Groups. */ DefaultTeamSSOParams: { /** * Budget Duration - * @description Default budget duration for new automatically created teams (e.g. 'daily', 'weekly', 'monthly') + * @description Default budget duration for new teams, when not explicitly provided (e.g. '24h', '7d', '30d') */ budget_duration?: string | null; /** * Max Budget - * @description Default maximum budget (in USD) for new automatically created teams + * @description Default maximum budget (in USD) for new teams, when not explicitly provided */ max_budget?: number | null; /** * Models - * @description Default list of models that new automatically created teams can access + * @description Default list of models for teams automatically created via SSO Groups * @default [] */ models: string[]; @@ -24457,7 +24459,7 @@ export interface components { organization_id?: string | null; /** * Rpm Limit - * @description Default rpm limit for new automatically created teams + * @description Default rpm limit for new teams, when not explicitly provided */ rpm_limit?: number | null; /** @@ -24467,7 +24469,7 @@ export interface components { team_member_permissions?: components["schemas"]["KeyManagementRoutes"][] | null; /** * Tpm Limit - * @description Default tpm limit for new automatically created teams + * @description Default tpm limit for new teams, when not explicitly provided */ tpm_limit?: number | null; };