diff --git a/.env.example b/.env.example index 648713bd..389320a2 100644 --- a/.env.example +++ b/.env.example @@ -26,11 +26,23 @@ DB_MAX_CONNECTIONS=10 DB_PASSWORD=__SECRET_HEX_16__ # dev: DATABASE_URL=postgres://${DB_USER}:${DB_PASSWORD}@localhost:5432/${DB_NAME} # prod: DATABASE_URL=postgres://${DB_USER}:${DB_PASSWORD}@postgres:5432/${DB_NAME}?sslmode=disable +# Postgres directly, or a pooler in session mode: never one in transaction +# mode (PgBouncer pool_mode = transaction). Start-up schema setup holds a +# session-level advisory lock (deploy/helm/think-watch/README.md). # --- Redis --- REDIS_PASSWORD=__SECRET_HEX_16__ # dev: REDIS_URL=redis://:${REDIS_PASSWORD}@localhost:6379 # prod: REDIS_URL=redis://:${REDIS_PASSWORD}@redis:6379 +# Redis Cluster: REDIS_URL=redis-cluster://:${REDIS_PASSWORD}@redis-0:6379?node=redis-1:6379 +# (see deploy/helm/think-watch/README.md, "Redis Cluster") +# TLS (managed Redis services usually require it): the rediss:// scheme, +# rediss-cluster:// for a cluster. The certificate is checked against the +# public CAs and the host name in the URL. +# REDIS_URL=rediss://:${REDIS_PASSWORD}@my-cache.example.com:6379 +# A Redis whose certificate a private CA signed: a PEM file with that CA, +# then trusted alone (see deploy/helm/think-watch/README.md, "Redis over TLS") +# REDIS_CA_CERT=/etc/thinkwatch/redis-ca.crt # --- Application --- JWT_SECRET=__SECRET_HEX_32__ diff --git a/CHANGELOG.md b/CHANGELOG.md index dd1dbfb9..6d1a009d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,157 @@ target. ## [Unreleased] +## [3.2.0] — 2026-10-05 + +Redis Cluster and Redis over TLS now work, so managed Redis services can be +used. Usage limits work as documented: token limits refuse, an API key's +limits count on that key, and a refused request counts against nothing. +Server instances can start together against one database, a restart no +longer clears captured bodies older than 30 days, and the Helm network +policy allows the ports the database URLs name. The thinkwatch-core crates +stay at v0.62.0. + +### Read before upgrading + +- **Rate-limit windows start empty.** Rate-limit counters move to new Redis + keys (one hash per counter, tagged so that Redis Cluster can run them), and + the counts from before the upgrade are not carried over: every window starts + empty and fills from the first request after the upgrade. The old keys + expire by themselves within two window lengths. Users' budget counters are + kept; API keys' are not (see below). +- **Route health starts fresh.** A route's samples, circuit breaker and + lifetime request count move to new keys for the same reason, so every route + starts closed with nothing counted. The old lifetime counters never expire; + `redis-cli --scan --pattern 'route_health:[0-9a-f]*' | xargs redis-cli del` + removes them (the new keys start `route_health:{`). +- **An API key's own budgets start again.** 3.1.0 counted a key's budget on + its owner's counter (`budget:user::…`); 3.2.0 counts it on the + key's own (`budget:api_key_lineage::…`), which starts at zero. + A key with a monthly budget of its own can spend all of it again in the + rest of the month. +- **Limits are looser while the rollout runs.** Pods of 3.1.0 and 3.2.0 + count rate limits and route health on different keys, so each sees only + its own version's requests: limits let more through, and circuit breakers + can disagree, until the last 3.1.0 pod is gone. +- **Token limits start refusing.** A `tokens` rate limit never refused a + request before. It now does once its window is full, so a deployment with + token limits will see `429`s where it saw none. +- **A key's limits no longer replace its owner's.** A rate limit or budget set + on an API key used to take the place of the owner's limit for the same + window or period. Both now apply, each on its own counter: a key's limits can + narrow what its owner may do through that key, never widen it. A key given a + higher limit than its owner to give it more room needs the owner's limit + raised instead. +- **`rediss://` connects over TLS.** 3.1.0 connected to a `rediss://` URL + over plain TCP. A `rediss://` URL pointing at a port without TLS now fails + at start: point it at the TLS port, or write `redis://`. + - The certificate must name the host in a subjectAltName. One that names + it only in its CN, which `redis-cli` accepts, is refused. + - Cluster nodes that announce IP addresses their certificates don't name + can't be reached. +- **`budget_unavailable` in dashboards.** With Redis down and + `security.rate_limit_fail_closed` on, a request that has a budget is + refused as `budget_unavailable`, not `rate_limiter_unavailable`: budgets + are checked first. +- **Upgrade with an ordinary rollout.** Earlier versions don't take the lock + that now makes instances set up the schema one at a time, so don't restart + instances of the old version while the first one of this version starts. +- **No transaction-mode pooler in front of Postgres.** Schema setup now holds a + Postgres session-level advisory lock, which a pooler in transaction mode + (PgBouncer `pool_mode = transaction`) can leave held, and every later start + then waits for it: point `DATABASE_URL` at Postgres itself or at a pooler in + session mode (the Helm chart's README has the details). + +### Fixed + +- **Token limits refuse requests.** A `tokens` rate limit never refused + anything, and stopped counting once a request would have taken it past its + limit. A request is now refused once the window's recorded usage reaches the + limit, and every request's tokens are recorded after it, even past the limit + — a window can overshoot by what was in flight when it filled. +- **Several request limits at once.** With two or more `requests` limits on a + user (per minute and per hour, say), every request that passed was counted + twice, and only one of the limits could refuse; with + `security.rate_limit_fail_closed` on, every request was refused as + `rate_limiter_unavailable`. +- **An API key's limits count on that key.** They were counted on its owner's + counter, which every key of the owner shared, and the usage the console + reads for a key (`/api/admin/limits/api_key/{id}/usage`) was always 0. Each + key now has counters of its own, rate limits and budgets alike, for the + gateway and the MCP gateway, and its usage shows what it used. +- **A refused request counts against nothing.** A request refused by a spent + budget, or by one rate limit after another had passed, was still counted + against the request limits. Budgets are now checked first and every rate + limit in one step, so a refused request leaves every counter as it was. +- **`Retry-After` says when to retry.** A `429` from the gateway's own limits + said `Retry-After: 30` whatever the limit. It now gives the seconds until the + window has room for another request, or until a spent budget's period ends + (the next midnight, Monday or 1st of the month, UTC). A spent budget also + sends `x-should-retry: false`, so the OpenAI and Anthropic SDKs don't retry it + by themselves. The body stays in the caller's API format. +- **Redis Cluster.** The rate-limit, route-health and quota scripts touched + keys of several hash slots, which a Redis Cluster refuses: on a cluster, rate + limits silently stopped applying (or refused every request with + `security.rate_limit_fail_closed`), circuit breakers never tripped, and cache + invalidation reached one node only. Every key a script touches now shares a + hash tag, pattern deletes scan every node, and the Helm chart's README + describes a `redis-cluster://` URL. +- **Redis over TLS.** The server was built without TLS for Redis: a + `rediss://` URL was used as plain TCP, so against a Redis that requires + TLS — ElastiCache with in-transit encryption, Upstash, Azure Cache for + Redis and Redis Cloud among them — the server did not start + (`Failed to connect to Redis: Protocol Error: Expected string.`). + `rediss://` and `rediss-cluster://` URLs now connect over TLS and check + the certificate against the system's CAs, as upstream HTTPS does. For a + Redis whose certificate a private CA signed, `REDIS_CA_CERT` names a PEM + file with that CA, which is then trusted alone; the Helm chart sets it + from a Secret given in `redis.caSecret`. The chart's README describes + both. +- **Several instances starting at once.** Server instances starting together + against one database — a Helm `replicaCount` above 1, a rolling upgrade, an + autoscaler adding pods — applied the schema side by side, and all but one + could exit with `Database migration failed: apply db/schema.sql: … deadlock + detected` (on an empty database: `duplicate key value violates unique + constraint "pg_extension_name_index"`). With ClickHouse, the rollups that an + instance fills from the logs when it finds them empty (`cost_rollup_hourly`, + `provider_health_5m`, `mcp_server_call_counts`) could be filled by each of + them, counting every request once per instance on the cost pages, the + dashboard and the MCP server list. Instances now set up Postgres and + ClickHouse one at a time, under Postgres advisory locks: the others wait, + logging `Another instance is setting up the database schema; waiting for it + to finish`, then find it done. An instance that dies holding a lock releases + it with its connection. The Helm chart's startup probe allows a pod 125 s + to start instead of 35 s, set in `startupProbe` (the chart's README says + when to raise it), and an attempt to reach a ClickHouse that doesn't answer + gives up after 5 s instead of the system's TCP timeout. +- **Captured bodies kept as long as configured.** With ClickHouse and + `audit.body_retention_days` above 30, every server start could clear the + captured request and response bodies older than 30 days + (`gateway_logs.request_body` / `response_body`, `mcp_logs.tool_arguments` + / `tool_result`; the rows themselves stayed). The start-up table setup set + those columns' TTL to 30 days each time, and ClickHouse applies a TTL to + the data already stored as soon as it is set, before the server put the + configured TTL back a moment later. The setup now gives these columns a + TTL only when it creates them, so a restart leaves the configured one in + place. Bodies already cleared cannot be recovered. The log tables' own + TTLs (`data.retention_days_*`) were not affected. +- **Helm network policy and databases on other ports.** With + `networkPolicy.enabled`, the server could reach PostgreSQL only on `5432`, + Redis on `6379` and ClickHouse on `8123`, whatever their `externalUrl` said, + so a database on another port was blocked — Azure Cache for Redis over TLS + (`6380`), a managed Postgres or a ClickHouse on a port of its own: the + server could not start, or started without writing to ClickHouse. The + allowed ports now follow `postgres.externalUrl`, `redis.externalUrl` and + `clickhouse.externalUrl`: every port a URL names, and the client's default + for its scheme where it names none. `networkPolicy.extraEgress` adds egress + rules as written, for ports no URL names (Redis Cluster nodes announcing + other ports, an upstream or MCP server on a port other than `443`). The + chart's README describes both. Port `9000`, ClickHouse's native protocol, + is no longer allowed: the server reaches ClickHouse over plain HTTP only + (an `http://` URL; HTTPS is not supported). An S3 endpoint on `9000` + (RustFS, MinIO) configured outside the chart needs a rule in + `networkPolicy.extraEgress`. + ## [3.1.0] — 2026-10-05 The thinkwatch-core crates move from v0.59.0 to v0.62.0. Two changes reach @@ -1180,7 +1331,8 @@ unreleased builds should: stop the gateway, run `db/schema.sql` against PostgreSQL, restart against this tag. The schema is idempotent end-to-end, so the apply is safe to repeat. -[Unreleased]: https://github.com/ThinkWatchProject/ThinkWatch/compare/v3.1.0...HEAD +[Unreleased]: https://github.com/ThinkWatchProject/ThinkWatch/compare/v3.2.0...HEAD +[3.2.0]: https://github.com/ThinkWatchProject/ThinkWatch/releases/tag/v3.2.0 [3.1.0]: https://github.com/ThinkWatchProject/ThinkWatch/releases/tag/v3.1.0 [3.0.0]: https://github.com/ThinkWatchProject/ThinkWatch/releases/tag/v3.0.0 [2.2.0]: https://github.com/ThinkWatchProject/ThinkWatch/releases/tag/v2.2.0 diff --git a/Cargo.lock b/Cargo.lock index 93be0587..837f41fa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1335,9 +1335,12 @@ dependencies = [ "parking_lot", "rand 0.8.5", "redis-protocol", + "rustls", + "rustls-native-certs", "semver", "socket2 0.5.10", "tokio", + "tokio-rustls", "tokio-stream", "tokio-util", "url", @@ -3220,6 +3223,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "758025cb5fccfd3bc2fd74708fd4682be41d99e5dff73c377c0646c6012c73a4" dependencies = [ "aws-lc-rs", + "log", "once_cell", "ring", "rustls-pki-types", @@ -4034,7 +4038,7 @@ checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" [[package]] name = "think-watch-auth" -version = "3.1.0" +version = "3.2.0" dependencies = [ "anyhow", "argon2", @@ -4064,7 +4068,7 @@ dependencies = [ [[package]] name = "think-watch-common" -version = "3.1.0" +version = "3.2.0" dependencies = [ "aes-gcm", "anyhow", @@ -4077,13 +4081,17 @@ dependencies = [ "clickhouse", "dotenvy", "fred", + "futures", "hex", "hmac 0.13.0", "http 1.4.0", + "hyper-util", "metrics", "rand 0.10.0", "reqwest 0.13.2", "rust_decimal", + "rustls", + "rustls-platform-verifier", "serde", "serde_json", "sha1 0.11.0", @@ -4103,7 +4111,7 @@ dependencies = [ [[package]] name = "think-watch-gateway" -version = "3.1.0" +version = "3.2.0" dependencies = [ "anyhow", "arc-swap", @@ -4142,7 +4150,7 @@ dependencies = [ [[package]] name = "think-watch-mcp-gateway" -version = "3.1.0" +version = "3.2.0" dependencies = [ "anyhow", "arc-swap", @@ -4171,7 +4179,7 @@ dependencies = [ [[package]] name = "think-watch-server" -version = "3.1.0" +version = "3.2.0" dependencies = [ "anyhow", "arc-swap", @@ -4222,7 +4230,7 @@ dependencies = [ [[package]] name = "think-watch-test-support" -version = "3.1.0" +version = "3.2.0" dependencies = [ "anyhow", "async-stream", diff --git a/Cargo.toml b/Cargo.toml index 8c0b5b75..9577d8e8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,7 +10,7 @@ members = [ ] [workspace.package] -version = "3.1.0" +version = "3.2.0" edition = "2024" # Pin the MSRV to the first stable rustc that ships edition 2024 (1.85, # released 2025-02-20). Without this, contributors on older toolchains @@ -90,7 +90,14 @@ rust_decimal = { version = "1", features = ["serde", "serde-with-str"] } clickhouse = { version = "0.13", features = ["time", "chrono"] } # Redis -fred = { version = "10", features = ["subscriber-client", "i-scripts"] } +# `enable-rustls` gives fred its rustls connector, for `rediss://` URLs. +# The connector itself is built in `think_watch_common::redis_config` +# with the same TLS stack the HTTP client uses (below): rustls on +# aws-lc-rs, the platform's roots through rustls-platform-verifier. No +# OpenSSL anywhere, so the static musl image needs nothing new. +fred = { version = "10", features = ["subscriber-client", "i-scripts", "enable-rustls"] } +rustls = { version = "0.23", default-features = false, features = ["aws_lc_rs", "std", "tls12"] } +rustls-platform-verifier = "0.6" # Auth # Crypto-critical crates are pinned with `=` so a silent minor-release diff --git a/README.md b/README.md index 6e8773fa..70e29e53 100644 --- a/README.md +++ b/README.md @@ -92,7 +92,9 @@ The gateway (port `3000`) is the only part that clients need to reach. The conso - A model's maximum output tokens, set on the Models page, caps `max_tokens` on every request to that model; it replaces the old output length guardrail. **Limits and budgets** -- Request-count limits are checked before the request; token limits and budgets are counted after the response, so one request can cross a budget before the next is refused. +- Every limit and budget is checked before the request, against what earlier requests used; tokens are counted after the response, so one request can cross a token limit or a budget before the next is refused. A refused request counts against nothing, and its `429` says in `Retry-After` when the limit frees. +- A limit set on an API key applies on top of its owner's, on a counter of its own. +- Redis can be a single node or a Redis Cluster. - If Redis is unavailable, limits fail open by default. Setting `security.rate_limit_fail_closed` refuses requests instead. - Budget alerts fire once per period at 50%, 80%, 95% and 100%. diff --git a/README.zh-CN.md b/README.zh-CN.md index a740754b..2d98bc6d 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -92,7 +92,9 @@ cd web && pnpm install && pnpm dev - 模型的「最大输出 token」在模型页设置，限制发往该模型的每个请求的 `max_tokens`，取代原来的输出长度护栏。 **限流与预算** -- 请求数限制在请求发出前检查；Token 限制与预算在响应返回后计入，因此单个请求可能越过预算，此后的请求才会被拒绝。 +- 所有限流与预算都在请求发出前按此前的用量检查；Token 在响应返回后计入，因此单个请求可能越过 Token 限制或预算，此后的请求才会被拒绝。被拒绝的请求不计入任何限制，其 `429` 响应以 `Retry-After` 说明限制何时解除。 +- 设置在 API Key 上的限制叠加在其所属用户的限制之上，单独计数。 +- Redis 可以是单节点，也可以是 Redis Cluster。 - Redis 不可用时，限流默认放行；设置 `security.rate_limit_fail_closed` 后改为拒绝请求。 - 预算提醒在每个周期内于 50%、80%、95% 和 100% 各触发一次。 diff --git a/crates/auth/src/rbac.rs b/crates/auth/src/rbac.rs index 4adac13e..9867d1c6 100644 --- a/crates/auth/src/rbac.rs +++ b/crates/auth/src/rbac.rs @@ -373,61 +373,26 @@ pub async fn compute_user_surface_constraints( Ok(apply_user_overrides(role_merged, override_constraints)) } -/// Like [`compute_user_surface_constraints`] but also folds in any -/// active rate-limit / budget overrides keyed on a specific -/// `api_key_id`. Per-key overrides REPLACE the user-derived value -/// in the matching `(surface, metric, window)` / `(surface, period)` -/// slot — the same merge semantics as user-scope overrides. -/// -/// Use this in the gateway hot path (where the auth middleware -/// already knows the api_key id). Other callers (analytics -/// dashboards, admin "what does this user see today" views) keep -/// using the user-only variant since they have no key context. -pub async fn compute_effective_surface_constraints( +/// The limits attached to one API key — its lineage's active +/// `rate_limit_rules` / `budget_caps` rows — on their own. They are not +/// merged into the owner's: the gateway counts them on the lineage's +/// counters and checks them on top of the owner's limits, so a key's +/// limits can narrow what its owner may do through it but never widen +/// it. Keyed on the lineage so they survive rotation. +pub async fn compute_key_surface_constraints( pool: &PgPool, - user_id: Uuid, - api_key_id: Uuid, + lineage_id: Uuid, ) -> Result { use think_watch_common::limits::{ - self, BudgetSubject, RateLimitSubject, apply_user_overrides, - list_enabled_caps_for_subjects, list_enabled_rules_for_subjects, side_table_as_constraints, + BudgetSubject, RateLimitSubject, list_enabled_caps_for_subjects, + list_enabled_rules_for_subjects, side_table_as_constraints, }; - - let user_merged = compute_user_surface_constraints(pool, user_id).await?; - - // Per-key overrides are stored against the key's `lineage_id` - // (subject_kind = 'api_key_lineage') so they survive rotation. - // Resolve api_key_id → lineage_id once and bind every lookup on - // the lineage. A non-existent api_key_id (never happens at - // runtime — the auth middleware just authenticated this id — - // but defend anyway) maps to an empty override set. - let lineage_id: Option = - sqlx::query_scalar("SELECT lineage_id FROM api_keys WHERE id = $1") - .bind(api_key_id) - .fetch_optional(pool) - .await?; - let Some(lineage_id) = lineage_id else { - return Ok(user_merged); - }; - - // The api_key-side overrides are loaded separately so the - // existing `compute_user_*` helper stays a pure function of - // user_id (used by analytics + admin views). Loading both kinds - // here is two queries instead of one — that's bounded and the - // hot path already does enough DB round-trips that it doesn't - // dominate latency. - let key_rules = + let rules = list_enabled_rules_for_subjects(pool, &[(RateLimitSubject::ApiKeyLineage, lineage_id)]) .await?; - let key_caps = + let caps = list_enabled_caps_for_subjects(pool, &[(BudgetSubject::ApiKeyLineage, lineage_id)]).await?; - let key_overrides = side_table_as_constraints(&key_rules, &key_caps); - - if key_overrides == limits::SurfaceConstraints::default() { - return Ok(user_merged); - } - - Ok(apply_user_overrides(user_merged, key_overrides)) + Ok(side_table_as_constraints(&rules, &caps)) } /// Compute the set of permissions that are explicitly denied to `user_id` diff --git a/crates/common/Cargo.toml b/crates/common/Cargo.toml index 8826a6ee..4c113971 100644 --- a/crates/common/Cargo.toml +++ b/crates/common/Cargo.toml @@ -10,6 +10,8 @@ tw-bedrock = { workspace = true } axum = { workspace = true } sqlx = { workspace = true } fred = { workspace = true } +rustls = { workspace = true } +rustls-platform-verifier = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } uuid = { workspace = true } @@ -30,7 +32,11 @@ sha2 = { workspace = true } hmac = { workspace = true } metrics = { workspace = true } clickhouse = { workspace = true } +# The ClickHouse client's HTTP connector, built here to give it a connect +# timeout (`clickhouse_client.rs`). +hyper-util = { workspace = true, features = ["client-legacy", "http1", "tokio"] } bytes = { workspace = true } +futures = { workspace = true } url = { workspace = true } # S3-compatible body offload (matches the same SigV4 + reqwest pattern # the Bedrock provider uses — no aws-sdk-s3 dependency, so the build diff --git a/crates/common/src/clickhouse_client.rs b/crates/common/src/clickhouse_client.rs index 282446d3..38382cbb 100644 --- a/crates/common/src/clickhouse_client.rs +++ b/crates/common/src/clickhouse_client.rs @@ -1,11 +1,41 @@ +use std::time::Duration; + +use hyper_util::client::legacy::Client as HyperClient; +use hyper_util::client::legacy::connect::HttpConnector; +use hyper_util::rt::TokioExecutor; + use crate::audit::AuditConfig; +/// How long opening a connection to ClickHouse may take. The `clickhouse` +/// crate's own connector has no limit, so a ClickHouse whose network +/// drops packets instead of refusing them held each attempt for the +/// operating system's TCP timeout (about two minutes on Linux) — at start, +/// while holding the ClickHouse setup lock that other instances queue +/// behind. +const CONNECT_TIMEOUT: Duration = Duration::from_secs(5); + +/// The crate's defaults, which a client built by hand has to restate. +const TCP_KEEPALIVE: Duration = Duration::from_secs(60); +/// Below ClickHouse's own keep-alive timeout (3 s before 23.11, 10 s +/// since), so the client never reuses a socket the server has closed. +const POOL_IDLE_TIMEOUT: Duration = Duration::from_secs(2); + /// Create a `clickhouse::Client` from our config. Returns `None` if ClickHouse /// is not configured (no URL). +/// +/// Plain HTTP only: the crate is built without TLS, and the connector +/// refuses an `https://` URL. pub fn create_client(config: &AuditConfig) -> Option { let url = config.clickhouse_url.as_deref()?; - let mut client = clickhouse::Client::default() + let mut connector = HttpConnector::new(); + connector.set_keepalive(Some(TCP_KEEPALIVE)); + connector.set_connect_timeout(Some(CONNECT_TIMEOUT)); + let http = HyperClient::builder(TokioExecutor::new()) + .pool_idle_timeout(POOL_IDLE_TIMEOUT) + .build(connector); + + let mut client = clickhouse::Client::with_http_client(http) .with_url(url) .with_database(&config.clickhouse_db) .with_product_info("think-watch", env!("CARGO_PKG_VERSION")); @@ -19,3 +49,57 @@ pub fn create_client(config: &AuditConfig) -> Option { Some(client) } + +#[cfg(test)] +mod tests { + use super::*; + + fn config(url: &str) -> AuditConfig { + AuditConfig { + clickhouse_url: Some(url.to_owned()), + ..Default::default() + } + } + + /// A ClickHouse that never answers the handshake fails the query + /// after the connect timeout, not the operating system's. + #[tokio::test] + async fn an_unanswered_connect_gives_up_after_the_timeout() { + // A listener whose accept queue is full drops further SYNs + // unanswered, as a firewall that drops packets does. Fill it. + let socket = tokio::net::TcpSocket::new_v4().unwrap(); + socket.bind("127.0.0.1:0".parse().unwrap()).unwrap(); + let listener = socket.listen(1).unwrap(); + let addr = listener.local_addr().unwrap(); + let mut held = Vec::new(); + while let Ok(stream) = tokio::time::timeout( + Duration::from_millis(500), + tokio::net::TcpStream::connect(addr), + ) + .await + { + held.push(stream.unwrap()); + assert!(held.len() < 64, "the accept queue never filled"); + } + + let client = create_client(&config(&format!("http://{addr}"))).unwrap(); + let started = std::time::Instant::now(); + let result = tokio::time::timeout( + CONNECT_TIMEOUT * 4, + client.query("SELECT 1").fetch_one::(), + ) + .await; + let elapsed = started.elapsed(); + assert!( + matches!(result, Ok(Err(_))), + "still waiting after {elapsed:?}, or answered" + ); + // The operating system gives up far later: about two minutes on + // Linux, some 8 s on macOS for a local address. + assert!( + elapsed < CONNECT_TIMEOUT + Duration::from_secs(1), + "{elapsed:?}" + ); + drop(held); + } +} diff --git a/crates/common/src/config.rs b/crates/common/src/config.rs index 3702f021..9afd363e 100644 --- a/crates/common/src/config.rs +++ b/crates/common/src/config.rs @@ -4,6 +4,11 @@ use serde::Deserialize; pub struct AppConfig { pub database_url: String, pub redis_url: String, + /// `REDIS_CA_CERT`: a PEM file with the CA certificate(s) to trust + /// for a `rediss://` `REDIS_URL`, in place of the system's roots — + /// for a Redis whose certificate a private CA signed. See + /// [`crate::redis_config`]. + pub redis_ca_cert: Option, pub jwt_secret: String, pub encryption_key: String, pub server_host: String, @@ -51,6 +56,9 @@ impl AppConfig { .map_err(|_| anyhow::anyhow!("DATABASE_URL environment variable is required"))?, redis_url: std::env::var("REDIS_URL") .map_err(|_| anyhow::anyhow!("REDIS_URL environment variable is required"))?, + redis_ca_cert: std::env::var_os("REDIS_CA_CERT") + .filter(|s| !s.is_empty()) + .map(Into::into), jwt_secret: std::env::var("JWT_SECRET") .map_err(|_| anyhow::anyhow!("JWT_SECRET environment variable is required"))?, encryption_key: std::env::var("ENCRYPTION_KEY") @@ -134,6 +142,13 @@ impl AppConfig { } } + if self.redis_ca_cert.is_some() && !crate::redis_config::uses_tls(&self.redis_url) { + tracing::warn!( + "REDIS_CA_CERT is set, but REDIS_URL does not use TLS — the connection is \ + unencrypted and the certificate unused; use a rediss:// URL" + ); + } + // ClickHouse auth warning if self.clickhouse_url.is_some() && self.clickhouse_password.is_none() { tracing::warn!( @@ -144,6 +159,13 @@ impl AppConfig { Ok(()) } + /// The fred config every Redis client of the server is built from: + /// `REDIS_URL`, with TLS when its scheme is `rediss` and + /// `REDIS_CA_CERT`'s roots when set. + pub fn redis_config(&self) -> anyhow::Result { + crate::redis_config::client_config(&self.redis_url, self.redis_ca_cert.as_deref()) + } + pub fn gateway_addr(&self) -> String { format!("{}:{}", self.server_host, self.gateway_port) } @@ -157,6 +179,7 @@ impl AppConfig { Self { database_url: "postgres://test".into(), redis_url: "redis://test".into(), + redis_ca_cert: None, jwt_secret: jwt_secret.into(), encryption_key: encryption_key.into(), server_host: "0.0.0.0".into(), diff --git a/crates/common/src/db.rs b/crates/common/src/db.rs index 25e70127..f92a8828 100644 --- a/crates/common/src/db.rs +++ b/crates/common/src/db.rs @@ -1,5 +1,6 @@ -use sqlx::PgPool; +use sqlx::pool::PoolConnection; use sqlx::postgres::PgPoolOptions; +use sqlx::{Executor, PgConnection, PgPool, Postgres}; pub async fn create_pool(database_url: &str) -> anyhow::Result { if !database_url.contains("sslmode=") { @@ -28,6 +29,64 @@ pub async fn create_pool(database_url: &str) -> anyhow::Result { Ok(pool) } +/// Advisory-lock key held while an instance applies the schema, the +/// seeds and the conversions ([`run_migrations`]): "twschema" in ASCII. +pub const SCHEMA_LOCK: i64 = i64::from_be_bytes(*b"twschema"); + +/// Advisory-lock key held while an instance sets up the ClickHouse tables +/// and backfills the rollups: "twchinit" in ASCII. +pub const CLICKHOUSE_SETUP_LOCK: i64 = i64::from_be_bytes(*b"twchinit"); + +/// A Postgres session-level advisory lock on a connection of its own. +/// +/// Start-up work that every instance runs, but that two instances must +/// not run at once, holds one: an instance that starts while another +/// holds the lock waits for it, then finds the work done. +/// +/// The lock lasts as long as the session. The connection never goes back +/// to the pool — [`release`](Self::release) closes it, and so does +/// dropping the guard on an error — so no later query can find itself +/// holding the lock. An instance that dies holding it loses its +/// connection, and Postgres releases the lock with the session. +pub struct StartupLock { + conn: PoolConnection, +} + +impl StartupLock { + /// Take the lock `key`, waiting while another instance holds it. + /// `work` says what the holder is doing, for the log line a waiting + /// instance writes. + pub async fn acquire(pool: &PgPool, key: i64, work: &str) -> anyhow::Result { + let mut conn = pool.acquire().await?; + conn.close_on_drop(); + let free: bool = sqlx::query_scalar("SELECT pg_try_advisory_lock($1)") + .bind(key) + .fetch_one(&mut *conn) + .await?; + if !free { + tracing::info!("Another instance is {work}; waiting for it to finish"); + sqlx::query("SELECT pg_advisory_lock($1)") + .bind(key) + .execute(&mut *conn) + .await?; + } + Ok(Self { conn }) + } + + /// The connection holding the lock. The work runs on it. + pub fn conn(&mut self) -> &mut PgConnection { + &mut self.conn + } + + /// Release the lock by ending the session. + pub async fn release(self) { + if let Err(e) = self.conn.close().await { + // The session is gone either way, and the lock with it. + tracing::debug!("closing the start-up lock's connection: {e}"); + } + } +} + /// Bring the database up to the schema declared in `db/schema.sql`, /// then idempotent-apply seeds from `db/seeds.sql`. Both files are /// embedded at compile time, so the running binary never needs disk @@ -45,20 +104,35 @@ pub async fn create_pool(database_url: &str) -> anyhow::Result { /// to carry them out by itself, a conversion that runs here in one /// transaction and finds nothing left to do on the next boot (the guard /// settings, [`crate::guard_policy::legacy`]). +/// +/// **One instance at a time.** Several replicas start together — a +/// first install with more than one, a rolling upgrade, an autoscaler — +/// and each runs this. Run side by side, the statements deadlock (each +/// session holding a lock on `api_keys` or its trigger that another +/// needs) or collide creating the same object on an empty database, and +/// the instances that lose exit. So all of it runs under +/// [`SCHEMA_LOCK`], on the connection that holds it: the next instance +/// waits, then applies a schema that is already there, seeds rows that +/// already exist, and finds nothing left to convert. pub async fn run_migrations(pool: &PgPool) -> anyhow::Result<()> { + let mut lock = + StartupLock::acquire(pool, SCHEMA_LOCK, "setting up the database schema").await?; + let conn = lock.conn(); let schema = include_str!("../../../db/schema.sql"); - sqlx::raw_sql(schema) - .execute(pool) + // `Executor::execute` rather than `RawSql::execute`: with a borrowed + // connection, the compiler cannot prove the latter's future `Send`, + // and a caller that spawns this would not compile. + conn.execute(sqlx::raw_sql(schema)) .await .map_err(|e| anyhow::anyhow!("apply db/schema.sql: {e}"))?; let seeds = include_str!("../../../db/seeds.sql"); - sqlx::raw_sql(seeds) - .execute(pool) + conn.execute(sqlx::raw_sql(seeds)) .await .map_err(|e| anyhow::anyhow!("apply db/seeds.sql: {e}"))?; - crate::guard_policy::legacy::upgrade(pool) + crate::guard_policy::legacy::upgrade(conn) .await .map_err(|e| anyhow::anyhow!("convert the previous guard settings: {e}"))?; + lock.release().await; tracing::info!("Database schema + seeds applied"); Ok(()) } diff --git a/crates/common/src/guard_policy/legacy.rs b/crates/common/src/guard_policy/legacy.rs index 7fe61d1f..568fdcd8 100644 --- a/crates/common/src/guard_policy/legacy.rs +++ b/crates/common/src/guard_policy/legacy.rs @@ -42,14 +42,16 @@ //! **Once, in one transaction, safe to run again.** The old keys and the //! old column go in the same transaction that writes what replaces them, //! so the next boot finds nothing to convert, and a failure leaves -//! everything as it was. An advisory lock keeps two replicas booting at -//! once from both converting. +//! everything as it was. Two replicas booting at once don't both convert: +//! it runs under the schema lock [`crate::db::run_migrations`] holds, and +//! under a lock of its own, which an instance of a version from before +//! the schema lock takes too. use std::collections::{BTreeMap, BTreeSet, HashSet}; use serde::Deserialize; use serde_json::Value; -use sqlx::PgPool; +use sqlx::{Connection, PgConnection}; use tw_guard::policy::{ ContentAction, ContentMatch, ContentPolicy, CustomContentRule, CustomRedactRule, CustomToolRule, DEFAULT_LABEL, Guard, LABEL_MAX, Mode, RedactPolicy, ToolAction, ToolPolicy, @@ -94,8 +96,8 @@ pub const MARKER: &str = "security.legacy_converted"; /// a rollback, or an old replica restarting, whose seeds write their /// defaults back. Converting them would overwrite the policies in force /// with those defaults, so they are only removed, with a warning. -pub async fn upgrade(pool: &PgPool) -> anyhow::Result<()> { - let mut tx = pool.begin().await?; +pub async fn upgrade(conn: &mut PgConnection) -> anyhow::Result<()> { + let mut tx = conn.begin().await?; sqlx::query("SELECT pg_advisory_xact_lock($1)") .bind(LOCK) .execute(&mut *tx) diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs index b3d1d944..334950bd 100644 --- a/crates/common/src/lib.rs +++ b/crates/common/src/lib.rs @@ -49,5 +49,7 @@ pub mod crypto; // AES-256-GCM envelope for secrets at rest pub mod fixed_window; pub mod json_secret; // `{"$enc": ...}` — a secret nested inside a JSONB column pub mod pii; // BlobRedactor — at-rest body redaction shared by gateway + mcp-gateway +pub mod redis_config; // the fred client config for REDIS_URL, with TLS for rediss:// +pub mod redis_keys; // pattern deletes that work on one node and on a Redis Cluster pub mod tasks; // supervised_spawn — panic-isolated background tasks pub mod validation; diff --git a/crates/common/src/lifecycle/mod.rs b/crates/common/src/lifecycle/mod.rs index 7723c8b3..d6ae2cd0 100644 --- a/crates/common/src/lifecycle/mod.rs +++ b/crates/common/src/lifecycle/mod.rs @@ -8,8 +8,8 @@ //! //! ```text //! Raw -//! → check_limits (rate-limit gate) -//! → check_budget (pre-call budget peek) +//! → check_budget (pre-call budget peek; 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) diff --git a/crates/common/src/lifecycle/stages/check_budget.rs b/crates/common/src/lifecycle/stages/check_budget.rs index fbd5e66f..48e8c609 100644 --- a/crates/common/src/lifecycle/stages/check_budget.rs +++ b/crates/common/src/lifecycle/stages/check_budget.rs @@ -16,26 +16,27 @@ //! (request allowed) unless the caller passes `fail_closed = true` //! — matching the [`super::check_limits`] semantics. //! -//! Ordering note: this stage runs AFTER `check_limits`. A request -//! that hits the per-window requests counter via `check_limits` and -//! then gets rejected here for being over-budget will have its -//! `requests` counter incremented anyway — rate-limit counters -//! measure "attempts", not "successes". Operators querying the -//! rate-limit metric will see budget-rejected requests reflected -//! there, by design. +//! 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. +use chrono::Utc; use fred::clients::Client; use crate::audit::AuditLogger; use crate::limits::{BudgetCap, budget}; use super::super::Surface; -use super::super::state::LimitsChecked; +use super::super::state::Raw; /// Run a read-only spend check against every supplied cap. On /// allow, returns the input state unchanged (passthrough — no /// extra type narrowing). On deny, emits a `"budget_exceeded"` -/// audit row + short-circuits with `S::budget_exceeded_response`. +/// audit row + short-circuits with `S::budget_exceeded_response`, +/// naming the spent cap whose period ends last — the request can't +/// succeed before then — and the seconds until it does. /// /// `fail_closed` controls behaviour on a Redis read error: /// - `false` (default): bumps `lifecycle_budget_fail_open_total` @@ -47,12 +48,12 @@ use super::super::state::LimitsChecked; fields(trace_id = %state.trace_id, cap_count = caps.len()), )] pub async fn check_budget( - state: LimitsChecked, + state: Raw, caps: &[BudgetCap], redis: &Client, fail_closed: bool, audit: &AuditLogger, -) -> Result, S::Response> { +) -> Result, S::Response> { if caps.is_empty() { return Ok(state); } @@ -84,27 +85,34 @@ pub async fn check_budget( // Walk caps + statuses in lockstep — `current_spend` preserves // input order so the pairing is positional. - for (cap, status) in caps.iter().zip(statuses.iter()) { - if status.current >= status.limit { - let label = budget_label(cap); - metrics::counter!("lifecycle_budget_exceeded_total").increment(1); - tracing::warn!( - trace_id = %state.trace_id, - cap = %label, - current = status.current, - limit = status.limit, - "budget cap exceeded" - ); - let entry = S::audit_entry(&state.identity, "budget_exceeded") - .trace_id(state.trace_id.clone()) - .detail(serde_json::json!({ - "limit": label, - "current": status.current, - "max": status.limit, - })); - audit.log(entry); - return Err(S::budget_exceeded_response(&label)); - } + let now = Utc::now(); + let spent = caps + .iter() + .zip(statuses.iter()) + .filter(|(_, status)| status.current >= status.limit) + .map(|(cap, status)| (cap, status, budget::secs_until_period_end(cap.period, now))) + .max_by_key(|(_, _, wait)| *wait); + if let Some((cap, status, retry_after_secs)) = spent { + let label = budget_label(cap); + metrics::counter!("lifecycle_budget_exceeded_total").increment(1); + tracing::warn!( + trace_id = %state.trace_id, + cap = %label, + current = status.current, + limit = status.limit, + retry_after_secs, + "budget cap exceeded" + ); + let entry = S::audit_entry(&state.identity, "budget_exceeded") + .trace_id(state.trace_id.clone()) + .detail(serde_json::json!({ + "limit": label, + "current": status.current, + "max": status.limit, + "retry_after_secs": retry_after_secs, + })); + audit.log(entry); + return Err(S::budget_exceeded_response(&label, retry_after_secs)); } Ok(state) @@ -125,7 +133,6 @@ fn budget_label(cap: &BudgetCap) -> String { #[cfg(test)] mod tests { - use super::super::super::state::{LimitCheckRecord, LimitsChecked}; use super::super::super::test_surface::{TestResponse, TestSurface, make_raw}; use super::*; use crate::limits::{BudgetPeriod, BudgetSubject}; @@ -142,26 +149,13 @@ mod tests { crate::audit::AuditLogger::test_drain() } - 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(), - }, - } - } - /// No caps configured → trivially pass through without touching /// Redis. The disconnected dummy client confirms the function /// never reached out. #[tokio::test] async fn passes_through_with_no_caps() { let user_id = Uuid::new_v4(); - let state = make_limits_checked(user_id); + let state = make_raw(user_id); let trace_id = state.trace_id.clone(); let started_at = state.started_at; @@ -194,7 +188,7 @@ mod tests { /// short-circuit Response type to be propagated via `?`-style /// match — same shape `check_limits` uses. #[allow(dead_code)] - async fn type_state_compiles(state: LimitsChecked) { + async fn type_state_compiles(state: Raw) { let _ = check_budget::(state, &[], &dummy_redis(), false, &dummy_audit()) .await .map_err(|r| match r { diff --git a/crates/common/src/lifecycle/stages/check_limits.rs b/crates/common/src/lifecycle/stages/check_limits.rs index 31d86f6e..210fb2ae 100644 --- a/crates/common/src/lifecycle/stages/check_limits.rs +++ b/crates/common/src/lifecycle/stages/check_limits.rs @@ -1,11 +1,17 @@ -//! `check_limits` — the first stage of the request lifecycle. -//! Charges the user's `requests` rate-limit counters and either: +//! `check_limits` — the rate-limit gate of the request lifecycle. +//! Checks every rate-limit rule the request is held to, in one atomic +//! step, and either: //! -//! * passes through to [`LimitsChecked`] (carrying the post-INCR -//! currents for downstream audit), or -//! * short-circuits with `S::rate_limited_response(label)` / -//! `S::rate_limiter_unavailable_response()` after emitting the -//! audit row for the deny. +//! * passes through to [`LimitsChecked`] (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)` +//! / `S::rate_limiter_unavailable_response()` after emitting the +//! audit row for the deny, having charged nothing. +//! +//! A `tokens` rule refuses the request once its window's recorded +//! usage has reached the limit; the tokens the call uses are added +//! after it, by the surface's `record_usage` hook. //! //! No surface-specific logic lives here — the stage takes //! pre-built `RateLimitRule`s. The surface's pipeline runner is @@ -13,17 +19,20 @@ //! before calling. use fred::clients::Client; +use uuid::Uuid; use crate::audit::AuditLogger; -use crate::limits::{RateLimitRule, RateMetric, sliding}; +use crate::limits::{RateLimitRule, sliding}; use super::super::Surface; use super::super::state::{LimitCheckRecord, LimitsChecked, Raw}; -/// Resolve and charge `requests` counters. See module-level docs. +/// Check the rules and charge the request's `requests` counters. See +/// module-level docs. `owner` is the user the request runs as: every +/// counter carries their hash tag (`sliding::counter_key`). /// /// `fail_closed` controls behaviour on Redis errors: -/// - `false` (default): bumps `lifecycle_rate_limiter_fail_open_total` +/// - `false` (default): bumps `gateway_rate_limiter_fail_open_total` /// and lets the request through. /// - `true`: emits the audit row and short-circuits with /// `S::rate_limiter_unavailable_response()`. Wired up from the @@ -35,16 +44,15 @@ use super::super::state::{LimitCheckRecord, LimitsChecked, Raw}; pub async fn check_limits( state: Raw, rules: &[RateLimitRule], + owner: Uuid, redis: &Client, fail_closed: bool, audit: &AuditLogger, ) -> Result, S::Response> { - let resolved = sliding::resolve_rules(rules, RateMetric::Requests); - // No rules configured ⇒ trivially pass. Skip Redis entirely so // a misconfigured surface (no rules attached) doesn't pay a // round-trip per request. - if resolved.is_empty() { + if rules.is_empty() { return Ok(LimitsChecked { identity: state.identity, trace_id: state.trace_id, @@ -56,63 +64,45 @@ pub async fn check_limits( }); } - let outcome = match sliding::check_and_record(redis, &resolved, 1, !fail_closed).await { + let outcome = match sliding::admit(redis, rules, owner, !fail_closed).await { Ok(o) => o, Err(e) => { - if fail_closed { - metrics::counter!("lifecycle_rate_limiter_unavailable_total").increment(1); - tracing::warn!( - error = %e, - trace_id = %state.trace_id, - "rate limiter unavailable; failing closed" - ); - let entry = S::audit_entry(&state.identity, "rate_limiter_unavailable") - .trace_id(state.trace_id.clone()); - audit.log(entry); - return Err(S::rate_limiter_unavailable_response()); - } - // Fail-open path: bump the counter, log at warn, and - // synthesise an "allowed" outcome with empty currents - // so downstream stages see an unobjectionable result. - metrics::counter!("lifecycle_rate_limiter_fail_open_total").increment(1); + // `admit` only errs when failing closed; it fails open by + // itself. + metrics::counter!("lifecycle_rate_limiter_unavailable_total").increment(1); tracing::warn!( error = %e, trace_id = %state.trace_id, - "rate limiter unavailable; failing open" + "rate limiter unavailable; failing closed" ); - sliding::CheckOutcome { - allowed: true, - exceeded_index: -1, - currents: Vec::new(), - } + let entry = S::audit_entry(&state.identity, "rate_limiter_unavailable") + .trace_id(state.trace_id.clone()); + audit.log(entry); + return Err(S::rate_limiter_unavailable_response()); } }; if !outcome.allowed { - // `exceeded_index` is into the `requests`-filtered slice; - // re-filter to find the rule that tripped so we can label - // the response with `subject:metric/window`. - let label = (outcome.exceeded_index >= 0) - .then(|| { - rules - .iter() - .filter(|r| r.metric == RateMetric::Requests) - .nth(outcome.exceeded_index as usize) - .map(sliding::rate_label) - }) - .flatten() + let label = outcome + .exceeded_index + .and_then(|i| rules.get(i)) + .map(sliding::rate_label) .unwrap_or_else(|| "rate limit".to_string()); metrics::counter!("lifecycle_rate_limited_total").increment(1); tracing::warn!( trace_id = %state.trace_id, limit = %label, + retry_after_secs = outcome.retry_after_secs, "rate limited" ); let entry = S::audit_entry(&state.identity, "rate_limited") .trace_id(state.trace_id.clone()) - .detail(serde_json::json!({ "limit": label })); + .detail(serde_json::json!({ + "limit": label, + "retry_after_secs": outcome.retry_after_secs, + })); audit.log(entry); - return Err(S::rate_limited_response(&label)); + return Err(S::rate_limited_response(&label, outcome.retry_after_secs)); } Ok(LimitsChecked { @@ -162,6 +152,7 @@ mod tests { let result = check_limits::( raw, &[], + user_id, &dummy_redis(), true, // fail_closed — irrelevant when no rules &dummy_audit(), diff --git a/crates/common/src/lifecycle/stages/mod.rs b/crates/common/src/lifecycle/stages/mod.rs index 0109e0c5..d0dd2ffb 100644 --- a/crates/common/src/lifecycle/stages/mod.rs +++ b/crates/common/src/lifecycle/stages/mod.rs @@ -1,7 +1,7 @@ //! Pipeline stages — each surface-agnostic stage is a plain //! `async fn` here. See [`super`] for the full pipeline shape. //! -//! Short-circuit stages (`check_limits`, `check_budget`, +//! 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 //! `Err(S::Response)`. The four post-invoke stages diff --git a/crates/common/src/lifecycle/surface.rs b/crates/common/src/lifecycle/surface.rs index 1fe02cb0..1d5a6420 100644 --- a/crates/common/src/lifecycle/surface.rs +++ b/crates/common/src/lifecycle/surface.rs @@ -70,8 +70,9 @@ pub trait Surface: Sized + Send + Sync + 'static { /// Render the surface's "you got rate limited" response. The /// `label` is a human string like `"user:requests/1m"` that /// the limit engine produces; the surface decides whether to - /// surface it verbatim or wrap it. - fn rate_limited_response(label: &str) -> Self::Response; + /// surface it verbatim or wrap it. `retry_after_secs` is how long + /// until that window has room again (HTTP `Retry-After`). + fn rate_limited_response(label: &str, retry_after_secs: u64) -> Self::Response; /// Render the surface's "rate-limiter unavailable, fail-closed" /// response. Distinct from `rate_limited_response` because the @@ -99,11 +100,10 @@ pub trait Surface: Sized + Send + Sync + 'static { /// Render the surface's "budget cap exhausted" response. The /// label is `":budget/"` (e.g. /// `"user:budget/monthly"`) so clients can tell which cap fired - /// without parsing prose. Maps to 429 on the wire so existing - /// rate-limit retry semantics apply — `GatewayError::LocalRateLimited`'s - /// docstring explicitly covers both rate and budget under that - /// status family. - fn budget_exceeded_response(label: &str) -> Self::Response; + /// without parsing prose. `retry_after_secs` is how long until the + /// cap's period ends — a budget does not free before that, so the + /// AI gateway also tells SDKs not to retry on their own. + fn budget_exceeded_response(label: &str, retry_after_secs: u64) -> Self::Response; /// Render the surface's "budget read backend unavailable" /// response (Redis outage + `fail_closed` enabled). Same wire diff --git a/crates/common/src/lifecycle/test_surface.rs b/crates/common/src/lifecycle/test_surface.rs index 6cc9758f..129e3374 100644 --- a/crates/common/src/lifecycle/test_surface.rs +++ b/crates/common/src/lifecycle/test_surface.rs @@ -30,6 +30,7 @@ pub enum TestResponse { Ok, RateLimited { label: String, + retry_after_secs: u64, }, RateLimiterUnavailable, AccessDenied { @@ -37,6 +38,7 @@ pub enum TestResponse { }, BudgetExceeded { label: String, + retry_after_secs: u64, }, BudgetUnavailable, } @@ -133,9 +135,10 @@ impl Surface for TestSurface { .audit(action) } - fn rate_limited_response(label: &str) -> Self::Response { + fn rate_limited_response(label: &str, retry_after_secs: u64) -> Self::Response { TestResponse::RateLimited { label: label.to_owned(), + retry_after_secs, } } @@ -156,9 +159,10 @@ impl Surface for TestSurface { } } - fn budget_exceeded_response(label: &str) -> Self::Response { + fn budget_exceeded_response(label: &str, retry_after_secs: u64) -> Self::Response { TestResponse::BudgetExceeded { label: label.to_owned(), + retry_after_secs, } } diff --git a/crates/common/src/limits/budget.rs b/crates/common/src/limits/budget.rs index 0e1d9346..4869ccfb 100644 --- a/crates/common/src/limits/budget.rs +++ b/crates/common/src/limits/budget.rs @@ -24,18 +24,21 @@ // Why a separate module from `sliding`: // - Period semantics are different from window semantics — a 5h // "today" budget would be confusing. -// - The check is post-hoc only (responses, not requests) so the -// all-or-nothing Lua dance is unnecessary; plain INCRBY + -// value-read is enough. +// - The pre-call check only reads (`current_spend`) and the debit +// comes after the response, so the all-or-nothing Lua dance is +// unnecessary; plain INCRBY + value-read is enough. // - The Redis key namespace is intentionally separate to avoid // collisions if a future migration changes one shape. +// +// Redis Cluster: every command here touches one key, so the counters +// need no hash tag. // ============================================================================ -use chrono::{DateTime, Datelike, Utc}; +use chrono::{DateTime, Datelike, Days, Months, NaiveTime, Utc}; use fred::clients::Client; use fred::interfaces::KeysInterface; -use super::{BudgetCap, BudgetSubject}; +use super::{BudgetCap, BudgetPeriod, BudgetSubject}; // ---------------------------------------------------------------------------- // Threshold alerting @@ -107,6 +110,29 @@ pub fn bucket_id(period: &str, now: DateTime) -> String { } } +/// When the period that contains `now` ends — and its counter stops +/// counting: the next midnight (daily), the next Monday 00:00 (weekly, +/// ISO weeks), the 1st of the next month 00:00 (monthly), all UTC like +/// [`bucket_id`]. +pub fn period_end(period: BudgetPeriod, now: DateTime) -> DateTime { + let today = now.date_naive(); + let first_day_after = match period { + BudgetPeriod::Daily => today + Days::new(1), + BudgetPeriod::Weekly => { + today + Days::new(7 - u64::from(today.weekday().num_days_from_monday())) + } + BudgetPeriod::Monthly => today.with_day(1).expect("every month has a 1st") + Months::new(1), + }; + first_day_after.and_time(NaiveTime::MIN).and_utc() +} + +/// Whole seconds from `now` until [`period_end`], rounded up and at least +/// one: what a refused request's `Retry-After` says. +pub fn secs_until_period_end(period: BudgetPeriod, now: DateTime) -> u64 { + let ms = (period_end(period, now) - now).num_milliseconds(); + super::sliding::retry_after_secs(ms) +} + /// TTL (in seconds) to set on the period counter when we INCRBY it. /// Picked at 2 × the period so a slow process can still find the /// key on the day after, but it's well gone before any chance of @@ -188,7 +214,7 @@ pub async fn current_spend( /// On Redis error this logs and returns `Ok(empty)` so the gateway /// keeps running with broken accounting (surfaced via the /// `gateway_budget_fail_open_total` metric). Same fail-open posture -/// as `sliding::check_and_record`. +/// as `sliding::admit` with `fail_open` set. pub async fn add_weighted_tokens( redis: &Client, caps: &[BudgetCap], @@ -300,6 +326,53 @@ mod tests { assert!(w.starts_with("2026-W"), "got {w}"); } + #[test] + fn a_period_ends_where_its_bucket_id_changes() { + // Wednesday 2026-04-08, mid-afternoon. + let t = Utc.with_ymd_and_hms(2026, 4, 8, 15, 30, 0).unwrap(); + let at = |y, m, d| Utc.with_ymd_and_hms(y, m, d, 0, 0, 0).unwrap(); + assert_eq!(period_end(BudgetPeriod::Daily, t), at(2026, 4, 9)); + assert_eq!(period_end(BudgetPeriod::Weekly, t), at(2026, 4, 13)); + assert_eq!(period_end(BudgetPeriod::Monthly, t), at(2026, 5, 1)); + // Boundaries: a Monday at midnight is the start of a week, not + // its end; December rolls into the next year. + assert_eq!( + period_end(BudgetPeriod::Weekly, at(2026, 4, 13)), + at(2026, 4, 20) + ); + assert_eq!( + period_end(BudgetPeriod::Monthly, at(2026, 12, 31)), + at(2027, 1, 1) + ); + for p in [ + BudgetPeriod::Daily, + BudgetPeriod::Weekly, + BudgetPeriod::Monthly, + ] { + let end = period_end(p, t); + assert_ne!( + bucket_id(p.as_str(), end - chrono::Duration::seconds(1)), + bucket_id(p.as_str(), end), + "{p:?}" + ); + assert_eq!( + bucket_id(p.as_str(), end - chrono::Duration::seconds(1)), + bucket_id(p.as_str(), t), + "{p:?}" + ); + } + } + + #[test] + fn a_spent_budget_waits_until_its_period_ends() { + let t = Utc.with_ymd_and_hms(2026, 4, 30, 23, 59, 59).unwrap() + + chrono::Duration::milliseconds(500); + assert_eq!(secs_until_period_end(BudgetPeriod::Daily, t), 1); + assert_eq!(secs_until_period_end(BudgetPeriod::Monthly, t), 1); + let t = Utc.with_ymd_and_hms(2026, 4, 30, 0, 0, 0).unwrap(); + assert_eq!(secs_until_period_end(BudgetPeriod::Daily, t), 86_400); + } + #[test] fn unknown_period_falls_back_to_monthly() { 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 9e748422..8ea3bf14 100644 --- a/crates/common/src/limits/mod.rs +++ b/crates/common/src/limits/mod.rs @@ -10,7 +10,7 @@ // // And the corresponding submodules in this folder: // -// sliding — bucketed Lua check_and_record over Redis +// sliding — bucketed Lua admit (before) / record (after) over Redis // budget — natural-period add_weighted_tokens / check_cap // weight — model_id → weighted token converter (LRU cached) // @@ -31,14 +31,22 @@ // Role- and team-level constraints are NOT their own subjects — they live // in `rbac_roles.policy_document` (and, if ever added, the analogous field // on teams) and fold into each member's merged policy at request time, -// materializing as `subject = User` rules. Redis counters therefore stay -// user-scoped and grouping membership never becomes a shared pool. +// materializing as `subject = User` rules. Grouping membership therefore +// never becomes a shared pool. // -// At request time the proxy resolves which subjects apply (user + -// api_key, plus the merged role/team constraint set attributed to that -// same user) and runs every matching enabled rule through -// `sliding::check_and_record`. Any single failure rejects the request — -// Lua handles the all-or-nothing INCR. +// At request time the proxy builds a [`RequestLimits`]: the user's rules +// and caps (role defaults with the user's overrides) on the user's +// counters, and the calling key's own rules and caps on the key +// lineage's counters. Both sets apply — a key's limits narrow what its +// owner may do through that key, they never widen the owner's. Every +// rule goes through one `sliding::admit` call, so a request refused by +// any of them charges none. +// +// Redis Cluster: every rate-limit counter of a request carries the hash +// tag `{user:}` (see `sliding::counter_key`), so the user's and +// the key's counters share a slot and one script can check them all. +// Budget counters are read and written one key per command and need no +// tag. // // See `plan.md` (limits chapter) for the full design. // ============================================================================ @@ -1138,6 +1146,104 @@ fn apply_block_overrides(target: &mut Option, overrides: Option, + pub caps: Vec, +} + +impl RequestLimits { + /// `user` is the owner's constraints — role defaults with the user's + /// overrides — and counts on the user's counters. `key`, when the + /// request came with an API key that has limits of its own, is the + /// key's lineage id and those limits; they count on the lineage's + /// counters. Both apply: a key's limits narrow its owner's, never + /// widen them. + pub fn for_request( + surface: Surface, + owner: Uuid, + user: &SurfaceConstraints, + key: Option<(Uuid, &SurfaceConstraints)>, + ) -> Self { + let mut out = Self { + owner, + ..Self::default() + }; + out.add(surface, RateLimitSubject::User, owner, user); + if let Some((lineage, constraints)) = key { + out.add( + surface, + RateLimitSubject::ApiKeyLineage, + lineage, + constraints, + ); + } + out + } + + fn add( + &mut self, + surface: Surface, + subject_kind: RateLimitSubject, + subject_id: Uuid, + constraints: &SurfaceConstraints, + ) { + let Some(block) = constraints.block(surface) else { + return; + }; + let budget_kind = match subject_kind { + RateLimitSubject::User => BudgetSubject::User, + RateLimitSubject::ApiKeyLineage => BudgetSubject::ApiKeyLineage, + }; + self.rules.extend( + block + .rules + .iter() + .filter(|r| r.enabled) + .map(|r| RateLimitRule { + id: Uuid::nil(), + subject_kind, + subject_id, + surface, + metric: r.metric, + window_secs: r.window_secs, + max_count: r.max_count, + enabled: true, + expires_at: None, + reason: None, + created_by: None, + }), + ); + self.caps.extend( + block + .budgets + .iter() + .filter(|b| b.enabled) + .map(|b| BudgetCap { + id: Uuid::nil(), + subject_kind: budget_kind, + subject_id, + period: b.period, + limit_tokens: b.limit_tokens, + enabled: true, + expires_at: None, + reason: None, + created_by: None, + }), + ); + } +} + // ---------------------------------------------------------------------------- // Cache-invalidation pubsub // @@ -1584,6 +1690,67 @@ mod tests { ); } + #[test] + fn a_request_is_held_to_its_users_limits_and_its_keys_each_on_their_own_counter() { + let user = Uuid::new_v4(); + let lineage = Uuid::new_v4(); + let block = |max: i64, budget: i64| SurfaceConstraints { + ai_gateway: Some(SurfaceBlock { + rules: vec![SurfaceRule { + metric: RateMetric::Requests, + window_secs: 60, + max_count: max, + enabled: true, + }], + budgets: vec![SurfaceBudget { + period: BudgetPeriod::Daily, + limit_tokens: budget, + enabled: true, + }], + }), + mcp_gateway: None, + }; + let limits = RequestLimits::for_request( + Surface::AiGateway, + user, + &block(1, 100), + Some((lineage, &block(5, 500))), + ); + assert_eq!(limits.owner, user); + // The key's rule of the same window does not replace the user's: + // both are checked. + let rules: Vec<_> = limits + .rules + .iter() + .map(|r| (r.subject_kind, r.subject_id, r.max_count)) + .collect(); + assert_eq!( + rules, + vec![ + (RateLimitSubject::User, user, 1), + (RateLimitSubject::ApiKeyLineage, lineage, 5), + ] + ); + let caps: Vec<_> = limits + .caps + .iter() + .map(|c| (c.subject_kind, c.subject_id, c.limit_tokens)) + .collect(); + assert_eq!( + caps, + vec![ + (BudgetSubject::User, user, 100), + (BudgetSubject::ApiKeyLineage, lineage, 500), + ] + ); + // Another surface's block is not this request's. + assert!( + RequestLimits::for_request(Surface::McpGateway, user, &block(1, 100), None) + .rules + .is_empty() + ); + } + #[test] fn window_to_secs_roundtrip() { for (s, expected) in [ diff --git a/crates/common/src/limits/sliding.rs b/crates/common/src/limits/sliding.rs index ef5cd682..fb340f13 100644 --- a/crates/common/src/limits/sliding.rs +++ b/crates/common/src/limits/sliding.rs @@ -4,9 +4,9 @@ // A pure-sliding window over a 1-week timespan would need ~50k members // in a single Redis ZSET (one per request) and would dominate memory // at any meaningful traffic. We approximate with **fixed-bucket -// sliding**: each window is split into 60 buckets, each bucket is one -// INCR counter, the "current value" is the sum of the last 60 buckets. -// Precision is ~1.6%, more than enough for rate limiting. +// sliding**: each window is split into 60 buckets, the "current value" +// is the sum of the last 60 buckets. Precision is ~1.6%, more than +// enough for rate limiting. // // Bucket sizing per window: // @@ -17,28 +17,43 @@ // 1d → 24m × 60 // 1w → 168m × 60 // -// All N rules for a request go through ONE Lua script invocation. The -// script: +// Storage: one Redis hash per counter, its fields the bucket ids +// (`now_secs / bucket_secs`), each holding what that bucket counted: // -// 1. For each rule, computes its `current` = sum(last 60 bucket counters). -// 2. Checks `current + cost <= max_count` for every rule. -// 3. If any rule would be exceeded → returns `{0, exceeded_rule_index}` -// and writes nothing (atomic all-or-nothing). -// 4. Otherwise INCRs every rule's current bucket by `cost` and refreshes -// the bucket TTL to `2 * window_secs` (so old buckets self-evict). -// 5. Returns `{1, ...}` on success. +// ratelimit:{user:}::::: // -// One round trip, one atomic decision. The script is loaded with -// SCRIPT LOAD on first use and reused via EVALSHA — see `script_sha`. +// `` is the user the request runs as — for a key-lineage counter, +// the key's owner. The braces are a Redis Cluster hash tag: every counter +// one request touches carries the same tag, so all of them sit in one +// slot and one script can read and write them together. Each counter is +// one declared key; the script never builds a key name of its own. +// +// A request goes through two scripts: +// +// * `admit`, before the call: refuses when any rule's window is full +// (`current >= max_count`) and otherwise charges 1 to every +// `requests` rule. `tokens` rules are only read — a request can't +// know its tokens yet. All or nothing: a refused request charges no +// counter at all. A refusal also says how long until the limiting +// window has room again. +// * `record`, after the call: adds the tokens it used to every +// `tokens` rule, unconditionally. Usage that overshoots a limit is +// still recorded — the overshoot is bounded by what was in flight +// when the window filled, and the next request is refused. +// +// Old buckets are pruned by `admit` as it reads them; the hash itself +// expires 2 × window after its last write. // ============================================================================ +use std::collections::HashMap; use std::sync::OnceLock; use fred::clients::Client; use fred::interfaces::LuaInterface; use sha1::{Digest, Sha1}; +use uuid::Uuid; -use super::{RateLimitRule, RateMetric}; +use super::{RateLimitRule, RateLimitSubject, RateMetric, Surface}; /// 60 buckets per window — fixed across every supported window size. /// Larger N tightens precision but balloons Redis memory linearly. @@ -53,80 +68,138 @@ pub fn bucket_secs(window_secs: i32) -> i32 { } // ---------------------------------------------------------------------------- -// The Lua script — see file header for the algorithm +// The Lua scripts — see file header for the algorithm // ---------------------------------------------------------------------------- -const LUA_CHECK_AND_RECORD: &str = r#" --- KEYS: one base key per rule (the prefix; bucket suffix is appended in-script) --- ARGV: [now_ms, cost, n_rules, b1_secs, m1_max, b2_secs, m2_max, ...] --- Returns: {ok, idx_of_first_exceeded_rule_or_-1, [debug_current_per_rule...]} - -local now_ms = tonumber(ARGV[1]) -local cost = tonumber(ARGV[2]) -local n_rules = tonumber(ARGV[3]) -local now_secs = math.floor(now_ms / 1000) - --- Phase 1: read every rule's current sum, decide pass/fail. -local currents = {} -for i = 1, n_rules do - local bucket_secs = tonumber(ARGV[3 + (i - 1) * 2 + 1]) - local max_count = tonumber(ARGV[3 + (i - 1) * 2 + 2]) - local base_key = KEYS[i] - - -- Window starts `bucket_secs * 60` ago. Sum the 60 buckets in - -- [now - 60*bucket_secs, now], one MGET-style loop. - local sum = 0 - local current_bucket = math.floor(now_secs / bucket_secs) - for b = 0, 59 do - local bucket_id = current_bucket - b - local v = redis.call("GET", base_key .. ":" .. bucket_id) - if v then sum = sum + tonumber(v) end +const LUA_ADMIT: &str = r#" +-- KEYS: one counter hash per rule +-- ARGV: [now_ms, then (bucket_secs, max_count, charge) per rule] +-- Returns: {allowed, limiting rule (1-based, 0 when allowed), +-- ms until it has room (0 when allowed), current per rule...} + +local now_ms = tonumber(ARGV[1]) +local now_secs = math.floor(now_ms / 1000) +local n = #KEYS + +local sums = {} +local limiting, wait_ms = 0, 0 +for i = 1, n do + local bucket_secs = tonumber(ARGV[2 + (i - 1) * 3]) + local max_count = tonumber(ARGV[3 + (i - 1) * 3]) + local current = math.floor(now_secs / bucket_secs) + local oldest = current - 59 + + local counts, stale, sum = {}, {}, 0 + local fields = redis.call('HGETALL', KEYS[i]) + for j = 1, #fields, 2 do + local b = tonumber(fields[j]) + if b == nil or b < oldest then + stale[#stale + 1] = fields[j] + elseif b <= current then + local v = tonumber(fields[j + 1]) or 0 + counts[b] = v + sum = sum + v + end + end + if #stale > 0 then + redis.call('HDEL', KEYS[i], unpack(stale)) end - currents[i] = sum + sums[i] = sum - if sum + cost > max_count then - return {0, i, currents} + if sum >= max_count then + -- The window has room for one more once enough of its oldest + -- buckets have left it. Bucket b leaves at (b + 60) * bucket_secs. + local need, freed, b = sum - max_count + 1, 0, oldest + while b < current do + freed = freed + (counts[b] or 0) + if freed >= need then break end + b = b + 1 + end + local w = (b + 60) * bucket_secs * 1000 - now_ms + if w > wait_ms then + wait_ms = w + limiting = i + end end end --- Phase 2: every rule passed → INCR each rule's current bucket. -for i = 1, n_rules do - local bucket_secs = tonumber(ARGV[3 + (i - 1) * 2 + 1]) - local base_key = KEYS[i] - local current_bucket = math.floor(now_secs / bucket_secs) - local k = base_key .. ":" .. current_bucket - redis.call("INCRBY", k, cost) - -- TTL = 2 × window so old buckets vanish before they're reused. - redis.call("EXPIRE", k, bucket_secs * 60 * 2) +if limiting > 0 then + local out = {0, limiting, wait_ms} + for i = 1, n do out[#out + 1] = sums[i] end + return out end -return {1, -1, currents} +for i = 1, n do + local charge = ARGV[4 + (i - 1) * 3] + if tonumber(charge) > 0 then + local bucket_secs = tonumber(ARGV[2 + (i - 1) * 3]) + redis.call('HINCRBY', KEYS[i], math.floor(now_secs / bucket_secs), charge) + redis.call('EXPIRE', KEYS[i], bucket_secs * 120) + sums[i] = sums[i] + tonumber(charge) + end +end +local out = {1, 0, 0} +for i = 1, n do out[#out + 1] = sums[i] end +return out "#; -fn script_sha() -> &'static String { - static SHA: OnceLock = OnceLock::new(); - SHA.get_or_init(|| { +const LUA_RECORD: &str = r#" +-- KEYS: one counter hash per rule +-- ARGV: [now_ms, amount, then bucket_secs per rule] +local now_secs = math.floor(tonumber(ARGV[1]) / 1000) +for i = 1, #KEYS do + local bucket_secs = tonumber(ARGV[2 + i]) + redis.call('HINCRBY', KEYS[i], math.floor(now_secs / bucket_secs), ARGV[2]) + redis.call('EXPIRE', KEYS[i], bucket_secs * 120) +end +return #KEYS +"#; + +fn sha(script: &'static str, cell: &'static OnceLock) -> &'static str { + cell.get_or_init(|| { let mut h = Sha1::new(); - h.update(LUA_CHECK_AND_RECORD.as_bytes()); + h.update(script.as_bytes()); hex::encode(h.finalize()) }) } +/// EVALSHA, falling back to EVAL when the node doesn't have the script +/// cached yet (a restart, a failover, a node new to the cluster). +async fn run( + redis: &Client, + script: &'static str, + cell: &'static OnceLock, + keys: Vec, + args: Vec, +) -> Result, fred::error::Error> { + match redis + .evalsha::, _, _, _>(sha(script, cell), keys.clone(), args.clone()) + .await + { + Ok(v) => Ok(v), + Err(e) if e.details().starts_with("NOSCRIPT") => redis.eval(script, keys, args).await, + Err(e) => Err(e), + } +} + +static ADMIT_SHA: OnceLock = OnceLock::new(); +static RECORD_SHA: OnceLock = OnceLock::new(); + // ---------------------------------------------------------------------------- // Public API // ---------------------------------------------------------------------------- /// Render a rule's identity as the canonical /// `:/` label (e.g. -/// `api_key:tokens/1h`, `user:requests/1m`). Used by: -/// - the `Retry-After`-style HTTP body for rate-limited responses, +/// `api_key_lineage:tokens/1h`, `user:requests/1m`). Used by: +/// - the HTTP body for rate-limited responses, /// - the `gateway_logs` / `mcp_logs` audit row's `limits.label`, /// - log scrapers that group by `{subject_kind, metric, window}`. /// /// Lives next to the engine because every surface that runs -/// `check_and_record` also needs to render the label that caused a -/// deny — keeping them apart bred two identical copies (one each in -/// the AI gateway and MCP gateway). One implementation, one format. +/// [`admit`] also needs to render the label that caused a deny — +/// keeping them apart bred two identical copies (one each in the AI +/// gateway and MCP gateway). One implementation, one format. pub fn rate_label(rule: &super::RateLimitRule) -> String { let window = match rule.window_secs { 60 => "1m".to_string(), @@ -145,202 +218,264 @@ pub fn rate_label(rule: &super::RateLimitRule) -> String { ) } -/// One rule resolved against the current request's metric, ready to -/// hand to `check_and_record`. The `cost` parameter on the helper -/// applies to all rules in the slice equally; callers MUST pre-filter -/// the slice to a single metric before calling. +/// The Redis Cluster hash tag every limit counter of `owner`'s requests +/// carries, so one script can touch all of them. +pub fn hash_tag(owner: Uuid) -> String { + format!("{{user:{owner}}}") +} + +/// The Redis key of one counter. `owner` is the user the requests run +/// as: the subject itself for a user rule, the key's owner for a key +/// lineage rule. +pub fn counter_key( + owner: Uuid, + surface: Surface, + subject_kind: RateLimitSubject, + subject_id: Uuid, + metric: RateMetric, + window_secs: i32, +) -> String { + format!( + "ratelimit:{}:{}:{}:{subject_id}:{}:{window_secs}", + hash_tag(owner), + surface.as_str(), + subject_kind.as_str(), + metric.as_str() + ) +} + +/// One rule's counter, ready for the scripts. #[derive(Debug, Clone)] pub struct ResolvedRule { - pub id: uuid::Uuid, - pub base_key: String, + pub id: Uuid, + pub key: String, pub bucket_secs: i32, pub max_count: i64, } +impl ResolvedRule { + pub fn new(rule: &RateLimitRule, owner: Uuid) -> Self { + Self { + id: rule.id, + key: counter_key( + owner, + rule.surface, + rule.subject_kind, + rule.subject_id, + rule.metric, + rule.window_secs, + ), + bucket_secs: bucket_secs(rule.window_secs), + max_count: rule.max_count, + } + } +} + #[derive(Debug, Clone)] pub struct CheckOutcome { pub allowed: bool, - /// Index into the `rules` slice of the first rule that would be - /// exceeded. -1 when allowed. - pub exceeded_index: i32, - /// Current sum (pre-INCR) for every rule in the slice, in input - /// order. Useful for surfacing "X / Y used" in the response or - /// for `gateway_rate_limit_remaining` headers. + /// Index into the rules handed to [`admit`] of the rule that + /// refused the request — of those whose windows were full, the one + /// that frees last. `None` when allowed. + pub exceeded_index: Option, + /// Seconds until that rule's window has room for the request again, + /// rounded up. 0 when allowed. + pub retry_after_secs: u64, + /// Each rule's count in its window, after this request's charge + /// when it was allowed, in input order. Empty when Redis failed + /// open. pub currents: Vec, } -/// Build a Redis key prefix for a rule. The bucket id is appended by -/// the Lua script. Format: -/// -/// ratelimit::::: -/// -/// Splitting on the metric AND window means a (subject, surface) pair -/// can hold multiple rules without colliding their counters. -pub fn build_base_key( - surface: &str, - subject_kind: &str, - subject_id: uuid::Uuid, - metric: RateMetric, - window_secs: i32, -) -> String { - format!( - "ratelimit:{surface}:{subject_kind}:{subject_id}:{}:{window_secs}", - metric.as_str() - ) -} - -/// Convert a slice of `RateLimitRule` rows into the engine-facing -/// `ResolvedRule` shape, dropping any rules whose metric doesn't match -/// `metric_filter`. Use this to split a "load every rule" result into -/// the requests-pass and tokens-pass batches. -pub fn resolve_rules(rules: &[RateLimitRule], metric_filter: RateMetric) -> Vec { - rules - .iter() - .filter(|r| r.metric == metric_filter) - .map(|r| ResolvedRule { - id: r.id, - base_key: build_base_key( - r.surface.as_str(), - r.subject_kind.as_str(), - r.subject_id, - r.metric, - r.window_secs, - ), - bucket_secs: bucket_secs(r.window_secs), - max_count: r.max_count, - }) - .collect() +impl CheckOutcome { + fn open() -> Self { + Self { + allowed: true, + exceeded_index: None, + retry_after_secs: 0, + currents: Vec::new(), + } + } } -/// Atomic "check then INCR-by-cost" for an entire batch of rules. +/// Before the call: refuse the request if any rule's window is full, +/// otherwise charge 1 to every `requests` rule. `tokens` rules are only +/// read; [`record`] charges them after the call. Atomic across every +/// rule — a refused request charges nothing. /// -/// `cost` is in metric units: `1` for `requests`, weighted-token count -/// for `tokens`. All rules in the slice MUST share a single metric and -/// must already be filtered through `resolve_rules`. +/// `owner` is the user the request runs as (see [`counter_key`]). /// /// On Redis error the policy is controlled by `fail_open`: /// /// * `fail_open = true` (default): returns `Ok(allowed = true)` with /// empty `currents`, bumps `gateway_rate_limiter_fail_open_total`, -/// and lets the request through. This is the historical behavior -/// and matches what most operators want — Redis should not be a -/// single point of failure for the AI control plane. +/// and lets the request through. Redis should not be a single point +/// of failure for the AI control plane. /// * `fail_open = false`: returns the underlying `fred` error so /// callers can refuse the request. Wired up via the -/// `security.rate_limit_fail_closed` system setting; flip it on -/// when the limits engine is the only thing standing between you -/// and a cost incident. -pub async fn check_and_record( +/// `security.rate_limit_fail_closed` system setting. +pub async fn admit( redis: &Client, - rules: &[ResolvedRule], - cost: i64, + rules: &[RateLimitRule], + owner: Uuid, fail_open: bool, ) -> Result { - if rules.is_empty() || cost <= 0 { - return Ok(CheckOutcome { - allowed: true, - exceeded_index: -1, - currents: Vec::new(), - }); - } - - let now_ms = chrono::Utc::now().timestamp_millis(); - let n_rules = rules.len(); - - // Build KEYS[] in the same order as the loop below appends ARGV - // pairs so the script's index math lines up. - let keys: Vec = rules.iter().map(|r| r.base_key.clone()).collect(); + admit_at( + redis, + rules, + owner, + chrono::Utc::now().timestamp_millis(), + fail_open, + ) + .await +} - // ARGV layout: [now_ms, cost, n_rules, then (bucket_secs, max_count) × n] - let mut args: Vec = Vec::with_capacity(3 + n_rules * 2); +/// [`admit`] at a given time, in Unix milliseconds — for tests. +pub async fn admit_at( + redis: &Client, + rules: &[RateLimitRule], + owner: Uuid, + 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()); - args.push(cost.to_string()); - args.push(n_rules.to_string()); - for r in rules { + 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(), + }); } - // Try EVALSHA first; on NOSCRIPT load and retry. fred's evalsha - // surfaces script-not-found as an error so we fall back manually. - let raw: Result, _> = redis - .evalsha(script_sha().as_str(), keys.clone(), args.clone()) - .await; - let result: Vec = match raw { + let reply = match run(redis, LUA_ADMIT, &ADMIT_SHA, keys, args).await { Ok(v) => v, + 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()); + } Err(e) => { - // Either NOSCRIPT or a transient redis problem. Try a - // direct EVAL once — Redis caches the loaded script for - // subsequent EVALSHA calls. - tracing::debug!("rate-limit evalsha failed ({e}); falling back to EVAL"); - match redis - .eval::, _, _, _>(LUA_CHECK_AND_RECORD, keys, args) - .await - { - Ok(v) => v, - Err(e) => { - if fail_open { - tracing::warn!("rate-limit EVAL failed: {e}; failing open"); - metrics::counter!("gateway_rate_limiter_fail_open_total").increment(1); - return Ok(CheckOutcome { - allowed: true, - exceeded_index: -1, - currents: Vec::new(), - }); - } else { - tracing::error!( - "rate-limit EVAL failed: {e}; failing closed per security.rate_limit_fail_closed" - ); - metrics::counter!("gateway_rate_limiter_fail_closed_total").increment(1); - return 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); } }; + Ok(parse_admit_reply(&reply)) +} - // Lua return shape: [ok, exceeded_index, currents...] - let ok = result.first().copied().unwrap_or(1); - let exceeded_index = result.get(1).copied().unwrap_or(-1) as i32; - let currents = result.iter().skip(2).copied().collect(); - - Ok(CheckOutcome { - allowed: ok == 1, - // Lua is 1-indexed; convert to 0-indexed for Rust callers. - exceeded_index: if exceeded_index >= 1 { - exceeded_index - 1 +/// Reply shape: `[allowed, limiting rule (1-based, 0 = none), wait_ms, +/// currents...]`. +fn parse_admit_reply(reply: &[i64]) -> CheckOutcome { + let allowed = reply.first().copied().unwrap_or(1) == 1; + let limiting = reply.get(1).copied().unwrap_or(0); + let wait_ms = reply.get(2).copied().unwrap_or(0).max(0); + CheckOutcome { + allowed, + exceeded_index: (!allowed && limiting >= 1).then(|| (limiting - 1) as usize), + retry_after_secs: if allowed { + 0 } else { - -1 + retry_after_secs(wait_ms) }, - currents, - }) + currents: reply.iter().skip(3).copied().collect(), + } +} + +/// Whole seconds a client should wait for `wait_ms` to pass: rounded up, +/// and at least one — `Retry-After: 0` reads as "now". +pub fn retry_after_secs(wait_ms: i64) -> u64 { + (wait_ms.max(0) as u64).div_ceil(1000).max(1) +} + +/// After the call: add `amount` to every rule of `metric`, whether or +/// not that takes it past its limit — the call has been made, and +/// hiding what it used would only let the next one through too. +pub async fn record( + redis: &Client, + rules: &[RateLimitRule], + owner: Uuid, + metric: RateMetric, + amount: i64, +) -> Result<(), fred::error::Error> { + record_at( + redis, + rules, + owner, + metric, + amount, + chrono::Utc::now().timestamp_millis(), + ) + .await +} + +/// [`record`] at a given time, in Unix milliseconds — for tests. +pub async fn record_at( + redis: &Client, + rules: &[RateLimitRule], + owner: Uuid, + metric: RateMetric, + amount: i64, + now_ms: i64, +) -> Result<(), fred::error::Error> { + let resolved: Vec = rules + .iter() + .filter(|r| r.metric == metric) + .map(|r| ResolvedRule::new(r, owner)) + .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(|_| ()) } /// Read-only "what's the current sum for this rule" helper. Used by -/// the limits CRUD's usage endpoint to render "X / Y used" without -/// any side effects. Walks the same 60 buckets the Lua script does -/// but issues plain GETs instead of an INCR. +/// the console's usage views to render "X / Y used" without any side +/// effects. /// /// Returns 0 on Redis error so the UI can fall back to "no data" /// rather than 500. Real failures are logged. pub async fn current_count(redis: &Client, rule: &ResolvedRule) -> i64 { - use fred::interfaces::KeysInterface; - let now_secs = chrono::Utc::now().timestamp(); - let bucket_secs = rule.bucket_secs as i64; + use fred::interfaces::HashesInterface; + match redis.hgetall::, _>(&rule.key).await { + Ok(buckets) => window_sum(&buckets, chrono::Utc::now().timestamp(), rule.bucket_secs), + Err(e) => { + tracing::warn!(key = %rule.key, "rate-limit usage read failed: {e}"); + 0 + } + } +} + +/// The sum of the buckets inside the window that ends at `now_secs`. +fn window_sum(buckets: &HashMap, now_secs: i64, bucket_secs: i32) -> i64 { + let bucket_secs = i64::from(bucket_secs); if bucket_secs <= 0 { return 0; } - let current_bucket = now_secs / bucket_secs; - let mut sum: i64 = 0; - for b in 0..BUCKETS_PER_WINDOW { - let bucket_id = current_bucket - b; - let key = format!("{}:{}", rule.base_key, bucket_id); - let v: Option = redis.get(&key).await.ok().flatten(); - if let Some(n) = v { - sum += n; - } - } - sum + let current = now_secs.div_euclid(bucket_secs); + let oldest = current - (BUCKETS_PER_WINDOW - 1); + buckets + .iter() + .filter_map(|(b, v)| b.parse::().ok().map(|b| (b, *v))) + .filter(|(b, _)| (oldest..=current).contains(b)) + .map(|(_, v)| v) + .sum() } #[cfg(test)] @@ -364,4 +499,108 @@ mod tests { assert_eq!(bucket_secs(86_400), 1_440); assert_eq!(bucket_secs(604_800), 10_080); } + + fn rule(kind: RateLimitSubject, subject: Uuid, metric: RateMetric) -> RateLimitRule { + RateLimitRule { + id: Uuid::nil(), + subject_kind: kind, + subject_id: subject, + surface: Surface::AiGateway, + metric, + window_secs: 60, + max_count: 1, + enabled: true, + expires_at: None, + reason: None, + created_by: None, + } + } + + /// Redis Cluster runs a script only when every key it declares hashes + /// to one slot. A request's counters — the user's and the key + /// lineage's, requests and tokens — must all land together. + #[test] + fn every_counter_of_a_request_hashes_to_one_slot() { + use fred::util::redis_keyslot; + let user = Uuid::new_v4(); + let lineage = Uuid::new_v4(); + let keys: Vec = [ + rule(RateLimitSubject::User, user, RateMetric::Requests), + rule(RateLimitSubject::User, user, RateMetric::Tokens), + rule( + RateLimitSubject::ApiKeyLineage, + lineage, + RateMetric::Requests, + ), + rule(RateLimitSubject::ApiKeyLineage, lineage, RateMetric::Tokens), + ] + .iter() + .map(|r| ResolvedRule::new(r, user).key) + .collect(); + let slot = redis_keyslot(keys[0].as_bytes()); + for k in &keys { + assert_eq!(redis_keyslot(k.as_bytes()), slot, "{k}"); + } + // The tag is what decides: the slot is the tag's own. + assert_eq!(slot, redis_keyslot(format!("user:{user}").as_bytes())); + } + + #[test] + fn a_key_lineage_counter_is_named_after_the_lineage_under_its_owners_tag() { + let user = Uuid::new_v4(); + let lineage = Uuid::new_v4(); + let key = ResolvedRule::new( + &rule(RateLimitSubject::ApiKeyLineage, lineage, RateMetric::Tokens), + user, + ) + .key; + assert_eq!( + key, + format!("ratelimit:{{user:{user}}}:ai_gateway:api_key_lineage:{lineage}:tokens:60") + ); + } + + #[test] + fn the_window_sum_counts_the_last_sixty_buckets_only() { + let now = 10_000; + let b = |id: i64, v: i64| (id.to_string(), v); + let buckets: HashMap = [ + b(now, 1), // current bucket + b(now - 59, 2), // oldest still inside + b(now - 60, 4), // just left + b(now + 1, 8), // a pod whose clock runs ahead + ] + .into_iter() + .chain([("garbage".to_string(), 16)]) + .collect(); + assert_eq!(window_sum(&buckets, now, 1), 3); + // Five-second buckets: bucket id = now / 5. + let five: HashMap = [b(now / 5, 7)].into_iter().collect(); + assert_eq!(window_sum(&five, now, 5), 7); + } + + #[test] + fn retry_after_rounds_up_and_never_says_now() { + assert_eq!(retry_after_secs(1), 1); + assert_eq!(retry_after_secs(1_000), 1); + assert_eq!(retry_after_secs(1_001), 2); + assert_eq!(retry_after_secs(59_999), 60); + assert_eq!(retry_after_secs(0), 1); + assert_eq!(retry_after_secs(-5), 1); + } + + #[test] + fn a_refusal_names_the_limiting_rule_and_its_wait() { + let refused = parse_admit_reply(&[0, 2, 4_500, 3, 10]); + assert!(!refused.allowed); + assert_eq!(refused.exceeded_index, Some(1)); + assert_eq!(refused.retry_after_secs, 5); + assert_eq!(refused.currents, vec![3, 10]); + + let allowed = parse_admit_reply(&[1, 0, 0, 4]); + assert!(allowed.allowed); + assert_eq!(allowed.exceeded_index, None); + assert_eq!(allowed.retry_after_secs, 0); + assert_eq!(allowed.currents, vec![4]); + } } diff --git a/crates/common/src/limits/weight.rs b/crates/common/src/limits/weight.rs index 67411571..132648d2 100644 --- a/crates/common/src/limits/weight.rs +++ b/crates/common/src/limits/weight.rs @@ -14,8 +14,8 @@ // premium. Unset cache weights follow the input weight — see // `Weights::resolve` for the ratios. // -// Used by the gateway hot path to feed `sliding::check_and_record` -// (tokens metric) and `budget::add_weighted_tokens`. The weights +// Used by the gateway hot path to feed `sliding::record` (tokens +// metric) and `budget::add_weighted_tokens`. The weights // are loaded from PG into a process-local cache the first time we // see a given model_id; subsequent calls hit the cache directly. // diff --git a/crates/common/src/redis_config.rs b/crates/common/src/redis_config.rs new file mode 100644 index 00000000..1c3932bf --- /dev/null +++ b/crates/common/src/redis_config.rs @@ -0,0 +1,194 @@ +//! The fred client config for a Redis URL, TLS included. +//! +//! Every Redis client the server builds — the command client and the +//! three Pub/Sub subscribers — goes through [`client_config`], so a +//! `rediss://` URL means the same thing for all of them. +//! +//! A `rediss` / `valkeys` scheme (with or without `-cluster` / +//! `-sentinel`) turns TLS on. The connector is built here rather than +//! left to fred, which would build its own on `ClientConfig::builder()`: +//! that picks the process-wide rustls crypto provider, and this binary +//! compiles in two (aws-lc-rs for reqwest 0.13, ring for the reqwest +//! 0.12 under openidconnect), so rustls cannot choose one and panics. +//! Here the provider is named — aws-lc-rs, as reqwest uses — and the +//! server certificate is checked the way the HTTP client checks an +//! upstream's: against the platform's roots, through +//! rustls-platform-verifier. Managed Redis services (ElastiCache, +//! Upstash, Azure Cache, Redis Cloud) present certificates from public +//! CAs, which those roots cover. +//! +//! A self-hosted Redis whose certificate a private CA signed needs that +//! CA: `REDIS_CA_CERT` names a PEM file holding it (one certificate or +//! several), and then only those certificates are trusted — the same as +//! `redis-cli --cacert`. + +use std::path::Path; +use std::sync::Arc; + +use anyhow::Context; +use fred::types::config::{Config, TlsConfig, TlsConnector, TlsHostMapping}; +use rustls::pki_types::CertificateDer; +use rustls::pki_types::pem::PemObject; + +/// The fred config for `url`, with a TLS connector when its scheme asks +/// for TLS. `ca_cert` is the PEM file of `REDIS_CA_CERT`; it only +/// matters for a TLS URL. +/// +/// Errors never quote `url`: it carries the Redis password. +pub fn client_config(url: &str, ca_cert: Option<&Path>) -> anyhow::Result { + let mut parsed = url::Url::parse(url).context("REDIS_URL is not a valid URL")?; + // Hand fred the plain-TCP twin of a TLS scheme, so that it parses + // the URL without building a connector of its own (see the module + // docs), and attach ours. + let tls = match plain_twin(parsed.scheme()) { + Some(plain) => { + parsed + .set_scheme(&plain) + .map_err(|()| anyhow::anyhow!("REDIS_URL has an unusable scheme"))?; + true + } + None => false, + }; + let mut config = Config::from_url(parsed.as_str()).context("REDIS_URL is not usable")?; + if tls { + config.tls = Some(TlsConfig { + connector: tls_connector(ca_cert)?, + // A cluster node is reached at the address it announces, and + // its certificate must name that address (host name or IP). + hostnames: TlsHostMapping::None, + }); + } + Ok(config) +} + +/// Whether `url` asks for TLS, by the same rule fred applies. +pub fn uses_tls(url: &str) -> bool { + url::Url::parse(url).is_ok_and(|u| plain_twin(u.scheme()).is_some()) +} + +/// `rediss` → `redis`, `rediss-cluster` → `redis-cluster`, `valkeys` → +/// `valkey`, …; `None` for a scheme without TLS. +fn plain_twin(scheme: &str) -> Option { + [("rediss", "redis"), ("valkeys", "valkey")] + .into_iter() + .find_map(|(tls, plain)| { + scheme + .strip_prefix(tls) + .map(|rest| format!("{plain}{rest}")) + }) +} + +fn tls_connector(ca_cert: Option<&Path>) -> anyhow::Result { + let provider = Arc::new(rustls::crypto::aws_lc_rs::default_provider()); + let builder = rustls::ClientConfig::builder_with_provider(provider.clone()) + .with_safe_default_protocol_versions() + .context("Redis TLS: no usable protocol version")?; + let config = match ca_cert { + Some(path) => builder.with_root_certificates(private_roots(path)?), + None => { + let verifier = rustls_platform_verifier::Verifier::new(provider) + .context("Redis TLS: could not load the system's CA certificates")?; + // `dangerous()` is only rustls's door to a verifier of one's + // own; this one verifies fully (chain and host name), as it + // does inside reqwest. + builder + .dangerous() + .with_custom_certificate_verifier(Arc::new(verifier)) + } + }; + Ok(config.with_no_client_auth().into()) +} + +/// The certificates of `REDIS_CA_CERT`, as the only roots trusted. +fn private_roots(path: &Path) -> anyhow::Result { + let shown = path.display(); + let mut roots = rustls::RootCertStore::empty(); + for cert in CertificateDer::pem_file_iter(path) + .with_context(|| format!("REDIS_CA_CERT: cannot read {shown}"))? + { + let cert = cert.with_context(|| format!("REDIS_CA_CERT: {shown} is not valid PEM"))?; + roots + .add(cert) + .with_context(|| format!("REDIS_CA_CERT: {shown} holds an unusable certificate"))?; + } + anyhow::ensure!( + !roots.is_empty(), + "REDIS_CA_CERT: {shown} holds no certificate" + ); + Ok(roots) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn plain_urls_have_no_tls() { + for url in [ + "redis://:pw@localhost:6379/1", + "redis-cluster://:pw@redis-0:6379?node=redis-1:6379", + "valkey://localhost", + ] { + let config = client_config(url, None).unwrap(); + assert!(config.tls.is_none(), "{url}"); + assert!(!uses_tls(url), "{url}"); + } + } + + #[test] + fn tls_urls_get_our_connector_and_keep_everything_else() { + let config = client_config("rediss://user:p%40ss@cache.example.com:6380/3", None).unwrap(); + assert!(config.uses_rustls()); + assert_eq!(config.username.as_deref(), Some("user")); + assert_eq!(config.password.as_deref(), Some("p@ss")); + assert_eq!(config.database, Some(3)); + let fred::types::config::ServerConfig::Centralized { server } = &config.server else { + panic!("not centralized: {:?}", config.server); + }; + assert_eq!((&*server.host, server.port), ("cache.example.com", 6380)); + + let config = client_config( + "rediss-cluster://:pw@redis-0:6379?node=redis-1:6379&node=redis-2:6379", + None, + ) + .unwrap(); + assert!(config.uses_rustls()); + assert!(config.server.is_clustered()); + assert_eq!(config.server.hosts().len(), 3); + + for url in [ + "valkeys://localhost", + "rediss-sentinel://localhost:26379/0?sentinelServiceName=m", + ] { + assert!(client_config(url, None).unwrap().uses_rustls(), "{url}"); + assert!(uses_tls(url), "{url}"); + } + } + + #[test] + fn a_ca_file_without_certificates_is_refused() { + let dir = std::env::temp_dir().join(format!("tw-redis-ca-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let empty = dir.join("empty.pem"); + std::fs::write(&empty, "not a certificate\n").unwrap(); + let err = client_config("rediss://localhost", Some(&empty)).unwrap_err(); + assert!( + format!("{err:#}").contains("holds no certificate"), + "{err:#}" + ); + + let missing = dir.join("missing.pem"); + let err = client_config("rediss://localhost", Some(&missing)).unwrap_err(); + assert!(format!("{err:#}").contains("cannot read"), "{err:#}"); + std::fs::remove_dir_all(&dir).unwrap(); + + // A plain URL never reads it. + assert!(client_config("redis://localhost", Some(&missing)).is_ok()); + } + + #[test] + fn errors_do_not_quote_the_url() { + let err = client_config("rediss://:hunter2@", None).unwrap_err(); + assert!(!format!("{err:#}").contains("hunter2"), "{err:#}"); + } +} diff --git a/crates/common/src/redis_keys.rs b/crates/common/src/redis_keys.rs new file mode 100644 index 00000000..b2066d4a --- /dev/null +++ b/crates/common/src/redis_keys.rs @@ -0,0 +1,264 @@ +//! Key-space helpers that work the same on one Redis node and on a +//! Redis Cluster. +//! +//! On a cluster, `SCAN` walks only the node it is sent to, and a `DEL` +//! naming keys of different hash slots is refused (`CROSSSLOT`). A +//! pattern delete therefore has to scan every primary and delete slot +//! by slot. + +use std::collections::BTreeMap; +use std::future::Future; + +use fred::clients::Client; +use fred::interfaces::{ClientLike, KeysInterface}; +use fred::types::Key; +use fred::types::scan::ScanType; +use futures::{Stream, StreamExt}; + +/// Keys deleted per `DEL` on a single node. +const BATCH: usize = 256; + +/// A pattern delete during which a `SCAN` or `DEL` failed. The keys +/// the other commands found were deleted all the same. +#[derive(Debug, thiserror::Error)] +#[error("{failures} SCAN/DEL command(s) failed, {deleted} key(s) deleted; first error: {first}")] +pub struct IncompleteDelete { + /// Keys deleted despite the failures. + pub deleted: usize, + /// Commands that failed. + pub failures: usize, + /// The first failure. + pub first: fred::error::Error, +} + +/// Delete every key matching `pattern` (a `SCAN MATCH` glob), of +/// `kind` when given. Returns how many were deleted. Keys written while +/// the scan runs may survive it, as with any `SCAN`. +/// +/// A failed `SCAN` (of one node) or `DEL` (of one batch or slot) is +/// logged and the rest of the delete goes on, so one unreachable node +/// or refused batch doesn't leave every later key in place; the error +/// comes at the end, with the count of what was deleted. +pub async fn delete_matching( + redis: &Client, + pattern: &str, + kind: Option, +) -> Result { + let clustered = redis.is_clustered(); + let keys = if clustered { + redis + .scan_cluster_buffered(pattern.to_owned(), Some(BATCH as u32), kind) + .boxed() + } else { + redis + .scan_buffered(pattern.to_owned(), Some(BATCH as u32), kind) + .boxed() + }; + delete_all(keys, clustered, pattern, |keys| async move { + redis.del::(keys).await + }) + .await +} + +/// [`delete_matching`] past the choice of client: `del` deletes one +/// batch of keys that a single `DEL` may name. +async fn delete_all( + mut keys: S, + clustered: bool, + pattern: &str, + mut del: D, +) -> Result +where + S: Stream> + Unpin, + D: FnMut(Vec) -> F, + F: Future>, +{ + let mut tally = Tally::default(); + let mut batch: Vec = Vec::with_capacity(BATCH); + while let Some(key) = keys.next().await { + match key { + Ok(key) => batch.push(key), + // On a cluster the other nodes' scans go on; on one node + // the stream ends after its error. + Err(e) => tally.failed(pattern, "SCAN", 0, e), + } + if batch.len() == BATCH { + delete( + &mut tally, + pattern, + std::mem::take(&mut batch), + clustered, + &mut del, + ) + .await; + } + } + delete(&mut tally, pattern, batch, clustered, &mut del).await; + match tally.first { + None => Ok(tally.deleted), + Some(first) => Err(IncompleteDelete { + deleted: tally.deleted, + failures: tally.failures, + first, + }), + } +} + +#[derive(Default)] +struct Tally { + deleted: usize, + failures: usize, + first: Option, +} + +impl Tally { + fn failed(&mut self, pattern: &str, command: &str, keys: usize, e: fred::error::Error) { + tracing::warn!( + pattern, command, keys, error = %e, + "Redis command failed during a pattern delete; deleting the rest" + ); + self.failures += 1; + self.first.get_or_insert(e); + } +} + +async fn delete( + tally: &mut Tally, + pattern: &str, + keys: Vec, + clustered: bool, + del: &mut D, +) where + D: FnMut(Vec) -> F, + F: Future>, +{ + if keys.is_empty() { + return; + } + let groups = if clustered { + by_slot(keys).into_values().collect() + } else { + vec![keys] + }; + for keys in groups { + let n = keys.len(); + match del(keys).await { + Ok(deleted) => tally.deleted += deleted as usize, + Err(e) => tally.failed(pattern, "DEL", n, e), + } + } +} + +/// Group keys by Redis Cluster hash slot: one `DEL` per group is one a +/// cluster accepts. +fn by_slot(keys: Vec) -> BTreeMap> { + let mut out: BTreeMap> = BTreeMap::new(); + for key in keys { + out.entry(fred::util::redis_keyslot(key.as_bytes())) + .or_default() + .push(key); + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + use fred::error::{Error, ErrorKind}; + + #[test] + fn keys_of_one_slot_are_deleted_together() { + let keys: Vec = ["{a}1", "{a}2", "{b}1", "c"] + .into_iter() + .map(Key::from) + .collect(); + let groups = by_slot(keys); + assert_eq!(groups.len(), 3); + let a = fred::util::redis_keyslot(b"a"); + assert_eq!(groups[&a].len(), 2); + for (slot, keys) in &groups { + for k in keys { + assert_eq!(fred::util::redis_keyslot(k.as_bytes()), *slot); + } + } + } + + fn keys(names: &[&str]) -> Vec> { + names.iter().map(|k| Ok(Key::from(*k))).collect() + } + + /// A `DEL` that fails for any batch naming `poison`, and records + /// every batch it was given. + fn del_failing_on<'a>( + poison: &'static str, + seen: &'a std::sync::Mutex>>, + ) -> impl FnMut(Vec) -> std::future::Ready> + 'a { + move |batch: Vec| { + let names: Vec = batch + .iter() + .map(|k| k.as_str().unwrap().to_owned()) + .collect(); + let fails = names.iter().any(|k| k == poison); + let n = names.len() as u64; + seen.lock().unwrap().push(names); + std::future::ready(if fails { + Err(Error::new(ErrorKind::IO, "connection reset")) + } else { + Ok(n) + }) + } + } + + #[tokio::test] + async fn a_failed_slot_does_not_stop_the_other_slots() { + let seen = std::sync::Mutex::new(Vec::new()); + let stream = futures::stream::iter(keys(&["{a}1", "{a}2", "{b}1", "{c}1"])); + let err = delete_all(stream, true, "*", del_failing_on("{b}1", &seen)) + .await + .unwrap_err(); + assert_eq!(seen.lock().unwrap().len(), 3, "every slot got its DEL"); + assert_eq!(err.deleted, 3); + assert_eq!(err.failures, 1); + } + + #[tokio::test] + async fn a_failed_batch_does_not_stop_the_later_batches() { + let seen = std::sync::Mutex::new(Vec::new()); + let names: Vec = (0..BATCH * 2 + 1).map(|i| format!("k{i}")).collect(); + let stream = futures::stream::iter(names.iter().map(|k| Ok(Key::from(k.as_str())))); + let err = delete_all(stream, false, "*", del_failing_on("k0", &seen)) + .await + .unwrap_err(); + assert_eq!(seen.lock().unwrap().len(), 3); + assert_eq!(err.deleted, BATCH + 1); + assert_eq!(err.failures, 1); + } + + #[tokio::test] + async fn a_failed_scan_still_deletes_what_the_others_found() { + let seen = std::sync::Mutex::new(Vec::new()); + let mut items = keys(&["a", "b"]); + items.insert(1, Err(Error::new(ErrorKind::IO, "node down"))); + let err = delete_all( + futures::stream::iter(items), + false, + "*", + del_failing_on("-", &seen), + ) + .await + .unwrap_err(); + assert_eq!(err.deleted, 2); + assert_eq!(err.failures, 1); + assert!(err.to_string().contains("node down"), "{err}"); + } + + #[tokio::test] + async fn nothing_failed_is_ok() { + let seen = std::sync::Mutex::new(Vec::new()); + let stream = futures::stream::iter(keys(&["{a}1", "{b}1"])); + let deleted = delete_all(stream, true, "*", del_failing_on("-", &seen)) + .await + .unwrap(); + assert_eq!(deleted, 2); + } +} diff --git a/crates/gateway/src/cache.rs b/crates/gateway/src/cache.rs index 68202ce0..b6bcedc5 100644 --- a/crates/gateway/src/cache.rs +++ b/crates/gateway/src/cache.rs @@ -147,34 +147,16 @@ impl ResponseCache { }) } - /// Invalidate all cached responses by deleting keys matching the cache prefix. - /// Uses Lua script for atomic pattern deletion. + /// Invalidate all cached responses by deleting keys matching the + /// cache prefix — on every node, when Redis is a cluster. pub async fn invalidate_all(&self) { - use fred::interfaces::LuaInterface; - // Use Lua EVAL to scan and delete in batches server-side - const LUA_INVALIDATE: &str = r#" -local cursor = '0' -local total = 0 -repeat - local result = redis.call('SCAN', cursor, 'MATCH', ARGV[1], 'COUNT', 100) - cursor = result[1] - local keys = result[2] - if #keys > 0 then - redis.call('DEL', unpack(keys)) - total = total + #keys - end -until cursor == '0' -return total -"#; - let deleted: i64 = self - .redis - .eval( - LUA_INVALIDATE, - Vec::::new(), - vec!["llm_cache:*".to_string()], - ) - .await - .unwrap_or(0); + let deleted = + think_watch_common::redis_keys::delete_matching(&self.redis, "llm_cache:*", None) + .await + .unwrap_or_else(|e| { + tracing::warn!("Cache invalidation incomplete: {e}"); + e.deleted + }); metrics::counter!("gateway_cache_invalidations_total").increment(1); tracing::info!(deleted, "Cache invalidated"); } diff --git a/crates/gateway/src/error.rs b/crates/gateway/src/error.rs index 38a48020..56d88a1a 100644 --- a/crates/gateway/src/error.rs +++ b/crates/gateway/src/error.rs @@ -46,12 +46,22 @@ pub enum GatewayError { /// names the account and the IAM principal. #[error("Authentication failed with upstream")] UpstreamAuthError { status: u16, message: String }, - /// Local rate limit / budget cap was hit. The String is the rule - /// label so the response body can tell the caller WHICH limit - /// fired (e.g. "user requests/5h", "api_key tokens/1d", - /// "monthly budget"). Maps to 429 in `IntoResponse`. - #[error("Rate limited: {0}")] - LocalRateLimited(String), + /// Local rate limit / budget cap was hit. `label` names the limit + /// so the response body can tell the caller WHICH one fired (e.g. + /// `user:requests/5h`, `api_key_lineage:tokens/1d`, + /// `user:budget/monthly`). Maps to 429 in `IntoResponse`, with + /// `Retry-After: retry_after_secs`. `retry` is false for a spent + /// budget: it frees only when its period ends, so SDKs that retry a + /// 429 by themselves are told not to (`x-should-retry: false`). + /// Build it with [`GatewayError::rate_limited`], + /// [`GatewayError::budget_exhausted`] or + /// [`GatewayError::limiter_unavailable`]. + #[error("Rate limited: {label}")] + LocalRateLimited { + label: String, + retry_after_secs: u32, + retry: bool, + }, /// Refused by the gateway's own policy — a content filter rule set to /// refuse matched what the caller sent, or a tool call the upstream /// returned matched a rule set to cut it. Not a malformed request (not @@ -77,7 +87,7 @@ impl GatewayError { GatewayError::ProviderInvalidResponse(_) => 502, GatewayError::TransformError(_) => 400, GatewayError::NetworkError(_) => 502, - GatewayError::UpstreamRateLimited { .. } | GatewayError::LocalRateLimited(_) => 429, + GatewayError::UpstreamRateLimited { .. } | GatewayError::LocalRateLimited { .. } => 429, GatewayError::UpstreamAuthError { .. } => 401, GatewayError::PolicyBlocked(_) => 403, } @@ -95,29 +105,75 @@ impl GatewayError { GatewayError::TransformError(_) => "TransformError", GatewayError::NetworkError(_) => "NetworkError", GatewayError::UpstreamRateLimited { .. } => "UpstreamRateLimited", - GatewayError::LocalRateLimited(_) => "LocalRateLimited", + GatewayError::LocalRateLimited { .. } => "LocalRateLimited", GatewayError::UpstreamAuthError { .. } => "UpstreamAuthError", GatewayError::PolicyBlocked(_) => "PolicyBlocked", } } + /// A rate-limit window is full. It has room again in + /// `retry_after_secs`, by itself. + pub fn rate_limited(label: impl Into, retry_after_secs: u64) -> Self { + GatewayError::LocalRateLimited { + label: label.into(), + retry_after_secs: clamp_secs(retry_after_secs), + retry: true, + } + } + + /// A budget is spent. It frees when its period ends, in + /// `retry_after_secs` — far too long for an SDK's automatic retries. + pub fn budget_exhausted(label: impl Into, retry_after_secs: u64) -> Self { + GatewayError::LocalRateLimited { + label: label.into(), + retry_after_secs: clamp_secs(retry_after_secs), + retry: false, + } + } + + /// The limit counters can't be read and the gateway fails closed. + /// Nothing says when they will be back; a short wait is a guess. + pub fn limiter_unavailable(label: impl Into) -> Self { + const UNAVAILABLE_RETRY_SECS: u32 = 30; + GatewayError::LocalRateLimited { + label: label.into(), + retry_after_secs: UNAVAILABLE_RETRY_SECS, + retry: true, + } + } + /// Hint, in seconds, for `Retry-After` on a 429 response. For - /// upstream limits we echo the upstream's own header when present; - /// for local limits we fall back to a conservative 30s so naive - /// clients don't spin into a tight retry loop while the bucket is - /// still refilling. Capped at one hour to keep the header sane - /// even when an upstream returns an absurd value. + /// upstream limits we echo the upstream's own header when present, + /// capped at one hour to keep the header sane even when an upstream + /// returns an absurd value. For local limits it is when the limit + /// lets a request through again — for a budget, the end of its + /// period, which can be weeks away. pub fn retry_after_secs(&self) -> Option { const HARD_CAP_SECS: u32 = 3600; - const LOCAL_DEFAULT_SECS: u32 = 30; match self { GatewayError::UpstreamRateLimited { retry_after_secs } => { retry_after_secs.map(|s| s.min(HARD_CAP_SECS)) } - GatewayError::LocalRateLimited(_) => Some(LOCAL_DEFAULT_SECS), + GatewayError::LocalRateLimited { + retry_after_secs, .. + } => Some(*retry_after_secs), _ => None, } } + + /// `Some(false)` when the client must not retry on its own: sent as + /// `x-should-retry`, which the OpenAI and Anthropic SDKs read before + /// retrying a 429. + pub fn should_retry(&self) -> Option { + match self { + GatewayError::LocalRateLimited { retry: false, .. } => Some(false), + _ => None, + } + } +} + +fn clamp_secs(secs: u64) -> u32 { + u32::try_from(secs).unwrap_or(u32::MAX).max(1) } /// Parse RFC 7231 `Retry-After` (delta-seconds form). HTTP-date is @@ -141,4 +197,25 @@ mod tests { assert_eq!(e.error_tag(), "PolicyBlocked"); assert_eq!(e.retry_after_secs(), None); } + + #[test] + fn a_full_window_says_when_it_frees_and_a_spent_budget_says_not_to_retry() { + let window = GatewayError::rate_limited("user:requests/1m", 17); + assert_eq!(window.status_code(), 429); + assert_eq!(window.error_tag(), "LocalRateLimited"); + assert_eq!(window.retry_after_secs(), Some(17)); + assert_eq!(window.should_retry(), None); + assert_eq!(window.to_string(), "Rate limited: user:requests/1m"); + + // A month away is not capped like an upstream's hint. + let budget = GatewayError::budget_exhausted("user:budget/monthly", 2_000_000); + assert_eq!(budget.status_code(), 429); + assert_eq!(budget.error_tag(), "LocalRateLimited"); + assert_eq!(budget.retry_after_secs(), Some(2_000_000)); + assert_eq!(budget.should_retry(), Some(false)); + + let down = GatewayError::limiter_unavailable("rate_limiter_unavailable"); + assert_eq!(down.retry_after_secs(), Some(30)); + assert_eq!(down.should_retry(), None); + } } diff --git a/crates/gateway/src/health.rs b/crates/gateway/src/health.rs index c26797c8..83a77503 100644 --- a/crates/gateway/src/health.rs +++ b/crates/gateway/src/health.rs @@ -27,6 +27,10 @@ //! express — operators tuning weights need to know whether a //! route has actually carried any requests at all. //! +//! The braces are literal: `{}` is a Redis Cluster hash tag, +//! so a route's three keys share a slot and one script (or one `DEL`) +//! can touch them together. +//! //! ### One round trip, two when the state changes //! //! Recording a completion is one Lua call: insert the sample, drop what @@ -183,11 +187,14 @@ impl CircuitBreakerConfig { } } +/// A route's samples, state and counters keys. `{}` is a Redis +/// Cluster hash tag: the scripts below declare two or three of them at +/// once, which a cluster accepts only when they share a slot. fn keys(route_id: Uuid) -> (String, String, String) { ( - format!("route_health:{route_id}:samples"), - format!("route_health:{route_id}:state"), - format!("route_health:{route_id}:counters"), + format!("route_health:{{{route_id}}}:samples"), + format!("route_health:{{{route_id}}}:state"), + format!("route_health:{{{route_id}}}:counters"), ) } @@ -401,9 +408,7 @@ impl HealthTracker { /// orphan keys after route churn. Best-effort: a Redis hiccup /// here is not worth failing the delete over. pub async fn forget(&self, route_id: Uuid) { - let samples_key = format!("route_health:{route_id}:samples"); - let state_key = format!("route_health:{route_id}:state"); - let counters_key = format!("route_health:{route_id}:counters"); + let (samples_key, state_key, counters_key) = keys(route_id); if let Err(e) = self .redis .del::(vec![samples_key, state_key, counters_key]) @@ -418,6 +423,14 @@ impl HealthTracker { mod tests { use super::*; + #[test] + fn a_routes_keys_share_one_cluster_slot() { + let (samples, state, counters) = keys(Uuid::new_v4()); + let slot = fred::util::redis_keyslot(samples.as_bytes()); + assert_eq!(fred::util::redis_keyslot(state.as_bytes()), slot); + assert_eq!(fred::util::redis_keyslot(counters.as_bytes()), slot); + } + fn cfg() -> CircuitBreakerConfig { CircuitBreakerConfig { enabled: true, diff --git a/crates/gateway/src/lifecycle/mod.rs b/crates/gateway/src/lifecycle/mod.rs index 35061314..0b72fb79 100644 --- a/crates/gateway/src/lifecycle/mod.rs +++ b/crates/gateway/src/lifecycle/mod.rs @@ -28,7 +28,7 @@ use rust_decimal::Decimal; use think_watch_common::audit::{AuditActor, AuditEntry, GatewayActor}; use think_watch_common::lifecycle::Surface; use think_watch_common::lifecycle::state::{CapturedView, Invoked, LimitCheckRecord}; -use think_watch_common::limits::{BudgetCap, RateLimitRule}; +use think_watch_common::limits::RequestLimits; use tw_dialect::ir::Dialect; use crate::guards::Guards; @@ -108,8 +108,7 @@ pub(crate) struct ChatRequestSnapshot { /// Pre-flight rule + cap lists, reused by the post-flight debit. pub(crate) struct ChatPreflightLists { - pub request_rules: Vec, - pub budget_caps: Vec, + pub limits: RequestLimits, } /// The route that actually served the request. @@ -472,20 +471,19 @@ impl Surface for ChatCompletionSurface { .audit(action) } - fn rate_limited_response(label: &str) -> Self::Response { + fn rate_limited_response(label: &str, retry_after_secs: u64) -> Self::Response { // `LocalRateLimited` so `status_code() == 429` on the wire. - // The label (`":/"`) is what the - // pre-migration `preflight_request_limits` already produced; - // keeping it intact lets clients diff exhausted windows. - ChatCompletionOutcome::ShortCircuit(GatewayError::LocalRateLimited(label.to_owned())) + // The label (`":/"`) lets clients + // diff exhausted windows; `Retry-After` says when this one + // has room again. + ChatCompletionOutcome::ShortCircuit(GatewayError::rate_limited(label, retry_after_secs)) } fn rate_limiter_unavailable_response() -> Self::Response { - // Matches the pre-migration `preflight_request_limits` - // fail-closed path — same `LocalRateLimited` variant with - // the sentinel label dashboards already filter on. - ChatCompletionOutcome::ShortCircuit(GatewayError::LocalRateLimited( - "rate_limiter_unavailable".to_owned(), + // Same `LocalRateLimited` variant with the sentinel label + // dashboards already filter on. + ChatCompletionOutcome::ShortCircuit(GatewayError::limiter_unavailable( + "rate_limiter_unavailable", )) } @@ -507,18 +505,16 @@ impl Surface for ChatCompletionSurface { ))) } - fn budget_exceeded_response(label: &str) -> Self::Response { - // `LocalRateLimited` per its docstring's explicit budget - // coverage — wire status 429, label carries which cap fired - // (e.g. `"user:budget/monthly"`) so dashboards can split - // budget exhaustion from rate-limit hits. - ChatCompletionOutcome::ShortCircuit(GatewayError::LocalRateLimited(label.to_owned())) + fn budget_exceeded_response(label: &str, retry_after_secs: u64) -> Self::Response { + // Wire status 429; the label carries which cap fired (e.g. + // `"user:budget/monthly"`) so dashboards can split budget + // exhaustion from rate-limit hits. `Retry-After` is the end of + // the cap's period, and SDKs are told not to retry by themselves. + ChatCompletionOutcome::ShortCircuit(GatewayError::budget_exhausted(label, retry_after_secs)) } fn budget_unavailable_response() -> Self::Response { - ChatCompletionOutcome::ShortCircuit(GatewayError::LocalRateLimited( - "budget_unavailable".to_owned(), - )) + ChatCompletionOutcome::ShortCircuit(GatewayError::limiter_unavailable("budget_unavailable")) } async fn record_outcome(deps: &Self::PostInvokeDeps, invoked: &Invoked) { @@ -574,8 +570,7 @@ impl Surface for ChatCompletionSurface { deps.state.weight_cache.clone(), deps.request.mapped_model.clone(), priced(&extract_usage(&invoked.view)), - deps.preflight.request_rules.clone(), - deps.preflight.budget_caps.clone(), + &deps.preflight.limits, 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 770372be..3d17dda9 100644 --- a/crates/gateway/src/proxy/accounting.rs +++ b/crates/gateway/src/proxy/accounting.rs @@ -1,12 +1,12 @@ //! Post-flight accounting + token resolution for streaming responses. //! -//! [`post_flight_account`] runs the token-metric sliding rules and -//! budget caps against 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). +//! [`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 +//! 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). //! //! All errors are logged and swallowed — by the time we get here the //! caller has already received their response, so refusing to account @@ -17,7 +17,7 @@ use std::sync::Arc; use sqlx::PgPool; use think_watch_common::dynamic_config::DynamicConfig; -use think_watch_common::limits::{self, BudgetCap, RateMetric, sliding, weight}; +use think_watch_common::limits::{self, RateMetric, RequestLimits, sliding, weight}; #[allow(clippy::too_many_arguments)] pub(crate) async fn post_flight_account( @@ -27,15 +27,11 @@ pub(crate) async fn post_flight_account( weight_cache: weight::WeightCache, model: String, tokens: weight::TokenCounts, - request_rules: Vec, - budget_caps: Vec, + request_limits: &RequestLimits, // Actor attribution for `budget.threshold_crossed` audit entries. - // Without these the crossing log carries only `cap_id`, and the - // cap's `subject_id` may be a team/role — operators investigating - // a 100 %-cross had to time-join against gateway_logs to find the - // user who pushed it over. Cloned in the streaming path (the - // tokio task moves owned values) and inlined in the non-streaming - // path; passing `Option` keeps both call shapes flat. + // Without these the crossing log carries only `cap_id`, and + // operators investigating a 100 %-cross had to time-join against + // gateway_logs to find the user who pushed it over. actor_user_id: Option, actor_user_email: Option, actor_api_key_id: Option, @@ -48,27 +44,29 @@ pub(crate) async fn post_flight_account( return; } - // Token-metric sliding rules — same rule set the pre-flight - // loaded, filtered to tokens. Post-flight always runs fail-open - // because the response has already been delivered: refusing to - // record the spend would just hide it from analytics without - // recovering anything. - let resolved_token_rules = sliding::resolve_rules(&request_rules, RateMetric::Tokens); - if !resolved_token_rules.is_empty() - && let Err(e) = - sliding::check_and_record(&redis, &resolved_token_rules, weighted, true).await + // Token-metric sliding rules — the user's and the key's, the same + // rules the pre-flight checked. Recorded whatever they come to: a + // window this request overshoots refuses the next one. Post-flight + // always runs fail-open because the response has already been + // delivered: refusing to record the spend would just hide it from + // analytics without recovering anything. + if let Err(e) = sliding::record( + &redis, + &request_limits.rules, + request_limits.owner, + RateMetric::Tokens, + weighted, + ) + .await { tracing::warn!("token rate-limit accounting failed: {e}"); } - // Natural-period budget caps derived from the user's merged - // role-inline constraints. `db` is unused here — kept on the - // signature so future per-user overrides can fold in cleanly. - let _ = db; - if !budget_caps.is_empty() { - let caps = budget_caps; + // Natural-period budget caps — the user's and the key's. + if !request_limits.caps.is_empty() { + let caps = &request_limits.caps; { - match limits::budget::add_weighted_tokens(&redis, &caps, weighted).await { + match limits::budget::add_weighted_tokens(&redis, caps, weighted).await { Ok((_statuses, crossings)) if !crossings.is_empty() => { // Emit one `budget.threshold_crossed` audit-log // entry per crossing. The action is namespaced diff --git a/crates/gateway/src/proxy/generate.rs b/crates/gateway/src/proxy/generate.rs index e391ece7..9e93eab6 100644 --- a/crates/gateway/src/proxy/generate.rs +++ b/crates/gateway/src/proxy/generate.rs @@ -849,8 +849,7 @@ async fn run( input_estimate, }, preflight: crate::lifecycle::ChatPreflightLists { - request_rules: preflight.request_rules.clone(), - budget_caps: preflight.budget_caps.clone(), + limits: preflight.limits.clone(), }, route: crate::lifecycle::ChatPickedRoute { provider_name: route.provider_name.clone(), diff --git a/crates/gateway/src/proxy/identity.rs b/crates/gateway/src/proxy/identity.rs index 1b152391..82747ef1 100644 --- a/crates/gateway/src/proxy/identity.rs +++ b/crates/gateway/src/proxy/identity.rs @@ -1,78 +1,87 @@ -//! Materialize merged surface constraints into rate-limit rule rows -//! and budget caps keyed by the authenticated user id. +//! Materialize the constraints the auth middleware resolved into the +//! rate-limit rules and budget caps one request is held to. use uuid::Uuid; use super::GatewayRequestIdentity; -use think_watch_common::limits::{ - BudgetCap, BudgetSubject, RateLimitRule, RateLimitSubject, Surface, -}; +use think_watch_common::limits::{RequestLimits, Surface}; -/// Materialize the merged surface constraints into `RateLimitRule` -/// rows keyed by the authenticated user id. Redis counters now live -/// at `ratelimit::user::...` — one set per user -/// regardless of how many roles they hold. Roles merged to empty -/// (no user_id, or no rules) produce an empty list. -pub(super) fn rules_for_ai_gateway(identity: &GatewayRequestIdentity) -> Vec { - let Some(user_id) = identity - .user_id - .as_deref() - .and_then(|s| Uuid::parse_str(s).ok()) - else { - return Vec::new(); +/// The user's limits (role defaults with the user's overrides) on the +/// user's counters, and the calling key's own limits on its lineage's +/// counters — both apply. A request with no user (never past the auth +/// middleware today) is held to nothing. +pub(super) fn limits_for_ai_gateway(identity: &GatewayRequestIdentity) -> RequestLimits { + let parse = |s: &Option| s.as_deref().and_then(|s| Uuid::parse_str(s).ok()); + let Some(user_id) = parse(&identity.user_id) else { + return RequestLimits::default(); }; - let Some(block) = identity.surface_constraints.block(Surface::AiGateway) else { - return Vec::new(); - }; - block - .rules - .iter() - .filter(|r| r.enabled) - .map(|r| RateLimitRule { - // Synthetic id — stable across a single request so the - // exceeded_index in `CheckOutcome` maps back to the same - // rule without needing a persistence layer. - id: Uuid::nil(), - subject_kind: RateLimitSubject::User, - subject_id: user_id, - surface: Surface::AiGateway, - metric: r.metric, - window_secs: r.window_secs, - max_count: r.max_count, - enabled: true, - // In-memory synthesis — override metadata lives on persisted rows only. - expires_at: None, - reason: None, - created_by: None, - }) - .collect() + let key = + parse(&identity.api_key_lineage_id).map(|lineage| (lineage, &identity.key_constraints)); + RequestLimits::for_request( + Surface::AiGateway, + user_id, + &identity.surface_constraints, + key, + ) } -pub(super) fn budgets_for_ai_gateway(identity: &GatewayRequestIdentity) -> Vec { - let Some(user_id) = identity - .user_id - .as_deref() - .and_then(|s| Uuid::parse_str(s).ok()) - else { - return Vec::new(); - }; - let Some(block) = identity.surface_constraints.block(Surface::AiGateway) else { - return Vec::new(); +#[cfg(test)] +mod tests { + use super::*; + use think_watch_common::limits::{ + RateLimitSubject, RateMetric, SurfaceBlock, SurfaceConstraints, SurfaceRule, }; - block - .budgets - .iter() - .filter(|b| b.enabled) - .map(|b| BudgetCap { - id: Uuid::nil(), - subject_kind: BudgetSubject::User, - subject_id: user_id, - period: b.period, - limit_tokens: b.limit_tokens, - enabled: true, - expires_at: None, - reason: None, - created_by: None, - }) - .collect() + + fn one_rule(max_count: i64) -> SurfaceConstraints { + SurfaceConstraints { + ai_gateway: Some(SurfaceBlock { + rules: vec![SurfaceRule { + metric: RateMetric::Requests, + window_secs: 60, + max_count, + enabled: true, + }], + budgets: vec![], + }), + mcp_gateway: None, + } + } + + #[test] + fn a_keys_rules_count_on_its_lineage_and_its_owners_on_the_user() { + let user = Uuid::new_v4(); + let lineage = Uuid::new_v4(); + let identity = GatewayRequestIdentity { + user_id: Some(user.to_string()), + api_key_id: Some(Uuid::new_v4().to_string()), + api_key_lineage_id: Some(lineage.to_string()), + surface_constraints: one_rule(10), + key_constraints: one_rule(2), + ..Default::default() + }; + let limits = limits_for_ai_gateway(&identity); + assert_eq!(limits.owner, user); + let subjects: Vec<_> = limits + .rules + .iter() + .map(|r| (r.subject_kind, r.subject_id, r.max_count)) + .collect(); + assert_eq!( + subjects, + vec![ + (RateLimitSubject::User, user, 10), + (RateLimitSubject::ApiKeyLineage, lineage, 2), + ] + ); + } + + #[test] + fn no_user_no_limits() { + let identity = GatewayRequestIdentity { + surface_constraints: one_rule(1), + ..Default::default() + }; + let limits = limits_for_ai_gateway(&identity); + assert!(limits.rules.is_empty() && limits.caps.is_empty()); + } } diff --git a/crates/gateway/src/proxy/mod.rs b/crates/gateway/src/proxy/mod.rs index 4bebb01d..45088b48 100644 --- a/crates/gateway/src/proxy/mod.rs +++ b/crates/gateway/src/proxy/mod.rs @@ -114,12 +114,17 @@ pub struct GatewayRequestIdentity { /// PG via `rotated_from_id`. pub api_key_lineage_id: Option, pub allowed_models: Option>, - /// Merged-across-roles inline limits (most restrictive per - /// surface+metric+window / surface+period). Computed once by the - /// auth middleware via `rbac::compute_user_surface_constraints` - /// and consumed directly here — no side-table lookups on the - /// hot path. + /// The user's limits: merged-across-roles inline limits (most + /// restrictive per surface+metric+window / surface+period) with the + /// user's own overrides on top. Computed once by the auth middleware + /// via `rbac::compute_user_surface_constraints` and consumed directly + /// here — no side-table lookups on the hot path. Counted on the + /// user's counters. pub surface_constraints: SurfaceConstraints, + /// The calling key's own limits (`rbac::compute_key_surface_constraints`), + /// counted on the key lineage's counters. They apply on top of + /// `surface_constraints`, never instead of them. + pub key_constraints: SurfaceConstraints, /// Resolved client IP (honours `client_ip_source` + `trusted_proxies`). /// Populated by the API-key middleware via `extract_client_ip` so /// every `gateway_logs` row carries it without each handler reading @@ -190,8 +195,8 @@ impl IntoResponse for GatewayErrorResponse { body, ) .into_response(); - // Echo the upstream's Retry-After (or our local default) so - // well-behaved clients back off the right amount instead of + // Echo the upstream's Retry-After (or when our own limit frees) + // so well-behaved clients back off the right amount instead of // burning quota with tight 3× retries that all hit the same // open window. if let Some(secs) = self.error.retry_after_secs() @@ -199,6 +204,14 @@ impl IntoResponse for GatewayErrorResponse { { response.headers_mut().insert(header::RETRY_AFTER, v); } + // A spent budget: the OpenAI and Anthropic SDKs retry a 429 by + // themselves unless told not to. + if self.error.should_retry() == Some(false) { + response.headers_mut().insert( + header::HeaderName::from_static("x-should-retry"), + HeaderValue::from_static("false"), + ); + } response } } @@ -239,7 +252,7 @@ mod helper_tests { }, 429, ), - (GatewayError::LocalRateLimited("rule".into()), 429), + (GatewayError::rate_limited("rule", 5), 429), ( GatewayError::UpstreamAuthError { status: 401, @@ -355,16 +368,37 @@ mod helper_tests { "no header when upstream didn't tell us — guessing would mislead clients" ); - // Local limit — we set our own conservative default so SDKs - // see a number instead of immediately retrying. - let resp = GatewayErrorResponse::from(GatewayError::LocalRateLimited("budget".into())) + // Local limit — when the window has room again, and nothing + // that stops an SDK's own retries. + let resp = GatewayErrorResponse::from(GatewayError::rate_limited("user:requests/1m", 42)) .into_response(); assert_eq!(resp.status().as_u16(), 429); - assert!( + assert_eq!( resp.headers() .get(axum::http::header::RETRY_AFTER) - .is_some(), - "local rate-limit must carry a Retry-After default" + .and_then(|v| v.to_str().ok()), + Some("42") + ); + assert!(resp.headers().get("x-should-retry").is_none()); + + // A spent budget — the end of its period, and no retries. + let resp = GatewayErrorResponse::from(GatewayError::budget_exhausted( + "user:budget/monthly", + 1_234_567, + )) + .into_response(); + assert_eq!(resp.status().as_u16(), 429); + assert_eq!( + resp.headers() + .get(axum::http::header::RETRY_AFTER) + .and_then(|v| v.to_str().ok()), + Some("1234567") + ); + assert_eq!( + resp.headers() + .get("x-should-retry") + .and_then(|v| v.to_str().ok()), + Some("false") ); // Non-429 responses must NOT carry Retry-After — would @@ -384,10 +418,11 @@ mod helper_tests { async fn the_error_body_is_in_the_callers_format() { use tw_dialect::ir::Dialect; async fn body(d: Dialect) -> serde_json::Value { - let resp = GatewayErrorResponse::from(GatewayError::LocalRateLimited("rule".into())) + let resp = GatewayErrorResponse::from(GatewayError::budget_exhausted("rule", 60)) .in_dialect(d) .into_response(); assert_eq!(resp.status().as_u16(), 429); + assert_eq!(resp.headers()["x-should-retry"], "false"); let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX) .await .unwrap(); diff --git a/crates/gateway/src/proxy/pipeline.rs b/crates/gateway/src/proxy/pipeline.rs index 6eb7de88..506f3175 100644 --- a/crates/gateway/src/proxy/pipeline.rs +++ b/crates/gateway/src/proxy/pipeline.rs @@ -12,7 +12,7 @@ //! Plus [`LogCtx::new`] (in `log_ctx.rs`) builds the audit context in //! one call instead of 12-field literals at every site. -use super::identity::{budgets_for_ai_gateway, rules_for_ai_gateway}; +use super::identity::limits_for_ai_gateway; use super::{GatewayErrorResponse, GatewayRequestIdentity, GatewayState}; use super::shaper::StreamShaper; @@ -24,19 +24,20 @@ use think_watch_common::lifecycle::stages::{ check_access, check_budget, check_limits, run_post_invoke, }; use think_watch_common::lifecycle::state::{CapturedView, Invoked, LimitCheckRecord, Raw}; -use think_watch_common::limits::{BudgetCap, RateLimitRule}; +use think_watch_common::limits::RequestLimits; use tw_dialect::ir::Dialect; /// Pre-flight result threaded through to `ChatPostInvokeDeps` later /// in the handler. Computed once by [`run_preflight_stages`] so the /// caller doesn't have to re-derive the rule/cap lists from identity. pub(super) struct Preflight { - pub(super) request_rules: Vec, - pub(super) budget_caps: Vec, + pub(super) limits: RequestLimits, } -/// Run the three shared pre-flight stages — `check_limits` → -/// `check_budget` → `check_access` — that each AI surface gates on. +/// 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. /// /// Returns the resolved rule + cap lists so the caller can feed them /// into [`ChatPostInvokeDeps`] without re-walking the identity. @@ -50,8 +51,7 @@ pub(super) async fn run_preflight_stages( trace_id: &str, model: &str, ) -> Result { - let request_rules = rules_for_ai_gateway(identity); - let budget_caps = budgets_for_ai_gateway(identity); + let limits = limits_for_ai_gateway(identity); let fail_closed = state.dynamic_config.rate_limit_fail_closed().await; let raw = Raw::::new( @@ -59,18 +59,19 @@ pub(super) async fn run_preflight_stages( trace_id.to_string(), identity.ip_address.clone(), ); - let limits_checked = check_limits::( + let raw = check_budget::( raw, - &request_rules, + &limits.caps, &state.redis, fail_closed, &state.audit, ) .await .map_err(short_circuit_to_response)?; - let limits_checked = check_budget::( - limits_checked, - &budget_caps, + let limits_checked = check_limits::( + raw, + &limits.rules, + limits.owner, &state.redis, fail_closed, &state.audit, @@ -81,10 +82,7 @@ pub(super) async fn run_preflight_stages( .await .map_err(short_circuit_to_response)?; - Ok(Preflight { - request_rules, - budget_caps, - }) + Ok(Preflight { limits }) } fn short_circuit_to_response(outcome: ChatCompletionOutcome) -> GatewayErrorResponse { diff --git a/crates/gateway/src/quota.rs b/crates/gateway/src/quota.rs index 59f2efdc..31d2ef50 100644 --- a/crates/gateway/src/quota.rs +++ b/crates/gateway/src/quota.rs @@ -37,13 +37,16 @@ impl QuotaManager { 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") + format!("quota:{{{key}}}:limit") } fn usage_key(key: &str) -> String { let month = Self::current_month(); - format!("quota:{key}:used:{month}") + format!("quota:{{{key}}}:used:{month}") } /// Check if user/team has enough quota. Returns remaining tokens. @@ -213,3 +216,22 @@ return limit - used - tokens }) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_quotas_limit_and_usage_share_one_cluster_slot() { + for key in [ + "8a7c0c1e-0000-0000-0000-000000000000:gpt-4o", + "k:{odd}:model", + ] { + assert_eq!( + fred::util::redis_keyslot(QuotaManager::limit_key(key).as_bytes()), + fred::util::redis_keyslot(QuotaManager::usage_key(key).as_bytes()), + "{key}" + ); + } + } +} diff --git a/crates/gateway/src/rate_limiter.rs b/crates/gateway/src/rate_limiter.rs index 521df847..d43d0a80 100644 --- a/crates/gateway/src/rate_limiter.rs +++ b/crates/gateway/src/rate_limiter.rs @@ -95,7 +95,9 @@ impl RateLimiter { let now_ms = chrono::Utc::now().timestamp_millis() as f64; let window_start = now_ms - 60_000.0; let member_id = uuid::Uuid::new_v4().to_string(); - let rpm_key = format!("ratelimit:rpm:{key}"); + // `{}` is a Redis Cluster hash tag: the combined script + // declares both keys, which a cluster runs only in one slot. + let rpm_key = format!("ratelimit:rpm:{{{key}}}"); let reset_at = chrono::Utc::now().timestamp() + 60; // window resets in ~60s // When BOTH RPM and TPM are configured, evaluate them atomically @@ -106,7 +108,7 @@ impl RateLimiter { if let (Some(tpm_limit), Some(tokens)) = (tpm_limit, estimated_tokens) && tokens > 0 { - let tpm_key = format!("ratelimit:tpm:{key}"); + let tpm_key = format!("ratelimit:tpm:{{{key}}}"); let member_with_tokens = format!("{member_id}:{tokens}"); let result: Vec = self .redis diff --git a/crates/mcp-gateway/src/cache.rs b/crates/mcp-gateway/src/cache.rs index 8f7374e1..a34ad8aa 100644 --- a/crates/mcp-gateway/src/cache.rs +++ b/crates/mcp-gateway/src/cache.rs @@ -254,40 +254,23 @@ impl McpResponseCache { user_id: Option, server_id: &Uuid, ) { - let mut cursor: String = "0".to_string(); - let mut deleted: usize = 0; - loop { - let page: Result<(String, Vec), _> = self - .redis - .scan_page( - cursor.clone(), - pattern.clone(), - Some(256), - Some(ScanType::String), - ) - .await; - let (next, keys) = match page { - Ok(p) => p, - Err(e) => { - tracing::warn!( - server = %server_id, user = ?user_id, scope, error = %e, - "MCP cache invalidate: SCAN failed; some stale entries may persist" - ); - return; - } - }; - if !keys.is_empty() { - let n: Result = self.redis.del(keys.clone()).await; - match n { - Ok(n) => deleted += n as usize, - Err(e) => tracing::warn!(error = %e, scope, "MCP cache invalidate: DEL failed"), - } - } - if next == "0" { - break; + // On every node, slot by slot, when Redis is a cluster. + let deleted = match think_watch_common::redis_keys::delete_matching( + &self.redis, + &pattern, + Some(ScanType::String), + ) + .await + { + Ok(n) => n, + Err(e) => { + tracing::warn!( + server = %server_id, user = ?user_id, scope, error = %e, + "MCP cache invalidate incomplete; some stale entries may persist" + ); + e.deleted } - cursor = next; - } + }; if deleted > 0 { tracing::info!( server = %server_id, diff --git a/crates/mcp-gateway/src/lifecycle/mod.rs b/crates/mcp-gateway/src/lifecycle/mod.rs index 4c58c17a..095f343f 100644 --- a/crates/mcp-gateway/src/lifecycle/mod.rs +++ b/crates/mcp-gateway/src/lifecycle/mod.rs @@ -18,7 +18,7 @@ use think_watch_common::lifecycle::Surface; use think_watch_common::lifecycle::state::{CapturedView, Invoked}; use think_watch_common::lifecycle::streaming::StreamOutcome; use think_watch_common::limits::{ - RateLimitRule, RateLimitSubject, Surface as LimitSurface, SurfaceConstraints, + RateMetric, RequestLimits, Surface as LimitSurface, SurfaceConstraints, }; use crate::access_control::is_tool_allowed; @@ -76,10 +76,10 @@ impl Surface for McpSurface { .audit(action) } - fn rate_limited_response(label: &str) -> Self::Response { + fn rate_limited_response(label: &str, _retry_after_secs: u64) -> Self::Response { // Bumps the existing operator-facing metric so dashboards // built around `mcp_rate_limited_total` keep working after - // the migration. + // the migration. A JSON-RPC error has no `Retry-After`. metrics::counter!("mcp_rate_limited_total").increment(1); err_response(None, INVALID_REQUEST, format!("Rate limited: {label}")) } @@ -104,7 +104,7 @@ impl Surface for McpSurface { err_response(None, INVALID_REQUEST, "Access denied for this tool") } - fn budget_exceeded_response(label: &str) -> Self::Response { + fn budget_exceeded_response(label: &str, _retry_after_secs: u64) -> Self::Response { // MCP doesn't currently wire budget caps into // `handle_tools_call` (the AI gateway is the only consumer // of `check_budget` today). The factory exists so the @@ -298,34 +298,18 @@ fn response_for_hooks(invoked: &Invoked, deps: &McpPostInvokeDeps) - } } -/// Build the MCP-surface `requests` rate-limit rules from a -/// materialised [`SurfaceConstraints`]. Mirrors the inline -/// extraction in the previous `handle_tools_call` — same shape, -/// just lifted out so the surface stage gets a clean -/// `&[RateLimitRule]` slice without re-implementing the -/// extraction at every call site. -pub fn rate_limit_rules(constraints: &SurfaceConstraints, user_id: Uuid) -> Vec { - constraints - .block(LimitSurface::McpGateway) - .map(|block| { - block - .rules - .iter() - .filter(|r| r.enabled) - .map(|r| RateLimitRule { - id: Uuid::nil(), - subject_kind: RateLimitSubject::User, - subject_id: user_id, - surface: LimitSurface::McpGateway, - metric: r.metric, - window_secs: r.window_secs, - max_count: r.max_count, - enabled: true, - expires_at: None, - reason: None, - created_by: None, - }) - .collect() - }) - .unwrap_or_default() +/// The MCP-surface rate limits one `tools/call` is held to: the user's +/// (role defaults with the user's overrides) on the user's counters, and +/// the calling key's own on its lineage's — both apply. `requests` rules +/// only: a tool call has no tokens to count, so a `tokens` rule on this +/// surface would only ever read zero. +pub fn rate_limits( + user_id: Uuid, + user: &SurfaceConstraints, + key: Option<(Uuid, &SurfaceConstraints)>, +) -> RequestLimits { + let mut limits = RequestLimits::for_request(LimitSurface::McpGateway, user_id, user, key); + limits.rules.retain(|r| r.metric == RateMetric::Requests); + limits.caps.clear(); + limits } diff --git a/crates/mcp-gateway/src/proxy.rs b/crates/mcp-gateway/src/proxy.rs index 1e0af486..d268bdcc 100644 --- a/crates/mcp-gateway/src/proxy.rs +++ b/crates/mcp-gateway/src/proxy.rs @@ -40,7 +40,12 @@ pub struct RequestContext<'a> { pub user_id: Uuid, pub user_email: &'a str, pub client_session_id: &'a str, + /// The user's limits, counted on the user's counters. pub surface_constraints: &'a SurfaceConstraints, + /// The calling key's lineage and its own limits, counted on the + /// lineage's counters on top of the user's. + pub api_key_lineage_id: Option, + pub key_constraints: &'a SurfaceConstraints, pub allowed_mcp_tools: Option<&'a [String]>, pub trace_id: &'a str, /// Per-server MCP account override JSON from the calling API key @@ -586,7 +591,12 @@ impl McpProxy { // 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 rules = crate::lifecycle::rate_limit_rules(surface_constraints, user_id); + let limits = crate::lifecycle::rate_limits( + user_id, + surface_constraints, + ctx.api_key_lineage_id + .map(|lineage| (lineage, ctx.key_constraints)), + ); let fail_closed = self.dynamic_config.rate_limit_fail_closed().await; let raw = think_watch_common::lifecycle::state::Raw::::new( crate::lifecycle::McpIdentity { @@ -601,7 +611,14 @@ impl McpProxy { ); let limits_checked = match think_watch_common::lifecycle::stages::check_limits::< crate::lifecycle::McpSurface, - >(raw, &rules, &self.redis, fail_closed, &self.audit) + >( + raw, + &limits.rules, + limits.owner, + &self.redis, + fail_closed, + &self.audit, + ) .await { Ok(s) => s, diff --git a/crates/mcp-gateway/src/transport/streamable_http.rs b/crates/mcp-gateway/src/transport/streamable_http.rs index 380a4d36..a0f86554 100644 --- a/crates/mcp-gateway/src/transport/streamable_http.rs +++ b/crates/mcp-gateway/src/transport/streamable_http.rs @@ -38,11 +38,17 @@ pub struct McpRequestIdentity { /// keys (no associated user) — those will be denied any tool /// that requires a role match. pub user_roles: Vec, - /// Effective (most-restrictive across roles) rate-limit rules and - /// budget caps for this user. Materialized by the parent crate - /// once per request so the MCP proxy can apply them without a DB - /// round-trip. + /// Effective (most-restrictive across roles, with the user's own + /// overrides) rate-limit rules and budget caps for this user. + /// Materialized by the parent crate once per request so the MCP + /// proxy can apply them without a DB round-trip. pub surface_constraints: think_watch_common::limits::SurfaceConstraints, + /// The calling key's lineage id. Limits attached to the key live + /// on it so they survive rotation. + pub api_key_lineage_id: Uuid, + /// The calling key's own limits, counted on the lineage's counters + /// on top of the user's. + pub key_constraints: think_watch_common::limits::SurfaceConstraints, /// MCP tool access patterns from role union. `None` = unrestricted. pub allowed_mcp_tools: Option>, /// Per-server account-label override map carried by the calling @@ -134,6 +140,8 @@ pub async fn handle_post( user_email: &identity.user_email, client_session_id: &session_id, surface_constraints: &identity.surface_constraints, + api_key_lineage_id: Some(identity.api_key_lineage_id), + key_constraints: &identity.key_constraints, allowed_mcp_tools: identity.allowed_mcp_tools.as_deref(), trace_id: &trace_id, mcp_account_overrides: &identity.mcp_account_overrides, diff --git a/crates/server/src/handlers/limits.rs b/crates/server/src/handlers/limits.rs index 1d49b64f..7a98f347 100644 --- a/crates/server/src/handlers/limits.rs +++ b/crates/server/src/handlers/limits.rs @@ -675,21 +675,27 @@ pub async fn get_usage( let rate_subject = parse_rate_subject(&kind)?; let storage_id = resolve_subject_id(&state.db, &kind, subject_id).await?; let rules = limits::list_rules(&state.db, rate_subject, storage_id).await?; + // The counters the gateway writes carry the requesting user's hash + // tag: for a key, its owner. A key without one has never been + // let through, so it has counted nothing. + let owner = match rate_subject { + RateLimitSubject::User => Some(subject_id), + RateLimitSubject::ApiKeyLineage => { + sqlx::query_scalar::<_, Option>("SELECT user_id FROM api_keys WHERE id = $1") + .bind(subject_id) + .fetch_optional(&state.db) + .await? + .flatten() + } + }; let mut rule_usage: Vec = Vec::with_capacity(rules.len()); for r in &rules { - let resolved = sliding::ResolvedRule { - id: r.id, - base_key: sliding::build_base_key( - r.surface.as_str(), - r.subject_kind.as_str(), - r.subject_id, - r.metric, - r.window_secs, - ), - bucket_secs: sliding::bucket_secs(r.window_secs), - max_count: r.max_count, + let current = match owner { + Some(owner) => { + sliding::current_count(&state.redis, &sliding::ResolvedRule::new(r, owner)).await + } + None => 0, }; - let current = sliding::current_count(&state.redis, &resolved).await; rule_usage.push(RuleUsage { rule_id: r.id, current, diff --git a/crates/server/src/handlers/user_limits.rs b/crates/server/src/handlers/user_limits.rs index 9d97baa8..fa9e76eb 100644 --- a/crates/server/src/handlers/user_limits.rs +++ b/crates/server/src/handlers/user_limits.rs @@ -217,9 +217,10 @@ async fn build_effective_rules( // shouldn't 500 the whole dashboard, just show 0. let resolved = sliding::ResolvedRule { id: ov.map(|o| o.id).unwrap_or(Uuid::nil()), - base_key: sliding::build_base_key( - surface.as_str(), - "user", + key: sliding::counter_key( + user_id, + surface, + RateLimitSubject::User, user_id, rule.metric, rule.window_secs, @@ -518,27 +519,22 @@ async fn reset_rule_counter( .window_secs .ok_or_else(|| AppError::BadRequest("window_secs is required for rule reset".into()))?; - let base_key = sliding::build_base_key(surface.as_str(), "user", user_id, metric, window_secs); - // Buckets are timestamp-derived (`now_secs / bucket_secs`). DEL the - // 60-window range plus a small safety margin in case a late write - // lands after we read the clock. - let bucket_secs = sliding::bucket_secs(window_secs) as i64; - if bucket_secs <= 0 { - return Ok(0); - } - let now = chrono::Utc::now().timestamp(); - let current_bucket = now / bucket_secs; - let mut deleted = 0usize; - for b in 0..(sliding::BUCKETS_PER_WINDOW + 2) { - let key = format!("{}:{}", base_key, current_bucket - b); - match state.redis.del::(&key).await { - Ok(n) => deleted += n as usize, - Err(e) => { - tracing::warn!("reset_rule_counter DEL failed for {key}: {e}"); - } + // One hash holds every bucket of the window. + let key = sliding::counter_key( + user_id, + surface, + RateLimitSubject::User, + user_id, + metric, + window_secs, + ); + match state.redis.del::(&key).await { + Ok(n) => Ok(n as usize), + Err(e) => { + tracing::warn!("reset_rule_counter DEL failed for {key}: {e}"); + Ok(0) } } - Ok(deleted) } async fn reset_cap_counter( diff --git a/crates/server/src/init.rs b/crates/server/src/init.rs index e5f5f8fd..c9dd5dfa 100644 --- a/crates/server/src/init.rs +++ b/crates/server/src/init.rs @@ -18,6 +18,7 @@ use sqlx::PgPool; use think_watch_auth::oidc::OidcManager; use think_watch_common::audit::{self, AuditConfig, AuditLogger}; use think_watch_common::config::AppConfig; +use think_watch_common::db; use think_watch_common::dynamic_config::{self, DynamicConfig}; use think_watch_common::tasks::supervise; @@ -43,33 +44,8 @@ pub async fn init_state( .await .context("persisted rate-limit / weight rows fail validation")?; - // ClickHouse tables. Same bounded retry as production but without - // the metrics counter (recorder is not installed in tests). - if ch_client.is_some() { - let mut attempt = 0u32; - loop { - match audit::ensure_clickhouse_tables(&ch_client).await { - Ok(()) => break, - Err(e) if attempt < 4 => { - let backoff_ms = 1_500u64 * 2u64.pow(attempt); - tracing::warn!( - attempt = attempt + 1, - backoff_ms, - "ClickHouse table init failed, retrying: {e}" - ); - tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)).await; - attempt += 1; - } - Err(e) => { - tracing::error!( - "ClickHouse table init failed after {} attempts: {e}", - attempt + 1 - ); - break; - } - } - } - } + // ClickHouse tables, when ClickHouse is configured. + ensure_clickhouse_schema(&pool, &ch_client).await?; let dynamic_config = Arc::new(DynamicConfig::load(pool.clone()).await?); @@ -174,6 +150,57 @@ pub async fn init_state( Ok(state) } +/// Create or bring up to date the ClickHouse tables, and backfill the +/// rollups that are still empty. Retried with a backoff; a ClickHouse +/// that stays unreachable is logged, not fatal. +/// +/// One instance at a time, under [`db::CLICKHOUSE_SETUP_LOCK`] in +/// Postgres: a rollup is backfilled when it is found empty, and two +/// instances starting together both found it empty and each copied the +/// whole log into it, counting every request twice. The lock is held for +/// an attempt, not across the backoff, so that instances starting while +/// ClickHouse is down don't wait out each other's retries — longer than +/// the chart's startup probe allows. +pub async fn ensure_clickhouse_schema( + pool: &PgPool, + ch_client: &Option, +) -> anyhow::Result<()> { + if ch_client.is_none() { + return Ok(()); + } + let mut attempt = 0u32; + loop { + let lock = db::StartupLock::acquire( + pool, + db::CLICKHOUSE_SETUP_LOCK, + "setting up the ClickHouse tables", + ) + .await?; + let result = audit::ensure_clickhouse_tables(ch_client).await; + lock.release().await; + match result { + Ok(()) => return Ok(()), + Err(e) if attempt < 4 => { + let backoff_ms = 1_500u64 * 2u64.pow(attempt); + tracing::warn!( + attempt = attempt + 1, + backoff_ms, + "ClickHouse table init failed, retrying: {e}" + ); + tokio::time::sleep(std::time::Duration::from_millis(backoff_ms)).await; + attempt += 1; + } + Err(e) => { + tracing::error!( + "ClickHouse table init failed after {} attempts: {e}", + attempt + 1 + ); + return Ok(()); + } + } + } +} + fn audit_config(config: &AppConfig) -> AuditConfig { config.audit_config() } @@ -221,14 +248,14 @@ pub fn install_cb_listener(state: &AppState) { /// pool whenever any instance flips a setting. pub async fn spawn_config_subscriber(state: &AppState) -> anyhow::Result<()> { // Multi-instance config sync (`system_settings.value` updates → Pub/Sub). - let sub_main = fred::types::config::Config::from_url(&state.config.redis_url)?; + let sub_main = state.config.redis_config()?; let sub_main_redis: fred::clients::SubscriberClient = Builder::from_config(sub_main).build_subscriber_client()?; sub_main_redis.init().await?; dynamic_config::spawn_config_subscriber(sub_main_redis, state.dynamic_config.clone()); // Hot-reload the per-state arc-swap handles on the same channel. - let sub_filters_cfg = fred::types::config::Config::from_url(&state.config.redis_url)?; + let sub_filters_cfg = state.config.redis_config()?; let sub_filters: fred::clients::SubscriberClient = Builder::from_config(sub_filters_cfg).build_subscriber_client()?; sub_filters.init().await?; @@ -342,7 +369,7 @@ pub async fn spawn_config_subscriber(state: &AppState) -> anyhow::Result<()> { // replica's subscriber rebuilds its local `ArcSwap`. // Without this, multi-instance deployments had per-replica stale // routers between CRUD time and the next process restart. - let sub_router_cfg = fred::types::config::Config::from_url(&state.config.redis_url)?; + let sub_router_cfg = state.config.redis_config()?; let sub_router: fred::clients::SubscriberClient = Builder::from_config(sub_router_cfg).build_subscriber_client()?; sub_router.init().await?; diff --git a/crates/server/src/main.rs b/crates/server/src/main.rs index 82b88575..864b2658 100644 --- a/crates/server/src/main.rs +++ b/crates/server/src/main.rs @@ -50,14 +50,19 @@ async fn main() -> anyhow::Result<()> { std::process::exit(1); } - let redis_config = Config::from_url(&config.redis_url).map_err(|e| { - tracing::error!("Invalid REDIS_URL: {e}"); - anyhow::anyhow!("Invalid REDIS_URL") - })?; + let redis_config = match config.redis_config() { + Ok(c) => c, + Err(e) => { + tracing::error!("Invalid Redis configuration: {e:#}"); + std::process::exit(1); + } + }; let redis = Builder::from_config(redis_config).build()?; if let Err(e) = redis.init().await { tracing::error!("Failed to connect to Redis: {e}"); - tracing::error!("Check REDIS_URL and ensure Redis is running"); + tracing::error!( + "Check REDIS_URL (and REDIS_CA_CERT for a rediss:// URL) and ensure Redis is running" + ); std::process::exit(1); } tracing::info!("Redis connected"); diff --git a/crates/server/src/middleware/api_key_auth.rs b/crates/server/src/middleware/api_key_auth.rs index 55eda7c0..1f13aad4 100644 --- a/crates/server/src/middleware/api_key_auth.rs +++ b/crates/server/src/middleware/api_key_auth.rs @@ -292,24 +292,24 @@ 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) = if let Some(uid) = row.user_id { + let (role_limits, user_roles, surface_constraints, key_constraints) = if let Some(uid) = row.user_id { let limits = rbac::compute_user_resource_limits(&state.db, uid) .await .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?; let names = rbac::load_user_role_names(&state.db, uid) .await .unwrap_or_default(); - // Use the api_key-aware variant so per-key - // `rate_limit_rules` / `budget_caps` rows fire on the - // gateway hot path. Falling back to the user-only - // function would silently drop api_key-scope - // overrides — the schema supports them but the - // gateway would never see them. - let constraints = - rbac::compute_effective_surface_constraints(&state.db, uid, row.id) + // 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) + (limits, names, constraints, key_constraints) } else { // A key without an owner (its user row was removed // and `user_id` set NULL) has no roles to grant @@ -320,6 +320,7 @@ pub fn require_api_key( rbac::UserResourceLimits::none(), Vec::new(), think_watch_common::limits::SurfaceConstraints::default(), + think_watch_common::limits::SurfaceConstraints::default(), ) }; @@ -382,6 +383,7 @@ pub fn require_api_key( api_key_lineage_id: Some(row.lineage_id.to_string()), allowed_models: merged_models.clone(), surface_constraints: surface_constraints.clone(), + key_constraints: key_constraints.clone(), ip_address: client_ip.clone(), }; @@ -414,6 +416,8 @@ pub fn require_api_key( user_email, user_roles, surface_constraints: surface_constraints.clone(), + api_key_lineage_id: row.lineage_id, + key_constraints, allowed_mcp_tools: merged_mcp_tools.clone(), mcp_account_overrides: row.mcp_account_overrides.clone(), ip_address: client_ip.clone(), diff --git a/crates/server/src/services/mcp_store_repository.rs b/crates/server/src/services/mcp_store_repository.rs index f445d384..9afeae54 100644 --- a/crates/server/src/services/mcp_store_repository.rs +++ b/crates/server/src/services/mcp_store_repository.rs @@ -21,6 +21,9 @@ use uuid::Uuid; /// Reserved advisory lock keys (keep this list current): /// * `MCP_STORE_INSTALL_LOCK_KEY` (here): template-install /// serialization in `create_server` when `template_slug` is set. +/// * `think_watch_common::db::SCHEMA_LOCK` ("twschema") and +/// `CLICKHOUSE_SETUP_LOCK` ("twchinit"): one instance at a time +/// through the start-up schema setup. const MCP_STORE_INSTALL_LOCK_KEY: i64 = 0x6D637053746F7265; /// A category and how many templates are in it. diff --git a/crates/test-support/src/lib.rs b/crates/test-support/src/lib.rs index e7eae132..fc0791f9 100644 --- a/crates/test-support/src/lib.rs +++ b/crates/test-support/src/lib.rs @@ -8,6 +8,10 @@ //! Required env (defaults match `make infra`): //! - `TEST_DATABASE_BASE_URL` (default `postgres://thinkwatch:thinkwatch@localhost:5432`) //! - `TEST_REDIS_URL` (default `redis://localhost:6379`) +//! - `TEST_REDIS_CA_CERT` (optional): the PEM CA file for a `rediss://` +//! Redis whose certificate a private CA signed — `REDIS_CA_CERT` for +//! every Redis the tests reach (`TEST_REDIS_URL`, the cluster and TLS +//! test URLs) //! //! Tests run against the same Redis instance. Each `TestApp` FLUSHDBs //! its logical DB on spawn, so concurrent tests need separate DBs: @@ -20,6 +24,7 @@ pub mod client; pub mod fixtures; pub mod mock_provider; pub mod pg; +pub mod redis_scripts; use std::sync::Arc; @@ -27,7 +32,6 @@ use anyhow::Context; use fred::clients::Client as RedisClient; use fred::interfaces::ClientLike; use fred::types::Builder; -use fred::types::config::Config as RedisConfig; use sqlx::PgPool; use think_watch_common::config::AppConfig; use think_watch_server::{app, init}; @@ -85,6 +89,10 @@ pub struct SpawnOptions { /// to exercise the offload path inject an in-memory store here /// without standing up a real S3 backend. pub blob_store: Option>, + /// Boot against this Redis instead of `TEST_REDIS_URL`'s per-slot + /// logical DB — e.g. a `redis-cluster://` URL. Nothing is flushed: + /// the test must use keys no other test touches. + pub redis_url: Option, } /// An SSRF guard that lets a `wiremock` on `127.0.0.1` through and still @@ -143,15 +151,18 @@ impl TestApp { pub async fn try_spawn_with(opts: SpawnOptions) -> anyhow::Result { init_test_tracing(); - let base_url = std::env::var("TEST_DATABASE_BASE_URL").unwrap_or_else(|_| { - "postgres://thinkwatch:7c3fe6307d00fe3f2f29f534e806ac71@localhost:5432".into() - }); + let base_url = database_base_url(); // Default to logical DB 1 so we never trample the dev Redis // (DB 0). Override via env when CI uses a dedicated instance. let redis_url = std::env::var("TEST_REDIS_URL").unwrap_or_else(|_| { "redis://:225b3facaf55212ff86ad6595e6d6471@localhost:6379/1".into() }); - let redis_url = redis_url_for_slot(&redis_url)?; + let shared_redis = opts.redis_url.is_some(); + let redis_ca_cert = test_redis_ca_cert(); + let redis_url = match &opts.redis_url { + Some(url) => url.clone(), + None => redis_url_for_slot(&redis_url)?, + }; // Per-test database with migrations applied. let db_owner = IsolatedDatabase::create(&base_url) @@ -166,11 +177,11 @@ impl TestApp { // Makefile target) or under nextest, where // `redis_url_for_slot` gives each concurrently running test // its own DB. - let redis = build_redis(&redis_url).await?; + let redis = build_redis(&redis_url, redis_ca_cert.as_deref()).await?; // fred 10 doesn't expose FLUSHDB directly (only FLUSHALL), // and we don't want to nuke the dev DB. Send the raw // command so we only clear the test logical DB. - { + if !shared_redis { use fred::interfaces::ClientLike; use fred::types::{ClusterHash, CustomCommand}; let cmd = CustomCommand::new("FLUSHDB", ClusterHash::FirstKey, false); @@ -182,14 +193,7 @@ impl TestApp { // Per-test ClickHouse — only when the test asked for it. let (ch_owner, ch_client, ch_url, ch_db, ch_user, ch_password) = if opts.clickhouse { - let url = std::env::var("TEST_CLICKHOUSE_URL") - .unwrap_or_else(|_| "http://localhost:8123".into()); - let user = std::env::var("TEST_CLICKHOUSE_USER") - .ok() - .or_else(|| Some("thinkwatch".into())); - let password = std::env::var("TEST_CLICKHOUSE_PASSWORD") - .ok() - .or_else(|| Some("c693ded3da8388c7b6a4288dac91a2ad".into())); + let (url, user, password) = clickhouse_env(); let owner = IsolatedClickHouseDatabase::create(&url, user.as_deref(), password.as_deref()) .await @@ -211,6 +215,7 @@ impl TestApp { let config = AppConfig { database_url: db_owner.url().to_string(), redis_url, + redis_ca_cert, jwt_secret: "test-jwt-secret-with-enough-entropy-aaa".into(), // 64 hex chars = 32 bytes; valid for AES-256-GCM. encryption_key: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef" @@ -433,8 +438,62 @@ fn redis_url_for_slot(redis_url: &str) -> anyhow::Result { Ok(url.to_string()) } -async fn build_redis(redis_url: &str) -> anyhow::Result { - let cfg = RedisConfig::from_url(redis_url).context("parse REDIS_URL")?; +/// `TEST_REDIS_CA_CERT`, the CA file the tests trust for a `rediss://` +/// Redis. +pub fn test_redis_ca_cert() -> Option { + std::env::var_os("TEST_REDIS_CA_CERT") + .filter(|s| !s.is_empty()) + .map(Into::into) +} + +/// The fred config for a test Redis URL, built the way the server builds +/// its own (TLS included), with `TEST_REDIS_CA_CERT`'s roots. +pub fn test_redis_config(redis_url: &str) -> fred::types::config::Config { + think_watch_common::redis_config::client_config(redis_url, test_redis_ca_cert().as_deref()) + .expect("a usable test Redis URL") +} + +/// Subscribe a fresh subscriber built from `config` to `config:changed` +/// and wait until a notice sent through `publisher` reaches it — what a +/// second instance's `init::spawn_config_subscriber` does. +/// +/// The notice is sent again until it arrives: fred's `subscribe` can +/// return while the subscription is still on its way to Redis (a +/// `PUBLISH` right after it has been seen to reach 0 receivers), and a +/// Pub/Sub message nobody is subscribed to yet is lost. +pub async fn assert_config_notice_arrives( + config: fred::types::config::Config, + publisher: &RedisClient, +) { + use fred::interfaces::{EventInterface, PubsubInterface}; + use std::time::{Duration, Instant}; + + let subscriber = Builder::from_config(config) + .build_subscriber_client() + .expect("build_subscriber_client"); + subscriber.init().await.expect("subscriber init"); + let mut rx = subscriber.message_rx(); + subscriber + .subscribe("config:changed") + .await + .expect("subscribe"); + let deadline = Instant::now() + Duration::from_secs(5); + loop { + think_watch_common::dynamic_config::notify_config_changed(publisher).await; + if let Ok(msg) = tokio::time::timeout(Duration::from_millis(200), rx.recv()).await { + assert_eq!(msg.expect("message channel").channel, "config:changed"); + return; + } + assert!(Instant::now() < deadline, "no config notice within 5 s"); + } +} + +async fn build_redis( + redis_url: &str, + ca_cert: Option<&std::path::Path>, +) -> anyhow::Result { + let cfg = think_watch_common::redis_config::client_config(redis_url, ca_cert) + .context("parse REDIS_URL")?; let client = Builder::from_config(cfg).build()?; client.init().await.context("redis init")?; Ok(client) @@ -455,6 +514,28 @@ fn init_test_tracing() { }); } +/// The Postgres server tests create their databases on +/// (`TEST_DATABASE_BASE_URL`), without a database name. +pub fn database_base_url() -> String { + std::env::var("TEST_DATABASE_BASE_URL").unwrap_or_else(|_| { + "postgres://thinkwatch:7c3fe6307d00fe3f2f29f534e806ac71@localhost:5432".into() + }) +} + +/// The ClickHouse server tests create their databases on: URL, user and +/// password (`TEST_CLICKHOUSE_URL`, `_USER`, `_PASSWORD`). +pub fn clickhouse_env() -> (String, Option, Option) { + let url = + std::env::var("TEST_CLICKHOUSE_URL").unwrap_or_else(|_| "http://localhost:8123".into()); + let user = std::env::var("TEST_CLICKHOUSE_USER") + .ok() + .or_else(|| Some("thinkwatch".into())); + let password = std::env::var("TEST_CLICKHOUSE_PASSWORD") + .ok() + .or_else(|| Some("c693ded3da8388c7b6a4288dac91a2ad".into())); + (url, user, password) +} + /// Convenience re-exports so test files only need one `use`. pub mod prelude { pub use crate::SpawnOptions; diff --git a/crates/test-support/src/pg.rs b/crates/test-support/src/pg.rs index ab77687d..77d37707 100644 --- a/crates/test-support/src/pg.rs +++ b/crates/test-support/src/pg.rs @@ -24,6 +24,17 @@ impl IsolatedDatabase { /// including) the database name — e.g. /// `postgres://user:pwd@localhost:5432`. pub async fn create(base_url: &str) -> anyhow::Result { + let db = Self::create_empty(base_url).await?; + // Use the workspace migrator from common so the schema stays + // in lockstep with production. + think_watch_common::db::run_migrations(db.pool()) + .await + .context("run migrations into per-test DB")?; + Ok(db) + } + + /// A fresh database with nothing in it: what a first start finds. + pub async fn create_empty(base_url: &str) -> anyhow::Result { let admin_url = format!("{}/postgres", base_url.trim_end_matches('/')); let name = format!("test_tw_{}", Uuid::new_v4().simple()); @@ -45,12 +56,6 @@ impl IsolatedDatabase { .await .context("connect to per-test DB")?; - // Use the workspace migrator from common so the schema stays - // in lockstep with production. - think_watch_common::db::run_migrations(&pool) - .await - .context("run migrations into per-test DB")?; - Ok(Self { base_url: base_url.trim_end_matches('/').to_string(), name, diff --git a/crates/test-support/src/redis_scripts.rs b/crates/test-support/src/redis_scripts.rs new file mode 100644 index 00000000..5109a280 --- /dev/null +++ b/crates/test-support/src/redis_scripts.rs @@ -0,0 +1,79 @@ +//! 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`. + +use uuid::Uuid; + +fn rule( + kind: think_watch_common::limits::RateLimitSubject, + subject: Uuid, + metric: think_watch_common::limits::RateMetric, + window_secs: i32, + max_count: i64, +) -> think_watch_common::limits::RateLimitRule { + think_watch_common::limits::RateLimitRule { + id: Uuid::nil(), + subject_kind: kind, + subject_id: subject, + surface: think_watch_common::limits::Surface::AiGateway, + metric, + window_secs, + max_count, + enabled: true, + expires_at: None, + reason: None, + created_by: None, + } +} + +/// The admit/record sequence both the single-node and the cluster test +/// run: windows that fill, free bucket by bucket, refuse without +/// charging, and record past a limit. +pub async fn exercise_the_limit_scripts(redis: &fred::clients::Client) { + use fred::interfaces::HashesInterface; + use think_watch_common::limits::{RateLimitSubject as S, RateMetric as M, sliding}; + + let owner = Uuid::new_v4(); + let lineage = Uuid::new_v4(); + let rules = [ + rule(S::User, owner, M::Requests, 60, 3), + rule(S::ApiKeyLineage, lineage, M::Tokens, 300, 100), + ]; + // A whole minute, so every bucket boundary below is exact. + let t0: i64 = 1_800_000_000_000; + let admit = |ms: i64| sliding::admit_at(redis, &rules, owner, t0 + ms, false); + + for (i, at) in [0, 10_000, 20_000].into_iter().enumerate() { + let o = admit(at).await.unwrap(); + assert!(o.allowed, "request {i}"); + assert_eq!(o.currents, vec![i as i64 + 1, 0]); + } + // Full. The first request's bucket leaves the window at t0 + 60 s. + let o = admit(30_000).await.unwrap(); + assert!(!o.allowed); + assert_eq!(o.exceeded_index, Some(0)); + assert_eq!(o.retry_after_secs, 30); + assert_eq!(o.currents, vec![3, 0], "a refusal charges nothing"); + + // Tokens are recorded in full, past the limit. + sliding::record_at(redis, &rules, owner, M::Tokens, 150, t0 + 30_000) + .await + .unwrap(); + // The requests window has room again; the tokens one does not, and + // frees when the 5-second bucket of t0 + 30 s leaves it, 300 s on. + let o = admit(61_000).await.unwrap(); + assert!(!o.allowed); + assert_eq!(o.exceeded_index, Some(1)); + assert_eq!(o.retry_after_secs, 330 - 61); + assert_eq!(o.currents, vec![2, 150]); + + // Both windows have moved past everything: one request counted, and + // the old buckets are gone from the hash. + let o = admit(400_000).await.unwrap(); + assert!(o.allowed); + assert_eq!(o.currents, vec![1, 0]); + let requests = sliding::ResolvedRule::new(&rules[0], owner); + let tokens = sliding::ResolvedRule::new(&rules[1], owner); + assert_eq!(redis.hlen::(&requests.key).await.unwrap(), 1); + assert_eq!(redis.hlen::(&tokens.key).await.unwrap(), 0); +} diff --git a/crates/test-support/tests/clickhouse_retention_restart.rs b/crates/test-support/tests/clickhouse_retention_restart.rs new file mode 100644 index 00000000..bbef5c5a --- /dev/null +++ b/crates/test-support/tests/clickhouse_retention_restart.rs @@ -0,0 +1,166 @@ +//! A restart never takes a ClickHouse retention below the configured one. +//! +//! The server sets up the ClickHouse tables on every start, then applies +//! the configured retentions: `data.retention_days_*` as each log table's +//! TTL, `audit.body_retention_days` as the TTL of the columns holding +//! request and response bodies. The table setup used to give the body +//! columns a 30-day TTL each time, and ClickHouse acts on a TTL as soon +//! as it is set: with a longer body retention configured, every restart +//! cleared the bodies older than 30 days, although the configured TTL was +//! back a moment later. + +use think_watch_server::handlers::admin::reconcile_clickhouse_ttls; +use think_watch_server::init::ensure_clickhouse_schema; +use think_watch_test_support::prelude::*; + +const BODY_COLUMNS: [(&str, &str); 4] = [ + ("gateway_logs", "request_body"), + ("gateway_logs", "response_body"), + ("mcp_logs", "tool_arguments"), + ("mcp_logs", "tool_result"), +]; + +/// The days in a `toIntervalDay(N)` on `line`. +fn interval_days(line: &str) -> Option { + let rest = &line[line.find("toIntervalDay(")? + "toIntervalDay(".len()..]; + rest[..rest.find(')')?].parse().ok() +} + +/// Every TTL that matters here, in days: each log table's and each body +/// column's, as `SHOW CREATE TABLE` gives them. +async fn ttls(ch: &clickhouse::Client) -> Vec<(String, Option)> { + let mut out = Vec::new(); + for table in ["gateway_logs", "mcp_logs"] { + let create: String = ch + .query(&format!("SHOW CREATE TABLE {table}")) + .fetch_one() + .await + .unwrap(); + let table_ttl = create + .lines() + .find(|l| l.starts_with("TTL ")) + .and_then(interval_days); + out.push((table.to_string(), table_ttl)); + for (_, column) in BODY_COLUMNS.iter().filter(|(t, _)| *t == table) { + let line = create + .lines() + .find(|l| l.trim_start().starts_with(&format!("`{column}` "))) + .unwrap_or_else(|| panic!("no {table}.{column} in {create}")); + out.push((format!("{table}.{column}"), interval_days(line))); + } + } + out +} + +fn configured() -> Vec<(String, Option)> { + vec![ + ("gateway_logs".into(), Some(365)), + ("gateway_logs.request_body".into(), Some(90)), + ("gateway_logs.response_body".into(), Some(90)), + ("mcp_logs".into(), Some(365)), + ("mcp_logs.tool_arguments".into(), Some(90)), + ("mcp_logs.tool_result".into(), Some(90)), + ] +} + +/// Wait until ClickHouse has carried out every mutation of this database, +/// the TTL it materializes after an `ALTER … TTL` among them. +async fn mutations_done(ch: &clickhouse::Client) { + tokio::time::timeout(std::time::Duration::from_secs(60), async { + loop { + let pending: u64 = ch + .query( + "SELECT count() FROM system.mutations \ + WHERE database = currentDatabase() AND NOT is_done", + ) + .fetch_one() + .await + .unwrap(); + if pending == 0 { + return; + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + }) + .await + .expect("ClickHouse finished its mutations"); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_restart_keeps_retentions_longer_than_the_defaults() { + let app = TestApp::spawn_with_clickhouse().await; + let ch = app.state.clickhouse.clone().expect("ClickHouse"); + + // A deployment that keeps bodies 90 days and logs a year, running. + for (key, days) in [ + ("audit.body_retention_days", 90), + ("data.retention_days_gateway", 365), + ("data.retention_days_mcp", 365), + ] { + fixtures::set_setting(&app.db, key, json!(days)) + .await + .unwrap(); + } + app.state.dynamic_config.reload().await.unwrap(); + reconcile_clickhouse_ttls(&app.state).await; + mutations_done(&ch).await; + assert_eq!(ttls(&ch).await, configured(), "as the deployment runs"); + + // Requests from 60 days ago, their bodies within the retention, and + // from 200 days ago, within the logs'. + ch.query( + "INSERT INTO gateway_logs (id, model_id, request_body, response_body, created_at) VALUES \ + ('60d', 'm', 'request', 'response', now64(3) - INTERVAL 60 DAY)", + ) + .execute() + .await + .unwrap(); + ch.query( + "INSERT INTO gateway_logs (id, model_id, created_at) VALUES \ + ('200d', 'm', now64(3) - INTERVAL 200 DAY)", + ) + .execute() + .await + .unwrap(); + ch.query( + "INSERT INTO mcp_logs (id, tool_arguments, tool_result, created_at) VALUES \ + ('60d', 'arguments', 'result', now64(3) - INTERVAL 60 DAY)", + ) + .execute() + .await + .unwrap(); + + // The restart, step by step as the server takes it: the tables, then + // the configured retentions. No step may shorten one. + ensure_clickhouse_schema(&app.db, &app.state.clickhouse) + .await + .unwrap(); + assert_eq!(ttls(&ch).await, configured(), "after the table setup"); + reconcile_clickhouse_ttls(&app.state).await; + assert_eq!(ttls(&ch).await, configured(), "after the retentions"); + + mutations_done(&ch).await; + let gateway: Vec = ch + .query( + "SELECT concat(id, ' ', ifNull(request_body, '-'), ' ', ifNull(response_body, '-')) \ + FROM gateway_logs ORDER BY id", + ) + .fetch_all() + .await + .unwrap(); + assert_eq!( + gateway, + ["200d - -", "60d request response"], + "the 60-day-old bodies and the 200-day-old row are still there" + ); + let mcp: Vec = ch + .query( + "SELECT concat(id, ' ', ifNull(tool_arguments, '-'), ' ', ifNull(tool_result, '-')) \ + FROM mcp_logs", + ) + .fetch_all() + .await + .unwrap(); + assert_eq!(mcp, ["60d arguments result"]); +} diff --git a/crates/test-support/tests/concurrent_startup.rs b/crates/test-support/tests/concurrent_startup.rs new file mode 100644 index 00000000..ca8d6057 --- /dev/null +++ b/crates/test-support/tests/concurrent_startup.rs @@ -0,0 +1,339 @@ +//! Several instances starting at the same moment against one database: +//! a Helm `replicaCount` above 1, a rolling upgrade, an autoscaler adding +//! pods. Every instance applies `db/schema.sql`, the seeds and the +//! guard-settings conversion, then sets up the ClickHouse tables and +//! backfills the rollups. Run side by side, the schema statements of one +//! instance deadlocked with another's, and one of them exited at its +//! migration; two backfills of the same empty rollup each copied the +//! whole log into it. +//! +//! Each round here starts the setup on [`INSTANCES`] tasks at once, each +//! with a connection pool of its own as separate processes would have: +//! every one succeeds, and the database ends up as one instance alone +//! leaves it. + +use std::sync::Arc; +use std::time::Duration; + +use serde_json::Value; +use sqlx::PgPool; +use sqlx::postgres::PgPoolOptions; +use think_watch_common::db::{SCHEMA_LOCK, StartupLock, run_migrations}; +use think_watch_test_support::{ + IsolatedClickHouseDatabase, IsolatedDatabase, clickhouse_env, database_base_url, +}; +use tokio::sync::Barrier; + +const INSTANCES: usize = 4; + +/// One pool per instance, against the same database. +async fn instance_pools(db: &IsolatedDatabase) -> Vec { + let mut pools = Vec::with_capacity(INSTANCES); + for _ in 0..INSTANCES { + pools.push( + PgPoolOptions::new() + .max_connections(4) + .connect(db.url()) + .await + .expect("connect an instance's pool"), + ); + } + pools +} + +/// Start the schema setup on every instance at the same moment and +/// return what each one returned. +async fn migrate_at_once(db: &IsolatedDatabase) -> Vec> { + let barrier = Arc::new(Barrier::new(INSTANCES)); + let mut tasks = Vec::with_capacity(INSTANCES); + for pool in instance_pools(db).await { + let barrier = barrier.clone(); + tasks.push(tokio::spawn(async move { + barrier.wait().await; + let result = run_migrations(&pool).await.map_err(|e| format!("{e:#}")); + pool.close().await; + result + })); + } + let mut results = Vec::with_capacity(INSTANCES); + for task in tasks { + results.push(task.await.expect("an instance's setup panicked")); + } + results +} + +fn assert_all_started(results: &[Result<(), String>]) { + let failed: Vec<&String> = results.iter().filter_map(|r| r.as_ref().err()).collect(); + assert!( + failed.is_empty(), + "{} of {INSTANCES} instances failed their migration: {failed:#?}", + failed.len() + ); +} + +/// Every object the schema setup creates, one line each, sorted. +async fn schema_of(pool: &PgPool) -> Vec { + sqlx::query_scalar( + "SELECT 'column ' || table_name || '.' || column_name || ' ' || data_type + || ' nullable=' || is_nullable || ' default=' || coalesce(column_default, '') + FROM information_schema.columns WHERE table_schema = 'public' + UNION ALL + SELECT 'index ' || indexdef FROM pg_indexes WHERE schemaname = 'public' + UNION ALL + SELECT 'constraint ' || conrelid::regclass::text || ' ' || conname || ' ' + || pg_get_constraintdef(oid) + FROM pg_constraint WHERE connamespace = 'public'::regnamespace + UNION ALL + SELECT 'trigger ' || pg_get_triggerdef(oid) FROM pg_trigger WHERE NOT tgisinternal + UNION ALL + SELECT 'function ' || oid::regprocedure::text + FROM pg_proc WHERE pronamespace = 'public'::regnamespace + UNION ALL + SELECT 'extension ' || extname FROM pg_extension + ORDER BY 1", + ) + .fetch_all(pool) + .await + .expect("read the schema") +} + +/// Every row the seeds write, one line each, sorted. The conversion's +/// record says when it was written, so only its presence is compared. +async fn seeds_of(pool: &PgPool) -> Vec { + sqlx::query_scalar( + "SELECT 'role ' || name || ' ' || policy_document::text FROM rbac_roles + UNION ALL + SELECT 'surface ' || name FROM api_key_surface_kinds + UNION ALL + SELECT 'pricing ' || id::text FROM platform_pricing + UNION ALL + SELECT 'setting ' || key || ' ' + || CASE WHEN key = 'security.legacy_converted' THEN 'present' ELSE value::text END + FROM system_settings + UNION ALL + SELECT 'template ' || slug FROM mcp_store_templates + ORDER BY 1", + ) + .fetch_all(pool) + .await + .expect("read the seeded rows") +} + +async fn setting(pool: &PgPool, key: &str) -> Option { + sqlx::query_scalar("SELECT value FROM system_settings WHERE key = $1") + .bind(key) + .fetch_optional(pool) + .await + .unwrap() +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn replicas_starting_together_on_a_fresh_database_all_start() { + let base = database_base_url(); + let alone = IsolatedDatabase::create(&base).await.unwrap(); + let db = IsolatedDatabase::create_empty(&base).await.unwrap(); + + // The first start of a deployment: an empty database. + assert_all_started(&migrate_at_once(&db).await); + assert_eq!(schema_of(db.pool()).await, schema_of(alone.pool()).await); + assert_eq!(seeds_of(db.pool()).await, seeds_of(alone.pool()).await); + + // All of them restarting at once (a rollout of the same version): + // everything is there already, and nothing changes. + assert_all_started(&migrate_at_once(&db).await); + assert_eq!(schema_of(db.pool()).await, schema_of(alone.pool()).await); + assert_eq!(seeds_of(db.pool()).await, seeds_of(alone.pool()).await); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn replicas_upgrading_together_apply_the_upgrade_once() { + let base = database_base_url(); + let alone = IsolatedDatabase::create(&base).await.unwrap(); + let db = IsolatedDatabase::create(&base).await.unwrap(); + + // The database as an earlier version leaves it: columns and an index + // the schema has added since are missing, and the guard settings are + // still the old ones, with nothing converted yet. + sqlx::raw_sql( + r#"ALTER TABLE models DROP COLUMN max_output_tokens; + ALTER TABLE model_routes DROP COLUMN upstream_protocol; + DROP INDEX idx_api_keys_user_not_deleted; + DELETE FROM system_settings WHERE key = 'security.legacy_converted'; + INSERT INTO system_settings (key, value, category, description) VALUES + ('security.hidden_text', '"block"', 'security', 'old'), + ('security.tool_inspection', + '{"mode": "enforce", "disabled": ["chmod-777"], "actions": {}, "custom": []}', + 'security', 'old'); + ALTER TABLE models ADD COLUMN output_guardrails JSONB NOT NULL DEFAULT '[]'::jsonb;"#, + ) + .execute(db.pool()) + .await + .unwrap(); + + assert_all_started(&migrate_at_once(&db).await); + + // The schema is complete, the old column is gone... + assert_eq!(schema_of(db.pool()).await, schema_of(alone.pool()).await); + // ...and the old settings were converted, by one of them. + for key in ["security.hidden_text", "security.tool_inspection"] { + assert_eq!(setting(db.pool(), key).await, None, "{key}"); + } + let marker = setting(db.pool(), "security.legacy_converted") + .await + .expect("the conversion was recorded"); + let mut converted: Vec<&str> = marker["converted"] + .as_array() + .expect("{marker}") + .iter() + .filter_map(Value::as_str) + .collect(); + converted.sort_unstable(); + assert_eq!( + converted, + ["security.hidden_text", "security.tool_inspection"], + "{marker}" + ); + assert_eq!( + setting(db.pool(), "security.inspect_tools").await, + Some(serde_json::json!({"mode": "enforce", "disable": ["chmod-777"]})), + ); +} + +/// Advisory locks held in `pool`'s database: `(granted, waiting)`. +async fn advisory_locks(pool: &PgPool) -> (i64, i64) { + sqlx::query_as( + "SELECT count(*) FILTER (WHERE granted), count(*) FILTER (WHERE NOT granted) + FROM pg_locks + WHERE locktype = 'advisory' + AND database = (SELECT oid FROM pg_database WHERE datname = current_database())", + ) + .fetch_one(pool) + .await + .unwrap() +} + +/// Wait until `pool`'s database shows `want` advisory locks. +async fn until_advisory_locks(pool: &PgPool, want: (i64, i64), what: &str) { + tokio::time::timeout(Duration::from_secs(10), async { + while advisory_locks(pool).await != want { + tokio::time::sleep(Duration::from_millis(20)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("{what}: the advisory locks never came to {want:?}")); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn the_schema_lock_never_outlives_its_holder() { + let db = IsolatedDatabase::create(&database_base_url()) + .await + .unwrap(); + + // An instance that dies holding the lock: its session ends, and the + // instance waiting for it goes ahead. + let mut holder = StartupLock::acquire(db.pool(), SCHEMA_LOCK, "testing") + .await + .unwrap(); + let pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(holder.conn()) + .await + .unwrap(); + let pool = db.pool().clone(); + let waiter = tokio::spawn(async move { run_migrations(&pool).await }); + until_advisory_locks(db.pool(), (1, 1), "the second instance waits").await; + sqlx::query("SELECT pg_terminate_backend($1)") + .bind(pid) + .execute(db.pool()) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(30), waiter) + .await + .expect("the waiting instance got the lock once its holder was gone") + .unwrap() + .unwrap(); + drop(holder); + until_advisory_locks(db.pool(), (0, 0), "released after the migration").await; + + // A setup that fails drops the lock without releasing it: its + // connection is closed, not handed back to the pool still holding it. + let lock = StartupLock::acquire(db.pool(), SCHEMA_LOCK, "testing") + .await + .unwrap(); + assert_eq!(advisory_locks(db.pool()).await, (1, 0)); + drop(lock); + until_advisory_locks(db.pool(), (0, 0), "released on drop").await; +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn replicas_starting_together_backfill_the_rollups_once() { + let pg = IsolatedDatabase::create(&database_base_url()) + .await + .unwrap(); + let (url, user, password) = clickhouse_env(); + let ch = IsolatedClickHouseDatabase::create(&url, user.as_deref(), password.as_deref()) + .await + .unwrap(); + let client = ch.client().clone(); + + // Requests logged before the rollups existed: the rows are in + // gateway_logs, and the rollups a new version adds start empty. + client + .query( + "INSERT INTO gateway_logs (id, model_id, provider, status_code, latency_ms, \ + input_tokens, output_tokens) \ + SELECT toString(number), 'm', 'p', 200, 10, 3, 4 FROM numbers(25)", + ) + .execute() + .await + .unwrap(); + for table in ["cost_rollup_hourly", "provider_health_5m"] { + client + .query(&format!("TRUNCATE TABLE {table}")) + .execute() + .await + .unwrap(); + } + + let barrier = Arc::new(Barrier::new(INSTANCES)); + let mut tasks = Vec::with_capacity(INSTANCES); + for pool in instance_pools(&pg).await { + let barrier = barrier.clone(); + let ch = Some(client.clone()); + tasks.push(tokio::spawn(async move { + barrier.wait().await; + let result = think_watch_server::init::ensure_clickhouse_schema(&pool, &ch) + .await + .map_err(|e| format!("{e:#}")); + pool.close().await; + result + })); + } + let mut results = Vec::with_capacity(INSTANCES); + for task in tasks { + results.push(task.await.expect("an instance's setup panicked")); + } + assert_all_started(&results); + + let requests: u64 = client + .query("SELECT toUInt64(sum(request_count)) FROM cost_rollup_hourly") + .fetch_one() + .await + .unwrap(); + assert_eq!(requests, 25, "cost_rollup_hourly counts each request once"); + let tokens: i64 = client + .query("SELECT toInt64(sum(output_tokens)) FROM cost_rollup_hourly") + .fetch_one() + .await + .unwrap(); + assert_eq!(tokens, 100); + let health: u64 = client + .query("SELECT toUInt64(sum(total_requests)) FROM provider_health_5m") + .fetch_one() + .await + .unwrap(); + assert_eq!(health, 25, "provider_health_5m counts each request once"); +} diff --git a/crates/test-support/tests/limits.rs b/crates/test-support/tests/limits.rs index 5565305a..0161533c 100644 --- a/crates/test-support/tests/limits.rs +++ b/crates/test-support/tests/limits.rs @@ -249,9 +249,8 @@ async fn rate_limit_window_validation_rejects_off_grid_seconds() { async fn api_key_scope_rate_limit_isolates_from_other_keys() { // Per-key rate-limit rules MUST fire on the gateway hot path — // schema supports `subject_kind='api_key_lineage'` and the auth - // middleware passes `api_key_id` through - // `compute_effective_surface_constraints`, which resolves it to - // `lineage_id` before binding. Two keys for the same user: one + // middleware loads the key lineage's rules through + // `compute_key_surface_constraints`. Two keys for the same user: one // carries a max_count=1 rule keyed on its lineage_id, the other // carries nothing. Each key must behave independently. let app = TestApp::spawn().await; @@ -445,3 +444,497 @@ async fn api_key_rate_limit_survives_rotation_via_lineage_id() { r.status ); } + +// ---------------------------------------------------------------------------- +// Enforcement: token limits, key counters, refusals that charge nothing, +// and the Retry-After a refusal carries. +// ---------------------------------------------------------------------------- + +fn chat() -> Json { + json!({"model": "gpt-test", "messages": [{"role": "user", "content": "x"}]}) +} + +/// `current` for every rule and cap on a subject, from the console's +/// usage endpoint — what an operator sees. +async fn usage(con: &TestClient, kind: &str, id: Uuid) -> (Vec, Vec) { + let r = con + .get(&format!("/api/admin/limits/{kind}/{id}/usage")) + .await + .unwrap(); + r.assert_ok(); + let v: Json = r.json().unwrap(); + let currents = |field: &str| -> Vec { + v[field] + .as_array() + .unwrap() + .iter() + .map(|x| x["current"].as_i64().unwrap()) + .collect() + }; + (currents("rules"), currents("caps")) +} + +fn retry_after(r: &think_watch_test_support::client::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)) +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_token_limit_refuses_once_the_window_is_used_up() { + // A request costs ~10 weighted tokens against the mock. A limit of 5 + // lets the first request through (nothing was used yet), records all + // of what it used even though that overshoots, and refuses the next. + let app = TestApp::spawn().await; + let (api_key, user_id) = seed_runtime(&app).await; + fixtures::create_rate_limit_rule(&app.db, "user", user_id, "ai_gateway", "tokens", 60, 5) + .await + .unwrap(); + let con = admin_session(&app).await; + + let gw = app.gateway_client(); + gw.set_bearer(&api_key); + gw.post("/v1/chat/completions", chat()) + .await + .unwrap() + .assert_ok(); + + let (rules, _) = usage(&con, "user", user_id).await; + assert!( + rules[0] > 5, + "the whole use is recorded, past the limit; got {rules:?}" + ); + + let r = gw.post("/v1/chat/completions", chat()).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}"); + assert_eq!( + usage(&con, "user", user_id).await.0, + rules, + "a refused request records nothing" + ); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn every_request_limit_applies_when_there_are_several() { + let app = TestApp::spawn().await; + let (api_key, user_id) = seed_runtime(&app).await; + fixtures::create_rate_limit_rule(&app.db, "user", user_id, "ai_gateway", "requests", 60, 2) + .await + .unwrap(); + fixtures::create_rate_limit_rule(&app.db, "user", user_id, "ai_gateway", "requests", 300, 5) + .await + .unwrap(); + + let gw = app.gateway_client(); + gw.set_bearer(&api_key); + for _ in 0..2 { + gw.post("/v1/chat/completions", chat()) + .await + .unwrap() + .assert_ok(); + } + let r = gw.post("/v1/chat/completions", chat()).await.unwrap(); + assert_eq!(r.status.as_u16(), 429, "body={}", r.text()); +} + +/// A key and its owner, each with limits of their own. +async fn second_key(app: &TestApp, user_id: Uuid) -> fixtures::SeededApiKey { + fixtures::create_api_key( + &app.db, + user_id, + &unique_name("limit-key"), + &["ai_gateway"], + None, + None, + ) + .await + .unwrap() +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_keys_limits_count_on_the_key_not_on_its_owner() { + let app = TestApp::spawn().await; + let (_, user_id) = seed_runtime(&app).await; + let key_a = second_key(&app, user_id).await; + let key_b = second_key(&app, user_id).await; + for key in [&key_a, &key_b] { + fixtures::create_rate_limit_rule( + &app.db, + "api_key_lineage", + key.row.lineage_id, + "ai_gateway", + "requests", + 60, + 1, + ) + .await + .unwrap(); + } + let con = admin_session(&app).await; + let gw = app.gateway_client(); + + gw.set_bearer(&key_a.plaintext); + gw.post("/v1/chat/completions", chat()) + .await + .unwrap() + .assert_ok(); + // Key B has a counter of its own: A's request does not use it up. + gw.set_bearer(&key_b.plaintext); + gw.post("/v1/chat/completions", chat()) + .await + .unwrap() + .assert_ok(); + gw.set_bearer(&key_a.plaintext); + let r = gw.post("/v1/chat/completions", chat()).await.unwrap(); + assert_eq!(r.status.as_u16(), 429, "body={}", r.text()); + + // The console reads the counter the gateway writes. + assert_eq!(usage(&con, "api_key", key_a.row.id).await.0, vec![1]); + assert_eq!(usage(&con, "api_key", key_b.row.id).await.0, vec![1]); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_key_limit_does_not_lift_its_owners() { + // The owner may make one request a minute; the key's own limit of + // five does not raise that. + let app = TestApp::spawn().await; + let (_, user_id) = seed_runtime(&app).await; + let key = second_key(&app, user_id).await; + fixtures::create_rate_limit_rule(&app.db, "user", user_id, "ai_gateway", "requests", 60, 1) + .await + .unwrap(); + fixtures::create_rate_limit_rule( + &app.db, + "api_key_lineage", + key.row.lineage_id, + "ai_gateway", + "requests", + 60, + 5, + ) + .await + .unwrap(); + + let gw = app.gateway_client(); + gw.set_bearer(&key.plaintext); + gw.post("/v1/chat/completions", chat()) + .await + .unwrap() + .assert_ok(); + let r = gw.post("/v1/chat/completions", chat()).await.unwrap(); + assert_eq!(r.status.as_u16(), 429, "body={}", r.text()); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_key_budget_counts_on_the_key() { + let app = TestApp::spawn().await; + let (_, user_id) = seed_runtime(&app).await; + let key = second_key(&app, user_id).await; + fixtures::create_budget_cap( + &app.db, + "api_key_lineage", + key.row.lineage_id, + "daily", + 1_000_000, + ) + .await + .unwrap(); + let con = admin_session(&app).await; + + let gw = app.gateway_client(); + gw.set_bearer(&key.plaintext); + gw.post("/v1/chat/completions", chat()) + .await + .unwrap() + .assert_ok(); + let (_, caps) = usage(&con, "api_key", key.row.id).await; + assert!( + caps[0] > 0, + "the key's budget counted the request: {caps:?}" + ); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_request_refused_by_a_budget_charges_no_request_limit() { + use fred::interfaces::KeysInterface; + use think_watch_common::limits::budget; + + let app = TestApp::spawn().await; + let (api_key, user_id) = seed_runtime(&app).await; + fixtures::create_rate_limit_rule(&app.db, "user", user_id, "ai_gateway", "requests", 60, 2) + .await + .unwrap(); + fixtures::create_budget_cap(&app.db, "user", user_id, "daily", 100) + .await + .unwrap(); + let spent = budget::build_key("user", user_id, "daily", chrono::Utc::now()); + let _: () = app + .state + .redis + .set(&spent, 105, None, None, false) + .await + .unwrap(); + let con = admin_session(&app).await; + + let gw = app.gateway_client(); + gw.set_bearer(&api_key); + for _ in 0..3 { + let r = gw.post("/v1/chat/completions", chat()).await.unwrap(); + assert_eq!(r.status.as_u16(), 429, "body={}", r.text()); + } + assert_eq!(usage(&con, "user", user_id).await.0, vec![0]); + + // With the budget freed, both requests the limit allows go through. + let _: () = app.state.redis.del(&spent).await.unwrap(); + for _ in 0..2 { + gw.post("/v1/chat/completions", chat()) + .await + .unwrap() + .assert_ok(); + } +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_request_refused_by_a_token_limit_charges_no_request_limit() { + let app = TestApp::spawn().await; + let (api_key, user_id) = seed_runtime(&app).await; + let requests = + fixtures::create_rate_limit_rule(&app.db, "user", user_id, "ai_gateway", "requests", 60, 5) + .await + .unwrap(); + fixtures::create_rate_limit_rule(&app.db, "user", user_id, "ai_gateway", "tokens", 300, 1) + .await + .unwrap(); + let con = admin_session(&app).await; + + let gw = app.gateway_client(); + gw.set_bearer(&api_key); + gw.post("/v1/chat/completions", chat()) + .await + .unwrap() + .assert_ok(); + for _ in 0..3 { + let r = gw.post("/v1/chat/completions", chat()).await.unwrap(); + assert_eq!(r.status.as_u16(), 429, "body={}", r.text()); + } + + let r: Json = con + .get(&format!("/api/admin/limits/user/{user_id}/usage")) + .await + .unwrap() + .json() + .unwrap(); + let requests_used = r["rules"] + .as_array() + .unwrap() + .iter() + .find(|x| x["rule_id"] == json!(requests)) + .map(|x| x["current"].as_i64().unwrap()); + assert_eq!(requests_used, Some(1), "usage: {r}"); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_rate_limit_says_when_its_window_frees() { + let app = TestApp::spawn().await; + let (api_key, user_id) = seed_runtime(&app).await; + fixtures::create_rate_limit_rule(&app.db, "user", user_id, "ai_gateway", "requests", 300, 1) + .await + .unwrap(); + + let gw = app.gateway_client(); + gw.set_bearer(&api_key); + gw.post("/v1/chat/completions", chat()) + .await + .unwrap() + .assert_ok(); + let r = gw.post("/v1/chat/completions", chat()).await.unwrap(); + assert_eq!(r.status.as_u16(), 429); + // The one request leaves the five-minute window in under five + // minutes, and not before the window has nearly run its course. + let secs = retry_after(&r); + assert!((290..=300).contains(&secs), "Retry-After {secs}"); + assert!( + r.headers.get("x-should-retry").is_none(), + "a window frees by itself; SDK retries are welcome" + ); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_spent_budget_says_when_its_period_ends_and_not_to_retry() { + use chrono::{Datelike, TimeZone, Utc}; + use fred::interfaces::KeysInterface; + use think_watch_common::limits::budget; + + let app = TestApp::spawn().await; + let (api_key, user_id) = seed_runtime(&app).await; + fixtures::create_budget_cap(&app.db, "user", user_id, "monthly", 100) + .await + .unwrap(); + let now = Utc::now(); + let spent = budget::build_key("user", user_id, "monthly", now); + let _: () = app + .state + .redis + .set(&spent, 100, None, None, false) + .await + .unwrap(); + let (y, m) = if now.month() == 12 { + (now.year() + 1, 1) + } else { + (now.year(), now.month() + 1) + }; + let next_month = Utc.with_ymd_and_hms(y, m, 1, 0, 0, 0).unwrap(); + let expected = (next_month - now).num_seconds(); + + let gw = app.gateway_client(); + gw.set_bearer(&api_key); + // The Anthropic surface, to see the refusal in the caller's format. + let r = gw + .post( + "/v1/messages", + json!({"model": "gpt-test", "max_tokens": 16, + "messages": [{"role": "user", "content": "x"}]}), + ) + .await + .unwrap(); + assert_eq!(r.status.as_u16(), 429, "body={}", r.text()); + let secs = retry_after(&r) as i64; + assert!( + (expected - 5..=expected + 5).contains(&secs), + "Retry-After {secs}, the month ends in {expected}" + ); + assert_eq!( + r.headers + .get("x-should-retry") + .and_then(|v| v.to_str().ok()), + Some("false") + ); + let body: Json = r.json().unwrap(); + assert_eq!(body["type"], "error", "{body}"); + assert_eq!(body["error"]["type"], "rate_limit_error", "{body}"); +} + +// ---------------------------------------------------------------------------- +// The scripts themselves, on a clock the test sets. +// ---------------------------------------------------------------------------- + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn the_limit_scripts_fill_free_and_record() { + let app = TestApp::spawn().await; + think_watch_test_support::redis_scripts::exercise_the_limit_scripts(&app.state.redis).await; +} + +// ---------------------------------------------------------------------------- +// MCP: a key's limits count on the key there too. +// ---------------------------------------------------------------------------- + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn a_keys_mcp_limits_count_on_the_key_not_on_its_owner() { + 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("limits-mcp"), + "lim", + &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 mut keys = Vec::new(); + for _ in 0..2 { + let key = fixtures::create_api_key( + &app.db, + user.user.id, + &unique_name("mcp-key"), + &["mcp_gateway"], + None, + None, + ) + .await + .unwrap(); + fixtures::create_rate_limit_rule( + &app.db, + "api_key_lineage", + key.row.lineage_id, + "mcp_gateway", + "requests", + 60, + 1, + ) + .await + .unwrap(); + keys.push(key); + } + let con = admin_session(&app).await; + + let gw = app.gateway_client(); + let call = |id: i64| { + json!({"jsonrpc": "2.0", "id": id, "method": "tools/call", + "params": {"name": "lim__anything", "arguments": {}}}) + }; + let mut answers = Vec::new(); + for (id, key) in [(1, &keys[0]), (2, &keys[1]), (3, &keys[0])] { + gw.set_bearer(&key.plaintext); + let r = gw.post("/mcp", call(id)).await.unwrap(); + r.assert_ok(); + answers.push(r.json::().unwrap()); + } + assert!(answers[0]["result"].is_object(), "{}", answers[0]); + assert!( + answers[1]["result"].is_object(), + "the second key has its own counter: {}", + answers[1] + ); + assert_eq!( + answers[2]["error"]["message"], "Rate limited: api_key_lineage:requests/1m", + "{}", + answers[2] + ); + assert_eq!(usage(&con, "api_key", keys[0].row.id).await.0, vec![1]); +} diff --git a/crates/test-support/tests/multi_instance_config_sync.rs b/crates/test-support/tests/multi_instance_config_sync.rs index b42336d6..40e56606 100644 --- a/crates/test-support/tests/multi_instance_config_sync.rs +++ b/crates/test-support/tests/multi_instance_config_sync.rs @@ -46,8 +46,7 @@ async fn instance_b_picks_up_instance_a_setting_change_via_pubsub() { .await .expect("instance B DynamicConfig::load"), ); - let sub_cfg = fred::types::config::Config::from_url(&app.state.config.redis_url) - .expect("parse redis_url"); + let sub_cfg = app.state.config.redis_config().expect("parse redis_url"); let subscriber: fred::clients::SubscriberClient = fred::types::Builder::from_config(sub_cfg) .build_subscriber_client() .expect("build_subscriber_client"); diff --git a/crates/test-support/tests/redis_cluster.rs b/crates/test-support/tests/redis_cluster.rs new file mode 100644 index 00000000..94b83dc9 --- /dev/null +++ b/crates/test-support/tests/redis_cluster.rs @@ -0,0 +1,274 @@ +//! The gateway's Lua scripts and multi-key commands against a real +//! Redis Cluster. +//! +//! A cluster runs a script only when every key it declares hashes to one +//! slot, refuses a `DEL` across slots (`CROSSSLOT`), and answers `SCAN` +//! for one node only. These tests point at a cluster given by +//! `TEST_REDIS_CLUSTER_URL` (e.g. `redis-cluster://127.0.0.1:37001`) and +//! skip without it, so CI — which has a single Redis — stays green. To +//! run them locally, start a three-primary cluster whose nodes announce +//! an address this process can reach, for example: +//! +//! ```text +//! docker run -d --rm --name tw-redis-cluster \ +//! -p 37001:37001 -p 37002:37002 -p 37003:37003 redis:8-alpine sh -c ' +//! for p in 37001 37002 37003; do +//! redis-server --port $p --cluster-enabled yes --cluster-config-file n-$p.conf \ +//! --cluster-announce-ip 127.0.0.1 --protected-mode no --save "" --daemonize yes --dir /tmp +//! done; sleep 1 +//! redis-cli --cluster create 127.0.0.1:37001 127.0.0.1:37002 127.0.0.1:37003 \ +//! --cluster-replicas 0 --cluster-yes; tail -f /dev/null' +//! TEST_REDIS_CLUSTER_URL=redis-cluster://127.0.0.1:37001 \ +//! cargo nextest run -p think-watch-test-support --test redis_cluster --run-ignored only +//! ``` +//! +//! Keys carry fresh UUIDs, so nothing is flushed and runs don't collide. +//! +//! A TLS cluster (`rediss-cluster://`) works the same way, with +//! `TEST_REDIS_CA_CERT` naming the CA of its certificates; see +//! `tests/redis_tls.rs` for how to start one. + +use fred::clients::Client; +use fred::interfaces::{ClientLike, KeysInterface}; +use fred::types::Builder; +use think_watch_test_support::prelude::*; +use think_watch_test_support::test_redis_config; + +fn cluster_url() -> Option { + let url = std::env::var("TEST_REDIS_CLUSTER_URL").ok(); + if url.is_none() { + eprintln!("TEST_REDIS_CLUSTER_URL not set — skipping the Redis Cluster test"); + } + url +} + +async fn cluster(url: &str) -> Client { + let client = Builder::from_config(test_redis_config(url)) + .build() + .unwrap(); + client.init().await.unwrap(); + assert!(client.is_clustered(), "{url} is not a cluster URL"); + client +} + +#[ignore = "integration test — needs TEST_REDIS_CLUSTER_URL"] +#[tokio::test] +async fn the_limit_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_limit_scripts(&redis).await; +} + +#[ignore = "integration test — needs TEST_REDIS_CLUSTER_URL"] +#[tokio::test] +async fn budgets_route_health_and_quotas_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; + + // Budgets: one key per command, any slot. + let cap = |kind, id| BudgetCap { + id: Uuid::nil(), + subject_kind: kind, + subject_id: id, + period: BudgetPeriod::Daily, + limit_tokens: 100, + enabled: true, + expires_at: None, + reason: None, + created_by: None, + }; + let caps = [ + cap(BudgetSubject::User, Uuid::new_v4()), + cap(BudgetSubject::ApiKeyLineage, Uuid::new_v4()), + ]; + budget::add_weighted_tokens(&redis, &caps, 60) + .await + .unwrap(); + let (statuses, crossings) = budget::add_weighted_tokens(&redis, &caps, 60) + .await + .unwrap(); + assert_eq!( + statuses.iter().map(|s| s.current).collect::>(), + [120, 120] + ); + assert_eq!(crossings.len(), 6, "80/95/100 % on both caps"); + let spent = budget::current_spend(&redis, &caps).await.unwrap(); + assert_eq!( + spent.iter().map(|s| s.current).collect::>(), + [120, 120] + ); + + // Route health: three keys in one script, then two, then a DEL of + // all three. + let health = HealthTracker::new(redis.clone()); + let route = Uuid::new_v4(); + let cfg = CircuitBreakerConfig { + enabled: true, + error_pct: 50, + min_samples: 2, + window_secs: 60, + open_secs: 30, + }; + health.record(route, 10, true, cfg).await; + let after = health.record(route, 10, true, cfg).await; + assert_eq!(after.total, 2, "the record script ran"); + assert_eq!(after.lifetime_requests, 2); + assert_eq!(format!("{:?}", after.state), "Open", "the state write ran"); + 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 a_pattern_delete_reaches_every_node_of_a_cluster() { + let Some(url) = cluster_url() else { return }; + let redis = cluster(&url).await; + + let run = Uuid::new_v4(); + let keys: Vec = (0..300).map(|i| format!("tw-test:{run}:{i}")).collect(); + for k in &keys { + let _: () = redis.set(k, 1, None, None, false).await.unwrap(); + } + let slots: std::collections::HashSet = keys + .iter() + .map(|k| fred::util::redis_keyslot(k.as_bytes())) + .collect(); + assert!(slots.len() > 100, "the keys spread over the cluster"); + + let deleted = + think_watch_common::redis_keys::delete_matching(&redis, &format!("tw-test:{run}:*"), None) + .await + .unwrap(); + assert_eq!(deleted, keys.len()); + for k in &keys { + assert_eq!(redis.exists::(k).await.unwrap(), 0, "{k}"); + } +} + +#[ignore = "integration test — needs TEST_REDIS_CLUSTER_URL"] +#[tokio::test] +async fn the_gateway_enforces_limits_on_a_cluster() { + let Some(url) = cluster_url() else { return }; + let app = TestApp::try_spawn_with(SpawnOptions { + redis_url: Some(url), + ..Default::default() + }) + .await + .unwrap(); + + let user = fixtures::create_random_user(&app.db).await.unwrap(); + let user_id = user.user.id; + let mock = MockProvider::openai_chat_ok("gpt-test").await; + let provider = fixtures::create_provider( + &app.db, + &unique_name("cluster-prov"), + "openai", + &mock.uri(), + None, + ) + .await + .unwrap(); + fixtures::create_model_and_route(&app.db, provider.id, "gpt-test") + .await + .unwrap(); + app.rebuild_gateway_router().await; + let key = fixtures::create_api_key( + &app.db, + user_id, + &unique_name("cluster-key"), + &["ai_gateway"], + None, + None, + ) + .await + .unwrap(); + // The user's request limit and the key's token limit and budget: one + // script checks counters of both subjects. + fixtures::create_rate_limit_rule(&app.db, "user", user_id, "ai_gateway", "requests", 60, 2) + .await + .unwrap(); + fixtures::create_rate_limit_rule( + &app.db, + "api_key_lineage", + key.row.lineage_id, + "ai_gateway", + "tokens", + 60, + 1_000_000, + ) + .await + .unwrap(); + fixtures::create_budget_cap( + &app.db, + "api_key_lineage", + key.row.lineage_id, + "daily", + 1_000_000, + ) + .await + .unwrap(); + let con = admin_session(&app).await; + + let gw = app.gateway_client(); + gw.set_bearer(&key.plaintext); + // A prompt of its own: the response cache lives in the same Redis, + // which nothing flushes, and a cached answer records no tokens. + let chat = json!({"model": "gpt-test", + "messages": [{"role": "user", "content": Uuid::new_v4().to_string()}]}); + for _ in 0..2 { + gw.post("/v1/chat/completions", chat.clone()) + .await + .unwrap() + .assert_ok(); + } + let r = gw.post("/v1/chat/completions", chat).await.unwrap(); + assert_eq!(r.status.as_u16(), 429, "body={}", r.text()); + assert!(r.headers.get("retry-after").is_some()); + + let usage: Json = con + .get(&format!("/api/admin/limits/api_key/{}/usage", key.row.id)) + .await + .unwrap() + .json() + .unwrap(); + assert!( + usage["rules"][0]["current"].as_i64().unwrap() > 0, + "the key's tokens were recorded: {usage}" + ); + assert!( + usage["caps"][0]["current"].as_i64().unwrap() > 0, + "the key's budget was debited: {usage}" + ); + let usage: Json = con + .get(&format!("/api/admin/limits/user/{user_id}/usage")) + .await + .unwrap() + .json() + .unwrap(); + assert_eq!(usage["rules"][0]["current"], 2, "{usage}"); +} + +#[ignore = "integration test — needs TEST_REDIS_CLUSTER_URL"] +#[tokio::test] +async fn config_change_notices_reach_a_subscriber_on_a_cluster() { + // What `init::spawn_config_subscriber` does with the same URL: a + // subscriber on one node hears a publish sent through another. + let Some(url) = cluster_url() else { return }; + let publisher = cluster(&url).await; + think_watch_test_support::assert_config_notice_arrives(test_redis_config(&url), &publisher) + .await; +} diff --git a/crates/test-support/tests/redis_tls.rs b/crates/test-support/tests/redis_tls.rs new file mode 100644 index 00000000..9419ffc0 --- /dev/null +++ b/crates/test-support/tests/redis_tls.rs @@ -0,0 +1,188 @@ +//! The server against a Redis that speaks TLS only (`rediss://`). +//! +//! Managed Redis services usually require TLS. These tests point at a +//! TLS-only Redis given by `TEST_REDIS_TLS_URL` (e.g. +//! `rediss://:pw@localhost:46380`) whose certificate the CA in +//! `TEST_REDIS_CA_CERT` signed, and skip without the URL, so CI — which +//! has a plain Redis — stays green. To run them locally, make a throwaway +//! CA and a certificate for `localhost` and `127.0.0.1`, and start Redis +//! with TLS on and its plain port off: +//! +//! ```text +//! mkdir -p /tmp/tw-tls && cd /tmp/tw-tls +//! openssl req -x509 -new -nodes -newkey rsa:2048 -keyout ca.key -out ca.crt \ +//! -days 2 -subj "/CN=ThinkWatch test CA" +//! openssl req -new -nodes -newkey rsa:2048 -keyout server.key -out server.csr \ +//! -subj "/CN=localhost" +//! printf 'subjectAltName=DNS:localhost,IP:127.0.0.1\n' > san.cnf +//! openssl x509 -req -in server.csr -CA ca.crt -CAkey ca.key -CAcreateserial \ +//! -days 2 -extfile san.cnf -out server.crt +//! chmod 644 server.key # readable by the container's redis user +//! docker run -d --rm --name tw-redis-tls -p 46380:6380 -v /tmp/tw-tls:/tls:ro \ +//! redis:8-alpine redis-server --port 0 --tls-port 6380 \ +//! --tls-cert-file /tls/server.crt --tls-key-file /tls/server.key \ +//! --tls-ca-cert-file /tls/ca.crt --tls-auth-clients no --requirepass pw +//! TEST_REDIS_TLS_URL=rediss://:pw@localhost:46380 TEST_REDIS_CA_CERT=/tmp/tw-tls/ca.crt \ +//! cargo nextest run -p think-watch-test-support --test redis_tls --run-ignored only +//! ``` +//! +//! The whole suite runs over TLS the same way: `TEST_REDIS_URL` set to +//! `rediss://:pw@localhost:46380/1` with `TEST_REDIS_CA_CERT`. +//! +//! A TLS Redis Cluster (`rediss-cluster://`) runs `tests/redis_cluster.rs`. +//! Its nodes announce `127.0.0.1`, which the certificate above names. The +//! cluster bus authenticates both of its ends with that certificate, so it +//! must not be limited to server use (the commands above set no +//! `extendedKeyUsage`): +//! +//! ```text +//! docker run -d --rm --name tw-redis-cluster-tls -v /tmp/tw-tls:/tls:ro \ +//! -p 37101:37101 -p 37102:37102 -p 37103:37103 redis:8-alpine sh -c ' +//! tls="--tls-cert-file /tls/server.crt --tls-key-file /tls/server.key +//! --tls-ca-cert-file /tls/ca.crt --tls-auth-clients no" +//! for p in 37101 37102 37103; do +//! redis-server --port 0 --tls-port $p --tls-cluster yes $tls \ +//! --cluster-enabled yes --cluster-config-file n-$p.conf \ +//! --cluster-announce-ip 127.0.0.1 --protected-mode no --save "" \ +//! --daemonize yes --dir /tmp +//! done; sleep 1 +//! redis-cli --tls --cacert /tls/ca.crt --cluster create 127.0.0.1:37101 \ +//! 127.0.0.1:37102 127.0.0.1:37103 --cluster-replicas 0 --cluster-yes +//! tail -f /dev/null' +//! TEST_REDIS_CLUSTER_URL=rediss-cluster://127.0.0.1:37101 TEST_REDIS_CA_CERT=/tmp/tw-tls/ca.crt \ +//! cargo nextest run -p think-watch-test-support --test redis_cluster --run-ignored only +//! ``` +//! +//! Keys carry fresh UUIDs, so nothing is flushed and runs don't collide. + +use std::time::Duration; + +use fred::clients::Client; +use fred::interfaces::ClientLike; +use fred::types::Builder; +use think_watch_test_support::prelude::*; +use think_watch_test_support::test_redis_config; + +fn tls_url() -> Option { + let url = std::env::var("TEST_REDIS_TLS_URL").ok(); + if url.is_none() { + eprintln!("TEST_REDIS_TLS_URL not set — skipping the Redis TLS test"); + } + url +} + +async fn connect(url: &str) -> Client { + let config = test_redis_config(url); + assert!(config.uses_rustls(), "{url} is not a TLS URL"); + let client = Builder::from_config(config).build().unwrap(); + client.init().await.unwrap(); + client +} + +#[ignore = "integration test — needs TEST_REDIS_TLS_URL"] +#[tokio::test] +async fn the_limit_scripts_run_over_tls() { + let Some(url) = tls_url() else { return }; + let redis = connect(&url).await; + think_watch_test_support::redis_scripts::exercise_the_limit_scripts(&redis).await; +} + +#[ignore = "integration test — needs TEST_REDIS_TLS_URL"] +#[tokio::test] +async fn a_certificate_the_roots_do_not_cover_is_refused() { + // The test server's certificate comes from a private CA. Without + // REDIS_CA_CERT only the system's roots are trusted, so the handshake + // must fail: the certificate is checked, not merely accepted. + let Some(url) = tls_url() else { return }; + let config = think_watch_common::redis_config::client_config(&url, None).unwrap(); + let client = Builder::from_config(config).build().unwrap(); + let init = tokio::time::timeout(Duration::from_secs(10), client.init()) + .await + .expect("the handshake ends within 10 s"); + let err = init.expect_err("a certificate from an untrusted CA was accepted"); + // fred reports rustls's verdict as an I/O error; the rustls error + // inside names the certificate. + assert!(err.details().contains("InvalidCertificate"), "{err:?}"); +} + +#[ignore = "integration test — needs TEST_REDIS_TLS_URL"] +#[tokio::test] +async fn the_gateway_runs_over_tls() { + let Some(url) = tls_url() else { return }; + let app = TestApp::try_spawn_with(SpawnOptions { + redis_url: Some(url), + ..Default::default() + }) + .await + .unwrap(); + assert!(app.state.config.redis_config().unwrap().uses_rustls()); + + let user = fixtures::create_random_user(&app.db).await.unwrap(); + let mock = MockProvider::openai_chat_ok("gpt-test").await; + let provider = fixtures::create_provider( + &app.db, + &unique_name("tls-prov"), + "openai", + &mock.uri(), + None, + ) + .await + .unwrap(); + fixtures::create_model_and_route(&app.db, provider.id, "gpt-test") + .await + .unwrap(); + app.rebuild_gateway_router().await; + + let gw = app.gateway_client(); + let ready: Json = gw + .get("/health/ready") + .await + .unwrap() + .assert_ok() + .json() + .unwrap(); + assert_eq!(ready["redis"], true, "{ready}"); + + // A request limit lives in Redis: two pass, the third is refused. + fixtures::create_rate_limit_rule( + &app.db, + "user", + user.user.id, + "ai_gateway", + "requests", + 60, + 2, + ) + .await + .unwrap(); + let key = fixtures::create_api_key( + &app.db, + user.user.id, + &unique_name("tls-key"), + &["ai_gateway"], + None, + None, + ) + .await + .unwrap(); + gw.set_bearer(&key.plaintext); + let chat = json!({"model": "gpt-test", + "messages": [{"role": "user", "content": Uuid::new_v4().to_string()}]}); + for _ in 0..2 { + gw.post("/v1/chat/completions", chat.clone()) + .await + .unwrap() + .assert_ok(); + } + let r = gw.post("/v1/chat/completions", chat).await.unwrap(); + assert_eq!(r.status.as_u16(), 429, "body={}", r.text()); + + // Config change notices between instances: a subscriber built the + // way `init::spawn_config_subscriber` builds its three hears a + // publish sent through the server's own client. + think_watch_test_support::assert_config_notice_arrives( + app.state.config.redis_config().unwrap(), + &app.state.redis, + ) + .await; +} diff --git a/db/schema.sql b/db/schema.sql index 3b21f660..877519f4 100644 --- a/db/schema.sql +++ b/db/schema.sql @@ -5,7 +5,10 @@ -- in place when the schema changes; the application calls -- `sqlx::raw_sql(include_str!("../../../db/schema.sql"))` on every -- boot, and every statement here is wrapped in `IF NOT EXISTS` / --- `OR REPLACE` so a re-run is a no-op on an up-to-date DB. +-- `OR REPLACE` so a re-run is a no-op on an up-to-date DB. Replicas +-- starting together apply it one at a time, under an advisory lock +-- (crates/common/src/db.rs::run_migrations): run side by side, its DDL +-- deadlocks. -- -- Limits of declarative apply: -- * column rename, type narrowing, or DROP COLUMN need an explicit @@ -90,7 +93,10 @@ CREATE TABLE IF NOT EXISTS teams ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), name VARCHAR(255) NOT NULL UNIQUE, description TEXT, - -- Budget caps live in `budget_caps` (subject_kind = 'team'). + -- Teams carry no limits or budgets of their own: `budget_caps` and + -- `rate_limit_rules` attach to users and API keys only (see the + -- note above `rate_limit_rules`). A team's members are limited + -- through the roles the team grants. created_at TIMESTAMPTZ NOT NULL DEFAULT now() ); diff --git a/deploy/clickhouse/initdb.d/01_init.sql b/deploy/clickhouse/initdb.d/01_init.sql index 20c153b0..a876095f 100644 --- a/deploy/clickhouse/initdb.d/01_init.sql +++ b/deploy/clickhouse/initdb.d/01_init.sql @@ -185,24 +185,24 @@ ALTER TABLE gateway_logs ADD PROJECTION IF NOT EXISTS proj_by_latency ( -- compression dominates the cold-storage footprint. Sit AFTER -- session_id so the existing column order is preserved and new -- deployments + upgraded ones converge on the same shape. -ALTER TABLE gateway_logs ADD COLUMN IF NOT EXISTS request_body Nullable(String) CODEC(ZSTD(6)) AFTER session_id; -ALTER TABLE gateway_logs ADD COLUMN IF NOT EXISTS response_body Nullable(String) CODEC(ZSTD(6)) AFTER request_body; +-- +-- Body columns get a SHORTER TTL than the row-level retention: when it +-- fires, the value is reset to NULL while the row stays around for the +-- full table TTL, so metadata queries remain whole after bodies have +-- aged out. The 30 days here is only the TTL a column is created with. +-- The server sets the operator's `audit.body_retention_days` once it is +-- up (`apply_body_column_ttls`, crates/server/src/handlers/admin/ +-- retention.rs), and this file runs on every start, so it must never +-- set the TTL of a column that exists: a 30-day TTL applied, even for +-- the moment until the server restores a longer one, makes ClickHouse +-- clear every body older than 30 days. +ALTER TABLE gateway_logs ADD COLUMN IF NOT EXISTS request_body Nullable(String) CODEC(ZSTD(6)) TTL toDateTime(created_at) + INTERVAL 30 DAY AFTER session_id; +ALTER TABLE gateway_logs ADD COLUMN IF NOT EXISTS response_body Nullable(String) CODEC(ZSTD(6)) TTL toDateTime(created_at) + INTERVAL 30 DAY AFTER request_body; ALTER TABLE gateway_logs ADD COLUMN IF NOT EXISTS request_body_bytes Nullable(UInt32) AFTER response_body; ALTER TABLE gateway_logs ADD COLUMN IF NOT EXISTS response_body_bytes Nullable(UInt32) AFTER request_body_bytes; -- 'captured' | 'truncated' | 'disabled' | 'from_cache' | 'error' ALTER TABLE gateway_logs ADD COLUMN IF NOT EXISTS body_capture_status LowCardinality(Nullable(String)) AFTER response_body_bytes; --- Body columns get a SHORTER TTL than the row-level retention. The --- Rust side (`apply_body_column_ttls` in handlers/admin.rs) re-issues --- these at startup against the operator-configurable --- `audit.body_retention_days` setting (default 30); this seed-default --- exists so a CH bootstrap that happens before the server ever runs --- still has the right shape. When the column TTL fires, the value is --- reset to NULL while the row stays around for the full table TTL — --- so metadata queries remain whole even after bodies have aged out. -ALTER TABLE gateway_logs MODIFY COLUMN request_body TTL toDateTime(created_at) + INTERVAL 30 DAY; -ALTER TABLE gateway_logs MODIFY COLUMN response_body TTL toDateTime(created_at) + INTERVAL 30 DAY; - -- Substring search across captured bodies — auditors searching -- "which conversations mentioned API key XYZ" or "which tool calls -- referenced /etc/passwd" otherwise have to scan every row. tokenbf @@ -254,16 +254,14 @@ ALTER TABLE mcp_logs ADD PROJECTION IF NOT EXISTS proj_by_duration ( -- by sanitize_detail); promote both arguments and the upstream result -- to first-class columns so audit queries don't have to JSON-parse on -- every row. Same ZSTD(6) trade-off as gateway_logs. -ALTER TABLE mcp_logs ADD COLUMN IF NOT EXISTS tool_arguments Nullable(String) CODEC(ZSTD(6)) AFTER detail; -ALTER TABLE mcp_logs ADD COLUMN IF NOT EXISTS tool_result Nullable(String) CODEC(ZSTD(6)) AFTER tool_arguments; +-- Their TTL as gateway_logs' body columns above: set when the column +-- is created, never again by this file. +ALTER TABLE mcp_logs ADD COLUMN IF NOT EXISTS tool_arguments Nullable(String) CODEC(ZSTD(6)) TTL toDateTime(created_at) + INTERVAL 30 DAY AFTER detail; +ALTER TABLE mcp_logs ADD COLUMN IF NOT EXISTS tool_result Nullable(String) CODEC(ZSTD(6)) TTL toDateTime(created_at) + INTERVAL 30 DAY AFTER tool_arguments; ALTER TABLE mcp_logs ADD COLUMN IF NOT EXISTS arguments_bytes Nullable(UInt32) AFTER tool_result; ALTER TABLE mcp_logs ADD COLUMN IF NOT EXISTS result_bytes Nullable(UInt32) AFTER arguments_bytes; ALTER TABLE mcp_logs ADD COLUMN IF NOT EXISTS body_capture_status LowCardinality(Nullable(String)) AFTER result_bytes; --- Same body-column TTL story as gateway_logs above; see comment there. -ALTER TABLE mcp_logs MODIFY COLUMN tool_arguments TTL toDateTime(created_at) + INTERVAL 30 DAY; -ALTER TABLE mcp_logs MODIFY COLUMN tool_result TTL toDateTime(created_at) + INTERVAL 30 DAY; - -- Same substring-search rationale as gateway_logs above. ALTER TABLE mcp_logs ADD INDEX IF NOT EXISTS idx_tool_arguments ifNull(tool_arguments, '') TYPE tokenbf_v1(512, 3, 0) GRANULARITY 4; ALTER TABLE mcp_logs ADD INDEX IF NOT EXISTS idx_tool_result ifNull(tool_result, '') TYPE tokenbf_v1(512, 3, 0) GRANULARITY 4; diff --git a/deploy/helm/think-watch/Chart.yaml b/deploy/helm/think-watch/Chart.yaml index 53ed6b54..480e3fb5 100644 --- a/deploy/helm/think-watch/Chart.yaml +++ b/deploy/helm/think-watch/Chart.yaml @@ -2,8 +2,8 @@ apiVersion: v2 name: think-watch description: Enterprise AI API Gateway & MCP Management Platform type: application -version: 3.1.0 -appVersion: "3.1.0" +version: 3.2.0 +appVersion: "3.2.0" keywords: - ai - gateway diff --git a/deploy/helm/think-watch/README.md b/deploy/helm/think-watch/README.md index bd84dd75..4d1d0036 100644 --- a/deploy/helm/think-watch/README.md +++ b/deploy/helm/think-watch/README.md @@ -72,6 +72,143 @@ helm upgrade --install thinkwatch deploy/helm/think-watch \ When `bundled=false` and `externalUrl` is empty the chart fails at install-time with an explicit message — no silent broken Secret. +### PostgreSQL behind a connection pooler + +Each server instance sets up the schema when it starts, holding a +Postgres session-level advisory lock so that instances starting together +take turns. The lock belongs to one Postgres session, so `externalUrl` +must reach Postgres directly or through a pooler in session mode, never +one in transaction mode (PgBouncer `pool_mode = transaction`, or the +transaction-mode port of a managed pooler such as Supabase's): there, +the lock can stay held on a server connection the pooler keeps after the +instance is done with it, and every later start waits for it for good. +The server runs its schema setup on the connections of `DATABASE_URL`; +there is no separate URL for it. + +### Several replicas starting together + +Instances set up the schema one at a time, so with several replicas +starting together each one waits for those before it, and with ClickHouse +unreachable each one also retries it for up to about a minute. The server +answers its startup probe only after that. The probe gives a pod +`5 + 3 × 40 = 125` seconds by default (`startupProbe.initialDelaySeconds`, +`periodSeconds`, `failureThreshold`); raise `startupProbe.failureThreshold` +when pods are restarted before they finish starting. + +### Redis Cluster + +The external Redis can be a Redis Cluster. Give its URL the +`redis-cluster://` scheme and name one node or more; the server finds the +rest of the cluster from them: + +```yaml +redis: + bundled: false + externalUrl: redis-cluster://:pass@redis-0.redis:6379?node=redis-1.redis:6379&node=redis-2.redis:6379 +``` + +- Every node must be reachable from the server pods at the address it + announces to the cluster (`cluster-announce-ip` / `-port`): the server + follows the cluster's redirects to it. With `networkPolicy.enabled`, + ports the URL doesn't name go in `networkPolicy.extraEgress` (see + [Network policy](#network-policy)). +- A cluster has only database 0, so the URL names no `/`. +- Nothing else is needed. Every script the server runs keeps its keys in + one hash slot — the counters of one user's request share the tag + `{user:}`, a route's health keys `{}` — and pattern + deletes (cache invalidation) scan every primary. + +Rate limits, budgets, route health, caches and the config change +notices between instances are tested against a three-primary Redis 8 +cluster (`crates/test-support/tests/redis_cluster.rs`). + +### Redis over TLS + +Managed Redis services usually require TLS: ElastiCache with in-transit +encryption, Upstash, Azure Cache for Redis, Redis Cloud. Give the URL +the `rediss://` scheme — `rediss-cluster://` for a cluster: + +```yaml +redis: + bundled: false + externalUrl: rediss://:pass@master.my-cache.abc123.use1.cache.amazonaws.com:6379 +``` + +- The server checks the Redis certificate against the public CAs, the + same roots it trusts for upstream HTTPS, and against the host name in + the URL. Use the endpoint name the service gives, not an IP address, + unless the certificate names that address. The name must be in the + certificate's subjectAltName: a certificate that names the host only in + its CN, which `redis-cli` accepts, is refused. +- In a cluster, each node is reached at the address it announces, and + its certificate must name that address — the host name, or the IP + address when the node announces one. A cluster whose nodes announce + IP addresses that their certificates do not name cannot be used over + TLS. +- The port is whatever the service uses for TLS (Azure Cache for Redis: + `6380`). Write it in the URL: a `rediss://` URL without one means + `6379`, as `redis://` does. + +A self-hosted Redis whose certificate a private CA signed needs that +CA. Put its PEM certificate in a Secret and name it; the server then +trusts only the certificates in it for Redis: + +```bash +kubectl -n thinkwatch create secret generic redis-ca --from-file=ca.crt=./ca.crt +``` + +```yaml +redis: + bundled: false + externalUrl: rediss://:pass@redis.internal:6379 + caSecret: + name: redis-ca + key: ca.crt +``` + +The chart mounts the key at `/etc/thinkwatch/redis-ca/` and sets +`REDIS_CA_CERT` to it. The server reads it at start, so restart the +server pods after changing the Secret. Outside the chart, set +`REDIS_CA_CERT` to the PEM file's path yourself. Client certificates +(mutual TLS) are not supported: give such a Redis `tls-auth-clients no` +and authenticate with the password. + +## Network policy + +`networkPolicy.enabled` limits what the server pods may reach: DNS, port +`443` (upstreams, the OIDC provider), and PostgreSQL, Redis and +ClickHouse on the ports the server connects to them on: + +- A bundled database: its service port (`5432`, `6379`, `8123`). +- An external one: every port its `externalUrl` names, so a database on + another port needs no setting of its own. That is the port of each host + in the URL, of each `node=` of a Redis Cluster or Sentinel URL, and a + Postgres `?port=`. +- A URL without a port: the client's default for the scheme, which is + `5432` for `postgres://`, `6379` for `redis://` and `rediss://` (TLS + does not change it), `26379` for a Sentinel and `6379` for the primary it + points to, and `80` for ClickHouse's `http://`. + +The server talks to ClickHouse over plain HTTP only, so ClickHouse's +native port (`9000`) is not allowed. HTTPS to ClickHouse is not supported: +`clickhouse.externalUrl` must be an `http://` URL. + +What no URL names goes in `networkPolicy.extraEgress`, rules added to the +server's egress as written: Redis Cluster nodes that announce ports the +URL doesn't list, a Sentinel's primary on a port other than `6379`, an +upstream, MCP server or S3 endpoint on a port other than `443` (RustFS and +MinIO listen on `9000`). + +```yaml +networkPolicy: + enabled: true + extraEgress: + - ports: + - port: 7000 + endPort: 7005 + protocol: TCP +``` + ## Rotating secrets `-secrets` is kept on `helm uninstall`. To rotate passwords: diff --git a/deploy/helm/think-watch/templates/_helpers.tpl b/deploy/helm/think-watch/templates/_helpers.tpl index 6735a849..964bed49 100644 --- a/deploy/helm/think-watch/templates/_helpers.tpl +++ b/deploy/helm/think-watch/templates/_helpers.tpl @@ -24,3 +24,70 @@ Resource names for bundled databases — stable across upgrades. {{- define "tw.redis.name" -}}{{ .Release.Name }}-redis{{- end -}} {{- define "tw.clickhouse.name" -}}{{ .Release.Name }}-clickhouse{{- end -}} {{- define "tw.secret.name" -}}{{ .Release.Name }}-secrets{{- end -}} + +{{/* +The TCP ports a URL reaches, space-separated: the port of each host in it +(comma-separated hosts included), of each `node=` in its query (the seed +nodes of a Redis Cluster or Sentinel) and of a `port=` in its query (which +Postgres honours); `default` for a host written without a port. Never +fails on a URL it cannot read: the default stands. +Usage: {{ include "tw.urlPorts" (dict "url" $url "default" 6379) }} +*/}} +{{- define "tw.urlPorts" -}} +{{- $rest := regexReplaceAll "^[A-Za-z][A-Za-z0-9+.-]*://" .url "" -}} +{{- $hosts := regexReplaceAll "^.*@" (regexFind "^[^/?#]*" $rest) "" -}} +{{- $query := "" -}} +{{- if contains "?" $rest -}} +{{- $query = regexReplaceAll "#.*$" (regexReplaceAll "^[^?]*[?]" $rest "") "" -}} +{{- end -}} +{{- $ports := list -}} +{{- range $host := splitList "," $hosts -}} +{{- $port := trimPrefix ":" (regexFind ":[0-9]+$" (regexReplaceAll "\\[[^]]*\\]" $host "")) -}} +{{- $ports = append $ports (default (toString $.default) $port) -}} +{{- end -}} +{{- range $param := splitList "&" $query -}} +{{- if hasPrefix "node=" $param -}} +{{- $node := trimPrefix "node=" $param | replace "%3A" ":" | replace "%3a" ":" -}} +{{- $port := trimPrefix ":" (regexFind ":[0-9]+$" (regexReplaceAll "\\[[^]]*\\]" $node "")) -}} +{{- if $port -}}{{- $ports = append $ports $port -}}{{- end -}} +{{- else if hasPrefix "port=" $param -}} +{{- $port := regexFind "^[0-9]+$" (trimPrefix "port=" $param) -}} +{{- if $port -}}{{- $ports = append $ports $port -}}{{- end -}} +{{- end -}} +{{- end -}} +{{- $ports | uniq | join " " -}} +{{- end -}} + +{{/* +The ports the server reaches each database on, for the network policy: +the bundled service's, or those of its externalUrl, with the client's +default for the scheme where the URL names none (Postgres 5432; Redis +6379, a Sentinel 26379 and 6379 for its primary; ClickHouse 80, for the +http:// URL the server requires — an https:// one, which it refuses, +would mean 443). +*/}} +{{- define "tw.postgres.ports" -}} +{{- if .Values.postgres.bundled -}}5432 +{{- else -}}{{ include "tw.urlPorts" (dict "url" .Values.postgres.externalUrl "default" 5432) }} +{{- end -}} +{{- end -}} + +{{- define "tw.redis.ports" -}} +{{- if .Values.redis.bundled -}}6379 +{{- else -}} +{{- $url := .Values.redis.externalUrl -}} +{{- $sentinel := or (regexMatch "^[A-Za-z]+-sentinel://" $url) (contains "sentinelServiceName=" $url) -}} +{{- $ports := splitList " " (include "tw.urlPorts" (dict "url" $url "default" (ternary 26379 6379 $sentinel))) -}} +{{- /* The primary a Sentinel points to: the URL can't say where. */ -}} +{{- if $sentinel -}}{{- $ports = append $ports "6379" -}}{{- end -}} +{{- $ports | uniq | join " " -}} +{{- end -}} +{{- end -}} + +{{- define "tw.clickhouse.ports" -}} +{{- if .Values.clickhouse.bundled -}}8123 +{{- else -}} +{{- $url := .Values.clickhouse.externalUrl -}} +{{- include "tw.urlPorts" (dict "url" $url "default" (ternary 443 80 (hasPrefix "https://" (lower $url)))) -}} +{{- end -}} +{{- end -}} diff --git a/deploy/helm/think-watch/templates/deployment.yaml b/deploy/helm/think-watch/templates/deployment.yaml index c2be4980..fc19ca03 100644 --- a/deploy/helm/think-watch/templates/deployment.yaml +++ b/deploy/helm/think-watch/templates/deployment.yaml @@ -95,13 +95,24 @@ spec: name: {{ .Release.Name }}-config - secretRef: name: {{ .Release.Name }}-secrets + {{- with .Values.redis.caSecret }} + {{- if .name }} + env: + - name: REDIS_CA_CERT + value: {{ printf "/etc/thinkwatch/redis-ca/%s" (required "redis.caSecret.key: name the key of Secret redis.caSecret.name that holds the CA certificate (e.g. ca.crt)" .key) | quote }} + volumeMounts: + - name: redis-ca + mountPath: /etc/thinkwatch/redis-ca + readOnly: true + {{- end }} + {{- end }} startupProbe: httpGet: path: /health/live port: gateway - initialDelaySeconds: 5 - periodSeconds: 3 - failureThreshold: 10 + initialDelaySeconds: {{ .Values.startupProbe.initialDelaySeconds }} + periodSeconds: {{ .Values.startupProbe.periodSeconds }} + failureThreshold: {{ .Values.startupProbe.failureThreshold }} livenessProbe: httpGet: path: /health/live @@ -118,6 +129,17 @@ spec: failureThreshold: 3 resources: {{- toYaml .Values.resources.server | nindent 12 }} + {{- with .Values.redis.caSecret }} + {{- if .name }} + volumes: + - name: redis-ca + secret: + secretName: {{ .name | quote }} + items: + - key: {{ .key | quote }} + path: {{ .key | quote }} + {{- end }} + {{- end }} --- apiVersion: apps/v1 kind: Deployment diff --git a/deploy/helm/think-watch/templates/networkpolicy.yaml b/deploy/helm/think-watch/templates/networkpolicy.yaml index d2c77e56..06955d37 100644 --- a/deploy/helm/think-watch/templates/networkpolicy.yaml +++ b/deploy/helm/think-watch/templates/networkpolicy.yaml @@ -1,6 +1,7 @@ {{- if .Values.networkPolicy.enabled }} # Server pod: allow ingress on gateway (3000) and console (3001) ports, -# and egress to PostgreSQL, Redis, ClickHouse, and DNS. +# and egress to PostgreSQL, Redis, ClickHouse, DNS, HTTPS and +# networkPolicy.extraEgress. apiVersion: networking.k8s.io/v1 kind: NetworkPolicy metadata: @@ -37,26 +38,24 @@ spec: protocol: UDP - port: 53 protocol: TCP - # PostgreSQL + # PostgreSQL, Redis and ClickHouse (HTTP), on the ports the server + # connects to them on (tw..ports in _helpers.tpl) + {{- range $db := list "postgres" "redis" "clickhouse" }} + # {{ $db }} - ports: - - port: 5432 - protocol: TCP - # Redis - - ports: - - port: 6379 - protocol: TCP - # ClickHouse HTTP - - ports: - - port: 8123 - protocol: TCP - # ClickHouse native - - ports: - - port: 9000 + {{- range splitList " " (include (printf "tw.%s.ports" $db) $) }} + - port: {{ . }} protocol: TCP + {{- end }} + {{- end }} # OIDC provider (HTTPS) - ports: - port: 443 protocol: TCP + {{- with .Values.networkPolicy.extraEgress }} + # networkPolicy.extraEgress + {{- toYaml . | nindent 4 }} + {{- end }} --- # Web pod: allow ingress on port 80, egress only to server console port and DNS. apiVersion: networking.k8s.io/v1 diff --git a/deploy/helm/think-watch/values-production.yaml.example b/deploy/helm/think-watch/values-production.yaml.example index cc13e9d0..f4886d47 100644 --- a/deploy/helm/think-watch/values-production.yaml.example +++ b/deploy/helm/think-watch/values-production.yaml.example @@ -35,7 +35,8 @@ config: rustLog: "info,think_watch=info" # Swap any of these to bundled: false + externalUrl to point at a managed -# service (RDS, ElastiCache, ClickHouse Cloud). Bundled=true gives you a +# service (RDS, ElastiCache; ClickHouse over plain http:// only, as the +# server does not support HTTPS to it). Bundled=true gives you a # single-replica in-cluster StatefulSet — fine for small prod, not for HA. postgres: @@ -45,6 +46,13 @@ postgres: redis: bundled: false externalUrl: "redis://:REDACTED@redis.internal:6379" + # A Redis that requires TLS: rediss:// (rediss-cluster:// for a cluster). + # externalUrl: "rediss://:REDACTED@redis.internal:6380" + # Its certificate signed by a private CA: a Secret holding the CA's PEM + # certificate under `key`, then trusted alone for Redis (see README.md). + # caSecret: + # name: redis-ca + # key: ca.crt clickhouse: bundled: false @@ -64,6 +72,18 @@ podDisruptionBudget: networkPolicy: enabled: true + # Egress for ports no externalUrl names: Redis Cluster nodes announcing + # other ports, an upstream, MCP server or S3 endpoint off port 443. + # extraEgress: + # - ports: + # - port: 7000 + # endPort: 7005 + # protocol: TCP + +# Room for replicas that start together and set up the schema in turn: +# 5 + 3 × failureThreshold seconds (see values.yaml). +# startupProbe: +# failureThreshold: 60 topologySpreadConstraints: enabled: true @@ -77,5 +97,7 @@ resources: # SSO/OIDC is configured at runtime via Admin > Settings > SSO in the # web console (no env vars / Helm values needed). -secrets: - # jwtSecret and encryptionKey left empty → auto-generated. +# jwtSecret and encryptionKey left out → auto-generated. +# secrets: +# jwtSecret: "" +# encryptionKey: "" diff --git a/deploy/helm/think-watch/values.yaml b/deploy/helm/think-watch/values.yaml index e44ac6ae..3c577399 100644 --- a/deploy/helm/think-watch/values.yaml +++ b/deploy/helm/think-watch/values.yaml @@ -76,7 +76,9 @@ postgres: image: postgres:18-alpine user: thinkwatch database: think_watch - # Used only when bundled=false; e.g. postgres://user:pass@host:5432/db + # Used only when bundled=false; e.g. postgres://user:pass@host:5432/db. + # Postgres itself or a pooler in session mode, never transaction mode + # (see README.md, "PostgreSQL behind a connection pooler"). externalUrl: "" storage: 10Gi storageClassName: "" @@ -91,8 +93,18 @@ postgres: redis: bundled: true image: redis:8-alpine - # Used only when bundled=false; e.g. redis://:pass@host:6379 + # Used only when bundled=false; e.g. redis://:pass@host:6379, or + # rediss://:pass@host:6379 for a Redis that requires TLS (see README.md) externalUrl: "" + # For a rediss:// externalUrl whose certificate a private CA signed: + # an existing Secret in the release namespace holding that CA's PEM + # certificate(s) under `key` (cert-manager keeps it under `ca.crt`; + # required when `name` is set). + # The server then trusts only those certificates for Redis. Empty = + # the public CAs, which is what managed services need. + caSecret: + name: "" + key: ca.crt storage: 2Gi storageClassName: "" resources: @@ -108,7 +120,8 @@ clickhouse: image: clickhouse/clickhouse-server:26.3-alpine user: thinkwatch database: think_watch - # Used only when bundled=false; e.g. http://host:8123 + # Used only when bundled=false; e.g. http://host:8123. Plain HTTP only: + # the server does not connect to ClickHouse over HTTPS (https:// URLs). externalUrl: "" storage: 20Gi storageClassName: "" @@ -126,6 +139,26 @@ clickhouse: networkPolicy: enabled: false + # The server may reach PostgreSQL, Redis and ClickHouse on the ports it + # connects to them on: a bundled service's port, or every port its + # externalUrl names (each host, each `node=` of a Redis Cluster or + # Sentinel URL, a Postgres `port=`). Where the URL names no port, the + # client's default for its scheme: postgres:// 5432; redis:// and + # rediss:// 6379; a Sentinel 26379, plus 6379 for the primary it points + # to; ClickHouse's http:// 80. A database on another port, such as Azure + # Cache for Redis over TLS (rediss://...:6380), needs nothing more than + # the port in its URL. + # + # Egress rules added as written, for ports that no URL names: Redis + # Cluster nodes announcing ports the URL doesn't list, a Sentinel's + # primary on a port other than 6379, an upstream, MCP server or S3 + # endpoint on a port other than 443 (RustFS and MinIO: 9000). ClickHouse + # is reached over plain HTTP only, so its native port 9000 is not allowed. + extraEgress: [] + # - ports: + # - port: 7000 + # endPort: 7005 + # protocol: TCP autoscaling: enabled: false @@ -139,6 +172,18 @@ podDisruptionBudget: terminationGracePeriodSeconds: 30 +# The server pod's startup probe. The server answers it once it has set up +# the database schema, which instances do one at a time: with several +# replicas starting together (an install, a rollout, an autoscaler adding +# pods), each waits for those before it, and with ClickHouse unreachable +# each retries it for up to about a minute. The defaults give a pod +# 5 + 3 × 40 = 125 s before the kubelet restarts it; raise +# failureThreshold for more replicas than that covers. +startupProbe: + initialDelaySeconds: 5 + periodSeconds: 3 + failureThreshold: 40 + podAntiAffinity: enabled: true type: preferred diff --git a/web/package.json b/web/package.json index 5259ad1a..147c56c1 100644 --- a/web/package.json +++ b/web/package.json @@ -1,7 +1,7 @@ { "name": "web", "private": true, - "version": "3.1.0", + "version": "3.2.0", "type": "module", "packageManager": "pnpm@11.0.0", "scripts": {