diff --git a/CHANGELOG.md b/CHANGELOG.md index 4b1e4763..29968f4a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,161 @@ target. ## [Unreleased] +## [3.5.0] — 2026-10-11 + +Limits hold where they used to leak: route caps are enforced, a limits +store that cannot be read follows `security.rate_limit_fail_closed`, a +refused model or tool no longer spends a request, a team-scoped role no +longer limits the gateway it does not open, and cached answers count their +tokens. A client can ask the gateway what room its API key has left +(`GET /v1/usage`), and an API key's own limits are edited in the console. +The legacy monthly token quota is gone, and the thinkwatch-core crates move +from v0.69.0 to v0.70.0. + +### Read before upgrading + +- **Route caps are enforced.** A route's RPM cap and TPM cap + (`model_routes.rpm_cap` / `tpm_cap`, set in the route editor) were saved + but never applied. They now are, per route over a sliding minute, the TPM + cap in weighted tokens like every other token limit (counted after the + response). A route at either cap is skipped and the model's next route + serves the request; when every route left is at its cap, the request gets + 429 with `Retry-After`, labelled `route:requests/1m` or + `route:tokens/1m`. Such a request has passed the caller's own limits and + counts on them, as a request every upstream fails does. Review the caps + already set on your routes: a cap that did nothing until now takes effect + on upgrade. Routes with a cap write new Redis keys, + `ratelimit:{route:}:ai_gateway:route:…`. +- **A limits store that cannot be read follows + `security.rate_limit_fail_closed`.** When the database failed while a + request's rate limits and budgets were being loaded, the request ran with + no limits at all; when Redis failed while a budget was being read, the + budget read as unspent. Both now follow the setting, as a Redis failure + in the rate limiter always has, and so do the route caps: off (the + default), the request goes on and a warning with the error is logged + (`gateway_limits_load_fail_open_total`, + `lifecycle_budget_fail_open_total`); on, the request is refused, labelled + `limits_unavailable`, `budget_unavailable` or `rate_limiter_unavailable` + (on the AI gateway, 429 with `Retry-After: 30`). With the setting off, a + Redis error used to fail every AI-gateway request anyway, with 502 + `Quota exceeded`, at the legacy quota check (see Changed); such a request + now goes on, as the setting says. Turn the setting on to keep refusing + requests while Redis is down. +- **A role granted at team scope no longer limits gateway requests.** Such + a grant administers that team and has never opened a model or an MCP tool + at the gateways, but its rate limits and budgets were applied to every + gateway request of the user who held it. A role's limits now apply + exactly where its gateway grant does: roles assigned globally, and roles + attached to a team the user is a member of. A user who held a limiting + role only at team scope is no longer held to its limits; assign the role + globally, attach it to the team, or set a user limit if they were meant. +- **Answers from the response cache count their tokens.** A cache hit + counted toward request limits but added no tokens to token limits and + budgets, so repeating a cached request got round them. It now adds the + 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 key's + limits or their counters could not be read, whichever way + `security.rate_limit_fail_closed` is set. 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 + +- **A request the key may not make no longer spends a request limit.** The + rate-limit check, which counts the request, ran before the model check on + the AI gateway and before the tool check on the MCP gateway, so a refused + model or tool used up the caller's `requests` limits. Access is now + checked first (after the budget check, which counts nothing). +- **Route caps are applied**, **limits and budgets that cannot be read + follow the fail-closed setting**, **team-scoped roles no longer limit** + and **cached answers count their tokens**: see Read before upgrading. +- Comments that described a limits cache, a `limits:changed` subscriber and + team-level budget caps, none of which exist, now say what happens: limits + are read from the database on every request, nothing subscribes to + `limits:changed`, and teams carry no limits of their own. + +### Changed + +- **The legacy monthly token quota is removed.** Every request counted its + tokens into `quota:{:}:used:` in Redis + against a limit nothing could set, so it never refused anything. The + gateway no longer reads or writes these keys; existing `quota:*` keys can + be deleted. Monthly limits are budgets. Its metric, + `gateway_quota_overflow_total`, is gone too: the Grafana overview panel + that plotted it shows routes at their cap (`gateway_route_capped_total`) + instead. +- The `security.rate_limit_fail_closed` hint in Settings names everything + the setting now covers. +- thinkwatch-core crates (tw-bedrock, tw-breaker, tw-dialect, tw-guard) + v0.69.0 → v0.70.0. The four crates do not change beyond their version. + + ## [3.4.0] — 2026-10-11 Conversations with reasoning models behind Chat-format upstreams keep their @@ -1768,7 +1923,8 @@ unreleased builds should: stop the gateway, run `db/schema.sql` against PostgreSQL, restart against this tag. The schema is idempotent end-to-end, so the apply is safe to repeat. -[Unreleased]: https://github.com/ThinkWatchProject/ThinkWatch/compare/v3.4.0...HEAD +[Unreleased]: https://github.com/ThinkWatchProject/ThinkWatch/compare/v3.5.0...HEAD +[3.5.0]: https://github.com/ThinkWatchProject/ThinkWatch/releases/tag/v3.5.0 [3.4.0]: https://github.com/ThinkWatchProject/ThinkWatch/releases/tag/v3.4.0 [3.3.1]: https://github.com/ThinkWatchProject/ThinkWatch/releases/tag/v3.3.1 [3.3.0]: https://github.com/ThinkWatchProject/ThinkWatch/releases/tag/v3.3.0 diff --git a/Cargo.lock b/Cargo.lock index 7f2b54fb..bd41fd36 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4189,7 +4189,7 @@ dependencies = [ [[package]] name = "think-watch-auth" -version = "3.4.0" +version = "3.5.0" dependencies = [ "anyhow", "argon2", @@ -4219,7 +4219,7 @@ dependencies = [ [[package]] name = "think-watch-common" -version = "3.4.0" +version = "3.5.0" dependencies = [ "aes-gcm", "anyhow", @@ -4262,7 +4262,7 @@ dependencies = [ [[package]] name = "think-watch-gateway" -version = "3.4.0" +version = "3.5.0" dependencies = [ "anyhow", "arc-swap", @@ -4301,7 +4301,7 @@ dependencies = [ [[package]] name = "think-watch-mcp-gateway" -version = "3.4.0" +version = "3.5.0" dependencies = [ "anyhow", "arc-swap", @@ -4330,7 +4330,7 @@ dependencies = [ [[package]] name = "think-watch-server" -version = "3.4.0" +version = "3.5.0" dependencies = [ "anyhow", "arc-swap", @@ -4381,7 +4381,7 @@ dependencies = [ [[package]] name = "think-watch-test-support" -version = "3.4.0" +version = "3.5.0" dependencies = [ "anyhow", "async-stream", @@ -4837,8 +4837,8 @@ dependencies = [ [[package]] name = "tw-bedrock" -version = "0.69.0" -source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.69.0#13d039d812f39445d44f59b46d33845715bb33ed" +version = "0.70.0" +source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.70.0#205390ee6400126286a86b09c8b751f56368c439" dependencies = [ "aws-credential-types", "aws-sigv4", @@ -4852,16 +4852,16 @@ dependencies = [ [[package]] name = "tw-breaker" -version = "0.69.0" -source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.69.0#13d039d812f39445d44f59b46d33845715bb33ed" +version = "0.70.0" +source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.70.0#205390ee6400126286a86b09c8b751f56368c439" dependencies = [ "serde", ] [[package]] name = "tw-dialect" -version = "0.69.0" -source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.69.0#13d039d812f39445d44f59b46d33845715bb33ed" +version = "0.70.0" +source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.70.0#205390ee6400126286a86b09c8b751f56368c439" dependencies = [ "memchr", "serde", @@ -4870,8 +4870,8 @@ dependencies = [ [[package]] name = "tw-guard" -version = "0.69.0" -source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.69.0#13d039d812f39445d44f59b46d33845715bb33ed" +version = "0.70.0" +source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.70.0#205390ee6400126286a86b09c8b751f56368c439" dependencies = [ "base64 0.22.1", "bytes", diff --git a/Cargo.toml b/Cargo.toml index bf8c74b3..f01ba91c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,7 +10,7 @@ members = [ ] [workspace.package] -version = "3.4.0" +version = "3.5.0" edition = "2024" # The MSRV: the newest `rust-version` among the locked dependencies # (sqlx 0.9 and the aws-smithy crates under aws-sigv4). Without it, @@ -70,10 +70,10 @@ opt-level = 3 # never re-exported through a local shim. And the reverse: something only # this side uses (the at-rest crypto, IMDSv2 credentials, the gateway error) # lives here, not in core. -tw-bedrock = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.69.0" } -tw-breaker = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.69.0" } -tw-dialect = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.69.0" } -tw-guard = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.69.0" } +tw-bedrock = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.70.0" } +tw-breaker = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.70.0" } +tw-dialect = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.70.0" } +tw-guard = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.70.0" } # Web framework axum = { version = "0.8", features = ["macros", "ws"] } diff --git a/README.md b/README.md index 70e29e53..a3c083c9 100644 --- a/README.md +++ b/README.md @@ -94,8 +94,10 @@ The gateway (port `3000`) is the only part that clients need to reach. The conso **Limits and budgets** - Every limit and budget is checked before the request, against what earlier requests used; tokens are counted after the response, so one request can cross a token limit or a budget before the next is refused. A refused request counts against nothing, and its `429` says in `Retry-After` when the limit frees. - A limit set on an API key applies on top of its owner's, on a counter of its own. +- A route's RPM and TPM caps count over a sliding minute. A route at its cap is skipped for the model's next route; when every route is at its cap, the request gets `429`. Such a request has passed the caller's own limits and counts on them. +- `GET /v1/usage` on the gateway port tells a client, with its API key, what it has used today and this month and what room its limits leave. - Redis can be a single node or a Redis Cluster. -- If Redis is unavailable, limits fail open by default. Setting `security.rate_limit_fail_closed` refuses requests instead. +- If Redis is unavailable, or the database fails while a request's limits are loaded, limits fail open by default. Setting `security.rate_limit_fail_closed` refuses those requests instead. - Budget alerts fire once per period at 50%, 80%, 95% and 100%. ## Documentation diff --git a/crates/auth/src/rbac.rs b/crates/auth/src/rbac.rs index 9867d1c6..25a9142d 100644 --- a/crates/auth/src/rbac.rs +++ b/crates/auth/src/rbac.rs @@ -15,6 +15,10 @@ use uuid::Uuid; // (a role granted at team scope administers that team only). // Rate limits: per (metric, window) take MIN MaxCount — most restrictive. // Budgets: per Period take MIN MaxTokens — most restrictive. +// Rate limits and budgets come from the same roles as the model / tool +// scope (`load_gateway_policy_documents`), and from the same statements +// (the `Constraints` of the Allow for `*_gateway:use`): a role's limits +// apply exactly when its gateway grant does. // Deny statements: win over Allow across all roles. // // `compute_user_permissions` is the single source of truth for the @@ -49,8 +53,9 @@ pub async fn compute_user_permissions( } /// Load all `policy_document` JSONB values from roles assigned to -/// `user_id` (direct + team-inherited). Shared helper for permission -/// extraction, resource limit computation, and constraint aggregation. +/// `user_id` (direct at any scope + team-inherited). The console's +/// permission checks start from these; the gateways do not (see +/// [`load_gateway_policy_documents`]). async fn load_user_policy_documents( pool: &PgPool, user_id: Uuid, @@ -72,14 +77,18 @@ async fn load_user_policy_documents( Ok(rows.into_iter().map(|(v,)| v).collect()) } -/// The policy documents that decide gateway access for `user_id`: -/// roles assigned at global scope, plus roles attached to a team the -/// user is a member of (a team's roles are its members' working roles). +/// The policy documents that apply to `user_id`'s gateway requests — +/// the one place that decides it, for access (models, MCP tools) and +/// for the role-level rate limits and budgets alike: roles assigned at +/// global scope, plus roles attached to a team the user is a member of +/// (a team's roles are its members' working roles). /// /// Roles granted at `scope_kind = 'team'` are left out. Such a grant /// lets the holder administer that team from the console; gateway /// requests carry no team, so honouring it here would turn a team -/// grant into platform-wide model and tool access. +/// grant into platform-wide model and tool access — or, for its limits, +/// hold the holder's every request to a role that grants them nothing +/// at the gateway. async fn load_gateway_policy_documents( pool: &PgPool, user_id: Uuid, @@ -128,26 +137,6 @@ pub async fn load_user_role_names( Ok(rows.into_iter().map(|(n,)| n).collect()) } -/// Load the distinct role IDs assigned to `user_id`, including -/// roles inherited through team membership. -/// Used by the rate-limit engine so role-level quotas apply to the user. -pub async fn load_user_role_ids(pool: &PgPool, user_id: Uuid) -> Result, sqlx::Error> { - let rows: Vec<(Uuid,)> = sqlx::query_as( - "SELECT DISTINCT role_id FROM ( \ - SELECT ra.role_id FROM rbac_role_assignments ra WHERE ra.user_id = $1 \ - UNION \ - SELECT tra.role_id \ - FROM team_members tm \ - JOIN team_role_assignments tra ON tra.team_id = tm.team_id \ - WHERE tm.user_id = $1 \ - ) combined", - ) - .bind(user_id) - .fetch_all(pool) - .await?; - Ok(rows.into_iter().map(|(id,)| id).collect()) -} - /// Load every `(role_id, scope_kind, scope_id)` row for `user_id`, /// including roles inherited through team membership (surfaced as /// global scope since team roles grant permissions platform-wide). @@ -315,11 +304,16 @@ fn resource_limits_from_documents(docs: &[serde_json::Value]) -> UserResourceLim /// `compute_user_surface_constraints` so callers that need both the /// pre-override and post-override view (the admin dashboard) can avoid /// re-running the role document load. +/// +/// From the roles that decide gateway access +/// ([`load_gateway_policy_documents`]), the same set +/// [`compute_user_resource_limits`] reads: a role granted at team scope +/// neither opens the gateway nor limits it. pub async fn compute_user_role_constraints( pool: &PgPool, user_id: Uuid, ) -> Result { - let docs = load_user_policy_documents(pool, user_id).await?; + let docs = load_gateway_policy_documents(pool, user_id).await?; let per_role: Vec = docs .iter() .map(think_watch_common::limits::extract_surface_constraints) @@ -329,9 +323,9 @@ pub async fn compute_user_role_constraints( )) } -/// Aggregate surface constraints from policy_documents across every -/// role assigned to `user_id` (direct + team-inherited) using "most -/// restrictive wins": per `(surface, metric, window_secs)` take the +/// Aggregate surface constraints from the policy_documents of the roles +/// that apply at the gateway (see [`compute_user_role_constraints`]) +/// using "most restrictive wins": per `(surface, metric, window_secs)` take the /// MIN `max_count`; per `(surface, period)` take the MIN /// `limit_tokens`. Disabled or non-positive entries are ignored, then /// any active side-table overrides for the user are applied on top. diff --git a/crates/common/src/lifecycle/mod.rs b/crates/common/src/lifecycle/mod.rs index d6ae2cd0..f06ddb8f 100644 --- a/crates/common/src/lifecycle/mod.rs +++ b/crates/common/src/lifecycle/mod.rs @@ -9,8 +9,8 @@ //! ```text //! Raw //! → check_budget (pre-call budget peek; charges nothing) +//! → check_access (allowed_models / allowed_tools; charges nothing) //! → check_limits (rate-limit gate; charges only when it passes) -//! → check_access (allowed_models / allowed_tools) //! → surface-specific (cache lookup, breaker, credential resolution, //! invoke_upstream → Invocation) //! → run_post_invoke @@ -29,7 +29,7 @@ //! `record_outcome` / `write_cache` / `record_usage` / //! `emit_audit` and short-circuit response factories. //! - [`state`] — per-stage state structs ([`state::Raw`], -//! [`state::LimitsChecked`], [`state::Authorized`], +//! [`state::AccessChecked`], [`state::Authorized`], //! [`state::Invoked`], [`state::Emitted`]) plus the //! [`state::Invocation`] / [`state::CapturedView`] enums that //! model the buffered/streaming fork at `invoke_upstream`. diff --git a/crates/common/src/lifecycle/stages/check_access.rs b/crates/common/src/lifecycle/stages/check_access.rs index a05c266d..6f426958 100644 --- a/crates/common/src/lifecycle/stages/check_access.rs +++ b/crates/common/src/lifecycle/stages/check_access.rs @@ -4,14 +4,18 @@ //! decision in the same audit + short-circuit pattern //! [`super::check_limits`] uses, so the deny path is uniform //! regardless of which surface invoked the stage. +//! +//! Runs before `check_limits`, which charges the `requests` counters: +//! a request the identity may not make is refused without counting +//! against any limit. Charges nothing itself. use crate::audit::AuditLogger; use super::super::Surface; -use super::super::state::{Authorized, LimitsChecked}; +use super::super::state::{AccessChecked, Raw}; /// Run the access check. On allow, transition to -/// [`Authorized`] carrying the `candidate` string for downstream +/// [`AccessChecked`] carrying the `candidate` string for downstream /// audit. On deny, emit a `"access_denied"` audit row tagged with /// the subject and short-circuit with /// [`Surface::access_denied_response`]. @@ -20,10 +24,10 @@ use super::super::state::{Authorized, LimitsChecked}; fields(trace_id = %state.trace_id, candidate = %candidate), )] pub async fn check_access( - state: LimitsChecked, + state: Raw, candidate: &str, audit: &AuditLogger, -) -> Result, S::Response> { +) -> Result, S::Response> { if !S::is_access_allowed(&state.identity, candidate) { metrics::counter!("lifecycle_access_denied_total").increment(1); tracing::warn!( @@ -38,42 +42,27 @@ pub async fn check_access( return Err(S::access_denied_response(candidate)); } - Ok(Authorized { + Ok(AccessChecked { identity: state.identity, trace_id: state.trace_id, started_at: state.started_at, client_ip: state.client_ip, - limit_check: state.limit_check, access_candidate: candidate.to_owned(), }) } #[cfg(test)] mod tests { - use super::super::super::state::{LimitCheckRecord, LimitsChecked}; use super::super::super::test_surface::{TestResponse, TestSurface, make_raw}; use super::*; use uuid::Uuid; - fn make_limits_checked(user_id: Uuid) -> LimitsChecked { - let raw = make_raw(user_id); - LimitsChecked { - identity: raw.identity, - trace_id: raw.trace_id, - started_at: raw.started_at, - client_ip: raw.client_ip, - limit_check: LimitCheckRecord { - currents: Vec::new(), - }, - } - } - #[tokio::test] async fn allow_passes_through() { // `TestSurface::is_access_allowed` returns true iff // candidate starts with "allowed_" — see test_surface.rs. let user_id = Uuid::new_v4(); - let state = make_limits_checked(user_id); + let state = make_raw(user_id); let result = check_access::( state, "allowed_tool", @@ -92,7 +81,7 @@ mod tests { #[tokio::test] async fn deny_short_circuits_with_response() { let user_id = Uuid::new_v4(); - let state = make_limits_checked(user_id); + let state = make_raw(user_id); let result = check_access::( state, "blocked_tool", diff --git a/crates/common/src/lifecycle/stages/check_budget.rs b/crates/common/src/lifecycle/stages/check_budget.rs index 48e8c609..34afe1d7 100644 --- a/crates/common/src/lifecycle/stages/check_budget.rs +++ b/crates/common/src/lifecycle/stages/check_budget.rs @@ -14,13 +14,16 @@ //! //! No mutation, so failure of the Redis read defaults to fail-open //! (request allowed) unless the caller passes `fail_closed = true` -//! — matching the [`super::check_limits`] semantics. +//! — matching the [`super::check_limits`] semantics. A failed read is +//! never taken for an unspent budget (`budget::current_spend` returns +//! the error). //! -//! Ordering note: this stage runs BEFORE `check_limits`, on the -//! [`Raw`] state. `check_limits` charges the `requests` counters, so -//! running it first would charge a request this stage then refuses. -//! Since this stage charges nothing, a request refused by either -//! charges nothing. +//! Ordering note: this stage runs first, on the [`Raw`] state, then +//! `check_access`, then `check_limits`. `check_limits` charges the +//! `requests` counters, so running it before either of the others +//! would charge a request they then refuse. Neither of them charges +//! anything, so a request refused before `check_limits` charges +//! nothing. use chrono::Utc; use fred::clients::Client; @@ -141,8 +144,16 @@ mod tests { use uuid::Uuid; fn dummy_redis() -> fred::clients::Client { + // Never connected, so a command waits for a connection that + // never comes; the timeout turns that into the error an outage + // gives. let cfg = RedisConfig::from_url("redis://127.0.0.1:6379").expect("parse url"); - Builder::from_config(cfg).build().expect("build client") + Builder::from_config(cfg) + .with_performance_config(|c| { + c.default_command_timeout = std::time::Duration::from_millis(50) + }) + .build() + .expect("build client") } fn dummy_audit() -> crate::audit::AuditLogger { @@ -168,6 +179,50 @@ mod tests { assert_eq!(passed.started_at, started_at); } + fn one_cap() -> BudgetCap { + BudgetCap { + id: Uuid::nil(), + subject_kind: BudgetSubject::User, + subject_id: Uuid::nil(), + period: BudgetPeriod::Daily, + limit_tokens: 1000, + enabled: true, + expires_at: None, + reason: None, + created_by: None, + } + } + + /// The budget can't be read and the gateway fails closed: refused + /// as unavailable. (The disconnected client fails every command.) + #[tokio::test] + async fn an_unreadable_budget_refuses_when_failing_closed() { + let result = check_budget::( + make_raw(Uuid::new_v4()), + &[one_cap()], + &dummy_redis(), + true, + &dummy_audit(), + ) + .await; + assert_eq!(result.err(), Some(TestResponse::BudgetUnavailable)); + } + + /// The budget can't be read and the gateway fails open: the request + /// goes on. + #[tokio::test] + async fn an_unreadable_budget_passes_when_failing_open() { + let result = check_budget::( + make_raw(Uuid::new_v4()), + &[one_cap()], + &dummy_redis(), + false, + &dummy_audit(), + ) + .await; + assert!(result.is_ok()); + } + #[test] fn label_format_is_subject_budget_period() { let cap = BudgetCap { diff --git a/crates/common/src/lifecycle/stages/check_limits.rs b/crates/common/src/lifecycle/stages/check_limits.rs index 210fb2ae..6e56bac7 100644 --- a/crates/common/src/lifecycle/stages/check_limits.rs +++ b/crates/common/src/lifecycle/stages/check_limits.rs @@ -2,7 +2,7 @@ //! Checks every rate-limit rule the request is held to, in one atomic //! step, and either: //! -//! * passes through to [`LimitsChecked`] (carrying the post-charge +//! * passes through to [`Authorized`] (carrying the post-charge //! currents for downstream audit), having charged 1 to every //! `requests` rule, or //! * short-circuits with `S::rate_limited_response(label, retry_after)` @@ -13,6 +13,10 @@ //! usage has reached the limit; the tokens the call uses are added //! after it, by the surface's `record_usage` hook. //! +//! It is the last pre-call gate and takes [`AccessChecked`]: only a +//! request the identity may make reaches it, so a refused model or +//! tool never spends a request. +//! //! No surface-specific logic lives here — the stage takes //! pre-built `RateLimitRule`s. The surface's pipeline runner is //! responsible for materialising rules from `SurfaceConstraints` @@ -25,7 +29,7 @@ use crate::audit::AuditLogger; use crate::limits::{RateLimitRule, sliding}; use super::super::Surface; -use super::super::state::{LimitCheckRecord, LimitsChecked, Raw}; +use super::super::state::{AccessChecked, Authorized, LimitCheckRecord}; /// Check the rules and charge the request's `requests` counters. See /// module-level docs. `owner` is the user the request runs as: every @@ -42,18 +46,18 @@ use super::super::state::{LimitCheckRecord, LimitsChecked, Raw}; fields(trace_id = %state.trace_id, rule_count = rules.len()), )] pub async fn check_limits( - state: Raw, + state: AccessChecked, rules: &[RateLimitRule], owner: Uuid, redis: &Client, fail_closed: bool, audit: &AuditLogger, -) -> Result, S::Response> { +) -> Result, S::Response> { // No rules configured ⇒ trivially pass. Skip Redis entirely so // a misconfigured surface (no rules attached) doesn't pay a // round-trip per request. if rules.is_empty() { - return Ok(LimitsChecked { + return Ok(Authorized { identity: state.identity, trace_id: state.trace_id, started_at: state.started_at, @@ -61,6 +65,7 @@ pub async fn check_limits( limit_check: LimitCheckRecord { currents: Vec::new(), }, + access_candidate: state.access_candidate, }); } @@ -105,7 +110,7 @@ pub async fn check_limits( return Err(S::rate_limited_response(&label, outcome.retry_after_secs)); } - Ok(LimitsChecked { + Ok(Authorized { identity: state.identity, trace_id: state.trace_id, started_at: state.started_at, @@ -113,23 +118,62 @@ pub async fn check_limits( limit_check: LimitCheckRecord { currents: outcome.currents, }, + access_candidate: state.access_candidate, }) } #[cfg(test)] mod tests { - use super::super::super::test_surface::{TestSurface, make_raw}; + use super::super::super::test_surface::{TestResponse, TestSurface, make_raw}; use super::*; + use crate::limits::{RateLimitSubject, RateMetric, Surface as LimitSurface}; use fred::types::Builder; use fred::types::config::Config as RedisConfig; use uuid::Uuid; - /// Build a disconnected fred client. The `no_rules_…` test - /// never reaches a Redis call so this is fine; tests that need - /// real Redis live in `crates/test-support/tests/limits*.rs`. + /// Build a disconnected fred client: it is never connected, so + /// every command it is given fails — a Redis outage, for the tests + /// that need one. Tests that need real Redis live in + /// `crates/test-support/tests/limits*.rs`. fn dummy_redis() -> fred::clients::Client { + // Never connected, so a command waits for a connection that + // never comes; the timeout turns that into the error an outage + // gives. let cfg = RedisConfig::from_url("redis://127.0.0.1:6379").expect("parse url"); - Builder::from_config(cfg).build().expect("build client") + Builder::from_config(cfg) + .with_performance_config(|c| { + c.default_command_timeout = std::time::Duration::from_millis(50) + }) + .build() + .expect("build client") + } + + /// What `check_access` hands on. + fn make_access_checked(user_id: Uuid) -> AccessChecked { + let raw = make_raw(user_id); + AccessChecked { + identity: raw.identity, + trace_id: raw.trace_id, + started_at: raw.started_at, + client_ip: raw.client_ip, + access_candidate: "allowed_model".to_owned(), + } + } + + fn one_rule(user_id: Uuid) -> RateLimitRule { + RateLimitRule { + id: Uuid::nil(), + subject_kind: RateLimitSubject::User, + subject_id: user_id, + surface: LimitSurface::AiGateway, + metric: RateMetric::Requests, + window_secs: 60, + max_count: 10, + enabled: true, + expires_at: None, + reason: None, + created_by: None, + } } fn dummy_audit() -> crate::audit::AuditLogger { @@ -142,10 +186,10 @@ mod tests { #[tokio::test] async fn passes_through_with_no_rules() { // No rules ⇒ stage skips Redis entirely and emits - // LimitsChecked with empty currents. Verifies the type + // Authorized with empty currents. Verifies the type // transition and field-preservation contract. let user_id = Uuid::new_v4(); - let raw = make_raw(user_id); + let raw = make_access_checked(user_id); let trace_id = raw.trace_id.clone(); let started_at = raw.started_at; @@ -167,5 +211,41 @@ mod tests { limits.limit_check.currents.is_empty(), "no rules ⇒ no currents to report" ); + assert_eq!(limits.access_candidate, "allowed_model"); + } + + /// Redis down, failing closed: the request is refused as + /// unavailable, not let through and not reported as over a limit. + #[tokio::test] + async fn an_unreachable_store_refuses_when_failing_closed() { + let user_id = Uuid::new_v4(); + let result = check_limits::( + make_access_checked(user_id), + &[one_rule(user_id)], + user_id, + &dummy_redis(), + true, + &dummy_audit(), + ) + .await; + assert_eq!(result.err(), Some(TestResponse::RateLimiterUnavailable)); + } + + /// Redis down, failing open: the request passes, with nothing + /// counted. + #[tokio::test] + async fn an_unreachable_store_passes_when_failing_open() { + let user_id = Uuid::new_v4(); + let result = check_limits::( + make_access_checked(user_id), + &[one_rule(user_id)], + user_id, + &dummy_redis(), + false, + &dummy_audit(), + ) + .await; + let passed = result.expect("fail-open lets the request through"); + assert!(passed.limit_check.currents.is_empty()); } } diff --git a/crates/common/src/lifecycle/stages/mod.rs b/crates/common/src/lifecycle/stages/mod.rs index d0dd2ffb..715ac84e 100644 --- a/crates/common/src/lifecycle/stages/mod.rs +++ b/crates/common/src/lifecycle/stages/mod.rs @@ -1,9 +1,9 @@ //! Pipeline stages — each surface-agnostic stage is a plain //! `async fn` here. See [`super`] for the full pipeline shape. //! -//! Short-circuit stages (`check_budget`, `check_limits`, -//! `check_access`) live as standalone fns that take the previous -//! state struct and return either the next state or +//! Short-circuit stages (`check_budget`, `check_access`, +//! `check_limits`, in that order) live as standalone fns that take the +//! previous state struct and return either the next state or //! `Err(S::Response)`. The four post-invoke stages //! (`record_outcome` → `write_cache` → `record_usage` → //! `emit_audit`) live in [`run_post_invoke`] and are dispatched in diff --git a/crates/common/src/lifecycle/state.rs b/crates/common/src/lifecycle/state.rs index 93b39a09..3788f923 100644 --- a/crates/common/src/lifecycle/state.rs +++ b/crates/common/src/lifecycle/state.rs @@ -1,21 +1,22 @@ //! Per-stage state structs. Each stage consumes one struct and //! returns the next, encoding the lifecycle progression in the -//! type system: you literally cannot call `check_access` against a -//! [`Raw`] state because the signature takes `LimitsChecked`. +//! type system: you literally cannot call `check_limits` against a +//! [`Raw`] state because the signature takes [`AccessChecked`] — a +//! request is checked against what it may do before it is counted. //! //! The transitions move ownership, not clone. Each subsequent //! struct destructures the previous one and adds the new field(s). //! Stage code looks like: //! //! ```ignore -//! let Raw { identity, trace_id, started_at, client_ip } = state; +//! let AccessChecked { identity, trace_id, started_at, client_ip, access_candidate } = state; //! let limit_check = run_limit_check(&identity, &redis, &rules).await?; -//! Ok(LimitsChecked { identity, trace_id, started_at, client_ip, limit_check }) +//! Ok(Authorized { identity, trace_id, started_at, client_ip, limit_check, access_candidate }) //! ``` //! //! Verbose by design — the destructure makes every field carried //! across the transition visible at the call site. Adding a field -//! to `Raw` that should also live in `LimitsChecked` is a localised +//! to `Raw` that should also live in `Authorized` is a localised //! edit (one struct + one transition fn). use std::pin::Pin; @@ -60,25 +61,25 @@ impl Raw { } } -/// Output of `check_limits`. Carries everything from [`Raw`] plus -/// the materialised limit-check result so downstream stages (in -/// particular `emit_audit`) can record what budgets / counters -/// this request charged. -pub struct LimitsChecked { +/// Output of `check_access`: everything from [`Raw`] plus the +/// `candidate` the access decision was made against (tool name for +/// MCP, model name for the AI gateway). Access is decided before the +/// limits are checked, so a request the identity may not make is +/// refused without charging any counter. +pub struct AccessChecked { pub identity: S::Identity, pub trace_id: String, pub started_at: Instant, pub client_ip: Option, - /// Outcome of the rate-limit check. For `check_limits` to - /// emit `LimitsChecked` at all, the check must have *passed*; - /// this field records the per-window currents the audit row - /// surfaces in its `limits` block. - pub limit_check: LimitCheckRecord, + /// Surface-specific candidate the access check ran against — + /// a tool name (`stream__test_tool`) for MCP, a model id + /// (`gpt-4o-mini`) for the AI gateway. + pub access_candidate: String, } /// Successful-check trace data. Mirrors the shape /// [`crate::limits::sliding::CheckOutcome`] returns, minus the -/// "allowed: bool" — if we're carrying this in a `LimitsChecked` +/// "allowed: bool" — if we're carrying this in an [`Authorized`] /// we already know it was allowed. #[derive(Debug, Clone)] pub struct LimitCheckRecord { @@ -88,17 +89,21 @@ pub struct LimitCheckRecord { pub currents: Vec, } -/// Output of `check_access`. Same field set as -/// [`LimitsChecked`] — the access stage doesn't add new data, it -/// just narrows the type so subsequent stages can't be reached -/// without it. Carries the `candidate` string the access decision -/// was made against (tool name for MCP, model name for the AI -/// gateway) so downstream audit can record what was authorized. +/// Output of `check_limits`, the last pre-call gate: access was +/// granted and every rate limit passed. Carries the materialised +/// limit-check result so downstream stages (in particular +/// `emit_audit`) can record what counters this request charged, and +/// the `candidate` string the access decision was made against so +/// downstream audit can record what was authorized. pub struct Authorized { pub identity: S::Identity, pub trace_id: String, pub started_at: Instant, pub client_ip: Option, + /// Outcome of the rate-limit check. For `check_limits` to emit + /// `Authorized` at all, the check must have *passed*; this field + /// records the per-window currents the audit row surfaces in its + /// `limits` block. pub limit_check: LimitCheckRecord, /// Surface-specific candidate the access check ran against — /// a tool name (`stream__test_tool`) for MCP, a model id diff --git a/crates/common/src/limits/budget.rs b/crates/common/src/limits/budget.rs index 4869ccfb..ebb2cb05 100644 --- a/crates/common/src/limits/budget.rs +++ b/crates/common/src/limits/budget.rs @@ -179,8 +179,10 @@ pub struct CapStatus { pub limit: i64, } -/// Read the current spend for every cap in the slice. No mutation — -/// used for "show me the dashboard". +/// Read the current spend for every cap in the slice, in input order. +/// No mutation. A counter that does not exist yet reads as 0; a read +/// that fails is an error, never a 0 — the pre-call gate +/// (`check_budget`) decides from it whether to fail open or closed. pub async fn current_spend( redis: &Client, caps: &[BudgetCap], @@ -188,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.ok(); - out.push(CapStatus { - cap_id: cap.id, - subject_kind: cap.subject_kind, - subject_id: cap.subject_id, - current: v.unwrap_or(0), - limit: cap.limit_tokens, - }); + 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 @@ -395,6 +425,35 @@ mod tests { assert_eq!(crossed_thresholds(0, 999, 0), Vec::::new()); } + /// A read that fails is an error, not an unspent budget: with + /// `.ok()` here the fail-closed branch of `check_budget` could never + /// run. A client that never connected times its commands out. + #[tokio::test] + async fn a_failed_read_is_an_error_not_zero() { + use fred::types::Builder; + use fred::types::config::Config; + let redis = Builder::from_config(Config::from_url("redis://127.0.0.1:1").unwrap()) + .with_performance_config(|c| { + c.default_command_timeout = std::time::Duration::from_millis(50) + }) + .build() + .unwrap(); + let cap = BudgetCap { + id: uuid::Uuid::nil(), + subject_kind: BudgetSubject::User, + subject_id: uuid::Uuid::nil(), + period: BudgetPeriod::Daily, + limit_tokens: 100, + enabled: true, + expires_at: None, + reason: None, + created_by: None, + }; + assert!(current_spend(&redis, &[cap]).await.is_err()); + // No caps, nothing to read. + assert!(current_spend(&redis, &[]).await.unwrap().is_empty()); + } + #[test] fn key_format_is_stable() { let t = Utc.with_ymd_and_hms(2026, 4, 8, 0, 0, 0).unwrap(); diff --git a/crates/common/src/limits/mod.rs b/crates/common/src/limits/mod.rs index 8ea3bf14..3041cd16 100644 --- a/crates/common/src/limits/mod.rs +++ b/crates/common/src/limits/mod.rs @@ -11,7 +11,7 @@ // And the corresponding submodules in this folder: // // sliding — bucketed Lua admit (before) / record (after) over Redis -// budget — natural-period add_weighted_tokens / check_cap +// budget — natural-period current_spend (before) / add_weighted_tokens (after) // weight — model_id → weighted token converter (LRU cached) // // Subject identification: every rule / cap is keyed by a (subject_kind, @@ -48,7 +48,11 @@ // Budget counters are read and written one key per command and need no // tag. // -// See `plan.md` (limits chapter) for the full design. +// Nothing caches these rows: the API-key middleware loads a request's +// rules and caps from Postgres on every request +// (`rbac::compute_user_surface_constraints` / +// `compute_key_surface_constraints`), so an edit applies to the next +// request. // ============================================================================ use serde::{Deserialize, Serialize}; @@ -57,6 +61,7 @@ use uuid::Uuid; pub mod budget; pub mod sliding; +pub mod usage; pub mod weight; // ---------------------------------------------------------------------------- @@ -737,8 +742,8 @@ pub struct BudgetCap { } // ---------------------------------------------------------------------------- -// CRUD — straight DB reads / writes. The hot read path goes through -// `LimitsCache` (below) instead of these helpers. +// CRUD — straight DB reads / writes. The request path uses them too: +// there is no cache in front of them (see the module header). // ---------------------------------------------------------------------------- /// Permitted sliding-window lengths (seconds). Anything outside this @@ -1158,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, } @@ -1177,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); @@ -1245,11 +1255,12 @@ impl RequestLimits { } // ---------------------------------------------------------------------------- -// Cache-invalidation pubsub +// Change notice // -// Same shape as `dynamic_config::notify_config_changed`, but on its -// own channel so a settings change doesn't force every gateway to -// drop its limits cache. +// Published on `limits:changed` after every rule / cap write. Nothing +// subscribes to it: no process caches limits (they are read per +// request, see the module header), so a change needs no notice to take +// effect. // ---------------------------------------------------------------------------- const LIMITS_CHANGED_CHANNEL: &str = "limits:changed"; diff --git a/crates/common/src/limits/sliding.rs b/crates/common/src/limits/sliding.rs index fb340f13..af6302b5 100644 --- a/crates/common/src/limits/sliding.rs +++ b/crates/common/src/limits/sliding.rs @@ -338,41 +338,99 @@ pub async fn admit_at( now_ms: i64, fail_open: bool, ) -> Result { - if rules.is_empty() { - return Ok(CheckOutcome::open()); - } - let mut keys = Vec::with_capacity(rules.len()); - let mut args = Vec::with_capacity(1 + rules.len() * 3); - args.push(now_ms.to_string()); - for rule in rules { - let r = ResolvedRule::new(rule, owner); - keys.push(r.key); - args.push(r.bucket_secs.to_string()); - args.push(r.max_count.to_string()); - args.push(match rule.metric { - RateMetric::Requests => "1".to_string(), - RateMetric::Tokens => "0".to_string(), - }); - } + let counters: Vec = rules + .iter() + .map(|rule| { + let r = ResolvedRule::new(rule, owner); + Counter { + key: r.key, + bucket_secs: r.bucket_secs, + max_count: r.max_count, + charge: match rule.metric { + RateMetric::Requests => 1, + RateMetric::Tokens => 0, + }, + } + }) + .collect(); - let reply = match run(redis, LUA_ADMIT, &ADMIT_SHA, keys, args).await { - Ok(v) => v, + match admit_counters_at(redis, &counters, now_ms).await { + Ok(outcome) => Ok(outcome), Err(e) if fail_open => { tracing::warn!("rate-limit check failed: {e}; failing open"); metrics::counter!("gateway_rate_limiter_fail_open_total").increment(1); - return Ok(CheckOutcome::open()); + Ok(CheckOutcome::open()) } Err(e) => { tracing::error!( "rate-limit check failed: {e}; failing closed per security.rate_limit_fail_closed" ); metrics::counter!("gateway_rate_limiter_fail_closed_total").increment(1); - return Err(e); + Err(e) } - }; + } +} + +/// One counter for [`admit_counters_at`]: its Redis key, its bucket +/// size, its limit, and what an admitted request adds to it (1 for a +/// request count, 0 for a token count, which is only read). +#[derive(Debug, Clone)] +pub struct Counter { + pub key: String, + pub bucket_secs: i32, + pub max_count: i64, + pub charge: i64, +} + +/// The `admit` script on counters the caller names itself — for limits +/// that are not a user's or a key's (a route's capacity). Same +/// semantics as [`admit_at`]: refuse when any window is full, otherwise +/// charge each counter, all or nothing. Every key must carry the same +/// Redis Cluster hash tag. A Redis error is returned as is; failing +/// open or closed is the caller's decision. +pub async fn admit_counters_at( + redis: &Client, + counters: &[Counter], + now_ms: i64, +) -> Result { + if counters.is_empty() { + return Ok(CheckOutcome::open()); + } + let mut keys = Vec::with_capacity(counters.len()); + let mut args = Vec::with_capacity(1 + counters.len() * 3); + args.push(now_ms.to_string()); + for c in counters { + keys.push(c.key.clone()); + args.push(c.bucket_secs.to_string()); + args.push(c.max_count.to_string()); + args.push(c.charge.to_string()); + } + let reply = run(redis, LUA_ADMIT, &ADMIT_SHA, keys, args).await?; Ok(parse_admit_reply(&reply)) } +/// The `record` script on counters the caller names itself: add +/// `amount` to each `(key, bucket_secs)`, past its limit or not. Every +/// key must carry the same hash tag. +pub async fn record_counters_at( + redis: &Client, + counters: &[(String, i32)], + amount: i64, + now_ms: i64, +) -> Result<(), fred::error::Error> { + if counters.is_empty() || amount <= 0 { + return Ok(()); + } + let mut args = Vec::with_capacity(2 + counters.len()); + args.push(now_ms.to_string()); + args.push(amount.to_string()); + args.extend(counters.iter().map(|(_, b)| b.to_string())); + let keys = counters.iter().map(|(k, _)| k.clone()).collect(); + run(redis, LUA_RECORD, &RECORD_SHA, keys, args) + .await + .map(|_| ()) +} + /// Reply shape: `[allowed, limiting rule (1-based, 0 = none), wait_ms, /// currents...]`. fn parse_admit_reply(reply: &[i64]) -> CheckOutcome { @@ -427,22 +485,15 @@ pub async fn record_at( amount: i64, now_ms: i64, ) -> Result<(), fred::error::Error> { - let resolved: Vec = rules + let counters: Vec<(String, i32)> = rules .iter() .filter(|r| r.metric == metric) - .map(|r| ResolvedRule::new(r, owner)) + .map(|r| { + let r = ResolvedRule::new(r, owner); + (r.key, r.bucket_secs) + }) .collect(); - if resolved.is_empty() || amount <= 0 { - return Ok(()); - } - let mut args = Vec::with_capacity(2 + resolved.len()); - args.push(now_ms.to_string()); - args.push(amount.to_string()); - args.extend(resolved.iter().map(|r| r.bucket_secs.to_string())); - let keys = resolved.into_iter().map(|r| r.key).collect(); - run(redis, LUA_RECORD, &RECORD_SHA, keys, args) - .await - .map(|_| ()) + record_counters_at(redis, &counters, amount, now_ms).await } /// Read-only "what's the current sum for this rule" helper. Used by @@ -452,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 @@ -462,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/common/src/limits/weight.rs b/crates/common/src/limits/weight.rs index 132648d2..0482a444 100644 --- a/crates/common/src/limits/weight.rs +++ b/crates/common/src/limits/weight.rs @@ -21,10 +21,11 @@ // // Cache invalidation: by TTL only (5 minutes). Weights change rarely // (admin tunes them in the model management page) so a brief -// staleness window is acceptable. We do NOT subscribe to the limits -// pubsub for this — the channel is for rule / cap changes, not -// model rows. If we need faster propagation later, the caller can -// call `WeightCache::invalidate_all` from the model PATCH handler. +// staleness window is acceptable. Nothing notifies the cache of a +// model edit (the `limits:changed` notice is about rules and caps, +// and nothing subscribes to it). If we need faster propagation later, +// the caller can call `WeightCache::invalidate_all` from the model +// PATCH handler. // // Cache shape: `RwLock>`. // 5-min TTL on entries. Bounded to 1024 distinct model_ids — beyond diff --git a/crates/common/src/models/team.rs b/crates/common/src/models/team.rs index cfc416b2..b4848b1b 100644 --- a/crates/common/src/models/team.rs +++ b/crates/common/src/models/team.rs @@ -8,7 +8,10 @@ pub struct Team { pub id: Uuid, pub name: String, pub description: Option, - // Budget caps live in `budget_caps` (subject_kind = 'team'). + // A team has no limits of its own: rate limits and budgets attach to + // users and API keys (`rate_limit_rules` / `budget_caps` take + // `subject_kind` 'user' or 'api_key_lineage'), and a team's members + // get the limits of the roles the team grants them. pub created_at: DateTime, } diff --git a/crates/gateway/src/cache.rs b/crates/gateway/src/cache.rs index b6bcedc5..e8e3b933 100644 --- a/crates/gateway/src/cache.rs +++ b/crates/gateway/src/cache.rs @@ -8,7 +8,8 @@ use xxhash_rust::xxh3::xxh3_128; /// A cached answer, in the caller's format with redaction placeholders intact. pub struct Cached { pub body: Vec, - /// Kept so a hit can debit quota the way the original call did. + /// Kept so a hit counts toward token limits and budgets the way the + /// original call did. pub prompt_tokens: u32, pub completion_tokens: u32, } diff --git a/crates/gateway/src/lib.rs b/crates/gateway/src/lib.rs index f6f18b0e..65cf3470 100644 --- a/crates/gateway/src/lib.rs +++ b/crates/gateway/src/lib.rs @@ -12,9 +12,8 @@ pub mod metrics_labels; pub mod model_mapping; pub mod protocol; pub mod proxy; -pub mod quota; -pub mod rate_limiter; pub mod redaction; +pub mod route_caps; pub mod router; pub mod strategy; pub mod tool_inspection; diff --git a/crates/gateway/src/lifecycle/mod.rs b/crates/gateway/src/lifecycle/mod.rs index 8ebda1be..48e6e15b 100644 --- a/crates/gateway/src/lifecycle/mod.rs +++ b/crates/gateway/src/lifecycle/mod.rs @@ -119,6 +119,9 @@ pub(crate) struct ChatPickedRoute { pub upstream_model: Option, /// Used by `finalize_health` inside `record_outcome`. pub sel_record: SelectionRecord, + /// The route's capacity caps: its token cap counts the answer's + /// tokens in `record_usage`. + pub caps: crate::route_caps::RouteCaps, } /// Everything the post-invoke hooks read. Built once before the upstream @@ -928,6 +931,7 @@ impl Surface for ChatCompletionSurface { deps.request.mapped_model.clone(), priced(&extract_usage(&invoked.view)), &deps.preflight.limits, + Some(&deps.route.caps), deps.request.identity.user_id.clone(), deps.request.identity.user_email.clone(), deps.request.identity.api_key_id.clone(), diff --git a/crates/gateway/src/proxy/accounting.rs b/crates/gateway/src/proxy/accounting.rs index 3d17dda9..c25041a0 100644 --- a/crates/gateway/src/proxy/accounting.rs +++ b/crates/gateway/src/proxy/accounting.rs @@ -1,9 +1,10 @@ //! Post-flight accounting + token resolution for streaming responses. //! //! [`post_flight_account`] adds what a request used to its token-metric -//! sliding rules and budget caps, from the token counts the upstream -//! returned — or the estimate, when it returned none — with cache reads -//! and writes weighted apart from plain input. Used from BOTH the +//! sliding rules, its budget caps and the token cap of the route that +//! answered, from the token counts the upstream returned — or the +//! estimate, when it returned none — with cache reads and writes +//! weighted apart from plain input. Used from BOTH the //! non-streaming branch (called inline after the upstream future //! resolves) and the streaming branch (called from the post-invoke //! pipeline after the SSE stream is drained). @@ -28,6 +29,9 @@ pub(crate) async fn post_flight_account( model: String, tokens: weight::TokenCounts, request_limits: &RequestLimits, + // The route that answered, for its token cap. `None` when no route + // did (an answer from the response cache). + route: Option<&crate::route_caps::RouteCaps>, // Actor attribution for `budget.threshold_crossed` audit entries. // Without these the crossing log carries only `cap_id`, and // operators investigating a 100 %-cross had to time-join against @@ -44,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 @@ -62,6 +70,11 @@ pub(crate) async fn post_flight_account( tracing::warn!("token rate-limit accounting failed: {e}"); } + // The answering route's tokens-per-minute cap, counted the same way. + if let Some(caps) = route { + crate::route_caps::record_tokens(&redis, caps, weighted).await; + } + // Natural-period budget caps — the user's and the key's. if !request_limits.caps.is_empty() { let caps = &request_limits.caps; @@ -111,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/generate.rs b/crates/gateway/src/proxy/generate.rs index 90aa7600..d14c680d 100644 --- a/crates/gateway/src/proxy/generate.rs +++ b/crates/gateway/src/proxy/generate.rs @@ -709,7 +709,7 @@ async fn run( request_started_at, ); - // 2. Rate limits, budget, model access — each stage writes its own + // 2. Budget, model access, rate limits — each stage writes its own // audit row on short-circuit. let preflight = run_preflight_stages(&state, &identity, &trace_id, &mapped_model).await?; @@ -756,7 +756,7 @@ async fn run( // request can be converted at all. One that cannot — it continues a // conversation kept on OpenAI's servers (`previous_response_id`), or // carries a compaction only OpenAI can read — can still be forwarded - // to an upstream of its own format, so that is where it goes (step 9). + // to an upstream of its own format, so that is where it goes (step 8). let decoded = tw_dialect::convert::decode(surface.dialect, &raw, path, internal_query(surface.dialect)); @@ -798,24 +798,10 @@ async fn run( identity.user_email.clone(), ); - // 7. Quota, keyed on the model the caller named — that is what their - // dashboards group by. - let quota_key = identity - .user_id - .as_deref() - .or(identity.api_key_id.as_deref()) - .map(|id| format!("{id}:{mapped_model}")) - .unwrap_or_else(|| mapped_model.clone()); - if let Err(e) = state.quota.check_quota("a_key).await { - tracing::warn!("Quota exceeded for {quota_key}: {e}"); - return Err(ctx - .emit(GatewayError::ProviderError(format!("Quota exceeded: {e}"))) - .into()); - } - - // 8. Cache. A hit debits quota like a real call would — otherwise a - // deterministic prompt amortises one upstream call across an - // unbounded quota window. + // 7. Cache. A hit counts its tokens toward the token limits and + // budgets like a real call would — the caller received them, and + // otherwise repeating a cached prompt would get round every token + // limit. let cache_fingerprint = if surface.caches { ResponseCache::fingerprint(&outbound_body).map(|mut fp| { // The cap is applied after this, per hop: an answer made under @@ -846,10 +832,25 @@ async fn run( ) { return Err(ctx.emit(e).into()); } - let total = cached.prompt_tokens + cached.completion_tokens; - if let Err(e) = state.quota.consume("a_key, total).await { - tracing::warn!(quota_key = %quota_key, tokens = total, "quota consume on cache hit failed: {e}"); - } + // The tokens the stored answer records, weighted for this model, + // to the same rules and caps the pre-flight checked. No route + // answered, so no route cap counts them. + super::post_flight_account( + state.db.clone(), + state.redis.clone(), + state.dynamic_config.clone(), + state.weight_cache.clone(), + mapped_model.clone(), + cached_tokens(&cached), + &preflight.limits, + None, + identity.user_id.clone(), + identity.user_email.clone(), + identity.api_key_id.clone(), + identity.ip_address.clone(), + state.audit.clone(), + ) + .await; // Same capture pipeline as a fresh request — redaction toggle, // byte cap and offload all apply — with the status marked. @@ -908,7 +909,7 @@ async fn run( metrics::counter!("gateway_cache_total", "result" => "miss").increment(1); } - // 9. Route. + // 8. Route. let router = state.router.load(); let routes = router.route(&mapped_model).ok_or_else(|| { ctx.emit(GatewayError::ProviderError(format!( @@ -972,6 +973,7 @@ async fn run( provider_name: route.provider_name.clone(), upstream_model: route.upstream_model.clone(), sel_record, + caps: crate::route_caps::RouteCaps::of(route), }, cache_enabled: surface.caches, }; @@ -1078,15 +1080,7 @@ async fn run( // Cache fill, audit, breaker and budget debit — the same hooks the // stream runs in its tail. The cache keeps the placeholder form. let deps = snapshot(entry, sel_record); - let completed = run_buffered_post_invoke(&deps, completed).await; - - let (prompt, completion) = tokens(&completed.usage); - let total = prompt + completion; - if total > 0 - && let Err(e) = state.quota.consume("a_key, total).await - { - tracing::warn!("Failed to consume quota: {e}"); - } + run_buffered_post_invoke(&deps, completed).await; tracing::info!( request_id = %metadata.request_id, @@ -1119,6 +1113,19 @@ pub(crate) fn tokens(u: &tw_dialect::usage::Usage) -> (u32, u32) { ) } +/// What a cache hit counts toward token limits and budgets: the tokens +/// its stored answer records. The store keeps only the input and output +/// totals, so all of the input counts as plain input. +fn cached_tokens(cached: &crate::cache::Cached) -> TokenCounts { + TokenCounts { + input: i64::from(cached.prompt_tokens), + cache_read: 0, + cache_write: 0, + cache_write_1h: false, + output: i64::from(cached.completion_tokens), + } +} + /// The same usage split the way it is priced: cache reads and writes /// apart from plain input (see `cost_tracker`). pub(crate) fn priced(u: &tw_dialect::usage::Usage) -> TokenCounts { diff --git a/crates/gateway/src/proxy/identity.rs b/crates/gateway/src/proxy/identity.rs index 82747ef1..7543c1a1 100644 --- a/crates/gateway/src/proxy/identity.rs +++ b/crates/gateway/src/proxy/identity.rs @@ -10,7 +10,7 @@ use think_watch_common::limits::{RequestLimits, Surface}; /// user's counters, and the calling key's own limits on its lineage's /// counters — both apply. A request with no user (never past the auth /// middleware today) is held to nothing. -pub(super) fn limits_for_ai_gateway(identity: &GatewayRequestIdentity) -> RequestLimits { +pub fn limits_for_ai_gateway(identity: &GatewayRequestIdentity) -> RequestLimits { let parse = |s: &Option| s.as_deref().and_then(|s| Uuid::parse_str(s).ok()); let Some(user_id) = parse(&identity.user_id) else { return RequestLimits::default(); diff --git a/crates/gateway/src/proxy/mod.rs b/crates/gateway/src/proxy/mod.rs index 6370147d..62567398 100644 --- a/crates/gateway/src/proxy/mod.rs +++ b/crates/gateway/src/proxy/mod.rs @@ -14,8 +14,6 @@ use crate::cost_tracker::CostTracker; use crate::error::GatewayError; use crate::health::HealthTracker; use crate::model_mapping::ModelMapper; -use crate::quota::QuotaManager; -use crate::rate_limiter::RateLimiter; use crate::router::ModelRouter; use think_watch_common::dynamic_config::DynamicConfig; use think_watch_common::limits::SurfaceConstraints; @@ -48,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; @@ -60,17 +59,15 @@ pub struct GatewayState { /// tool-call inspection. Hot-swapped whole when an admin changes a /// policy; each request runs on the snapshot it took on arrival. pub guards: Arc>, - pub quota: 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. pub db: PgPool, /// Redis client used by the bucketed sliding-window engine and the - /// natural-period budget counters. Same connection used by `quota`, - /// `cache`, and the rest of the gateway. + /// natural-period budget counters. Same connection used by `cache` + /// and the rest of the gateway. pub redis: fred::clients::Client, /// LRU cache mapping `model_id → (input_weight, output_weight)`. /// Looked up once per request to convert raw token counts into @@ -131,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 @@ -169,7 +170,7 @@ impl From for GatewayErrorResponse { impl GatewayErrorResponse { /// Answer in `client`'s format. - pub(crate) fn in_dialect(mut self, client: tw_dialect::ir::Dialect) -> Self { + pub fn in_dialect(mut self, client: tw_dialect::ir::Dialect) -> Self { self.client = client; self } diff --git a/crates/gateway/src/proxy/pipeline.rs b/crates/gateway/src/proxy/pipeline.rs index 506f3175..eea8871c 100644 --- a/crates/gateway/src/proxy/pipeline.rs +++ b/crates/gateway/src/proxy/pipeline.rs @@ -5,7 +5,7 @@ //! are surface-specific. What's factored out here is the boilerplate //! that was duplicated verbatim across all three: //! -//! * [`run_preflight_stages`] — rate-limit + budget + access tower +//! * [`run_preflight_stages`] — budget + access + rate-limit tower //! * [`launch_stream_pump`] — build_chat_pump + detached post-invoke task //! * [`run_buffered_post_invoke`] — Invoked construction + post-invoke + unwrap //! @@ -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 @@ -35,9 +35,10 @@ pub(super) struct Preflight { } /// Run the three shared pre-flight stages — `check_budget` → -/// `check_limits` → `check_access` — that each AI surface gates on. -/// The budget peek goes first because it charges nothing: a request a -/// spent budget refuses must not have used up a request limit. +/// `check_access` → `check_limits` — that each AI surface gates on. +/// `check_limits` charges the request counters, so it goes last: a +/// request a spent budget or the key's model list refuses must not +/// have used up a request limit. The other two charge nothing. /// /// Returns the resolved rule + cap lists so the caller can feed them /// into [`ChatPostInvokeDeps`] without re-walking the identity. @@ -68,8 +69,11 @@ pub(super) async fn run_preflight_stages( ) .await .map_err(short_circuit_to_response)?; - let limits_checked = check_limits::( - raw, + let access_checked = check_access::(raw, model, &state.audit) + .await + .map_err(short_circuit_to_response)?; + let _authorized = check_limits::( + access_checked, &limits.rules, limits.owner, &state.redis, @@ -78,9 +82,16 @@ pub(super) async fn run_preflight_stages( ) .await .map_err(short_circuit_to_response)?; - let _authorized = check_access::(limits_checked, model, &state.audit) - .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/proxy/routing.rs b/crates/gateway/src/proxy/routing.rs index 56fec4fe..183ab6e6 100644 --- a/crates/gateway/src/proxy/routing.rs +++ b/crates/gateway/src/proxy/routing.rs @@ -1,5 +1,6 @@ //! Multi-route selection: strategy-driven weights + circuit-breaker -//! filter + per-model session affinity (none / provider / route). +//! filter + per-route capacity caps (`crate::route_caps`) + per-model +//! session affinity (none / provider / route). use rand::RngExt; use std::str::FromStr; @@ -9,6 +10,7 @@ use super::GatewayState; use crate::call_ctx::CallCtx; use crate::error::GatewayError; use crate::health::{CircuitBreakerConfig, RouteHealth}; +use crate::route_caps::{self, Admission, CappedRoutes, RouteCaps}; use crate::router::{AffinityMode, RouteEntry}; use crate::strategy::{self, RoutingStrategy}; @@ -121,6 +123,8 @@ pub(super) struct SelectionCtx<'a> { pub(super) affinity_ttl_secs: u32, pub(super) latency_k: f64, pub(super) breaker: CircuitBreakerConfig, + /// `security.rate_limit_fail_closed`, for the route caps. + pub(super) fail_closed: bool, pub(super) state: &'a GatewayState, } @@ -133,6 +137,7 @@ pub(super) async fn build_selection_ctx<'a>( resolve_routing_config(state, model_id).await; let latency_k = state.dynamic_config.latency_strategy_k().await; let breaker = resolve_breaker_config(state).await; + let fail_closed = state.dynamic_config.rate_limit_fail_closed().await; SelectionCtx { model_id, user_id, @@ -141,16 +146,19 @@ pub(super) async fn build_selection_ctx<'a>( affinity_ttl_secs, latency_k, breaker, + fail_closed, state, } } /// One selection attempt over the candidate set: snapshot health, -/// drop circuit-broken / already-tried candidates, compute strategy -/// weights, pick. +/// drop circuit-broken / already-tried (`tried`: provider ids) / +/// capped (`capped`: route ids) candidates, compute strategy weights, +/// pick. async fn pick_with_strategy<'a>( group: &[&'a RouteEntry], tried: &[Uuid], + capped: &[Uuid], ctx: &SelectionCtx<'_>, ) -> Option<&'a RouteEntry> { if group.is_empty() { @@ -178,7 +186,9 @@ async fn pick_with_strategy<'a>( let mut excluded: Vec = Vec::with_capacity(group.len()); for (i, entry) in group.iter().enumerate() { let h = &healths[i]; - let excl = h.state == tw_breaker::State::Open || tried.contains(&entry.provider_id); + let excl = h.state == tw_breaker::State::Open + || tried.contains(&entry.provider_id) + || capped.contains(&entry.route_id); excluded.push(excl); let success_rate = if h.total > 0 { Some((1.0 - h.error_pct / 100.0).clamp(0.0, 1.0)) @@ -316,11 +326,47 @@ pub(super) type Answered<'a> = ( SelectionRecord, ); +/// Pick a candidate the request may be sent to: one the strategy picks +/// among those not excluded, and under its caps — the pick is counted +/// against them. A pick at a cap is added to `capped` and another is +/// picked. `None` when no candidate is left. +async fn pick_admitted<'a>( + group: &[&'a RouteEntry], + tried: &[Uuid], + capped: &mut CappedRoutes, + ctx: &SelectionCtx<'_>, +) -> Result, GatewayError> { + // Each pass excludes one more route, so this many passes see them all. + for _ in 0..group.len() { + let Some(entry) = pick_with_strategy(group, tried, capped.routes(), ctx).await else { + return Ok(None); + }; + match route_caps::admit(&ctx.state.redis, &RouteCaps::of(entry), ctx.fail_closed).await? { + Admission::Admitted => return Ok(Some(entry)), + Admission::Capped { + label, + retry_after_secs, + } => { + tracing::info!( + provider = %entry.provider_name, + route_id = %entry.route_id, + cap = %label, + retry_after_secs, + "Route at its cap, trying next" + ); + capped.add(entry.route_id, label, retry_after_secs); + } + } + } + Ok(None) +} + /// Non-streaming selection + failover. All routes are peers (no /// priority tier in v2): `pick_with_strategy` picks one healthy -/// candidate, the proxy calls it, and when the upstream fails (see -/// [`is_upstream_failure`]) tries another candidate from the remaining -/// set until exhausted. +/// candidate under its caps, the proxy calls it, and when the upstream +/// fails (see [`is_upstream_failure`]) tries another candidate from the +/// remaining set until exhausted. When every candidate left is at its +/// cap and none failed, the answer is the caps' 429. pub(super) async fn select_route_with_failover<'a>( routes: &'a [RouteEntry], outbound: &super::generate::Outbound, @@ -333,9 +379,10 @@ pub(super) async fn select_route_with_failover<'a>( let mut last_error: Option = None; let mut tried: Vec = Vec::new(); + let mut capped = CappedRoutes::default(); for _ in 0..candidates.len() { - let Some(entry) = pick_with_strategy(&candidates, &tried, ctx).await else { + let Some(entry) = pick_admitted(&candidates, &tried, &mut capped, ctx).await? else { break; }; tried.push(entry.provider_id); @@ -427,24 +474,28 @@ pub(super) async fn select_route_with_failover<'a>( } } - Err(last_error.unwrap_or_else(|| { + // An upstream that failed is the answer's cause even when the rest + // were capped; with no failure, the caps are. + Err(last_error.or_else(|| capped.refusal()).unwrap_or_else(|| { GatewayError::ProviderError(format!("All routes failed for model: {}", ctx.model_id)) })) } -/// Streaming variant: pick via strategy + health filter, but don't -/// call the provider — caller wires up the SSE stream and once the +/// Streaming variant: pick via strategy + health filter + caps, but +/// don't call the provider — caller wires up the SSE stream and once the /// first chunk lands a retry is no longer possible. Returns the /// chosen entry plus a `SelectionRecord` so the streaming caller can -/// record health on stream completion. +/// record health on stream completion. Every candidate at its cap: the +/// caps' 429. pub(super) async fn select_route_for_stream<'a>( routes: &'a [RouteEntry], ctx: &SelectionCtx<'_>, ) -> Result<(&'a RouteEntry, SelectionRecord), GatewayError> { let started_at = std::time::Instant::now(); let candidates: Vec<&RouteEntry> = routes.iter().collect(); + let mut capped = CappedRoutes::default(); - if let Some(entry) = pick_with_strategy(&candidates, &[], ctx).await { + if let Some(entry) = pick_admitted(&candidates, &[], &mut capped, ctx).await? { return Ok(( entry, SelectionRecord { @@ -458,8 +509,7 @@ pub(super) async fn select_route_for_stream<'a>( )); } - Err(GatewayError::ProviderError(format!( - "No provider found for model: {}", - ctx.model_id - ))) + Err(capped.refusal().unwrap_or_else(|| { + GatewayError::ProviderError(format!("No provider found for model: {}", ctx.model_id)) + })) } diff --git a/crates/gateway/src/quota.rs b/crates/gateway/src/quota.rs deleted file mode 100644 index 31d2ef50..00000000 --- a/crates/gateway/src/quota.rs +++ /dev/null @@ -1,237 +0,0 @@ -use fred::clients::Client; -use fred::interfaces::{KeysInterface, LuaInterface}; -use think_watch_common::errors::AppError; - -/// Per-subject token quota information. Callers today key by user_id or -/// api_key_id — team/role groupings don't own their own counters (see -/// `crates/common/src/limits/mod.rs` for the design rationale). -#[derive(Debug, Clone)] -pub struct QuotaInfo { - /// Monthly token limit (0 means unlimited). - pub limit: u64, - /// Tokens used this month. - pub used: u64, - /// Remaining tokens (limit - used), clamped to 0. - pub remaining: u64, - /// The month key (e.g. "2026-04"). - pub period: String, -} - -/// Per-subject token quota with hard limits, backed by Redis. The `key` -/// parameter on every method is a free-form string by convention, but -/// all in-tree callers compose it from user_id or api_key_id — never -/// team or role id. Keep it that way so a single operator can't -/// accidentally create a shared counter. -#[derive(Clone)] -pub struct QuotaManager { - redis: Client, -} - -impl QuotaManager { - pub fn new(redis: Client) -> Self { - Self { redis } - } - - /// Returns the current month string for Redis keys (e.g. "2026-04"). - fn current_month() -> String { - chrono::Utc::now().format("%Y-%m").to_string() - } - - // `{}` is a Redis Cluster hash tag: `consume` and - // `check_and_consume` declare both keys in one script, which a - // cluster runs only when they share a slot. - fn limit_key(key: &str) -> String { - format!("quota:{{{key}}}:limit") - } - - fn usage_key(key: &str) -> String { - let month = Self::current_month(); - format!("quota:{{{key}}}:used:{month}") - } - - /// Check if user/team has enough quota. Returns remaining tokens. - /// Returns an error if quota is exceeded. - pub async fn check_quota(&self, key: &str) -> Result { - let limit: Option = self.redis.get(Self::limit_key(key)).await.map_err(|e| { - tracing::warn!("Quota limit read failed: {e}"); - AppError::Internal(anyhow::anyhow!("Quota check failed")) - })?; - - // No limit set means unlimited - let limit = match limit { - Some(l) => l, - None => return Ok(u64::MAX), - }; - - if limit == 0 { - return Ok(u64::MAX); - } - - let used: u64 = self.redis.get(Self::usage_key(key)).await.unwrap_or(0); - - if used >= limit { - return Err(AppError::BadRequest(format!( - "Token quota exceeded: used {used}/{limit} tokens this month" - ))); - } - - Ok(limit.saturating_sub(used)) - } - - /// Consume tokens after a successful request. Returns the new - /// remaining count. - /// - /// Atomically reads the limit, INCRs usage, refreshes TTL, and - /// reports whether the new total has crossed the limit — all in - /// a single Redis round-trip via Lua. This closes the - /// check-then-incr TOCTOU window where two concurrent requests - /// could both pass `check_quota` and double-spend. - /// - /// Crossing the limit during consume bumps - /// `gateway_quota_overflow_total` and logs a warning, but the - /// request is **not** rejected — the upstream has already - /// returned its response by this point, and refusing here would - /// just hide the spend without recovering the cost. Use the - /// metric to alert on operators who need to tighten limits. - pub async fn consume(&self, key: &str, tokens: u32) -> Result { - const LUA_CONSUME: &str = r#" -local limit_key = KEYS[1] -local usage_key = KEYS[2] -local tokens = tonumber(ARGV[1]) -local ttl_secs = tonumber(ARGV[2]) - -local limit = tonumber(redis.call('GET', limit_key) or '0') -local new_used = redis.call('INCRBY', usage_key, tokens) -redis.call('EXPIRE', usage_key, ttl_secs) - --- Returns: {new_used, limit, overflowed_flag} -local overflowed = 0 -if limit > 0 and new_used > limit then - overflowed = 1 -end -return {new_used, limit, overflowed} -"#; - - let limit_key = Self::limit_key(key); - let usage_key = Self::usage_key(key); - - let result: Vec = self - .redis - .eval( - LUA_CONSUME, - vec![limit_key, usage_key], - vec![tokens.to_string(), (32 * 86400).to_string()], - ) - .await - .map_err(|e| { - tracing::warn!("Quota consume failed: {e}"); - AppError::Internal(anyhow::anyhow!("Quota consume failed")) - })?; - - let new_used = result.first().copied().unwrap_or(0).max(0) as u64; - let limit = result.get(1).copied().unwrap_or(0).max(0) as u64; - let overflowed = result.get(2).copied().unwrap_or(0) == 1; - - if overflowed { - metrics::counter!("gateway_quota_overflow_total").increment(1); - tracing::warn!( - key = %key, - used = new_used, - limit = limit, - "quota overflow on consume — request was already served, soft cap exceeded" - ); - } - - Ok(limit.saturating_sub(new_used)) - } - - /// Get current usage for a quota key. - pub async fn get_usage(&self, key: &str) -> Result { - let limit: u64 = self.redis.get(Self::limit_key(key)).await.unwrap_or(0); - - let period = Self::current_month(); - let used: u64 = self.redis.get(Self::usage_key(key)).await.unwrap_or(0); - - let remaining = limit.saturating_sub(used); - - Ok(QuotaInfo { - limit, - used, - remaining, - period, - }) - } - - /// Atomically check quota and consume tokens in a single Redis round-trip. - /// Returns remaining tokens after consumption, or an error if quota would be exceeded. - pub async fn check_and_consume(&self, key: &str, tokens: u32) -> Result { - const LUA_QUOTA_CHECK_AND_CONSUME: &str = r#" -local limit_key = KEYS[1] -local usage_key = KEYS[2] -local tokens = tonumber(ARGV[1]) - -local limit = tonumber(redis.call('GET', limit_key) or '0') -if limit <= 0 then return -1 end - -local used = tonumber(redis.call('GET', usage_key) or '0') -if used + tokens > limit then return 0 end - -redis.call('INCRBY', usage_key, tokens) -return limit - used - tokens -"#; - - let limit_key = Self::limit_key(key); - let usage_key = Self::usage_key(key); - - let result: i64 = self - .redis - .eval::( - LUA_QUOTA_CHECK_AND_CONSUME, - vec![limit_key, usage_key], - vec![tokens.to_string()], - ) - .await - .map_err(|e| { - tracing::warn!("Atomic quota check failed: {e}"); - AppError::Internal(anyhow::anyhow!("Quota check failed")) - })?; - - match result { - -1 => Ok(u64::MAX), // No limit set (unlimited) - 0 => Err(AppError::BadRequest( - "Token quota exceeded for this month".into(), - )), - remaining => Ok(remaining as u64), - } - } - - /// Set/update quota limit for a key. - pub async fn set_limit(&self, key: &str, monthly_limit: u64) -> Result<(), AppError> { - self.redis - .set::<(), _, _>(Self::limit_key(key), monthly_limit, None, None, false) - .await - .map_err(|e| { - tracing::warn!("Quota set_limit failed: {e}"); - AppError::Internal(anyhow::anyhow!("Failed to set quota limit")) - }) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn a_quotas_limit_and_usage_share_one_cluster_slot() { - for key in [ - "8a7c0c1e-0000-0000-0000-000000000000:gpt-4o", - "k:{odd}:model", - ] { - assert_eq!( - fred::util::redis_keyslot(QuotaManager::limit_key(key).as_bytes()), - fred::util::redis_keyslot(QuotaManager::usage_key(key).as_bytes()), - "{key}" - ); - } - } -} 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/gateway/src/route_caps.rs b/crates/gateway/src/route_caps.rs new file mode 100644 index 00000000..e8a94a0c --- /dev/null +++ b/crates/gateway/src/route_caps.rs @@ -0,0 +1,342 @@ +//! Per-route capacity caps: `model_routes.rpm_cap` (requests per minute) +//! and `tpm_cap` (weighted tokens per minute). NULL = unlimited. +//! +//! A route at either cap is skipped for the request the way a +//! circuit-broken route is: the next candidate is tried. A request that +//! finds every remaining candidate at its cap is refused with 429 and a +//! `Retry-After` of when the first of them has room again. +//! +//! Counted with the rate-limit engine (`limits::sliding`): a one-minute +//! sliding window of 60 one-second buckets per counter, in Redis. +//! +//! * Requests: [`admit`] checks both caps before a request is sent to +//! the route and, when neither is reached, counts the request — one +//! atomic step, so concurrent requests on several gateway replicas +//! cannot all slip under the cap. Every attempt sent counts, a +//! failover's failed attempts included: the upstream received them. +//! * Tokens: [`record_tokens`] adds what the request used, weighted +//! like every other token limit (`limits::weight`), after the +//! response. A window an answer overshoots refuses the next request. +//! +//! The counters of one route share the Redis Cluster hash tag +//! `{route:}` so one script reads and writes both. +//! +//! When Redis fails, `security.rate_limit_fail_closed` decides, as for +//! the users' and keys' limits: failing open, the cap is not applied; +//! failing closed, the request is refused as `rate_limiter_unavailable`. + +use fred::clients::Client; +use think_watch_common::limits::RateMetric; +use think_watch_common::limits::sliding::{self, Counter}; +use uuid::Uuid; + +use crate::error::GatewayError; +use crate::router::RouteEntry; + +/// The caps count over one minute. +pub const WINDOW_SECS: i32 = 60; + +/// The Redis Cluster hash tag every counter of `route_id` carries. +pub fn hash_tag(route_id: Uuid) -> String { + format!("{{route:{route_id}}}") +} + +/// The Redis key of one of a route's counters. +pub fn counter_key(route_id: Uuid, metric: RateMetric) -> String { + format!( + "ratelimit:{}:ai_gateway:route:{route_id}:{}:{WINDOW_SECS}", + hash_tag(route_id), + metric.as_str() + ) +} + +/// One route's caps, as the selection and the post-response accounting +/// read them. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct RouteCaps { + pub route_id: Uuid, + /// Requests per minute. + pub rpm: Option, + /// Weighted tokens per minute. + pub tpm: Option, +} + +impl RouteCaps { + pub fn of(entry: &RouteEntry) -> Self { + Self { + route_id: entry.route_id, + rpm: entry.rpm_cap, + tpm: entry.tpm_cap, + } + } + + /// The counters to check before a request: the request count, which + /// an admitted request adds one to, and the token count, only read. + fn counters(&self) -> Vec<(RateMetric, Counter)> { + let bucket_secs = sliding::bucket_secs(WINDOW_SECS); + [ + (RateMetric::Requests, self.rpm, 1), + (RateMetric::Tokens, self.tpm, 0), + ] + .into_iter() + .filter_map(|(metric, cap, charge)| { + cap.map(|cap| { + ( + metric, + Counter { + key: counter_key(self.route_id, metric), + bucket_secs, + max_count: i64::from(cap), + charge, + }, + ) + }) + }) + .collect() + } +} + +/// What [`admit`] decided for a route. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Admission { + /// Under its caps (or uncapped); the request has been counted. + Admitted, + /// At a cap: skip the route. `label` names the cap + /// (`route:requests/1m`, `route:tokens/1m`), `retry_after_secs` is + /// when it has room again. + Capped { + label: String, + retry_after_secs: u64, + }, +} + +fn label(metric: RateMetric) -> String { + format!("route:{}/1m", metric.as_str()) +} + +/// Before a request is sent to the route: refuse it if the route is at +/// either cap, otherwise count it. A route with no caps costs no Redis +/// call. +pub async fn admit( + redis: &Client, + caps: &RouteCaps, + fail_closed: bool, +) -> Result { + admit_at( + redis, + caps, + fail_closed, + chrono::Utc::now().timestamp_millis(), + ) + .await +} + +/// [`admit`] at a given time, in Unix milliseconds — for tests. +pub async fn admit_at( + redis: &Client, + caps: &RouteCaps, + fail_closed: bool, + now_ms: i64, +) -> Result { + let counters = caps.counters(); + if counters.is_empty() { + return Ok(Admission::Admitted); + } + let plain: Vec = counters.iter().map(|(_, c)| c.clone()).collect(); + match sliding::admit_counters_at(redis, &plain, now_ms).await { + Ok(outcome) if outcome.allowed => Ok(Admission::Admitted), + Ok(outcome) => { + let metric = outcome + .exceeded_index + .and_then(|i| counters.get(i)) + .map_or(RateMetric::Requests, |(m, _)| *m); + metrics::counter!("gateway_route_capped_total", "metric" => metric.as_str()) + .increment(1); + Ok(Admission::Capped { + label: label(metric), + retry_after_secs: outcome.retry_after_secs, + }) + } + Err(e) if fail_closed => { + metrics::counter!("gateway_rate_limiter_fail_closed_total").increment(1); + tracing::error!( + route_id = %caps.route_id, + "route cap check failed: {e}; failing closed per security.rate_limit_fail_closed" + ); + Err(GatewayError::limiter_unavailable( + "rate_limiter_unavailable", + )) + } + Err(e) => { + metrics::counter!("gateway_rate_limiter_fail_open_total").increment(1); + tracing::warn!( + route_id = %caps.route_id, + "route cap check failed: {e}; failing open" + ); + Ok(Admission::Admitted) + } + } +} + +/// After the response: add the request's weighted tokens to the route's +/// token counter, when the route has a token cap. Errors are logged — +/// the answer has already gone out. +pub async fn record_tokens(redis: &Client, caps: &RouteCaps, weighted_tokens: i64) { + record_tokens_at( + redis, + caps, + weighted_tokens, + chrono::Utc::now().timestamp_millis(), + ) + .await; +} + +/// [`record_tokens`] at a given time, in Unix milliseconds — for tests. +pub async fn record_tokens_at(redis: &Client, caps: &RouteCaps, weighted_tokens: i64, now_ms: i64) { + if caps.tpm.is_none() || weighted_tokens <= 0 { + return; + } + let counter = ( + counter_key(caps.route_id, RateMetric::Tokens), + sliding::bucket_secs(WINDOW_SECS), + ); + if let Err(e) = sliding::record_counters_at(redis, &[counter], weighted_tokens, now_ms).await { + tracing::warn!(route_id = %caps.route_id, "route token accounting failed: {e}"); + } +} + +/// The routes a selection skipped at their caps, and what to answer when +/// nothing else could take the request. +#[derive(Debug, Default)] +pub(crate) struct CappedRoutes { + routes: Vec, + /// The cap that frees first, and when. + soonest: Option<(String, u64)>, +} + +impl CappedRoutes { + pub(crate) fn add(&mut self, route_id: Uuid, label: String, retry_after_secs: u64) { + self.routes.push(route_id); + if self + .soonest + .as_ref() + .is_none_or(|(_, secs)| retry_after_secs < *secs) + { + self.soonest = Some((label, retry_after_secs)); + } + } + + pub(crate) fn routes(&self) -> &[Uuid] { + &self.routes + } + + /// 429 when routes were skipped at their caps: try again when the + /// first of them has room. `None` when none was. + pub(crate) fn refusal(&self) -> Option { + self.soonest + .as_ref() + .map(|(label, secs)| GatewayError::rate_limited(label.clone(), *secs)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn caps(rpm: Option, tpm: Option) -> RouteCaps { + RouteCaps { + route_id: Uuid::new_v4(), + rpm, + tpm, + } + } + + /// Redis Cluster runs a script only when all its keys share a slot: + /// a route's two counters do, by their tag. + #[test] + fn a_routes_counters_share_one_slot() { + use fred::util::redis_keyslot; + let route = Uuid::new_v4(); + let requests = counter_key(route, RateMetric::Requests); + let tokens = counter_key(route, RateMetric::Tokens); + assert_eq!( + requests, + format!("ratelimit:{{route:{route}}}:ai_gateway:route:{route}:requests:60") + ); + assert_eq!( + redis_keyslot(requests.as_bytes()), + redis_keyslot(tokens.as_bytes()) + ); + assert_eq!( + redis_keyslot(requests.as_bytes()), + redis_keyslot(format!("route:{route}").as_bytes()) + ); + } + + /// Only the caps a route has are counted; an admitted request adds + /// one to the request count and nothing to the token count. + #[test] + fn only_set_caps_are_counted() { + assert!(caps(None, None).counters().is_empty()); + + let both = caps(Some(5), Some(1_000)).counters(); + assert_eq!(both.len(), 2); + let (metric, requests) = &both[0]; + assert_eq!(*metric, RateMetric::Requests); + assert_eq!((requests.max_count, requests.charge), (5, 1)); + assert_eq!(requests.bucket_secs, 1); + let (metric, tokens) = &both[1]; + assert_eq!(*metric, RateMetric::Tokens); + assert_eq!((tokens.max_count, tokens.charge), (1_000, 0)); + + let tokens_only = caps(None, Some(10)).counters(); + assert_eq!(tokens_only.len(), 1); + assert_eq!(tokens_only[0].0, RateMetric::Tokens); + } + + /// A never-connected client times every command out: Redis down. + fn unreachable_redis() -> Client { + use fred::types::Builder; + use fred::types::config::Config; + Builder::from_config(Config::from_url("redis://127.0.0.1:1").unwrap()) + .with_performance_config(|c| { + c.default_command_timeout = std::time::Duration::from_millis(50) + }) + .build() + .unwrap() + } + + #[tokio::test] + async fn an_uncapped_route_is_admitted_without_redis() { + let admitted = admit(&unreachable_redis(), &caps(None, None), true).await; + assert_eq!(admitted.unwrap(), Admission::Admitted); + } + + #[tokio::test] + async fn an_unreachable_store_follows_the_fail_closed_setting() { + let capped = caps(Some(1), None); + let open = admit(&unreachable_redis(), &capped, false).await; + assert_eq!(open.unwrap(), Admission::Admitted); + + let closed = admit(&unreachable_redis(), &capped, true) + .await + .unwrap_err(); + assert_eq!(closed.status_code(), 429); + assert_eq!(closed.to_string(), "Rate limited: rate_limiter_unavailable"); + } + + /// Every candidate at its cap: 429, retry when the first frees. + #[test] + fn the_refusal_waits_for_the_first_route_to_free() { + let mut skipped = CappedRoutes::default(); + assert!(skipped.refusal().is_none()); + let (a, b) = (Uuid::new_v4(), Uuid::new_v4()); + skipped.add(a, label(RateMetric::Requests), 40); + skipped.add(b, label(RateMetric::Tokens), 12); + assert_eq!(skipped.routes(), &[a, b]); + let refusal = skipped.refusal().unwrap(); + assert_eq!(refusal.status_code(), 429); + assert_eq!(refusal.retry_after_secs(), Some(12)); + assert_eq!(refusal.to_string(), "Rate limited: route:tokens/1m"); + } +} diff --git a/crates/mcp-gateway/src/lifecycle/stages/mod.rs b/crates/mcp-gateway/src/lifecycle/stages/mod.rs index ffdb5d09..8662379a 100644 --- a/crates/mcp-gateway/src/lifecycle/stages/mod.rs +++ b/crates/mcp-gateway/src/lifecycle/stages/mod.rs @@ -1,6 +1,6 @@ //! MCP-specific lifecycle stages. Each stage takes //! `Authorized` (the state every stage past -//! [`common::lifecycle::stages::check_access`] sees) and either +//! [`common::lifecycle::stages::check_limits`] sees) and either //! returns it unchanged on a continue, or returns //! `Err(JsonRpcResponse)` on a short-circuit (cache hit, //! breaker-open, etc.). Audit emission for short-circuits is diff --git a/crates/mcp-gateway/src/proxy.rs b/crates/mcp-gateway/src/proxy.rs index d268bdcc..762c358f 100644 --- a/crates/mcp-gateway/src/proxy.rs +++ b/crates/mcp-gateway/src/proxy.rs @@ -545,9 +545,10 @@ impl McpProxy { let allowed_mcp_tools = ctx.allowed_mcp_tools; let trace_id = ctx.trace_id; // Resolve params + tool target up front so we know which MCP - // server this call belongs to. The rate-limit subjects need - // the server id, so the rate-limit gate runs AFTER the - // server lookup but BEFORE access control + the real call. + // server this call belongs to. Then access control, then the + // rate-limit gate, then the real call: the gate charges the + // request counters, so a tool the key may not call is refused + // before it counts. let params = match &request.params { Some(p) => p, @@ -584,13 +585,6 @@ impl McpProxy { } }; - // Rate-limit pre-flight — runs the lifecycle's shared - // `check_limits` stage. The stage emits its own audit row - // on short-circuit and the response shape comes from - // `McpSurface::rate_limited_response` so the deny payload - // stays identical to the pre-migration version. We bind - // the `request.id` onto the response after the fact - // because the stage doesn't know the wire-level id. let limits = crate::lifecycle::rate_limits( user_id, surface_constraints, @@ -609,35 +603,41 @@ impl McpProxy { trace_id.to_owned(), ctx.ip_address.map(|s| s.to_owned()), ); - let limits_checked = match think_watch_common::lifecycle::stages::check_limits::< + + // Access control via the shared `check_access` stage. The + // candidate is the namespaced tool name we extracted at the + // top of the handler; the surface impl + // (`McpSurface::is_access_allowed`) calls `is_tool_allowed` + // with the identity's `allowed_mcp_tools` patterns. Each stage + // emits its own audit row on short-circuit; we bind the + // inbound JSON-RPC `id` onto the response after the fact so + // the client can correlate the deny to its request — the stage + // doesn't know the wire-level id. + let access_checked = match think_watch_common::lifecycle::stages::check_access::< crate::lifecycle::McpSurface, - >( - raw, - &limits.rules, - limits.owner, - &self.redis, - fail_closed, - &self.audit, - ) + >(raw, namespaced_name, &self.audit) .await { Ok(s) => s, Err(mut resp) => { - // Bind the inbound JSON-RPC `id` so the client - // can correlate the deny response to its request. resp.id = request.id; return HandleOutcome::Buffered(resp); } }; - // Access control via the shared `check_access` stage. The - // candidate is the namespaced tool name we extracted at the - // top of the handler; the surface impl - // (`McpSurface::is_access_allowed`) calls `is_tool_allowed` - // with the identity's `allowed_mcp_tools` patterns. - let authorized = match think_watch_common::lifecycle::stages::check_access::< + // Rate-limit pre-flight — the lifecycle's shared `check_limits` + // stage, on a call the key may make. The response shape comes + // from `McpSurface::rate_limited_response`. + let authorized = match think_watch_common::lifecycle::stages::check_limits::< crate::lifecycle::McpSurface, - >(limits_checked, namespaced_name, &self.audit) + >( + access_checked, + &limits.rules, + limits.owner, + &self.redis, + fail_closed, + &self.audit, + ) .await { Ok(s) => s, diff --git a/crates/server/src/app.rs b/crates/server/src/app.rs index b46b509c..6358fb7c 100644 --- a/crates/server/src/app.rs +++ b/crates/server/src/app.rs @@ -22,7 +22,6 @@ use think_watch_common::dynamic_config::DynamicConfig; use think_watch_gateway::cache::ResponseCache; use think_watch_gateway::model_mapping::ModelMapper; use think_watch_gateway::proxy::{self as gateway_proxy, GatewayState}; -use think_watch_gateway::quota::QuotaManager; use think_watch_gateway::router::{ModelRouter, RouteEntry}; use think_watch_mcp_gateway::proxy::McpProxy; use think_watch_mcp_gateway::session::SessionManager; @@ -234,7 +233,6 @@ pub async fn create_gateway_app(_config: &AppConfig, state: AppState) -> anyhow: model_mapper: Arc::new(ModelMapper::new()), // Share the hot-swappable guards with the gateway state. guards: state.guards.clone(), - quota: Arc::new(QuotaManager::new(state.redis.clone())), cache: Arc::new(ResponseCache::new( state.redis.clone(), state.dynamic_config.clone(), @@ -251,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, @@ -292,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 @@ -400,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..3482061d --- /dev/null +++ b/crates/server/src/handlers/key_usage.rs @@ -0,0 +1,581 @@ +// ============================================================================ +// 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}"); + unavailable_response() +} + +/// 503: the key's limits or their counters could not be read. Also what +/// `require_api_key_to_read` answers when the limits cannot be loaded — +/// an empty `limits` list would say the key has none. +pub(crate) fn unavailable_response() -> Response { + let body = tw_dialect::convert::error_body( + tw_dialect::ir::Dialect::Chat, + StatusCode::SERVICE_UNAVAILABLE.as_u16(), + "The key's limits and usage 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 1f13aad4..1016bba0 100644 --- a/crates/server/src/middleware/api_key_auth.rs +++ b/crates/server/src/middleware/api_key_auth.rs @@ -10,6 +10,8 @@ use think_watch_auth::rbac; use think_watch_gateway::proxy::GatewayRequestIdentity; use think_watch_mcp_gateway::transport::streamable_http::McpRequestIdentity; +use think_watch_common::limits::SurfaceConstraints; + use crate::app::AppState; /// Intersect a per-API-key allow-list with a per-role allow-list. @@ -73,13 +75,27 @@ fn mcp_entry_covers(general: &str, specific: &str) -> bool { is_tool_allowed(Some(&[general.to_string()]), specific) } +/// The format an AI-gateway caller reads errors in, from the path it +/// called. +fn client_dialect(path: &str) -> tw_dialect::ir::Dialect { + use tw_dialect::ir::Dialect; + if path == "/v1/messages" { + Dialect::Anthropic + } else if path.starts_with("/v1beta/") || path.starts_with("/v1/models/") { + Dialect::Gemini + } else if path == "/v1/responses" { + Dialect::Responses + } else { + Dialect::Chat + } +} + /// 403 for a key whose owner holds no role granting `surface`'s /// `*_gateway:use`. The body is in the shape the caller's SDK reads: /// the protocol's error object on the AI gateway (every one of them /// carries `error.message`), a JSON-RPC error on the MCP gateway. fn gateway_use_refused(surface: &str, path: &str) -> Response { use axum::response::IntoResponse; - use tw_dialect::ir::Dialect; let permission = format!("{surface}:use"); let message = @@ -92,18 +108,13 @@ fn gateway_use_refused(surface: &str, path: &str) -> Response { }); ("application/json", body.to_string().into_bytes()) } else { - let client = if path == "/v1/messages" { - Dialect::Anthropic - } else if path.starts_with("/v1beta/") || path.starts_with("/v1/models/") { - Dialect::Gemini - } else if path == "/v1/responses" { - Dialect::Responses - } else { - Dialect::Chat - }; ( "application/json", - tw_dialect::convert::error_body(client, StatusCode::FORBIDDEN.as_u16(), &message), + tw_dialect::convert::error_body( + client_dialect(path), + StatusCode::FORBIDDEN.as_u16(), + &message, + ), ) }; ( @@ -114,6 +125,90 @@ fn gateway_use_refused(surface: &str, path: &str) -> Response { .into_response() } +/// The limits a request is held to: its owner's (role limits with the +/// user's overrides) and the calling key's own, counted apart. +async fn load_request_limits( + db: &sqlx::PgPool, + user_id: uuid::Uuid, + lineage_id: uuid::Uuid, +) -> Result<(SurfaceConstraints, SurfaceConstraints), sqlx::Error> { + let user = rbac::compute_user_surface_constraints(db, user_id).await?; + let key = rbac::compute_key_surface_constraints(db, lineage_id).await?; + Ok((user, key)) +} + +/// The label a request refused for want of its limits carries — the +/// limits could not be loaded, so they could not be checked. +const LIMITS_UNAVAILABLE: &str = "limits_unavailable"; + +/// What a request goes on with when loading its limits failed, decided +/// by `security.rate_limit_fail_closed` — the setting the rate limiter +/// and the budget gate follow when Redis fails. Failing closed, `None`: +/// the request is refused ([`limits_unavailable`]). Failing open, no +/// limits for this request, and a warning that says why. +fn limits_or_refusal( + loaded: Result<(SurfaceConstraints, SurfaceConstraints), sqlx::Error>, + fail_closed: bool, + api_key_id: uuid::Uuid, +) -> Option<(SurfaceConstraints, SurfaceConstraints)> { + let error = match loaded { + Ok(limits) => return Some(limits), + Err(e) => e, + }; + if fail_closed { + metrics::counter!("gateway_limits_load_fail_closed_total").increment(1); + tracing::error!( + %api_key_id, + error = %error, + "loading the request's limits failed; failing closed per security.rate_limit_fail_closed" + ); + None + } else { + metrics::counter!("gateway_limits_load_fail_open_total").increment(1); + tracing::warn!( + %api_key_id, + error = %error, + "loading the request's limits failed; failing open, the request runs without limits" + ); + Some(Default::default()) + } +} + +/// The refusal for a request whose limits could not be loaded while +/// failing closed: the same answer the gateways give when the rate +/// limiter itself is unreachable — 429 with `Retry-After`, labelled +/// [`LIMITS_UNAVAILABLE`], in the caller's format. +fn limits_unavailable(surface: &str, path: &str) -> Response { + use axum::response::IntoResponse; + let refusal = think_watch_gateway::error::GatewayError::limiter_unavailable(LIMITS_UNAVAILABLE); + if surface == "mcp_gateway" { + let retry_after = refusal.retry_after_secs().unwrap_or(30); + let body = serde_json::json!({ + "jsonrpc": "2.0", + "id": null, + "error": { + "code": think_watch_mcp_gateway::proxy::INVALID_REQUEST, + "message": refusal.to_string(), + }, + }); + return ( + StatusCode::TOO_MANY_REQUESTS, + [ + ( + axum::http::header::CONTENT_TYPE, + "application/json".to_string(), + ), + (axum::http::header::RETRY_AFTER, retry_after.to_string()), + ], + body.to_string(), + ) + .into_response(); + } + think_watch_gateway::proxy::GatewayErrorResponse::from(refusal) + .in_dialect(client_dialect(path)) + .into_response() +} + /// The key a client presented, wherever its SDK puts it. /// /// `Authorization: Bearer` (OpenAI's SDKs, Claude Code with @@ -154,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`. @@ -166,6 +283,27 @@ 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. Limits that can't be loaded make the +/// read a 503 whatever `security.rate_limit_fail_closed` says: failing +/// open lets a call run without them, but a read would report none. +pub fn require_api_key_to_read( + surface: &'static str, +) -> impl Fn(State, Request, Next) -> AuthFuture + Clone { + authenticate(surface, KeyUse::Read) +} + +fn authenticate( + surface: &'static str, + key_use: KeyUse, ) -> impl Fn(State, Request, Next) -> AuthFuture + Clone { move |State(state): State, mut request: Request, next: Next| { Box::pin(async move { @@ -257,7 +395,7 @@ pub fn require_api_key( // records its own). Every return below produces a response // or an auth refusal, so the guard is disarmed after all of // them; only a dropped future leaves it armed. - let cancel = (surface == "ai_gateway").then(|| { + let cancel = (surface == "ai_gateway" && key_use == KeyUse::Call).then(|| { think_watch_gateway::proxy::EarlyCancel::arm( state.audit.clone(), GatewayRequestIdentity { @@ -271,18 +409,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 @@ -292,7 +432,7 @@ pub fn require_api_key( // can gate per-tool access without re-querying the DB, and // the aggregated `surface_constraints` JSON so the gateway // hot path has rate limits + budgets without further lookups. - let (role_limits, user_roles, surface_constraints, key_constraints) = if let Some(uid) = row.user_id { + let (role_limits, user_roles, loaded_limits) = if let Some(uid) = row.user_id { let limits = rbac::compute_user_resource_limits(&state.db, uid) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; @@ -301,15 +441,10 @@ pub fn require_api_key( .unwrap_or_default(); // The user's limits and the key's own, kept apart: // the gateway counts each on its own counters and - // checks both. - let constraints = rbac::compute_user_surface_constraints(&state.db, uid) - .await - .unwrap_or_default(); - let key_constraints = - rbac::compute_key_surface_constraints(&state.db, row.lineage_id) - .await - .unwrap_or_default(); - (limits, names, constraints, key_constraints) + // checks both. A failure to load them is decided + // below, once the request is known to be allowed. + let loaded = load_request_limits(&state.db, uid, row.lineage_id).await; + (limits, names, loaded) } else { // A key without an owner (its user row was removed // and `user_id` set NULL) has no roles to grant @@ -319,8 +454,7 @@ pub fn require_api_key( ( rbac::UserResourceLimits::none(), Vec::new(), - think_watch_common::limits::SurfaceConstraints::default(), - think_watch_common::limits::SurfaceConstraints::default(), + Ok(Default::default()), ) }; @@ -342,6 +476,51 @@ pub fn require_api_key( return Ok(gateway_use_refused(surface, request.uri().path())); } + // A read about the key reports the limits that hold its + // requests. Without them it has nothing true to report — + // an empty list would say the key has none — so it is + // unavailable, whatever the setting below lets a call do. + if key_use == KeyUse::Read + && let Err(e) = &loaded_limits + { + tracing::warn!( + api_key_id = %row.id, + error = %e, + "loading the key's limits failed; its usage can't be reported" + ); + return Ok(crate::handlers::key_usage::unavailable_response()); + } + + // Limits that could not be loaded are not "no limits": + // the request is refused or let through without them as + // `security.rate_limit_fail_closed` says. + let fail_closed = match &loaded_limits { + Ok(_) => false, + Err(_) => state.dynamic_config.rate_limit_fail_closed().await, + }; + let Some((surface_constraints, key_constraints)) = + limits_or_refusal(loaded_limits, fail_closed, row.id) + else { + let user_id = row.user_id.map(|u| u.to_string()); + let api_key_id = row.id.to_string(); + let lineage_id = row.lineage_id.to_string(); + state.audit.log( + think_watch_common::audit::AuditActor::audit( + &think_watch_common::audit::GatewayActor { + user_id: user_id.as_deref(), + user_email: None, + api_key_id: Some(&api_key_id), + api_key_lineage_id: Some(&lineage_id), + ip: None, + session_id: None, + }, + LIMITS_UNAVAILABLE, + ) + .detail(serde_json::json!({ "surface": surface })), + ); + return Ok(limits_unavailable(surface, request.uri().path())); + }; + let merged_models = intersect_allowlists( row.allowed_models.clone(), role_limits.allowed_models, @@ -385,6 +564,10 @@ pub fn require_api_key( surface_constraints: surface_constraints.clone(), key_constraints: key_constraints.clone(), ip_address: client_ip.clone(), + key_expires_at: stops_authenticating_at( + row.expires_at, + row.grace_period_ends_at, + ), }; // The MCP transport handlers expect their own typed @@ -492,6 +675,81 @@ mod tests { ); } + /// Limits that failed to load are not "no limits": failing closed + /// refuses the request, failing open runs it without them; limits + /// that loaded are used as they are either way. + #[test] + fn a_failed_limits_load_follows_the_fail_closed_setting() { + let key = uuid::Uuid::new_v4(); + let failed = || Err(sqlx::Error::PoolTimedOut); + assert_eq!(limits_or_refusal(failed(), true, key), None); + assert_eq!( + limits_or_refusal(failed(), false, key), + Some((SurfaceConstraints::default(), SurfaceConstraints::default())) + ); + + let user = SurfaceConstraints { + ai_gateway: Some(think_watch_common::limits::SurfaceBlock::default()), + mcp_gateway: None, + }; + let loaded = || Ok((user.clone(), SurfaceConstraints::default())); + for fail_closed in [true, false] { + assert_eq!( + limits_or_refusal(loaded(), fail_closed, key), + Some((user.clone(), SurfaceConstraints::default())) + ); + } + } + + /// The refusal is the limiter-unavailable answer: 429, a + /// `Retry-After`, the label, in the caller's format. + #[tokio::test] + async fn the_refusal_is_a_429_in_the_callers_format() { + async fn read(resp: Response) -> (u16, Option, serde_json::Value) { + let status = resp.status().as_u16(); + let retry = resp + .headers() + .get(axum::http::header::RETRY_AFTER) + .map(|v| v.to_str().unwrap().to_string()); + let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + (status, retry, serde_json::from_slice(&bytes).unwrap()) + } + + let (status, retry, body) = + read(limits_unavailable("ai_gateway", "/v1/chat/completions")).await; + assert_eq!((status, retry.as_deref()), (429, Some("30"))); + assert_eq!(body["error"]["message"], "Rate limited: limits_unavailable"); + + let (status, _, body) = read(limits_unavailable("ai_gateway", "/v1/messages")).await; + assert_eq!(status, 429); + assert_eq!(body["type"], "error"); + assert_eq!(body["error"]["message"], "Rate limited: limits_unavailable"); + + let (status, retry, body) = read(limits_unavailable("mcp_gateway", "/mcp")).await; + assert_eq!((status, retry.as_deref()), (429, Some("30"))); + assert_eq!(body["jsonrpc"], "2.0"); + 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/src/redis_scripts.rs b/crates/test-support/src/redis_scripts.rs index 5109a280..1d16d2d8 100644 --- a/crates/test-support/src/redis_scripts.rs +++ b/crates/test-support/src/redis_scripts.rs @@ -1,6 +1,7 @@ //! The rate-limit scripts, run against whatever Redis a test hands in — //! one node in `tests/limits.rs`, a Redis Cluster in -//! `tests/redis_cluster.rs`. +//! `tests/redis_cluster.rs`. The route caps run on the same scripts +//! with their own keys. use uuid::Uuid; @@ -77,3 +78,41 @@ pub async fn exercise_the_limit_scripts(redis: &fred::clients::Client) { assert_eq!(redis.hlen::(&requests.key).await.unwrap(), 1); assert_eq!(redis.hlen::(&tokens.key).await.unwrap(), 0); } + +/// A route's caps on the same scripts: the request cap fills and counts +/// nothing more once full, recorded tokens close the token cap, and a +/// route's two counters run in one script (one hash slot). +pub async fn exercise_the_route_cap_scripts(redis: &fred::clients::Client) { + use think_watch_gateway::route_caps::{self, Admission, RouteCaps}; + + let caps = RouteCaps { + route_id: Uuid::new_v4(), + rpm: Some(2), + tpm: Some(100), + }; + let t0: i64 = 1_800_000_000_000; + let admit = |ms: i64| route_caps::admit_at(redis, &caps, true, t0 + ms); + + assert_eq!(admit(0).await.unwrap(), Admission::Admitted); + assert_eq!(admit(10_000).await.unwrap(), Admission::Admitted); + // Full: the first request leaves the window at t0 + 60 s. + assert_eq!( + admit(20_000).await.unwrap(), + Admission::Capped { + label: "route:requests/1m".into(), + retry_after_secs: 40, + } + ); + // Room again for one request; the tokens then fill the token cap. + assert_eq!(admit(60_000).await.unwrap(), Admission::Admitted); + route_caps::record_tokens_at(redis, &caps, 150, t0 + 60_000).await; + assert_eq!( + admit(65_000).await.unwrap(), + Admission::Capped { + label: "route:tokens/1m".into(), + retry_after_secs: 55, + } + ); + // A minute on, both windows are clear. + assert_eq!(admit(121_000).await.unwrap(), Admission::Admitted); +} 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/limit_enforcement.rs b/crates/test-support/tests/limit_enforcement.rs new file mode 100644 index 00000000..23ab93bf --- /dev/null +++ b/crates/test-support/tests/limit_enforcement.rs @@ -0,0 +1,663 @@ +//! Where the limits used to leak, pinned shut: +//! +//! - a limits store that cannot be read follows +//! `security.rate_limit_fail_closed` — the database the limits are +//! loaded from, and the Redis a budget is read from — instead of +//! reading as "no limits", and `GET /v1/usage` does not report +//! limits it could not load as none; +//! - a route's `rpm_cap` / `tpm_cap` is enforced: a route at its cap +//! is skipped, and with every route capped the request gets 429; +//! - a request the key may not make (its model, its MCP tool) is +//! refused before it counts against a request limit; +//! - an answer from the response cache counts its tokens. + +use fred::interfaces::{HashesInterface, KeysInterface}; +use think_watch_common::limits::{ + RateLimitSubject, RateMetric, Surface as LimitSurface, budget, sliding, +}; +use think_watch_test_support::client::TestResponse; +use think_watch_test_support::prelude::*; + +const MODEL: &str = "gpt-test"; + +fn chat(model: &str) -> Json { + json!({"model": model, "messages": [{"role": "user", "content": "x"}]}) +} + +/// A request the response cache does not answer (it keeps only +/// deterministic ones), so every one reaches route selection. +fn uncached(model: &str) -> Json { + json!({"model": model, "temperature": 1, + "messages": [{"role": "user", "content": "x"}]}) +} + +/// A user with gateway access, an OpenAI mock routed as `MODEL`, and a +/// key. Returns the key and the user. +async fn seed(app: &TestApp, allowed_models: Option<&[&str]>) -> (String, Uuid) { + let user = fixtures::create_random_user(&app.db).await.unwrap(); + let mock = MockProvider::openai_chat_ok(MODEL).await; + let provider = fixtures::create_provider( + &app.db, + &unique_name("enf-prov"), + "openai", + &mock.uri(), + None, + ) + .await + .unwrap(); + Box::leak(Box::new(mock)); + fixtures::create_model_and_route(&app.db, provider.id, MODEL) + .await + .unwrap(); + app.rebuild_gateway_router().await; + let key = fixtures::create_api_key( + &app.db, + user.user.id, + &unique_name("enf-key"), + &["ai_gateway"], + allowed_models, + None, + ) + .await + .unwrap(); + (key.plaintext, user.user.id) +} + +fn gateway(app: &TestApp, key: &str) -> TestClient { + let gw = app.gateway_client(); + gw.set_bearer(key); + gw +} + +fn retry_after(r: &TestResponse) -> u64 { + r.headers + .get("retry-after") + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.parse().ok()) + .unwrap_or_else(|| panic!("a 429 carries Retry-After, headers: {:?}", r.headers)) +} + +/// What a user's `requests` rule on `surface` has counted in its window. +async fn user_requests(app: &TestApp, user_id: Uuid, surface: LimitSurface) -> i64 { + let rule = think_watch_common::limits::RateLimitRule { + id: Uuid::nil(), + subject_kind: RateLimitSubject::User, + subject_id: user_id, + surface, + metric: RateMetric::Requests, + window_secs: 60, + max_count: 1, + enabled: true, + expires_at: None, + reason: None, + created_by: None, + }; + sliding::current_count( + &app.state.redis, + &sliding::ResolvedRule::new(&rule, user_id), + ) + .await +} + +// ---------------------------------------------------------------------------- +// The limits can't be loaded +// ---------------------------------------------------------------------------- + +/// Loading the limits fails: the rule table is gone from under the +/// running server. +async fn break_the_limits_table(app: &TestApp) { + sqlx::query("DROP TABLE rate_limit_rules CASCADE") + .execute(&app.db) + .await + .unwrap(); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn limits_that_cannot_be_loaded_refuse_when_failing_closed() { + let app = TestApp::spawn().await; + let (key, _) = seed(&app, None).await; + app.set_setting("security.rate_limit_fail_closed", json!(true)) + .await; + let gw = gateway(&app, &key); + gw.post("/v1/chat/completions", chat(MODEL)) + .await + .unwrap() + .assert_ok(); + + break_the_limits_table(&app).await; + let r = gw.post("/v1/chat/completions", chat(MODEL)).await.unwrap(); + assert_eq!(r.status.as_u16(), 429, "body={}", r.text()); + assert_eq!(retry_after(&r), 30); + let body: Json = r.json().unwrap(); + assert_eq!(body["error"]["message"], "Rate limited: limits_unavailable"); + + // An Anthropic caller gets it in its own format. + let r = gw + .post( + "/v1/messages", + json!({"model": MODEL, "max_tokens": 8, + "messages": [{"role": "user", "content": "x"}]}), + ) + .await + .unwrap(); + assert_eq!(r.status.as_u16(), 429, "body={}", r.text()); + let body: Json = r.json().unwrap(); + assert_eq!(body["type"], "error"); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn limits_that_cannot_be_loaded_let_the_request_through_when_failing_open() { + let app = TestApp::spawn().await; + let (key, _) = seed(&app, None).await; + app.set_setting("security.rate_limit_fail_closed", json!(false)) + .await; + break_the_limits_table(&app).await; + gateway(&app, &key) + .post("/v1/chat/completions", chat(MODEL)) + .await + .unwrap() + .assert_ok(); +} + +/// `GET /v1/usage` reports the limits that hold the key. Limits that +/// can't be loaded leave it nothing true to report — an empty list would +/// say the key has none — so it is a 503 whichever way the setting goes, +/// while a model request follows the setting. +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn v1_usage_with_limits_that_cannot_be_loaded_is_unavailable() { + let app = TestApp::spawn().await; + let (key, _) = seed(&app, None).await; + let gw = gateway(&app, &key); + let before = gw.get("/v1/usage").await.unwrap(); + before.assert_ok(); + + break_the_limits_table(&app).await; + for fail_closed in [false, true] { + app.set_setting("security.rate_limit_fail_closed", json!(fail_closed)) + .await; + let r = gw.get("/v1/usage").await.unwrap(); + assert_eq!( + r.status.as_u16(), + 503, + "fail_closed={fail_closed}: {}", + r.text() + ); + let body: Json = r.json().unwrap(); + assert_eq!( + body["error"]["message"], + "The key's limits and usage are unavailable." + ); + let model = gw.post("/v1/chat/completions", chat(MODEL)).await.unwrap(); + let expected = if fail_closed { 429 } else { 200 }; + assert_eq!(model.status.as_u16(), expected, "body={}", model.text()); + } +} + +// ---------------------------------------------------------------------------- +// The budget can't be read +// ---------------------------------------------------------------------------- + +/// A daily budget whose counter Redis refuses to read as a number: the +/// key holds a hash, so GET fails with WRONGTYPE. +async fn unreadable_budget(app: &TestApp, user_id: Uuid) { + fixtures::create_budget_cap(&app.db, "user", user_id, "daily", 1_000_000) + .await + .unwrap(); + let key = budget::build_key("user", user_id, "daily", chrono::Utc::now()); + let _: i64 = app.state.redis.hset(&key, ("x", 1)).await.unwrap(); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_budget_that_cannot_be_read_refuses_when_failing_closed() { + let app = TestApp::spawn().await; + let (key, user_id) = seed(&app, None).await; + app.set_setting("security.rate_limit_fail_closed", json!(true)) + .await; + unreadable_budget(&app, user_id).await; + let r = gateway(&app, &key) + .post("/v1/chat/completions", chat(MODEL)) + .await + .unwrap(); + assert_eq!(r.status.as_u16(), 429, "body={}", r.text()); + assert_eq!(retry_after(&r), 30); + let body: Json = r.json().unwrap(); + assert_eq!(body["error"]["message"], "Rate limited: budget_unavailable"); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_budget_that_cannot_be_read_lets_the_request_through_when_failing_open() { + let app = TestApp::spawn().await; + let (key, user_id) = seed(&app, None).await; + app.set_setting("security.rate_limit_fail_closed", json!(false)) + .await; + unreadable_budget(&app, user_id).await; + gateway(&app, &key) + .post("/v1/chat/completions", chat(MODEL)) + .await + .unwrap() + .assert_ok(); +} + +// ---------------------------------------------------------------------------- +// Route caps +// ---------------------------------------------------------------------------- + +async fn set_route_caps(app: &TestApp, provider_id: Uuid, rpm: Option, tpm: Option) { + sqlx::query("UPDATE model_routes SET rpm_cap = $2, tpm_cap = $3 WHERE provider_id = $1") + .bind(provider_id) + .bind(rpm) + .bind(tpm) + .execute(&app.db) + .await + .unwrap(); +} + +async fn route_id(app: &TestApp, provider_id: Uuid) -> Uuid { + sqlx::query_scalar("SELECT id FROM model_routes WHERE provider_id = $1") + .bind(provider_id) + .fetch_one(&app.db) + .await + .unwrap() +} + +/// One provider serving `model` through one route, with a key. +async fn one_route(app: &TestApp, model: &str) -> (Uuid, String) { + let user = fixtures::create_random_user(&app.db).await.unwrap(); + let mock = MockProvider::openai_chat_ok(model).await; + let provider = fixtures::create_provider( + &app.db, + &unique_name("cap-prov"), + "openai", + &mock.uri(), + None, + ) + .await + .unwrap(); + Box::leak(Box::new(mock)); + fixtures::create_model_and_route(&app.db, provider.id, model) + .await + .unwrap(); + let key = fixtures::create_api_key( + &app.db, + user.user.id, + &unique_name("cap-key"), + &["ai_gateway"], + None, + None, + ) + .await + .unwrap(); + (provider.id, key.plaintext) +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_route_at_its_request_cap_is_skipped_for_the_next() { + let app = TestApp::spawn().await; + let model = "cap-model"; + let capped = MockProvider::openai_chat_ok(model).await; + let spare = MockProvider::openai_chat_ok(model).await; + let p_capped = fixtures::create_provider( + &app.db, + &unique_name("capped"), + "openai", + &capped.uri(), + None, + ) + .await + .unwrap(); + let p_spare = + fixtures::create_provider(&app.db, &unique_name("spare"), "openai", &spare.uri(), None) + .await + .unwrap(); + // The spare route carries no traffic of its own (weight 0): it + // serves only when the other cannot. + fixtures::create_model_route(&app.db, p_capped.id, model, 100) + .await + .unwrap(); + fixtures::create_model_route(&app.db, p_spare.id, model, 0) + .await + .unwrap(); + set_route_caps(&app, p_capped.id, Some(1), None).await; + app.rebuild_gateway_router().await; + let user = fixtures::create_random_user(&app.db).await.unwrap(); + let key = fixtures::create_api_key( + &app.db, + user.user.id, + &unique_name("cap-key"), + &["ai_gateway"], + None, + None, + ) + .await + .unwrap(); + let gw = gateway(&app, &key.plaintext); + + for _ in 0..3 { + gw.post("/v1/chat/completions", uncached(model)) + .await + .unwrap() + .assert_ok(); + } + assert_eq!( + capped.received_requests().await.len(), + 1, + "the capped route takes one request a minute" + ); + assert_eq!( + spare.received_requests().await.len(), + 2, + "the next route serves the rest" + ); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn every_route_at_its_cap_is_a_429_with_retry_after() { + let app = TestApp::spawn().await; + let model = "cap-all-model"; + let (provider, key) = one_route(&app, model).await; + set_route_caps(&app, provider, Some(1), None).await; + app.rebuild_gateway_router().await; + let gw = gateway(&app, &key); + + gw.post("/v1/chat/completions", uncached(model)) + .await + .unwrap() + .assert_ok(); + + let r = gw + .post("/v1/chat/completions", uncached(model)) + .await + .unwrap(); + assert_eq!(r.status.as_u16(), 429, "body={}", r.text()); + let secs = retry_after(&r); + assert!((1..=60).contains(&secs), "Retry-After {secs}"); + let body: Json = r.json().unwrap(); + assert_eq!(body["error"]["message"], "Rate limited: route:requests/1m"); + + // A stream is refused the same way, before it starts. + let mut streamed = uncached(model); + streamed["stream"] = json!(true); + let r = gw.post("/v1/chat/completions", streamed).await.unwrap(); + assert_eq!(r.status.as_u16(), 429, "body={}", r.text()); + retry_after(&r); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_token_cap_counts_weighted_tokens_after_the_response() { + let app = TestApp::spawn().await; + let model = "cap-tokens-model"; + let (provider, key) = one_route(&app, model).await; + // The mock answers 7 prompt + 3 completion tokens: 2×7 + 3×3 = 23 + // weighted, over a cap of 20 that the raw 10 would stay under. + sqlx::query("UPDATE models SET input_weight = 2, output_weight = 3 WHERE model_id = $1") + .bind(model) + .execute(&app.db) + .await + .unwrap(); + set_route_caps(&app, provider, None, Some(20)).await; + app.rebuild_gateway_router().await; + let gw = gateway(&app, &key); + + gw.post("/v1/chat/completions", uncached(model)) + .await + .unwrap() + .assert_ok(); + let route = route_id(&app, provider).await; + let counter = think_watch_gateway::route_caps::counter_key(route, RateMetric::Tokens); + let buckets: std::collections::HashMap = + app.state.redis.hgetall(&counter).await.unwrap(); + assert_eq!(buckets.values().sum::(), 23, "{buckets:?}"); + let requests = think_watch_gateway::route_caps::counter_key(route, RateMetric::Requests); + let exists: i64 = app.state.redis.exists(&requests).await.unwrap(); + assert_eq!( + exists, 0, + "a route without a request cap counts no requests" + ); + + let r = gw + .post("/v1/chat/completions", uncached(model)) + .await + .unwrap(); + assert_eq!(r.status.as_u16(), 429, "body={}", r.text()); + let body: Json = r.json().unwrap(); + assert_eq!(body["error"]["message"], "Rate limited: route:tokens/1m"); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn the_route_cap_scripts_fill_free_and_record() { + let app = TestApp::spawn().await; + think_watch_test_support::redis_scripts::exercise_the_route_cap_scripts(&app.state.redis).await; +} + +// ---------------------------------------------------------------------------- +// Access before limits +// ---------------------------------------------------------------------------- + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_model_the_key_may_not_use_does_not_count_a_request() { + let app = TestApp::spawn().await; + let (key, user_id) = seed(&app, Some(&[MODEL])).await; + fixtures::create_rate_limit_rule(&app.db, "user", user_id, "ai_gateway", "requests", 60, 1) + .await + .unwrap(); + let gw = gateway(&app, &key); + + for _ in 0..3 { + let r = gw + .post("/v1/chat/completions", chat("not-for-this-key")) + .await + .unwrap(); + assert_eq!(r.status.as_u16(), 400, "body={}", r.text()); + assert!(r.text().contains("is not allowed"), "{}", r.text()); + } + assert_eq!( + user_requests(&app, user_id, LimitSurface::AiGateway).await, + 0 + ); + + gw.post("/v1/chat/completions", chat(MODEL)) + .await + .unwrap() + .assert_ok(); + assert_eq!( + user_requests(&app, user_id, LimitSurface::AiGateway).await, + 1 + ); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn an_mcp_tool_the_key_may_not_call_does_not_count_a_request() { + use axum::{Json as AxumJson, Router, routing::post}; + + let app = TestApp::spawn().await; + let user = fixtures::create_random_user(&app.db).await.unwrap(); + let upstream = Router::new().route( + "/mcp", + post(|AxumJson(req): AxumJson| async move { + AxumJson(json!({"jsonrpc": "2.0", "id": req["id"], "result": {"content": []}})) + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + let _ = axum::serve(listener, upstream).await; + }); + let server_id = fixtures::create_mcp_server_with( + &app.db, + &unique_name("order-mcp"), + "ord", + &format!("http://{addr}/mcp"), + fixtures::McpServerOpts::default(), + ) + .await + .unwrap(); + let row = sqlx::query_as::<_, think_watch_common::models::McpServer>( + "SELECT * FROM mcp_servers WHERE id = $1", + ) + .bind(server_id) + .fetch_one(&app.db) + .await + .unwrap(); + let registered = think_watch_server::mcp_runtime::build_registered_server( + &app.db, + &row, + &app.state.config.encryption_key, + ) + .await + .unwrap(); + app.state.mcp_registry.register(registered).await; + + let key = fixtures::create_api_key( + &app.db, + user.user.id, + &unique_name("order-key"), + &["mcp_gateway"], + None, + None, + ) + .await + .unwrap(); + sqlx::query("UPDATE api_keys SET allowed_mcp_tools = ARRAY['ord__allowed'] WHERE id = $1") + .bind(key.row.id) + .execute(&app.db) + .await + .unwrap(); + fixtures::create_rate_limit_rule( + &app.db, + "user", + user.user.id, + "mcp_gateway", + "requests", + 60, + 1, + ) + .await + .unwrap(); + let gw = gateway(&app, &key.plaintext); + let call = |id: i64, tool: &str| { + json!({"jsonrpc": "2.0", "id": id, "method": "tools/call", + "params": {"name": tool, "arguments": {}}}) + }; + + for id in 1..=3 { + let r: Json = gw + .post("/mcp", call(id, "ord__forbidden")) + .await + .unwrap() + .json() + .unwrap(); + assert_eq!(r["error"]["message"], "Access denied for this tool", "{r}"); + } + assert_eq!( + user_requests(&app, user.user.id, LimitSurface::McpGateway).await, + 0 + ); + let r: Json = gw + .post("/mcp", call(4, "ord__allowed")) + .await + .unwrap() + .json() + .unwrap(); + assert!(r["result"].is_object(), "{r}"); + assert_eq!( + user_requests(&app, user.user.id, LimitSurface::McpGateway).await, + 1 + ); +} + +// ---------------------------------------------------------------------------- +// Cache hits +// ---------------------------------------------------------------------------- + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_cached_answer_counts_its_tokens() { + let app = TestApp::spawn().await; + let (key, user_id) = seed(&app, None).await; + let rule = fixtures::create_rate_limit_rule( + &app.db, + "user", + user_id, + "ai_gateway", + "tokens", + 60, + 1_000_000, + ) + .await + .unwrap(); + fixtures::create_budget_cap(&app.db, "user", user_id, "daily", 1_000_000) + .await + .unwrap(); + let gw = gateway(&app, &key); + let deterministic = json!({ + "model": MODEL, + "messages": [{"role": "user", "content": "the same every time"}], + "temperature": 0, + }); + + let tokens = || async { + let rule = think_watch_common::limits::RateLimitRule { + id: rule, + subject_kind: RateLimitSubject::User, + subject_id: user_id, + surface: LimitSurface::AiGateway, + metric: RateMetric::Tokens, + window_secs: 60, + max_count: 1_000_000, + enabled: true, + expires_at: None, + reason: None, + created_by: None, + }; + let window = sliding::current_count( + &app.state.redis, + &sliding::ResolvedRule::new(&rule, user_id), + ) + .await; + let spent: Option = app + .state + .redis + .get(budget::build_key( + "user", + user_id, + "daily", + chrono::Utc::now(), + )) + .await + .unwrap(); + (window, spent.unwrap_or(0)) + }; + + let first = gw + .post("/v1/chat/completions", deterministic.clone()) + .await + .unwrap(); + first.assert_ok(); + assert_eq!(first.headers["x-cache"], "MISS"); + let (window, spent) = tokens().await; + assert!(window > 0, "the upstream answer counted"); + assert_eq!(window, spent); + + let hit = gw + .post("/v1/chat/completions", deterministic.clone()) + .await + .unwrap(); + hit.assert_ok(); + assert_eq!(hit.headers["x-cache"], "HIT"); + assert_eq!( + tokens().await, + (2 * window, 2 * spent), + "the cached answer counts the tokens it records, like the call that made it" + ); +} 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/crates/test-support/tests/redis_cluster.rs b/crates/test-support/tests/redis_cluster.rs index 94b83dc9..0e3d78ae 100644 --- a/crates/test-support/tests/redis_cluster.rs +++ b/crates/test-support/tests/redis_cluster.rs @@ -61,10 +61,9 @@ async fn the_limit_scripts_run_on_a_cluster() { #[ignore = "integration test — needs TEST_REDIS_CLUSTER_URL"] #[tokio::test] -async fn budgets_route_health_and_quotas_run_on_a_cluster() { +async fn budgets_and_route_health_run_on_a_cluster() { use think_watch_common::limits::{BudgetCap, BudgetPeriod, BudgetSubject, budget}; use think_watch_gateway::health::{CircuitBreakerConfig, HealthTracker}; - use think_watch_gateway::quota::QuotaManager; let Some(url) = cluster_url() else { return }; let redis = cluster(&url).await; @@ -121,15 +120,14 @@ async fn budgets_route_health_and_quotas_run_on_a_cluster() { assert_eq!(format!("{:?}", health.state(route, cfg).await), "Open"); health.forget(route).await; assert_eq!(health.snapshot(route, cfg).await.lifetime_requests, 0); +} - // Quotas: limit and usage in one script. - let quota = QuotaManager::new(redis.clone()); - let key = format!("{}:gpt-test", Uuid::new_v4()); - quota.set_limit(&key, 100).await.unwrap(); - assert_eq!(quota.consume(&key, 30).await.unwrap(), 70); - assert_eq!(quota.check_and_consume(&key, 20).await.unwrap(), 50); - assert!(quota.check_and_consume(&key, 60).await.is_err()); - assert_eq!(quota.get_usage(&key).await.unwrap().used, 50); +#[ignore = "integration test — needs TEST_REDIS_CLUSTER_URL"] +#[tokio::test] +async fn the_route_cap_scripts_run_on_a_cluster() { + let Some(url) = cluster_url() else { return }; + let redis = cluster(&url).await; + think_watch_test_support::redis_scripts::exercise_the_route_cap_scripts(&redis).await; } #[ignore = "integration test — needs TEST_REDIS_CLUSTER_URL"] @@ -260,6 +258,18 @@ async fn the_gateway_enforces_limits_on_a_cluster() { .json() .unwrap(); assert_eq!(usage["rules"][0]["current"], 2, "{usage}"); + + // The key's and its owner's day and month usage: one script over four + // hashes (the owner's hash tag), read back by `GET /v1/usage`. The + // refused request counted nothing. + let reported: Json = gw.get("/v1/usage").await.unwrap().json().unwrap(); + assert_eq!(reported["scope"], "key", "{reported}"); + assert_eq!(reported["usage"]["requests_today"], 2, "{reported}"); + assert_eq!(reported["usage"]["requests_month"], 2, "{reported}"); + assert!( + reported["usage"]["tokens_today"].as_i64().unwrap() > 0, + "the key's tokens were counted: {reported}" + ); } #[ignore = "integration test — needs TEST_REDIS_CLUSTER_URL"] diff --git a/crates/test-support/tests/team_rbac.rs b/crates/test-support/tests/team_rbac.rs index b09cfbaa..b7b733f9 100644 --- a/crates/test-support/tests/team_rbac.rs +++ b/crates/test-support/tests/team_rbac.rs @@ -501,3 +501,123 @@ async fn role_attached_to_a_team_gives_its_members_gateway_access() { .unwrap() .assert_status(403); } + +// --------------------------------------------------------------------------- +// A role's limits apply exactly where its gateway grant does +// --------------------------------------------------------------------------- + +/// A role granting model-a at the AI gateway, held to one request a +/// minute — the limit sits in the same Allow statement as the grant. +async fn model_a_once_a_minute(db: &sqlx::PgPool) -> Uuid { + sqlx::query_scalar( + r#"INSERT INTO rbac_roles (name, is_system, policy_document) + VALUES ($1, FALSE, '{"Version":"2024-01-01","Statement":[{"Effect":"Allow", + "Action":["ai_gateway:use"],"Resource":["model:tm-model-a"], + "Constraints":{"RateLimits":[{"Metric":"requests","Window":"1m","MaxCount":1}]}}]}') + RETURNING id"#, + ) + .bind(unique_name("model-a-once")) + .fetch_one(db) + .await + .unwrap() +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_role_attached_to_the_team_grants_and_limits() { + // Attached to the user's team, the role is a working role: it opens + // model-a, and its limit holds the member to one request a minute. + let app = TestApp::spawn().await; + let _upstream = two_routed_models(&app).await; + let role = model_a_once_a_minute(&app.db).await; + let team = make_team(&app.db, "tm-lim-attached").await; + sqlx::query("INSERT INTO team_role_assignments (team_id, role_id) VALUES ($1, $2)") + .bind(team) + .bind(role) + .execute(&app.db) + .await + .unwrap(); + let member = fixtures::create_user(&app.db, &unique_email(), "Member", "MemberPwd_12345!") + .await + .unwrap(); + add_to_team(&app.db, team, member.user.id).await; + let gw = gateway_key(&app, member.user.id).await; + + gw.post("/v1/chat/completions", call(GW_MODEL_A)) + .await + .unwrap() + .assert_ok(); + let limited = gw + .post("/v1/chat/completions", call(GW_MODEL_A)) + .await + .unwrap(); + assert_eq!(limited.status.as_u16(), 429, "{}", limited.text()); + assert!( + limited.text().contains("user:requests/1m"), + "{}", + limited.text() + ); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_role_granted_at_team_scope_neither_grants_nor_limits() { + // Granted at the scope of the user's team, the same role administers + // the team only: it opens no model at the gateway, and its limit + // holds none of the requests another role lets through. + let app = TestApp::spawn().await; + let _upstream = two_routed_models(&app).await; + let role = model_a_once_a_minute(&app.db).await; + let open_b: Uuid = sqlx::query_scalar( + r#"INSERT INTO rbac_roles (name, is_system, policy_document) + VALUES ($1, FALSE, '{"Version":"2024-01-01","Statement":[{"Effect":"Allow", + "Action":["ai_gateway:use"],"Resource":["model:tm-model-b"]}]}') + RETURNING id"#, + ) + .bind(unique_name("model-b-open")) + .fetch_one(&app.db) + .await + .unwrap(); + let user = fixtures::create_user(&app.db, &unique_email(), "Scoped", "ScopedPwd_12345!") + .await + .unwrap(); + sqlx::query( + "INSERT INTO rbac_role_assignments (user_id, role_id, scope_kind, assigned_by) + VALUES ($1, $2, 'global', $1)", + ) + .bind(user.user.id) + .bind(open_b) + .execute(&app.db) + .await + .unwrap(); + let team = make_team(&app.db, "tm-lim-scoped").await; + add_to_team(&app.db, team, user.user.id).await; + sqlx::query( + "INSERT INTO rbac_role_assignments (user_id, role_id, scope_kind, scope_id, assigned_by) + VALUES ($1, $2, 'team', $3, $1)", + ) + .bind(user.user.id) + .bind(role) + .bind(team) + .execute(&app.db) + .await + .unwrap(); + let gw = gateway_key(&app, user.user.id).await; + + let refused = gw + .post("/v1/chat/completions", call(GW_MODEL_A)) + .await + .unwrap(); + assert!( + !refused.status.is_success() && refused.status.as_u16() != 429, + "model-a is not granted by a team-scoped role: {} {}", + refused.status, + refused.text() + ); + for _ in 0..3 { + gw.post("/v1/chat/completions", call(GW_MODEL_B)) + .await + .unwrap() + .assert_ok(); + } +} diff --git a/db/seeds.sql b/db/seeds.sql index 77448c15..596ef90b 100644 --- a/db/seeds.sql +++ b/db/seeds.sql @@ -78,7 +78,7 @@ INSERT INTO system_settings (key, value, category, description) VALUES ('security.signature_nonce_ttl_secs', '600', 'security', 'Request signature nonce TTL in seconds'), ('security.signature_drift_secs', '300', 'security', 'Maximum allowed clock skew for signatures'), ('security.totp_required', 'false', 'security', 'Require all users to enable TOTP two-factor authentication'), -('security.rate_limit_fail_closed', 'false', 'security', 'When true the rate-limit engine refuses requests on Redis outage instead of failing open'), +('security.rate_limit_fail_closed', 'false', 'security', 'When true, a request whose limits cannot be checked (Redis outage, or a database error loading them) is refused instead of let through'), ('security.client_ip_source', '"connection"', 'security', 'Client IP source: "connection", "xff", or "x-real-ip"'), ('security.client_ip_xff_position', '"left"', 'security', 'XFF pick direction: "left" (first) or "right" (last)'), ('security.client_ip_xff_depth', '1', 'security', 'Position depth (1-based) from chosen XFF direction'), diff --git a/deploy/grafana/dashboards/README.md b/deploy/grafana/dashboards/README.md index ed5150d1..b1e037e8 100644 --- a/deploy/grafana/dashboards/README.md +++ b/deploy/grafana/dashboards/README.md @@ -38,7 +38,8 @@ own boards without grepping the source. |---|---| | `gateway_cache_total{result}` | `result=hit` / `miss`. Cache hit rate = `hit / (hit + miss)`. | | `gateway_rate_limited_total` | Per-key rate-limit denials at the proxy. | -| `gateway_quota_overflow_total` | Token quota exceeded. | +| `gateway_route_capped_total{metric}` | A route was skipped at its RPM (`metric=requests`) or TPM (`metric=tokens`) cap. Sustained = the route's caps are below its traffic. | +| `gateway_limits_load_fail_open_total` / `..._fail_closed_total` | A request's limits could not be loaded from Postgres; it ran without them (open) or was refused (closed, `security.rate_limit_fail_closed`). Should be 0. | | `gateway_budget_fail_open_total` | Budget check couldn't reach Redis and fell open. Should be 0 in steady state. | | `gateway_rate_limiter_fail_open_total` / `..._fail_closed_total` | Mirror, for the per-key rate limiter. | | `gateway_sse_buffer_overflow_total` | An SSE event exceeded 8 MiB without a delimiter — stream torn down. Non-zero = a misbehaving upstream provider; investigate. | diff --git a/deploy/grafana/dashboards/think-watch-overview.json b/deploy/grafana/dashboards/think-watch-overview.json index 62316373..eedd28e6 100644 --- a/deploy/grafana/dashboards/think-watch-overview.json +++ b/deploy/grafana/dashboards/think-watch-overview.json @@ -59,7 +59,7 @@ { "id": 3, "type": "timeseries", - "title": "Rate-limited + quota-exceeded (per sec)", + "title": "Rate-limited + routes at cap (per sec)", "datasource": { "type": "prometheus", "uid": "${datasource}" }, "gridPos": { "x": 0, "y": 7, "w": 12, "h": 7 }, "targets": [ @@ -69,8 +69,8 @@ "refId": "A" }, { - "expr": "rate(gateway_quota_overflow_total[1m])", - "legendFormat": "quota overflow", + "expr": "sum by (metric) (rate(gateway_route_capped_total[1m]))", + "legendFormat": "route at {{metric}} cap", "refId": "B" } ], diff --git a/deploy/helm/think-watch/Chart.yaml b/deploy/helm/think-watch/Chart.yaml index 7506d2cf..b6f42727 100644 --- a/deploy/helm/think-watch/Chart.yaml +++ b/deploy/helm/think-watch/Chart.yaml @@ -2,8 +2,8 @@ apiVersion: v2 name: think-watch description: Enterprise AI API Gateway & MCP Management Platform type: application -version: 3.4.0 -appVersion: "3.4.0" +version: 3.5.0 +appVersion: "3.5.0" keywords: - ai - gateway diff --git a/web/package.json b/web/package.json index 90048baa..0780cd31 100644 --- a/web/package.json +++ b/web/package.json @@ -1,7 +1,7 @@ { "name": "web", "private": true, - "version": "3.4.0", + "version": "3.5.0", "type": "module", "packageManager": "pnpm@12.9.1", "scripts": { 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 1897f426..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.", @@ -1777,8 +1796,8 @@ "xffDepthHint": "1-based position from the chosen direction", "rateLimiter": { "title": "Rate Limiter", - "failClosed": "Fail closed on Redis outage", - "failClosedHint": "Default is fail-open: a Redis blip lets requests through and bumps gateway_rate_limiter_fail_open_total. With this on, the gateway returns 429 (rate_limiter_unavailable) for any request the engine couldn't check, and bumps gateway_rate_limiter_fail_closed_total. Turn on if missing accounting is worse than denying traffic." + "failClosed": "Fail closed when limits can't be checked", + "failClosedHint": "Default is fail-open: when a request's limits can't be checked — Redis is unreachable, or the database fails while the limits are loaded — the request goes through and a warning is logged (gateway_rate_limiter_fail_open_total and related metrics). With this on, such a request is refused with 429 (rate_limiter_unavailable, budget_unavailable or limits_unavailable), and gateway_rate_limiter_fail_closed_total and related metrics count it. Turn on if missing accounting is worse than denying traffic." } }, "logs": { @@ -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 aac5beeb..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": "当前真实生效的额度 —— 角色合并后再叠加所有生效的用户覆盖。", @@ -1777,8 +1796,8 @@ "xffDepthHint": "从选定方向起的第 N 个（从 1 开始）", "rateLimiter": { "title": "限速引擎", - "failClosed": "Redis 故障时拒绝请求", - "failClosedHint": "默认为 fail-open:Redis 抖动时请求照常放行,并增加 gateway_rate_limiter_fail_open_total 指标。打开后,引擎无法检查的请求会返回 429 (rate_limiter_unavailable),并增加 gateway_rate_limiter_fail_closed_total。如果漏算比拒绝流量更糟,就把它打开。" + "failClosed": "限额无法检查时拒绝请求", + "failClosedHint": "默认为 fail-open：请求的限额无法检查时（Redis 不可达，或加载限额时数据库出错），请求照常放行并记录警告（gateway_rate_limiter_fail_open_total 等指标）。打开后，这类请求会返回 429（rate_limiter_unavailable、budget_unavailable 或 limits_unavailable），并计入 gateway_rate_limiter_fail_closed_total 等指标。如果漏算比拒绝流量更糟，就把它打开。" } }, "logs": { @@ -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')} +