feat(ptu): daily rollup job writes flat cost to LiteLLM_DailyTeamSpend

[yucheng-berri]

Relevant issues

Linear ticket

Resolves LIT-1697 (stage 2 of 5)

Stacks on #33130 (stage 1). Merge that first.

Pre-Submission checklist

• I have added meaningful tests
• My PR passes all CI/CD checks (e.g., lint, format, unit tests)
• My PR's scope is as isolated as possible; it only solves 1 specific problem
• I have received a Greptile Confidence Score of at least 4/5 before requesting a maintainer review

Screenshots / Proof of Fix

Stage 1 shipped the reservation storage and endpoints; stage 2 makes the daily rollup actually write. To exercise end-to-end against a live proxy on this branch's HEAD, with the feature flag on and a manual PTU reservation:

$ MASTER_KEY=sk-1234
$ curl -s -X POST http://localhost:4000/config/general_settings \
    -H "Authorization: Bearer $MASTER_KEY" -H "Content-Type: application/json" \
    -d '{"enable_ptu_cost_attribution": true}' > /dev/null
$ TEAM=<real team id>

# admin creates a reservation
$ curl -s -X POST http://localhost:4000/ptu_reservation/new \
    -H "Authorization: Bearer $MASTER_KEY" -H "Content-Type: application/json" \
    -d "{\"team_id\":\"$TEAM\",\"model\":\"gpt-4\",\"ptu_count\":1,\"cost_per_ptu\":200,\"effective_from\":\"2026-07-01T00:00:00Z\"}" | jq '.id'
"<res-id>"

# run backfill for a specific day
$ .venv/bin/python scripts/ptu_reservation_backfill.py --date 2026-07-12
[2026-07-12] reservations=1 rows_written=1
total rows written: 1

# check the row landed with sentinel api_key and expected math (200 / 31)
$ psql $DATABASE_URL -c "SELECT team_id, date, api_key, model, ptu_flat_cost, ptu_reservation_id, spend
                          FROM \"LiteLLM_DailyTeamSpend\"
                          WHERE date='2026-07-12' AND team_id='$TEAM';"
 team_id | date       | api_key                | model | ptu_flat_cost      | ptu_reservation_id | spend
---------+------------+------------------------+-------+--------------------+--------------------+-------
 team_x  | 2026-07-12 | __ptu_reservation__    | gpt-4 | 6.4516129032258064 | <res-id>           |   0.0

# rerun backfill — same row, no duplicate (idempotent)
$ .venv/bin/python scripts/ptu_reservation_backfill.py --date 2026-07-12
[2026-07-12] reservations=1 rows_written=1
total rows written: 1

$ psql $DATABASE_URL -c "SELECT COUNT(*) FROM \"LiteLLM_DailyTeamSpend\"
                          WHERE date='2026-07-12' AND team_id='$TEAM'
                            AND api_key='__ptu_reservation__' AND ptu_reservation_id='<res-id>';"
 count
-------
     1

# range backfill for a whole month, Feb 2026 (28 days) at $200/PTU should produce 200/28 per day
$ .venv/bin/python scripts/ptu_reservation_backfill.py --date-range 2026-02-01:2026-02-28
[2026-02-01] reservations=1 rows_written=1
...
[2026-02-28] reservations=1 rows_written=1
total rows written: 28

The rollup does not fire when the flag is off:

$ curl -s -X POST http://localhost:4000/config/general_settings \
    -H "Authorization: Bearer $MASTER_KEY" -H "Content-Type: application/json" \
    -d '{"enable_ptu_cost_attribution": false}' > /dev/null
$ .venv/bin/python scripts/ptu_reservation_backfill.py --date 2026-07-13
[2026-07-13] reservations=0 rows_written=0 (flag off, skipped)
total rows written: 0

Per-request rows are untouched by the rollup; a team that also has real Azure gpt-4 traffic on the same day has its own row with a real hashed api_key and the sentinel PTU row alongside it. Stage 3 will teach /team/daily/activity to sum the two.

