From c04599db3d9fdda39ff1729768dcdff3feaad8fb Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Sat, 10 Oct 2026 22:18:17 +0800 Subject: [PATCH 1/3] feat(gateway): GET /v1/usage reports a key's usage and the limits binding it A client holding a gateway API key had no way to ask how much room the key has left before a request is refused. GET /v1/usage answers that on the gateway port, authenticated exactly as a model request is. `limits` is built with `limits_for_ai_gateway`, the function the pre-flight stages use, so it lists the same rules and caps a request is held to: the key's own (scope "key", counted on the lineage) and the owner's effective limits (scope "user": roles, including team-granted roles, merged most-restrictive, then overrides; counted on the owner's counters, shared by all of the owner's keys). `used` is read from the same Redis counters the same way, without charging them, and the list is ordered by the share left, least first. A Redis read error answers 503 rather than zeros. `usage` needs a count for keys that have no limits, which no existing counter provides (limit counters exist only for configured limits; the request log is optional and holds raw tokens). The gateway now counts each key lineage's requests and weighted tokens per UTC day and month in Redis, at the points the limits engine counts them: a request once it passes the rate limits, its tokens in post-flight accounting. The writes fail open. The month's cost comes from gateway_logs, as the cost reports read it, and is null without ClickHouse. The endpoint sits behind a read-only variant of the API-key middleware: same lookup, checks and refusals, but it neither updates last_used_at (a poller must not keep an idle key past its inactivity timeout) nor arms the early-cancel request row. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 54 +++ crates/common/src/limits/budget.rs | 57 ++- crates/common/src/limits/key_usage.rs | 168 +++++++ crates/common/src/limits/mod.rs | 6 + crates/common/src/limits/sliding.rs | 18 +- crates/gateway/src/proxy/accounting.rs | 9 + crates/gateway/src/proxy/identity.rs | 2 +- crates/gateway/src/proxy/mod.rs | 5 + crates/gateway/src/proxy/pipeline.rs | 10 +- crates/server/src/app.rs | 12 + crates/server/src/handlers/key_usage.rs | 456 ++++++++++++++++++ crates/server/src/handlers/mod.rs | 1 + crates/server/src/middleware/api_key_auth.rs | 90 +++- .../tests/analytics_clickhouse.rs | 67 +++ crates/test-support/tests/limits.rs | 451 +++++++++++++++++ web/src/i18n/en.json | 1 + web/src/i18n/zh.json | 1 + web/src/routes/guide.tsx | 4 + 18 files changed, 1380 insertions(+), 32 deletions(-) create mode 100644 crates/common/src/limits/key_usage.rs create mode 100644 crates/server/src/handlers/key_usage.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 4b1e4763..ab91ef8f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,60 @@ target. ## [Unreleased] +### Read before upgrading + +- **Each API key gets day and month usage counters in Redis.** An + AI-gateway request made with a key now also counts on + `key_usage:{api_key_lineage:}:daily:` and + `…:monthly:`, one hash each that expires two periods after its last + write: one more Redis script per request once it passes the rate limits, + and one when its tokens are counted. The counters start empty, so the + `usage` that `GET /v1/usage` reports counts from the upgrade on. No + setting, database schema or Helm value changes. + +### Added + +- **`GET /v1/usage` tells a client what its API key has used and what room + it has left.** Called on the gateway port with a gateway API key, as a + model request is, it answers with what the key has done today and this + month, every limit that binds its model requests with what has been used + of it, and when the key stops working: + + ```json + { + "usage": {"requests_today": 12, "tokens_today": 48210, + "requests_month": 340, "tokens_month": 1290455, + "cost_usd_month": 3.82}, + "limits": [ + {"scope": "key", "kind": "tokens", "window": "daily", "window_secs": null, + "limit": 100000, "used": 48210, "resets_at": "2026-10-12T00:00:00Z"}, + {"scope": "user", "kind": "requests", "window": "1m", "window_secs": 60, + "limit": 60, "used": 4, "resets_at": null} + ], + "expires_at": "2026-12-31T00:00:00Z" + } + ``` + + `usage` counts this key alone, across its rotations: the requests the + rate limits let through and the weighted tokens limits count, for the + UTC day and month, and its cost this month from the request log (`null` + without ClickHouse). `limits` lists the key's own limits (`scope: key`, + counted for the key) and its owner's effective limits (`scope: user`: + roles, including those a team grants, merged most-restrictive, then the + user's overrides, counted for everything the owner does through any key), + each with `used` read from the counter that refuses requests, the one + with the least left first. `window` is a rate limit's sliding window + (`1m`, `5m`, `1h`, `5h`, `1d`, `1w`, with its length in `window_secs`) or + a budget's calendar period (`daily`, `weekly`, `monthly`, with its end in + `resets_at`, UTC). Limits on the MCP gateway are not listed. `expires_at` + is the key's expiry or the end of its rotation grace period, whichever + comes first. A key a model request would refuse gets the same `401` or + `403`, and `503` means the counters could not be read. Calling it charges + no limit, writes no request log row and is not a use of the key: + `last_used_at` stays as it was, so polling does not keep an idle key from + its inactivity timeout. The console's Configuration Guide lists it with + the other gateway endpoints. + ## [3.4.0] — 2026-10-11 Conversations with reasoning models behind Chat-format upstreams keep their diff --git a/crates/common/src/limits/budget.rs b/crates/common/src/limits/budget.rs index 4869ccfb..e9e5dcb3 100644 --- a/crates/common/src/limits/budget.rs +++ b/crates/common/src/limits/budget.rs @@ -188,24 +188,53 @@ pub async fn current_spend( let now = Utc::now(); let mut out = Vec::with_capacity(caps.len()); for cap in caps { - let key = build_key( - cap.subject_kind.as_str(), - cap.subject_id, - cap.period.as_str(), - now, - ); - let v: Option = redis.get(&key).await.ok(); - out.push(CapStatus { - cap_id: cap.id, - subject_kind: cap.subject_kind, - subject_id: cap.subject_id, - current: v.unwrap_or(0), - limit: cap.limit_tokens, - }); + let current = spend(redis, cap, now).await.unwrap_or(0); + out.push(cap_status(cap, current)); } Ok(out) } +/// [`current_spend`] at `now`, except that a Redis error is returned +/// rather than read as nothing spent. +pub async fn read_spend( + redis: &Client, + caps: &[BudgetCap], + now: DateTime, +) -> Result, fred::error::Error> { + let mut out = Vec::with_capacity(caps.len()); + for cap in caps { + out.push(cap_status(cap, spend(redis, cap, now).await?)); + } + Ok(out) +} + +/// The counter of `cap`'s period containing `now`; a period nothing has +/// been spent in yet has no counter. +async fn spend( + redis: &Client, + cap: &BudgetCap, + now: DateTime, +) -> Result { + let key = build_key( + cap.subject_kind.as_str(), + cap.subject_id, + cap.period.as_str(), + now, + ); + let v: Option = redis.get(&key).await?; + Ok(v.unwrap_or(0)) +} + +fn cap_status(cap: &BudgetCap, current: i64) -> CapStatus { + CapStatus { + cap_id: cap.id, + subject_kind: cap.subject_kind, + subject_id: cap.subject_id, + current, + limit: cap.limit_tokens, + } +} + /// Add `weighted_tokens` to every cap counter in the slice and /// return the new running totals so the caller can fire alerts on /// crossings. Soft cap: this never blocks the request — it only diff --git a/crates/common/src/limits/key_usage.rs b/crates/common/src/limits/key_usage.rs new file mode 100644 index 00000000..40c608c9 --- /dev/null +++ b/crates/common/src/limits/key_usage.rs @@ -0,0 +1,168 @@ +// ============================================================================ +// Per-key calendar usage +// +// What each API key (lineage) has been let through and used today and +// this month, UTC, whether or not it has limits of its own: the rate-limit +// and budget counters only exist for the limits that are configured, so +// they can't say what a key without limits has done. `GET /v1/usage` +// reads these. +// +// They count what the limits count, at the moments the limits engine +// counts it: a request once it has passed the rate limits (where +// `sliding::admit` charges `requests` rules), and its weighted tokens +// after the call (where `sliding::record` and +// `budget::add_weighted_tokens` add them). A request the limits refuse +// counts nothing. +// +// Storage (Redis): one hash per lineage per period, fields `requests` +// and `tokens`: +// +// key_usage:{api_key_lineage:}:daily:2026-10-11 +// key_usage:{api_key_lineage:}:monthly:2026-10 +// +// The braces are a Redis Cluster hash tag, so one script writes both. +// Each hash expires 2 × its period after its last write, as the budget +// counters do. +// ============================================================================ + +use chrono::{DateTime, Utc}; +use fred::clients::Client; +use fred::interfaces::{HashesInterface, LuaInterface}; +use uuid::Uuid; + +use super::budget::bucket_id; + +const DAY_TTL_SECS: i64 = 2 * 86_400; +const MONTH_TTL_SECS: i64 = 2 * 31 * 86_400; + +/// What one key has used today and this month (UTC). +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct KeyUsageCounts { + pub requests_today: i64, + pub tokens_today: i64, + pub requests_month: i64, + pub tokens_month: i64, +} + +/// The two hashes of `lineage` for the day and the month containing `now`. +pub fn keys(lineage: Uuid, now: DateTime) -> [String; 2] { + ["daily", "monthly"].map(|period| { + format!( + "key_usage:{{api_key_lineage:{lineage}}}:{period}:{}", + bucket_id(period, now) + ) + }) +} + +const LUA_ADD: &str = r#" +-- KEYS: day hash, month hash +-- ARGV: field, amount, day ttl, month ttl +redis.call('HINCRBY', KEYS[1], ARGV[1], ARGV[2]) +redis.call('EXPIRE', KEYS[1], ARGV[3]) +redis.call('HINCRBY', KEYS[2], ARGV[1], ARGV[2]) +redis.call('EXPIRE', KEYS[2], ARGV[4]) +return 1 +"#; + +async fn add( + redis: &Client, + lineage: Uuid, + field: &str, + amount: i64, + now: DateTime, +) -> Result<(), fred::error::Error> { + if amount <= 0 { + return Ok(()); + } + let _: i64 = redis + .eval( + LUA_ADD, + keys(lineage, now).to_vec(), + vec![ + field.to_string(), + amount.to_string(), + DAY_TTL_SECS.to_string(), + MONTH_TTL_SECS.to_string(), + ], + ) + .await?; + Ok(()) +} + +/// Count one request the rate limits let through. +pub async fn record_request( + redis: &Client, + lineage: Uuid, + now: DateTime, +) -> Result<(), fred::error::Error> { + add(redis, lineage, "requests", 1, now).await +} + +/// Add a call's weighted tokens — the amount its limits were charged. +pub async fn record_tokens( + redis: &Client, + lineage: Uuid, + weighted_tokens: i64, + now: DateTime, +) -> Result<(), fred::error::Error> { + add(redis, lineage, "tokens", weighted_tokens, now).await +} + +/// Read `lineage`'s counts for the day and month containing `now`. A +/// period nothing was counted in reads as zero; a Redis error is +/// returned. +pub async fn read( + redis: &Client, + lineage: Uuid, + now: DateTime, +) -> Result { + let [day, month] = keys(lineage, now); + let fields = || vec!["requests", "tokens"]; + let d: Vec> = redis.hmget(&day, fields()).await?; + let m: Vec> = redis.hmget(&month, fields()).await?; + let at = |v: &[Option], i: usize| v.get(i).copied().flatten().unwrap_or(0); + Ok(KeyUsageCounts { + requests_today: at(&d, 0), + tokens_today: at(&d, 1), + requests_month: at(&m, 0), + tokens_month: at(&m, 1), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::TimeZone; + + #[test] + fn a_keys_day_and_month_share_one_cluster_slot() { + use fred::util::redis_keyslot; + let lineage = Uuid::new_v4(); + let now = Utc.with_ymd_and_hms(2026, 10, 11, 23, 59, 59).unwrap(); + let [day, month] = keys(lineage, now); + assert_eq!( + day, + format!("key_usage:{{api_key_lineage:{lineage}}}:daily:2026-10-11") + ); + assert_eq!( + month, + format!("key_usage:{{api_key_lineage:{lineage}}}:monthly:2026-10") + ); + assert_eq!( + redis_keyslot(day.as_bytes()), + redis_keyslot(month.as_bytes()) + ); + } + + #[test] + fn the_buckets_turn_over_at_utc_midnight_and_the_first() { + let lineage = Uuid::nil(); + let before = Utc.with_ymd_and_hms(2026, 10, 31, 23, 59, 59).unwrap(); + let after = Utc.with_ymd_and_hms(2026, 11, 1, 0, 0, 0).unwrap(); + let [d1, m1] = keys(lineage, before); + let [d2, m2] = keys(lineage, after); + assert_ne!(d1, d2); + assert_ne!(m1, m2); + assert!(d2.ends_with(":daily:2026-11-01") && m2.ends_with(":monthly:2026-11")); + } +} diff --git a/crates/common/src/limits/mod.rs b/crates/common/src/limits/mod.rs index 8ea3bf14..b8c5d7a1 100644 --- a/crates/common/src/limits/mod.rs +++ b/crates/common/src/limits/mod.rs @@ -56,6 +56,7 @@ use sqlx::{PgPool, Row}; use uuid::Uuid; pub mod budget; +pub mod key_usage; pub mod sliding; pub mod weight; @@ -1158,6 +1159,10 @@ pub struct RequestLimits { /// The user the request runs as. Every rate-limit counter below /// carries their Redis Cluster hash tag (`sliding::counter_key`). pub owner: Uuid, + /// The lineage of the API key the request came with, whether or not + /// it has limits of its own: its `key_usage` counters count the + /// request. + pub key_lineage: Option, pub rules: Vec, pub caps: Vec, } @@ -1177,6 +1182,7 @@ impl RequestLimits { ) -> Self { let mut out = Self { owner, + key_lineage: key.map(|(lineage, _)| lineage), ..Self::default() }; out.add(surface, RateLimitSubject::User, owner, user); diff --git a/crates/common/src/limits/sliding.rs b/crates/common/src/limits/sliding.rs index fb340f13..f760b253 100644 --- a/crates/common/src/limits/sliding.rs +++ b/crates/common/src/limits/sliding.rs @@ -452,9 +452,8 @@ pub async fn record_at( /// Returns 0 on Redis error so the UI can fall back to "no data" /// rather than 500. Real failures are logged. pub async fn current_count(redis: &Client, rule: &ResolvedRule) -> i64 { - use fred::interfaces::HashesInterface; - match redis.hgetall::, _>(&rule.key).await { - Ok(buckets) => window_sum(&buckets, chrono::Utc::now().timestamp(), rule.bucket_secs), + match read_count(redis, rule, chrono::Utc::now().timestamp()).await { + Ok(count) => count, Err(e) => { tracing::warn!(key = %rule.key, "rate-limit usage read failed: {e}"); 0 @@ -462,6 +461,19 @@ pub async fn current_count(redis: &Client, rule: &ResolvedRule) -> i64 { } } +/// What the window that ends at `now_secs` holds — the sum [`admit`] +/// compares with the limit — without changing it. A Redis error is +/// returned, not read as an empty window. +pub async fn read_count( + redis: &Client, + rule: &ResolvedRule, + now_secs: i64, +) -> Result { + use fred::interfaces::HashesInterface; + let buckets = redis.hgetall::, _>(&rule.key).await?; + Ok(window_sum(&buckets, now_secs, rule.bucket_secs)) +} + /// The sum of the buckets inside the window that ends at `now_secs`. fn window_sum(buckets: &HashMap, now_secs: i64, bucket_secs: i32) -> i64 { let bucket_secs = i64::from(bucket_secs); diff --git a/crates/gateway/src/proxy/accounting.rs b/crates/gateway/src/proxy/accounting.rs index 3d17dda9..e9f61f0b 100644 --- a/crates/gateway/src/proxy/accounting.rs +++ b/crates/gateway/src/proxy/accounting.rs @@ -44,6 +44,15 @@ pub(crate) async fn post_flight_account( return; } + // The key's own day and month, counted whether or not it has limits. + if let Some(lineage) = request_limits.key_lineage + && let Err(e) = + limits::key_usage::record_tokens(&redis, lineage, weighted, chrono::Utc::now()).await + { + metrics::counter!("gateway_key_usage_fail_open_total").increment(1); + tracing::warn!("key usage token count failed: {e}"); + } + // Token-metric sliding rules — the user's and the key's, the same // rules the pre-flight checked. Recorded whatever they come to: a // window this request overshoots refuses the next one. Post-flight diff --git a/crates/gateway/src/proxy/identity.rs b/crates/gateway/src/proxy/identity.rs index 82747ef1..7543c1a1 100644 --- a/crates/gateway/src/proxy/identity.rs +++ b/crates/gateway/src/proxy/identity.rs @@ -10,7 +10,7 @@ use think_watch_common::limits::{RequestLimits, Surface}; /// user's counters, and the calling key's own limits on its lineage's /// counters — both apply. A request with no user (never past the auth /// middleware today) is held to nothing. -pub(super) fn limits_for_ai_gateway(identity: &GatewayRequestIdentity) -> RequestLimits { +pub fn limits_for_ai_gateway(identity: &GatewayRequestIdentity) -> RequestLimits { let parse = |s: &Option| s.as_deref().and_then(|s| Uuid::parse_str(s).ok()); let Some(user_id) = parse(&identity.user_id) else { return RequestLimits::default(); diff --git a/crates/gateway/src/proxy/mod.rs b/crates/gateway/src/proxy/mod.rs index 6370147d..9e036db4 100644 --- a/crates/gateway/src/proxy/mod.rs +++ b/crates/gateway/src/proxy/mod.rs @@ -48,6 +48,7 @@ pub use early_cancel::{EarlyCancel, EarlyCancelSlot}; pub use generate::{ proxy_anthropic_messages, proxy_chat_completion, proxy_gemini, proxy_responses, }; +pub use identity::limits_for_ai_gateway; pub use models::{list_gemini_models_handler, list_models_handler}; pub use responses_ws::proxy_responses_ws; @@ -131,6 +132,10 @@ pub struct GatewayRequestIdentity { /// every `gateway_logs` row carries it without each handler reading /// headers themselves. `None` only if extraction failed. pub ip_address: Option, + /// When the calling key stops authenticating: its expiry, or the end + /// of its rotation grace period, whichever comes first. `None` when + /// it has neither. + pub key_expires_at: Option>, } /// Thin wrapper kept for call-site readability; delegates to diff --git a/crates/gateway/src/proxy/pipeline.rs b/crates/gateway/src/proxy/pipeline.rs index 506f3175..a44ff572 100644 --- a/crates/gateway/src/proxy/pipeline.rs +++ b/crates/gateway/src/proxy/pipeline.rs @@ -24,7 +24,7 @@ use think_watch_common::lifecycle::stages::{ check_access, check_budget, check_limits, run_post_invoke, }; use think_watch_common::lifecycle::state::{CapturedView, Invoked, LimitCheckRecord, Raw}; -use think_watch_common::limits::RequestLimits; +use think_watch_common::limits::{RequestLimits, key_usage}; use tw_dialect::ir::Dialect; /// Pre-flight result threaded through to `ChatPostInvokeDeps` later @@ -78,6 +78,14 @@ pub(super) async fn run_preflight_stages( ) .await .map_err(short_circuit_to_response)?; + // Past the rate limits, the request counts on its key, as it has on + // every `requests` rule. Fail-open: the counts are for reading. + if let Some(lineage) = limits.key_lineage + && let Err(e) = key_usage::record_request(&state.redis, lineage, chrono::Utc::now()).await + { + metrics::counter!("gateway_key_usage_fail_open_total").increment(1); + tracing::warn!("key usage request count failed: {e}"); + } let _authorized = check_access::(limits_checked, model, &state.audit) .await .map_err(short_circuit_to_response)?; diff --git a/crates/server/src/app.rs b/crates/server/src/app.rs index b46b509c..07581abe 100644 --- a/crates/server/src/app.rs +++ b/crates/server/src/app.rs @@ -292,6 +292,17 @@ pub async fn create_gateway_app(_config: &AppConfig, state: AppState) -> anyhow: )) .with_state(gateway_state); + // Reads about the calling key itself: the same key checks as the + // routes above, but not a use of the key — nothing is charged, and + // neither `last_used_at` nor the request log moves. + let key_routes = Router::new() + .route("/v1/usage", get(handlers::key_usage::get_key_usage)) + .layer(axum::middleware::from_fn_with_state( + state.clone(), + crate::middleware::api_key_auth::require_api_key_to_read("ai_gateway"), + )) + .with_state(state.clone()); + // MCP Gateway: /mcp // // Use the shared registry + circuit breakers from AppState so that the @@ -400,6 +411,7 @@ pub async fn create_gateway_app(_config: &AppConfig, state: AppState) -> anyhow: let app = Router::new() .merge(health) .merge(ai_routes) + .merge(key_routes) .merge(mcp_routes) .layer(RequestBodyLimitLayer::new(10 * 1024 * 1024)) // 10MB for large prompts .layer(TimeoutLayer::with_status_code( diff --git a/crates/server/src/handlers/key_usage.rs b/crates/server/src/handlers/key_usage.rs new file mode 100644 index 00000000..3ac0e781 --- /dev/null +++ b/crates/server/src/handlers/key_usage.rs @@ -0,0 +1,456 @@ +// ============================================================================ +// GET /v1/usage (gateway port) +// +// What the calling API key has used today and this month, every limit +// that binds its AI-gateway requests with what has been used of it, and +// when the key stops working. +// +// The limits are the ones enforcement builds for a request with this key +// (`limits_for_ai_gateway`): the key's own, on its lineage's counters, and +// its owner's effective limits — roles, including those a team grants, +// merged most-restrictive, then the user's overrides — on the owner's +// counters, which count every key the owner has. They are read from the +// same Redis counters, the same way, without charging them. MCP-surface +// limits don't apply to model requests and are left out. +// +// Mounted behind `require_api_key_to_read`: the same key checks as a model +// request, but calling it is not a use of the key. +// ============================================================================ + +use std::cmp::Ordering; + +use axum::Json; +use axum::extract::State; +use axum::http::{StatusCode, header}; +use axum::response::{IntoResponse, Response}; +use chrono::{DateTime, Datelike, NaiveTime, Utc}; +use rust_decimal::prelude::ToPrimitive; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use think_watch_common::limits::{ + BudgetCap, BudgetPeriod, BudgetSubject, RateLimitRule, RateLimitSubject, RequestLimits, budget, + key_usage, secs_to_window, sliding, +}; +use think_watch_gateway::proxy::{GatewayRequestIdentity, limits_for_ai_gateway}; + +use crate::app::AppState; + +#[derive(Debug, Serialize, PartialEq)] +pub struct KeyUsageResponse { + pub usage: Usage, + pub limits: Vec, + /// When the key stops authenticating; `null` when it does not. + pub expires_at: Option>, +} + +/// What this key (its lineage, across rotations) has done, UTC day and +/// month. Tokens are weighted tokens, as limits count them. +#[derive(Debug, Serialize, PartialEq)] +pub struct Usage { + pub requests_today: i64, + pub tokens_today: i64, + pub requests_month: i64, + pub tokens_month: i64, + /// The key's cost this month from the request log, as the cost + /// reports sum it; `null` without ClickHouse or when it can't be read. + pub cost_usd_month: Option, +} + +/// One limit that binds the key's requests. +#[derive(Debug, Serialize, PartialEq)] +pub struct LimitUsage { + /// `key`: set on the key, counted for it alone. `user`: its owner's, + /// counted for everything the owner does. + pub scope: &'static str, + /// `requests` or `tokens` (weighted). + pub kind: &'static str, + /// A rate limit's sliding window (`1m` … `1w`) or a budget's calendar + /// period (`daily`, `weekly`, `monthly`). + pub window: String, + /// A sliding window's length; `null` for a calendar period. + pub window_secs: Option, + pub limit: i64, + /// What enforcement compares with `limit`. + pub used: i64, + /// A calendar period's end (UTC); `null` for a sliding window. + pub resets_at: Option>, +} + +/// GET /v1/usage +pub async fn get_key_usage( + State(state): State, + axum::Extension(identity): axum::Extension, +) -> Response { + let limits = limits_for_ai_gateway(&identity); + let now = Utc::now(); + + let counts = match limits.key_lineage { + Some(lineage) => match key_usage::read(&state.redis, lineage, now).await { + Ok(c) => c, + Err(e) => return unavailable(&e), + }, + None => key_usage::KeyUsageCounts::default(), + }; + let mut rule_used = Vec::with_capacity(limits.rules.len()); + for rule in &limits.rules { + let resolved = sliding::ResolvedRule::new(rule, limits.owner); + match sliding::read_count(&state.redis, &resolved, now.timestamp()).await { + Ok(n) => rule_used.push(n), + Err(e) => return unavailable(&e), + } + } + let cap_used: Vec = match budget::read_spend(&state.redis, &limits.caps, now).await { + Ok(statuses) => statuses.into_iter().map(|s| s.current).collect(), + Err(e) => return unavailable(&e), + }; + let cost = match limits.key_lineage { + Some(lineage) => cost_this_month(&state, lineage, now).await, + None => None, + }; + + Json(response( + &limits, + &rule_used, + &cap_used, + counts, + cost, + identity.key_expires_at, + now, + )) + .into_response() +} + +/// The answer for a request's limits and what each has used +/// (`rule_used` pairs with `limits.rules`, `cap_used` with +/// `limits.caps`), the key's counts and cost. Limits are ordered by the +/// share of them left, least first. +fn response( + limits: &RequestLimits, + rule_used: &[i64], + cap_used: &[i64], + counts: key_usage::KeyUsageCounts, + cost_usd_month: Option, + expires_at: Option>, + now: DateTime, +) -> KeyUsageResponse { + let mut entries: Vec = limits + .rules + .iter() + .zip(rule_used) + .map(|(rule, &used)| rule_entry(rule, used)) + .chain( + limits + .caps + .iter() + .zip(cap_used) + .map(|(cap, &used)| cap_entry(cap, used, now)), + ) + .collect(); + entries.sort_by(by_share_left); + KeyUsageResponse { + usage: Usage { + requests_today: counts.requests_today, + tokens_today: counts.tokens_today, + requests_month: counts.requests_month, + tokens_month: counts.tokens_month, + cost_usd_month, + }, + limits: entries, + expires_at, + } +} + +fn rule_entry(rule: &RateLimitRule, used: i64) -> LimitUsage { + LimitUsage { + scope: match rule.subject_kind { + RateLimitSubject::ApiKeyLineage => "key", + RateLimitSubject::User => "user", + }, + kind: rule.metric.as_str(), + window: secs_to_window(rule.window_secs) + .map(str::to_string) + .unwrap_or_else(|| format!("{}s", rule.window_secs)), + window_secs: Some(i64::from(rule.window_secs)), + limit: rule.max_count, + used, + resets_at: None, + } +} + +fn cap_entry(cap: &BudgetCap, used: i64, now: DateTime) -> LimitUsage { + LimitUsage { + scope: match cap.subject_kind { + BudgetSubject::ApiKeyLineage => "key", + BudgetSubject::User => "user", + }, + kind: "tokens", + window: cap.period.as_str().to_string(), + window_secs: None, + limit: cap.limit_tokens, + used, + resets_at: Some(budget::period_end(cap.period, now)), + } +} + +/// Least of its limit left first; on a tie, the key's before its +/// owner's, sliding windows shortest first before calendar periods, and +/// requests before tokens. +fn by_share_left(a: &LimitUsage, b: &LimitUsage) -> Ordering { + // (limit - used) / limit, compared without dividing; limits are > 0. + let left = |e: &LimitUsage| i128::from(e.limit) - i128::from(e.used); + let (la, lb) = (i128::from(a.limit.max(1)), i128::from(b.limit.max(1))); + (left(a) * lb) + .cmp(&(left(b) * la)) + .then_with(|| (a.scope != "key").cmp(&(b.scope != "key"))) + .then_with(|| window_rank(a).cmp(&window_rank(b))) + .then_with(|| (a.kind != "requests").cmp(&(b.kind != "requests"))) +} + +fn window_rank(e: &LimitUsage) -> i64 { + e.window_secs.unwrap_or_else(|| { + let period = BudgetPeriod::parse(&e.window); + i64::MAX + - match period { + Some(BudgetPeriod::Daily) => 3, + Some(BudgetPeriod::Weekly) => 2, + Some(BudgetPeriod::Monthly) | None => 1, + } + }) +} + +/// The key's cost since the 1st of this month, UTC, summed from the +/// request log as the cost reports sum it. `None` without ClickHouse; +/// a failed read is logged and also `None` — the cost is informational, +/// the counters above are what limits use. +async fn cost_this_month(state: &AppState, lineage: Uuid, now: DateTime) -> Option { + #[derive(clickhouse::Row, Deserialize)] + struct Cost { + cost: i128, + } + let ch = state.clickhouse.as_ref()?; + let month_start = now + .date_naive() + .with_day(1) + .expect("every month has a 1st") + .and_time(NaiveTime::MIN) + .and_utc(); + let row = ch + .query( + "SELECT sum(ifNull(cost_usd, 0)) AS cost FROM gateway_logs \ + PREWHERE created_at >= parseDateTimeBestEffort(?) \ + AND api_key_lineage_id = ?", + ) + .bind(month_start.format("%Y-%m-%d %H:%M:%S").to_string()) + .bind(lineage.to_string()) + .fetch_one::() + .await; + match row { + Ok(r) => think_watch_common::cost_decimal::decode_i128(r.cost).to_f64(), + Err(e) => { + tracing::warn!(error = %e, "key usage: monthly cost read failed"); + None + } + } +} + +/// The counters could not be read. Answering zero would say the key has +/// room it may not have. +fn unavailable(e: &fred::error::Error) -> Response { + tracing::warn!("key usage read failed: {e}"); + let body = tw_dialect::convert::error_body( + tw_dialect::ir::Dialect::Chat, + StatusCode::SERVICE_UNAVAILABLE.as_u16(), + "Usage counters are unavailable.", + ); + ( + StatusCode::SERVICE_UNAVAILABLE, + [(header::CONTENT_TYPE, "application/json")], + body, + ) + .into_response() +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::TimeZone; + use think_watch_common::limits::{ + RateMetric, Surface, SurfaceBlock, SurfaceBudget, SurfaceConstraints, SurfaceRule, + }; + + fn at(y: i32, m: u32, d: u32, h: u32) -> DateTime { + Utc.with_ymd_and_hms(y, m, d, h, 0, 0).unwrap() + } + + fn block(rules: &[(RateMetric, i32, i64)], budgets: &[(BudgetPeriod, i64)]) -> SurfaceBlock { + SurfaceBlock { + rules: rules + .iter() + .map(|&(metric, window_secs, max_count)| SurfaceRule { + metric, + window_secs, + max_count, + enabled: true, + }) + .collect(), + budgets: budgets + .iter() + .map(|&(period, limit_tokens)| SurfaceBudget { + period, + limit_tokens, + enabled: true, + }) + .collect(), + } + } + + fn counts() -> key_usage::KeyUsageCounts { + key_usage::KeyUsageCounts { + requests_today: 3, + tokens_today: 300, + requests_month: 40, + tokens_month: 4_000, + } + } + + #[test] + fn a_key_bound_by_nothing_still_reports_its_usage() { + let limits = RequestLimits::for_request( + Surface::AiGateway, + Uuid::new_v4(), + &SurfaceConstraints::default(), + Some((Uuid::new_v4(), &SurfaceConstraints::default())), + ); + let body = response(&limits, &[], &[], counts(), None, None, at(2026, 10, 14, 9)); + assert_eq!( + serde_json::to_value(&body).unwrap(), + serde_json::json!({ + "usage": {"requests_today": 3, "tokens_today": 300, + "requests_month": 40, "tokens_month": 4000, + "cost_usd_month": null}, + "limits": [], + "expires_at": null, + }) + ); + } + + #[test] + fn the_keys_and_its_owners_limits_map_to_the_contract_shape() { + // Wednesday 2026-10-14, 09:00 UTC. The owner has a requests/min + // rule and a monthly budget (MCP rules don't bind model requests); + // the key a 5h token window and a daily budget. + let user = SurfaceConstraints { + ai_gateway: Some(block( + &[(RateMetric::Requests, 60, 10)], + &[(BudgetPeriod::Monthly, 1_000_000)], + )), + mcp_gateway: Some(block(&[(RateMetric::Requests, 60, 1)], &[])), + }; + let key = SurfaceConstraints { + ai_gateway: Some(block( + &[(RateMetric::Tokens, 18_000, 50_000)], + &[(BudgetPeriod::Daily, 100_000)], + )), + mcp_gateway: None, + }; + let limits = RequestLimits::for_request( + Surface::AiGateway, + Uuid::new_v4(), + &user, + Some((Uuid::new_v4(), &key)), + ); + // Input order: user rule, key rule; user cap, key cap. + let body = response( + &limits, + &[5, 45_000], + &[250_000, 100_500], + counts(), + Some(1.25), + Some(at(2026, 12, 31, 0)), + at(2026, 10, 14, 9), + ); + assert_eq!( + serde_json::to_value(&body).unwrap(), + serde_json::json!({ + "usage": {"requests_today": 3, "tokens_today": 300, + "requests_month": 40, "tokens_month": 4000, + "cost_usd_month": 1.25}, + "limits": [ + // Overshot: -0.5 % left. + {"scope": "key", "kind": "tokens", "window": "daily", "window_secs": null, + "limit": 100000, "used": 100500, "resets_at": "2026-10-15T00:00:00Z"}, + // 10 % left. + {"scope": "key", "kind": "tokens", "window": "5h", "window_secs": 18000, + "limit": 50000, "used": 45000, "resets_at": null}, + // 50 % left. + {"scope": "user", "kind": "requests", "window": "1m", "window_secs": 60, + "limit": 10, "used": 5, "resets_at": null}, + // 75 % left. + {"scope": "user", "kind": "tokens", "window": "monthly", "window_secs": null, + "limit": 1000000, "used": 250000, "resets_at": "2026-11-01T00:00:00Z"}, + ], + "expires_at": "2026-12-31T00:00:00Z", + }) + ); + } + + #[test] + fn ties_put_the_key_first_then_the_shorter_window_then_requests() { + let e = |scope, kind, window: &str, window_secs| LimitUsage { + scope, + kind, + window: window.to_string(), + window_secs, + limit: 10, + used: 0, + resets_at: None, + }; + let mut v = [ + e("user", "requests", "1m", Some(60)), + e("key", "tokens", "monthly", None), + e("key", "tokens", "weekly", None), + e("key", "tokens", "1w", Some(604_800)), + e("key", "tokens", "1m", Some(60)), + e("key", "requests", "1m", Some(60)), + ]; + v.sort_by(by_share_left); + let order: Vec<(&str, &str, &str)> = v + .iter() + .map(|e| (e.scope, e.kind, e.window.as_str())) + .collect(); + assert_eq!( + order, + [ + ("key", "requests", "1m"), + ("key", "tokens", "1m"), + ("key", "tokens", "1w"), + ("key", "tokens", "weekly"), + ("key", "tokens", "monthly"), + ("user", "requests", "1m"), + ] + ); + } + + #[test] + fn every_allowed_window_has_its_label() { + let r = |secs| RateLimitRule { + id: Uuid::nil(), + subject_kind: RateLimitSubject::User, + subject_id: Uuid::nil(), + surface: Surface::AiGateway, + metric: RateMetric::Requests, + window_secs: secs, + max_count: 1, + enabled: true, + expires_at: None, + reason: None, + created_by: None, + }; + let labels: Vec = [60, 300, 3_600, 18_000, 86_400, 604_800] + .into_iter() + .map(|s| rule_entry(&r(s), 0).window) + .collect(); + assert_eq!(labels, ["1m", "5m", "1h", "5h", "1d", "1w"]); + } +} diff --git a/crates/server/src/handlers/mod.rs b/crates/server/src/handlers/mod.rs index b9903cda..d1656294 100644 --- a/crates/server/src/handlers/mod.rs +++ b/crates/server/src/handlers/mod.rs @@ -12,6 +12,7 @@ pub mod cost_forecast; pub mod dashboard; pub mod gateway_logs; pub mod health; +pub mod key_usage; pub mod limits; pub mod limits_bulk; pub mod log_forwarders; diff --git a/crates/server/src/middleware/api_key_auth.rs b/crates/server/src/middleware/api_key_auth.rs index 1f13aad4..51042107 100644 --- a/crates/server/src/middleware/api_key_auth.rs +++ b/crates/server/src/middleware/api_key_auth.rs @@ -154,6 +154,28 @@ fn presented_key<'a>( type AuthFuture = std::pin::Pin> + Send>>; +/// When the key stops authenticating: its expiry, or the end of its +/// rotation grace period (set on a key that a rotation replaced), +/// whichever comes first. +fn stops_authenticating_at( + expires_at: Option>, + grace_period_ends_at: Option>, +) -> Option> { + match (expires_at, grace_period_ends_at) { + (Some(a), Some(b)) => Some(a.min(b)), + (a, b) => a.or(b), + } +} + +/// What the request does with the key. +#[derive(Clone, Copy, PartialEq, Eq)] +enum KeyUse { + /// Calls the gateway: a model, a model listing, an MCP tool. + Call, + /// Reads about the key itself (`GET /v1/usage`). + Read, +} + /// Build a middleware that authenticates requests via `tw-` prefixed /// API keys and additionally requires the key's `surfaces` array to /// contain `surface`. @@ -166,6 +188,25 @@ type AuthFuture = /// JWT for interactive admin users. pub fn require_api_key( surface: &'static str, +) -> impl Fn(State, Request, Next) -> AuthFuture + Clone { + authenticate(surface, KeyUse::Call) +} + +/// [`require_api_key`] for reads about the calling key itself +/// (`GET /v1/usage`): the same lookup, checks and refusals, but reading +/// about a key is not a use of it. `last_used_at`, which the inactivity +/// cutoff reads, stays as it was — a client polling its usage does not +/// keep an otherwise idle key alive — and a client that leaves early +/// leaves no `gateway_logs` row. +pub fn require_api_key_to_read( + surface: &'static str, +) -> impl Fn(State, Request, Next) -> AuthFuture + Clone { + authenticate(surface, KeyUse::Read) +} + +fn authenticate( + surface: &'static str, + key_use: KeyUse, ) -> impl Fn(State, Request, Next) -> AuthFuture + Clone { move |State(state): State, mut request: Request, next: Next| { Box::pin(async move { @@ -257,7 +298,7 @@ pub fn require_api_key( // records its own). Every return below produces a response // or an auth refusal, so the guard is disarmed after all of // them; only a dropped future leaves it armed. - let cancel = (surface == "ai_gateway").then(|| { + let cancel = (surface == "ai_gateway" && key_use == KeyUse::Call).then(|| { think_watch_gateway::proxy::EarlyCancel::arm( state.audit.clone(), GatewayRequestIdentity { @@ -271,18 +312,20 @@ pub fn require_api_key( }); let result: Result = async { // Update last_used_at (best-effort, don't block on failure) - let db = state.db.clone(); - let key_id = row.id; - tokio::spawn(async move { - if let Err(e) = - sqlx::query("UPDATE api_keys SET last_used_at = now() WHERE id = $1") - .bind(key_id) - .execute(&db) - .await - { - tracing::warn!("Failed to update api_key last_used_at: {e}"); - } - }); + if key_use == KeyUse::Call { + let db = state.db.clone(); + let key_id = row.id; + tokio::spawn(async move { + if let Err(e) = + sqlx::query("UPDATE api_keys SET last_used_at = now() WHERE id = $1") + .bind(key_id) + .execute(&db) + .await + { + tracing::warn!("Failed to update api_key last_used_at: {e}"); + } + }); + } // Compute the user's role-derived constraints and intersect // with the API-key allow-list. The role union is loaded once @@ -385,6 +428,10 @@ pub fn require_api_key( surface_constraints: surface_constraints.clone(), key_constraints: key_constraints.clone(), ip_address: client_ip.clone(), + key_expires_at: stops_authenticating_at( + row.expires_at, + row.grace_period_ends_at, + ), }; // The MCP transport handlers expect their own typed @@ -492,6 +539,23 @@ mod tests { ); } + #[test] + fn a_key_stops_at_its_expiry_or_its_grace_end_whichever_is_first() { + use chrono::TimeZone; + let d = |day| { + Some( + chrono::Utc + .with_ymd_and_hms(2026, 10, day, 0, 0, 0) + .unwrap(), + ) + }; + assert_eq!(stops_authenticating_at(None, None), None); + assert_eq!(stops_authenticating_at(d(20), None), d(20)); + assert_eq!(stops_authenticating_at(None, d(12)), d(12)); + assert_eq!(stops_authenticating_at(d(20), d(12)), d(12)); + assert_eq!(stops_authenticating_at(d(11), d(12)), d(11)); + } + #[test] fn a_key_is_read_where_each_sdk_puts_it() { assert_eq!( diff --git a/crates/test-support/tests/analytics_clickhouse.rs b/crates/test-support/tests/analytics_clickhouse.rs index 48530fb8..28b1847b 100644 --- a/crates/test-support/tests/analytics_clickhouse.rs +++ b/crates/test-support/tests/analytics_clickhouse.rs @@ -872,3 +872,70 @@ async fn cached_input_is_billed_at_the_cache_prices() { // 0.000002 × (100 + 5 000 + 2 500) + 0.0004 assert_eq!(second, Decimal::from_str("0.0156").unwrap()); } + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn v1_usage_reports_the_keys_cost_this_month_from_the_request_log() { + // Two keys of one owner: each key's month counts its own requests. + let app = TestApp::spawn_with_clickhouse().await; + let (key_a, user_id) = seed_runtime(&app).await; + let key_b = fixtures::create_api_key( + &app.db, + user_id, + &unique_name("ck-key-b"), + &["ai_gateway"], + None, + None, + ) + .await + .unwrap(); + let ch = app.state.clickhouse.as_ref().expect("clickhouse wired up"); + + let gw_a = app.gateway_client(); + gw_a.set_bearer(&key_a); + let gw_b = app.gateway_client(); + gw_b.set_bearer(&key_b.plaintext); + for (gw, n) in [(&gw_a, 1), (&gw_b, 2), (&gw_b, 3)] { + gw.post( + "/v1/chat/completions", + json!({"model": "gpt-test", "messages": [{"role": "user", "content": format!("x{n}")}]}), + ) + .await + .unwrap() + .assert_ok(); + } + let (one_call, _, _) = wait_for_gateway_log(ch, user_id).await; + let one_call: f64 = one_call.to_string().parse().unwrap(); + // Wait until all three rows have landed. + for _ in 0..200 { + let n: u64 = ch + .query("SELECT count() FROM gateway_logs WHERE user_id = ? AND cost_usd IS NOT NULL") + .bind(user_id.to_string()) + .fetch_one() + .await + .unwrap(); + if n == 3 { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + + let cost = |gw: &TestClient| { + let gw = gw.clone(); + async move { + let r = gw.get("/v1/usage").await.unwrap(); + r.assert_ok(); + let v: Value = r.json().unwrap(); + v["usage"]["cost_usd_month"] + .as_f64() + .unwrap_or_else(|| panic!("cost_usd_month is a number: {v:#}")) + } + }; + let (a, b) = (cost(&gw_a).await, cost(&gw_b).await); + assert!(one_call > 0.0); + assert!((a - one_call).abs() < 1e-9, "a={a} one call={one_call}"); + assert!( + (b - 2.0 * one_call).abs() < 1e-9, + "b={b} one call={one_call}" + ); +} diff --git a/crates/test-support/tests/limits.rs b/crates/test-support/tests/limits.rs index 0161533c..5480cdd1 100644 --- a/crates/test-support/tests/limits.rs +++ b/crates/test-support/tests/limits.rs @@ -938,3 +938,454 @@ async fn a_keys_mcp_limits_count_on_the_key_not_on_its_owner() { ); assert_eq!(usage(&con, "api_key", keys[0].row.id).await.0, vec![1]); } + +// ---------------------------------------------------------------------------- +// GET /v1/usage: what the calling key has used, and every limit binding it +// ---------------------------------------------------------------------------- + +/// `GET /v1/usage` with the client's key, as JSON. +async fn key_usage(gw: &TestClient) -> Json { + let r = gw.get("/v1/usage").await.unwrap(); + r.assert_ok(); + r.json().unwrap() +} + +/// The `limits` entries as `(scope, kind, window, limit, used)`, in the +/// order the endpoint gives them. +fn limit_rows(v: &Json) -> Vec<(String, String, String, i64, i64)> { + v["limits"] + .as_array() + .unwrap_or_else(|| panic!("no limits array: {v:#}")) + .iter() + .map(|l| { + ( + l["scope"].as_str().unwrap().to_string(), + l["kind"].as_str().unwrap().to_string(), + l["window"].as_str().unwrap().to_string(), + l["limit"].as_i64().unwrap(), + l["used"].as_i64().unwrap(), + ) + }) + .collect() +} + +fn row( + scope: &str, + kind: &str, + window: &str, + limit: i64, + used: i64, +) -> (String, String, String, i64, i64) { + (scope.into(), kind.into(), window.into(), limit, used) +} + +/// A role whose statement carries `constraints`, granted to `user_id` +/// through a new team (`via_team`) or directly at global scope. +async fn role_with_constraints(app: &TestApp, user_id: Uuid, constraints: Json, via_team: bool) { + let doc = json!({ + "Version": "2024-01-01", + "Statement": [{ + "Effect": "Allow", + "Action": ["ai_gateway:use"], + "Resource": ["*"], + "Constraints": constraints, + }], + }); + let role: Uuid = sqlx::query_scalar( + "INSERT INTO rbac_roles (name, is_system, policy_document) VALUES ($1, FALSE, $2) RETURNING id", + ) + .bind(unique_name("usage-role")) + .bind(doc) + .fetch_one(&app.db) + .await + .unwrap(); + if via_team { + let team: Uuid = sqlx::query_scalar( + "INSERT INTO teams (name, description) VALUES ($1, '') RETURNING id", + ) + .bind(unique_name("usage-team")) + .fetch_one(&app.db) + .await + .unwrap(); + sqlx::query("INSERT INTO team_role_assignments (team_id, role_id) VALUES ($1, $2)") + .bind(team) + .bind(role) + .execute(&app.db) + .await + .unwrap(); + sqlx::query("INSERT INTO team_members (team_id, user_id) VALUES ($1, $2)") + .bind(team) + .bind(user_id) + .execute(&app.db) + .await + .unwrap(); + } else { + sqlx::query( + "INSERT INTO rbac_role_assignments (user_id, role_id, scope_kind, assigned_by) \ + VALUES ($1, $2, 'global', $1)", + ) + .bind(user_id) + .bind(role) + .execute(&app.db) + .await + .unwrap(); + } +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn v1_usage_reports_a_keys_own_limits_on_its_own_counters() { + let app = TestApp::spawn().await; + let (_, user_id) = seed_runtime(&app).await; + let key = fixtures::create_api_key( + &app.db, + user_id, + &unique_name("usage-key"), + &["ai_gateway"], + None, + Some(chrono::Utc::now() + chrono::Duration::days(30)), + ) + .await + .unwrap(); + let lineage = key.row.lineage_id; + fixtures::create_budget_cap(&app.db, "api_key_lineage", lineage, "daily", 1_000_000) + .await + .unwrap(); + for (metric, window, max) in [("tokens", 18_000, 1_000_000), ("requests", 60, 100)] { + fixtures::create_rate_limit_rule( + &app.db, + "api_key_lineage", + lineage, + "ai_gateway", + metric, + window, + max, + ) + .await + .unwrap(); + } + // An MCP-surface limit binds no model request. + fixtures::create_rate_limit_rule( + &app.db, + "api_key_lineage", + lineage, + "mcp_gateway", + "requests", + 60, + 7, + ) + .await + .unwrap(); + let con = admin_session(&app).await; + + let gw = app.gateway_client(); + gw.set_bearer(&key.plaintext); + for _ in 0..2 { + gw.post("/v1/chat/completions", chat()) + .await + .unwrap() + .assert_ok(); + } + + let v = key_usage(&gw).await; + // The console reads the same counters: the AI rules ordered requests + // then tokens (the MCP rule last), then the cap. + let (rules, caps) = usage(&con, "api_key", key.row.id).await; + assert_eq!(rules[0], 2); + let tokens = rules[1]; + assert!(tokens > 0 && caps[0] == tokens, "{rules:?} {caps:?}"); + // Least left first: requests 98 %, then the token window and the + // budget, equal, sliding before calendar. + assert_eq!( + limit_rows(&v), + vec![ + row("key", "requests", "1m", 100, 2), + row("key", "tokens", "5h", 1_000_000, tokens), + row("key", "tokens", "daily", 1_000_000, tokens), + ] + ); + let midnight = think_watch_common::limits::budget::period_end( + think_watch_common::limits::BudgetPeriod::Daily, + chrono::Utc::now(), + ); + assert_eq!(v["limits"][0]["window_secs"], 60); + assert_eq!(v["limits"][0]["resets_at"], Json::Null); + assert_eq!(v["limits"][2]["window_secs"], Json::Null); + assert_eq!(v["limits"][2]["resets_at"], json!(midnight)); + assert_eq!( + v["usage"], + json!({"requests_today": 2, "tokens_today": tokens, + "requests_month": 2, "tokens_month": tokens, + "cost_usd_month": null}) + ); + let reported: chrono::DateTime = + serde_json::from_value(v["expires_at"].clone()).unwrap(); + assert_eq!(Some(reported), key.row.expires_at); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn v1_usage_reports_the_owners_effective_limits_on_the_owners_counters() { + // The owner's limits from a role, from a role its team grants, and + // from an override, counting what every key of the owner does. The + // key's own usage counts only its own requests. + let app = TestApp::spawn().await; + let (key_a, user_id) = seed_runtime(&app).await; + let key_b = second_key(&app, user_id).await; + role_with_constraints( + &app, + user_id, + json!({"Budgets": [{"Period": "monthly", "MaxTokens": 9_000_000}]}), + false, + ) + .await; + role_with_constraints( + &app, + user_id, + json!({"Budgets": [{"Period": "daily", "MaxTokens": 5_000_000}], + "RateLimits": [{"Metric": "tokens", "Window": "1h", "MaxCount": 4_000_000}, + {"Metric": "requests", "Window": "1m", "MaxCount": 20}]}), + true, + ) + .await; + // The override replaces the team role's requests/min. + fixtures::create_rate_limit_rule(&app.db, "user", user_id, "ai_gateway", "requests", 60, 50) + .await + .unwrap(); + + let gw_a = app.gateway_client(); + gw_a.set_bearer(&key_a); + let gw_b = app.gateway_client(); + gw_b.set_bearer(&key_b.plaintext); + // Distinct prompts: a cached answer charges no tokens. + let prompt = |n: u32| json!({"model": "gpt-test", "messages": [{"role": "user", "content": format!("x{n}")}]}); + gw_a.post("/v1/chat/completions", prompt(1)) + .await + .unwrap() + .assert_ok(); + for n in 2..4 { + gw_b.post("/v1/chat/completions", prompt(n)) + .await + .unwrap() + .assert_ok(); + } + + let a = key_usage(&gw_a).await; + let b = key_usage(&gw_b).await; + let (a_tokens, b_tokens) = ( + a["usage"]["tokens_today"].as_i64().unwrap(), + b["usage"]["tokens_today"].as_i64().unwrap(), + ); + assert!(a_tokens > 0 && b_tokens > a_tokens, "{a:#} {b:#}"); + assert_eq!(a["usage"]["requests_today"], 1); + assert_eq!(b["usage"]["requests_today"], 2); + let all = a_tokens + b_tokens; + let expected = vec![ + row("user", "requests", "1m", 50, 3), + row("user", "tokens", "1h", 4_000_000, all), + row("user", "tokens", "daily", 5_000_000, all), + row("user", "tokens", "monthly", 9_000_000, all), + ]; + assert_eq!(limit_rows(&a), expected); + assert_eq!( + limit_rows(&b), + expected, + "both keys are bound by the same limits" + ); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn v1_usage_lists_the_keys_and_the_owners_limits_least_left_first() { + let app = TestApp::spawn().await; + let (_, user_id) = seed_runtime(&app).await; + let key = second_key(&app, user_id).await; + fixtures::create_rate_limit_rule(&app.db, "user", user_id, "ai_gateway", "requests", 60, 50) + .await + .unwrap(); + fixtures::create_rate_limit_rule( + &app.db, + "api_key_lineage", + key.row.lineage_id, + "ai_gateway", + "requests", + 300, + 5, + ) + .await + .unwrap(); + + let gw = app.gateway_client(); + gw.set_bearer(&key.plaintext); + gw.post("/v1/chat/completions", chat()) + .await + .unwrap() + .assert_ok(); + assert_eq!( + limit_rows(&key_usage(&gw).await), + vec![ + row("key", "requests", "5m", 5, 1), + row("user", "requests", "1m", 50, 1), + ] + ); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn v1_usage_of_a_key_bound_by_nothing_still_reports_its_usage() { + let app = TestApp::spawn().await; + let (api_key, _) = seed_runtime(&app).await; + + // Anthropic's SDKs send the key in x-api-key. + let gw = app.gateway_client(); + gw.set_header("x-api-key", &api_key); + assert_eq!( + key_usage(&gw).await, + json!({"usage": {"requests_today": 0, "tokens_today": 0, + "requests_month": 0, "tokens_month": 0, + "cost_usd_month": null}, + "limits": [], "expires_at": null}) + ); + gw.post("/v1/chat/completions", chat()) + .await + .unwrap() + .assert_ok(); + let v = key_usage(&gw).await; + assert_eq!(v["limits"], json!([])); + assert_eq!(v["usage"]["requests_today"], 1); + assert_eq!(v["usage"]["requests_month"], 1); + let tokens = v["usage"]["tokens_today"].as_i64().unwrap(); + assert!(tokens > 0, "{v:#}"); + assert_eq!(v["usage"]["tokens_month"], tokens); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn v1_usage_refuses_a_key_a_model_request_refuses() { + let app = TestApp::spawn().await; + let (_, user_id) = seed_runtime(&app).await; + let expired = fixtures::create_api_key( + &app.db, + user_id, + &unique_name("usage-expired"), + &["ai_gateway"], + None, + Some(chrono::Utc::now() - chrono::Duration::minutes(1)), + ) + .await + .unwrap(); + let mcp_only = fixtures::create_api_key( + &app.db, + user_id, + &unique_name("usage-mcp"), + &["mcp_gateway"], + None, + None, + ) + .await + .unwrap(); + + let gw = app.gateway_client(); + let unknown = format!("tw-{}", Uuid::new_v4().simple()); + for (key, status) in [ + (None, 401), + (Some(unknown.as_str()), 401), + (Some("sk-not-a-gateway-key"), 401), + (Some(expired.plaintext.as_str()), 401), + (Some(mcp_only.plaintext.as_str()), 403), + ] { + match key { + Some(k) => gw.set_bearer(k), + None => gw.clear_bearer(), + } + let usage = gw.get("/v1/usage").await.unwrap(); + let model = gw.post("/v1/chat/completions", chat()).await.unwrap(); + assert_eq!(usage.status.as_u16(), status, "{key:?}: {}", usage.text()); + assert_eq!(model.status.as_u16(), status, "{key:?}: {}", model.text()); + assert_eq!(usage.body, model.body, "{key:?}"); + } +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn calling_v1_usage_moves_no_counter() { + let app = TestApp::spawn().await; + let (_, user_id) = seed_runtime(&app).await; + let key = second_key(&app, user_id).await; + let lineage = key.row.lineage_id; + for (metric, max) in [("requests", 2), ("tokens", 1_000_000)] { + fixtures::create_rate_limit_rule( + &app.db, + "api_key_lineage", + lineage, + "ai_gateway", + metric, + 60, + max, + ) + .await + .unwrap(); + } + fixtures::create_budget_cap(&app.db, "api_key_lineage", lineage, "daily", 1_000_000) + .await + .unwrap(); + fixtures::create_rate_limit_rule(&app.db, "user", user_id, "ai_gateway", "requests", 60, 2) + .await + .unwrap(); + let con = admin_session(&app).await; + + let gw = app.gateway_client(); + gw.set_bearer(&key.plaintext); + gw.post("/v1/chat/completions", chat()) + .await + .unwrap() + .assert_ok(); + + // The model request's `last_used_at` write is spawned; wait for it, + // then pin the column to a known value. + let last_used = || async { + sqlx::query_scalar::<_, Option>>( + "SELECT last_used_at FROM api_keys WHERE id = $1", + ) + .bind(key.row.id) + .fetch_one(&app.db) + .await + .unwrap() + }; + for _ in 0..50 { + if last_used().await.is_some() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + let pinned: chrono::DateTime = sqlx::query_scalar( + "UPDATE api_keys SET last_used_at = date_trunc('second', now()) - interval '1 hour' \ + WHERE id = $1 RETURNING last_used_at", + ) + .bind(key.row.id) + .fetch_one(&app.db) + .await + .unwrap(); + + let key_before = usage(&con, "api_key", key.row.id).await; + let user_before = usage(&con, "user", user_id).await; + let reported = key_usage(&gw).await; + assert_eq!(reported["usage"]["requests_today"], 1); + for _ in 0..5 { + assert_eq!(key_usage(&gw).await, reported); + } + tokio::time::sleep(std::time::Duration::from_millis(300)).await; + assert_eq!(usage(&con, "api_key", key.row.id).await, key_before); + assert_eq!(usage(&con, "user", user_id).await, user_before); + assert_eq!(last_used().await, Some(pinned)); + + // The second model request the limits allow still goes through. + gw.post("/v1/chat/completions", chat()) + .await + .unwrap() + .assert_ok(); + let r = gw.post("/v1/chat/completions", chat()).await.unwrap(); + assert_eq!(r.status.as_u16(), 429, "body={}", r.text()); + // The refused request counted nothing. + assert_eq!(key_usage(&gw).await["usage"]["requests_today"], 2); +} diff --git a/web/src/i18n/en.json b/web/src/i18n/en.json index 1897f426..2af13915 100644 --- a/web/src/i18n/en.json +++ b/web/src/i18n/en.json @@ -1942,6 +1942,7 @@ "responsesEndpoint": "OpenAI Responses (new): /v1/responses", "responsesWsEndpoint": "OpenAI Responses over WebSocket: /v1/responses (one response.create frame per turn)", "modelsEndpoint": "List models: /v1/models", + "usageEndpoint": "Usage and limits of the calling key: /v1/usage", "claudeDesktop": "Claude Desktop", "claudeDesktopDesc": "Configure Claude Desktop to use MCP tools through ThinkWatch.", "claudeCodeMcpDesc": "Register ThinkWatch as an MCP server directly from the Claude Code CLI.", diff --git a/web/src/i18n/zh.json b/web/src/i18n/zh.json index aac5beeb..dae1203f 100644 --- a/web/src/i18n/zh.json +++ b/web/src/i18n/zh.json @@ -1942,6 +1942,7 @@ "responsesEndpoint": "OpenAI Responses（新版）：/v1/responses", "responsesWsEndpoint": "OpenAI Responses（WebSocket）：/v1/responses（每轮发一个 response.create 帧）", "modelsEndpoint": "模型列表：/v1/models", + "usageEndpoint": "调用密钥的用量与限额：/v1/usage", "claudeDesktop": "Claude Desktop", "claudeDesktopDesc": "配置 Claude Desktop 通过 ThinkWatch 使用 MCP 工具。", "claudeCodeMcpDesc": "通过 Claude Code CLI 将 ThinkWatch 注册为 MCP 服务器。", diff --git a/web/src/routes/guide.tsx b/web/src/routes/guide.tsx index 7d53a1e9..0166bbab 100644 --- a/web/src/routes/guide.tsx +++ b/web/src/routes/guide.tsx @@ -785,6 +785,10 @@ export function GuidePage() { GET {t('guide.modelsEndpoint')}

