feat(ptu): gate PTU flat-cost attribution behind an opt-in env var (#36138)

LITELLM_ENABLE_PTU_COST_ATTRIBUTION, read through get_secret_bool and defaulting to
false, makes the whole PTU flat-cost feature inert unless an operator opts in. The
daily rollup cron is not registered at all, so no sentinel row is ever written;
/model/new and /model/{id}/update reject a request that carries any PTU model_info
field with a 400 naming the fields and the env var rather than dropping them; the
daily activity read path reports zero flat cost; and the model add and edit forms
hide the four PTU inputs.

The read gate lives where flat cost enters SpendMetrics rather than in the aggregated
SQL select. /team/daily/activity, the endpoint the Usage page reads, is served by the
paginated find_many path and never runs that query, so forcing the select to a
constant zero would have left the reporting surface that matters still showing flat
cost.

Sentinel row filtering is deliberately not gated. An operator can enable the flag,
accrue rows under the __ptu_flat_cost__ api_key, then disable it, and those rows stay
in LiteLLM_DailyTeamSpend; gating the filter too would surface the sentinel as a bogus
api_key and mint a provider bucket for its empty provider. Response fields keep their
shape and report 0.0, so typed clients are unaffected, and the migration and the
ModelInfo field declarations are untouched.

The write gate reads the incoming request rather than the merged deployment, so a
model configured during an earlier opt-in stays editable, and the edit form drops the
PTU keys from the payload instead of sending nulls that would clear stored config.

The dashboard reads the flag from a read-only enable_ptu_cost_attribution key on
/get/ui_settings, computed from the environment on every read. It is deliberately not
an allowlisted persisted setting, and PATCH /update/ui_settings rejects it with a 400,
so an admin cannot flip an env-gated feature from the UI.

Two review findings on the gate itself. The PTU clear loop now runs only when the
feature is enabled: the write gate rejects a value but lets an explicit null through,
and a client round-tripping a model_info blob sends the PTU keys as nulls, so a
disabled proxy would have quietly erased a billing configuration set up during an
earlier opt-in. Disabling pauses PTU rather than discarding its setup. And the
dashboard flag is re-read every thirty seconds instead of the hour the other UI settings
use, since those are persisted records while this one tracks the proxy process; a
restart that flips the variable would otherwise leave the model form offering inputs
the backend now rejects. The flag is polled rather than only marked stale, since a form
that stays mounted and focused never refetches on its own.

The read gate checks the row before the flag. It runs once per metric accumulation and a
record fans out across roughly a dozen breakdowns, while 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 reaches it.
This commit is contained in:
yucheng-berri 2026-08-10 12:23:20 -07:00 committed by GitHub
parent 9de3315dad
commit e014b341c8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
21 changed files with 1513 additions and 184 deletions

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

@ -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 == []

View file

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

View file

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

View file

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

View file

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

View file

@ -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<Record<string, any>>({
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,
});
};

View file

