feat: key budgets, spend limits, and per-key spend visibility #4
Loading…
Add table
Add a link
Reference in a new issue
No description provided.
Delete branch "trevin/feat/key-budgets-spend-visibility"
Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?
Summary
Adds key spend limits (budgets) and per-key spend/usage visibility to the admin UI.
Key management
403 aor_key_disabled(not 401, since the credential is valid).Budgets (spend limits)
429 aor_budget_exceededwhen exceeded;403 aor_model_unpricedwhen a budgeted key calls an unpriced model.Pricing
None(unpriced).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
stream_options.include_usage=trueupstream (OpenAI passthrough) and forcesinclude_usage=Trueon the translator (translated streams) so usage is always present; a meter wraps the stream and fills aTokenCountsholder as bytes flow to the client unchanged. The streaming background task back-fills the usage row and charges budgets after the stream drains.Config example
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 passedNew test files:
tests/test_budgets.py,tests/test_keys.py,tests/test_pricing.py. New tests intest_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)
[DONE]but are logged as success — billed at full captured-token cost (inherent to no-mid-stream-retry).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.
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/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.- 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.