From e3d4b9528e568e328ce54d56b9109a887784bec3 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Sat, 10 Oct 2026 23:00:51 +0800 Subject: [PATCH] fix(limits): enforce route caps and close the places limits leaked Each of these was a way past a limit, or a limit that did nothing: - Loading a request's limits from Postgres fell back to "no limits" on any error, and a budget read swallowed Redis errors as 0, so the budget gate's fail-closed branch could never run. Both now follow security.rate_limit_fail_closed like the rate limiter does: open runs the request and logs the error, closed answers the existing limiter-unavailable 429 (limits_unavailable / budget_unavailable). - model_routes.rpm_cap / tpm_cap were loaded and never read. They are now enforced as the schema and the route editor describe: a route at its cap is skipped like a broken one, and with every candidate capped the request gets 429 with Retry-After. Counted on the sliding-window engine under a {route:} hash tag; tokens are weighted and added after the response. - The rate-limit stage, which charges the request counters, ran before the model / tool access check on both gateways, so a refused request spent a request. The order is now budget, access, limits, and the lifecycle type states enforce it. - Role-level limits were merged from every role assignment while gateway access reads only global and team-attached roles. A team-scoped grant administers the team and grants nothing at the gateway, as the UI says, so its limits no longer apply either: one loader decides both. - A cache hit counted a request but no tokens, so repeating a cached request got round token limits and budgets. It now counts the tokens the cached answer records. - The legacy monthly quota counted into quota:* keys against a limit nothing could set. Removed with its metric; budgets cover it. Also corrects comments that described a limits cache, a limits:changed subscriber and team budget caps, none of which exist. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 72 ++ crates/auth/src/rbac.rs | 54 +- crates/common/src/lifecycle/mod.rs | 4 +- .../src/lifecycle/stages/check_access.rs | 33 +- .../src/lifecycle/stages/check_budget.rs | 69 +- .../src/lifecycle/stages/check_limits.rs | 106 ++- crates/common/src/lifecycle/stages/mod.rs | 6 +- crates/common/src/lifecycle/state.rs | 49 +- crates/common/src/limits/budget.rs | 37 +- crates/common/src/limits/mod.rs | 21 +- crates/common/src/limits/sliding.rs | 119 +++- crates/common/src/limits/weight.rs | 9 +- crates/common/src/models/team.rs | 5 +- crates/gateway/src/cache.rs | 3 +- crates/gateway/src/lib.rs | 2 +- crates/gateway/src/lifecycle/mod.rs | 4 + crates/gateway/src/proxy/accounting.rs | 15 +- crates/gateway/src/proxy/generate.rs | 75 ++- crates/gateway/src/proxy/mod.rs | 8 +- crates/gateway/src/proxy/pipeline.rs | 19 +- crates/gateway/src/proxy/routing.rs | 84 ++- crates/gateway/src/quota.rs | 237 ------- crates/gateway/src/route_caps.rs | 342 ++++++++++ .../mcp-gateway/src/lifecycle/stages/mod.rs | 2 +- crates/mcp-gateway/src/proxy.rs | 56 +- crates/server/src/app.rs | 2 - crates/server/src/middleware/api_key_auth.rs | 223 ++++++- crates/test-support/src/redis_scripts.rs | 41 +- .../test-support/tests/limit_enforcement.rs | 627 ++++++++++++++++++ crates/test-support/tests/redis_cluster.rs | 18 +- crates/test-support/tests/team_rbac.rs | 120 ++++ db/seeds.sql | 2 +- deploy/grafana/dashboards/README.md | 3 +- .../dashboards/think-watch-overview.json | 6 +- web/src/i18n/en.json | 4 +- web/src/i18n/zh.json | 4 +- 36 files changed, 1951 insertions(+), 530 deletions(-) delete mode 100644 crates/gateway/src/quota.rs create mode 100644 crates/gateway/src/route_caps.rs create mode 100644 crates/test-support/tests/limit_enforcement.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 4b1e4763..25bae8c4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,78 @@ target. ## [Unreleased] +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. The legacy monthly token quota is gone. + +### 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`. 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 with 429, + `Retry-After: 30`, labelled `limits_unavailable`, `budget_unavailable` or + `rate_limiter_unavailable`. +- **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. + +### 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. + ## [3.4.0] — 2026-10-11 Conversations with reasoning models behind Chat-format upstreams keep their 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..eec124e1 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], @@ -194,7 +196,7 @@ pub async fn current_spend( cap.period.as_str(), now, ); - let v: Option = redis.get(&key).await.ok(); + let v: Option = redis.get(&key).await?; out.push(CapStatus { cap_id: cap.id, subject_kind: cap.subject_kind, @@ -395,6 +397,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..ecbee959 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}; @@ -737,8 +741,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 @@ -1245,11 +1249,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..803b80bc 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 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..f1f5ab03 100644 --- a/crates/gateway/src/lib.rs +++ b/crates/gateway/src/lib.rs @@ -12,9 +12,9 @@ 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..b36e7201 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 @@ -62,6 +66,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; 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/mod.rs b/crates/gateway/src/proxy/mod.rs index 6370147d..f5cfe5b5 100644 --- a/crates/gateway/src/proxy/mod.rs +++ b/crates/gateway/src/proxy/mod.rs @@ -14,7 +14,6 @@ use crate::cost_tracker::CostTracker; use crate::error::GatewayError; use crate::health::HealthTracker; use crate::model_mapping::ModelMapper; -use crate::quota::QuotaManager; use crate::rate_limiter::RateLimiter; use crate::router::ModelRouter; use think_watch_common::dynamic_config::DynamicConfig; @@ -60,7 +59,6 @@ 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, @@ -69,8 +67,8 @@ pub struct GatewayState { /// 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 @@ -169,7 +167,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..b4119a2c 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 //! @@ -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,6 @@ 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)?; 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/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..08343e71 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(), diff --git a/crates/server/src/middleware/api_key_auth.rs b/crates/server/src/middleware/api_key_auth.rs index 1f13aad4..4af5492d 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 @@ -292,7 +387,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 +396,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 +409,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 +431,36 @@ pub fn require_api_key( return Ok(gateway_use_refused(surface, request.uri().path())); } + // 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, @@ -492,6 +611,64 @@ 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_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/limit_enforcement.rs b/crates/test-support/tests/limit_enforcement.rs new file mode 100644 index 00000000..faa308eb --- /dev/null +++ b/crates/test-support/tests/limit_enforcement.rs @@ -0,0 +1,627 @@ +//! 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"; +//! - 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(); +} + +// ---------------------------------------------------------------------------- +// 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/redis_cluster.rs b/crates/test-support/tests/redis_cluster.rs index 94b83dc9..9eeb9cd8 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"] 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/web/src/i18n/en.json b/web/src/i18n/en.json index 1897f426..d7f7e17a 100644 --- a/web/src/i18n/en.json +++ b/web/src/i18n/en.json @@ -1777,8 +1777,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": { diff --git a/web/src/i18n/zh.json b/web/src/i18n/zh.json index aac5beeb..cd1b172d 100644 --- a/web/src/i18n/zh.json +++ b/web/src/i18n/zh.json @@ -1777,8 +1777,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": {