diff --git a/CHANGELOG.md b/CHANGELOG.md index 25bae8c4..7e731190 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -54,6 +54,74 @@ tokens. The legacy monthly token quota is gone. tokens the cached answer records, weighted like any answer (all of its input as plain input). The cost a cache hit reports stays 0. Callers that lean on the cache reach their token limits and budgets sooner. +- **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 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": "key", "kind": "requests", "window": "1m", "window_secs": 60, + "limit": 60, "used": 4, "resets_at": null} + ], + "expires_at": "2026-12-31T00:00:00Z" + } + ``` + + `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. +- **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. ### Fixed @@ -83,6 +151,7 @@ tokens. The legacy monthly token quota is gone. - The `security.rate_limit_fail_closed` hint in Settings names everything the setting now covers. + ## [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 eec124e1..ebb2cb05 100644 --- a/crates/common/src/limits/budget.rs +++ b/crates/common/src/limits/budget.rs @@ -190,24 +190,52 @@ 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?; - 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, - }); + out.push(cap_status(cap, spend(redis, cap, now).await?)); } 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/mod.rs b/crates/common/src/limits/mod.rs index ecbee959..3041cd16 100644 --- a/crates/common/src/limits/mod.rs +++ b/crates/common/src/limits/mod.rs @@ -61,6 +61,7 @@ use uuid::Uuid; pub mod budget; pub mod sliding; +pub mod usage; pub mod weight; // ---------------------------------------------------------------------------- @@ -1162,6 +1163,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 `usage` counters, and its owner's, + /// count the request. + pub key_lineage: Option, pub rules: Vec, pub caps: Vec, } @@ -1181,6 +1186,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 803b80bc..af6302b5 100644 --- a/crates/common/src/limits/sliding.rs +++ b/crates/common/src/limits/sliding.rs @@ -503,9 +503,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 @@ -513,6 +512,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/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/lib.rs b/crates/gateway/src/lib.rs index f1f5ab03..65cf3470 100644 --- a/crates/gateway/src/lib.rs +++ b/crates/gateway/src/lib.rs @@ -12,7 +12,6 @@ pub mod metrics_labels; pub mod model_mapping; pub mod protocol; pub mod proxy; -pub mod rate_limiter; pub mod redaction; pub mod route_caps; pub mod router; diff --git a/crates/gateway/src/proxy/accounting.rs b/crates/gateway/src/proxy/accounting.rs index b36e7201..c25041a0 100644 --- a/crates/gateway/src/proxy/accounting.rs +++ b/crates/gateway/src/proxy/accounting.rs @@ -48,6 +48,10 @@ pub(crate) async fn post_flight_account( return; } + // 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 // window this request overshoots refuses the next one. Post-flight @@ -120,3 +124,28 @@ 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}"); + } +} 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 f5cfe5b5..62567398 100644 --- a/crates/gateway/src/proxy/mod.rs +++ b/crates/gateway/src/proxy/mod.rs @@ -14,7 +14,6 @@ use crate::cost_tracker::CostTracker; use crate::error::GatewayError; use crate::health::HealthTracker; use crate::model_mapping::ModelMapper; -use crate::rate_limiter::RateLimiter; use crate::router::ModelRouter; use think_watch_common::dynamic_config::DynamicConfig; use think_watch_common::limits::SurfaceConstraints; @@ -47,6 +46,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; @@ -61,7 +61,6 @@ pub struct GatewayState { pub guards: Arc>, pub cache: Arc, pub cost_tracker: Arc, - pub rate_limiter: Arc, /// PG pool — used to query enabled rate-limit rules and budget caps /// per request. Cached above the proxy via `WeightCache` for the /// model weights; raw rules go through a separate cache later. @@ -129,6 +128,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 b4119a2c..eea8871c 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, usage}; use tw_dialect::ir::Dialect; /// Pre-flight result threaded through to `ChatPostInvokeDeps` later @@ -82,6 +82,16 @@ pub(super) async fn run_preflight_stages( ) .await .map_err(short_circuit_to_response)?; + // 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) = + usage::record_request(&state.redis, limits.owner, lineage, chrono::Utc::now()).await + { + metrics::counter!("gateway_usage_count_fail_open_total").increment(1); + tracing::warn!("usage request count failed: {e}"); + } Ok(Preflight { limits }) } diff --git a/crates/gateway/src/rate_limiter.rs b/crates/gateway/src/rate_limiter.rs deleted file mode 100644 index d43d0a80..00000000 --- a/crates/gateway/src/rate_limiter.rs +++ /dev/null @@ -1,182 +0,0 @@ -use fred::clients::Client; -use fred::interfaces::LuaInterface; -use think_watch_common::errors::AppError; - -/// Sliding-window rate limiter backed by Redis sorted sets + Lua scripts. -/// -/// All check-and-record operations are atomic via EVAL to prevent race conditions. -#[derive(Clone)] -pub struct RateLimiter { - redis: Client, -} - -/// Atomic RPM check: trim window → count → conditionally record → set TTL. -/// Returns {allowed (0/1), current_count, ttl_ms}. -const LUA_RPM_CHECK: &str = r#" -local key = KEYS[1] -local window_start = tonumber(ARGV[1]) -local now_ms = tonumber(ARGV[2]) -local member = ARGV[3] -local limit = tonumber(ARGV[4]) - -redis.call('ZREMRANGEBYSCORE', key, '-inf', window_start) -local count = redis.call('ZCARD', key) -if count >= limit then - return {0, count, 60000} -end -redis.call('ZADD', key, now_ms, member) -redis.call('EXPIRE', key, 120) -return {1, count + 1, 60000} -"#; - -/// Atomic combined RPM + TPM check. Evaluates BOTH limits and only -/// records (in both sets) if BOTH pass. The split-call shape -/// (`RPM_CHECK` then `TPM_CHECK`) had a bug: RPM recorded its hit -/// *before* TPM ran, so a TPM-blocked 429 still burned the user's -/// RPM budget. Returns `{allowed, rpm_count, denied_by}` where -/// `denied_by` is `0` (allowed), `1` (rpm), or `2` (tpm). -const LUA_COMBINED_CHECK: &str = r#" -local rpm_key = KEYS[1] -local tpm_key = KEYS[2] -local window_start = tonumber(ARGV[1]) -local now_ms = tonumber(ARGV[2]) -local rpm_member = ARGV[3] -local rpm_limit = tonumber(ARGV[4]) -local tpm_member = ARGV[5] -local tokens = tonumber(ARGV[6]) -local tpm_limit = tonumber(ARGV[7]) - -redis.call('ZREMRANGEBYSCORE', rpm_key, '-inf', window_start) -local rpm_count = redis.call('ZCARD', rpm_key) -if rpm_count >= rpm_limit then - return {0, rpm_count, 1} -end - -redis.call('ZREMRANGEBYSCORE', tpm_key, '-inf', window_start) -local members = redis.call('ZRANGEBYSCORE', tpm_key, window_start, '+inf') -local current_tokens = 0 -for _, m in ipairs(members) do - local t = m:match(':(%d+)$') - if t then current_tokens = current_tokens + tonumber(t) end -end -if current_tokens + tokens > tpm_limit then - return {0, rpm_count, 2} -end - -redis.call('ZADD', rpm_key, now_ms, rpm_member) -redis.call('EXPIRE', rpm_key, 120) -redis.call('ZADD', tpm_key, now_ms, tpm_member) -redis.call('EXPIRE', tpm_key, 120) -return {1, rpm_count + 1, 0} -"#; - -/// Rate limit check result with metadata for response headers. -#[derive(Debug, Clone)] -pub struct RateLimitInfo { - pub limit: u32, - pub remaining: u32, - pub reset_at: i64, // Unix timestamp -} - -impl RateLimiter { - pub fn new(redis: Client) -> Self { - Self { redis } - } - - /// Check (and atomically record) a request against the sliding window limits. - /// Returns `RateLimitInfo` on success for setting response headers. - pub async fn check_rate_limit( - &self, - key: &str, - rpm_limit: u32, - tpm_limit: Option, - estimated_tokens: Option, - ) -> Result { - let now_ms = chrono::Utc::now().timestamp_millis() as f64; - let window_start = now_ms - 60_000.0; - let member_id = uuid::Uuid::new_v4().to_string(); - // `{}` is a Redis Cluster hash tag: the combined script - // declares both keys, which a cluster runs only in one slot. - let rpm_key = format!("ratelimit:rpm:{{{key}}}"); - let reset_at = chrono::Utc::now().timestamp() + 60; // window resets in ~60s - - // When BOTH RPM and TPM are configured, evaluate them atomically - // in one Lua call. Previously RPM was checked + recorded BEFORE - // TPM ran; if TPM then rejected, the user's RPM slot was already - // burned on a 429'd request. The combined script only records - // when both pass. - if let (Some(tpm_limit), Some(tokens)) = (tpm_limit, estimated_tokens) - && tokens > 0 - { - let tpm_key = format!("ratelimit:tpm:{{{key}}}"); - let member_with_tokens = format!("{member_id}:{tokens}"); - let result: Vec = self - .redis - .eval( - LUA_COMBINED_CHECK, - vec![rpm_key.as_str(), tpm_key.as_str()], - vec![ - window_start.to_string(), - now_ms.to_string(), - member_id, - rpm_limit.to_string(), - member_with_tokens, - tokens.to_string(), - tpm_limit.to_string(), - ], - ) - .await - .map_err(|e| { - tracing::warn!("Combined rate limit check failed: {e}"); - AppError::Internal(anyhow::anyhow!("Rate limit check failed")) - })?; - let allowed = result.first().copied().unwrap_or(1); - let current = result.get(1).copied().unwrap_or(0) as u32; - let denied_by = result.get(2).copied().unwrap_or(0); - if allowed == 0 { - let label = if denied_by == 2 { "tpm" } else { "rpm" }; - metrics::counter!("gateway_rate_limited_total", "type" => label).increment(1); - return Err(AppError::RateLimited); - } - return Ok(RateLimitInfo { - limit: rpm_limit, - remaining: rpm_limit.saturating_sub(current), - reset_at, - }); - } - - // RPM-only path (no TPM rule). Single Lua call records the hit - // when allowed. - let result: Vec = self - .redis - .eval( - LUA_RPM_CHECK, - vec![rpm_key.as_str()], - vec![ - window_start.to_string(), - now_ms.to_string(), - member_id, - rpm_limit.to_string(), - ], - ) - .await - .map_err(|e| { - tracing::warn!("Rate limit RPM check failed: {e}"); - AppError::Internal(anyhow::anyhow!("Rate limit check failed")) - })?; - - let allowed = result.first().copied().unwrap_or(1); - let current = result.get(1).copied().unwrap_or(0) as u32; - - if allowed == 0 { - metrics::counter!("gateway_rate_limited_total", "type" => "rpm").increment(1); - return Err(AppError::RateLimited); - } - - Ok(RateLimitInfo { - limit: rpm_limit, - remaining: rpm_limit.saturating_sub(current), - reset_at, - }) - } -} diff --git a/crates/server/src/app.rs b/crates/server/src/app.rs index 08343e71..6358fb7c 100644 --- a/crates/server/src/app.rs +++ b/crates/server/src/app.rs @@ -249,9 +249,6 @@ pub async fn create_gateway_app(_config: &AppConfig, state: AppState) -> anyhow: // this the proxy would PUT into one configuration and the // viewer would try to GET from another. blob_store: state.blob_store.clone(), - rate_limiter: Arc::new(think_watch_gateway::rate_limiter::RateLimiter::new( - state.redis.clone(), - )), db: state.db.clone(), redis: state.redis.clone(), weight_cache, @@ -290,6 +287,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 @@ -398,6 +406,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..9964a05e --- /dev/null +++ b/crates/server/src/handlers/key_usage.rs @@ -0,0 +1,574 @@ +// ============================================================================ +// GET /v1/usage (gateway port) +// +// 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`), 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, + secs_to_window, sliding, + usage::{self, UsageCounts, UsageSubject}, +}; +use think_watch_gateway::proxy::{GatewayRequestIdentity, limits_for_ai_gateway}; + +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>, +} + +#[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 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, +} + +/// One limit that binds the key's requests. +#[derive(Debug, Serialize, PartialEq)] +pub struct LimitUsage { + /// 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 + /// 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 (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 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 { + 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 = cost_this_month(&state, limits.owner, subject, now).await; + + Json(response( + scope, + &limits, + &rule_used, + &cap_used, + counts, + cost, + identity.key_expires_at, + now, + )) + .into_response() +} + +/// 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`), 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: UsageCounts, + 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 { + scope, + 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: rule_scope(rule), + 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: cap_scope(cap), + 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, 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(|| 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 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, + } + 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 (column, id) = match subject { + UsageSubject::ApiKeyLineage(lineage) => ("api_key_lineage_id", lineage), + UsageSubject::User => ("user_id", owner), + }; + let row = ch + .query(&format!( + "SELECT sum(ifNull(cost_usd, 0)) AS cost FROM gateway_logs \ + PREWHERE created_at >= parseDateTimeBestEffort(?) \ + AND {column} = ?" + )) + .bind(month_start.format("%Y-%m-%d %H:%M:%S").to_string()) + .bind(id.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() -> UsageCounts { + UsageCounts { + requests_today: 3, + tokens_today: 300, + requests_month: 40, + tokens_month: 4_000, + } + } + + /// 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(), + )); + assert_eq!(scope, Scope::User); + let body = response( + scope, + &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}, + "limits": [], + "expires_at": null, + }) + ); + } + + #[test] + 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)], + &[(BudgetPeriod::Daily, 100_000)], + )), + mcp_gateway: None, + }; + let (scope, limits) = scoped(limits_for(&owners(), &key)); + assert_eq!(scope, Scope::Key); + let body = response( + scope, + &limits, + &[45_000], + &[100_500], + counts(), + 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": null}, + "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}, + ], + "expires_at": null, + }) + ); + } + + #[test] + 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, + limit: 10, + used: 0, + resets_at: None, + }; + let mut v = [ + 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)> = v.iter().map(|e| (e.kind, e.window.as_str())).collect(); + assert_eq!( + order, + [ + ("requests", "1m"), + ("tokens", "1m"), + ("tokens", "1w"), + ("tokens", "weekly"), + ("tokens", "monthly"), + ] + ); + } + + #[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 4af5492d..1b49065d 100644 --- a/crates/server/src/middleware/api_key_auth.rs +++ b/crates/server/src/middleware/api_key_auth.rs @@ -249,6 +249,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`. @@ -261,6 +283,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 { @@ -352,7 +393,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 { @@ -366,18 +407,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 @@ -504,6 +547,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 @@ -669,6 +716,23 @@ mod tests { assert_eq!(body["error"]["message"], "Rate limited: limits_unavailable"); } + #[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..ba1e0296 100644 --- a/crates/test-support/tests/analytics_clickhouse.rs +++ b/crates/test-support/tests/analytics_clickhouse.rs @@ -872,3 +872,88 @@ 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_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( + &app.db, + user_id, + &unique_name("ck-key-b"), + &["ai_gateway"], + None, + None, + ) + .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(); + 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 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(); + let cost = v["usage"]["cost_usd_month"] + .as_f64() + .unwrap_or_else(|| panic!("cost_usd_month is a number: {v:#}")); + (v["scope"].as_str().unwrap().to_string(), cost) + } + }; + 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 - 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 0161533c..9cc093fd 100644 --- a/crates/test-support/tests/limits.rs +++ b/crates/test-support/tests/limits.rs @@ -938,3 +938,528 @@ 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(); + } +} + +/// 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_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 (other_key, 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(); + 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 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; + 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_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; + 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(); + // 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); + 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", distinct_chat(n)) + .await + .unwrap() + .assert_ok(); + } + + let a = key_usage(&gw_a).await; + let b = key_usage(&gw_b).await; + 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), + row("user", "tokens", "daily", 5_000_000, all), + row("user", "tokens", "monthly", 9_000_000, all), + ]; + assert_eq!(limit_rows(&a), expected); + assert_eq!( + 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_with_no_limits_anywhere_answers_for_the_owner() { + let app = TestApp::spawn().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!({"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", 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"], 2); + assert_eq!(v["usage"]["requests_month"], 2); + let tokens = v["usage"]["tokens_today"].as_i64().unwrap(); + 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`"] +#[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); +} + +#[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/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 d7f7e17a..80a20942 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.", @@ -1942,6 +1961,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: /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 cd1b172d..c8878723 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": "当前真实生效的额度 —— 角色合并后再叠加所有生效的用户覆盖。", @@ -1942,6 +1961,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/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() + }) }) 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')} +