@ -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(
<AdvancedSettings
showAdvancedSettings={true}
setShowAdvancedSettings={() => {}}
guardrailsList={[]}
tagsList={{}}
accessToken="test-token"
/>,
);
describe("AdvancedSettings", () => {
beforeEach(() => {
vi.clearAllMocks();
mockUsePtuCostAttributionEnabled.mockReturnValue(false);
});
it("should render", () => {
render(
<AdvancedSettings
showAdvancedSettings={true}
setShowAdvancedSettings={() => {}}
guardrailsList={[]}
tagsList={{}}
accessToken="test-token"
/>,
);
renderAdvancedSettings();
});
it("should render tags list", async () => {
const { getByText } = render(
<AdvancedSettings
showAdvancedSettings={true}
setShowAdvancedSettings={() => {}}
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(
<AdvancedSettings
showAdvancedSettings={true}
setShowAdvancedSettings={() => {}}
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();
});
});

View file

@ -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<AdvancedSettingsProps> = ({
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<AdvancedSettingsProps> = ({
/>
</Form.Item>
<Form.Item
label="PTU Count"
name={PTU_COUNT_FIELD}
dependencies={[PTU_RATE_FIELD]}
rules={[{ validator: validateNumber }, ...ptuCountRules, ptuPairRule(PTU_RATE_FIELD)]}
tooltip="Provisioned throughput units for this deployment. Set together with Cost per PTU / Hour and a Team to attribute a flat daily cost."
className="mb-4"
>
<TextInput placeholder="e.g. 15" />
</Form.Item>
{ptuCostAttributionEnabled && (
<>
<Form.Item
label="PTU Count"
name={PTU_COUNT_FIELD}
dependencies={[PTU_RATE_FIELD]}
rules={[{ validator: validateNumber }, ...ptuCountRules, ptuPairRule(PTU_RATE_FIELD)]}
tooltip="Provisioned throughput units for this deployment. Set together with Cost per PTU / Hour and a Team to attribute a flat daily cost."
className="mb-4"
>
<TextInput placeholder="e.g. 15" />
</Form.Item>
<Form.Item
label="Calculated Cost per PTU / Hour (USD)"
name={PTU_RATE_FIELD}
dependencies={[PTU_COUNT_FIELD]}
rules={[{ validator: validateNumber }, ...ptuRateRules, ptuPairRule(PTU_COUNT_FIELD)]}
tooltip="Flat cost = PTU count * this rate * active hours, attributed to the deployment's team."
className="mb-4"
>
<TextInput placeholder="e.g. 2.00" />
</Form.Item>
<Form.Item
label="Calculated Cost per PTU / Hour (USD)"
name={PTU_RATE_FIELD}
dependencies={[PTU_COUNT_FIELD]}
rules={[{ validator: validateNumber }, ...ptuRateRules, ptuPairRule(PTU_COUNT_FIELD)]}
tooltip="Flat cost = PTU count * this rate * active hours, attributed to the deployment's team."
className="mb-4"
>
<TextInput placeholder="e.g. 2.00" />
</Form.Item>
<Form.Item
label="PTU Effective From (UTC)"
name={PTU_START_FIELD}
dependencies={[PTU_COUNT_FIELD, PTU_END_FIELD]}
rules={[ptuStartRequiredRule(PTU_COUNT_FIELD), ptuWindowOrderRule(PTU_END_FIELD, "start")]}
tooltip="Start of the PTU window, required when PTU Count is set. Flat cost accrues by the hour within the window; a window opening at 23:00 charges one hour that day."
className="mb-4"
>
<DatePicker showTime style={{ width: "100%" }} />
</Form.Item>
<Form.Item
label="PTU Effective From (UTC)"
name={PTU_START_FIELD}
dependencies={[PTU_COUNT_FIELD, PTU_END_FIELD]}
rules={[ptuStartRequiredRule(PTU_COUNT_FIELD), ptuWindowOrderRule(PTU_END_FIELD, "start")]}
tooltip="Start of the PTU window, required when PTU Count is set. Flat cost accrues by the hour within the window; a window opening at 23:00 charges one hour that day."
className="mb-4"
>
<DatePicker showTime style={{ width: "100%" }} />
</Form.Item>
<Form.Item
label="PTU Effective To (UTC)"
name={PTU_END_FIELD}
dependencies={[PTU_START_FIELD]}
rules={[ptuWindowOrderRule(PTU_START_FIELD, "end")]}
tooltip="Optional end of the PTU window (exclusive). Leave blank for open-ended."
className="mb-4"
>
<DatePicker showTime style={{ width: "100%" }} />
</Form.Item>
<Form.Item
label="PTU Effective To (UTC)"
name={PTU_END_FIELD}
dependencies={[PTU_START_FIELD]}
rules={[ptuWindowOrderRule(PTU_START_FIELD, "end")]}
tooltip="Optional end of the PTU window (exclusive). Leave blank for open-ended."
className="mb-4"
>
<DatePicker showTime style={{ width: "100%" }} />
</Form.Item>
</>
)}
{customPricing && (
<div className="ml-6 pl-4 border-l-2 border-gray-200">

View file

@ -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(<ModelInfoView {...DEFAULT_ADMIN_PROPS} />, { 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,

View file

@ -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({
)}
</div>
{PTU_EDIT_FIELDS.map((ptuField) => {
const { name, label, input, placeholder, isCount, isRate, isStart, pairedWith } = ptuField;
const { windowPeer, bound } = ptuField;
return (
<div key={name}>
<Text className="font-medium">{label}</Text>
{isEditing ? (
<Form.Item
name={name}
className="mb-0"
dependencies={ptuFieldDependencies(ptuField)}
rules={[
...(isCount ? ptuCountRules : []),
...(isRate ? ptuRateRules : []),
...(isStart ? [ptuStartRequiredRule(PTU_COUNT_FIELD)] : []),
...(pairedWith ? [ptuPairRule(pairedWith)] : []),
...(windowPeer && bound ? [ptuWindowOrderRule(windowPeer, bound)] : []),
]}
>
{input === "number" ? (
<NumericalInput
placeholder={placeholder}
step={isCount ? 1 : undefined}
min={isCount ? 1 : 0}
/>
) : (
<DatePicker showTime style={{ width: "100%" }} />
)}
</Form.Item>
) : (
<div className="mt-1 p-2 bg-gray-50 rounded-sm">
{(input === "datetime"
? formatPtuUtcDisplay(localModelData?.model_info?.[name])
: localModelData?.model_info?.[name]) ?? "Not Set"}
</div>
)}
</div>
);
})}
{ptuCostAttributionEnabled &&
PTU_EDIT_FIELDS.map((ptuField) => {
const { name, label, input, placeholder, isCount, isRate, isStart, pairedWith } = ptuField;
const { windowPeer, bound } = ptuField;
return (
<div key={name}>
<Text className="font-medium">{label}</Text>
{isEditing ? (
<Form.Item
name={name}
className="mb-0"
dependencies={ptuFieldDependencies(ptuField)}
rules={[
...(isCount ? ptuCountRules : []),
...(isRate ? ptuRateRules : []),
...(isStart ? [ptuStartRequiredRule(PTU_COUNT_FIELD)] : []),
...(pairedWith ? [ptuPairRule(pairedWith)] : []),
...(windowPeer && bound ? [ptuWindowOrderRule(windowPeer, bound)] : []),
]}
>
{input === "number" ? (
<NumericalInput
placeholder={placeholder}
step={isCount ? 1 : undefined}
min={isCount ? 1 : 0}
/>
) : (
<DatePicker showTime style={{ width: "100%" }} />
)}
</Form.Item>
) : (
<div className="mt-1 p-2 bg-gray-50 rounded-sm">
{(input === "datetime"
? formatPtuUtcDisplay(localModelData?.model_info?.[name])
: localModelData?.model_info?.[name]) ?? "Not Set"}
</div>
)}
</div>
);
})}
<div>
<Text className="font-medium">Cache Read Cost (per 1M tokens)</Text>

View file

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

View file

@ -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<string, unknown>,
values: PtuFormValues,
enabled: boolean,
): Record<string, unknown> => {
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),
};
};