Type

New Feature

Changes

Adds the daily rollup that writes flat PTU cost onto LiteLLM_DailyTeamSpend. Stage 1 gave admins a way to register reservations; stage 2 makes those reservations show up in per-day spend rows without touching the per-request pipeline.

Storage. Two additive columns on LiteLLM_DailyTeamSpend: ptu_flat_cost (float, default 0.0) and ptu_reservation_id (nullable text). New index on ptu_reservation_id. No change to the existing unique constraint; PTU rows share it via a sentinel api_key = "__ptu_reservation__" that also namespaces them away from real per-request rows in every existing filter that keys on api_key.

Rollup job. One scheduled APScheduler cron job at 00:15 UTC daily, registered inside initialize_scheduled_background_jobs behind the enable_ptu_cost_attribution config flag. The job body also re-reads the flag at execution time and no-ops if the flag has been toggled off since registration. For each active reservation at the day's start (effective_from <= 00:00 UTC < effective_to), it computes ptu_count * cost_per_ptu / days_in_that_calendar_month and upserts a sentinel-api_key row for (team_id, date, model). Calendar-month proration is the deliberate choice: Feb 28 days yields ~$7.14/day for a $200/PTU reservation, Jul 31 days yields ~$6.45/day, matching how Azure amortizes reservation cost.

Idempotency. The upsert uses LiteLLM_DailyTeamSpend's composite unique constraint. Re-running the job for the same day overwrites ptu_flat_cost and ptu_reservation_id with the same values; no duplicate rows. Failures on one reservation don't stop the batch; the loop logs and continues so a bad row can't block the rest.

Backfill CLI. scripts/ptu_reservation_backfill.py calls the same rollup function for a specific date or an inclusive YYYY-MM-DD:YYYY-MM-DD range. Same idempotency guarantees. Operators use this after a proxy outage or when onboarding Evernorth's already-elapsed months (if requested; see stage 5).

Mid-day edit semantics. Whichever reservation is active at 00:00 UTC of a given day owns that day's full flat cost. If an admin closes a reservation at 09:00 UTC and creates a replacement, the old one gets that day; the new one starts contributing the following day. No hour-level proration.

Deviation from the admin-entity pattern

Stage 1 already documented the "close old + create new" edit flow. Stage 2 does not introduce further deviations from established patterns. The scheduled-job registration follows the same shape as spend_log_cleanup_job and the tag-spend batch job registered in the same initialize_scheduled_background_jobs block.

Behavior changes

• New rows appear on LiteLLM_DailyTeamSpend. Once the flag is on, teams with active reservations will start seeing rows with api_key = "__ptu_reservation__" from the day after the flag is enabled. Any external SQL consumer of LiteLLM_DailyTeamSpend that already filters on api_key in a specific set will naturally exclude these; consumers filtering on api_key = ANY will see them and should special-case the sentinel value if the desired output is per-request-only.
• Two new columns on LiteLLM_DailyTeamSpend: ptu_flat_cost and ptu_reservation_id. Both default-safe (0.0 and null); no code path outside this feature reads them yet. Stage 3 will surface them in /team/daily/activity.
• No change to any existing spend, prompt_tokens, api_requests, etc. numeric column
• No change to the per-request write path: _PROXY_track_cost_callback, db_spend_update_writer.update_database, _batch_database_updates, _commit_spend_updates_to_db, common_daily_activity.py, LiteLLM_TeamTable.spend, budget enforcement are all untouched
• Feature is off by default; enabling it registers the daily job on the next proxy restart and starts writing rows at the next 00:15 UTC boundary
• No migration for existing customers required; the two new columns default-safe

Boilerplate note

None new. The rollup module is a first-of-its-kind cost writer that runs on a schedule; no existing template to copy. The CLI helper is small and self-contained.

Files changed

