feat: maintain transactional inference and service spend counters

[henrypark133]

Add transactional spending counters as the first counter rollout slice: separate inference/service totals on organization_balance and a per-key api_key_spend table. Both existing posting transactions increment counters only after a new usage insert succeeds, so duplicate deliveries do not double-count and failures roll back usage, allocation, balances, and counters together.

Readers remain unchanged in this PR. Historical totals require a separate backfill and readiness check before key lists, admission checks, or billing summaries can use these counters. This PR does not yet remove the lifetime scans or change spending-limit behavior.

Validation:

• Baseline counter tests failed because api_key_spend did not exist.
• cargo nextest run --test e2e_all -E 'test(spend_counters::) | test(/^usage_recording::test_record_usage_idempotent_duplicate$/) | test(/^usage_provider_attribution::duplicate_usage_preserves_original_provider_attribution$/)' --no-fail-fast: 4 passed, 799 skipped, against dedicated PostgreSQL 16 with the repository's E2E bootstrap.
• Coverage includes service-first counter creation, subsequent inference/service additions, duplicate deliveries for both usage types, and a forced bigint overflow that preserves the previous balance/allocation totals and leaves no new usage row.
• cargo fmt --all -- --check and git diff --check passed. Independent source review found no correctness issues in the migration or posting changes.

[ironloopai]

Review · Status

🟩 Completed

IronLoop completed the review and posted it to GitHub.

Result

Open submitted review →

Run details

• Run: b109401f-42cb-477e-b209-3c01329eed84
• Base: main at 7f22cbd
• Head: codex/analytics-spend-counters at a9789f2
• Created: 2026-09-22 17:20 UTC
• Updated: 2026-09-22 17:24 UTC

_{Automatic trigger · attempt 1 of 3 · completed in 3m 34s}

[ironloopai]

Review · Summary

🟢 No actionable findings

No actionable issues found in the reviewed change.

Validation

• ✅ Transactional counter review — Static inspection confirms inference and service posting update their split counters within the existing transaction only after a new usage row is inserted; duplicate paths roll back before counter writes.

Review details

• Run: b109401f-42cb-477e-b209-3c01329eed84
• Attempts: 1

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[claude]

Review — feat: maintain transactional inference and service spend counters

Verified the change statically against the two posting paths, the migration, and every existing organization_balance writer. Agreeing with the earlier IronLoop review: no critical issues found.

What I confirmed

• Positional params are correct in both upserts. organization_usage.rs:201 maps to (org, total_spent=$2, inference_spent=$2, service_spent=0, last_usage_at=$3, total_requests=1, total_tokens=$4, updated_at=$5); organization_service_usage.rs:366 maps to (org, total_spent=$2, inference_spent=0, service_spent=$2, ...). Both column lists line up with their bind arrays, and each DO UPDATE branch increments the matching split counter.
• Counters are genuinely transactional. increment_api_key_spend takes &Transaction and runs after the usage insert returned a row and after allocate_usage, before commit(). The duplicate branch rollback()s before reaching it, so re-delivery of the same inference_id cannot double-count.
• These are the only two production writers of organization_balance.total_spent, so the split counters cannot drift from total_spent going forward. credit_allocation.rs:206/428 only touch unresolved_unfunded_amount.
• No new lock-ordering risk. Both paths already hold the per-org SELECT ... FOR UPDATE from lock_organization_accounting, which serializes all postings for an org, and the api_keys row is already KEY SHARE-locked by the existing usage-log FK. The new upsert adds no lock class and no new deadlock edge.
• FK is safe. organization_usage_log.api_key_id and organization_service_usage_log.api_key_id are both NOT NULL REFERENCES api_keys(id) without cascade, so the parent row always exists by the time the counter is written.
• Migration is rolling-update safe. ADD COLUMN ... NOT NULL DEFAULT 0 is metadata-only on PG 11+, both callers name columns explicitly, and readers use named row.get, so old pods and struct mappers are unaffected. Style matches V0078/V0080.
• Overflow is not a new failure mode. Per-key totals are bounded by the org total, so api_key_spend cannot overflow before organization_balance.total_spent already would.

Non-blocking notes for the follow-up slices

• Gate the backfill on rollout completion. During a rolling update, old replicas post usage without touching api_key_spend or the split columns, so counters silently undercount for the rollout window. The planned backfill should run after every replica is on the new build, and should be a full recompute from organization_usage_log / organization_service_usage_log rather than a delta — otherwise the rollout gap persists permanently. Worth encoding in the readiness check the description mentions.
• Readers must tolerate a missing row. api_key_spend only materializes on first spend, so key-list and admission queries need LEFT JOIN ... COALESCE(..., 0), not query_one.
• ON DELETE CASCADE vs. lifetime totals. Keys are soft-deleted today (api_key.rs:379) and the usage-log FK would block a hard delete anyway, so this is inert. If a hard-delete path is ever added, per-key counters would vanish while organization_balance keeps the spend; a reconciliation check comparing SUM(api_key_spend) against the org counters would catch that.
• Test hygiene: spend_counters.rs is UUID-scoped (own org, workspace, key, and services row) and asserts through test-owned IDs, consistent with the E2E conventions in CLAUDE.md. The exact-string assertion on Database error (22003): bigint out of range is brittle if map_db_error formatting changes — matching the SQLSTATE alone would be more durable.

