diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index fb302e87bd9..d1542b38996 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -10,6 +10,7 @@ from typing_extensions import TypedDict from litellm._logging import verbose_proxy_logger from litellm.constants import PTU_SENTINEL_API_KEY from litellm.proxy._types import CommonProxyErrors +from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled from litellm.proxy.utils import PrismaClient from litellm.repositories.table_repositories import DeletedVerificationTokenRepository from litellm.repositories.verification_token_repository import ( @@ -141,6 +142,28 @@ class _GroupingSetsRow(SimpleNamespace): failed_requests: int | None +def _reported_flat_cost(record: DailySpendRecord | _GroupingSetsRow) -> float: + """Flat cost a daily row reports, which is zero unless PTU cost attribution is enabled. + + Both read paths funnel through here: the paginated path reads the ``ptu_flat_cost`` + column straight off the row, and the aggregated path reads the SUM() alias. Rows an + operator accrued during an earlier opt-in stay in the table, so the gate lives on the + read rather than on the query that produced the rows. + + The row is checked before the flag because this runs once per metric accumulation, and + a record fans out across roughly a dozen breakdowns. The flag reads through the secret + manager, uncached, so consulting it for every accumulation put thousands of lookups on + a shared endpoint that made none before. Only a row actually carrying flat cost, which + is a sentinel row, reaches it now. + """ + raw: Final = getattr(record, "ptu_flat_cost", None) or 0.0 + if not raw: + return 0.0 + if not is_ptu_cost_attribution_enabled(): + return 0.0 + return raw + + def update_metrics(existing_metrics: SpendMetrics, record: DailySpendRecord) -> SpendMetrics: """Update metrics with new record data. @@ -151,7 +174,7 @@ def update_metrics(existing_metrics: SpendMetrics, record: DailySpendRecord) -> prompt_tokens: Final = record.prompt_tokens or 0 completion_tokens: Final = record.completion_tokens or 0 existing_metrics.spend += record.spend or 0.0 - existing_metrics.flat_cost += getattr(record, "ptu_flat_cost", None) or 0.0 + existing_metrics.flat_cost += _reported_flat_cost(record) existing_metrics.prompt_tokens += prompt_tokens existing_metrics.completion_tokens += completion_tokens existing_metrics.total_tokens += prompt_tokens + completion_tokens @@ -784,7 +807,7 @@ def _record_to_spend_metrics(record: _GroupingSetsRow) -> SpendMetrics: completion_tokens: Final = record.completion_tokens or 0 return SpendMetrics( spend=record.spend or 0.0, - flat_cost=getattr(record, "ptu_flat_cost", None) or 0.0, + flat_cost=_reported_flat_cost(record), prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens, diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index cf6af66d5a9..8a52b0d1abb 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -56,6 +56,10 @@ from litellm.proxy.management_endpoints.team_endpoints import ( update_team as _legacy_update_team, ) from litellm.proxy.management_helpers.audit_logs import create_object_audit_log +from litellm.proxy.spend_tracking.ptu_feature_flag import ( + PTU_COST_ATTRIBUTION_ENV_VAR, + is_ptu_cost_attribution_enabled, +) from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.model_repository import ModelRepository from litellm.repositories.table_repositories import ModelTableRepository @@ -239,8 +243,12 @@ _PTU_MODEL_INFO_FIELDS: Final = ("ptu_count", "cost_per_ptu_per_hour", "ptu_effe def _explicitly_cleared_ptu_fields(model_info: ModelInfo | None) -> frozenset[str]: - """The PTU fields a patch sends as an explicit null, which update_db_model drops.""" - if model_info is None: + """The PTU fields a patch sends as an explicit null, which update_db_model drops. + + Empty while the feature is off, so disabling pauses PTU rather than letting a client + that round-trips a model_info blob erase a configuration set up during an earlier opt-in. + """ + if model_info is None or not is_ptu_cost_attribution_enabled(): return frozenset() return frozenset( field @@ -262,6 +270,32 @@ def _merged_ptu_model_info(*, db_model: Deployment, patch_data: updateDeployment return MappingProxyType({k: v for k, v in {**stored, **incoming}.items() if k not in cleared}) +def _raise_if_ptu_cost_attribution_disabled(incoming_model_info: Mapping[str, object]) -> None: + """Reject PTU model_info fields unless the operator opted into PTU cost attribution. + + Takes the incoming request's model_info rather than the merged deployment, so an + unrelated patch of a model that still stores PTU config from an earlier opt-in is + left alone. The fields are rejected rather than dropped so a caller never believes + a flat cost was configured while the rollup that would price it is not running. + + Only a value is rejected. An explicit null reaches the clear loop, which is gated on + the same flag, so a disabled proxy neither writes PTU config nor erases what an + earlier opt-in stored. Disabling pauses the feature rather than discarding its setup. + """ + if is_ptu_cost_attribution_enabled(): + return + supplied: Final = tuple(field for field in _PTU_MODEL_INFO_FIELDS if incoming_model_info.get(field) is not None) + if not supplied: + return + raise HTTPException( + status_code=400, + detail=( + f"PTU cost attribution is disabled, so {', '.join(supplied)} cannot be set. " + f"Set {PTU_COST_ATTRIBUTION_ENV_VAR}=true to enable it." + ), + ) + + def _validate_ptu_model_info(model_info: Mapping[str, object]) -> None: """Enforce the PTU cross-field invariant on the effective model_info. @@ -326,6 +360,8 @@ def _coerce_ptu_datetime(value: object) -> datetime.datetime | None: def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> PrismaCompatibleUpdateDBModel: + if updated_patch.model_info is not None: + _raise_if_ptu_cost_attribution_disabled(updated_patch.model_info.model_dump(exclude_none=True)) merged_model_name: Final = updated_patch.model_name or db_model.model_name merged_litellm_params: Final = db_model.litellm_params.model_dump(exclude_none=True) merged_model_info: Final = db_model.model_info.model_dump(exclude_none=True) @@ -821,6 +857,7 @@ async def _update_team_model_in_db( # raising the rate on a configured model carries no ptu_effective_from, which the # stored row supplies. if patch_data.model_info is not None: + _raise_if_ptu_cost_attribution_disabled(patch_data.model_info.model_dump(exclude_none=True)) _validate_ptu_model_info(_merged_ptu_model_info(db_model=db_model, patch_data=patch_data)) patch_team_id: Final = patch_data.model_info.team_id if patch_data.model_info else None @@ -1531,7 +1568,9 @@ async def add_new_model( model_response: LiteLLM_ProxyModelTable | None = None # update DB - _validate_ptu_model_info(model_params.model_info.model_dump(exclude_none=True)) + incoming_model_info: Final = model_params.model_info.model_dump(exclude_none=True) + _raise_if_ptu_cost_attribution_disabled(incoming_model_info) + _validate_ptu_model_info(incoming_model_info) if store_model_in_db is True: """ diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 641d294f592..3a4896dca9e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8471,40 +8471,45 @@ class ProxyStartupEvent: await cls._initialize_spend_tracking_background_jobs(scheduler=scheduler) ### PTU DAILY ROLLUP ### - from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import ( - PTU_ROLLUP_JOB_ID, - run_scheduled_ptu_rollup, + from litellm.proxy.spend_tracking.ptu_feature_flag import ( + is_ptu_cost_attribution_enabled, ) - async def _alert_ptu_rollup_failure(message: str) -> None: - await proxy_logging_obj.alerting_handler( - message=message, - level="High", - alert_type=AlertType.failed_tracking_spend, + if is_ptu_cost_attribution_enabled(): + from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import ( + PTU_ROLLUP_JOB_ID, + run_scheduled_ptu_rollup, ) - async def _scheduled_ptu_rollup() -> None: - # Reuse the PodLockManager from db_spend_update_writer so only one pod - # reconciles a day; a multi-pod race could prune another pod's fresh rows - await run_scheduled_ptu_rollup( - prisma_client, - pod_lock_manager=proxy_logging_obj.db_spend_update_writer.pod_lock_manager, - alert=_alert_ptu_rollup_failure, - ) + async def _alert_ptu_rollup_failure(message: str) -> None: + await proxy_logging_obj.alerting_handler( + message=message, + level="High", + alert_type=AlertType.failed_tracking_spend, + ) - scheduler.add_job( - _scheduled_ptu_rollup, - "cron", - hour=0, - minute=15, - timezone="UTC", - id=PTU_ROLLUP_JOB_ID, - replace_existing=True, - misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, - ) - verbose_proxy_logger.info( - "PTU rollup job scheduled at 00:15 UTC daily (only models with PTU config accrue flat cost)" - ) + async def _scheduled_ptu_rollup() -> None: + # Reuse the PodLockManager from db_spend_update_writer so only one pod + # reconciles a day; a multi-pod race could prune another pod's fresh rows + await run_scheduled_ptu_rollup( + prisma_client, + pod_lock_manager=proxy_logging_obj.db_spend_update_writer.pod_lock_manager, + alert=_alert_ptu_rollup_failure, + ) + + scheduler.add_job( + _scheduled_ptu_rollup, + "cron", + hour=0, + minute=15, + timezone="UTC", + id=PTU_ROLLUP_JOB_ID, + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + verbose_proxy_logger.info( + "PTU rollup job scheduled at 00:15 UTC daily (only models with PTU config accrue flat cost)" + ) ### SPEND LOG CLEANUP ### if ( diff --git a/litellm/proxy/spend_tracking/ptu_feature_flag.py b/litellm/proxy/spend_tracking/ptu_feature_flag.py new file mode 100644 index 00000000000..9078079b676 --- /dev/null +++ b/litellm/proxy/spend_tracking/ptu_feature_flag.py @@ -0,0 +1,18 @@ +"""Opt-in flag for PTU (provisioned throughput unit) flat-cost attribution. + +The whole feature is inert unless an operator sets +``LITELLM_ENABLE_PTU_COST_ATTRIBUTION``: the daily rollup is not scheduled, the +model endpoints reject PTU config, the daily activity read path reports zero flat +cost, and the model form hides the PTU inputs. +""" + +from typing import Final + +from litellm.secret_managers.main import get_secret_bool + +PTU_COST_ATTRIBUTION_ENV_VAR: Final = "LITELLM_ENABLE_PTU_COST_ATTRIBUTION" + + +def is_ptu_cost_attribution_enabled() -> bool: + """Report whether this deployment opted into PTU flat-cost attribution.""" + return get_secret_bool(PTU_COST_ATTRIBUTION_ENV_VAR, False) is True diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py index a77c3e29cfa..029648f7901 100644 --- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py +++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py @@ -27,6 +27,7 @@ from litellm.constants import ( PTU_ROLLUP_MAX_BACKFILL_DAYS, PTU_SENTINEL_API_KEY, ) +from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled from litellm.types.router import ModelInfo if TYPE_CHECKING: @@ -512,7 +513,15 @@ async def run_scheduled_ptu_rollup( duplicate work rather than correctness: the upserts are idempotent on the sentinel key and the prune reads only the row's own timestamp, so a second pod arriving mid-run cannot corrupt the day. + + Returns None without touching the database when PTU cost attribution is off. Proxy + startup already skips scheduling the cron, so this guards the function itself rather + than its one caller, and a deployment that never opted in accrues nothing whatever + reaches it. """ + if not is_ptu_cost_attribution_enabled(): + return None + if pod_lock_manager is None or pod_lock_manager.redis_cache is None: return await _run_and_alert(prisma_client, target_date=target_date, alert=alert, may_prune=False) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 08bb8698cac..b3feb5bd8d6 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -21,6 +21,7 @@ from litellm.proxy.config_resolvers.sso import ( SSO_SECRET_FIELDS, resolve_sso_config, ) +from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled from litellm.proxy.utils import invalidate_config_param from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.organization_repository import OrganizationRepository @@ -307,6 +308,27 @@ ALLOWED_UI_SETTINGS_FIELDS: Final = { "enable_chat_ui", } +ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: Final = "enable_ptu_cost_attribution" + +# UI settings derived from the deployment environment. Deliberately kept out of +# ALLOWED_UI_SETTINGS_FIELDS: they are read-only, never persisted, and PATCH +# rejects them so an admin cannot flip an env-gated feature at runtime. +_DERIVED_UI_SETTINGS_FIELDS: Final[frozenset[str]] = frozenset({ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING}) + + +def _derived_ui_setting_value(key: str) -> object: + """The environment-derived value GET reports for ``key``. + + PATCH compares against this rather than rejecting the key outright, so the body GET + hands back is still a valid PATCH body. Rejecting on presence broke read-modify-write: + a client that edited one setting and sent the rest back unchanged got a 400 and lost + the edit it actually wanted. + """ + if key == ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: + return is_ptu_cost_attribution_enabled() + return None + + # Flags that must be synced from the persisted UISettings into # general_settings at runtime (on both read and write). _RUNTIME_GENERAL_SETTINGS_FLAGS: Final = [ @@ -1345,21 +1367,15 @@ async def get_ui_settings(): detail={"error": "Database not connected. Please connect a database."}, ) - ui_settings: Mapping[str, JsonValue] = {} - db_record: Final = await _ui_settings_db(UISettingsRepository(prisma_client)).find_unique( where={"id": "ui_settings"} ) - if db_record and db_record.ui_settings: - ui_settings_json: Final = db_record.ui_settings - if isinstance(ui_settings_json, str): - ui_settings = json.loads(ui_settings_json) - else: - ui_settings = dict(ui_settings_json) + stored: Final = (db_record.ui_settings if db_record else None) or "{}" + parsed: Final = json.loads(stored) if isinstance(stored, str) else stored # Sanitize any unexpected keys from persisted config before returning - ui_settings = {k: v for k, v in ui_settings.items() if k in ALLOWED_UI_SETTINGS_FIELDS} + ui_settings: Final = {k: v for k, v in parsed.items() if k in ALLOWED_UI_SETTINGS_FIELDS} # Sync runtime flags into general_settings so the proxy picks them up # at runtime (covers server restart scenarios). @@ -1377,11 +1393,18 @@ async def get_ui_settings(): # Build config-like object for schema helper config: Final[dict[str, object]] = {"litellm_settings": {"ui_settings": ui_settings}} - return await _get_settings_with_schema( + settings: Final = await _get_settings_with_schema( settings_key="ui_settings", settings_class=_get_effective_ui_settings_class(), config=config, ) + return UISettingsResponse( + values={ + **settings["values"], + ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: is_ptu_cost_attribution_enabled(), + }, + field_schema=settings["field_schema"], + ) @router.patch( @@ -1418,6 +1441,20 @@ async def update_ui_settings( detail={"error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature."}, ) + conflicting_keys: Final = sorted( + key + for key, value in settings_body.items() + if key in _DERIVED_UI_SETTINGS_FIELDS and value != _derived_ui_setting_value(key) + ) + if conflicting_keys: + raise HTTPException( + status_code=400, + detail=( + f"Setting(s) {conflicting_keys} are derived from the deployment environment " + "and cannot be changed from the UI." + ), + ) + # Validate against the same effective class GET advertises, so # enterprise-registered fields are typed consistently on both sides. effective_cls: Final = _get_effective_ui_settings_class() diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index a388e7aaf09..ab9b4bc3922 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -7,6 +7,8 @@ from unittest.mock import AsyncMock, MagicMock import pytest +from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + sys.path.insert(0, os.path.abspath("../../../..")) # Adds the parent directory to the system path from litellm.proxy.management_endpoints.common_daily_activity import ( @@ -1142,6 +1144,11 @@ class TestEverySavingsDriverSurvivesTheReadPath: ) +@pytest.fixture +def ptu_cost_attribution_enabled(monkeypatch): + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + + def _spend_record(api_key, *, model="gpt-4o-mini-ptu", spend=0.0, ptu_flat_cost=0.0): return SimpleNamespace( api_key=api_key, @@ -1167,13 +1174,13 @@ def _spend_record(api_key, *, model="gpt-4o-mini-ptu", spend=0.0, ptu_flat_cost= ) -def test_update_metrics_accumulates_ptu_flat_cost(): +def test_update_metrics_accumulates_ptu_flat_cost(ptu_cost_attribution_enabled): metrics = update_metrics(SpendMetrics(), _spend_record("real-key", spend=1.0, ptu_flat_cost=240.0)) assert metrics.flat_cost == 240.0 assert metrics.spend == 1.0 -def test_ptu_sentinel_excluded_from_key_breakdown_but_flat_cost_aggregates(): +def test_ptu_sentinel_excluded_from_key_breakdown_but_flat_cost_aggregates(ptu_cost_attribution_enabled): from litellm.constants import PTU_SENTINEL_API_KEY from litellm.proxy.management_endpoints.common_daily_activity import update_breakdown_metrics from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics @@ -1230,7 +1237,7 @@ def _grouping_row( ) -def test_grouping_sets_dispatcher_excludes_ptu_sentinel_from_key_breakdowns(): +def test_grouping_sets_dispatcher_excludes_ptu_sentinel_from_key_breakdowns(ptu_cost_attribution_enabled): """The GROUPING SETS path must mirror the per-row path: the flat-cost sentinel aggregates into the date/model/total metrics but never surfaces as an api_key.""" from litellm.constants import PTU_SENTINEL_API_KEY @@ -1267,7 +1274,7 @@ def test_grouping_sets_dispatcher_excludes_ptu_sentinel_from_key_breakdowns(): assert "real-key" in model_bucket.api_key_breakdown -def test_grouping_sets_dispatcher_populates_every_breakdown_level(): +def test_grouping_sets_dispatcher_populates_every_breakdown_level(ptu_cost_attribution_enabled): """Every GROUPING SETS level lands in its bucket, and the flat-cost sentinel is kept out of the model_group and provider api_key sub-breakdowns too.""" from litellm.constants import PTU_SENTINEL_API_KEY @@ -1359,7 +1366,7 @@ def test_grouping_sets_dispatcher_keeps_a_real_provider_row_that_shares_the_sent assert unknown.metrics.flat_cost == 0.0 -def test_update_breakdown_metrics_covers_mcp_endpoint_and_entity(): +def test_update_breakdown_metrics_covers_mcp_endpoint_and_entity(ptu_cost_attribution_enabled): """A full request record fans out into the mcp, endpoint, provider and entity breakdowns, while the flat-cost sentinel stays out of the entity api_key sub-map.""" from litellm.constants import PTU_SENTINEL_API_KEY @@ -1432,6 +1439,10 @@ class TestSentinelRowsDisplayTheirModelName: """A sentinel row keys on the deployment id so a rename cannot move it. The usage views render the breakdown key directly as a label, so the read path has to show the name.""" + @pytest.fixture(autouse=True) + def _enabled(self, ptu_cost_attribution_enabled): + """Flat cost is gated off by default, and these assert on the amounts.""" + @staticmethod def _breakdown(records): from litellm.proxy.management_endpoints.common_daily_activity import update_breakdown_metrics @@ -1485,3 +1496,226 @@ class TestSentinelRowsDisplayTheirModelName: models = self._breakdown([self._sentinel(model_id="dep-1", model_group=None)]).models assert models["dep-1"].metrics.flat_cost == pytest.approx(480.0) + + +def _daily_team_row(api_key, *, spend=0.0, ptu_flat_cost=0.0): + """A LiteLLM_DailyTeamSpend row as the paginated read path receives it from find_many.""" + base: Final = _spend_record(api_key, spend=spend, ptu_flat_cost=ptu_flat_cost) + return SimpleNamespace(**{**base.__dict__, "date": "2026-07-01", "team_id": "team-1"}) + + +class TestPtuCostAttributionDisabled: + """With LITELLM_ENABLE_PTU_COST_ATTRIBUTION unset, both read paths report zero flat + cost, while the sentinel filtering that keeps ``__ptu_flat_cost__`` out of the + breakdowns keeps running. + + Filtering is deliberately not gated: an operator can enable the flag, accrue + sentinel rows, then disable it, and those rows stay in LiteLLM_DailyTeamSpend + forever. Gating the filter too would surface the sentinel as a bogus api_key and + mint a provider bucket for its empty provider. + """ + + @pytest.fixture(autouse=True) + def _flag_off(self, monkeypatch): + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + + def test_paginated_path_reports_zero_flat_cost(self): + metrics = update_metrics(SpendMetrics(), _spend_record("real-key", spend=1.0, ptu_flat_cost=240.0)) + + assert metrics.flat_cost == 0.0 + assert metrics.spend == 1.0 + + def test_aggregated_path_reports_zero_flat_cost(self): + from litellm.proxy.management_endpoints.common_daily_activity import _GROUP_GRAND_TOTAL + + metrics = _record_to_spend_metrics(_grouping_row(_GROUP_GRAND_TOTAL, spend=5.0, ptu_flat_cost=240.0)) + + assert metrics.flat_cost == 0.0 + assert metrics.spend == 5.0 + + def test_aggregated_totals_and_buckets_report_zero_flat_cost(self): + from litellm.constants import PTU_SENTINEL_API_KEY + from litellm.proxy.management_endpoints.common_daily_activity import ( + _GROUP_DATE_API_KEY, + _GROUP_DATE_MODEL, + _GROUP_GRAND_TOTAL, + _aggregate_grouping_sets_records_sync, + ) + + records = [ + _grouping_row(_GROUP_DATE_API_KEY, api_key=PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0), + _grouping_row(_GROUP_DATE_MODEL, model="gpt-4o-mini-ptu", spend=5.0, ptu_flat_cost=240.0), + _grouping_row(_GROUP_GRAND_TOTAL, spend=5.0, ptu_flat_cost=240.0), + ] + + aggregated = _aggregate_grouping_sets_records_sync(records=records, api_key_metadata={}) + + assert aggregated["totals"].flat_cost == 0.0 + assert aggregated["totals"].spend == 5.0 + assert aggregated["results"][0].breakdown.models["gpt-4o-mini-ptu"].metrics.flat_cost == 0.0 + + def test_sentinel_still_excluded_from_the_api_key_breakdown(self): + from litellm.constants import PTU_SENTINEL_API_KEY + from litellm.proxy.management_endpoints.common_daily_activity import update_breakdown_metrics + from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics + + breakdown = BreakdownMetrics() + update_breakdown_metrics(breakdown, _spend_record("real-key", spend=5.0), {}, {}, {}) + update_breakdown_metrics( + breakdown, _spend_record(PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0), {}, {}, {}, entity_id_field="team_id" + ) + + assert PTU_SENTINEL_API_KEY not in breakdown.api_keys + assert PTU_SENTINEL_API_KEY not in breakdown.models["gpt-4o-mini-ptu"].api_key_breakdown + assert "real-key" in breakdown.models["gpt-4o-mini-ptu"].api_key_breakdown + + def test_sentinel_still_excluded_from_the_provider_breakdown(self): + from litellm.constants import PTU_SENTINEL_API_KEY + from litellm.proxy.management_endpoints.common_daily_activity import update_breakdown_metrics + from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics + + breakdown = BreakdownMetrics() + update_breakdown_metrics(breakdown, _spend_record(PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0), {}, {}, {}) + + assert breakdown.providers == {} + + def test_grouping_sets_sentinel_still_excluded_from_breakdowns(self): + from litellm.constants import PTU_SENTINEL_API_KEY + from litellm.proxy.management_endpoints.common_daily_activity import ( + _GROUP_DATE_API_KEY, + _GROUP_DATE_MODEL, + _GROUP_DATE_MODEL_API_KEY, + _GROUP_DATE_PROVIDER, + _aggregate_grouping_sets_records_sync, + ) + + records = [ + _grouping_row(_GROUP_DATE_API_KEY, api_key=PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0), + _grouping_row(_GROUP_DATE_MODEL, model="gpt-4o-mini-ptu", spend=5.0, ptu_flat_cost=240.0), + _grouping_row( + _GROUP_DATE_MODEL_API_KEY, model="gpt-4o-mini-ptu", api_key=PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0 + ), + _grouping_row(_GROUP_DATE_PROVIDER, custom_llm_provider="", ptu_flat_cost=240.0), + ] + + day = _aggregate_grouping_sets_records_sync(records=records, api_key_metadata={})["results"][0] + + assert PTU_SENTINEL_API_KEY not in day.breakdown.api_keys + assert PTU_SENTINEL_API_KEY not in day.breakdown.models["gpt-4o-mini-ptu"].api_key_breakdown + assert sum(bucket.metrics.flat_cost for bucket in day.breakdown.providers.values()) == 0.0 + + @pytest.mark.asyncio + async def test_team_daily_activity_endpoint_reports_zero_flat_cost(self): + """/team/daily/activity reads rows with find_many rather than the aggregated SQL, so + forcing the SQL select to a constant zero would leave this path reporting flat cost.""" + from litellm.constants import PTU_SENTINEL_API_KEY + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_table = MagicMock() + mock_table.count = AsyncMock(return_value=2) + mock_table.find_many = AsyncMock( + return_value=[ + _daily_team_row("real-key", spend=5.0), + _daily_team_row(PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0), + ] + ) + mock_prisma.db.litellm_verificationtoken = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_dailyteamspend = mock_table + + result = await get_daily_activity( + prisma_client=mock_prisma, + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id="team-1", + entity_metadata_field=None, + start_date="2026-07-01", + end_date="2026-07-01", + model=None, + api_key=None, + page=1, + page_size=50, + ) + + assert result.metadata.total_flat_cost == 0.0 + assert result.metadata.total_spend == 5.0 + assert PTU_SENTINEL_API_KEY not in result.results[0].breakdown.api_keys + + @pytest.mark.asyncio + async def test_team_daily_activity_endpoint_reports_flat_cost_once_enabled(self, monkeypatch): + from litellm.constants import PTU_SENTINEL_API_KEY + + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_table = MagicMock() + mock_table.count = AsyncMock(return_value=2) + mock_table.find_many = AsyncMock( + return_value=[ + _daily_team_row("real-key", spend=5.0), + _daily_team_row(PTU_SENTINEL_API_KEY, ptu_flat_cost=240.0), + ] + ) + mock_prisma.db.litellm_verificationtoken = MagicMock() + mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma.db.litellm_dailyteamspend = mock_table + + result = await get_daily_activity( + prisma_client=mock_prisma, + table_name="litellm_dailyteamspend", + entity_id_field="team_id", + entity_id="team-1", + entity_metadata_field=None, + start_date="2026-07-01", + end_date="2026-07-01", + model=None, + api_key=None, + page=1, + page_size=50, + ) + + assert result.metadata.total_flat_cost == 240.0 + assert PTU_SENTINEL_API_KEY not in result.results[0].breakdown.api_keys + + +class TestFlagIsNotReadOnTheHotPath: + """update_metrics runs once per accumulation and a record fans out across roughly a + dozen breakdowns, so a flag that reads through the secret manager must not be consulted + for rows that carry no flat cost at all.""" + + @staticmethod + def _count_flag_reads(records): + import litellm.proxy.management_endpoints.common_daily_activity as cda + from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics + + reads = [] + real = cda.is_ptu_cost_attribution_enabled + + def counted(): + reads.append(1) + return real() + + cda.is_ptu_cost_attribution_enabled = counted + try: + breakdown = BreakdownMetrics() + for record in records: + cda.update_breakdown_metrics(breakdown, record, {}, {}, {}) + finally: + cda.is_ptu_cost_attribution_enabled = real + return len(reads) + + def test_a_request_row_never_reads_the_flag(self): + reads = self._count_flag_reads([_spend_record("real-key", spend=5.0, ptu_flat_cost=0.0)]) + assert reads == 0, f"{reads} secret-manager lookups for a row with no flat cost" + + def test_a_page_of_request_rows_never_reads_the_flag(self): + rows = [_spend_record(f"key-{i}", spend=1.0, ptu_flat_cost=0.0) for i in range(50)] + assert self._count_flag_reads(rows) == 0 + + def test_a_sentinel_row_still_consults_the_flag(self): + from litellm.constants import PTU_SENTINEL_API_KEY + + reads = self._count_flag_reads([_spend_record(PTU_SENTINEL_API_KEY, spend=0.0, ptu_flat_cost=240.0)]) + assert reads > 0 diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py index e8131854acf..30fe78d93c7 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py @@ -1,34 +1,57 @@ -import datetime -import json - """Tests for PTU config on the model deployment (v1 model-settings design).""" -from unittest.mock import AsyncMock, MagicMock +import datetime +import json +from contextlib import ExitStack +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException +from litellm.proxy._types import LiteLLM_ProxyModelTable, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.management_endpoints.model_management_endpoints import ( _merged_ptu_model_info, + _raise_if_ptu_cost_attribution_disabled, _validate_ptu_model_info, + add_new_model, + update_db_model, ) +from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment def test_model_info_accepts_valid_ptu_fields(): - info = ModelInfo(id="x", team_id="t", ptu_count=5, cost_per_ptu_per_hour=2.0) + info = ModelInfo( + id="x", + team_id="t", + ptu_count=5, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + ) assert info.ptu_count == 5 assert info.cost_per_ptu_per_hour == 2.0 def test_model_info_rejects_non_positive_count(): with pytest.raises(ValueError): - ModelInfo(id="x", team_id="t", ptu_count=0, cost_per_ptu_per_hour=2.0) + ModelInfo( + id="x", + team_id="t", + ptu_count=0, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + ) def test_model_info_rejects_negative_rate(): with pytest.raises(ValueError): - ModelInfo(id="x", team_id="t", ptu_count=5, cost_per_ptu_per_hour=-1.0) + ModelInfo( + id="x", + team_id="t", + ptu_count=5, + cost_per_ptu_per_hour=-1.0, + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + ) def test_model_info_rejects_a_count_beyond_the_cap(): @@ -223,6 +246,11 @@ class TestPartialPtuEditsUseTheMergedView: """A PTU invariant holds over the deployment as it will exist, not over whichever subset of fields a caller sent. Validating the patch alone rejected an ordinary edit.""" + @pytest.fixture(autouse=True) + def _enabled(self, monkeypatch): + """PTU writes are gated off by default; these are about the validator, not the gate.""" + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + @staticmethod def _configured(): return Deployment( @@ -310,6 +338,11 @@ class TestTeamModelUpdateValidatesBeforeWriting: """Drives the endpoint path itself, not the helpers. The validator sits above the team ACL write, which autocommits, so what it validates has to be right at that call site.""" + @pytest.fixture(autouse=True) + def _enabled(self, monkeypatch): + """PTU writes are gated off by default; these are about the validator, not the gate.""" + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + @staticmethod async def _run(db_model, patch_data, monkeypatch, touched=None): import litellm.proxy.management_endpoints.model_management_endpoints as mme @@ -352,6 +385,33 @@ class TestTeamModelUpdateValidatesBeforeWriting: assert "ptu_effective_from is required" in exc.value.detail + @pytest.mark.asyncio + async def test_the_gate_refuses_before_the_team_write(self, monkeypatch): + """The gate lived inside update_db_model, which runs after the team ACL write, so a + rejected edit still moved the model between teams.""" + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + db_model = Deployment( + model_name="gpt-4o", + litellm_params=LiteLLM_Params(model="openai/gpt-4o"), + model_info=ModelInfo(id="dep-0", team_id="team-A"), + ) + patch = updateDeployment( + model_info=ModelInfo( + id="dep-0", + team_id="team-B", + ptu_count=15, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2026, 8, 1, tzinfo=datetime.timezone.utc), + ) + ) + touched = [] + + with pytest.raises(HTTPException) as exc: + await self._run(db_model, patch, monkeypatch, touched) + + assert PTU_COST_ATTRIBUTION_ENV_VAR in exc.value.detail + assert touched == [] + @pytest.mark.asyncio async def test_clearing_half_the_pair_is_refused_before_the_team_write(self, monkeypatch): """The write drops the nulled field, so validating against the stored one let a @@ -380,3 +440,270 @@ class TestTeamModelUpdateValidatesBeforeWriting: stored = json.loads(result["model_info"]) assert "ptu_count" not in stored assert "cost_per_ptu_per_hour" not in stored + + +class TestPtuCostAttributionGate: + """PTU config is only writable once an operator sets LITELLM_ENABLE_PTU_COST_ATTRIBUTION. + + The fields are rejected rather than dropped: a silent accept-and-drop would let a + caller believe a flat cost was configured while the rollup that prices it is not + even scheduled. + """ + + @pytest.fixture(autouse=True) + def _flag_off(self, monkeypatch): + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + + @pytest.fixture + def flag_on(self, monkeypatch): + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + + @pytest.mark.parametrize( + "model_info", + [ + {"team_id": "t", "ptu_count": 5, "cost_per_ptu_per_hour": 2.0}, + {"ptu_count": 5}, + {"cost_per_ptu_per_hour": 2.0}, + {"ptu_effective_from": "2026-08-01T00:00:00Z"}, + {"ptu_effective_to": "2026-08-02T00:00:00Z"}, + ], + ) + def test_rejects_any_ptu_field_while_disabled(self, model_info): + with pytest.raises(HTTPException) as exc: + _raise_if_ptu_cost_attribution_disabled(model_info) + assert exc.value.status_code == 400 + assert PTU_COST_ATTRIBUTION_ENV_VAR in exc.value.detail + + def test_names_every_offending_field(self): + with pytest.raises(HTTPException) as exc: + _raise_if_ptu_cost_attribution_disabled({"ptu_count": 5, "cost_per_ptu_per_hour": 2.0}) + assert "ptu_count" in exc.value.detail + assert "cost_per_ptu_per_hour" in exc.value.detail + + def test_allows_a_request_without_ptu_fields_while_disabled(self): + _raise_if_ptu_cost_attribution_disabled({"team_id": "t", "access_groups": ["a"]}) + + def test_allows_every_ptu_field_once_enabled(self, flag_on): + _raise_if_ptu_cost_attribution_disabled( + { + "team_id": "t", + "ptu_count": 5, + "cost_per_ptu_per_hour": 2.0, + "ptu_effective_from": "2026-08-01T00:00:00Z", + "ptu_effective_to": "2026-08-02T00:00:00Z", + } + ) + + +def _deployment_without_ptu() -> Deployment: + return Deployment( + model_name="gpt-4o", + litellm_params=LiteLLM_Params(model="openai/gpt-4o"), + model_info=ModelInfo(id="dep-0", team_id="t"), + ) + + +def _deployment_with_stored_ptu() -> Deployment: + return Deployment( + model_name="gpt-4o", + litellm_params=LiteLLM_Params(model="openai/gpt-4o"), + model_info=ModelInfo( + id="dep-0", + team_id="t", + ptu_count=15, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + ), + ) + + +class TestUpdateDbModelPtuGate: + @pytest.fixture(autouse=True) + def _flag_off(self, monkeypatch): + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + + def test_patch_carrying_ptu_config_is_rejected(self): + with pytest.raises(HTTPException) as exc: + update_db_model( + db_model=_deployment_without_ptu(), + updated_patch=updateDeployment(model_info=ModelInfo(id="dep-0", team_id="t", ptu_count=15)), + ) + assert exc.value.status_code == 400 + + def test_patch_that_touches_nothing_ptu_still_succeeds(self): + result = update_db_model( + db_model=_deployment_without_ptu(), + updated_patch=updateDeployment(model_info=ModelInfo(id="dep-0", access_groups=["a"])), + ) + assert json.loads(result["model_info"])["access_groups"] == ["a"] + + def test_unrelated_patch_of_a_model_that_stores_ptu_config_is_not_blocked(self): + """A deployment configured during an earlier opt-in stays editable: the gate reads the + incoming patch, not the merged deployment, so the stored config is left in place.""" + result = update_db_model( + db_model=_deployment_with_stored_ptu(), + updated_patch=updateDeployment(model_name="gpt-4o-renamed"), + ) + assert result["model_name"] == "gpt-4o-renamed" + + def test_explicit_nulls_do_not_erase_stored_ptu_config_while_disabled(self): + """A client round-tripping a model_info blob sends the PTU keys as nulls. While the + feature is disabled those nulls must not reach the clear loop: disabling pauses PTU, + it does not silently discard a billing configuration the operator set up earlier.""" + result = update_db_model( + db_model=_deployment_with_stored_ptu(), + updated_patch=updateDeployment( + model_info=ModelInfo(id="dep-0", ptu_count=None, cost_per_ptu_per_hour=None) + ), + ) + stored = json.loads(result["model_info"]) + assert stored["ptu_count"] == 15 + assert stored["cost_per_ptu_per_hour"] == 2.0 + + def test_the_merged_view_agrees_with_the_write_while_disabled(self): + """The validator sees what the write will store. If the merged view honoured a null the + clear loop ignores, a round-tripped blob would 400 on a half-set pair that never forms.""" + merged = _merged_ptu_model_info( + db_model=_deployment_with_stored_ptu(), + patch_data=updateDeployment(model_info=ModelInfo(id="dep-0", ptu_count=None)), + ) + assert merged["ptu_count"] == 15 + _validate_ptu_model_info(merged) + + def test_explicit_nulls_still_clear_once_enabled(self, monkeypatch): + """Clearing remains available to an operator who opted in, which is how PTU config is + removed from a deployment.""" + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + result = update_db_model( + db_model=_deployment_with_stored_ptu(), + updated_patch=updateDeployment( + model_info=ModelInfo(id="dep-0", ptu_count=None, cost_per_ptu_per_hour=None) + ), + ) + stored = json.loads(result["model_info"]) + assert "ptu_count" not in stored + assert "cost_per_ptu_per_hour" not in stored + + def test_patch_carrying_ptu_config_is_accepted_once_enabled(self, monkeypatch): + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + result = update_db_model( + db_model=_deployment_without_ptu(), + updated_patch=updateDeployment( + model_info=ModelInfo( + id="dep-0", + team_id="t", + ptu_count=15, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + ) + ), + ) + stored = json.loads(result["model_info"]) + assert stored["ptu_count"] == 15 + assert stored["cost_per_ptu_per_hour"] == 2.0 + + +class TestAddNewModelPtuGate: + @pytest.fixture(autouse=True) + def _flag_off(self, monkeypatch): + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + + @staticmethod + def _patched_proxy(model_id: str): + """Patch everything /model/new touches except the PTU gate, and hand back the DB writers.""" + db_row = LiteLLM_ProxyModelTable( + model_id=model_id, + model_name="ptu-model", + litellm_params={"model": "openai/gpt-4.1-nano"}, + model_info={"id": model_id}, + created_by="test-admin", + updated_by="test-admin", + ) + add_model_to_db = AsyncMock(return_value=db_row) + add_team_model_to_db = AsyncMock(return_value=db_row) + + mock_proxy_config = MagicMock() + mock_proxy_config.add_deployment = AsyncMock(return_value=None) + + mock_router = MagicMock() + mock_router.get_model_ids.return_value = [model_id] + + proxy_server = "litellm.proxy.proxy_server" + endpoints = "litellm.proxy.management_endpoints.model_management_endpoints" + return (add_model_to_db, add_team_model_to_db), [ + patch(f"{proxy_server}.prisma_client", MagicMock()), + patch(f"{proxy_server}.store_model_in_db", True), + patch(f"{proxy_server}.proxy_config", mock_proxy_config), + patch(f"{proxy_server}.proxy_logging_obj", MagicMock()), + patch(f"{proxy_server}.general_settings", {}), + patch(f"{proxy_server}.premium_user", True), + patch(f"{proxy_server}.llm_router", mock_router), + patch( + f"{endpoints}.ModelManagementAuthChecks.can_user_make_model_call", + AsyncMock(return_value=True), + ), + patch(f"{endpoints}._add_model_to_db", add_model_to_db), + patch(f"{endpoints}._add_team_model_to_db", add_team_model_to_db), + ] + + @staticmethod + def _ptu_deployment(model_id: str) -> Deployment: + return Deployment( + model_name="ptu-model", + litellm_params=LiteLLM_Params(model="openai/gpt-4.1-nano", api_key="fake-key"), + model_info=ModelInfo( + id=model_id, + team_id="team-1", + ptu_count=15, + cost_per_ptu_per_hour=2.0, + ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + ), + ) + + @pytest.mark.asyncio + async def test_model_new_rejects_ptu_config_while_disabled(self): + (add_model_to_db, add_team_model_to_db), patches = self._patched_proxy("ptu-gate-model") + admin = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + with ExitStack() as stack: + for active_patch in patches: + stack.enter_context(active_patch) + with pytest.raises(Exception) as exc: + await add_new_model(model_params=self._ptu_deployment("ptu-gate-model"), user_api_key_dict=admin) + + assert PTU_COST_ATTRIBUTION_ENV_VAR in str(exc.value) + add_model_to_db.assert_not_called() + add_team_model_to_db.assert_not_called() + + @pytest.mark.asyncio + async def test_model_new_accepts_a_deployment_without_ptu_config_while_disabled(self): + _, patches = self._patched_proxy("plain-model") + admin = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + with ExitStack() as stack: + for active_patch in patches: + stack.enter_context(active_patch) + result = await add_new_model( + model_params=Deployment( + model_name="ptu-model", + litellm_params=LiteLLM_Params(model="openai/gpt-4.1-nano", api_key="fake-key"), + model_info=ModelInfo(id="plain-model"), + ), + user_api_key_dict=admin, + ) + + assert result.model_id == "plain-model" + + @pytest.mark.asyncio + async def test_model_new_accepts_ptu_config_once_enabled(self, monkeypatch): + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + (_, add_team_model_to_db), patches = self._patched_proxy("ptu-gate-model") + admin = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + with ExitStack() as stack: + for active_patch in patches: + stack.enter_context(active_patch) + result = await add_new_model(model_params=self._ptu_deployment("ptu-gate-model"), user_api_key_dict=admin) + + assert result.model_id == "ptu-gate-model" + add_team_model_to_db.assert_called_once() diff --git a/tests/test_litellm/proxy/spend_tracking/test_ptu_feature_flag.py b/tests/test_litellm/proxy/spend_tracking/test_ptu_feature_flag.py new file mode 100644 index 00000000000..7f4bd935a2b --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_ptu_feature_flag.py @@ -0,0 +1,33 @@ +"""Tests for the opt-in flag that gates PTU flat-cost attribution.""" + +import pytest + +from litellm.proxy.spend_tracking.ptu_feature_flag import ( + PTU_COST_ATTRIBUTION_ENV_VAR, + is_ptu_cost_attribution_enabled, +) + + +def test_disabled_when_env_var_is_unset(monkeypatch): + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + assert is_ptu_cost_attribution_enabled() is False + + +@pytest.mark.parametrize("value", ["true", "True", "TRUE", " true "]) +def test_enabled_for_the_values_the_house_helper_recognises(monkeypatch, value): + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, value) + assert is_ptu_cost_attribution_enabled() is True + + +@pytest.mark.parametrize("value", ["false", "False", "0", "1", "", "yes", "off", "maybe"]) +def test_disabled_for_everything_else(monkeypatch, value): + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, value) + assert is_ptu_cost_attribution_enabled() is False + + +def test_reads_the_env_var_on_every_call(monkeypatch): + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + assert is_ptu_cost_attribution_enabled() is False + + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + assert is_ptu_cost_attribution_enabled() is True diff --git a/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py index 261e3f4ef76..d17f6293cc3 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py +++ b/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py @@ -8,6 +8,7 @@ import pytest import litellm.proxy.spend_tracking.ptu_flat_cost_rollup as ptu_rollup from litellm.constants import PTU_ROLLUP_MAX_BACKFILL_DAYS, PTU_SENTINEL_API_KEY +from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR from litellm.types.router import ModelInfo from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import ( PTUModel, @@ -29,6 +30,13 @@ TODAY = date(2026, 7, 31) _DEFAULT_PTU_START = "2020-01-01T00:00:00Z" +@pytest.fixture(autouse=True) +def _ptu_enabled(monkeypatch): + """PTU is gated off by default. These cover the rollup's mechanics, not the gate, so + they run with it on; the gate itself is covered by its own test below.""" + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + + _VALID_PTU = {"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"} @@ -1479,3 +1487,18 @@ async def test_the_prune_cutoff_allows_for_clock_skew_between_hosts(): "a charge written 30s ago by a lagging pod was swept" ) assert ("t", DAY.isoformat(), PTU_SENTINEL_API_KEY, "dep-stale") not in table.rows + + +@pytest.mark.asyncio +async def test_scheduled_rollup_writes_nothing_when_ptu_attribution_is_disabled(monkeypatch): + """Startup already skips scheduling the cron, so this guards the function itself: a + deployment that never opted in accrues nothing whatever route reaches the rollup.""" + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + table = _FakeSentinelTable() + prisma = _prisma_for([_model_row(model_info=_VALID_PTU)], table) + + result = await run_scheduled_ptu_rollup(prisma, pod_lock_manager=None, alert=None) + + assert result is None + assert table.rows == {} + assert table.upsert_keys == [] diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 74120beb7b9..87c8c180d9e 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -11280,14 +11280,8 @@ async def test_setup_prisma_client_returns_none_when_connect_itself_fails(monkey assert mock_client.health_check.await_count == 0 -@pytest.mark.asyncio -async def test_ptu_rollup_job_registered_at_startup(monkeypatch): - """The PTU rollup cron is registered at startup; only models with PTU config accrue flat cost (asserted in test_ptu_flat_cost_rollup.py).""" - monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) +async def _run_scheduled_background_jobs(): from litellm.proxy.proxy_server import ProxyStartupEvent - from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import ( - PTU_ROLLUP_JOB_ID, - ) from litellm.proxy.utils import ProxyLogging mock_prisma_client = MagicMock() @@ -11311,7 +11305,41 @@ async def test_ptu_rollup_job_registered_at_startup(monkeypatch): proxy_logging_obj=mock_proxy_logging, ) - import litellm.proxy.proxy_server as ps + 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 + assert ps.scheduler is not None + return ps.scheduler + + +@pytest.mark.asyncio +async def test_ptu_rollup_job_registered_at_startup(monkeypatch): + """The PTU rollup cron is registered once an operator opts in; only models with PTU config accrue flat cost (asserted in test_ptu_flat_cost_rollup.py).""" + monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) + from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import ( + PTU_ROLLUP_JOB_ID, + ) + + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + + scheduler = await _run_scheduled_background_jobs() + + assert scheduler.get_job(PTU_ROLLUP_JOB_ID) is not None + + +@pytest.mark.asyncio +async def test_ptu_rollup_job_not_registered_without_opt_in(monkeypatch): + """Without LITELLM_ENABLE_PTU_COST_ATTRIBUTION the rollup never runs, so no sentinel row + is ever written. This is the gate that keeps the whole feature inert by default.""" + monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) + from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + from litellm.proxy.spend_tracking.ptu_flat_cost_rollup import ( + PTU_ROLLUP_JOB_ID, + ) + + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + + scheduler = await _run_scheduled_background_jobs() + + assert scheduler.get_job(PTU_ROLLUP_JOB_ID) is None + assert len(scheduler.get_jobs()) > 0 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 1075bffbeb2..8ee4e92ca9b 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 @@ -2928,3 +2928,143 @@ def test_update_mcp_semantic_filter_settings_requires_proxy_admin(monkeypatch): assert "proxy admin" in resp.json()["detail"].lower() finally: app.dependency_overrides.pop(user_api_key_auth, None) + + +class TestPtuCostAttributionUISetting: + """``enable_ptu_cost_attribution`` is derived from the environment on every GET. + + It is deliberately not an allowlisted, persisted setting: the point of gating PTU + flat cost on an env var is that an admin cannot flip it at runtime from the UI. + """ + + @staticmethod + def _mock_prisma(monkeypatch, stored=None): + from unittest.mock import AsyncMock, MagicMock + + mock_prisma = MagicMock() + mock_record = None + if stored is not None: + mock_record = MagicMock() + mock_record.ui_settings = stored + mock_prisma.db.litellm_uisettings.find_unique = AsyncMock(return_value=mock_record) + mock_prisma.db.litellm_uisettings.upsert = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma) + return mock_prisma + + def test_reported_false_when_the_env_var_is_unset(self, mock_auth, monkeypatch): + from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + self._mock_prisma(monkeypatch) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + assert response.json()["values"]["enable_ptu_cost_attribution"] is False + + def test_reported_true_once_the_env_var_is_set(self, mock_auth, monkeypatch): + from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + + monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true") + self._mock_prisma(monkeypatch) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + assert response.json()["values"]["enable_ptu_cost_attribution"] is True + + def test_a_persisted_true_cannot_forge_the_derived_value(self, mock_auth, monkeypatch): + """A row written before the allowlist existed must not be able to turn the feature on.""" + from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + self._mock_prisma(monkeypatch, stored={"enable_ptu_cost_attribution": True}) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + assert response.json()["values"]["enable_ptu_cost_attribution"] is False + + def test_is_not_an_allowlisted_persisted_setting(self): + from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( + ALLOWED_UI_SETTINGS_FIELDS, + ) + + assert "enable_ptu_cost_attribution" not in ALLOWED_UI_SETTINGS_FIELDS + + def test_the_body_get_returns_is_a_valid_patch_body(self, mock_auth, monkeypatch): + """Read-modify-write is how a client edits one setting. GET injects the derived key, + so rejecting it on presence made GET's own output an invalid PATCH body: the caller + got a 400 and silently lost the edit it actually wanted.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + mock_prisma = self._mock_prisma(monkeypatch) + + try: + round_tripped = client.get("/get/ui_settings").json()["values"] + assert "enable_ptu_cost_attribution" in round_tripped + response = client.patch("/update/ui_settings", json=round_tripped) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 200 + assert mock_prisma.db.litellm_uisettings.upsert.called + + def test_a_co_submitted_setting_still_applies_alongside_the_derived_key(self, mock_auth, monkeypatch): + """The derived key riding along must not discard the caller's real edit.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + mock_prisma = self._mock_prisma(monkeypatch) + + try: + response = client.patch( + "/update/ui_settings", + json={"enable_ptu_cost_attribution": False, "enable_chat_ui": True}, + ) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 200 + upsert_data = mock_prisma.db.litellm_uisettings.upsert.call_args.kwargs["data"] + persisted = json.loads(upsert_data["create"]["ui_settings"]) + assert persisted["enable_chat_ui"] is True + assert "enable_ptu_cost_attribution" not in persisted + + def test_patch_rejects_the_derived_setting(self, mock_auth, monkeypatch): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + mock_prisma = self._mock_prisma(monkeypatch) + + try: + response = client.patch( + "/update/ui_settings", + json={"enable_ptu_cost_attribution": True}, + ) + finally: + app.dependency_overrides.clear() + + assert response.status_code == 400 + assert "enable_ptu_cost_attribution" in str(response.json()["detail"]) + assert not mock_prisma.db.litellm_uisettings.upsert.called diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled.test.ts new file mode 100644 index 00000000000..2215817b618 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled.test.ts @@ -0,0 +1,135 @@ +import { getUiSettings } from "@/components/networking"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { renderHook, waitFor } from "@testing-library/react"; +import React, { ReactNode } from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { PTU_FLAG_REFRESH_MS, usePtuCostAttributionEnabled } from "./usePtuCostAttributionEnabled"; +import { useUISettings } from "./useUISettings"; + +vi.mock("@/components/networking", () => ({ + getUiSettings: vi.fn(), +})); + +describe("usePtuCostAttributionEnabled", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + vi.clearAllMocks(); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + /** Read the flag alongside the query it derives from, so assertions wait for a settled fetch. */ + const renderSettledFlag = async (settings: unknown) => { + (getUiSettings as any).mockResolvedValue(settings); + const { result } = renderHook(() => ({ enabled: usePtuCostAttributionEnabled(), query: useUISettings() }), { + wrapper, + }); + await waitFor(() => { + expect(result.current.query.isSuccess).toBe(true); + }); + return result; + }; + + it("is true only when the proxy reports the flag as enabled", async () => { + const result = await renderSettledFlag({ values: { enable_ptu_cost_attribution: true } }); + expect(result.current.enabled).toBe(true); + }); + + it("is false when the proxy reports the flag as disabled", async () => { + const result = await renderSettledFlag({ values: { enable_ptu_cost_attribution: false } }); + expect(result.current.enabled).toBe(false); + }); + + it("is false when the proxy omits the flag entirely", async () => { + const result = await renderSettledFlag({ values: { enable_chat_ui: true } }); + expect(result.current.enabled).toBe(false); + }); + + it("is false when the proxy returns no values at all", async () => { + const result = await renderSettledFlag({}); + expect(result.current.enabled).toBe(false); + }); + + it("does not treat a truthy non-boolean as enabled", async () => { + const result = await renderSettledFlag({ values: { enable_ptu_cost_attribution: "false" } }); + expect(result.current.enabled).toBe(false); + }); + + it("does not treat the string 'true' as enabled, since the proxy sends a real boolean", async () => { + const result = await renderSettledFlag({ values: { enable_ptu_cost_attribution: "true" } }); + expect(result.current.enabled).toBe(false); + }); + + it("is false before the settings request resolves", () => { + (getUiSettings as any).mockReturnValue(new Promise(() => {})); + const { result } = renderHook(() => usePtuCostAttributionEnabled(), { wrapper }); + expect(result.current).toBe(false); + }); + + it("is false when the settings request fails", async () => { + (getUiSettings as any).mockRejectedValue(new Error("boom")); + const { result } = renderHook(() => ({ enabled: usePtuCostAttributionEnabled(), query: useUISettings() }), { + wrapper, + }); + await waitFor(() => { + expect(result.current.query.isError).toBe(true); + }); + expect(result.current.enabled).toBe(false); + }); +}); + +describe("staleness", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + vi.clearAllMocks(); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("polls the flag once it is on, so an already-open dashboard notices it going off", async () => { + (getUiSettings as any).mockResolvedValue({ values: { enable_ptu_cost_attribution: true } }); + const { result } = renderHook(() => usePtuCostAttributionEnabled(), { wrapper }); + await waitFor(() => { + expect(result.current).toBe(true); + }); + + const observers = queryClient.getQueryCache().getAll()[0].observers; + const polling = observers.filter((o: any) => o.options.refetchInterval === PTU_FLAG_REFRESH_MS); + expect(polling.length).toBeGreaterThan(0); + expect(polling[0].options.staleTime).toBe(PTU_FLAG_REFRESH_MS); + expect(PTU_FLAG_REFRESH_MS).toBeLessThan(60 * 60 * 1000); + }); + + it("does not poll while the flag is off, which is every deployment that never opted in", async () => { + // The hook cannot gate on the flag before reading it, so it starts on the shared + // one-hour cache and only escalates once it has seen the feature enabled. Polling + // unconditionally made a disabled deployment re-fetch settings 120x more often. + (getUiSettings as any).mockResolvedValue({ values: { enable_ptu_cost_attribution: false } }); + const { result } = renderHook(() => usePtuCostAttributionEnabled(), { wrapper }); + await waitFor(() => { + expect(result.current).toBe(false); + }); + + const observers = queryClient.getQueryCache().getAll()[0].observers; + expect(observers.every((o: any) => o.options.refetchInterval === undefined)).toBe(true); + expect(observers.every((o: any) => o.options.staleTime === 60 * 60 * 1000)).toBe(true); + }); + + it("leaves the default alone for every other settings consumer", async () => { + (getUiSettings as any).mockResolvedValue({ values: {} }); + const { result } = renderHook(() => useUISettings(), { wrapper }); + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + const observers = queryClient.getQueryCache().getAll()[0].observers; + expect(observers[0].options.staleTime).toBe(60 * 60 * 1000); + expect(observers[0].options.refetchInterval).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled.ts new file mode 100644 index 00000000000..e9b5afac562 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled.ts @@ -0,0 +1,26 @@ +import { useUISettings } from "./useUISettings"; + +export const PTU_COST_ATTRIBUTION_SETTING_KEY = "enable_ptu_cost_attribution"; + +/** + * Whether the proxy opted into PTU flat-cost attribution. + * + * Derived on the proxy from LITELLM_ENABLE_PTU_COST_ATTRIBUTION and returned read-only on + * /get/ui_settings, so it is not editable from the UI. Anything other than an explicit + * true (including a settings fetch that has not resolved) counts as off. + * + * Polled only once the flag has been seen on. This tracks the proxy process rather than a + * persisted setting, so an already-open dashboard has to notice a restart that turns the + * feature off, and a form that stays mounted and focused never refetches on staleTime + * alone. A deployment that never opts in is the common case and gets the shared one-hour + * cache, so the poll costs nothing where the feature is unused; the trade is that turning + * it on reaches an open dashboard on the next natural refetch rather than within 30s. + */ +export const PTU_FLAG_REFRESH_MS = 30 * 1000; + +export const usePtuCostAttributionEnabled = (): boolean => { + const { data } = useUISettings(); + const enabled = data?.values?.[PTU_COST_ATTRIBUTION_SETTING_KEY] === true; + useUISettings(enabled ? { staleTime: PTU_FLAG_REFRESH_MS, refetchInterval: PTU_FLAG_REFRESH_MS } : undefined); + return enabled; +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.ts index 14c6c5e3888..749fc98c0d8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.ts @@ -4,11 +4,21 @@ import { createQueryKeys } from "../common/queryKeysFactory"; const uiSettingsKeys = createQueryKeys("uiSettings"); -export const useUISettings = () => { +/** + * UI settings, cached for an hour by default because they rarely change. + * + * Both options are per observer in react-query, so a caller reading a value that tracks + * proxy process state, rather than a persisted setting, can refresh it on its own cadence + * without changing how long every other caller caches. `staleTime` alone only marks the + * cached copy stale; a screen that stays mounted and focused never refetches on its own, + * so a caller that needs to notice a change also has to poll. + */ +export const useUISettings = (options?: { staleTime?: number; refetchInterval?: number }) => { return useQuery>({ queryKey: uiSettingsKeys.list({}), queryFn: async () => await getUiSettings(), - staleTime: 60 * 60 * 1000, // 1 hour - data rarely changes + staleTime: options?.staleTime ?? 60 * 60 * 1000, // 1 hour - data rarely changes gcTime: 60 * 60 * 1000, // 1 hour - keep in cache for 1 hour + refetchInterval: options?.refetchInterval, }); }; diff --git a/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx b/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx index 9fe36e13998..4e5c5f25374 100644 --- a/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/advanced_settings.test.tsx @@ -2,32 +2,37 @@ import { act, fireEvent, render, waitFor } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import AdvancedSettings from "./advanced_settings"; +const mockUsePtuCostAttributionEnabled = vi.fn(); + +vi.mock("@/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled", () => ({ + usePtuCostAttributionEnabled: () => mockUsePtuCostAttributionEnabled(), +})); + +const PTU_LABELS = ["PTU Count", "Calculated Cost per PTU / Hour (USD)", "PTU Effective From (UTC)"]; + +const renderAdvancedSettings = () => + render( + {}} + guardrailsList={[]} + tagsList={{}} + accessToken="test-token" + />, + ); + describe("AdvancedSettings", () => { beforeEach(() => { vi.clearAllMocks(); + mockUsePtuCostAttributionEnabled.mockReturnValue(false); }); + it("should render", () => { - render( - {}} - guardrailsList={[]} - tagsList={{}} - accessToken="test-token" - />, - ); + renderAdvancedSettings(); }); it("should render tags list", async () => { - const { getByText } = render( - {}} - guardrailsList={[]} - tagsList={{}} - accessToken="test-token" - />, - ); + const { getByText } = renderAdvancedSettings(); fireEvent.click(getByText("Advanced Settings")); await waitFor(() => { expect(getByText("Tags")).toBeInTheDocument(); @@ -35,15 +40,7 @@ describe("AdvancedSettings", () => { }); it("should render the litellm params", async () => { - const { getByText } = render( - {}} - guardrailsList={[]} - tagsList={{}} - accessToken="test-token" - />, - ); + const { getByText } = renderAdvancedSettings(); act(() => { fireEvent.click(getByText("Advanced Settings")); }); @@ -51,4 +48,35 @@ describe("AdvancedSettings", () => { expect(getByText("LiteLLM Params")).toBeInTheDocument(); }); }); + + it("hides every PTU field when PTU cost attribution is disabled", async () => { + const { getByText, queryByText } = renderAdvancedSettings(); + act(() => { + fireEvent.click(getByText("Advanced Settings")); + }); + await waitFor(() => { + expect(getByText("Tags")).toBeInTheDocument(); + }); + + for (const label of PTU_LABELS) { + expect(queryByText(label)).not.toBeInTheDocument(); + } + expect(queryByText("PTU Effective To (UTC)")).not.toBeInTheDocument(); + }); + + it("shows every PTU field when PTU cost attribution is enabled", async () => { + mockUsePtuCostAttributionEnabled.mockReturnValue(true); + const { getByText } = renderAdvancedSettings(); + act(() => { + fireEvent.click(getByText("Advanced Settings")); + }); + + await waitFor(() => { + expect(getByText("PTU Count")).toBeInTheDocument(); + }); + for (const label of PTU_LABELS) { + expect(getByText(label)).toBeInTheDocument(); + } + expect(getByText("PTU Effective To (UTC)")).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx b/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx index 9bc78783a57..5d7196a84fe 100644 --- a/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx +++ b/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx @@ -20,6 +20,7 @@ import { ptuWindowOrderRule, PTU_END_FIELD, } from "../../utils/ptuValidation"; +import { usePtuCostAttributionEnabled } from "@/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled"; const { Link } = Typography; interface AdvancedSettingsProps { @@ -43,6 +44,7 @@ const AdvancedSettings: React.FC = ({ const [customPricing, setCustomPricing] = React.useState(false); const [pricingModel, setPricingModel] = React.useState<"per_token" | "per_second">("per_token"); const [showCacheControl, setShowCacheControl] = React.useState(false); + const ptuCostAttributionEnabled = usePtuCostAttributionEnabled(); // Add validation function for numbers const validateNumber = (_: any, value: string) => { @@ -193,49 +195,53 @@ const AdvancedSettings: React.FC = ({ /> - - - + {ptuCostAttributionEnabled && ( + <> + + + - - - + + + - - - + + + - - - + + + + + )} {customPricing && (
diff --git a/ui/litellm-dashboard/src/components/model_info_view.test.tsx b/ui/litellm-dashboard/src/components/model_info_view.test.tsx index 326b49ff896..a3ce40494eb 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.test.tsx @@ -47,6 +47,11 @@ vi.mock("@/app/(dashboard)/hooks/models/useModelCostMap", () => ({ useModelCostMap: (...args: any[]) => mockUseModelCostMap(...args), })); +const mockUsePtuCostAttributionEnabled = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled", () => ({ + usePtuCostAttributionEnabled: () => mockUsePtuCostAttributionEnabled(), +})); + const mockNotificationsManager = vi.mocked(NotificationsManager); const mockModelInfoV1Call = vi.mocked(networking.modelInfoV1Call); const mockCredentialGetCall = vi.mocked(networking.credentialGetCall); @@ -99,6 +104,7 @@ describe("ModelInfoView", () => { }, }); vi.clearAllMocks(); + mockUsePtuCostAttributionEnabled.mockReturnValue(false); mockUseModelsInfo.mockReturnValue({ data: { @@ -608,6 +614,100 @@ describe("ModelInfoView", () => { expect(updatePayload.litellm_params).not.toHaveProperty("vector_store_ids"); }); + describe("PTU cost attribution gate", () => { + const ptuModelData = { + ...defaultModelData, + model_info: { + ...defaultModelData.model_info, + team_id: "team-1", + ptu_count: 15, + cost_per_ptu_per_hour: 2, + ptu_effective_from: "2026-07-01T00:00:00+00:00", + ptu_effective_to: "2026-08-01T00:00:00+00:00", + }, + }; + + const renderWithPtuModel = () => { + mockUseModelsInfo.mockReturnValue({ data: { data: [ptuModelData] }, isLoading: false, error: null }); + mockModelInfoV1Call.mockResolvedValue({ data: [ptuModelData] }); + return render(, { wrapper }); + }; + + it("hides the PTU fields when disabled, even for a model that already stores PTU config", async () => { + renderWithPtuModel(); + + await waitFor(() => { + expect(screen.getByText("Model Settings")).toBeInTheDocument(); + }); + + expect(screen.queryByText("PTU Count")).not.toBeInTheDocument(); + expect(screen.queryByText("Cost per PTU / Hour (USD)")).not.toBeInTheDocument(); + expect(screen.queryByText("PTU Effective From (UTC)")).not.toBeInTheDocument(); + expect(screen.queryByText("PTU Effective To (UTC)")).not.toBeInTheDocument(); + }); + + it("shows the PTU fields when enabled", async () => { + mockUsePtuCostAttributionEnabled.mockReturnValue(true); + renderWithPtuModel(); + + await waitFor(() => { + expect(screen.getByText("PTU Count")).toBeInTheDocument(); + }); + expect(screen.getByText("Cost per PTU / Hour (USD)")).toBeInTheDocument(); + expect(screen.getByText("PTU Effective From (UTC)")).toBeInTheDocument(); + expect(screen.getByText("PTU Effective To (UTC)")).toBeInTheDocument(); + }); + + it("omits PTU fields from the save payload when disabled, so an unrelated edit cannot clear stored config", async () => { + const user = userEvent.setup(); + renderWithPtuModel(); + + 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.getByRole("button", { name: /save changes/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(mockModelPatchUpdateCall).toHaveBeenCalled(); + }); + + const modelInfo = mockModelPatchUpdateCall.mock.calls[0][1].model_info; + expect(modelInfo).not.toHaveProperty("ptu_count"); + expect(modelInfo).not.toHaveProperty("cost_per_ptu_per_hour"); + expect(modelInfo).not.toHaveProperty("ptu_effective_from"); + expect(modelInfo).not.toHaveProperty("ptu_effective_to"); + }); + + it("sends the PTU fields on save when enabled", async () => { + mockUsePtuCostAttributionEnabled.mockReturnValue(true); + const user = userEvent.setup(); + renderWithPtuModel(); + + 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.getByRole("button", { name: /save changes/i })).toBeInTheDocument(); + }); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(mockModelPatchUpdateCall).toHaveBeenCalled(); + }); + + const modelInfo = mockModelPatchUpdateCall.mock.calls[0][1].model_info; + expect(modelInfo.ptu_count).toBe(15); + expect(modelInfo.cost_per_ptu_per_hour).toBe(2); + }); + }); + it("should not include input_cost_per_token or output_cost_per_token in update payload when user does not touch cost fields", async () => { // Regression: editing a model without touching cost fields used to inject // input_cost_per_token: 0 and output_cost_per_token: 0 into litellm_params, diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx index bd8da531a8f..e1a4d4311ab 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.tsx @@ -18,7 +18,9 @@ import { Button as TremorButton, } from "@tremor/react"; import { Button, DatePicker, Form, Input, Modal, Select, Tooltip } from "antd"; -import { formatPtuUtcDisplay, ptuPickerToUtcIso, utcIsoToPickerValue } from "../utils/ptuDatetime"; +import { formatPtuUtcDisplay, utcIsoToPickerValue } from "../utils/ptuDatetime"; +import { applyPtuModelInfo } from "../utils/ptuModelInfo"; +import { usePtuCostAttributionEnabled } from "@/app/(dashboard)/hooks/uiSettings/usePtuCostAttributionEnabled"; import { PTU_COUNT_FIELD, PTU_RATE_FIELD, @@ -224,6 +226,7 @@ export default function ModelInfoView({ const { data: modelCostMapData } = useModelCostMap(); const { data: modelHubData } = useModelHub(); const { data: teams } = useTeams(); + const ptuCostAttributionEnabled = usePtuCostAttributionEnabled(); // Transform the model data const getProviderFromModel = (model: string) => { @@ -495,15 +498,7 @@ export default function ModelInfoView({ health_check_model: values.health_check_model, }; } - const ptuNumber = (val: string | number | null | undefined): number | null => - val !== undefined && val !== null && val !== "" ? Number(val) : null; - updatedModelInfo = { - ...updatedModelInfo, - ptu_count: ptuNumber(values.ptu_count), - cost_per_ptu_per_hour: ptuNumber(values.cost_per_ptu_per_hour), - ptu_effective_from: ptuPickerToUtcIso(values.ptu_effective_from), - ptu_effective_to: ptuPickerToUtcIso(values.ptu_effective_to), - }; + updatedModelInfo = applyPtuModelInfo(updatedModelInfo, values, ptuCostAttributionEnabled); } catch (e) { NotificationsManager.fromBackend("Invalid JSON in Model Info"); return; @@ -953,45 +948,46 @@ export default function ModelInfoView({ )}
- {PTU_EDIT_FIELDS.map((ptuField) => { - const { name, label, input, placeholder, isCount, isRate, isStart, pairedWith } = ptuField; - const { windowPeer, bound } = ptuField; - return ( -
- {label} - {isEditing ? ( - - {input === "number" ? ( - - ) : ( - - )} - - ) : ( -
- {(input === "datetime" - ? formatPtuUtcDisplay(localModelData?.model_info?.[name]) - : localModelData?.model_info?.[name]) ?? "Not Set"} -
- )} -
- ); - })} + {ptuCostAttributionEnabled && + PTU_EDIT_FIELDS.map((ptuField) => { + const { name, label, input, placeholder, isCount, isRate, isStart, pairedWith } = ptuField; + const { windowPeer, bound } = ptuField; + return ( +
+ {label} + {isEditing ? ( + + {input === "number" ? ( + + ) : ( + + )} + + ) : ( +
+ {(input === "datetime" + ? formatPtuUtcDisplay(localModelData?.model_info?.[name]) + : localModelData?.model_info?.[name]) ?? "Not Set"} +
+ )} +
+ ); + })}
Cache Read Cost (per 1M tokens) diff --git a/ui/litellm-dashboard/src/utils/ptuModelInfo.test.ts b/ui/litellm-dashboard/src/utils/ptuModelInfo.test.ts new file mode 100644 index 00000000000..36b0e0eac4e --- /dev/null +++ b/ui/litellm-dashboard/src/utils/ptuModelInfo.test.ts @@ -0,0 +1,68 @@ +import dayjs from "dayjs"; +import utc from "dayjs/plugin/utc"; +import { describe, expect, it } from "vitest"; +import { applyPtuModelInfo, PTU_MODEL_INFO_FIELDS } from "./ptuModelInfo"; + +dayjs.extend(utc); + +const storedModelInfo = () => ({ + id: "model-1", + team_id: "team-1", + ptu_count: 15, + cost_per_ptu_per_hour: 2, + ptu_effective_from: "2026-07-01T00:00:00.000Z", + ptu_effective_to: "2026-08-01T00:00:00.000Z", +}); + +describe("applyPtuModelInfo", () => { + it("folds the form values into model_info when PTU cost attribution is enabled", () => { + const result = applyPtuModelInfo( + { id: "model-1", team_id: "team-1" }, + { + ptu_count: "20", + cost_per_ptu_per_hour: "3.5", + ptu_effective_from: dayjs.utc("2026-09-01T00:00:00.000Z"), + ptu_effective_to: null, + }, + true, + ); + + expect(result).toEqual({ + id: "model-1", + team_id: "team-1", + ptu_count: 20, + cost_per_ptu_per_hour: 3.5, + ptu_effective_from: "2026-09-01T00:00:00.000Z", + ptu_effective_to: null, + }); + }); + + it("sends an explicit null for a field the operator cleared while enabled", () => { + const result = applyPtuModelInfo(storedModelInfo(), { ptu_count: "", cost_per_ptu_per_hour: "" }, true); + + expect(result.ptu_count).toBeNull(); + expect(result.cost_per_ptu_per_hour).toBeNull(); + }); + + it("strips every PTU field from the payload when PTU cost attribution is disabled", () => { + const result = applyPtuModelInfo(storedModelInfo(), { ptu_count: "20", cost_per_ptu_per_hour: "3.5" }, false); + + for (const field of PTU_MODEL_INFO_FIELDS) { + expect(Object.keys(result)).not.toContain(field); + } + expect(result).toEqual({ id: "model-1", team_id: "team-1" }); + }); + + it("never sends a null PTU field when disabled, so an unrelated save cannot clear stored config", () => { + const result = applyPtuModelInfo(storedModelInfo(), {}, false); + + expect(Object.values(result)).not.toContain(null); + expect("ptu_count" in result).toBe(false); + }); + + it("leaves non-PTU model_info untouched when disabled", () => { + const result = applyPtuModelInfo({ id: "model-1", access_groups: ["a"], health_check_model: "gpt-5.2" }, {}, false); + + expect(result).toEqual({ id: "model-1", access_groups: ["a"], health_check_model: "gpt-5.2" }); + }); +}); diff --git a/ui/litellm-dashboard/src/utils/ptuModelInfo.ts b/ui/litellm-dashboard/src/utils/ptuModelInfo.ts new file mode 100644 index 00000000000..5cb1452a6b5 --- /dev/null +++ b/ui/litellm-dashboard/src/utils/ptuModelInfo.ts @@ -0,0 +1,44 @@ +import { Dayjs } from "dayjs"; +import { ptuPickerToUtcIso } from "./ptuDatetime"; +import { PTU_COUNT_FIELD, PTU_RATE_FIELD } from "./ptuValidation"; + +export const PTU_MODEL_INFO_FIELDS: readonly string[] = [ + PTU_COUNT_FIELD, + PTU_RATE_FIELD, + "ptu_effective_from", + "ptu_effective_to", +]; + +export interface PtuFormValues { + ptu_count?: string | number | null; + cost_per_ptu_per_hour?: string | number | null; + ptu_effective_from?: Dayjs | null; + ptu_effective_to?: Dayjs | null; +} + +const ptuNumber = (value: string | number | null | undefined): number | null => + value !== undefined && value !== null && value !== "" ? Number(value) : null; + +/** + * Fold the PTU form values into the model_info an edit is about to save. + * + * When PTU cost attribution is off the four fields are stripped rather than sent as null: + * the form does not render them, so a null would be an explicit clear of config the operator + * never saw, and any PTU field present in the payload is rejected by the proxy. + */ +export const applyPtuModelInfo = ( + modelInfo: Record, + values: PtuFormValues, + enabled: boolean, +): Record => { + if (!enabled) { + return Object.fromEntries(Object.entries(modelInfo).filter(([key]) => !PTU_MODEL_INFO_FIELDS.includes(key))); + } + return { + ...modelInfo, + ptu_count: ptuNumber(values.ptu_count), + cost_per_ptu_per_hour: ptuNumber(values.cost_per_ptu_per_hour), + ptu_effective_from: ptuPickerToUtcIso(values.ptu_effective_from), + ptu_effective_to: ptuPickerToUtcIso(values.ptu_effective_to), + }; +};