• schema.prisma: ptu_flat_cost (Float default 0.0) and ptu_reservation_id (String?) added to LiteLLM_DailyTeamSpend; new @@index([ptu_reservation_id])
• litellm-proxy-extras/.../20260713010000_add_ptu_columns_to_daily_team_spend/migration.sql: ADD COLUMN IF NOT EXISTS + CREATE INDEX IF NOT EXISTS
• litellm/proxy/spend_tracking/ptu_reservation_rollup.py: rollup job body, proration function, upsert helper, sentinel constant, dataclass result, module __all__
• litellm/proxy/proxy_server.py: PTU rollup job registered inside initialize_scheduled_background_jobs when enable_ptu_cost_attribution is truthy at startup; uses the string-based "cron" trigger form to avoid a stub-tracking import delta
• scripts/ptu_reservation_backfill.py: --date / --date-range CLI wrapping the same rollup function
• tests/test_litellm/proxy/spend_tracking/test_ptu_reservation_rollup.py: 19 tests

Tests

19 tests in test_ptu_reservation_rollup.py. Suite covers:

Proration math:

• _days_in_month returns 31/28/29/30 across the four calendar shapes (including leap Feb 2024)
• Flat cost for 31-day month, 28-day month, leap Feb — each asserted with numeric approx
• Flat cost scales linearly with ptu_count (100 PTU = 100 × 1 PTU)
• Flat cost is 0 when cost_source != "manual" (forward-compat with azure_billing)
• Flat cost is 0 when manual fields are missing

Job body:

• Flag off short-circuits before any DB read (asserts find_many never awaited)
• prisma_client is None returns zero-row result
• Single active reservation writes one upsert with the correct where clause and create / update payloads
• cost_source="azure_billing" reservation is walked but skipped (rows_written == 0)
• Three reservations across two teams and two models each write one row
• target_date defaults to yesterday UTC
• find_many is called with the correct half-open (effective_from <= day_start, effective_to > day_start OR null) where clause
• Two consecutive runs on the same day write to the same composite key (idempotent)
• Per-reservation upsert failure isolated: bad row logged, remaining reservations still processed
• Sentinel api_key = "__ptu_reservation__" used (not a real token)
• PTU rows carry only ptu_flat_cost and ptu_reservation_id on their payload; do not touch spend, prompt_tokens, completion_tokens, api_requests
• Reservation with effective_from exactly at 00:00 UTC of the target day is treated as active (inclusive left boundary)

$ .venv/bin/python -m pytest tests/test_litellm/proxy/spend_tracking/test_ptu_reservation_rollup.py -q
19 passed

Combined stage 1 + stage 2 suite:

$ .venv/bin/python -m pytest tests/test_litellm/proxy/spend_tracking/test_ptu_reservation_rollup.py tests/test_litellm/proxy/management_endpoints/test_ptu_reservation_endpoints.py -q
47 passed

Neighboring spend_tracking + budget tests all still pass (276 passed).

---

Note

Medium Risk
Touches daily team spend storage and a new scheduled writer; feature-flagged and additive, but SQL/API consumers of LiteLLM_DailyTeamSpend may see new sentinel rows until stage 3 aggregates them.

Overview
Adds stage 2 PTU cost attribution: a scheduled daily rollup that writes prorated prepaid PTU cost into LiteLLM_DailyTeamSpend, separate from per-request spend.

Schema & migration. LiteLLM_DailyTeamSpend gains ptu_flat_cost (default 0) and ptu_reservation_id, plus an index on ptu_reservation_id. Prisma schemas and a proxy-extras migration use IF NOT EXISTS for safe rollout.

Rollup behavior. New ptu_reservation_rollup loads active reservations at UTC day start, prorates manual reservations as (ptu_count × cost_per_ptu) / days_in_calendar_month, and idempotently upserts rows keyed by team/date/model with sentinel api_key = "__ptu_reservation__" so they sit beside real API-key rows without changing spend or token counters. Non-manual (azure_billing) reservations are skipped; per-reservation upsert failures are logged and the batch continues.

