From 1b85426a6ad2612aa5b9b5e2d889be1f8443d68d Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Wed, 15 Jul 2026 12:54:26 -0700 Subject: [PATCH] feat(ptu): runtime UI flag for enable_ptu_cost_attribution + always-registered rollup Add enable_ptu_cost_attribution to UISettings so admins can flip it from the UI's Admin Settings page without a proxy restart. Register the daily rollup job unconditionally at startup; it already no-ops on flag-off, so the runtime toggle takes effect on the next scheduled fire without any restart. Existing per-request guards in the /ptu_reservation endpoints and in run_ptu_reservation_rollup remain the source of truth. Persisted settings sync into general_settings on both PATCH and GET paths as well as on server startup via _sync_ui_settings_to_general_settings. --- litellm/proxy/proxy_server.py | 33 +++++------ .../proxy_setting_endpoints.py | 7 +++ tests/test_litellm/proxy/test_proxy_server.py | 38 +++++++++++++ .../test_proxy_setting_endpoints.py | 55 +++++++++++++++++++ .../UISettings/UISettings.test.tsx | 33 +++++++++++ .../AdminSettings/UISettings/UISettings.tsx | 33 +++++++++++ 6 files changed, 183 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b0099481de5..085ea7342fd 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7930,23 +7930,24 @@ async def initialize_scheduled_background_jobs( await cls._initialize_spend_tracking_background_jobs(scheduler=scheduler) ### PTU RESERVATION DAILY ROLLUP ### - if general_settings.get("enable_ptu_cost_attribution", False): - from litellm.proxy.spend_tracking.ptu_reservation_rollup import ( - PTU_ROLLUP_JOB_ID, - run_ptu_reservation_rollup, - ) + from litellm.proxy.spend_tracking.ptu_reservation_rollup import ( + PTU_ROLLUP_JOB_ID, + run_ptu_reservation_rollup, + ) - scheduler.add_job( - run_ptu_reservation_rollup, - "cron", - hour=0, - minute=15, - args=[prisma_client], - id=PTU_ROLLUP_JOB_ID, - replace_existing=True, - misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, - ) - verbose_proxy_logger.info("PTU reservation rollup job scheduled at 00:15 UTC daily") + scheduler.add_job( + run_ptu_reservation_rollup, + "cron", + hour=0, + minute=15, + args=[prisma_client], + id=PTU_ROLLUP_JOB_ID, + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + verbose_proxy_logger.info( + "PTU reservation rollup job scheduled at 00:15 UTC daily (no-ops until enable_ptu_cost_attribution is set)" + ) ### SPEND LOG CLEANUP ### if general_settings.get("maximum_spend_logs_retention_period") is not None: diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index a8926d26047..80d5f6c4c44 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -182,6 +182,11 @@ class UISettings(BaseModel): description="If true, shows the Chat page in the UI sidebar, letting users chat with an LLM and connect their own MCP server credentials via OAuth.", ) + enable_ptu_cost_attribution: bool = Field( + default=False, + description="If true, enables admin-registered PTU reservations and daily flat-cost attribution on team daily spend. Governs the /ptu_reservation CRUD endpoints, the daily rollup job, the PTU Reservations UI page, and the Flat Cost column on the Usage page.", + ) + class UISettingsResponse(SettingsResponse): """Response model for UI settings""" @@ -206,6 +211,7 @@ class UISettingsResponse(SettingsResponse): "disable_custom_api_keys", "disable_key_generate_for_org_admin", "enable_chat_ui", + "enable_ptu_cost_attribution", } # Flags that must be synced from the persisted UISettings into @@ -219,6 +225,7 @@ class UISettingsResponse(SettingsResponse): "disable_vector_stores_for_internal_users", "allow_vector_stores_for_team_admins", "disable_key_generate_for_org_admin", + "enable_ptu_cost_attribution", ] # Extension point: packages outside OSS (e.g. litellm_enterprise) can diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 2f0924e9192..64235e2a306 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9718,3 +9718,41 @@ async def test_startup_survives_database_read_failure_for_coordination_redis(): ) assert result is None + + +@pytest.mark.asyncio +async def test_ptu_rollup_job_registered_regardless_of_flag(monkeypatch): + """The PTU rollup cron must be registered at startup so a UI-time flip of enable_ptu_cost_attribution takes effect without a proxy restart. The job itself no-ops when the flag is off (asserted in test_ptu_reservation_rollup.py).""" + monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) + from litellm.proxy.proxy_server import ProxyStartupEvent + from litellm.proxy.spend_tracking.ptu_reservation_rollup import ( + PTU_ROLLUP_JOB_ID, + ) + from litellm.proxy.utils import ProxyLogging + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=None) + + mock_proxy_logging = MagicMock(spec=ProxyLogging) + mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_config = AsyncMock() + + with ( + patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch("litellm.proxy.proxy_server.get_secret_bool", return_value=True), + ): + # Flag OFF at startup — job must still be registered + await ProxyStartupEvent.initialize_scheduled_background_jobs( + general_settings={"enable_ptu_cost_attribution": False}, + prisma_client=mock_prisma_client, + proxy_budget_rescheduler_min_time=1, + proxy_budget_rescheduler_max_time=2, + proxy_batch_write_at=5, + proxy_logging_obj=mock_proxy_logging, + ) + + import litellm.proxy.proxy_server as ps + + assert ps.scheduler is not None + assert ps.scheduler.get_job(PTU_ROLLUP_JOB_ID) is not None diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 69845ec59c2..6b3ab154fe7 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -1336,6 +1336,61 @@ def test_update_ui_settings_persists_and_syncs_disable_key_generate_for_org_admi # Synced into general_settings so the enforcement helper sees it assert general_settings.get(flag_name) is True + def test_update_ui_settings_persists_and_syncs_enable_ptu_cost_attribution( + self, mock_auth, monkeypatch + ): + """enable_ptu_cost_attribution must be allowlisted, persisted, and synced to general_settings so PTU endpoints and rollup pick it up without a proxy restart.""" + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + + general_settings: dict = {} + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", general_settings + ) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_uisettings.upsert = AsyncMock() + mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + + flag_name = "enable_ptu_cost_attribution" + payload = {flag_name: True} + + try: + response = client.patch("/update/ui_settings", json=payload) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 200 + data = response.json() + assert data["settings"][flag_name] is True + + call_args = mock_prisma.db.litellm_uisettings.upsert.call_args + stored_settings = json.loads(call_args.kwargs["data"]["create"]["ui_settings"]) + assert stored_settings[flag_name] is True + + assert general_settings.get(flag_name) is True + + def test_enable_ptu_cost_attribution_registered_as_runtime_flag(self): + """The runtime-sync list must include the PTU flag so a UI flip lands in general_settings without a restart.""" + from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( + _RUNTIME_GENERAL_SETTINGS_FLAGS, + ALLOWED_UI_SETTINGS_FIELDS, + ) + + assert "enable_ptu_cost_attribution" in _RUNTIME_GENERAL_SETTINGS_FLAGS + assert "enable_ptu_cost_attribution" in ALLOWED_UI_SETTINGS_FIELDS + def test_get_sso_settings_from_database( self, mock_proxy_config, mock_auth, monkeypatch ): diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.test.tsx index 639564bbd35..9d472292bc6 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.test.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.test.tsx @@ -41,12 +41,16 @@ const buildSettingsResponse = (overrides?: Partial>) => require_auth_for_public_ai_hub: { description: "Require authentication for public AI Hub", }, + enable_ptu_cost_attribution: { + description: "Enable PTU cost attribution", + }, }, }, values: { disable_model_add_for_internal_users: false, disable_team_admin_delete_team_user: false, require_auth_for_public_ai_hub: false, + enable_ptu_cost_attribution: false, }, }, isLoading: false, @@ -162,4 +166,33 @@ describe("UISettings", () => { ); expect(NotificationManager.success).toHaveBeenCalledWith("UI settings updated successfully"); }); + + it("should toggle enable PTU cost attribution setting and call update with page-refresh notice", () => { + const mutateMock = vi.fn((_settings, options) => { + options?.onSuccess?.(); + }); + + mockUseUpdateUISettings.mockReturnValue({ + mutate: mutateMock, + isPending: false, + error: null, + }); + + render(); + + const toggle = screen.getByRole("switch", { name: "Enable PTU cost attribution" }); + + act(() => { + fireEvent.click(toggle); + }); + + expect(mutateMock).toHaveBeenCalledWith( + { enable_ptu_cost_attribution: true }, + expect.objectContaining({ + onSuccess: expect.any(Function), + onError: expect.any(Function), + }), + ); + expect(NotificationManager.success).toHaveBeenCalledWith("UI settings updated successfully. Refreshing page..."); + }); }); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx index ec970c34873..7ea5b6ee903 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx @@ -20,6 +20,7 @@ export default function UISettings() { const forwardLLMProviderAuthHeadersProperty = schema?.properties?.forward_llm_provider_auth_headers; const enableProjectsUIProperty = schema?.properties?.enable_projects_ui; const enableChatUIProperty = schema?.properties?.enable_chat_ui; + const enablePtuCostAttributionProperty = schema?.properties?.enable_ptu_cost_attribution; const enabledPagesProperty = schema?.properties?.enabled_ui_pages_internal_users; const disableAgentsProperty = schema?.properties?.disable_agents_for_internal_users; const allowAgentsTeamAdminsProperty = schema?.properties?.allow_agents_for_team_admins; @@ -130,6 +131,21 @@ export default function UISettings() { ); }; + const handleToggleEnablePtuCostAttribution = (checked: boolean) => { + updateSettings( + { enable_ptu_cost_attribution: checked }, + { + onSuccess: () => { + NotificationManager.success("UI settings updated successfully. Refreshing page..."); + setTimeout(() => window.location.reload(), 1000); + }, + onError: (error) => { + NotificationManager.fromBackend(error); + }, + }, + ); + }; + const handleToggleRequireAuthForPublicAIHub = (checked: boolean) => { updateSettings( { require_auth_for_public_ai_hub: checked }, @@ -368,6 +384,23 @@ export default function UISettings() { + + + + Enable PTU cost attribution (page will refresh) + + {enablePtuCostAttributionProperty?.description ?? + "If enabled, unlocks admin-registered PTU reservations and daily flat-cost attribution on team daily spend."} + + + + {/* Agents access control */}