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.
This commit is contained in:
Yucheng Zhu 2026-07-15 12:54:26 -07:00
parent 5808dcf079
commit 1b85426a6a
6 changed files with 183 additions and 16 deletions

View file

@ -7930,23 +7930,24 @@ class ProxyStartupEvent:
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:

View file

@ -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 @@ ALLOWED_UI_SETTINGS_FIELDS = {
"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 @@ _RUNTIME_GENERAL_SETTINGS_FLAGS = [
"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

View file

@ -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

View file

@ -1336,6 +1336,61 @@ class TestProxySettingEndpoints:
# 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
):

View file

@ -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...");
});
});

View file

@ -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() {
</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 */}