Scheduling & ops. When enable_ptu_cost_attribution is on at proxy startup, APScheduler registers a 00:15 UTC cron job; the job also no-ops at runtime if the flag is off. scripts/ptu_reservation_backfill.py re-runs the same logic for one day or an inclusive date range with force=True (no flag required).

Tests. Broad unit coverage for proration math, flag gating, upsert shape, idempotency, and the backfill CLI.

^{Reviewed by Cursor Bugbot for commit 1c317df. Bugbot is set up for automated code reviews on this repo. Configure here.}

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
1 out of 2 committers have signed the CLA.

✅ yucheng-berri
❌ github-actions[bot]
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[codecov]

Codecov Report

❌ Patch coverage is 94.82759% with 3 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
litellm/proxy/proxy_server.py	25.00%	3 Missing ⚠️

📢 Thoughts on this report? Let us know!

[greptile-apps]

Greptile Summary

This PR implements stage 2 of the PTU reservation cost attribution feature: a nightly APScheduler cron job that writes prorated flat PTU costs into LiteLLM_DailyTeamSpend using a sentinel api_key = \"__ptu_reservation__\" to keep these rows distinct from per-request traffic. The previously-flagged P1 (backfill CLI importing prisma_client from the proxy module globals, which are always None when running standalone) is correctly resolved — the script now bootstraps its own PrismaClient from DATABASE_URL and passes force=True to bypass the feature-flag check.

• Schema: Two additive, default-safe columns (ptu_flat_cost Float @default(0.0), ptu_reservation_id String?) plus a new index on LiteLLM_DailyTeamSpend, with an idempotent ADD COLUMN IF NOT EXISTS migration.
• Rollup logic: run_ptu_reservation_rollup re-reads general_settings at execution time so the flag can be toggled off without a restart; force=True lets the CLI backfill without the scheduler running.
• Tests: 19 rollup unit tests + 8 backfill CLI tests, all pure-mock with no real network calls, covering proration math, idempotency, fault isolation, flag bypass, and lifecycle.

Confidence Score: 5/5

Safe to merge — additive schema changes with safe defaults, the rollup is off by default, and the previously broken backfill CLI is now correctly self-contained.

The only previously blocking issue (backfill CLI always exiting with zero rows due to importing proxy module globals that are never populated when running standalone) is correctly fixed: the script creates its own PrismaClient from DATABASE_URL and uses force=True to bypass the flag check. Rollup logic, idempotency, fault isolation, schema migration, and test coverage are all solid.

No files require special attention.

Important Files Changed

Filename	Overview
litellm/proxy/spend_tracking/ptu_reservation_rollup.py	New rollup module: proration math, idempotent upsert, feature-flag re-check at runtime, force bypass for CLI — all logic is correct and well-isolated
scripts/ptu_reservation_backfill.py	Backfill CLI now correctly bootstraps its own PrismaClient from DATABASE_URL and passes force=True to bypass the feature-flag check; P1 from prior review is resolved
litellm/proxy/proxy_server.py	Cron job registered at 00:15 UTC behind the enable_ptu_cost_attribution flag, following the same pattern as existing scheduled jobs; prisma_client captured at startup time matches the established convention
litellm-proxy-extras/litellm_proxy_extras/migrations/20260713010000_add_ptu_columns_to_daily_team_spend/migration.sql	Additive migration using ADD COLUMN IF NOT EXISTS and CREATE INDEX IF NOT EXISTS — fully safe for rolling deploys and idempotent re-runs
tests/test_litellm/proxy/spend_tracking/test_ptu_reservation_rollup.py	19 pure-mock tests covering proration math, flag bypass, force=True path, idempotency, per-reservation fault isolation, and sentinel payload invariants
tests/test_scripts/test_ptu_reservation_backfill.py	8 new tests covering date parsing, range validation, DATABASE_URL guard, force=True forwarding, connect/disconnect lifecycle — all properly mocked, no real network calls
schema.prisma	Adds ptu_flat_cost (Float default 0.0) and ptu_reservation_id (String?) to LiteLLM_DailyTeamSpend plus a new index; both columns are default-safe and additive