+

+ GET + {t('guide.usageEndpoint')} +

From 8c9b07d78e97e24be32c07cb338e73cfc0b2cafc Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Sat, 10 Oct 2026 22:58:26 +0800 Subject: [PATCH 2/3] feat(gateway): GET /v1/usage answers for the key or for its owner The answer used to mix the key's own limits with its owner's, next to a usage count for the key alone, so a client could not tell whose room the numbers described. It now picks one subject and says which in a top-level `scope`: - "key" when the key has limits of its own on the AI gateway: `limits` lists only those, counted for the key, and `usage` is the key's own. - "user" when it has none: `limits` lists the owner's effective limits (roles, including team-granted ones, merged most-restrictive, then overrides), counted for everything the owner does, and `usage` is the owner's total over all of their keys, `cost_usd_month` included. With no limits on the owner either, `limits` is empty. The limits still come from `limits_for_ai_gateway`, filtered to the chosen subject, and `used` is read as before. A key whose only limit is on the MCP gateway answers for its owner. The owner's totals need a counter of their own: the user-subject counters that exist are limit counters, kept only for configured limits and only since they were configured, so they cannot give a day or month total. The per-key counters become per-key and per-user ones in one module (`limits::usage`, was `key_usage`): every counted request and its tokens land on the key's and the owner's day and month hashes, all four under the owner's `{user:}` hash tag, one script per count, at the same two points as before. An answer from the response cache now counts its tokens on these counters too, weighted as a call's are, so the usage agrees with the limits once those count cache hits as well. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 75 ++-- crates/common/src/limits/key_usage.rs | 168 --------- crates/common/src/limits/mod.rs | 6 +- crates/common/src/limits/usage.rs | 209 +++++++++++ crates/gateway/src/proxy/accounting.rs | 63 +++- crates/gateway/src/proxy/generate.rs | 10 + crates/gateway/src/proxy/pipeline.rs | 14 +- crates/server/src/handlers/key_usage.rs | 330 ++++++++++++------ .../tests/analytics_clickhouse.rs | 32 +- crates/test-support/tests/limits.rs | 206 +++++++---- web/src/i18n/en.json | 2 +- web/src/i18n/zh.json | 2 +- 12 files changed, 718 insertions(+), 399 deletions(-) delete mode 100644 crates/common/src/limits/key_usage.rs create mode 100644 crates/common/src/limits/usage.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index ab91ef8f..6366f224 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,57 +13,66 @@ target. ### Read before upgrading -- **Each API key gets day and month usage counters in Redis.** An - AI-gateway request made with a key now also counts on - `key_usage:{api_key_lineage:}:daily:` and - `…:monthly:`, one hash each that expires two periods after its last - write: one more Redis script per request once it passes the rate limits, - and one when its tokens are counted. The counters start empty, so the - `usage` that `GET /v1/usage` reports counts from the upgrade on. No - setting, database schema or Helm value changes. +- **Each API key and each user get day and month usage counters in Redis.** + An AI-gateway request made with a key now also counts on the key's + `usage:{user:}:api_key_lineage::daily:` and + `…:monthly:` and on its owner's + `usage:{user:}:user::daily:` and + `…:monthly:`, hashes that expire two periods after their last write. + That is one more Redis script per request once it passes the rate limits, + and one when its tokens are counted, after the call or when it is answered + from the response cache. The counters start empty, so the `usage` that + `GET /v1/usage` reports counts from the upgrade on. No setting, database + schema or Helm value changes. ### Added -- **`GET /v1/usage` tells a client what its API key has used and what room - it has left.** Called on the gateway port with a gateway API key, as a - model request is, it answers with what the key has done today and this - month, every limit that binds its model requests with what has been used - of it, and when the key stops working: +- **`GET /v1/usage` tells a client what room its API key has left.** Called + on the gateway port with a gateway API key, as a model request is, it + answers about one of two subjects, named in `scope`: + + - `"key"` when the key has limits of its own on the AI gateway: `limits` + lists only the key's limits, each with `used` counted for the key, and + `usage` is the key's own (across its rotations). + - `"user"` when it has none: `limits` lists its owner's effective limits + (roles, including those a team grants, merged most-restrictive, then + the user's overrides), each with `used` counted for everything the + owner does, and `usage` is the owner's total over all of their keys. + With no limits on the owner either, `limits` is empty. ```json { + "scope": "key", "usage": {"requests_today": 12, "tokens_today": 48210, "requests_month": 340, "tokens_month": 1290455, "cost_usd_month": 3.82}, "limits": [ {"scope": "key", "kind": "tokens", "window": "daily", "window_secs": null, "limit": 100000, "used": 48210, "resets_at": "2026-10-12T00:00:00Z"}, - {"scope": "user", "kind": "requests", "window": "1m", "window_secs": 60, + {"scope": "key", "kind": "requests", "window": "1m", "window_secs": 60, "limit": 60, "used": 4, "resets_at": null} ], "expires_at": "2026-12-31T00:00:00Z" } ``` - `usage` counts this key alone, across its rotations: the requests the - rate limits let through and the weighted tokens limits count, for the - UTC day and month, and its cost this month from the request log (`null` - without ClickHouse). `limits` lists the key's own limits (`scope: key`, - counted for the key) and its owner's effective limits (`scope: user`: - roles, including those a team grants, merged most-restrictive, then the - user's overrides, counted for everything the owner does through any key), - each with `used` read from the counter that refuses requests, the one - with the least left first. `window` is a rate limit's sliding window - (`1m`, `5m`, `1h`, `5h`, `1d`, `1w`, with its length in `window_secs`) or - a budget's calendar period (`daily`, `weekly`, `monthly`, with its end in - `resets_at`, UTC). Limits on the MCP gateway are not listed. `expires_at` - is the key's expiry or the end of its rotation grace period, whichever - comes first. A key a model request would refuse gets the same `401` or - `403`, and `503` means the counters could not be read. Calling it charges - no limit, writes no request log row and is not a use of the key: - `last_used_at` stays as it was, so polling does not keep an idle key from - its inactivity timeout. The console's Configuration Guide lists it with - the other gateway endpoints. + `usage` holds the requests the rate limits let through and the weighted + tokens limits count, answers from the response cache included, for the + UTC day and month, and the cost this month from the request log (`null` + without ClickHouse). Every limit's `used` is read from the counter that + refuses requests, and the one with the least left comes first. `window` + is a rate limit's sliding window (`1m`, `5m`, `1h`, `5h`, `1d`, `1w`, with + its length in `window_secs`) or a budget's calendar period (`daily`, + `weekly`, `monthly`, with its end in `resets_at`, UTC). Limits on the MCP + gateway are not listed. `expires_at` is the key's expiry or the end of + its rotation grace period, whichever comes first. A key a model request + would refuse gets the same `401` or `403`, and `503` means the counters + could not be read. Calling it charges no limit, writes no request log row + and is not a use of the key: `last_used_at` stays as it was, so polling + does not keep an idle key from its inactivity timeout. Whoever holds a key + without limits of its own sees its owner's totals and limits; a key handed + to someone else should carry limits of its own. The console's + Configuration Guide lists the endpoint with the other gateway endpoints. ## [3.4.0] — 2026-10-11 diff --git a/crates/common/src/limits/key_usage.rs b/crates/common/src/limits/key_usage.rs deleted file mode 100644 index 40c608c9..00000000 --- a/crates/common/src/limits/key_usage.rs +++ /dev/null @@ -1,168 +0,0 @@ -// ============================================================================ -// Per-key calendar usage -// -// What each API key (lineage) has been let through and used today and -// this month, UTC, whether or not it has limits of its own: the rate-limit -// and budget counters only exist for the limits that are configured, so -// they can't say what a key without limits has done. `GET /v1/usage` -// reads these. -// -// They count what the limits count, at the moments the limits engine -// counts it: a request once it has passed the rate limits (where -// `sliding::admit` charges `requests` rules), and its weighted tokens -// after the call (where `sliding::record` and -// `budget::add_weighted_tokens` add them). A request the limits refuse -// counts nothing. -// -// Storage (Redis): one hash per lineage per period, fields `requests` -// and `tokens`: -// -// key_usage:{api_key_lineage:}:daily:2026-10-11 -// key_usage:{api_key_lineage:}:monthly:2026-10 -// -// The braces are a Redis Cluster hash tag, so one script writes both. -// Each hash expires 2 × its period after its last write, as the budget -// counters do. -// ============================================================================ - -use chrono::{DateTime, Utc}; -use fred::clients::Client; -use fred::interfaces::{HashesInterface, LuaInterface}; -use uuid::Uuid; - -use super::budget::bucket_id; - -const DAY_TTL_SECS: i64 = 2 * 86_400; -const MONTH_TTL_SECS: i64 = 2 * 31 * 86_400; - -/// What one key has used today and this month (UTC). -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] -pub struct KeyUsageCounts { - pub requests_today: i64, - pub tokens_today: i64, - pub requests_month: i64, - pub tokens_month: i64, -} - -/// The two hashes of `lineage` for the day and the month containing `now`. -pub fn keys(lineage: Uuid, now: DateTime) -> [String; 2] { - ["daily", "monthly"].map(|period| { - format!( - "key_usage:{{api_key_lineage:{lineage}}}:{period}:{}", - bucket_id(period, now) - ) - }) -} - -const LUA_ADD: &str = r#" --- KEYS: day hash, month hash --- ARGV: field, amount, day ttl, month ttl -redis.call('HINCRBY', KEYS[1], ARGV[1], ARGV[2]) -redis.call('EXPIRE', KEYS[1], ARGV[3]) -redis.call('HINCRBY', KEYS[2], ARGV[1], ARGV[2]) -redis.call('EXPIRE', KEYS[2], ARGV[4]) -return 1 -"#; - -async fn add( - redis: &Client, - lineage: Uuid, - field: &str, - amount: i64, - now: DateTime, -) -> Result<(), fred::error::Error> { - if amount <= 0 { - return Ok(()); - } - let _: i64 = redis - .eval( - LUA_ADD, - keys(lineage, now).to_vec(), - vec![ - field.to_string(), - amount.to_string(), - DAY_TTL_SECS.to_string(), - MONTH_TTL_SECS.to_string(), - ], - ) - .await?; - Ok(()) -} - -/// Count one request the rate limits let through. -pub async fn record_request( - redis: &Client, - lineage: Uuid, - now: DateTime, -) -> Result<(), fred::error::Error> { - add(redis, lineage, "requests", 1, now).await -} - -/// Add a call's weighted tokens — the amount its limits were charged. -pub async fn record_tokens( - redis: &Client, - lineage: Uuid, - weighted_tokens: i64, - now: DateTime, -) -> Result<(), fred::error::Error> { - add(redis, lineage, "tokens", weighted_tokens, now).await -} - -/// Read `lineage`'s counts for the day and month containing `now`. A -/// period nothing was counted in reads as zero; a Redis error is -/// returned. -pub async fn read( - redis: &Client, - lineage: Uuid, - now: DateTime, -) -> Result { - let [day, month] = keys(lineage, now); - let fields = || vec!["requests", "tokens"]; - let d: Vec> = redis.hmget(&day, fields()).await?; - let m: Vec> = redis.hmget(&month, fields()).await?; - let at = |v: &[Option], i: usize| v.get(i).copied().flatten().unwrap_or(0); - Ok(KeyUsageCounts { - requests_today: at(&d, 0), - tokens_today: at(&d, 1), - requests_month: at(&m, 0), - tokens_month: at(&m, 1), - }) -} - -#[cfg(test)] -mod tests { - use super::*; - use chrono::TimeZone; - - #[test] - fn a_keys_day_and_month_share_one_cluster_slot() { - use fred::util::redis_keyslot; - let lineage = Uuid::new_v4(); - let now = Utc.with_ymd_and_hms(2026, 10, 11, 23, 59, 59).unwrap(); - let [day, month] = keys(lineage, now); - assert_eq!( - day, - format!("key_usage:{{api_key_lineage:{lineage}}}:daily:2026-10-11") - ); - assert_eq!( - month, - format!("key_usage:{{api_key_lineage:{lineage}}}:monthly:2026-10") - ); - assert_eq!( - redis_keyslot(day.as_bytes()), - redis_keyslot(month.as_bytes()) - ); - } - - #[test] - fn the_buckets_turn_over_at_utc_midnight_and_the_first() { - let lineage = Uuid::nil(); - let before = Utc.with_ymd_and_hms(2026, 10, 31, 23, 59, 59).unwrap(); - let after = Utc.with_ymd_and_hms(2026, 11, 1, 0, 0, 0).unwrap(); - let [d1, m1] = keys(lineage, before); - let [d2, m2] = keys(lineage, after); - assert_ne!(d1, d2); - assert_ne!(m1, m2); - assert!(d2.ends_with(":daily:2026-11-01") && m2.ends_with(":monthly:2026-11")); - } -} diff --git a/crates/common/src/limits/mod.rs b/crates/common/src/limits/mod.rs index b8c5d7a1..f974db2d 100644 --- a/crates/common/src/limits/mod.rs +++ b/crates/common/src/limits/mod.rs @@ -56,8 +56,8 @@ use sqlx::{PgPool, Row}; use uuid::Uuid; pub mod budget; -pub mod key_usage; pub mod sliding; +pub mod usage; pub mod weight; // ---------------------------------------------------------------------------- @@ -1160,8 +1160,8 @@ pub struct RequestLimits { /// carries their Redis Cluster hash tag (`sliding::counter_key`). pub owner: Uuid, /// The lineage of the API key the request came with, whether or not - /// it has limits of its own: its `key_usage` counters count the - /// request. + /// it has limits of its own: its `usage` counters, and its owner's, + /// count the request. pub key_lineage: Option, pub rules: Vec, pub caps: Vec, diff --git a/crates/common/src/limits/usage.rs b/crates/common/src/limits/usage.rs new file mode 100644 index 00000000..b075271c --- /dev/null +++ b/crates/common/src/limits/usage.rs @@ -0,0 +1,209 @@ +// ============================================================================ +// Calendar usage per API key and per user +// +// What each API key (lineage) and each user has been let through and used +// today and this month, UTC, whether or not they have limits: the +// rate-limit and budget counters only exist for the limits that are +// configured, and only from when they were, so they can't say what a key +// or a user without limits has done. `GET /v1/usage` reads these. +// +// A request made with a key counts on the key and on the key's owner, so +// a user's count is the sum over all their keys. They count what the +// limits count, at the moments the limits engine counts it: a request +// once it has passed the rate limits (where `sliding::admit` charges +// `requests` rules), and its weighted tokens after the call (where +// `sliding::record` and `budget::add_weighted_tokens` add them) or when +// it is answered from the response cache. A request the limits refuse +// counts nothing. +// +// Storage (Redis): one hash per subject per period, fields `requests` +// and `tokens`: +// +// usage:{user:}:user::daily:2026-10-11 +// usage:{user:}:user::monthly:2026-10 +// usage:{user:}:api_key_lineage::daily:2026-10-11 +// usage:{user:}:api_key_lineage::monthly:2026-10 +// +// `` is the user the request runs as — for a key, its owner. The +// braces are a Redis Cluster hash tag, as on the rate-limit counters, so +// one script writes all four. Each hash expires 2 × its period after its +// last write, as the budget counters do. +// ============================================================================ + +use chrono::{DateTime, Utc}; +use fred::clients::Client; +use fred::interfaces::{HashesInterface, LuaInterface}; +use uuid::Uuid; + +use super::budget::bucket_id; + +const DAY_TTL_SECS: i64 = 2 * 86_400; +const MONTH_TTL_SECS: i64 = 2 * 31 * 86_400; + +/// What one key or one user has used today and this month (UTC). +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct UsageCounts { + pub requests_today: i64, + pub tokens_today: i64, + pub requests_month: i64, + pub tokens_month: i64, +} + +/// Whose counts: a user's, over all their keys, or one key's (its +/// lineage, across rotations). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum UsageSubject { + User, + ApiKeyLineage(Uuid), +} + +/// The day and month hashes of `subject`, a key of `owner` or `owner` +/// itself, for the day and month containing `now`. +pub fn keys(owner: Uuid, subject: UsageSubject, now: DateTime) -> [String; 2] { + let (kind, id) = match subject { + UsageSubject::User => ("user", owner), + UsageSubject::ApiKeyLineage(lineage) => ("api_key_lineage", lineage), + }; + ["daily", "monthly"].map(|period| { + format!( + "usage:{{user:{owner}}}:{kind}:{id}:{period}:{}", + bucket_id(period, now) + ) + }) +} + +const LUA_ADD: &str = r#" +-- KEYS: day hash, month hash, per subject +-- ARGV: field, amount, day ttl, month ttl +for i = 1, #KEYS do + redis.call('HINCRBY', KEYS[i], ARGV[1], ARGV[2]) + if i % 2 == 1 then + redis.call('EXPIRE', KEYS[i], ARGV[3]) + else + redis.call('EXPIRE', KEYS[i], ARGV[4]) + end +end +return 1 +"#; + +/// Add `amount` to `field` on the key's and its owner's day and month. +async fn add( + redis: &Client, + owner: Uuid, + lineage: Uuid, + field: &str, + amount: i64, + now: DateTime, +) -> Result<(), fred::error::Error> { + if amount <= 0 { + return Ok(()); + } + let mut hashes = keys(owner, UsageSubject::ApiKeyLineage(lineage), now).to_vec(); + hashes.extend(keys(owner, UsageSubject::User, now)); + let _: i64 = redis + .eval( + LUA_ADD, + hashes, + vec![ + field.to_string(), + amount.to_string(), + DAY_TTL_SECS.to_string(), + MONTH_TTL_SECS.to_string(), + ], + ) + .await?; + Ok(()) +} + +/// Count one request the rate limits let through, made by `owner` with +/// the key of `lineage`. +pub async fn record_request( + redis: &Client, + owner: Uuid, + lineage: Uuid, + now: DateTime, +) -> Result<(), fred::error::Error> { + add(redis, owner, lineage, "requests", 1, now).await +} + +/// Add a call's weighted tokens — the amount its limits are charged — +/// on the key of `lineage` and its owner. +pub async fn record_tokens( + redis: &Client, + owner: Uuid, + lineage: Uuid, + weighted_tokens: i64, + now: DateTime, +) -> Result<(), fred::error::Error> { + add(redis, owner, lineage, "tokens", weighted_tokens, now).await +} + +/// Read `subject`'s counts for the day and month containing `now`. A +/// period nothing was counted in reads as zero; a Redis error is +/// returned. +pub async fn read( + redis: &Client, + owner: Uuid, + subject: UsageSubject, + now: DateTime, +) -> Result { + let [day, month] = keys(owner, subject, now); + let fields = || vec!["requests", "tokens"]; + let d: Vec> = redis.hmget(&day, fields()).await?; + let m: Vec> = redis.hmget(&month, fields()).await?; + let at = |v: &[Option], i: usize| v.get(i).copied().flatten().unwrap_or(0); + Ok(UsageCounts { + requests_today: at(&d, 0), + tokens_today: at(&d, 1), + requests_month: at(&m, 0), + tokens_month: at(&m, 1), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::TimeZone; + + #[test] + fn a_key_and_its_owner_share_one_cluster_slot() { + use fred::util::redis_keyslot; + let owner = Uuid::new_v4(); + let lineage = Uuid::new_v4(); + let now = Utc.with_ymd_and_hms(2026, 10, 11, 23, 59, 59).unwrap(); + let [key_day, key_month] = keys(owner, UsageSubject::ApiKeyLineage(lineage), now); + let [user_day, user_month] = keys(owner, UsageSubject::User, now); + assert_eq!( + key_day, + format!("usage:{{user:{owner}}}:api_key_lineage:{lineage}:daily:2026-10-11") + ); + assert_eq!( + key_month, + format!("usage:{{user:{owner}}}:api_key_lineage:{lineage}:monthly:2026-10") + ); + assert_eq!( + user_day, + format!("usage:{{user:{owner}}}:user:{owner}:daily:2026-10-11") + ); + assert_eq!( + user_month, + format!("usage:{{user:{owner}}}:user:{owner}:monthly:2026-10") + ); + let slot = redis_keyslot(key_day.as_bytes()); + for k in [&key_month, &user_day, &user_month] { + assert_eq!(redis_keyslot(k.as_bytes()), slot, "{k}"); + } + } + + #[test] + fn the_buckets_turn_over_at_utc_midnight_and_the_first() { + let owner = Uuid::nil(); + let before = Utc.with_ymd_and_hms(2026, 10, 31, 23, 59, 59).unwrap(); + let after = Utc.with_ymd_and_hms(2026, 11, 1, 0, 0, 0).unwrap(); + let [d1, m1] = keys(owner, UsageSubject::User, before); + let [d2, m2] = keys(owner, UsageSubject::User, after); + assert_ne!(d1, d2); + assert_ne!(m1, m2); + assert!(d2.ends_with(":daily:2026-11-01") && m2.ends_with(":monthly:2026-11")); + } +} diff --git a/crates/gateway/src/proxy/accounting.rs b/crates/gateway/src/proxy/accounting.rs index e9f61f0b..725e42d9 100644 --- a/crates/gateway/src/proxy/accounting.rs +++ b/crates/gateway/src/proxy/accounting.rs @@ -44,14 +44,9 @@ pub(crate) async fn post_flight_account( return; } - // The key's own day and month, counted whether or not it has limits. - if let Some(lineage) = request_limits.key_lineage - && let Err(e) = - limits::key_usage::record_tokens(&redis, lineage, weighted, chrono::Utc::now()).await - { - metrics::counter!("gateway_key_usage_fail_open_total").increment(1); - tracing::warn!("key usage token count failed: {e}"); - } + // The key's and its owner's day and month, counted whether or not + // they have limits. + record_usage_tokens(&redis, request_limits, weighted).await; // Token-metric sliding rules — the user's and the key's, the same // rules the pre-flight checked. Recorded whatever they come to: a @@ -120,3 +115,55 @@ pub(crate) async fn post_flight_account( } } } + +/// Add `weighted` tokens to the day and month usage counters of the +/// request's key and its owner (`limits::usage`). Fail-open: the counts +/// are for reading. +async fn record_usage_tokens( + redis: &fred::clients::Client, + request_limits: &RequestLimits, + weighted: i64, +) { + let Some(lineage) = request_limits.key_lineage else { + return; + }; + if let Err(e) = limits::usage::record_tokens( + redis, + request_limits.owner, + lineage, + weighted, + chrono::Utc::now(), + ) + .await + { + metrics::counter!("gateway_usage_count_fail_open_total").increment(1); + tracing::warn!("usage token count failed: {e}"); + } +} + +/// A request answered from the response cache: the tokens its stored +/// answer records, weighted for `model` as a call's are, on the usage +/// counters of its key and owner. The caller received those tokens. The +/// store keeps only the input and output totals, so all of the input +/// counts as plain input. +pub(crate) async fn count_cache_hit_usage( + db: &PgPool, + redis: &fred::clients::Client, + weight_cache: &weight::WeightCache, + model: &str, + prompt_tokens: u32, + completion_tokens: u32, + request_limits: &RequestLimits, +) { + let tokens = weight::TokenCounts { + input: i64::from(prompt_tokens), + cache_read: 0, + cache_write: 0, + cache_write_1h: false, + output: i64::from(completion_tokens), + }; + let weighted = weight::weighted_tokens(&tokens, weight_cache.get(db, model).await); + if weighted > 0 { + record_usage_tokens(redis, request_limits, weighted).await; + } +} diff --git a/crates/gateway/src/proxy/generate.rs b/crates/gateway/src/proxy/generate.rs index 90aa7600..e0fca742 100644 --- a/crates/gateway/src/proxy/generate.rs +++ b/crates/gateway/src/proxy/generate.rs @@ -850,6 +850,16 @@ async fn run( if let Err(e) = state.quota.consume("a_key, total).await { tracing::warn!(quota_key = %quota_key, tokens = total, "quota consume on cache hit failed: {e}"); } + super::accounting::count_cache_hit_usage( + &state.db, + &state.redis, + &state.weight_cache, + &mapped_model, + cached.prompt_tokens, + cached.completion_tokens, + &preflight.limits, + ) + .await; // Same capture pipeline as a fresh request — redaction toggle, // byte cap and offload all apply — with the status marked. diff --git a/crates/gateway/src/proxy/pipeline.rs b/crates/gateway/src/proxy/pipeline.rs index a44ff572..ae0f4838 100644 --- a/crates/gateway/src/proxy/pipeline.rs +++ b/crates/gateway/src/proxy/pipeline.rs @@ -24,7 +24,7 @@ use think_watch_common::lifecycle::stages::{ check_access, check_budget, check_limits, run_post_invoke, }; use think_watch_common::lifecycle::state::{CapturedView, Invoked, LimitCheckRecord, Raw}; -use think_watch_common::limits::{RequestLimits, key_usage}; +use think_watch_common::limits::{RequestLimits, usage}; use tw_dialect::ir::Dialect; /// Pre-flight result threaded through to `ChatPostInvokeDeps` later @@ -78,13 +78,15 @@ pub(super) async fn run_preflight_stages( ) .await .map_err(short_circuit_to_response)?; - // Past the rate limits, the request counts on its key, as it has on - // every `requests` rule. Fail-open: the counts are for reading. + // Past the rate limits, the request counts on its key and its owner, + // as it has on every `requests` rule. Fail-open: the counts are for + // reading. if let Some(lineage) = limits.key_lineage - && let Err(e) = key_usage::record_request(&state.redis, lineage, chrono::Utc::now()).await + && let Err(e) = + usage::record_request(&state.redis, limits.owner, lineage, chrono::Utc::now()).await { - metrics::counter!("gateway_key_usage_fail_open_total").increment(1); - tracing::warn!("key usage request count failed: {e}"); + metrics::counter!("gateway_usage_count_fail_open_total").increment(1); + tracing::warn!("usage request count failed: {e}"); } let _authorized = check_access::(limits_checked, model, &state.audit) .await diff --git a/crates/server/src/handlers/key_usage.rs b/crates/server/src/handlers/key_usage.rs index 3ac0e781..9964a05e 100644 --- a/crates/server/src/handlers/key_usage.rs +++ b/crates/server/src/handlers/key_usage.rs @@ -1,17 +1,22 @@ // ============================================================================ // GET /v1/usage (gateway port) // -// What the calling API key has used today and this month, every limit -// that binds its AI-gateway requests with what has been used of it, and -// when the key stops working. +// What the calling API key may still do on the AI gateway, about one of +// two subjects: +// +// * `scope: "key"` — the key has limits of its own on the AI gateway. +// The answer is about the key alone: its limits, on its lineage's +// counters, and its own usage. +// * `scope: "user"` — the key has none. The answer is about its owner +// as a whole: the owner's effective limits — roles, including those a +// team grants, merged most-restrictive, then the user's overrides — on +// the owner's counters, and the owner's usage over all their keys. +// With no limits on either, `limits` is empty. // // The limits are the ones enforcement builds for a request with this key -// (`limits_for_ai_gateway`): the key's own, on its lineage's counters, and -// its owner's effective limits — roles, including those a team grants, -// merged most-restrictive, then the user's overrides — on the owner's -// counters, which count every key the owner has. They are read from the -// same Redis counters, the same way, without charging them. MCP-surface -// limits don't apply to model requests and are left out. +// (`limits_for_ai_gateway`), read from the same Redis counters, the same +// way, without charging them. MCP-surface limits don't apply to model +// requests and are left out. // // Mounted behind `require_api_key_to_read`: the same key checks as a model // request, but calling it is not a use of the key. @@ -30,7 +35,8 @@ use uuid::Uuid; use think_watch_common::limits::{ BudgetCap, BudgetPeriod, BudgetSubject, RateLimitRule, RateLimitSubject, RequestLimits, budget, - key_usage, secs_to_window, sliding, + secs_to_window, sliding, + usage::{self, UsageCounts, UsageSubject}, }; use think_watch_gateway::proxy::{GatewayRequestIdentity, limits_for_ai_gateway}; @@ -38,21 +44,34 @@ use crate::app::AppState; #[derive(Debug, Serialize, PartialEq)] pub struct KeyUsageResponse { + /// What the answer is about: the key, or its owner as a whole. + pub scope: Scope, pub usage: Usage, + /// The limits of `scope` that bind the key's model requests. pub limits: Vec, /// When the key stops authenticating; `null` when it does not. pub expires_at: Option>, } -/// What this key (its lineage, across rotations) has done, UTC day and -/// month. Tokens are weighted tokens, as limits count them. +#[derive(Debug, Clone, Copy, Serialize, PartialEq, Eq)] +#[serde(rename_all = "lowercase")] +pub enum Scope { + /// The key has limits of its own: the key alone (its lineage, across + /// rotations). + Key, + /// The key has none: its owner, over all their keys. + User, +} + +/// What the subject has done, UTC day and month. Tokens are weighted +/// tokens, as limits count them. #[derive(Debug, Serialize, PartialEq)] pub struct Usage { pub requests_today: i64, pub tokens_today: i64, pub requests_month: i64, pub tokens_month: i64, - /// The key's cost this month from the request log, as the cost + /// The subject's cost this month from the request log, as the cost /// reports sum it; `null` without ClickHouse or when it can't be read. pub cost_usd_month: Option, } @@ -60,9 +79,8 @@ pub struct Usage { /// One limit that binds the key's requests. #[derive(Debug, Serialize, PartialEq)] pub struct LimitUsage { - /// `key`: set on the key, counted for it alone. `user`: its owner's, - /// counted for everything the owner does. - pub scope: &'static str, + /// Whose limit it is; always the answer's `scope`. + pub scope: Scope, /// `requests` or `tokens` (weighted). pub kind: &'static str, /// A rate limit's sliding window (`1m` … `1w`) or a budget's calendar @@ -82,15 +100,16 @@ pub async fn get_key_usage( State(state): State, axum::Extension(identity): axum::Extension, ) -> Response { - let limits = limits_for_ai_gateway(&identity); + let (scope, limits) = scoped(limits_for_ai_gateway(&identity)); let now = Utc::now(); + let subject = match (scope, limits.key_lineage) { + (Scope::Key, Some(lineage)) => UsageSubject::ApiKeyLineage(lineage), + _ => UsageSubject::User, + }; - let counts = match limits.key_lineage { - Some(lineage) => match key_usage::read(&state.redis, lineage, now).await { - Ok(c) => c, - Err(e) => return unavailable(&e), - }, - None => key_usage::KeyUsageCounts::default(), + let counts = match usage::read(&state.redis, limits.owner, subject, now).await { + Ok(c) => c, + Err(e) => return unavailable(&e), }; let mut rule_used = Vec::with_capacity(limits.rules.len()); for rule in &limits.rules { @@ -104,12 +123,10 @@ pub async fn get_key_usage( Ok(statuses) => statuses.into_iter().map(|s| s.current).collect(), Err(e) => return unavailable(&e), }; - let cost = match limits.key_lineage { - Some(lineage) => cost_this_month(&state, lineage, now).await, - None => None, - }; + let cost = cost_this_month(&state, limits.owner, subject, now).await; Json(response( + scope, &limits, &rule_used, &cap_used, @@ -121,15 +138,53 @@ pub async fn get_key_usage( .into_response() } -/// The answer for a request's limits and what each has used +/// Which subject the answer is about, and the limits of that subject +/// alone: the key's own when it has any on the AI gateway, its owner's +/// otherwise. +fn scoped(mut limits: RequestLimits) -> (Scope, RequestLimits) { + let key_has_limits = limits + .rules + .iter() + .any(|r| r.subject_kind == RateLimitSubject::ApiKeyLineage) + || limits + .caps + .iter() + .any(|c| c.subject_kind == BudgetSubject::ApiKeyLineage); + let scope = if key_has_limits { + Scope::Key + } else { + Scope::User + }; + limits.rules.retain(|r| rule_scope(r) == scope); + limits.caps.retain(|c| cap_scope(c) == scope); + (scope, limits) +} + +fn rule_scope(rule: &RateLimitRule) -> Scope { + match rule.subject_kind { + RateLimitSubject::ApiKeyLineage => Scope::Key, + RateLimitSubject::User => Scope::User, + } +} + +fn cap_scope(cap: &BudgetCap) -> Scope { + match cap.subject_kind { + BudgetSubject::ApiKeyLineage => Scope::Key, + BudgetSubject::User => Scope::User, + } +} + +/// The answer for one scope's limits and what each has used /// (`rule_used` pairs with `limits.rules`, `cap_used` with -/// `limits.caps`), the key's counts and cost. Limits are ordered by the -/// share of them left, least first. +/// `limits.caps`), and that scope's counts and cost. Limits are ordered +/// by the share of them left, least first. +#[allow(clippy::too_many_arguments)] fn response( + scope: Scope, limits: &RequestLimits, rule_used: &[i64], cap_used: &[i64], - counts: key_usage::KeyUsageCounts, + counts: UsageCounts, cost_usd_month: Option, expires_at: Option>, now: DateTime, @@ -149,6 +204,7 @@ fn response( .collect(); entries.sort_by(by_share_left); KeyUsageResponse { + scope, usage: Usage { requests_today: counts.requests_today, tokens_today: counts.tokens_today, @@ -163,10 +219,7 @@ fn response( fn rule_entry(rule: &RateLimitRule, used: i64) -> LimitUsage { LimitUsage { - scope: match rule.subject_kind { - RateLimitSubject::ApiKeyLineage => "key", - RateLimitSubject::User => "user", - }, + scope: rule_scope(rule), kind: rule.metric.as_str(), window: secs_to_window(rule.window_secs) .map(str::to_string) @@ -180,10 +233,7 @@ fn rule_entry(rule: &RateLimitRule, used: i64) -> LimitUsage { fn cap_entry(cap: &BudgetCap, used: i64, now: DateTime) -> LimitUsage { LimitUsage { - scope: match cap.subject_kind { - BudgetSubject::ApiKeyLineage => "key", - BudgetSubject::User => "user", - }, + scope: cap_scope(cap), kind: "tokens", window: cap.period.as_str().to_string(), window_secs: None, @@ -193,16 +243,14 @@ fn cap_entry(cap: &BudgetCap, used: i64, now: DateTime) -> LimitUsage { } } -/// Least of its limit left first; on a tie, the key's before its -/// owner's, sliding windows shortest first before calendar periods, and -/// requests before tokens. +/// Least of its limit left first; on a tie, sliding windows shortest +/// first before calendar periods, and requests before tokens. fn by_share_left(a: &LimitUsage, b: &LimitUsage) -> Ordering { // (limit - used) / limit, compared without dividing; limits are > 0. let left = |e: &LimitUsage| i128::from(e.limit) - i128::from(e.used); let (la, lb) = (i128::from(a.limit.max(1)), i128::from(b.limit.max(1))); (left(a) * lb) .cmp(&(left(b) * la)) - .then_with(|| (a.scope != "key").cmp(&(b.scope != "key"))) .then_with(|| window_rank(a).cmp(&window_rank(b))) .then_with(|| (a.kind != "requests").cmp(&(b.kind != "requests"))) } @@ -219,11 +267,17 @@ fn window_rank(e: &LimitUsage) -> i64 { }) } -/// The key's cost since the 1st of this month, UTC, summed from the -/// request log as the cost reports sum it. `None` without ClickHouse; -/// a failed read is logged and also `None` — the cost is informational, -/// the counters above are what limits use. -async fn cost_this_month(state: &AppState, lineage: Uuid, now: DateTime) -> Option { +/// The subject's cost since the 1st of this month, UTC, summed from the +/// request log as the cost reports sum it: the key's lineage, or every +/// request of the owner. `None` without ClickHouse; a failed read is +/// logged and also `None` — the cost is informational, the counters +/// above are what limits use. +async fn cost_this_month( + state: &AppState, + owner: Uuid, + subject: UsageSubject, + now: DateTime, +) -> Option { #[derive(clickhouse::Row, Deserialize)] struct Cost { cost: i128, @@ -235,14 +289,18 @@ async fn cost_this_month(state: &AppState, lineage: Uuid, now: DateTime) -> .expect("every month has a 1st") .and_time(NaiveTime::MIN) .and_utc(); + let (column, id) = match subject { + UsageSubject::ApiKeyLineage(lineage) => ("api_key_lineage_id", lineage), + UsageSubject::User => ("user_id", owner), + }; let row = ch - .query( + .query(&format!( "SELECT sum(ifNull(cost_usd, 0)) AS cost FROM gateway_logs \ PREWHERE created_at >= parseDateTimeBestEffort(?) \ - AND api_key_lineage_id = ?", - ) + AND {column} = ?" + )) .bind(month_start.format("%Y-%m-%d %H:%M:%S").to_string()) - .bind(lineage.to_string()) + .bind(id.to_string()) .fetch_one::() .await; match row { @@ -305,8 +363,8 @@ mod tests { } } - fn counts() -> key_usage::KeyUsageCounts { - key_usage::KeyUsageCounts { + fn counts() -> UsageCounts { + UsageCounts { requests_today: 3, tokens_today: 300, requests_month: 40, @@ -314,18 +372,48 @@ mod tests { } } - #[test] - fn a_key_bound_by_nothing_still_reports_its_usage() { - let limits = RequestLimits::for_request( + /// The owner has a requests/min rule and a monthly budget; MCP rules + /// don't bind model requests. + fn owners() -> SurfaceConstraints { + SurfaceConstraints { + ai_gateway: Some(block( + &[(RateMetric::Requests, 60, 10)], + &[(BudgetPeriod::Monthly, 1_000_000)], + )), + mcp_gateway: Some(block(&[(RateMetric::Requests, 60, 1)], &[])), + } + } + + fn limits_for(user: &SurfaceConstraints, key: &SurfaceConstraints) -> RequestLimits { + RequestLimits::for_request( Surface::AiGateway, Uuid::new_v4(), + user, + Some((Uuid::new_v4(), key)), + ) + } + + #[test] + fn neither_the_key_nor_its_owner_has_limits() { + let (scope, limits) = scoped(limits_for( + &SurfaceConstraints::default(), &SurfaceConstraints::default(), - Some((Uuid::new_v4(), &SurfaceConstraints::default())), + )); + assert_eq!(scope, Scope::User); + let body = response( + scope, + &limits, + &[], + &[], + counts(), + None, + None, + at(2026, 10, 14, 9), ); - let body = response(&limits, &[], &[], counts(), None, None, at(2026, 10, 14, 9)); assert_eq!( serde_json::to_value(&body).unwrap(), serde_json::json!({ + "scope": "user", "usage": {"requests_today": 3, "tokens_today": 300, "requests_month": 40, "tokens_month": 4000, "cost_usd_month": null}, @@ -336,17 +424,49 @@ mod tests { } #[test] - fn the_keys_and_its_owners_limits_map_to_the_contract_shape() { - // Wednesday 2026-10-14, 09:00 UTC. The owner has a requests/min - // rule and a monthly budget (MCP rules don't bind model requests); - // the key a 5h token window and a daily budget. - let user = SurfaceConstraints { - ai_gateway: Some(block( - &[(RateMetric::Requests, 60, 10)], - &[(BudgetPeriod::Monthly, 1_000_000)], - )), + fn a_key_without_limits_of_its_own_answers_for_its_owner() { + // An MCP-only rule on the key is not a limit on its model requests. + let key = SurfaceConstraints { + ai_gateway: None, mcp_gateway: Some(block(&[(RateMetric::Requests, 60, 1)], &[])), }; + let (scope, limits) = scoped(limits_for(&owners(), &key)); + assert_eq!(scope, Scope::User); + // Wednesday 2026-10-14, 09:00 UTC. Input order: rule, cap. + let body = response( + scope, + &limits, + &[5], + &[250_000], + counts(), + Some(1.25), + Some(at(2026, 12, 31, 0)), + at(2026, 10, 14, 9), + ); + assert_eq!( + serde_json::to_value(&body).unwrap(), + serde_json::json!({ + "scope": "user", + "usage": {"requests_today": 3, "tokens_today": 300, + "requests_month": 40, "tokens_month": 4000, + "cost_usd_month": 1.25}, + "limits": [ + // 50 % left. + {"scope": "user", "kind": "requests", "window": "1m", "window_secs": 60, + "limit": 10, "used": 5, "resets_at": null}, + // 75 % left. + {"scope": "user", "kind": "tokens", "window": "monthly", "window_secs": null, + "limit": 1000000, "used": 250000, "resets_at": "2026-11-01T00:00:00Z"}, + ], + "expires_at": "2026-12-31T00:00:00Z", + }) + ); + } + + #[test] + fn a_key_with_limits_of_its_own_answers_for_itself_alone() { + // The owner's limits still bind the key's requests, but the answer + // is about the key: its limits only. let key = SurfaceConstraints { ai_gateway: Some(block( &[(RateMetric::Tokens, 18_000, 50_000)], @@ -354,28 +474,25 @@ mod tests { )), mcp_gateway: None, }; - let limits = RequestLimits::for_request( - Surface::AiGateway, - Uuid::new_v4(), - &user, - Some((Uuid::new_v4(), &key)), - ); - // Input order: user rule, key rule; user cap, key cap. + let (scope, limits) = scoped(limits_for(&owners(), &key)); + assert_eq!(scope, Scope::Key); let body = response( + scope, &limits, - &[5, 45_000], - &[250_000, 100_500], + &[45_000], + &[100_500], counts(), - Some(1.25), - Some(at(2026, 12, 31, 0)), + None, + None, at(2026, 10, 14, 9), ); assert_eq!( serde_json::to_value(&body).unwrap(), serde_json::json!({ + "scope": "key", "usage": {"requests_today": 3, "tokens_today": 300, "requests_month": 40, "tokens_month": 4000, - "cost_usd_month": 1.25}, + "cost_usd_month": null}, "limits": [ // Overshot: -0.5 % left. {"scope": "key", "kind": "tokens", "window": "daily", "window_secs": null, @@ -383,22 +500,28 @@ mod tests { // 10 % left. {"scope": "key", "kind": "tokens", "window": "5h", "window_secs": 18000, "limit": 50000, "used": 45000, "resets_at": null}, - // 50 % left. - {"scope": "user", "kind": "requests", "window": "1m", "window_secs": 60, - "limit": 10, "used": 5, "resets_at": null}, - // 75 % left. - {"scope": "user", "kind": "tokens", "window": "monthly", "window_secs": null, - "limit": 1000000, "used": 250000, "resets_at": "2026-11-01T00:00:00Z"}, ], - "expires_at": "2026-12-31T00:00:00Z", + "expires_at": null, }) ); } #[test] - fn ties_put_the_key_first_then_the_shorter_window_then_requests() { - let e = |scope, kind, window: &str, window_secs| LimitUsage { - scope, + fn a_key_budget_alone_makes_the_answer_the_keys() { + let key = SurfaceConstraints { + ai_gateway: Some(block(&[], &[(BudgetPeriod::Weekly, 7)])), + mcp_gateway: None, + }; + let (scope, limits) = scoped(limits_for(&owners(), &key)); + assert_eq!(scope, Scope::Key); + assert!(limits.rules.is_empty()); + assert_eq!(limits.caps.len(), 1); + } + + #[test] + fn ties_put_the_shorter_window_first_then_requests() { + let e = |kind, window: &str, window_secs| LimitUsage { + scope: Scope::Key, kind, window: window.to_string(), window_secs, @@ -407,27 +530,22 @@ mod tests { resets_at: None, }; let mut v = [ - e("user", "requests", "1m", Some(60)), - e("key", "tokens", "monthly", None), - e("key", "tokens", "weekly", None), - e("key", "tokens", "1w", Some(604_800)), - e("key", "tokens", "1m", Some(60)), - e("key", "requests", "1m", Some(60)), + e("tokens", "monthly", None), + e("tokens", "weekly", None), + e("tokens", "1w", Some(604_800)), + e("tokens", "1m", Some(60)), + e("requests", "1m", Some(60)), ]; v.sort_by(by_share_left); - let order: Vec<(&str, &str, &str)> = v - .iter() - .map(|e| (e.scope, e.kind, e.window.as_str())) - .collect(); + let order: Vec<(&str, &str)> = v.iter().map(|e| (e.kind, e.window.as_str())).collect(); assert_eq!( order, [ - ("key", "requests", "1m"), - ("key", "tokens", "1m"), - ("key", "tokens", "1w"), - ("key", "tokens", "weekly"), - ("key", "tokens", "monthly"), - ("user", "requests", "1m"), + ("requests", "1m"), + ("tokens", "1m"), + ("tokens", "1w"), + ("tokens", "weekly"), + ("tokens", "monthly"), ] ); } diff --git a/crates/test-support/tests/analytics_clickhouse.rs b/crates/test-support/tests/analytics_clickhouse.rs index 28b1847b..ba1e0296 100644 --- a/crates/test-support/tests/analytics_clickhouse.rs +++ b/crates/test-support/tests/analytics_clickhouse.rs @@ -875,8 +875,10 @@ async fn cached_input_is_billed_at_the_cache_prices() { #[ignore = "integration test — run via `make test-it`"] #[tokio::test] -async fn v1_usage_reports_the_keys_cost_this_month_from_the_request_log() { - // Two keys of one owner: each key's month counts its own requests. +async fn v1_usage_reports_the_cost_this_month_from_the_request_log() { + // Two keys of one owner. Key b has a limit of its own, so it answers + // with its own cost; key a has none, so it answers with the owner's, + // over both keys. let app = TestApp::spawn_with_clickhouse().await; let (key_a, user_id) = seed_runtime(&app).await; let key_b = fixtures::create_api_key( @@ -889,6 +891,17 @@ async fn v1_usage_reports_the_keys_cost_this_month_from_the_request_log() { ) .await .unwrap(); + fixtures::create_rate_limit_rule( + &app.db, + "api_key_lineage", + key_b.row.lineage_id, + "ai_gateway", + "requests", + 60, + 100, + ) + .await + .unwrap(); let ch = app.state.clickhouse.as_ref().expect("clickhouse wired up"); let gw_a = app.gateway_client(); @@ -920,20 +933,25 @@ async fn v1_usage_reports_the_keys_cost_this_month_from_the_request_log() { tokio::time::sleep(std::time::Duration::from_millis(50)).await; } - let cost = |gw: &TestClient| { + let usage = |gw: &TestClient| { let gw = gw.clone(); async move { let r = gw.get("/v1/usage").await.unwrap(); r.assert_ok(); let v: Value = r.json().unwrap(); - v["usage"]["cost_usd_month"] + let cost = v["usage"]["cost_usd_month"] .as_f64() - .unwrap_or_else(|| panic!("cost_usd_month is a number: {v:#}")) + .unwrap_or_else(|| panic!("cost_usd_month is a number: {v:#}")); + (v["scope"].as_str().unwrap().to_string(), cost) } }; - let (a, b) = (cost(&gw_a).await, cost(&gw_b).await); + let ((scope_a, a), (scope_b, b)) = (usage(&gw_a).await, usage(&gw_b).await); + assert_eq!((scope_a.as_str(), scope_b.as_str()), ("user", "key")); assert!(one_call > 0.0); - assert!((a - one_call).abs() < 1e-9, "a={a} one call={one_call}"); + assert!( + (a - 3.0 * one_call).abs() < 1e-9, + "a={a} one call={one_call}" + ); assert!( (b - 2.0 * one_call).abs() < 1e-9, "b={b} one call={one_call}" diff --git a/crates/test-support/tests/limits.rs b/crates/test-support/tests/limits.rs index 5480cdd1..9cc093fd 100644 --- a/crates/test-support/tests/limits.rs +++ b/crates/test-support/tests/limits.rs @@ -1032,11 +1032,21 @@ async fn role_with_constraints(app: &TestApp, user_id: Uuid, constraints: Json, } } +/// A chat request whose prompt no other request in the test repeats, so +/// it is never answered from the response cache. +fn distinct_chat(n: u32) -> Json { + json!({"model": "gpt-test", "messages": [{"role": "user", "content": format!("x{n}")}]}) +} + #[ignore = "integration test — run via `make test-it`"] #[tokio::test] -async fn v1_usage_reports_a_keys_own_limits_on_its_own_counters() { +async fn v1_usage_of_a_key_with_limits_answers_for_the_key_alone() { + // The key has limits of its own; its owner has limits too, one with + // less left than any of the key's, and another key of the owner is + // used as well. The answer is the key's limits on the key's counters + // and the key's own usage, nothing of the owner's. let app = TestApp::spawn().await; - let (_, user_id) = seed_runtime(&app).await; + let (other_key, user_id) = seed_runtime(&app).await; let key = fixtures::create_api_key( &app.db, user_id, @@ -1076,18 +1086,32 @@ async fn v1_usage_reports_a_keys_own_limits_on_its_own_counters() { ) .await .unwrap(); + fixtures::create_rate_limit_rule(&app.db, "user", user_id, "ai_gateway", "requests", 60, 5) + .await + .unwrap(); + fixtures::create_budget_cap(&app.db, "user", user_id, "monthly", 9_000_000) + .await + .unwrap(); let con = admin_session(&app).await; let gw = app.gateway_client(); gw.set_bearer(&key.plaintext); - for _ in 0..2 { - gw.post("/v1/chat/completions", chat()) + for n in 0..2 { + gw.post("/v1/chat/completions", distinct_chat(n)) .await .unwrap() .assert_ok(); } + let other = app.gateway_client(); + other.set_bearer(&other_key); + other + .post("/v1/chat/completions", distinct_chat(2)) + .await + .unwrap() + .assert_ok(); let v = key_usage(&gw).await; + assert_eq!(v["scope"], "key"); // The console reads the same counters: the AI rules ordered requests // then tokens (the MCP rule last), then the cap. let (rules, caps) = usage(&con, "api_key", key.row.id).await; @@ -1125,10 +1149,11 @@ async fn v1_usage_reports_a_keys_own_limits_on_its_own_counters() { #[ignore = "integration test — run via `make test-it`"] #[tokio::test] -async fn v1_usage_reports_the_owners_effective_limits_on_the_owners_counters() { - // The owner's limits from a role, from a role its team grants, and - // from an override, counting what every key of the owner does. The - // key's own usage counts only its own requests. +async fn v1_usage_of_a_key_without_limits_answers_for_its_owner() { + // Neither key has limits of its own. The owner's limits come from a + // role, from a role its team grants, and from an override; both keys + // answer with them, counted over everything the owner does, and with + // the owner's usage over both keys. let app = TestApp::spawn().await; let (key_a, user_id) = seed_runtime(&app).await; let key_b = second_key(&app, user_id).await; @@ -1152,19 +1177,34 @@ async fn v1_usage_reports_the_owners_effective_limits_on_the_owners_counters() { fixtures::create_rate_limit_rule(&app.db, "user", user_id, "ai_gateway", "requests", 60, 50) .await .unwrap(); + // A limit only on the MCP gateway does not make the answer the key's. + fixtures::create_rate_limit_rule( + &app.db, + "api_key_lineage", + key_b.row.lineage_id, + "mcp_gateway", + "requests", + 60, + 7, + ) + .await + .unwrap(); let gw_a = app.gateway_client(); gw_a.set_bearer(&key_a); let gw_b = app.gateway_client(); gw_b.set_bearer(&key_b.plaintext); - // Distinct prompts: a cached answer charges no tokens. - let prompt = |n: u32| json!({"model": "gpt-test", "messages": [{"role": "user", "content": format!("x{n}")}]}); - gw_a.post("/v1/chat/completions", prompt(1)) + gw_a.post("/v1/chat/completions", distinct_chat(1)) .await .unwrap() .assert_ok(); + // Every request against the mock comes to the same tokens. + let one = key_usage(&gw_a).await["usage"]["tokens_today"] + .as_i64() + .unwrap(); + assert!(one > 0); for n in 2..4 { - gw_b.post("/v1/chat/completions", prompt(n)) + gw_b.post("/v1/chat/completions", distinct_chat(n)) .await .unwrap() .assert_ok(); @@ -1172,14 +1212,8 @@ async fn v1_usage_reports_the_owners_effective_limits_on_the_owners_counters() { let a = key_usage(&gw_a).await; let b = key_usage(&gw_b).await; - let (a_tokens, b_tokens) = ( - a["usage"]["tokens_today"].as_i64().unwrap(), - b["usage"]["tokens_today"].as_i64().unwrap(), - ); - assert!(a_tokens > 0 && b_tokens > a_tokens, "{a:#} {b:#}"); - assert_eq!(a["usage"]["requests_today"], 1); - assert_eq!(b["usage"]["requests_today"], 2); - let all = a_tokens + b_tokens; + assert_eq!(a["scope"], "user"); + let all = 3 * one; let expected = vec![ row("user", "requests", "1m", 50, 3), row("user", "tokens", "1h", 4_000_000, all), @@ -1188,75 +1222,58 @@ async fn v1_usage_reports_the_owners_effective_limits_on_the_owners_counters() { ]; assert_eq!(limit_rows(&a), expected); assert_eq!( - limit_rows(&b), - expected, - "both keys are bound by the same limits" - ); -} - -#[ignore = "integration test — run via `make test-it`"] -#[tokio::test] -async fn v1_usage_lists_the_keys_and_the_owners_limits_least_left_first() { - let app = TestApp::spawn().await; - let (_, user_id) = seed_runtime(&app).await; - let key = second_key(&app, user_id).await; - fixtures::create_rate_limit_rule(&app.db, "user", user_id, "ai_gateway", "requests", 60, 50) - .await - .unwrap(); - fixtures::create_rate_limit_rule( - &app.db, - "api_key_lineage", - key.row.lineage_id, - "ai_gateway", - "requests", - 300, - 5, - ) - .await - .unwrap(); - - let gw = app.gateway_client(); - gw.set_bearer(&key.plaintext); - gw.post("/v1/chat/completions", chat()) - .await - .unwrap() - .assert_ok(); - assert_eq!( - limit_rows(&key_usage(&gw).await), - vec![ - row("key", "requests", "5m", 5, 1), - row("user", "requests", "1m", 50, 1), - ] + a["usage"], + json!({"requests_today": 3, "tokens_today": all, + "requests_month": 3, "tokens_month": all, + "cost_usd_month": null}) ); + assert_eq!(a, b, "both keys answer for the same owner"); } #[ignore = "integration test — run via `make test-it`"] #[tokio::test] -async fn v1_usage_of_a_key_bound_by_nothing_still_reports_its_usage() { +async fn v1_usage_with_no_limits_anywhere_answers_for_the_owner() { let app = TestApp::spawn().await; - let (api_key, _) = seed_runtime(&app).await; + let (api_key, user_id) = seed_runtime(&app).await; + let other = second_key(&app, user_id).await; // Anthropic's SDKs send the key in x-api-key. let gw = app.gateway_client(); gw.set_header("x-api-key", &api_key); assert_eq!( key_usage(&gw).await, - json!({"usage": {"requests_today": 0, "tokens_today": 0, + json!({"scope": "user", + "usage": {"requests_today": 0, "tokens_today": 0, "requests_month": 0, "tokens_month": 0, "cost_usd_month": null}, "limits": [], "expires_at": null}) ); - gw.post("/v1/chat/completions", chat()) + gw.post("/v1/chat/completions", distinct_chat(1)) .await .unwrap() .assert_ok(); + let one = key_usage(&gw).await["usage"]["tokens_today"] + .as_i64() + .unwrap(); + assert!(one > 0); + // The owner's other key counts too. + let gw_other = app.gateway_client(); + gw_other.set_bearer(&other.plaintext); + gw_other + .post("/v1/chat/completions", distinct_chat(2)) + .await + .unwrap() + .assert_ok(); + let v = key_usage(&gw).await; + assert_eq!(v["scope"], "user"); assert_eq!(v["limits"], json!([])); - assert_eq!(v["usage"]["requests_today"], 1); - assert_eq!(v["usage"]["requests_month"], 1); + assert_eq!(v["usage"]["requests_today"], 2); + assert_eq!(v["usage"]["requests_month"], 2); let tokens = v["usage"]["tokens_today"].as_i64().unwrap(); - assert!(tokens > 0, "{v:#}"); + assert_eq!(tokens, 2 * one, "{v:#}"); assert_eq!(v["usage"]["tokens_month"], tokens); + assert_eq!(key_usage(&gw_other).await, v); } #[ignore = "integration test — run via `make test-it`"] @@ -1389,3 +1406,60 @@ async fn calling_v1_usage_moves_no_counter() { // The refused request counted nothing. assert_eq!(key_usage(&gw).await["usage"]["requests_today"], 2); } + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn v1_usage_counts_the_tokens_of_an_answer_from_the_cache() { + // A key with a limit of its own answers for itself; the owner's other + // key, without one, for the owner. An answer from the response cache + // counts as a request and its tokens on both. + let app = TestApp::spawn().await; + let (other_key, user_id) = seed_runtime(&app).await; + let key = second_key(&app, user_id).await; + fixtures::create_rate_limit_rule( + &app.db, + "api_key_lineage", + key.row.lineage_id, + "ai_gateway", + "requests", + 60, + 100, + ) + .await + .unwrap(); + + let gw = app.gateway_client(); + gw.set_bearer(&key.plaintext); + let x_cache = |r: &think_watch_test_support::client::TestResponse| { + r.headers + .get("x-cache") + .and_then(|v| v.to_str().ok()) + .unwrap_or("") + .to_string() + }; + let first = gw.post("/v1/chat/completions", chat()).await.unwrap(); + first.assert_ok(); + assert_eq!(x_cache(&first), "MISS"); + let one = key_usage(&gw).await["usage"]["tokens_today"] + .as_i64() + .unwrap(); + assert!(one > 0); + let second = gw.post("/v1/chat/completions", chat()).await.unwrap(); + second.assert_ok(); + assert_eq!(x_cache(&second), "HIT"); + + let v = key_usage(&gw).await; + assert_eq!(v["scope"], "key"); + assert_eq!(limit_rows(&v), vec![row("key", "requests", "1m", 100, 2)]); + assert_eq!( + v["usage"], + json!({"requests_today": 2, "tokens_today": 2 * one, + "requests_month": 2, "tokens_month": 2 * one, + "cost_usd_month": null}) + ); + let owner = app.gateway_client(); + owner.set_bearer(&other_key); + let o = key_usage(&owner).await; + assert_eq!(o["scope"], "user"); + assert_eq!(o["usage"], v["usage"]); +} diff --git a/web/src/i18n/en.json b/web/src/i18n/en.json index 2af13915..bbeda4b7 100644 --- a/web/src/i18n/en.json +++ b/web/src/i18n/en.json @@ -1942,7 +1942,7 @@ "responsesEndpoint": "OpenAI Responses (new): /v1/responses", "responsesWsEndpoint": "OpenAI Responses over WebSocket: /v1/responses (one response.create frame per turn)", "modelsEndpoint": "List models: /v1/models", - "usageEndpoint": "Usage and limits of the calling key: /v1/usage", + "usageEndpoint": "Usage and limits: /v1/usage (the calling key's when it has limits of its own, otherwise its owner's)", "claudeDesktop": "Claude Desktop", "claudeDesktopDesc": "Configure Claude Desktop to use MCP tools through ThinkWatch.", "claudeCodeMcpDesc": "Register ThinkWatch as an MCP server directly from the Claude Code CLI.", diff --git a/web/src/i18n/zh.json b/web/src/i18n/zh.json index dae1203f..0b1f25c4 100644 --- a/web/src/i18n/zh.json +++ b/web/src/i18n/zh.json @@ -1942,7 +1942,7 @@ "responsesEndpoint": "OpenAI Responses（新版）：/v1/responses", "responsesWsEndpoint": "OpenAI Responses（WebSocket）：/v1/responses（每轮发一个 response.create 帧）", "modelsEndpoint": "模型列表：/v1/models", - "usageEndpoint": "调用密钥的用量与限额：/v1/usage", + "usageEndpoint": "用量与限额：/v1/usage（密钥有单独的限额时为该密钥，否则为所属用户）", "claudeDesktop": "Claude Desktop", "claudeDesktopDesc": "配置 Claude Desktop 通过 ThinkWatch 使用 MCP 工具。", "claudeCodeMcpDesc": "通过 Claude Code CLI 将 ThinkWatch 注册为 MCP 服务器。", From 997ec7398ac7b764c52a045d10c6cb15d58ac3c2 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Sat, 10 Oct 2026 22:58:33 +0800 Subject: [PATCH 3/3] feat(web): edit an API key's own limits in its dialog MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The backend has long held rate-limit rules and budgets on a key lineage and enforced them, but the console could only set limits on users and roles. The key's edit dialog now has a Limits tab: the key's rules (requests or weighted tokens over 1m … 1w) and budgets (daily, weekly, monthly weighted tokens), each with what it has used, read through the existing `/api/admin/limits/api_key/{id}/rules|budgets|usage` endpoints and added, changed and removed through the existing POST and DELETE ones. The tab reuses the user limits tab's parts rather than copying them: the usage meter, expiry cell and rule label are exported from it, and its override drawer becomes `LimitDrawer`, taking the subject to write to. For a key the drawer defaults to no expiry and, when editing, prefills the row's expiry and reason. In an edit the slot fields (type, gateway, metric, window, period) are locked for both users and keys: the row is stored by that slot, so changing one created a second row instead of editing the first. Permissions are the endpoints' own: the tab shows to holders of `rate_limits:read` and disappears when the endpoints answer 403 (a team-scoped grant never covers the caller's own keys); without `rate_limits:write` it is read-only. Everyone else sees the dialog as before. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 8 + web/src/components/limits/key-limits-query.ts | 82 ++++++ .../components/limits/key-limits-tab.test.tsx | 193 ++++++++++++ web/src/components/limits/key-limits-tab.tsx | 274 ++++++++++++++++++ web/src/components/limits/user-limits-tab.tsx | 214 ++++++++------ web/src/i18n/en.json | 21 +- web/src/i18n/zh.json | 21 +- web/src/routes/api-key-dialogs.tsx | 221 ++++++++------ web/src/routes/api-keys.test.tsx | 81 +++++- 9 files changed, 934 insertions(+), 181 deletions(-) create mode 100644 web/src/components/limits/key-limits-query.ts create mode 100644 web/src/components/limits/key-limits-tab.test.tsx create mode 100644 web/src/components/limits/key-limits-tab.tsx diff --git a/CHANGELOG.md b/CHANGELOG.md index 6366f224..7a127dfe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -73,6 +73,14 @@ target. without limits of its own sees its owner's totals and limits; a key handed to someone else should carry limits of its own. The console's Configuration Guide lists the endpoint with the other gateway endpoints. +- **An API key's own limits are edited in the console.** The key's edit + dialog has a Limits tab: rate limits (requests or weighted tokens over + 1m, 5m, 1h, 5h, 1d or 1w) and budgets (daily, weekly or monthly weighted + tokens), each with what it has used, added, changed and removed through + the existing limits endpoints. They follow the key across rotations. + Reading needs `rate_limits:read` and changing `rate_limits:write`, in a + scope that covers the key, as before; without write access the tab is + read-only, and without read access the dialog is unchanged. ## [3.4.0] — 2026-10-11 diff --git a/web/src/components/limits/key-limits-query.ts b/web/src/components/limits/key-limits-query.ts new file mode 100644 index 00000000..74a04c01 --- /dev/null +++ b/web/src/components/limits/key-limits-query.ts @@ -0,0 +1,82 @@ +// An API key's own rate-limit rules and budgets with what each has used, +// from the limits endpoints with `kind = api_key` (they resolve the key id +// to its lineage, so the limits follow the key across rotations): +// +// GET /api/admin/limits/api_key/{id}/rules +// GET /api/admin/limits/api_key/{id}/budgets +// GET /api/admin/limits/api_key/{id}/usage +// +// All three need `rate_limits:read` in a scope that covers the key. + +import { api } from '@/lib/api'; + +export type KeyLimitSurface = 'ai_gateway' | 'mcp_gateway'; +export type KeyLimitMetric = 'requests' | 'tokens'; +export type KeyLimitPeriod = 'daily' | 'weekly' | 'monthly'; + +export interface KeyRule { + id: string; + surface: KeyLimitSurface; + metric: KeyLimitMetric; + window_secs: number; + max_count: number; + enabled: boolean; + expires_at?: string | null; + reason?: string | null; + /** The window's count now. */ + current: number; +} + +export interface KeyCap { + id: string; + period: KeyLimitPeriod; + limit_tokens: number; + enabled: boolean; + expires_at?: string | null; + reason?: string | null; + /** The period's weighted tokens so far. */ + current: number; +} + +export interface KeyLimits { + rules: KeyRule[]; + caps: KeyCap[]; +} + +interface Items { + items: T[]; +} + +interface UsageResponse { + rules: { rule_id: string; current: number }[]; + caps: { cap_id: string; current: number }[]; +} + +export function keyLimitsQueryKey(keyId: string) { + return ['admin', 'limits', 'api_key', keyId] as const; +} + +export function keyLimitsQuery(keyId: string) { + return { + queryKey: keyLimitsQueryKey(keyId), + queryFn: async ({ signal }: { signal: AbortSignal }): Promise => { + const base = `/api/admin/limits/api_key/${keyId}`; + const [rules, caps, usage] = await Promise.all([ + api>>(`${base}/rules`, { signal }), + api>>(`${base}/budgets`, { signal }), + api(`${base}/usage`, { signal }), + ]); + const ruleUsed = new Map(usage.rules.map((u) => [u.rule_id, u.current])); + const capUsed = new Map(usage.caps.map((u) => [u.cap_id, u.current])); + return { + rules: rules.items.map((r) => ({ ...r, current: ruleUsed.get(r.id) ?? 0 })), + caps: caps.items.map((c) => ({ ...c, current: capUsed.get(c.id) ?? 0 })), + }; + }, + }; +} + +/** The caller holds `rate_limits:read`, but not in a scope covering this key. */ +export function isForbidden(error: unknown): boolean { + return (error as { status?: number } | null)?.status === 403; +} diff --git a/web/src/components/limits/key-limits-tab.test.tsx b/web/src/components/limits/key-limits-tab.test.tsx new file mode 100644 index 00000000..72250d4b --- /dev/null +++ b/web/src/components/limits/key-limits-tab.test.tsx @@ -0,0 +1,193 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest' +import { screen, waitFor, within } from '@testing-library/react' +import userEvent from '@testing-library/user-event' +import { renderWithQueryClient } from '@/test/render' +import { KeyLimitsTab } from './key-limits-tab' + +vi.mock('@/lib/api', () => ({ + api: vi.fn(), + apiPost: vi.fn(), + apiDelete: vi.fn(), +})) + +import { api, apiDelete, apiPost } from '@/lib/api' + +const mockApi = vi.mocked(api) +const mockPost = vi.mocked(apiPost) +const mockDelete = vi.mocked(apiDelete) + +const BASE = '/api/admin/limits/api_key/key-1' + +// One rule and one budget on the key, each with its counter. +function mockKeyLimits({ + rules = [ + { + id: 'r1', + subject_kind: 'api_key_lineage', + subject_id: 'lin-1', + surface: 'ai_gateway', + metric: 'requests', + window_secs: 60, + max_count: 100, + enabled: true, + }, + ], + caps = [ + { + id: 'c1', + subject_kind: 'api_key_lineage', + subject_id: 'lin-1', + period: 'daily', + limit_tokens: 1_000_000, + enabled: true, + }, + ], + usage = { + rules: [{ rule_id: 'r1', current: 40, limit: 100 }], + caps: [{ cap_id: 'c1', current: 250_000, limit: 1_000_000 }], + }, +}: { + rules?: Record[] + caps?: Record[] + usage?: Record +} = {}) { + mockApi.mockImplementation(async (url: string) => { + if (url === `${BASE}/rules`) return { items: rules } + if (url === `${BASE}/budgets`) return { items: caps } + if (url === `${BASE}/usage`) return usage + throw new Error(`unexpected ${url}`) + }) +} + +beforeEach(() => { + vi.clearAllMocks() +}) + +describe('KeyLimitsTab', () => { + it("lists the key's rules and budgets with what each has used", async () => { + mockKeyLimits() + renderWithQueryClient() + + await waitFor(() => { + expect(screen.getByText('40%')).toBeInTheDocument() + }) + // 100 - 40 and 1,000,000 - 250,000 (en-US grouping). + expect(screen.getByText('60 remaining')).toBeInTheDocument() + expect(screen.getByText('750,000 remaining')).toBeInTheDocument() + expect(screen.getByText('25%')).toBeInTheDocument() + expect(screen.getByText('budget · Daily')).toBeInTheDocument() + expect(screen.getByRole('button', { name: /add limit/i })).toBeInTheDocument() + expect(screen.getAllByRole('button', { name: /edit/i })).toHaveLength(2) + }) + + it('says when the key has no limits of its own', async () => { + mockKeyLimits({ rules: [], caps: [], usage: { rules: [], caps: [] } }) + renderWithQueryClient() + + await waitFor(() => { + expect(screen.getByText('This key has no limits of its own.')).toBeInTheDocument() + }) + }) + + it('is read-only without rate_limits:write', async () => { + mockKeyLimits() + renderWithQueryClient() + + await waitFor(() => { + expect(screen.getByText('40%')).toBeInTheDocument() + }) + expect(screen.queryByRole('button', { name: /add limit/i })).not.toBeInTheDocument() + expect(screen.queryByRole('button', { name: /edit/i })).not.toBeInTheDocument() + expect(screen.queryByRole('button', { name: /delete/i })).not.toBeInTheDocument() + }) + + it("adds a rule to the key's lineage, lasting until removed", async () => { + mockKeyLimits({ rules: [], caps: [], usage: { rules: [], caps: [] } }) + mockPost.mockResolvedValueOnce({}) + const user = userEvent.setup() + renderWithQueryClient() + + await user.click(await screen.findByRole('button', { name: /add limit/i })) + await user.type(screen.getByRole('spinbutton'), '500') + await user.click(screen.getByRole('button', { name: 'Save' })) + + await waitFor(() => { + expect(mockPost).toHaveBeenCalledWith(`${BASE}/rules`, { + surface: 'ai_gateway', + metric: 'requests', + window_secs: 3600, + max_count: 500, + enabled: true, + expires_at: null, + reason: null, + }) + }) + }) + + it('edits a rule in its slot, keeping its expiry and reason', async () => { + const expires = '2026-11-01T12:00:00.000Z' + mockKeyLimits({ + rules: [ + { + id: 'r1', + surface: 'ai_gateway', + metric: 'tokens', + window_secs: 18000, + max_count: 50_000, + enabled: true, + expires_at: expires, + reason: 'launch week', + }, + ], + caps: [], + usage: { rules: [{ rule_id: 'r1', current: 0, limit: 50_000 }], caps: [] }, + }) + mockPost.mockResolvedValueOnce({}) + const user = userEvent.setup() + renderWithQueryClient() + + await user.click(await screen.findByRole('button', { name: /edit/i })) + expect(screen.getByText('Edit limit')).toBeInTheDocument() + // Type, gateway, metric and window are the slot the row is stored by, + // so an edit can't change them; the expiry stays open. + const selects = screen.getAllByRole('combobox') + expect(selects).toHaveLength(5) + expect(selects.filter((s) => (s as HTMLButtonElement).disabled)).toHaveLength(4) + const value = screen.getByRole('spinbutton') + await user.clear(value) + await user.type(value, '80000') + await user.click(screen.getByRole('button', { name: 'Save' })) + + await waitFor(() => { + expect(mockPost).toHaveBeenCalledWith(`${BASE}/rules`, { + surface: 'ai_gateway', + metric: 'tokens', + window_secs: 18000, + max_count: 80_000, + enabled: true, + // Prefilled to the minute in local time, sent back as UTC. + expires_at: expires, + reason: 'launch week', + }) + }) + }) + + it('deletes a budget from the key after confirmation', async () => { + mockKeyLimits() + mockDelete.mockResolvedValueOnce(undefined) + const user = userEvent.setup() + renderWithQueryClient() + + await screen.findByText('budget · Daily') + const deletes = screen.getAllByRole('button', { name: /delete/i }) + // Rules come first, then budgets. + await user.click(deletes[1]) + const dialog = await screen.findByRole('dialog') + expect(within(dialog).getByText('Delete this limit?')).toBeInTheDocument() + await user.click(within(dialog).getByRole('button', { name: 'Delete' })) + + await waitFor(() => { + expect(mockDelete).toHaveBeenCalledWith(`${BASE}/budgets/c1`) + }) + }) +}) diff --git a/web/src/components/limits/key-limits-tab.tsx b/web/src/components/limits/key-limits-tab.tsx new file mode 100644 index 00000000..0ce573bf --- /dev/null +++ b/web/src/components/limits/key-limits-tab.tsx @@ -0,0 +1,274 @@ +// ============================================================================ +// API key limits tab — the "Limits" tab inside the API key edit dialog. +// +// The key's own rate-limit rules and budgets, each with what it has used, +// read live from the counters enforcement uses. Requests made with the key +// are held to these on top of its owner's limits (the user page's limits +// tab). The rules and budgets are stored against the key's lineage, so +// they follow the key across rotations. +// +// Reading needs `rate_limits:read`, changing `rate_limits:write`, each in a +// scope that covers the key; without write the tab is read-only. Adding +// and editing reuse the user tab's drawer with `kind = api_key`: +// +// POST /api/admin/limits/api_key/{id}/rules (+ budgets) +// DELETE /api/admin/limits/api_key/{id}/rules/{rule_id} (+ budgets/{cap_id}) +// ============================================================================ + +import { useState, type ReactNode } from 'react'; +import { useTranslation } from 'react-i18next'; +import { useQuery, useQueryClient } from '@tanstack/react-query'; +import { AlertCircle, Pencil, Plus, Trash2 } from 'lucide-react'; +import { toast } from 'sonner'; +import { Button } from '@/components/ui/button'; +import { Badge } from '@/components/ui/badge'; +import { Label } from '@/components/ui/label'; +import { ConfirmDialog } from '@/components/confirm-dialog'; +import { apiDelete } from '@/lib/api'; +import { + ExpiryCell, + LimitDrawer, + RuleScopeLabel, + UsageMeter, + type DrawerInit, +} from './user-limits-tab'; +import { + keyLimitsQuery, + keyLimitsQueryKey, + type KeyCap, + type KeyRule, +} from './key-limits-query'; + +interface KeyLimitsTabProps { + keyId: string; + /** The caller may add, edit and delete (`rate_limits:write`). */ + canEdit: boolean; +} + +export function KeyLimitsTab({ keyId, canEdit }: KeyLimitsTabProps) { + const { t } = useTranslation(); + const queryClient = useQueryClient(); + const query = useQuery(keyLimitsQuery(keyId)); + const [drawer, setDrawer] = useState(null); + const [deleting, setDeleting] = useState< + { kind: 'rule'; id: string } | { kind: 'cap'; id: string } | null + >(null); + + const reload = () => queryClient.invalidateQueries({ queryKey: keyLimitsQueryKey(keyId) }); + + const handleDelete = async () => { + if (!deleting) return; + const path = deleting.kind === 'rule' ? 'rules' : 'budgets'; + try { + await apiDelete(`/api/admin/limits/api_key/${keyId}/${path}/${deleting.id}`); + toast.success(t('keyLimits.deleted')); + setDeleting(null); + await reload(); + } catch (e) { + toast.error(e instanceof Error ? e.message : t('common.operationFailed')); + } + }; + + const data = query.data; + const empty = !!data && data.rules.length === 0 && data.caps.length === 0; + + return ( +
+
+
+ +