✅

[henrypark133]

Code Review (multi-agent)

Intent: Introduce transactional spend counters for inference/service usage on organization_balance and api_key_spend, incrementing atomically with usage while leaving existing readers unchanged.

Stats: 2 findings (from 2 raw, 2 after filter, 2 after dedup) across 2 files. Reviewers run: correctness, security, performance, design, coverage. Reviewers failed: none. Body-only: 0

regression-escape

• High Service-path counter overflow rollback is untested (crates/api/tests/e2e_all/spend_counters.rs:156-236, confidence 85) — anchor: crates/api/tests/e2e_all/spend_counters.rs:156
  The bigint-overflow/rollback guarantee for increment_api_key_spend is only tested for the inference path (spend_counters_ignore_duplicates_and_failed_posts drives OrganizationUsageRepository::record_usage). No test drives OrganizationServiceUsageRepository::record_usage into an api_key_spend/organization_balance overflow. Adversarial verification confirmed the code paths are structurally identical (transaction rolled back on drop before commit in both), so this is a real coverage gap rather than a live correctness bug -- but the duplicated control flow between the two repositories means a future edit to one path silently diverging would ship undetected.

concurrency

• Medium New api_key_spend UPSERT extends hold time of per-org FOR UPDATE lock (crates/database/src/repositories/organization_usage.rs:100-236, confidence 60) (candidate — validate claim) — anchor: crates/database/src/repositories/organization_usage.rs:229
  record_usage() (and the mirrored path in organization_service_usage.rs) takes an exclusive per-organization lock via lock_organization_accounting() before doing any writes, serializing all usage-recording transactions for that organization. This PR adds a 7th sequential statement (the increment_api_key_spend UPSERT) inside that already-serialized critical section, on the hot path hit by every chat completion / service call, directly adding latency to the org-wide bottleneck -- most relevant for high-volume orgs with several active API keys, which is exactly the case this feature targets.

[github-actions]

🔍 OpenCodeReview found 5 issue(s) in this PR.

• ✅ 5 posted as inline comment(s)
• 📝 0 posted as summary

⚠️ 1 warning(s) occurred during review.

[henrypark133]

Review pass: removed the extra database round trip from each accounting transaction and deleted the orphan helper. No new posting mode or generic spend abstraction was added.

The staged reader follow-up is #1119 → #1120. Backfill computes full raw-history totals and subtracts counters from the same repeatable-read snapshot; adding that correction to the current counters preserves concurrent new postings. It is not an incremental “since rollout” scan. Missing counter rows are handled as zero by #1120; hard-deleting keys is not introduced.

Test-file changes: spend_counters.rs adds service-key and organization-service overflow rollback checks, parses the fixture UUID once, documents MAX cleanup, and checks the exact SQLSTATE prefix instead of unstable database diagnostic wording. All duplicate, raw-row, counter, balance, and unfunded-state assertions remain intact.

Local aggregate validation (all four PRs plus review fixes):

cargo nextest run --test e2e_all --test spend_counter_backfill --test counter_readers --bin backfill-spend-counters --no-fail-fast
Summary: 819 passed, 0 failed, 12 existing skips

Formatting and diff checks passed. The final run used DATABASE_* configuration without PG* overrides. This local aggregate validation is not a production backfill or deployment.

GitHub CI verified on 2997c87e: Test Suite, lint, unit, integration, E2E, script tests, cargo audit, and cargo deny all succeeded. Build release was skipped by the workflow. No CI reruns were needed.

GitHub still reports REVIEW_REQUIRED; a maintainer approval is needed before merge.

[lloydmak99]

This PR maintains transactional inference and service spend counters by folding the per-key counter into the balance upsert via a single data-modifying CTE. I verified the positional params, CTE semantics, transactional rollback behavior, and migration safety, and found no merge-blocking issues.

• Positional params line up in both upserts — inference path organization_usage.rs:200-238 and service path organization_service_usage.rs:363-393 (column lists, VALUES, DO UPDATE branches, and bind arrays all match).
• balance_upsert is a data-modifying CTE with ON CONFLICT DO UPDATE ... RETURNING, so it always returns one row and the downstream INSERT ... SELECT ... WHERE TRUE always inserts exactly one api_key_spend row — no double-counting or partial writes.
• Both upserts run after allocate_usage and before commit(); duplicate deliveries roll back first, and an overflow surfaces DatabaseError(22003) which drops the transaction.
• Migration V0081 is rollout-safe: ADD COLUMN NOT NULL DEFAULT 0 is metadata-only on modern PG, both organization_balance writers are updated, and the reader (row_to_balance, organization_usage.rs:694) selects explicit named columns. Rolling-deploy undercount during the window is expected and handled by the follow-up backfill.

Checks: cargo fmt --all -- --check and git diff --check passed; cargo check -p database could not run (environment lacks the cc linker). Full build/e2e skipped as slow — CI and targeted spend_counters e2e were green.