_{Reviews (2): Last reviewed commit: "fix(ptu): backfill CLI bootstraps its ow..." | Re-trigger Greptile}

[greptile-apps]

Sentinel constant should live in constants.py

PTU_SENTINEL_API_KEY = "__ptu_reservation__" is a sentinel string defined inline in ptu_reservation_rollup.py. The team's convention (see existing sentinels such as _NEGATIVE_TEAM_SENTINEL) is to place all sentinel-like variables in constants.py so they have a single authoritative home. Stage 3 will need to reference this value when filtering spend rows, and any consumer that copies the string literal rather than importing it will silently drift. PTU_ROLLUP_JOB_ID is less critical but follows the same pattern.

Rule Used: What: Require sentinel-like variables (e.g., `_NEG... (source)

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[yucheng-berri]

@greptileai please re-review — pushed P1 fix for the backfill CLI. It now bootstraps its own PrismaClient from DATABASE_URL and passes force=True to run_ptu_reservation_rollup (new kwarg; default False keeps the scheduler contract intact). 8 new tests: 6 CLI-side, 2 rollup-side for the force branch. Base also rebased onto the updated stage 1 branch.

[yucheng-berri]

bugbot run

[cursor]

Cursor Bugbot has reviewed your changes using high effort and found 4 potential issues.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 1c317df. Configure here.}

[cursor]

Cron schedule is not UTC

Medium Severity

The PTU reservation rollup job schedules for 00:15 without an explicit timezone. This defaults to the host's local timezone, conflicting with the 00:15 UTC stated in logs and the rollup's UTC-based date logic. On non-UTC hosts, this means the job fires at the wrong wall time for the UTC day it attributes.

 

^{Reviewed by Cursor Bugbot for commit 1c317df. Configure here.}

[cursor]

Outages skip rollup days

High Severity

The daily job inherits misfire_grace_time of one hour and only ever rolls yesterday UTC when it runs. If the proxy is down or late by more than that window around the cron fire, that UTC day is never written automatically; the next successful run attributes a different day, so flat PTU cost rows are permanently missing until a manual backfill.

Additional Locations (1)

• litellm/proxy/spend_tracking/ptu_reservation_rollup.py#L87-L120

 

^{Reviewed by Cursor Bugbot for commit 1c317df. Configure here.}

[cursor]

Flag kill-switch never updates

Medium Severity

The rollup re-reads general_settings["enable_ptu_cost_attribution"] intending to no-op after a runtime toggle-off, but live DB config sync only copies a fixed allowlist of keys into the in-memory dict and that flag is not among them. Once the job is registered at startup, turning the flag off does not stop rollup writes without a process restart.

Additional Locations (1)

• litellm/proxy/proxy_server.py#L7839-L7841

 

^{Reviewed by Cursor Bugbot for commit 1c317df. Configure here.}

[cursor]

Stale PTU rows persist

Medium Severity

The rollup only upserts reservations still active at day_start. It never clears or zeros an existing sentinel row when that reservation is later closed, backdated, or no longer active for the target day. Re-running the job or backfill CLI after fixing reservation windows therefore leaves orphaned ptu_flat_cost / ptu_reservation_id values that stage 3 activity sums would still pick up.

 

^{Reviewed by Cursor Bugbot for commit 1c317df. Configure here.}

[yucheng-berri]

Closing as superseded. This is part of the v1 PTU design, which stored PTU config in a separate reservation table. The shipped design puts that config on the model deployment instead, merged as #35341, #35343, #35391, #35393 and #36829.

The rollup landed as #35343, keyed off the model deployment rather than a reservation row.

The branch is kept, so nothing here is lost.