{t('keyLimits.hint')}

+
+ {canEdit && ( + + )} +
+ + {query.error && ( +

+ + {query.error.message} +

+ )} + + {query.isPending ? ( +

{t('common.loading')}

+ ) : empty ? ( +

+ {t('keyLimits.none')} +

+ ) : data ? ( +
+ + + + + + + + {canEdit && ( + + )} + + + + {data.rules.map((r) => ( + + } + limit={r.max_count} + current={r.current} + enabled={r.enabled} + expiresAt={r.expires_at} + canEdit={canEdit} + onEdit={() => setDrawer(ruleInit(r))} + onDelete={() => setDeleting({ kind: 'rule', id: r.id })} + /> + ))} + {data.caps.map((c) => ( + setDrawer(capInit(c))} + onDelete={() => setDeleting({ kind: 'cap', id: c.id })} + /> + ))} + +
{t('userLimitsTab.col.scope')}{t('userLimitsTab.col.limit')}{t('userLimitsTab.col.usage')}{t('userLimitsTab.col.expires')} + {t('common.actions')} +
+
+ ) : null} + + {drawer && ( + setDrawer(null)} + onApplied={() => { + setDrawer(null); + reload(); + }} + /> + )} + + !v && setDeleting(null)} + title={t('keyLimits.confirmDeleteTitle')} + description={t('keyLimits.confirmDeleteBody')} + variant="destructive" + confirmLabel={t('common.delete')} + onConfirm={handleDelete} + /> +
+ ); +} + +function ruleInit(r: KeyRule): DrawerInit { + return { + kind: 'rule', + surface: r.surface, + metric: r.metric, + window_secs: r.window_secs, + current_value: r.max_count, + editing_override_id: r.id, + expires_at: r.expires_at ?? null, + reason: r.reason ?? null, + }; +} + +function capInit(c: KeyCap): DrawerInit { + return { + kind: 'cap', + period: c.period, + current_value: c.limit_tokens, + editing_override_id: c.id, + expires_at: c.expires_at ?? null, + reason: c.reason ?? null, + }; +} + +function LimitRow({ + scope, + limit, + current, + enabled, + expiresAt, + canEdit, + onEdit, + onDelete, +}: { + scope: ReactNode; + limit: number; + current: number; + enabled: boolean; + expiresAt?: string | null; + canEdit: boolean; + onEdit: () => void; + onDelete: () => void; +}) { + const { t } = useTranslation(); + return ( + + + {scope} + {!enabled && ( + + {t('keyLimits.disabled')} + + )} + + {limit.toLocaleString()} + + + + + + + {canEdit && ( + +
+ + +
+ + )} + + ); +} diff --git a/web/src/components/limits/user-limits-tab.tsx b/web/src/components/limits/user-limits-tab.tsx index 4aeaf14c..3dd08344 100644 --- a/web/src/components/limits/user-limits-tab.tsx +++ b/web/src/components/limits/user-limits-tab.tsx @@ -29,6 +29,11 @@ // POST /api/admin/users/{id}/limits/reset // POST /api/admin/limits/user/{id}/rules (+ budgets) — single // POST /api/admin/limits/bulk/{rules|budgets}/{disable|delete} — bulk +// +// The usage meter, expiry cell, rule label and the create / edit drawer +// are exported for the API key's limits tab (`key-limits-tab.tsx`), which +// edits a key's own rules and budgets through the same endpoints with +// `kind = api_key`. // ============================================================================ import { useMemo, useState } from 'react'; @@ -160,7 +165,7 @@ interface UserLimitsTabProps { userId: string; } -interface DrawerInit { +export interface DrawerInit { kind: 'rule' | 'cap'; surface?: Surface; metric?: Metric; @@ -171,8 +176,14 @@ interface DrawerInit { editing_override_id?: string; // Non-null when overriding a role default (for copy in create form) role_default?: number; + // Prefill when editing a key's limit: its expiry (null = none) and reason. + expires_at?: string | null; + reason?: string | null; } +/// Whose limits the drawer writes: `POST /api/admin/limits/{kind}/{id}/…`. +export type LimitSubject = { kind: 'user' | 'api_key'; id: string }; + export function UserLimitsTab({ userId }: UserLimitsTabProps) { const { t } = useTranslation(); const queryClient = useQueryClient(); @@ -392,8 +403,8 @@ export function UserLimitsTab({ userId }: UserLimitsTabProps) { {/* ------------------ Drawers / confirms ------------------ */} {drawer && ( - setDrawer(null)} onApplied={() => { @@ -586,18 +597,8 @@ function RuleRow({ onEdit: () => void; onReset: () => void; }) { - const { t, i18n } = useTranslation(); - const locale = i18n.language === 'zh' ? 'zh-CN' : 'en-US'; - const pct = row.max_count > 0 ? Math.min(100, (row.current / row.max_count) * 100) : 0; - const toneClass = - pct >= 100 - ? '[&>[data-slot=progress-indicator]]:bg-destructive' - : pct >= 80 - ? '[&>[data-slot=progress-indicator]]:bg-yellow-500' - : ''; + const { t } = useTranslation(); const deltaChip = deltaLabel(row.role_default_max_count, row.max_count); - const overCap = row.max_count > 0 && row.current >= row.max_count; - const remaining = row.max_count - row.current; return ( @@ -611,9 +612,7 @@ function RuleRow({ - {t(`limits.surfaceShort_${row.surface}` as const)} ·{' '} - {t(`limits.metric_${row.metric}` as const)} ·{' '} - {secsLabel(row.window_secs, t)} + {row.max_count.toLocaleString()} @@ -622,24 +621,7 @@ function RuleRow({ )} -
- - {row.current.toLocaleString()} - - - - {pct.toFixed(0)}% - - {overCap ? ( - - {t('limits.exceeded')} - - ) : remaining > 0 ? ( - - {t('limits.remainingLabel', { count: remaining.toLocaleString(locale) })} - - ) : null} -
+ @@ -671,21 +653,8 @@ function CapRow({ onEdit: () => void; onReset: () => void; }) { - const { t, i18n } = useTranslation(); - const locale = i18n.language === 'zh' ? 'zh-CN' : 'en-US'; - const pct = - row.limit_tokens > 0 - ? Math.min(100, (row.current / row.limit_tokens) * 100) - : 0; - const toneClass = - pct >= 100 - ? '[&>[data-slot=progress-indicator]]:bg-destructive' - : pct >= 80 - ? '[&>[data-slot=progress-indicator]]:bg-yellow-500' - : ''; + const { t } = useTranslation(); const deltaChip = deltaLabel(row.role_default_limit_tokens, row.limit_tokens); - const overCap = row.limit_tokens > 0 && row.current >= row.limit_tokens; - const remaining = row.limit_tokens - row.current; return ( @@ -710,24 +679,7 @@ function CapRow({ )} -
- - {row.current.toLocaleString()} - - - - {pct.toFixed(0)}% - - {overCap ? ( - - {t('limits.exceeded')} - - ) : remaining > 0 ? ( - - {t('limits.remainingLabel', { count: remaining.toLocaleString(locale) })} - - ) : null} -
+ @@ -744,6 +696,57 @@ function CapRow({ ); } +/// What a counter holds against its limit: the count, a bar, the share, +/// and what is left or that it is exceeded. +export function UsageMeter({ current, limit }: { current: number; limit: number }) { + const { t, i18n } = useTranslation(); + const locale = i18n.language === 'zh' ? 'zh-CN' : 'en-US'; + const pct = limit > 0 ? Math.min(100, (current / limit) * 100) : 0; + const toneClass = + pct >= 100 + ? '[&>[data-slot=progress-indicator]]:bg-destructive' + : pct >= 80 + ? '[&>[data-slot=progress-indicator]]:bg-yellow-500' + : ''; + const overCap = limit > 0 && current >= limit; + const remaining = limit - current; + return ( +
+ {current.toLocaleString()} + + {pct.toFixed(0)}% + {overCap ? ( + + {t('limits.exceeded')} + + ) : remaining > 0 ? ( + + {t('limits.remainingLabel', { count: remaining.toLocaleString(locale) })} + + ) : null} +
+ ); +} + +/// A rate-limit rule's slot: gateway · metric · window. +export function RuleScopeLabel({ + surface, + metric, + windowSecs, +}: { + surface: Surface; + metric: Metric; + windowSecs: number; +}) { + const { t } = useTranslation(); + return ( + <> + {t(`limits.surfaceShort_${surface}` as const)} · {t(`limits.metric_${metric}` as const)} ·{' '} + {secsLabel(windowSecs, t)} + + ); +} + function RowActions({ source, onOverride, @@ -789,7 +792,7 @@ function RowActions({ ); } -function ExpiryCell({ at }: { at?: string | null }) { +export function ExpiryCell({ at }: { at?: string | null }) { const { t } = useTranslation(); // Label reads in hours and days, so a minute is a fine resolution. const now = useNow(60_000); @@ -923,21 +926,33 @@ function summarizeDetail(d: Record): string { } // ---------------------------------------------------------------------------- -// Create / edit override drawer +// Create / edit drawer — a user's overrides, or an API key's own limits // ---------------------------------------------------------------------------- -function OverrideDrawer({ - userId, +/// `datetime-local` value (local time, minutes) for an ISO instant. +function toLocalInput(iso: string): string { + const d = new Date(iso); + const pad = (n: number) => String(n).padStart(2, '0'); + return `${d.getFullYear()}-${pad(d.getMonth() + 1)}-${pad(d.getDate())}T${pad(d.getHours())}:${pad(d.getMinutes())}`; +} + +export function LimitDrawer({ + subject, init, onClose, onApplied, }: { - userId: string; + subject: LimitSubject; init: DrawerInit; onClose: () => void; onApplied: () => void; }) { const { t } = useTranslation(); + const forKey = subject.kind === 'api_key'; + // The row a rule or budget is stored as is keyed by its slot — gateway, + // metric and window, or period — so an edit keeps the slot and changes + // the rest. + const editing = !!init.editing_override_id; const [kind, setKind] = useState<'rule' | 'cap'>(init.kind); const [surface, setSurface] = useState(init.surface ?? 'ai_gateway'); const [metric, setMetric] = useState(init.metric ?? 'requests'); @@ -946,12 +961,18 @@ function OverrideDrawer({ const [value, setValue] = useState( init.current_value !== undefined ? String(init.current_value) : '', ); - const [expiryPreset, setExpiryPreset] = useState('7d'); - const [customExpiry, setCustomExpiry] = useState(''); - const [reason, setReason] = useState(''); + // A user override is temporary by default; a key's limit lasts. + const [expiryPreset, setExpiryPreset] = useState( + init.expires_at ? 'custom' : forKey ? 'permanent' : '7d', + ); + const [customExpiry, setCustomExpiry] = useState( + init.expires_at ? toLocalInput(init.expires_at) : '', + ); + const [reason, setReason] = useState(init.reason ?? ''); const [busy, setBusy] = useState(false); const roleDefault = init.role_default; + const base = `/api/admin/limits/${subject.kind}/${subject.id}`; const resolveExpiry = (): string | null | 'invalid' => { if (expiryPreset === 'permanent') return null; @@ -980,7 +1001,7 @@ function OverrideDrawer({ setBusy(true); try { if (kind === 'rule') { - await apiPost(`/api/admin/limits/user/${userId}/rules`, { + await apiPost(`${base}/rules`, { surface, metric, window_secs: windowSecs, @@ -990,7 +1011,7 @@ function OverrideDrawer({ reason: reason.trim() || null, }); } else { - await apiPost(`/api/admin/limits/user/${userId}/budgets`, { + await apiPost(`${base}/budgets`, { period, limit_tokens: n, enabled: true, @@ -998,7 +1019,7 @@ function OverrideDrawer({ reason: reason.trim() || null, }); } - toast.success(t('userLimitOverrides.added')); + toast.success(forKey ? t('keyLimits.saved') : t('userLimitOverrides.added')); onApplied(); } catch (e) { toast.error(e instanceof Error ? e.message : t('common.operationFailed')); @@ -1021,9 +1042,13 @@ function OverrideDrawer({ >

