feat(usage): add AI budget usage caps

Merge usage-cap budget controls with conflict resolution against dev.
This commit is contained in:
Duy /zuey/ authored and GitHub committed 2026-05-24 13:57:31 +07:00
commit 2bb736cf18
108 files changed
+5562 -387

No files matched your search

+23
View File
@@ -54,6 +54,29 @@ Streaming fallback is conservative: backup models are tried only if the stream f
---
## Usage Cap Pricing Enforcement
Standard edition can enforce AI budget caps before billable provider dispatch. API-key providers use OpenRouter `/models` pricing as the catalog source, with optional tenant/provider/model overrides in the dashboard.
Excluded provider classes:
- `chatgpt_oauth`, `claude_cli`, and `bailian` are treated as subscription/non-API pricing in round one.
- local/no-key subprocess providers such as `acp` and `ollama` are skipped unless a future feature explicitly enables pricing for them.
Runtime flow:
1. Resolve the stored provider by name and skip non-billable provider classes.
2. Load matching policies for tenant, agent, provider, provider type, and model.
3. Resolve custom pricing override first, then OpenRouter catalog pricing when a matching policy has a cost ceiling. Native provider model IDs are mapped to OpenRouter prefixes for common providers such as OpenAI, Anthropic, and Gemini.
4. Reserve estimated tokens and cost atomically before each dispatch attempt.
5. Reconcile reserved counters after the provider returns usage or after a failed call.
Token-only policies do not require catalog pricing. Model fallback routes reserve against the actual candidate provider/model before each attempt. Cached input is separated from uncached input for OpenAI-compatible usage accounting. Partial stream failures keep the estimate, or actual provider usage when available, instead of clearing billed output to zero. Internal LLM calls for memory flush, compaction, media reading tools (`read_image`, `read_document`, `read_audio`, `read_video`), and subagents use the same preflight/reconcile path.
The legacy agent-level `budget_monthly_cents` field is treated as a generated monthly agent USD cap. Existing values are backfilled during migration, and later agent budget edits update or remove the generated cap policy.
Supported price units: input, output, cache read, cache write, reasoning, request, image, and web search.
---
## 2. Supported Providers
### Six Core Provider Types
+28
View File
@@ -45,6 +45,7 @@ The `Stores` struct is the top-level container holding all PostgreSQL-backed sto
| ContactStore | `PGContactStore` | Channel contacts (auto-collected), cross-channel deduplication, merge |
| ActivityStore | `PGActivityStore` | Audit logs, action tracking, compliance |
| SnapshotStore | `PGSnapshotStore` | Hourly usage snapshots, cost aggregation, time series queries |
| UsageCapStore | `PGUsageCapStore` | OpenRouter pricing catalog, pricing overrides, cap policies, reservations, counters, events |
| SecureCLIStore | `PGSecureCLIStore` | CLI binary configs with encrypted credential injection |
| APIKeyStore | `PGAPIKeyStore` | Gateway API keys, scopes, expiration, revocation |
| HookStore | `PGHookStore` | Lifecycle hook definitions (event, handler type, matcher, config), execution audit log |
@@ -85,6 +86,33 @@ Migration versions:
---
## Usage Cap Storage
Usage cap enforcement is Standard/PostgreSQL-only in round one. The `UsageCapStore` is wired on the PostgreSQL store factory and left nil in SQLite/Lite builds.
Tables:
- `usage_pricing_catalog`: OpenRouter model catalog prices, raw upstream model payload, sync time.
- `usage_pricing_overrides`: tenant/provider/model override prices for custom billing assumptions.
- `usage_cap_policies`: cap definitions scoped by tenant, agent, provider, provider type, model, `window_key`, and `source`.
- `usage_cap_counters`: current window used and reserved token/cost counters.
- `usage_cap_reservations`: preflight reservations keyed by LLM call attempt.
- `usage_cap_events`: allow/block/reconcile/skip audit events.
Reservation updates are atomic: counters are updated only when `used + reserved + estimate` remains below configured token and cost ceilings.
Reservation keys are idempotent per policy, so a retry using the same key does not double-increment reserved counters.
Policy `agent_id` references must belong to the same tenant as the policy. Policy and pricing override `provider_id` references may belong to the same tenant or the master tenant for default provider fallback, but not another non-master tenant.
Catalog and override price fields are nullable decimal strings with non-negative validation in the store layer and database checks.
Pricing resolution checks exact override/catalog model IDs first, then provider-derived OpenRouter aliases for native unprefixed model IDs.
Agent `budget_monthly_cents` values are bridged into `usage_cap_policies` with `source = 'agent_budget_monthly_cents'`, an agent scope, `window_key = 'month'`, and `max_cost_micros = budget_monthly_cents * 10000`. Updating or clearing the agent budget keeps that generated policy in sync; manual cap policies continue to use `source = 'manual'`.
Migration versions:
- PostgreSQL: `000070_usage_caps_pricing`, `000071_usage_cap_policies`, `000072_agent_budget_usage_cap_bridge`.
- SQLite: no schema change; feature is not active in Lite.
---
## 3. Session Caching
The session store uses an in-memory write-behind cache to minimize database I/O during the agent tool loop. All reads and writes happen in memory; data is flushed to the persistent backend only when `Save()` is called at the end of a run.
+53
View File
@@ -390,6 +390,36 @@ LLM provider management. API keys are encrypted with AES-256-GCM in the database
| `GET` | `/v1/embedding/status` | Check global embedding availability |
| `GET` | `/v1/providers/claude-cli/auth-status` | Check Claude CLI login status |
### Model Pricing
| Method | Path | Description |
|--------|------|-------------|
| `POST` | `/v1/model-pricing/sync-openrouter` | Sync OpenRouter `/models` pricing catalog (master scope) |
| `GET` | `/v1/model-pricing` | Search catalog models by `model` query |
| `GET` | `/v1/model-pricing/overrides` | List tenant pricing overrides, optionally by `provider_id` |
| `PUT` | `/v1/model-pricing/overrides` | Upsert provider/model custom pricing |
| `DELETE` | `/v1/model-pricing/overrides/{id}` | Delete pricing override |
Override body:
```json
{
"provider_id": "0193a5b0-7000-7000-8000-000000000123",
"provider_type": "openrouter",
"model_id": "anthropic/claude-sonnet-4-5",
"pricing": {
"input": "0.000003",
"output": "0.000015",
"cache_read": "0.0000003",
"cache_write": "0.00000375",
"reasoning": "0.000015",
"request": "0",
"image": "0",
"web_search": "0"
}
}
```
**Supported types:** `anthropic_native`, `openai_compat`, `chatgpt_oauth`, `gemini_native`, `dashscope`, `bailian`, `minimax`, `claude_cli`, `acp`
Reconnect response:
@@ -1271,11 +1301,34 @@ Follow response:
| `GET` | `/v1/usage/timeseries` | Time-series usage points |
| `GET` | `/v1/usage/breakdown` | Breakdown by provider/model/channel |
| `GET` | `/v1/usage/summary` | Summary with period comparison |
| `GET` | `/v1/usage-caps/policies` | List usage cap policies |
| `POST` | `/v1/usage-caps/policies` | Create token/cost cap policy |
| `PATCH` | `/v1/usage-caps/policies/{id}` | Update cap policy |
| `DELETE` | `/v1/usage-caps/policies/{id}` | Delete cap policy |
| `GET` | `/v1/usage-caps/utilization` | Current-window used and reserved counters |
| `GET` | `/v1/usage-caps/events` | Recent allow/block/reconcile/skip events |
**Query params:** `from`, `to` (RFC 3339), `agent_id`, `provider`, `model`, `channel`, `group_by`
**Periods:** `24h`, `today`, `7d`, `30d`
Usage cap policy body:
```json
{
"agent_id": "optional-agent-uuid",
"provider_id": "optional-provider-uuid",
"provider_type": "openrouter",
"model_id": "anthropic/claude-sonnet-4-5",
"window": "day",
"max_tokens": 500000,
"max_cost_usd": 25,
"enabled": true
}
```
Policy responses include read-only `source`. `source="agent_budget_monthly_cents"` means the policy is generated from the agent monthly budget field and must be changed there; direct policy updates/deletes return `409 Conflict`.
---
## 24. Activity & Audit
+20
View File
@@ -42,6 +42,26 @@ Significant changes, features, and fixes in reverse chronological order.
## 2026-05-22
### Usage Cap budget controls
**Features**
- Added Standard/PostgreSQL usage caps for AI budget control by hour, day, week, or month.
- Caps support token and USD cost ceilings at tenant, agent, provider, provider type, and model scopes.
- Added OpenRouter catalog sync plus tenant/provider/model pricing overrides for input, output, cache read/write, reasoning, request, image, and web search units.
- Enforced caps in agent, fallback model, subagent, memory flush, compaction, and media reading tools (`read_image`, `read_document`, `read_audio`, `read_video`) with preflight reservation and post-call reconciliation.
- Added non-negative validation for catalog and override pricing fields.
- Added OpenRouter alias resolution for native model IDs, cached-input accounting normalization, and partial-stream failure reconciliation.
- Bridged legacy `budget_monthly_cents` into generated monthly agent USD cap policies, including migration backfill and save-time sync.
- Added web dashboard controls on Usage and Provider detail pages.
- Added Usage page editing for manual cap policies, including enable/disable, scope clearing, and token/USD limit clearing while keeping generated agent-budget caps read-only.
- Added usage-cap decision metadata to LLM spans, including allow/skip/block reason, policy IDs, estimates, actuals, and reconcile status.
**Tests**
- Added pricing and cap service coverage.
- Verified Go builds, SQLite build compatibility, full Go test suite, integration race suite, and Web UI production build.
### Messaging debounce hardening
**Features**