feat(ptu): runtime UI flag for enable_ptu_cost_attribution + always-registered rollup

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:

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

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

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
     ):

ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.test.tsx

@@ -41,12 +41,16 @@ const buildSettingsResponse = (overrides?: Partial<Record<string, unknown>>) =>
         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(<UISettings />);
+
+    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...");
+  });
 });

ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx

@@ -7,7 +7,7 @@
 import PageVisibilitySettings from "./PageVisibilitySettings";
 import { Alert, Card, Divider, Skeleton, Space, Switch, Typography } from "antd";
 
 export default function UISettings() {
   const { accessToken } = useAuthorized();
   const { data, isLoading, isError, error } = useUISettings();
   const { mutate: updateSettings, isPending: isUpdating, error: updateError } = useUpdateUISettings(accessToken);
@@ -20,6 +20,7 @@
   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 @@
     );
   };
 
+  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 @@
             </Space>
           </Space>
 
+          <Space align="start" size="middle">
+            <Switch
+              checked={Boolean(values.enable_ptu_cost_attribution)}
+              disabled={isUpdating}
+              loading={isUpdating}
+              onChange={handleToggleEnablePtuCostAttribution}
+              aria-label={enablePtuCostAttributionProperty?.description ?? "Enable PTU cost attribution"}
+            />
+            <Space direction="vertical" size={4}>
+              <Typography.Text strong>Enable PTU cost attribution (page will refresh)</Typography.Text>
+              <Typography.Text type="secondary">
+                {enablePtuCostAttributionProperty?.description ??
+                  "If enabled, unlocks admin-registered PTU reservations and daily flat-cost attribution on team daily spend."}
+              </Typography.Text>
+            </Space>
+          </Space>
+
           <Divider />
 
           {/* Agents access control */}