- {init.editing_override_id - ? t('userLimitsTab.drawer.editTitle') - : t('userLimitsTab.drawer.createTitle')} + {forKey + ? editing + ? t('keyLimits.drawer.editTitle') + : t('keyLimits.drawer.createTitle') + : editing + ? t('userLimitsTab.drawer.editTitle') + : t('userLimitsTab.drawer.createTitle')}

{roleDefault != null && (

@@ -1036,7 +1061,11 @@ function OverrideDrawer({

- setKind(v as 'rule' | 'cap')} + disabled={editing} + > @@ -1053,7 +1082,11 @@ function OverrideDrawer({ - setSurface(v as Surface)} + disabled={editing} + > @@ -1074,7 +1107,7 @@ function OverrideDrawer({ setWindowSecs(Number(v))} + disabled={editing} > @@ -1113,7 +1147,11 @@ function OverrideDrawer({ - setPeriod(v as Period)} + disabled={editing} + > @@ -1195,7 +1233,11 @@ function OverrideDrawer({ {t('common.cancel')}
diff --git a/web/src/i18n/en.json b/web/src/i18n/en.json index bbeda4b7..c982c527 100644 --- a/web/src/i18n/en.json +++ b/web/src/i18n/en.json @@ -704,7 +704,11 @@ "expired": "Expired", "rotated": "Rotated", "gracePeriod": "Grace Period", - "inactiveTab": "Inactive" + "inactiveTab": "Inactive", + "tab": { + "settings": "Settings", + "limits": "Limits" + } }, "models": { "upstreamRefuses": "upstream refuses", @@ -1283,6 +1287,21 @@ "reason": "Reason" } }, + "keyLimits": { + "title": "Limits of this key", + "hint": "Requests made with this key are held to these limits as well as to its owner's.", + "none": "This key has no limits of its own.", + "add": "Add limit", + "saved": "Limit saved", + "deleted": "Limit deleted", + "disabled": "Disabled", + "confirmDeleteTitle": "Delete this limit?", + "confirmDeleteBody": "Requests made with this key are no longer held to it.", + "drawer": { + "createTitle": "Add limit", + "editTitle": "Edit limit" + } + }, "userLimitsTab": { "effectiveTitle": "Effective limits", "effectiveHint": "What's actually enforced for this user right now — role defaults merged with any active per-user overrides.", diff --git a/web/src/i18n/zh.json b/web/src/i18n/zh.json index 0b1f25c4..9abfec52 100644 --- a/web/src/i18n/zh.json +++ b/web/src/i18n/zh.json @@ -704,7 +704,11 @@ "expired": "已过期", "rotated": "已轮换", "gracePeriod": "过渡期", - "inactiveTab": "已失效" + "inactiveTab": "已失效", + "tab": { + "settings": "设置", + "limits": "限额" + } }, "models": { "upstreamRefuses": "上游不提供", @@ -1283,6 +1287,21 @@ "reason": "说明" } }, + "keyLimits": { + "title": "此密钥的限额", + "hint": "使用此密钥的请求同时受这些限额和所属用户的限额约束。", + "none": "此密钥没有单独的限额。", + "add": "添加限额", + "saved": "限额已保存", + "deleted": "限额已删除", + "disabled": "已停用", + "confirmDeleteTitle": "删除此限额？", + "confirmDeleteBody": "使用此密钥的请求将不再受此限额约束。", + "drawer": { + "createTitle": "添加限额", + "editTitle": "编辑限额" + } + }, "userLimitsTab": { "effectiveTitle": "生效额度", "effectiveHint": "当前真实生效的额度 —— 角色合并后再叠加所有生效的用户覆盖。", diff --git a/web/src/routes/api-key-dialogs.tsx b/web/src/routes/api-key-dialogs.tsx index 72fc9672..ca063d3f 100644 --- a/web/src/routes/api-key-dialogs.tsx +++ b/web/src/routes/api-key-dialogs.tsx @@ -1,5 +1,6 @@ import { useEffect, useState, type FormEvent } from 'react'; import { useTranslation } from 'react-i18next'; +import { useQuery } from '@tanstack/react-query'; import { toast } from 'sonner'; import { Button } from '@/components/ui/button'; import { Input } from '@/components/ui/input'; @@ -14,9 +15,12 @@ import { DialogFooter, } from '@/components/ui/dialog'; import { Alert, AlertDescription } from '@/components/ui/alert'; +import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs'; import { Copy, Check, AlertCircle } from 'lucide-react'; -import { api, apiPost, apiPatch, apiDelete } from '@/lib/api'; +import { api, apiPost, apiPatch, apiDelete, hasPermission } from '@/lib/api'; import { ConfirmDialog } from '@/components/confirm-dialog'; +import { KeyLimitsTab } from '@/components/limits/key-limits-tab'; +import { isForbidden, keyLimitsQuery } from '@/components/limits/key-limits-query'; import { ScopeDropdown, ToolScopeDropdown, @@ -347,6 +351,18 @@ export function EditApiKeyDialog({ const [editMcpOverrides, setEditMcpOverrides] = useState>({}); const [editSubmitting, setEditSubmitting] = useState(false); const [editError, setEditError] = useState(''); + const [tab, setTab] = useState<'settings' | 'limits'>('settings'); + + // The key's own limits, for whoever the limits endpoints let read them: + // `rate_limits:read` in a scope that covers the key. Anyone else gets + // the settings form alone, as before. + const canReadLimits = hasPermission('rate_limits:read'); + const limitsQuery = useQuery({ + ...keyLimitsQuery(apiKey?.id ?? ''), + enabled: open && !!apiKey && canReadLimits, + }); + const showLimits = canReadLimits && !!apiKey && !isForbidden(limitsQuery.error); + const activeTab = showLimits ? tab : 'settings'; // Sync local state when the dialog opens with a new key const [lastKeyId, setLastKeyId] = useState(null); @@ -369,6 +385,7 @@ export function EditApiKeyDialog({ setEditCostCenter(apiKey.cost_center ?? ''); setEditMcpOverrides(apiKey.mcp_account_overrides ?? {}); setEditError(''); + setTab('settings'); } // Reset tracking when dialog closes @@ -428,6 +445,98 @@ export function EditApiKeyDialog({ } }; + const settingsForm = ( +
+ {editError && ( + + + {editError} + + )} +
+ +
+ {ALL_SURFACES.map((s) => { + const checked = editSurfaces.includes(s); + return ( +
+ + {checked && s === 'ai_gateway' && ( +
+ +
+ )} + {checked && s === 'mcp_gateway' && ( +
+ +
+ )} +
+ ); + })} +
+
+
+ + setEditExpiresInDays(e.target.value)} placeholder="90" min={1} /> +

{t('apiKeys.expiresInHint')}

+
+
+ + setEditRotationPeriod(e.target.value)} placeholder="0" min={0} /> +
+
+ + setEditInactivityTimeout(e.target.value)} placeholder="0" min={0} /> +
+
+ + setEditCostCenter(e.target.value)} + placeholder={t('apiKeys.costCenterPlaceholder')} + maxLength={64} + /> + + {costCenterOptions.map((opt) => ( + +

{t('apiKeys.costCenterHint')}

+
+ + + + + + ); + return ( @@ -435,95 +544,27 @@ export function EditApiKeyDialog({ {t('apiKeys.editKey')} {apiKey?.name ?? ''} -
- {editError && ( - - - {editError} - - )} -
- -
- {ALL_SURFACES.map((s) => { - const checked = editSurfaces.includes(s); - return ( -
- - {checked && s === 'ai_gateway' && ( -
- -
- )} - {checked && s === 'mcp_gateway' && ( -
- -
- )} -
- ); - })} -
-
-
- - setEditExpiresInDays(e.target.value)} placeholder="90" min={1} /> -

{t('apiKeys.expiresInHint')}

-
-
- - setEditRotationPeriod(e.target.value)} placeholder="0" min={0} /> -
-
- - setEditInactivityTimeout(e.target.value)} placeholder="0" min={0} /> -
-
- - setEditCostCenter(e.target.value)} - placeholder={t('apiKeys.costCenterPlaceholder')} - maxLength={64} - /> - - {costCenterOptions.map((opt) => ( - -

{t('apiKeys.costCenterHint')}

-
- - - - - + {showLimits ? ( + setTab(v as 'settings' | 'limits')}> + + {t('apiKeys.tab.settings')} + {t('apiKeys.tab.limits')} + + + {settingsForm} + + + {apiKey && ( + + )} + + + ) : ( + settingsForm + )}
); diff --git a/web/src/routes/api-keys.test.tsx b/web/src/routes/api-keys.test.tsx index caf3dd37..9819e3fd 100644 --- a/web/src/routes/api-keys.test.tsx +++ b/web/src/routes/api-keys.test.tsx @@ -1,5 +1,5 @@ import { describe, it, expect, vi, beforeEach } from 'vitest' -import { screen, waitFor } from '@testing-library/react' +import { screen, waitFor, within } from '@testing-library/react' import userEvent from '@testing-library/user-event' import { renderWithQueryClient } from '@/test/render' import { ApiKeysPage } from './api-keys' @@ -22,9 +22,10 @@ vi.mock('@/lib/paginated-fetch', () => ({ fetchAllPaginated: vi.fn().mockResolvedValue([]), })) -import { api } from '@/lib/api' +import { api, hasPermission } from '@/lib/api' const mockApi = vi.mocked(api) +const mockHasPermission = vi.mocked(hasPermission) const makeKey = (overrides: Record = {}) => ({ id: 'key-1', @@ -52,17 +53,22 @@ const makeKey = (overrides: Record = {}) => ({ // Different endpoints have different response shapes — route by URL so // /api/keys returns the paginated key list while /api/keys/cost-centers // and /api/keys/policy-scope get their own appropriate shapes. -function mockKeysFetch(keys: ReturnType[]) { +function mockKeysFetch( + keys: ReturnType[], + limits: (url: string) => unknown = () => ({ items: [], rules: [], caps: [] }), +) { mockApi.mockImplementation(async (url: string) => { if (url.startsWith('/api/keys/cost-centers')) return [] as string[] if (url.startsWith('/api/keys/policy-scope')) return { allowed_models: null, allowed_mcp_tools: null } + if (url.startsWith('/api/admin/limits/')) return limits(url) return { data: keys, total: keys.length, page: 1, page_size: 20 } }) } beforeEach(() => { vi.clearAllMocks() + mockHasPermission.mockImplementation(() => true) }) describe('ApiKeysPage', () => { @@ -143,4 +149,73 @@ describe('ApiKeysPage', () => { expect(screen.getByText('Generate a new API key for the gateway.')).toBeInTheDocument() }) }) + + it("edits the key's own limits in a tab of the edit dialog", async () => { + mockKeysFetch([makeKey()], (url) => + url.endsWith('/usage') + ? { rules: [{ rule_id: 'r1', current: 3, limit: 10 }], caps: [] } + : url.endsWith('/rules') + ? { + items: [ + { + id: 'r1', + surface: 'ai_gateway', + metric: 'requests', + window_secs: 60, + max_count: 10, + enabled: true, + }, + ], + } + : { items: [] }, + ) + const user = userEvent.setup() + renderWithQueryClient() + + await user.click(await screen.findByRole('button', { name: 'Edit' })) + await user.click(await screen.findByRole('tab', { name: 'Limits' })) + + expect(await screen.findByText('30%')).toBeInTheDocument() + expect(screen.getByText('7 remaining')).toBeInTheDocument() + expect(mockApi).toHaveBeenCalledWith( + '/api/admin/limits/api_key/key-1/rules', + expect.anything(), + ) + }) + + it('shows no limits tab without rate_limits:read', async () => { + mockHasPermission.mockImplementation((perm) => perm !== 'rate_limits:read') + mockKeysFetch([makeKey()]) + const user = userEvent.setup() + renderWithQueryClient() + + await user.click(await screen.findByRole('button', { name: 'Edit' })) + expect(await screen.findByText('Edit Key')).toBeInTheDocument() + expect(screen.queryByRole('tab', { name: 'Limits' })).not.toBeInTheDocument() + expect(mockApi).not.toHaveBeenCalledWith( + expect.stringContaining('/api/admin/limits/'), + expect.anything(), + ) + }) + + it("shows no limits tab when the key is outside the caller's scope", async () => { + mockKeysFetch([makeKey()], () => { + throw Object.assign(new Error('Forbidden'), { status: 403 }) + }) + const user = userEvent.setup() + renderWithQueryClient() + + await user.click(await screen.findByRole('button', { name: 'Edit' })) + await waitFor(() => { + expect(mockApi).toHaveBeenCalledWith( + '/api/admin/limits/api_key/key-1/rules', + expect.anything(), + ) + }) + await waitFor(() => { + expect(screen.queryByRole('tab', { name: 'Limits' })).not.toBeInTheDocument() + }) + // The settings form is there as before. + expect(within(screen.getByRole('dialog')).getByText('Gateway access')).toBeInTheDocument() + }) })