feat: key budgets, spend limits, and per-key spend visibility #4

Merged
devtrev merged 9 commits from trevin/feat/key-budgets-spend-visibility into dev 2026-07-10 21:32:25 -06:00
Owner

Summary

Adds key spend limits (budgets) and per-key spend/usage visibility to the admin UI.

Key management

  • Rename keys (was create-only).
  • Enable/disable toggle — a disabled key gets 403 aor_key_disabled (not 401, since the credential is valid).
  • Revoke now cascade-deletes the key's budgets.

Budgets (spend limits)

  • A key may have zero or more budgets; no budgets = unlimited.
  • Each budget is a rolling window of arbitrary length (hours, days, whatever) with a dollar cap.
  • Lazy start: the window doesn't start ticking until the first cost-bearing event, then resets when it elapses (resets on the next request, not on a wall-clock schedule).
  • Multiple budgets are AND-enforced: a request is blocked if ANY budget is exceeded.
  • Pre-check blocks before any upstream call: 429 aor_budget_exceeded when exceeded; 403 aor_model_unpriced when a budgeted key calls an unpriced model.
  • Post-attempt charge accrues spend only on successful 2xx attempts (never on errors).

Pricing

  • Per-model and per-provider pricing in the config (USD per million tokens across input/output/cache_read/cache_write).
  • Resolution: provider override → model default → None (unpriced).
  • A model with no pricing block is unpriced: budgeted keys are blocked from it (infinite cost). Set explicit zeros to mark a model as free but still metered.
  • For mixed-pricing models (some providers priced, some not), budgeted keys are filtered to priced providers only — the walk never lands on an unpriced one.

Spend visibility (per-key detail page)

  • /admin/keys/{id}: usage summary (attempts, ok/errors, input/output/cache tokens, total $), daily spend table, per-model token/cost breakdown, budgets table with progress bars + window state, add-budget form, recent attempts with token/cost columns.

Token capture & metering

  • Buffered: parses the 2xx body for usage (OpenAI/Anthropic/Google native shapes, or OpenAI-shape after translation).
  • Streaming: injects stream_options.include_usage=true upstream (OpenAI passthrough) and forces include_usage=True on the translator (translated streams) so usage is always present; a meter wraps the stream and fills a TokenCounts holder as bytes flow to the client unchanged. The streaming background task back-fills the usage row and charges budgets after the stream drains.
  • Truncated streams (client disconnect): partial token capture, partial cost (best-effort).
  • Spend precision: cent-millis (1/1000 cent) internally so sub-cent requests still accrue (integer cents would silently round to 0 and budgets would never grow).

Config example

[models."glm-5.2"]
pricing = { input_per_million = 0.10, output_per_million = 0.30,
           cache_read_per_million = 0.01, cache_write_per_million = 0.05 }

[models."claude-sonnet-5"]
  [models."claude-sonnet-5".providers."zen-anthropic"]
  pricing = { input_per_million = 3.0, output_per_million = 15.0 }

Design doc

Full design (decisions, enforcement flow, race-condition acceptance, known limitations) in docs/requirements-key-budgets.md.

Testing

  • uv run ruff check .
  • uv run ruff format --check .
  • uv run pytest -q → 280 passed

New test files: tests/test_budgets.py, tests/test_keys.py, tests/test_pricing.py. New tests in test_proxy.py (token capture, budget enforcement, mixed-pricing filter, end-to-end accrual), test_admin.py (detail page, rename, toggle, budget CRUD), test_translate_anthropic.py (cache fields surfaced), test_db.py (migrations), test_config.py (pricing resolution).

Known limitations (documented in the design doc)

  • TOCTOU overshoot: budgets can overshoot by ~one in-flight request per concurrent slot (accepted soft limit, matches the existing backoff module's posture).
  • Mid-stream translation errors: yield [DONE] but are logged as success — billed at full captured-token cost (inherent to no-mid-stream-retry).
  • Truncated streams: partial token capture (best-effort).
  • Pricing accuracy: admin-maintained; may not match upstream's current real pricing.
  • No CSRF protection: admin UI mutations inherit the "single shared password" posture.

Commits

The branch is split into focused commits: schema migrations → pricing config → token extraction/cost → proxy metering → budget enforcement + disabled-key 403 → admin UI → docs → review fixes.

## Summary Adds key spend limits (budgets) and per-key spend/usage visibility to the admin UI. ### Key management - **Rename** keys (was create-only). - **Enable/disable toggle** — a disabled key gets `403 aor_key_disabled` (not 401, since the credential is valid). - **Revoke** now cascade-deletes the key's budgets. ### Budgets (spend limits) - A key may have **zero or more budgets**; no budgets = unlimited. - Each budget is a **rolling window of arbitrary length** (hours, days, whatever) with a **dollar cap**. - **Lazy start**: the window doesn't start ticking until the first cost-bearing event, then resets when it elapses (resets on the next request, not on a wall-clock schedule). - **Multiple budgets** are AND-enforced: a request is blocked if ANY budget is exceeded. - **Pre-check** blocks before any upstream call: `429 aor_budget_exceeded` when exceeded; `403 aor_model_unpriced` when a budgeted key calls an unpriced model. - **Post-attempt charge** accrues spend only on successful 2xx attempts (never on errors). ### Pricing - Per-model and per-provider pricing in the config (USD per million tokens across input/output/cache_read/cache_write). - Resolution: provider override → model default → `None` (unpriced). - A model with no pricing block is **unpriced**: budgeted keys are blocked from it (infinite cost). Set explicit zeros to mark a model as free but still metered. - For **mixed-pricing** models (some providers priced, some not), budgeted keys are filtered to priced providers only — the walk never lands on an unpriced one. ### Spend visibility (per-key detail page) - `/admin/keys/{id}`: usage summary (attempts, ok/errors, input/output/cache tokens, total $), daily spend table, per-model token/cost breakdown, budgets table with progress bars + window state, add-budget form, recent attempts with token/cost columns. ### Token capture & metering - **Buffered**: parses the 2xx body for usage (OpenAI/Anthropic/Google native shapes, or OpenAI-shape after translation). - **Streaming**: injects `stream_options.include_usage=true` upstream (OpenAI passthrough) and forces `include_usage=True` on the translator (translated streams) so usage is always present; a meter wraps the stream and fills a `TokenCounts` holder as bytes flow to the client unchanged. The streaming background task back-fills the usage row and charges budgets after the stream drains. - **Truncated streams** (client disconnect): partial token capture, partial cost (best-effort). - **Spend precision**: cent-millis (1/1000 cent) internally so sub-cent requests still accrue (integer cents would silently round to 0 and budgets would never grow). ### Config example ```toml [models."glm-5.2"] pricing = { input_per_million = 0.10, output_per_million = 0.30, cache_read_per_million = 0.01, cache_write_per_million = 0.05 } [models."claude-sonnet-5"] [models."claude-sonnet-5".providers."zen-anthropic"] pricing = { input_per_million = 3.0, output_per_million = 15.0 } ``` ## Design doc Full design (decisions, enforcement flow, race-condition acceptance, known limitations) in `docs/requirements-key-budgets.md`. ## Testing - `uv run ruff check .` ✓ - `uv run ruff format --check .` ✓ - `uv run pytest -q` → 280 passed New test files: `tests/test_budgets.py`, `tests/test_keys.py`, `tests/test_pricing.py`. New tests in `test_proxy.py` (token capture, budget enforcement, mixed-pricing filter, end-to-end accrual), `test_admin.py` (detail page, rename, toggle, budget CRUD), `test_translate_anthropic.py` (cache fields surfaced), `test_db.py` (migrations), `test_config.py` (pricing resolution). ## Known limitations (documented in the design doc) - **TOCTOU overshoot**: budgets can overshoot by ~one in-flight request per concurrent slot (accepted soft limit, matches the existing backoff module's posture). - **Mid-stream translation errors**: yield `[DONE]` but are logged as success — billed at full captured-token cost (inherent to no-mid-stream-retry). - **Truncated streams**: partial token capture (best-effort). - **Pricing accuracy**: admin-maintained; may not match upstream's current real pricing. - **No CSRF protection**: admin UI mutations inherit the "single shared password" posture. ## Commits The branch is split into focused commits: schema migrations → pricing config → token extraction/cost → proxy metering → budget enforcement + disabled-key 403 → admin UI → docs → review fixes.
Schema foundations for key spend limits and visibility:

- New budgets table: rolling-window spend caps per client key. window_start
  is NULL until the first cost-bearing event (lazy start). spend tracked in
  cent-millis (1/1000 cent) to avoid sub-cent truncation. window_end is
  computed (window_start + window_seconds), not stored, to avoid drift.
- usage table gains input_tokens, output_tokens, cache_read_tokens,
  cache_write_tokens, cost_cm — all nullable (NULL for error attempts or
  unmetered requests). New idx_usage_key_ts index for the per-key detail
  page (spend-over-time queries filtered by key, ordered by ts).
- client_keys gains enabled (default 1) for the disable-without-revoke
  toggle. revoke_key will cascade-delete budgets (wired in a later commit).
- Migrations follow the existing PRAGMA-table_info-guarded ALTER TABLE
  pattern and are idempotent.
Pricing is optional in USD per million tokens across four token classes:
input, output, cache_read, cache_write. Resolution is provider override to
model-level default to None (unpriced).

- A model with no pricing block resolves to None (unpriced): budgeted
  keys are blocked from unpriced models (infinite cost), while unbudgeted
  keys are unaffected. Set explicit zeros to mark a model as free but
  still meter its tokens.
- A per-provider override under
  [models.<id>.providers.<name>.pricing] wins over the model-level
  default for that provider only.
- The ModelProviderOverride.priority is now optional (nullable) so an
  override can set only pricing without forcing priority to 0; when
  priority is None the provider default priority is used.
- config.example.toml demonstrates both model-level and per-provider
  pricing, with a note on the unpriced-vs-free distinction.
aor/pricing.py provides:
- TokenCounts dataclass (mutable for streaming meter fills)
- compute_cost_cm: tokens x per-million pricing -> cent-millis (1/1000
  cent) to avoid sub-cent truncation that would silently zero budgets
- extract_usage_{openai,anthropic,google}: pull token counts from a
  native buffered 2xx body. OpenAI extraction also handles
  translator-produced bodies (already normalized to OpenAI shape).
  Covers cache_read/cache_write where the upstream exposes them.

No adapter Protocol change yet; these are standalone functions. The proxy
will call them by adapter name in a subsequent commit.
Metering foundations for budget enforcement:

- log_attempt now accepts input/output/cache_read/cache_write tokens and
  cost_cm, and returns the inserted row id (for streaming back-fill).
- New update_usage_tokens_cost back-fills a usage row after a stream drains.
- _Outcome carries tokens, cost_cm, pricing, and stream_resp/stream_client
  so handle_proxy can wire post-stream accounting.
- Buffered 2xx: extract usage from the body (OpenAI/Anthropic/Google
  native, or OpenAI-shape after translation) and compute cost_cm inline.
- Streaming 2xx: wrap the body iterator in a meter that parses SSE
  data: lines for usage without altering bytes sent to the client; a
  TokenCounts holder is filled as the stream drains.
  - Passthrough: _metered_passthrough parses native usage shapes per adapter.
  - Translated: _safe_translate_stream re-parses the OpenAI-shaped output
    for a usage chunk (the translator already emits one with include_usage).
- Inject stream_options.include_usage=true on OpenAI streaming upstream
  requests so the upstream always emits a usage chunk for the meter.
- Streaming background task (_stream_accounting) closes the upstream,
  back-fills the usage row, and charges budgets (best-effort, suppressed).
- Buffered success charges budgets inline (best-effort).
- Budget charging is lazy and suppressed until the budgets module lands.
- aor/budgets.py: rolling-window budget module. create/list/delete/check/
  charge. Windows start lazily (window_start NULL until first cost-bearing
  event), reset on expiry, and accrue spend in cent-millis to avoid
  sub-cent truncation. check() raises BudgetExceeded or ModelUnpriced;
  charge() updates all budgets in a single transaction. A key with no
  budgets is unlimited.
- aor/keys.py: update_key (rename, enable/disable), get_key, and
  lookup_key now treats disabled keys as absent. revoke_key cascades
  deletes of the key's budgets.
- aor/auth.py: authenticate_client returns an AuthResult distinguishing
  invalid (401) from disabled (403) keys, so disabled keys surface as
  403 aor_key_disabled rather than 401.
- aor/app.py: catch-all route maps auth outcomes to 401/403 and passes
  the key_hash to handle_proxy.
- aor/proxy.py: budget pre-check in handle_proxy after model lookup.
  A model is "priced" if its model-level default or any serving provider
  has a pricing block; budgeted keys are 403-blocked from unpriced
  models. Exceeded budgets 429 with aor_budget_exceeded.
- Buffered and streaming successes charge budgets best-effort (suppressed
  errors match the backoff module's soft-limit posture).
- TOCTOU overshoot (bounded by concurrency) is documented and accepted.
- aor/stats.py: key_summary, key_spend_over_time (per-day cost buckets),
  key_usage_by_model (per-model token/cost breakdown), key_recent_attempts
  scoped to one client key.
- aor/admin.py: new routes
  - GET  /admin/keys/{id}        detail page (summary, spend chart, per-model,
    budgets, recent attempts)
  - POST /admin/keys/{id}/rename
  - POST /admin/keys/{id}/toggle  enable/disable
  - POST /admin/keys/{id}/budgets  add budget (window hours/days + $ limit)
  - POST /admin/keys/{id}/budgets/{bid}/delete
  Budget create accepts either window_seconds (JS path) or window_hours
  (JS-less fallback).
- templates/key_detail.html: usage summary grid, daily spend table,
  per-model breakdown, budgets table with progress bars + window state,
  add-budget form, recent-attempts table with token/cost columns.
- templates/keys.html: name is a link to the detail page; new Status column
  with enable/disable toggle; row shows enabled/disabled inline.

All new routes require the admin session (redirect to login when unset),
matching the existing posture.
Captures the design decisions for the key budgets & spend visibility
feature: pricing config resolution, cent-millis spend precision, lazy-
start rolling budgets, token capture (buffered + streaming), enforcement
flow, accepted race conditions, and known limitations. AGENTS.md module
list updated to reference budgets.py, pricing.py, the new usage columns,
and the AuthResult status distinction.
Two HIGH-severity enforcement gaps found in review, plus a token-capture
accuracy fix:

- Translated streaming now forces include_usage=True on the translator
  (not just the openai-passthrough path). Without it, a budgeted key
  streaming via an Anthropic upstream would never accrue spend (the
  translator only emits a usage chunk when include_usage is set, which
  was the client's opt-in). The meter now sees usage for every translated
  stream. (HIGH)
- Budgeted keys are now filtered to priced providers only: a model with
  mixed pricing (some providers priced, others unpriced) would let the
  walk land on an unpriced provider and serve the request without
  charging any budget. The candidate sequence is now filtered to
  prov.pricing is not None when the key has budgets; if that empties the
  sequence, the pre-check blocks with aor_model_unpriced. Unbudgeted keys
  see the full sequence. (HIGH)
- Translated Anthropic responses now surface cache_read/cache_write via
  the OpenAI prompt_tokens_details/completion_tokens_details convention
  (buffered _convert_usage and the streaming usage chunk), so the proxy's
  OpenAI usage extractor recovers them and cache pricing actually accrues.
  _drain_sse_usage reads completion_tokens_details.reasoning_tokens for
  cache_write. (MEDIUM)

Also: end-to-end tests that a successful (buffered + streaming) request
accrues spend against a budget; a mixed-pricing test that a budgeted key
skips the unpriced provider; and a test that an all-unpriced model blocks
budgeted keys. LOW cleanups: _key_hash_for_id uses db.connect; dead
budgets_view fields dropped; two sloppy test lines fixed.
- window_ts now recognizes "30d" (previously fell through to return 0,
  the all-time case). The key detail page defaults to "30d" and showed
  all-time data while highlighting "30d" as active. Add a 30d=30*86400
  branch. (MEDIUM)
- extract_usage_openai now reads completion_tokens_details.reasoning_tokens
  into cache_write. The Anthropic translator was modified to surface
  cache_creation_input_tokens via the OpenAI completion_tokens_details
  convention, and the streaming meter (_drain_sse_usage) was already
  reading it, but the buffered extractor was not — cache_write tokens
  for translated Anthropic responses were silently dropped, under-accruing
  budgets (cache-write pricing is typically 1.25x input). (MEDIUM)
- _stream_accounting no longer skips the usage-row token back-fill when
  pricing is None. Streaming unpriced models previously got NULL token
  columns while the buffered path populated them. The stream now back-fills
  tokens always; cost compute + budget charge still gated on pricing.
  (LOW)

Tests: window_ts("30d") maps to 30 days; extract_usage_openai with
completion_tokens_details surfaces cache_write; streaming an unpriced
model back-fills tokens.
Sign in to join this conversation.
No reviewers
No labels
No milestone
No project
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
devtrev/actually-open-router!4
No description provided.