diff --git a/Cargo.lock b/Cargo.lock index bf96756..c50f39f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -858,6 +858,7 @@ dependencies = [ name = "gw-config" version = "0.0.0" dependencies = [ + "aho-corasick", "gw-consts", "regex", "serde", diff --git a/Cargo.toml b/Cargo.toml index b066b2a..cd44121 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -92,6 +92,7 @@ hmac = "0.12" hex = "0.4" chacha20poly1305 = "0.10" regex = "1" +aho-corasick = "1" # AWS EventStream framing (already in the lock via sqlx/reqwest) crc = "3" percent-encoding = "2" diff --git a/control-plane/go.mod b/control-plane/go.mod index b876530..4ea24b4 100644 --- a/control-plane/go.mod +++ b/control-plane/go.mod @@ -32,7 +32,7 @@ require ( golang.org/x/sys v0.47.0 // indirect golang.org/x/text v0.41.0 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260825221802-da73d73af1c5 // indirect - google.golang.org/grpc v1.83.1 // indirect + google.golang.org/grpc v1.83.2 // indirect google.golang.org/protobuf v1.36.12 // indirect gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect ) diff --git a/control-plane/go.sum b/control-plane/go.sum index ef93424..0114d2c 100644 --- a/control-plane/go.sum +++ b/control-plane/go.sum @@ -113,8 +113,8 @@ golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8T golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= google.golang.org/genproto/googleapis/rpc v0.0.0-20260825221802-da73d73af1c5 h1:1VUiZAXyC+zmiFYi+WLtBzr68Cj8wOofHjjrA/kkizc= google.golang.org/genproto/googleapis/rpc v0.0.0-20260825221802-da73d73af1c5/go.mod h1:DjtHYE8FKJLivXcBEjGwndXfIC23G0VpXiXKqG179uA= -google.golang.org/grpc v1.83.1 h1:HIO0+BEtBP6soyqvqC8sNUjZ7bTs+0hFQuFF+RAy++Y= -google.golang.org/grpc v1.83.1/go.mod h1:kDyl6SKsiHKt0uylY5gtn5cEjkrIOhQOGDgIc4JGwzQ= +google.golang.org/grpc v1.83.2 h1:EManeRomTObA0BU7I8vXgg/78uE5MJ9M8B39EX2WscU= +google.golang.org/grpc v1.83.2/go.mod h1:YPI1hK3kDked6iHvgX3tR0y+nX/qpMFKhPgFsokw1S8= google.golang.org/protobuf v1.36.12 h1:pJOKDDOyeXErUroCihFAd5LQuwXBSpVnKGrj5o/fwxc= google.golang.org/protobuf v1.36.12/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= diff --git a/control-plane/web/package-lock.json b/control-plane/web/package-lock.json index caa00fe..e2e11fa 100644 --- a/control-plane/web/package-lock.json +++ b/control-plane/web/package-lock.json @@ -326,9 +326,9 @@ } }, "node_modules/@jridgewell/sourcemap-codec": { - "version": "1.5.5", - "resolved": "https://registry.npmjs.org/@jridgewell/sourcemap-codec/-/sourcemap-codec-1.5.5.tgz", - "integrity": "sha512-cYQ9310grqxueWbl+WuIUIaiUaDcj7WOq5fVhEljNVgRfOUhY9fy2zTvfoqWsnebh8Sl70VScFbICvJnLKB0Og==", + "version": "1.6.0", + "resolved": "https://registry.npmjs.org/@jridgewell/sourcemap-codec/-/sourcemap-codec-1.6.0.tgz", + "integrity": "sha512-T7jf+5zgsZHwNJ4lvQ7/aezbyk0nNX+zJVWpmHA7VYsEx7a7qr5Rg5IbtJFqkgze5Y2sruq1RUY8Q837Od7iFw==", "dev": true, "license": "MIT" }, @@ -1173,16 +1173,16 @@ } }, "node_modules/@vitest/expect": { - "version": "4.1.10", - "resolved": "https://registry.npmjs.org/@vitest/expect/-/expect-4.1.10.tgz", - "integrity": "sha512-YsCn+qAk1GWjQOWFEsEcL2gNQ0zmVmQu3T03qP6UyjhtmdtwtbuI+DASn/7iQB3HGTXkdBwGddzxPlmiql5vlA==", + "version": "4.1.11", + "resolved": "https://registry.npmjs.org/@vitest/expect/-/expect-4.1.11.tgz", + "integrity": "sha512-VX2x5vNJXET47KAFzwERI+KRMtTTCSWTfSMKsW7JsUsXV4psq++e3DvZpuTDOpHcxytiDs6p2nhVb2tVDiiUYw==", "dev": true, "license": "MIT", "dependencies": { "@standard-schema/spec": "^1.1.0", "@types/chai": "^5.2.2", - "@vitest/spy": "4.1.10", - "@vitest/utils": "4.1.10", + "@vitest/spy": "4.1.11", + "@vitest/utils": "4.1.11", "chai": "^6.2.2", "tinyrainbow": "^3.1.0" }, @@ -1191,13 +1191,13 @@ } }, "node_modules/@vitest/mocker": { - "version": "4.1.10", - "resolved": "https://registry.npmjs.org/@vitest/mocker/-/mocker-4.1.10.tgz", - "integrity": "sha512-v0xaezt+DKEmKfaxg133ldzADrwLGd7Ze1MfQQTYfvs8OqZIwbxyxaYURivwV7sWy5fqn3rH5uOrSp07bp44Ow==", + "version": "4.1.11", + "resolved": "https://registry.npmjs.org/@vitest/mocker/-/mocker-4.1.11.tgz", + "integrity": "sha512-2XJVD55d1o5AZous5CCGKS74g/riOj9odEt2bQpCVZeblHyHdnMeFl4jl0XjU21stf4mbjUkew2eXQZt65g5CQ==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/spy": "4.1.10", + "@vitest/spy": "4.1.11", "estree-walker": "^3.0.3", "magic-string": "^0.30.21" }, @@ -1218,9 +1218,9 @@ } }, "node_modules/@vitest/pretty-format": { - "version": "4.1.10", - "resolved": "https://registry.npmjs.org/@vitest/pretty-format/-/pretty-format-4.1.10.tgz", - "integrity": "sha512-W1HsjSH4MXQ9YfmmhLAoIYf1HRfekQCGngeIgcei6MP5QQGWUe0gkopdZQaVCFO+JDJMrAJGwa5pRpNpvy4P8Q==", + "version": "4.1.11", + "resolved": "https://registry.npmjs.org/@vitest/pretty-format/-/pretty-format-4.1.11.tgz", + "integrity": "sha512-yiZzPbGTS9Sr/JpFl8zHrcIkAofNbFV6k21vIgQN/cY/oxZeXhJv5sc/MBJ5jFKWmWs+oJHw0UXLZjmf931+Vw==", "dev": true, "license": "MIT", "dependencies": { @@ -1231,13 +1231,13 @@ } }, "node_modules/@vitest/runner": { - "version": "4.1.10", - "resolved": "https://registry.npmjs.org/@vitest/runner/-/runner-4.1.10.tgz", - "integrity": "sha512-IKI6kpIH+LmpROplyLwBBaCfMgOZOMsygVa6BARD6ahA04VRuJSa6OaVG7kRvSEMD870Vd91rSSw0eegtWyLGg==", + "version": "4.1.11", + "resolved": "https://registry.npmjs.org/@vitest/runner/-/runner-4.1.11.tgz", + "integrity": "sha512-LztvUgdwMNJMIkj3hQnnxiC2Xy1zNxq928W/xhjCLaNCzqTZOudjwbQf6v9IntZGPw132i2Lq2rgTRZHD3JHNw==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/utils": "4.1.10", + "@vitest/utils": "4.1.11", "pathe": "^2.0.3" }, "funding": { @@ -1245,14 +1245,14 @@ } }, "node_modules/@vitest/snapshot": { - "version": "4.1.10", - "resolved": "https://registry.npmjs.org/@vitest/snapshot/-/snapshot-4.1.10.tgz", - "integrity": "sha512-xRkfOT1qpTAi/Ti4Y1LtfRc3kEuqxGw59eN2jN9pRWMtS/XDevekhcFSqvQqjUNGksfjMJu3Y+oJ+4Ypn2OaJw==", + "version": "4.1.11", + "resolved": "https://registry.npmjs.org/@vitest/snapshot/-/snapshot-4.1.11.tgz", + "integrity": "sha512-pN7ikn1ON7h8ee4gIAp4AzyK+zBtJPzVbqOgu5LCEh4VaJVbPQcgYQYJIMGQPXVeJJq1fnfazis7a5pFNPahog==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/pretty-format": "4.1.10", - "@vitest/utils": "4.1.10", + "@vitest/pretty-format": "4.1.11", + "@vitest/utils": "4.1.11", "magic-string": "^0.30.21", "pathe": "^2.0.3" }, @@ -1261,9 +1261,9 @@ } }, "node_modules/@vitest/spy": { - "version": "4.1.10", - "resolved": "https://registry.npmjs.org/@vitest/spy/-/spy-4.1.10.tgz", - "integrity": "sha512-PLf/Ugvoq5wO/b4rwYCR1h2PSIdXz7wnkQFMiUpLdtM7l6pqVFcQIBEHyT1+l+cj7mNwAfZHzqXqDyjvOuwbDw==", + "version": "4.1.11", + "resolved": "https://registry.npmjs.org/@vitest/spy/-/spy-4.1.11.tgz", + "integrity": "sha512-apNa/prQy2qCeywhnixOHPRCgGNhvg7T4Dapfl1GahLp/R+uhBm5cPyFoNVyqsNd2h1nJxL6BqqdIjiABL60YA==", "dev": true, "license": "MIT", "funding": { @@ -1271,13 +1271,13 @@ } }, "node_modules/@vitest/utils": { - "version": "4.1.10", - "resolved": "https://registry.npmjs.org/@vitest/utils/-/utils-4.1.10.tgz", - "integrity": "sha512-fy9am/HWxbaGt/Sawrp90vt6Y6jQwf1RX77cz3uwoJwJVMli/e1IEwRPnMNJ7vKfPTwo0diXifkpPvwH9v7nGA==", + "version": "4.1.11", + "resolved": "https://registry.npmjs.org/@vitest/utils/-/utils-4.1.11.tgz", + "integrity": "sha512-zTCVGpyFsGWBhllOyKlTw/vnr6D9qxsfSDyfbyZmTyjHw5N/VuvzHpHoQjm2ZJzn4RJgx5w4r7V0er69CmLgPQ==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/pretty-format": "4.1.10", + "@vitest/pretty-format": "4.1.11", "convert-source-map": "^2.0.0", "tinyrainbow": "^3.1.0" }, @@ -2316,9 +2316,9 @@ } }, "node_modules/tinyrainbow": { - "version": "3.1.0", - "resolved": "https://registry.npmjs.org/tinyrainbow/-/tinyrainbow-3.1.0.tgz", - "integrity": "sha512-Bf+ILmBgretUrdJxzXM0SgXLZ3XfiaUuOj/IKQHuTXip+05Xn+uyEYdVg0kYDipTBcLrCVyUzAPz7QmArb0mmw==", + "version": "3.1.1", + "resolved": "https://registry.npmjs.org/tinyrainbow/-/tinyrainbow-3.1.1.tgz", + "integrity": "sha512-yau8yJdTt989Mm0Bd/236QnzEiPf2xLLTqUZRUJOo/3CB078LSwzei343DgtJVmfJKJE3TMINY1u42SQsP6mXw==", "dev": true, "license": "MIT", "engines": { @@ -2518,19 +2518,19 @@ } }, "node_modules/vitest": { - "version": "4.1.10", - "resolved": "https://registry.npmjs.org/vitest/-/vitest-4.1.10.tgz", - "integrity": "sha512-R9jUTe5S4Qb0HCd4TNqpC7oGcrMssMRGXLW80ubjWsW9VH5GF8y1Y0SFLY9AbqSk6nt0PnOx4H4WNJYZ13GUPw==", + "version": "4.1.11", + "resolved": "https://registry.npmjs.org/vitest/-/vitest-4.1.11.tgz", + "integrity": "sha512-fhACrNXUidIbGSBr5FlbuBkO7VWC1ZyLl0DO4CU2DrQoAPxX84Ysxs+HeGQpii5lZWV1Q4gBZTTu49mF+A6Edw==", "dev": true, "license": "MIT", "dependencies": { - "@vitest/expect": "4.1.10", - "@vitest/mocker": "4.1.10", - "@vitest/pretty-format": "4.1.10", - "@vitest/runner": "4.1.10", - "@vitest/snapshot": "4.1.10", - "@vitest/spy": "4.1.10", - "@vitest/utils": "4.1.10", + "@vitest/expect": "4.1.11", + "@vitest/mocker": "4.1.11", + "@vitest/pretty-format": "4.1.11", + "@vitest/runner": "4.1.11", + "@vitest/snapshot": "4.1.11", + "@vitest/spy": "4.1.11", + "@vitest/utils": "4.1.11", "es-module-lexer": "^2.0.0", "expect-type": "^1.3.0", "magic-string": "^0.30.21", @@ -2558,12 +2558,12 @@ "@edge-runtime/vm": "*", "@opentelemetry/api": "^1.9.0", "@types/node": "^20.0.0 || ^22.0.0 || >=24.0.0", - "@vitest/browser-playwright": "4.1.10", - "@vitest/browser-preview": "4.1.10", - "@vitest/browser-webdriverio": "4.1.10", - "@vitest/coverage-istanbul": "4.1.10", - "@vitest/coverage-v8": "4.1.10", - "@vitest/ui": "4.1.10", + "@vitest/browser-playwright": "4.1.11", + "@vitest/browser-preview": "4.1.11", + "@vitest/browser-webdriverio": "4.1.11", + "@vitest/coverage-istanbul": "4.1.11", + "@vitest/coverage-v8": "4.1.11", + "@vitest/ui": "4.1.11", "happy-dom": "*", "jsdom": "*", "vite": "^6.0.0 || ^7.0.0 || ^8.0.0" diff --git a/crates/config/Cargo.toml b/crates/config/Cargo.toml index f3aa6dd..a992634 100644 --- a/crates/config/Cargo.toml +++ b/crates/config/Cargo.toml @@ -13,6 +13,7 @@ serde-saphyr = { workspace = true } sha2 = { workspace = true } thiserror = { workspace = true } regex = { workspace = true } +aho-corasick = { workspace = true } tracing = { workspace = true } [lints] diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 1177a71..a0c4e99 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -229,11 +229,13 @@ pub struct ModelConf { /// capacity. #[serde(default)] pub variants: Vec, + #[serde(skip)] + kind: Option, } impl ModelConf { pub fn protocol(&self) -> Option { - Protocol::from_wire(&self.protocol) + self.kind } } @@ -388,6 +390,9 @@ pub struct SecurityConf { /// Blocklist terms; normalized to lower-case (empties dropped) at load. #[serde(default)] pub blocklist: Vec, + /// The blocklist as one automaton; `None` when the list is empty. + #[serde(skip)] + pub blocklist_matcher: Option, /// What a blocklist hit does (default: block). #[serde(default)] pub blocklist_action: Action, @@ -413,6 +418,12 @@ pub struct SecurityConf { } impl SecurityConf { + /// The policy as load leaves it: blocklist lower-cased and compiled, regexes built. + pub fn compiled(mut self) -> Self { + compile_security(&mut self); + self + } + /// Whether responses must be redacted before leaving — the one predicate /// the outbound-DLP masking AND the stream-buffering boundary share, so a /// secrets-only tenant can't stream raw deltas past the masking. @@ -764,6 +775,8 @@ pub struct GatewayConfig { product_idx: HashMap, #[serde(skip)] tenant_idx: HashMap, + #[serde(skip)] + account_idx: HashMap, } impl GatewayConfig { @@ -781,6 +794,7 @@ impl GatewayConfig { self.model_idx = index_by(&self.models, |m| &m.name); self.product_idx = index_by(&self.products, |p| &p.name); self.tenant_idx = index_by(&self.tenants, |t| &t.name); + self.account_idx = index_by(&self.accounts, |a| &a.name); } /// Expand provider presets: fill each model's default wire type and @@ -870,6 +884,9 @@ impl GatewayConfig { compile_security(sec); } } + for m in &mut self.models { + m.kind = Protocol::from_wire(&m.protocol); + } Ok(()) } @@ -1278,6 +1295,10 @@ impl GatewayConfig { self.models.get(*self.model_idx.get(name)?) } + pub fn find_account(&self, name: &str) -> Option<&AccountConf> { + self.accounts.get(*self.account_idx.get(name)?) + } + /// Pricing for a public model name; zero if unlisted. fn prices_for(&self, name: &str) -> (i64, i64) { self.find_model(name) @@ -1344,7 +1365,7 @@ fn provider_preset(kind: &str) -> Option { wires: &["gemini", "realtime"], default_model_wire: "gemini", }, - // OpenAI-protocol vendors: same wire shape, different base URL. + // vendors on the OpenAI protocol share the wire shape and differ by base URL "deepseek" => ProviderPreset { endpoint: "https://api.deepseek.com", wires: &["openai-chat"], @@ -1469,6 +1490,19 @@ fn compile_security(sec: &mut SecurityConf) { .filter(|w| !w.is_empty()) .map(|w| w.to_lowercase()) .collect(); + sec.blocklist_matcher = (!sec.blocklist.is_empty()) + .then(|| { + aho_corasick::AhoCorasickBuilder::new() + .ascii_case_insensitive(true) + .build(&sec.blocklist) + }) + .and_then(|built| match built { + Ok(matcher) => Some(matcher), + Err(e) => { + tracing::error!(error = %e, "blocklist did not compile; the list is ignored"); + None + } + }); sec.regexes = sec .regex_rules .iter() diff --git a/crates/consts/src/error_class.rs b/crates/consts/src/error_class.rs index 1b8f8df..d880fdc 100644 --- a/crates/consts/src/error_class.rs +++ b/crates/consts/src/error_class.rs @@ -4,8 +4,7 @@ use crate::ErrCode; -/// One classification from the contract's closed set. `ModelStreamError` is -/// in-stream only and never renders at the HTTP phase. +/// One classification from the contract's closed set. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ErrClass { Validation, @@ -73,7 +72,7 @@ impl ErrClass { ErrClass::ModelTimeout => 408, ErrClass::Conflict => 409, ErrClass::RequestEntityTooLarge => 413, - // ModelStreamError never renders at the HTTP phase; 424 nominal + // the in-stream ModelStreamError never renders at the HTTP phase; 424 nominal ErrClass::ModelError | ErrClass::ModelStreamError => 424, ErrClass::Throttling => 429, ErrClass::InternalServer => 500, diff --git a/crates/dag/src/context.rs b/crates/dag/src/context.rs index 2dfb2b8..46f93dd 100644 --- a/crates/dag/src/context.rs +++ b/crates/dag/src/context.rs @@ -7,7 +7,7 @@ use std::sync::Arc; use gw_config::GatewayConfig; use gw_engines::{EngineOutcome, SharedTransport}; -use gw_models::GatewayRequest; +use gw_models::{GResult, GatewayError, GatewayRequest, ModelParamV2}; use gw_state::{AkInfo, GatewayState}; pub struct DagContext { @@ -26,6 +26,9 @@ pub struct DagContext { /// Request-level cache hit (downstream nodes short-circuit on this and skip /// account/engine/billing). pub cache_hit: bool, + /// The handler will retry an upstream fault on a fallback model, so this + /// attempt's failure is not the client-visible outcome. + pub fallback_ahead: bool, /// This request's cache key (computed by cache_lookup, reused by cache_store). pub cache_key: Option, /// Governance key for the (AK, model) daily counter — set by model_quota @@ -62,6 +65,7 @@ impl DagContext { outcome: None, decisions: Vec::new(), cache_hit: false, + fallback_ahead: false, cache_key: None, model_quota_key: None, quota_reserved: None, @@ -78,6 +82,14 @@ impl DagContext { /// The effective end user: the key's `owner` (authoritative) else request /// metadata; `""` when neither is present. Resolution lives on [`AkInfo`] so /// REST and realtime can't diverge on an empty owner. + /// The resolved model param; a node reaching here before resolve_model is a broken plan. + pub fn model_param(&self) -> GResult<&ModelParamV2> { + self.request + .model_param_v2 + .as_ref() + .ok_or_else(|| GatewayError::internal("model param missing after resolve_model")) + } + pub fn effective_user_id(&self) -> &str { self.ak .attributed_user(self.request.user_id.as_deref().unwrap_or_default()) diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 8aec472..e12d5f2 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -123,12 +123,7 @@ impl DagNode for TenantEntitlement { "tenant_entitlement" } async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { - let name = ctx - .request - .model_param_v2 - .as_ref() - .map(|p| p.model_name.as_str()) - .unwrap_or_default(); + let name = &ctx.model_param()?.model_name; if !ctx.cfg.tenant_allows_model(&ctx.ak.tenant, name) { return Err(GatewayError::new( ErrCode::PERMISSION_CHECK, @@ -154,9 +149,7 @@ impl DagNode for VariantSelect { "variant_select" } async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { - let Some(param) = ctx.request.model_param_v2.as_ref() else { - return Ok(()); - }; + let param = ctx.model_param()?; if param.fallback_from.is_some() { return Ok(()); } @@ -197,9 +190,7 @@ impl DagNode for CacheLookup { "cache_lookup" } async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { - let Some(param) = ctx.request.model_param_v2.as_ref() else { - return Ok(()); - }; + let param = ctx.model_param()?; // batch items bypass: a free (unbilled) hit would break their per-item billing if !ctx.request.buffered_online() { return Ok(()); @@ -318,9 +309,7 @@ impl DagNode for SelectAccount { .await; let Some(account) = account else { // unsampled, an exhausted pool would read no_data forever - ctx.state - .avail - .record(requested_model(ctx.request.model_param_v2.as_ref()), false); + note_unavailable(ctx); return Err(GatewayError::new( ErrCode::SYSTEM_ERROR, 503, @@ -388,12 +377,7 @@ impl DagNode for ModelQpmLimit { "model_qpm" } async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { - // fail loud: silently skipping would waive the limit on a broken plan - let param = ctx - .request - .model_param_v2 - .as_ref() - .ok_or_else(|| GatewayError::internal("model_qpm before resolve_model"))?; + let param = ctx.model_param()?; admission::check_model_qpm(ctx.state.governance.as_ref(), &ctx.cfg, ¶m.model_name) .await .map_err(limit_denied) @@ -496,9 +480,7 @@ impl DagNode for CallEngine { ) .await; let Some(next) = next else { - ctx.state - .avail - .record(requested_model(ctx.request.model_param_v2.as_ref()), false); + note_unavailable(ctx); return Err(named(first_err, ctx)); }; let spillover = failed.is_ptu() && !next.is_ptu(); @@ -520,9 +502,7 @@ impl DagNode for CallEngine { Ok(()) } Err(e) => { - ctx.state - .avail - .record(requested_model(ctx.request.model_param_v2.as_ref()), false); + note_unavailable(ctx); note_failure(ctx, &next.name).await; Err(named(e, ctx)) } @@ -541,6 +521,15 @@ fn latency_clock(ctx: &DagContext) -> Option { .then(std::time::Instant::now) } +/// A failed attempt is unavailable unless the handler still has a fallback model to try. +fn note_unavailable(ctx: &DagContext) { + if !ctx.fallback_ahead { + ctx.state + .avail + .record(requested_model(ctx.request.model_param_v2.as_ref()), false); + } +} + async fn note_engine_outcome( ctx: &mut DagContext, outcome: &gw_engines::EngineOutcome, @@ -702,8 +691,15 @@ impl BillTokens { /// paths without a usage payload (estimates, malformed usage) must price /// identically to the happy path or a cut stream changes effective pricing. fn weighted(prompt: i64, completion: i64, rate: &gw_models::TokenRate) -> Self { - let (billable_prompt, billable_completion) = - gw_models::weighted_pair(prompt, completion, rate); + let input = gw_models::TokenInput { + prompt, + completion, + ..Default::default() + }; + let (billable_prompt, billable_completion) = ( + gw_models::weighted_prompt(&input, rate), + gw_models::weighted_completion(&input, rate), + ); Self { prompt, completion, @@ -759,6 +755,7 @@ pub async fn settle_deferred_stream(ctx: &mut DagContext, delivery: StreamDelive return Ok(()); } ctx.billing_deferred = false; + // a request blocked before quota_check reserved nothing and bills nothing if ctx.quota_reserved.is_none() && ctx.tpm_reserved.is_none() { return Ok(()); } @@ -837,7 +834,7 @@ async fn bill(ctx: &mut DagContext, mut tokens: BillTokens, estimated: bool) -> ctx.tpm_reserved.take(), ctx.model_quota_key.take(), ); - let record = admission::settle_and_bill( + let settled = admission::settle_and_bill( ctx.state.as_ref(), &ctx.cfg, admission::SettleInput { @@ -875,15 +872,15 @@ async fn bill(ctx: &mut DagContext, mut tokens: BillTokens, estimated: bool) -> &ctx.cfg, &ctx.ak, ctx.effective_user_id(), - record.total_tokens, - record.cost_micros, + settled.total_tokens, + settled.cost_micros, ) .await; ctx.decide( "cost_calc", format!( "tokens={} cost_micros={}", - record.total_tokens, record.cost_micros + settled.total_tokens, settled.cost_micros ), ); Ok(()) diff --git a/crates/dag/src/token_estimate.rs b/crates/dag/src/token_estimate.rs index 8325063..538429b 100644 --- a/crates/dag/src/token_estimate.rs +++ b/crates/dag/src/token_estimate.rs @@ -12,13 +12,13 @@ pub trait TokenEncoder: Send + Sync { } /// Real tiktoken `cl100k_base` BPE — the tokenizer OpenAI's models use. -pub struct TiktokenEncoder { +struct TiktokenEncoder { bpe: tiktoken_rs::CoreBPE, } impl TiktokenEncoder { /// Fails only if the embedded vocabulary fails to load. - pub fn new() -> Result { + fn new() -> Result { let bpe = tiktoken_rs::cl100k_base().map_err(|e| format!("load cl100k_base: {e}"))?; Ok(Self { bpe }) } @@ -33,8 +33,7 @@ impl TokenEncoder for TiktokenEncoder { /// Approximation of cl100k_base counting (NOT tiktoken): ASCII letters ~1 token /// per 4 chars, digits per 3, punctuation 1 each, non-ASCII 1 per char; /// whitespace folds into the following word. -#[derive(Debug, Default, Clone, Copy)] -pub struct HeuristicEncoder; +struct HeuristicEncoder; impl HeuristicEncoder { const LETTERS_PER_TOKEN: usize = 4; @@ -128,7 +127,6 @@ pub fn estimate_prompt_tokens( if let Some(id) = &msg.tool_call_id { num += enc.encode_len(id); } - // assistant tool_calls: each call adds overhead (+3) plus encoded name and args. if let Some(Value::Array(calls)) = &msg.tool_calls { for call in calls { num += 3; @@ -185,9 +183,10 @@ fn message_text(msg: &ChatMsg) -> std::borrow::Cow<'_, str> { #[cfg(test)] mod tests { - use super::*; use serde_json::json; + use super::*; + #[test] fn heuristic_classifies_runs() { let e = HeuristicEncoder; diff --git a/crates/engines/src/bedrock.rs b/crates/engines/src/bedrock.rs index d7aba68..171baf3 100644 --- a/crates/engines/src/bedrock.rs +++ b/crates/engines/src/bedrock.rs @@ -52,7 +52,7 @@ pub(crate) fn aws_headers( ("host", host.into()), ("x-amz-date", amz_date.into()), ("authorization", authorization), - // InvokeModel requires accept; content-type is unsigned and added by the caller + // the InvokeModel call requires accept; content-type is unsigned and added by the caller ("accept", "application/json".into()), ] } @@ -174,8 +174,8 @@ pub(crate) async fn bedrock_stream( where F: FnMut(Value) -> GResult>, { - let model = base.model_name()?.to_owned(); - let reply = bedrock_send(base, &model, body).await?; + let uri = invoke_uri(base.model_name()?, base.request.stream); + let reply = bedrock_send_uri(base, &uri, body).await?; let status = reply.status; crate::pump::reject_json_error("bedrock", status, &reply.body)?; let r = diff --git a/crates/engines/src/bespoke.rs b/crates/engines/src/bespoke.rs index f1f2c85..c53e9c6 100644 --- a/crates/engines/src/bespoke.rs +++ b/crates/engines/src/bespoke.rs @@ -306,7 +306,7 @@ impl DashScopeEngine { fn headers(&self, stream: bool) -> Headers { let mut h = self.base.bearer_headers(); if stream { - // DashScope streams only when this header is present + // streaming needs this DashScope header h.push(("X-DashScope-SSE", "enable".into())); } h @@ -331,7 +331,7 @@ impl DashScopeEngine { "dashscope", reply.body, self.base.request.stream_tx.clone(), - |v| dashscope_apply_frame(&v, status, &mut resp, &mut full), + |v| dashscope_apply_frame(v, status, &mut resp, &mut full), ) .await?; resp.message = full; @@ -372,25 +372,25 @@ impl ModelEngine for DashScopeEngine { /// Apply one DashScope SSE frame: running frames carry the literal "null" /// finish_reason and cumulative usage (last frame wins). fn dashscope_apply_frame( - v: &Value, + mut v: Value, status: u16, resp: &mut GatewayResponse, full: &mut String, ) -> GResult> { - if let Some(err) = crate::engine::vendor_error(status, v) { + if let Some(err) = crate::engine::vendor_error(status, &v) { return Err(err); } let mut chunks = Vec::new(); - let choice = &v["output"]["choices"][0]; - if let Some(t) = choice["message"]["content"].as_str() + if let Some(t) = crate::engine::take_string(&mut v, "/output/choices/0/message/content") && !t.is_empty() { - full.push_str(t); + full.push_str(&t); chunks.push(StreamChunk { - delta: t.to_owned(), + delta: t, ..Default::default() }); } + let choice = &v["output"]["choices"][0]; if let Some(fr) = choice["finish_reason"].as_str() && !fr.is_empty() && fr != "null" diff --git a/crates/engines/src/claude_engine.rs b/crates/engines/src/claude_engine.rs index 7b9d1d1..4a0e2d4 100644 --- a/crates/engines/src/claude_engine.rs +++ b/crates/engines/src/claude_engine.rs @@ -44,7 +44,7 @@ impl ClaudeEngine { } else { "user" }; - // Claude wants replayed thinking blocks ahead of the tool_use they produced + // replayed thinking blocks go ahead of the tool_use they produced, as Claude expects let thinking: Vec = match m.reasoning_details { Some(Value::Array(details)) => details .into_iter() @@ -142,7 +142,7 @@ impl ClaudeEngine { if let Some(tc) = p.tool_choice { body.insert("tool_choice".into(), normalize_tool_choice_anthropic(tc)); } - // Anthropic takes `stop_sequences` (array); OpenAI's `stop` may be a string + // `stop_sequences` is an array here; OpenAI's `stop` may be a string if let Some(stop) = p.stop { let stop = match stop { Value::String(_) => Value::Array(vec![stop]), @@ -186,7 +186,7 @@ impl ClaudeEngine { let mut headers = vec![ ("content-type", "application/json".into()), ("x-api-key", self.base.api_key()), - // Anthropic API mandates this header; a real call 400s without it. + // the API mandates this header; a real call 400s without it ("anthropic-version", "2023-06-01".into()), ]; if let Some(betas) = self.base.request.anthropic_beta.take() { @@ -267,7 +267,13 @@ impl ClaudeEngine { }; let mut outcome = EngineOutcome::with_status(resp, status); if self.base.request.stream && self.base.request.preserve_anthropic_wire { - outcome.chunks = anthropic_native_chunks(&outcome.response, self.base.model_override()); + let content = outcome + .response + .anthropic_content + .clone() + .unwrap_or(Value::Null); + outcome.chunks = + anthropic_native_chunks(&outcome.response, content, self.base.model_override()); } Ok(outcome) } @@ -591,6 +597,7 @@ impl<'a> SseState<'a> { /// upstreams ignore `stream:true`), so thinking proof survives the streaming surface. pub fn anthropic_native_chunks( response: &GatewayResponse, + content: Value, model_override: Option<&str>, ) -> Vec { let stream_model = model_override.unwrap_or(&response.model); @@ -612,53 +619,52 @@ pub fn anthropic_native_chunks( "usage":start_usage } })]; - if let Some(blocks) = response - .anthropic_content - .as_ref() - .and_then(Value::as_array) - { - for (index, block) in blocks.iter().enumerate() { - let mut start = block.clone(); + if let Value::Array(blocks) = content { + for (index, mut start) in blocks.into_iter().enumerate() { let mut deltas = Vec::new(); - match block.get("type").and_then(Value::as_str) { - Some("thinking") => { - if let Some(object) = start.as_object_mut() { - object.insert("thinking".to_owned(), "".into()); - object.insert("signature".to_owned(), "".into()); - } - if let Some(thinking) = block.get("thinking").and_then(Value::as_str) - && !thinking.is_empty() - { - deltas.push(json!({"type":"thinking_delta","thinking":thinking})); - } - if let Some(signature) = block.get("signature").and_then(Value::as_str) - && !signature.is_empty() - { - deltas.push(json!({"type":"signature_delta","signature":signature})); - } + if start["type"] == "thinking" { + if let Some(Value::String(thinking)) = start.get_mut("thinking").map(Value::take) + && !thinking.is_empty() + { + deltas.push(object([ + ("type", "thinking_delta".into()), + ("thinking", Value::String(thinking)), + ])); } - Some("text") => { - if let Some(object) = start.as_object_mut() { - object.insert("text".to_owned(), "".into()); - } - if let Some(text) = block.get("text").and_then(Value::as_str) - && !text.is_empty() - { - deltas.push(json!({"type":"text_delta","text":text})); - } + if let Some(Value::String(signature)) = start.get_mut("signature").map(Value::take) + && !signature.is_empty() + { + deltas.push(object([ + ("type", "signature_delta".into()), + ("signature", Value::String(signature)), + ])); } - Some("tool_use") => { - if let Some(object) = start.as_object_mut() { - object.insert("input".to_owned(), json!({})); - } - if let Some(input) = block.get("input") { - deltas.push(object([ - ("type", "input_json_delta".into()), - ("partial_json", input.to_string().into()), - ])); - } + if let Some(object) = start.as_object_mut() { + object.insert("thinking".to_owned(), "".into()); + object.insert("signature".to_owned(), "".into()); + } + } else if start["type"] == "text" { + if let Some(Value::String(text)) = start.get_mut("text").map(Value::take) + && !text.is_empty() + { + deltas.push(object([ + ("type", "text_delta".into()), + ("text", Value::String(text)), + ])); + } + if let Some(object) = start.as_object_mut() { + object.insert("text".to_owned(), "".into()); + } + } else if start["type"] == "tool_use" { + if let Some(input) = start.get_mut("input").map(Value::take) { + deltas.push(object([ + ("type", "input_json_delta".into()), + ("partial_json", input.to_string().into()), + ])); + } + if let Some(object) = start.as_object_mut() { + object.insert("input".to_owned(), json!({})); } - _ => {} } // built by hand: json! would deep-copy the moved blocks let mut event = Map::with_capacity(3); @@ -1309,7 +1315,8 @@ mod tests { anthropic_content: Some(json!([{"type":"text","text":"answer"}])), ..Default::default() }; - let chunks = anthropic_native_chunks(&response, None); + let content = response.anthropic_content.clone().unwrap_or(Value::Null); + let chunks = anthropic_native_chunks(&response, content, None); let events: Vec<_> = chunks .iter() .filter_map(|chunk| chunk.native_event.as_ref()) diff --git a/crates/engines/src/converse.rs b/crates/engines/src/converse.rs index ea981e0..31c299c 100644 --- a/crates/engines/src/converse.rs +++ b/crates/engines/src/converse.rs @@ -61,7 +61,9 @@ impl Events { json!({"type": "text", "text": ""}), typed("text_delta", "text", text.take()), ) - } else if let Some(input) = delta["toolUse"].get_mut("input") { + } else if let Some(input) = + delta.get_mut("toolUse").and_then(|t| t.get_mut("input")) + { ( Value::Null, typed("input_json_delta", "partial_json", input.take()), @@ -198,8 +200,13 @@ pub(crate) fn request(mut body: Map, claude: bool) -> Value { /// A buffered Converse reply as a Messages reply. pub(crate) fn reply(mut v: Value, model: &str) -> Value { - let content: Vec = match v["output"]["message"]["content"].take() { - Value::Array(blocks) => blocks.into_iter().filter_map(anthropic_block).collect(), + let content: Vec = match v + .get_mut("output") + .and_then(|o| o.get_mut("message")) + .and_then(|m| m.get_mut("content")) + .map(Value::take) + { + Some(Value::Array(blocks)) => blocks.into_iter().filter_map(anthropic_block).collect(), _ => Vec::new(), }; let usage = usage(&mut v["usage"]); @@ -254,7 +261,13 @@ fn content_block(mut block: Value) -> Vec { .unwrap_or("png"); let image = object([ ("format", format.into()), - ("source", object([("bytes", source["data"].take())])), + ( + "source", + object([( + "bytes", + source.get_mut("data").map(Value::take).unwrap_or_default(), + )]), + ), ]); object([("image", image)]) } @@ -313,7 +326,10 @@ fn content_block(mut block: Value) -> Vec { fn tool_spec(mut tool: Value, claude: bool) -> Vec { let cache_control = tool.get_mut("cache_control").map(Value::take); let mut spec = Map::with_capacity(4); - spec.insert("name".into(), tool["name"].take()); + spec.insert( + "name".into(), + tool.get_mut("name").map(Value::take).unwrap_or_default(), + ); if let Some(d) = tool.get_mut("description").filter(|d| !d.is_null()) { spec.insert("description".into(), d.take()); } diff --git a/crates/engines/src/families.rs b/crates/engines/src/families.rs index cbb99a8..fabad09 100644 --- a/crates/engines/src/families.rs +++ b/crates/engines/src/families.rs @@ -18,30 +18,29 @@ use crate::transport::{Headers, SharedTransport, Transport, UpstreamBody, Upstre /// Gemini `parts` from a unified message: text and data-URI images (`inlineData`); /// remote image URLs cannot be inlined without a fetch and are skipped. -fn gemini_parts(m: &gw_models::ChatMsg) -> Vec { - if let Some(Value::Array(parts)) = &m.parts { +fn gemini_parts(mut m: gw_models::ChatMsg) -> Vec { + if let Some(Value::Array(parts)) = m.parts.take() { let mut out = Vec::new(); - for p in parts { - match p["type"].as_str() { - Some("text") => { - if let Some(t) = p["text"].as_str() { - out.push(json!({"text": t})); - } - } - Some("image_url") => { - let url = p["image_url"]["url"].as_str().unwrap_or_default(); - if let Some((mime, data)) = parse_data_uri(url) { - out.push(json!({"inlineData": {"mimeType": mime, "data": data}})); - } + for mut p in parts { + if p["type"] == "text" { + if let Some(Value::String(t)) = p.get_mut("text").map(Value::take) { + out.push(object([("text", Value::String(t))])); } - _ => {} + } else if p["type"] == "image_url" + && let Some(Value::String(url)) = p + .get_mut("image_url") + .and_then(|u| u.get_mut("url")) + .map(Value::take) + && let Some((mime, data)) = parse_data_uri(&url) + { + out.push(json!({"inlineData": {"mimeType": mime, "data": data}})); } } if !out.is_empty() { return out; } } - vec![json!({"text": m.content})] + vec![object([("text", Value::String(m.content))])] } /// Parse a `data:;base64,` URI into `(mime, payload)`. @@ -67,13 +66,11 @@ impl VertexEngine { ] } - fn build_body(&self) -> Value { - // Gemini has no system role: system turns go to systemInstruction, never contents - let contents: Vec = self - .base - .request - .message - .iter() + fn build_body(&mut self) -> Value { + let system = self.base.system_text(); + // there is no system role on Gemini: system turns go to systemInstruction, never contents + let contents: Vec = std::mem::take(&mut self.base.request.message) + .into_iter() .filter(|m| m.role != gw_consts::role::SYSTEM) .map(|m| { let role = if m.role == gw_consts::role::AI { @@ -89,7 +86,6 @@ impl VertexEngine { .collect(); let mut body = json!({}); body["contents"] = Value::Array(contents); - let system = self.base.system_text(); if !system.is_empty() { let part = object([("text", system.into())]); body["systemInstruction"] = object([("parts", Value::Array(vec![part]))]); @@ -115,7 +111,7 @@ impl VertexEngine { /// Native Gemini streaming: `:streamGenerateContent?alt=sse` frames decoded /// as they arrive and forwarded through `stream_tx` (the live-pump contract). - async fn run_stream(&self) -> GResult { + async fn run_stream(&mut self) -> GResult { let body = self.build_body(); let url = format!( "{}/v1beta/models/{}:streamGenerateContent?alt=sse", @@ -309,14 +305,14 @@ impl ModelEngine for EmbeddingsEngine { /// finished at "stop". fn family_outcome( message: String, - model: &str, + model: String, v: serde_json::Value, status: u16, ) -> EngineOutcome { EngineOutcome::with_status( GatewayResponse { message, - model: model.to_owned(), + model, response_v2: Some(v), finish_reason: "stop".to_owned(), ..Default::default() @@ -410,7 +406,7 @@ impl ModelEngine for ImageEngine { crate::engine::tok(&v["usage"]["input_tokens"]), crate::engine::tok(&v["usage"]["output_tokens"]), ); - let mut outcome = family_outcome(format!("{count} image(s) {verb}"), &model, v, status); + let mut outcome = family_outcome(format!("{count} image(s) {verb}"), model, v, status); outcome.response.prompt_tokens = input; outcome.response.completion_tokens = output; outcome.response.billed_units = count as i64; @@ -565,7 +561,7 @@ impl ModelEngine for AudioEngine { .or_else(|| local_seconds.map(|s| s.ceil() as i64)) .unwrap_or(0); } - let mut outcome = family_outcome(message, &model, v, status); + let mut outcome = family_outcome(message, model, v, status); outcome.response.prompt_tokens = input; outcome.response.completion_tokens = output; outcome.response.billed_units = units; @@ -575,7 +571,7 @@ impl ModelEngine for AudioEngine { } /// A vendor duration as whole billed seconds (fractions round up). -pub fn whole_seconds(v: &Value) -> Option { +fn whole_seconds(v: &Value) -> Option { v.as_i64() .or_else(|| v.as_f64().map(|f| f.ceil() as i64)) .or_else(|| { @@ -673,7 +669,7 @@ impl ModelEngine for VideoEngine { let dialect = video_dialect(self.base.provider(), self.base.wire_kind()); let model = self.base.model_name()?; let mut body = Map::new(); - // Kling names the field model_name and takes no inline image on this path + // the field is model_name on Kling, which takes no inline image on this path if dialect == VideoDialect::Kling { if p.image.is_some() { return Err(GatewayError::bad_request( @@ -742,7 +738,7 @@ impl ModelEngine for VideoEngine { vec![ ("duration", p.duration_seconds.map(|d| d.to_string().into())), ("aspect_ratio", p.aspect_ratio.map(Value::from)), - // Kling has quality modes, not resolutions: 1080p rides mode=pro + // quality modes replace resolutions on Kling: 1080p rides mode=pro ( "mode", p.resolution.map(|r| { @@ -812,7 +808,7 @@ fn video_outcome(model: &str, v: Value, status: u16) -> EngineOutcome { .unwrap_or_default() .to_owned(); let units = whole_seconds(&v["video"]["duration"]).unwrap_or(0); - let mut out = family_outcome(message, model, v, status); + let mut out = family_outcome(message, model.to_owned(), v, status); out.response.step = step; out.response.billed_units = units; out @@ -1064,16 +1060,17 @@ impl ModelEngine for SearchEngine { (status, v, "/results") } }; - let titles: Vec = v + let summary = v .pointer(results) .and_then(Value::as_array) .map(|rs| { rs.iter() - .filter_map(|r| r["title"].as_str().map(str::to_owned)) - .collect() + .filter_map(|r| r["title"].as_str()) + .collect::>() + .join("; ") }) .unwrap_or_default(); - let mut out = family_outcome(titles.join("; "), ¶m.model_name, v, status); + let mut out = family_outcome(summary, param.model_name.clone(), v, status); out.response.billed_units = 1; Ok(out) } @@ -1113,7 +1110,7 @@ impl ModelEngine for ModerationsEngine { .unwrap_or(0); Ok(family_outcome( format!("{flagged} flagged"), - model, + model.to_owned(), v, status, )) @@ -1151,7 +1148,7 @@ impl ModelEngine for RerankEngine { let n = v["results"].as_array().map(Vec::len).unwrap_or(0); let tokens = rerank_tokens(&v); let units = crate::engine::tok(&v["meta"]["billed_units"]["search_units"]); - let mut out = family_outcome(format!("{n} results"), &model, v, status); + let mut out = family_outcome(format!("{n} results"), model, v, status); out.response.prompt_tokens = tokens; out.response.total_tokens = tokens; out.response.billed_units = units; @@ -1190,7 +1187,7 @@ impl ModelEngine for PassthroughEngine { } else { "error" }; - Ok(family_outcome(message.to_owned(), &model, v, status)) + Ok(family_outcome(message.to_owned(), model, v, status)) } } diff --git a/crates/engines/src/lib.rs b/crates/engines/src/lib.rs index d796a9e..38009b1 100644 --- a/crates/engines/src/lib.rs +++ b/crates/engines/src/lib.rs @@ -32,8 +32,8 @@ pub use claude_engine::{ClaudeEngine, anthropic_native_chunks}; pub use engine::{EngineOutcome, ModelEngine, StreamChunk}; pub use factory::get_engine; pub use families::{ - AudioEngine, AudioKind, CompletionsEngine, EmbeddingsEngine, ImageEngine, ModerationsEngine, - PassthroughEngine, RerankEngine, ResponsesEngine, SearchEngine, VertexEngine, VideoEngine, + AudioEngine, AudioKind, CompletionsEngine, EmbeddingsEngine, ImageEngine, ResponsesEngine, + SearchEngine, VertexEngine, VideoEngine, }; pub use openai_engine::{OpenAiEngine, merge_tool_call_fragments}; pub use sse::SseDecoder; diff --git a/crates/engines/src/mock_transport.rs b/crates/engines/src/mock_transport.rs index fdea77f..7ff0d2c 100644 --- a/crates/engines/src/mock_transport.rs +++ b/crates/engines/src/mock_transport.rs @@ -11,8 +11,7 @@ use crate::transport::{ HeaderMap, MOCK_B64, MOCK_CREATED, Transport, UpstreamBody, UpstreamRequest, UpstreamResponse, }; -/// Deterministic fake vendor: parses the engine-built body and answers in the -/// vendor's wire shape; an account named "…down…" gets a 503 (the failover trigger). +/// Deterministic fake vendor: parses the engine-built body and answers in the vendor's wire shape. #[derive(Debug, Default)] pub struct MockTransport; @@ -30,11 +29,12 @@ impl MockTransport { } } + fn last_user(messages: &[Value]) -> Option<&Value> { + messages.iter().rev().find(|m| m["role"] == "user") + } + fn last_user_text(messages: &[Value]) -> String { - messages - .iter() - .rev() - .find(|m| m["role"] == "user") + Self::last_user(messages) .and_then(|m| { m["content"].as_str().map(str::to_owned).or_else(|| { m["content"].as_array().map(|blocks| { @@ -65,10 +65,7 @@ impl MockTransport { } fn image_count(messages: &[Value]) -> usize { - messages - .iter() - .rev() - .find(|m| m["role"] == "user") + Self::last_user(messages) .and_then(|m| m["content"].as_array()) .map(|parts| parts.iter().filter(|p| p["type"] == "image_url").count()) .unwrap_or(0) @@ -321,7 +318,7 @@ impl MockTransport { .map(String::into_bytes) .collect(), UpstreamBody::Json(bytes) => { - let mut v: Value = serde_json::from_slice(&bytes) + let v: Value = serde_json::from_slice(&bytes) .map_err(|e| GatewayError::internal("mock bedrock reply").with_source(e))?; if v.get("generation").is_some() { let text = v["generation"].as_str().unwrap_or_default().to_owned(); @@ -333,18 +330,6 @@ impl MockTransport { "generation_token_count": v["generation_token_count"], "stop_reason": "stop"}), ] - } else if v.get("text").is_some() { - let text = v["text"].as_str().unwrap_or_default().to_owned(); - let (a, b) = Self::split_half(&text); - vec![ - json!({"is_finished": false, "event_type": "text-generation", "text": a}), - json!({"is_finished": false, "event_type": "text-generation", "text": b}), - json!({"is_finished": true, "event_type": "stream-end", - "finish_reason": "COMPLETE", - "amazon-bedrock-invocationMetrics": { - "inputTokenCount": v["meta"]["tokens"]["input_tokens"].take(), - "outputTokenCount": v["meta"]["tokens"]["output_tokens"].take()}}), - ] } else { // a JSON answer to a stream request (the tool-use reply) stays JSON return Ok(UpstreamResponse { diff --git a/crates/engines/src/openai_engine.rs b/crates/engines/src/openai_engine.rs index fe9e856..2a4d7b7 100644 --- a/crates/engines/src/openai_engine.rs +++ b/crates/engines/src/openai_engine.rs @@ -28,7 +28,7 @@ impl OpenAiEngine { parts => { let mut msg = Map::new(); msg.insert("role".into(), m.role.into()); - // OpenAI: assistant tool-call turns carry content: null + // assistant tool-call turns carry content: null on the OpenAI wire let content = match (parts, m.content) { (Some(parts), _) => parts, (None, c) if c.is_empty() && m.tool_calls.is_some() => Value::Null, @@ -110,11 +110,9 @@ impl OpenAiEngine { let mut system = Map::with_capacity(2); system.insert("role".into(), "system".into()); system.insert("content".into(), Value::String(s)); - let mut msgs = vec![Value::Object(system)]; - if let Some(Value::Array(existing)) = body.remove("messages") { - msgs.extend(existing); + if let Some(Value::Array(msgs)) = body.get_mut("messages") { + msgs.insert(0, Value::Object(system)); } - body.insert("messages".into(), Value::Array(msgs)); } } let raw = self.base.take_raw(); @@ -278,7 +276,6 @@ fn apply_sse_event( full.push_str(&text); chunks.push(StreamChunk { delta: text, - finish_reason: None, ..Default::default() }); } @@ -313,7 +310,6 @@ fn apply_sse_event( } resp.finish_reason = fr.to_owned(); chunks.push(StreamChunk { - delta: String::new(), finish_reason: Some(fr.to_owned()), ..Default::default() }); diff --git a/crates/engines/src/pump.rs b/crates/engines/src/pump.rs index 2e1e5b9..4380593 100644 --- a/crates/engines/src/pump.rs +++ b/crates/engines/src/pump.rs @@ -131,7 +131,7 @@ where } else { out.aborted = true; } - out.streamed_live = tx.is_some(); + out.streamed_live = true; return Ok(out); } Err(e) => return Err(e), diff --git a/crates/engines/src/realtime.rs b/crates/engines/src/realtime.rs index f2d0f95..0533fe7 100644 --- a/crates/engines/src/realtime.rs +++ b/crates/engines/src/realtime.rs @@ -4,11 +4,6 @@ use serde_json::Value; -/// Whether a client frame is the OpenAI-dialect generation trigger. -pub fn is_response_create(frame: &Value) -> bool { - frame["type"] == "response.create" -} - /// The client frame that starts a generation — the admission point. OpenAI /// signals it as `response.create`, Gemini Live as a completed client turn. pub fn is_client_turn(provider: &str, frame: &Value) -> bool { @@ -16,7 +11,7 @@ pub fn is_client_turn(provider: &str, frame: &Value) -> bool { return frame["clientContent"]["turnComplete"] == Value::Bool(true) || frame["client_content"]["turn_complete"] == Value::Bool(true); } - is_response_create(frame) + frame["type"] == "response.create" } /// Visit a realtime frame's text-bearing string leaves with a visitor that may @@ -99,7 +94,6 @@ pub fn is_realtime_turn_started(provider: &str, frame: &Value) -> bool { /// Delivered output in one frame: OpenAI deltas yield text (or audio quanta), /// Gemini `modelTurn` parts count as byte-estimated opaque units. pub fn realtime_output_delta(frame: &Value) -> (Option<&str>, usize) { - // Gemini Live: delivered output rides serverContent.modelTurn parts if let Some(parts) = frame["serverContent"]["modelTurn"]["parts"].as_array() { let opaque = parts .iter() diff --git a/crates/engines/src/sigv4.rs b/crates/engines/src/sigv4.rs index 91dc9bf..7764c5d 100644 --- a/crates/engines/src/sigv4.rs +++ b/crates/engines/src/sigv4.rs @@ -27,7 +27,7 @@ pub struct SigV4Params<'a> { } fn hmac(key: &[u8], data: &[u8]) -> Vec { - // HMAC-SHA256 accepts any key length (RFC 2104), so this cannot fire + // any key length is valid for HMAC-SHA256 (RFC 2104), so this cannot fire #[allow(clippy::expect_used)] let mut mac = HmacSha256::new_from_slice(key).expect("hmac accepts any key length"); mac.update(data); @@ -39,7 +39,7 @@ fn sha256_hex(data: &[u8]) -> String { } /// The derived signing key: kSecret → kDate → kRegion → kService → kSigning. -pub fn signing_key(secret: &str, date: &str, region: &str, service: &str) -> Vec { +fn signing_key(secret: &str, date: &str, region: &str, service: &str) -> Vec { let k_date = hmac(format!("AWS4{secret}").as_bytes(), date.as_bytes()); let k_region = hmac(&k_date, region.as_bytes()); let k_service = hmac(&k_region, service.as_bytes()); diff --git a/crates/engines/tests/http_transport_wire.rs b/crates/engines/tests/http_transport_wire.rs index 2c3ed03..c156646 100644 --- a/crates/engines/tests/http_transport_wire.rs +++ b/crates/engines/tests/http_transport_wire.rs @@ -18,10 +18,12 @@ use gw_models::{ChatMsg, GatewayRequest, ModelParamV2}; use serde_json::{Value, json}; async fn serve_router(app: Router) -> std::net::SocketAddr { - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind"); + let addr = listener.local_addr().expect("addr"); tokio::spawn(async move { - axum::serve(listener, app).await.unwrap(); + axum::serve(listener, app).await.expect("serve"); }); addr } @@ -46,8 +48,7 @@ async fn spawn_vendor() -> String { ); axum::response::Response::builder() .header("content-type", "text/event-stream") - .body(axum::body::Body::from(sse)) - .unwrap() + .body(axum::body::Body::from(sse)).expect("response") } else { let payload = json!({ "id":"srv-1","object":"chat.completion","model":"srv", @@ -68,7 +69,7 @@ async fn spawn_vendor() -> String { #[tokio::test] async fn http_transport_json_over_real_socket() { let base = spawn_vendor().await; - let transport = HttpTransport::new(Duration::from_secs(5)).unwrap(); + let transport = HttpTransport::new(Duration::from_secs(5)).expect("transport"); let req = UpstreamRequest { protocol: Protocol::OpenaiChat, method: "POST", @@ -110,7 +111,7 @@ async fn spawn_stalled_json_vendor() -> String { axum::response::Response::builder() .header("content-type", "application/json") .body(axum::body::Body::from_stream(body)) - .unwrap() + .expect("response") }), ); let addr = serve_router(app).await; @@ -119,7 +120,7 @@ async fn spawn_stalled_json_vendor() -> String { #[tokio::test] async fn non_stream_body_timeout_classifies_as_model_timeout() { - let transport = HttpTransport::new(Duration::from_millis(300)).unwrap(); + let transport = HttpTransport::new(Duration::from_millis(300)).expect("transport"); let err = transport .send(UpstreamRequest { protocol: Protocol::OpenaiChat, @@ -153,7 +154,7 @@ async fn spawn_breaking_json_vendor() -> String { axum::response::Response::builder() .header("content-type", "application/json") .body(axum::body::Body::from_stream(body)) - .unwrap() + .expect("response") }), ); let addr = serve_router(app).await; @@ -162,7 +163,7 @@ async fn spawn_breaking_json_vendor() -> String { #[tokio::test] async fn non_stream_body_break_classifies_as_model_error() { - let transport = HttpTransport::new(Duration::from_secs(5)).unwrap(); + let transport = HttpTransport::new(Duration::from_secs(5)).expect("transport"); let err = transport .send(UpstreamRequest { protocol: Protocol::OpenaiChat, @@ -187,7 +188,7 @@ async fn non_stream_body_break_classifies_as_model_error() { #[tokio::test] async fn http_transport_sse_over_real_socket() { let base = spawn_vendor().await; - let transport = HttpTransport::new(Duration::from_secs(5)).unwrap(); + let transport = HttpTransport::new(Duration::from_secs(5)).expect("transport"); let req = UpstreamRequest { protocol: Protocol::OpenaiChat, method: "POST", @@ -215,7 +216,7 @@ async fn http_transport_sse_over_real_socket() { #[tokio::test] async fn dispatch_routes_mock_scheme_in_process_and_real_urls_over_http() { let base = spawn_vendor().await; - let transport = DispatchTransport::new(Duration::from_secs(5)).unwrap(); + let transport = DispatchTransport::new(Duration::from_secs(5)).expect("transport"); let req = |url: String| UpstreamRequest { protocol: Protocol::OpenaiChat, @@ -261,7 +262,7 @@ async fn dispatch_routes_mock_scheme_in_process_and_real_urls_over_http() { async fn engine_through_real_http_transport_end_to_end() { let base = spawn_vendor().await; let transport: gw_engines::SharedTransport = - Arc::new(HttpTransport::new(Duration::from_secs(5)).unwrap()); + Arc::new(HttpTransport::new(Duration::from_secs(5)).expect("transport")); let account = gw_models::Account { name: "real-local".into(), @@ -297,7 +298,8 @@ async fn per_account_policy_and_connect_retry() { connect_retries: 2, }, ); - let transport = HttpTransport::with_policies(UpstreamPolicy::default(), per_account).unwrap(); + let transport = + HttpTransport::with_policies(UpstreamPolicy::default(), per_account).expect("transport"); assert_eq!(transport.policy_for("tight").connect_retries, 2); assert_eq!(transport.policy_for("other").connect_retries, 1); @@ -322,8 +324,8 @@ async fn per_account_policy_and_connect_retry() { ); let closed = std::iter::repeat_with(|| { - let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); - l.local_addr().unwrap() + let l = std::net::TcpListener::bind("127.0.0.1:0").expect("bind"); + l.local_addr().expect("addr") }) .find(|addr| std::net::TcpStream::connect_timeout(addr, Duration::from_millis(50)).is_err()) .unwrap(); @@ -368,7 +370,7 @@ async fn spawn_paced_vendor(frames: usize, gap: Duration) -> String { axum::response::Response::builder() .header("content-type", "text/event-stream") .body(axum::body::Body::from_stream(sse)) - .unwrap() + .expect("response") }), ); let addr = serve_router(app).await; @@ -391,7 +393,7 @@ fn paced_req(url: String) -> UpstreamRequest { #[tokio::test] async fn slow_stream_outlives_the_total_policy_timeout() { let url = spawn_paced_vendor(5, Duration::from_millis(300)).await; - let transport = HttpTransport::new(Duration::from_secs(1)).unwrap(); + let transport = HttpTransport::new(Duration::from_secs(1)).expect("transport"); let resp = transport.send(paced_req(url)).await.unwrap(); let resp = resp .buffered() @@ -410,7 +412,7 @@ async fn slow_stream_outlives_the_total_policy_timeout() { #[tokio::test] async fn stalled_stream_errors_at_the_idle_gap_instead_of_hanging() { let url = spawn_paced_vendor(2, Duration::from_secs(20)).await; - let transport = HttpTransport::new(Duration::from_millis(300)).unwrap(); + let transport = HttpTransport::new(Duration::from_millis(300)).expect("transport"); let started = std::time::Instant::now(); let err = match transport.send(paced_req(url)).await { Ok(resp) => resp @@ -604,12 +606,11 @@ async fn an_error_status_under_an_sse_content_type_is_a_body_not_a_stream() { .header("content-type", "text/event-stream") .body(axum::body::Body::from( r#"{"error":{"code":404,"message":"models/nope is not found","status":"NOT_FOUND"}}"#, - )) - .unwrap() + )).expect("response") }), ); let addr = serve_router(app).await; - let transport = HttpTransport::new(Duration::from_secs(5)).unwrap(); + let transport = HttpTransport::new(Duration::from_secs(5)).expect("transport"); let resp = transport .send(UpstreamRequest { protocol: Protocol::Gemini, diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 615126f..c7dcf84 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -127,13 +127,12 @@ impl OnlineHandler { ); ctx.billing_deferred = dlp && ctx.request.is_online && ctx.request.stream; // every fired rule is recorded (block/flag/shadow alike); only a block-action hit denies - let rows: Vec = scan - .hits - .iter() - .map(|hit| security_event(&ctx, &hit.rule, hit.action.as_str(), hit.count)) - .collect(); + deferred.extend( + scan.hits + .iter() + .map(|hit| security_event(&ctx, &hit.rule, hit.action.as_str(), hit.count)), + ); if let Some(block) = scan.block { - deferred.extend(rows); ctx.decide( "security_check", format!("blocked (code {})", block.err_code), @@ -141,7 +140,6 @@ impl OnlineHandler { ctx.outcome = Some(content_filter_outcome(block)); return Ok(ctx); } - deferred.extend(rows); // pre-DLP text, computed once for moderation and the retained prompt let inbound = @@ -258,6 +256,8 @@ impl OnlineHandler { let mut tried = 0; loop { + ctx.fallback_ahead = !ctx.request.replays_reasoning_output() + && next_fallback(&snap.cfg, &ctx, tried).is_some(); // a panicking node must refund too; the refund reads only whole-written ctx fields let ran = std::panic::AssertUnwindSafe(gw_dag::run(&self.layers, &mut ctx)) .catch_unwind() @@ -305,7 +305,7 @@ impl OnlineHandler { // raw response pre-outbound-DLP, only when full retention can store it (key present) let capture_raw = matches!(retention, Some(r) if r.content == gw_config::ContentLevel::Full) - && gw_state::sealing_available(); + && gw_state::can_seal(); let raw_response = capture_raw .then(|| ctx.outcome.as_ref().map(|o| o.response.message.clone())) .flatten(); @@ -448,21 +448,19 @@ enum Moderation { } struct TerminalSubject { + ak: Arc, request_id: String, - ak: String, user_id: String, - tenant: String, } impl TerminalSubject { - fn new(ak: &AkInfo, request: &GatewayRequest) -> Self { + fn new(ak: &Arc, request: &GatewayRequest) -> Self { Self { + ak: Arc::clone(ak), request_id: request.request_id.clone(), - ak: ak.ak.clone(), user_id: ak .attributed_user(request.user_id.as_deref().unwrap_or_default()) .to_owned(), - tenant: ak.tenant.clone(), } } } @@ -673,9 +671,9 @@ async fn persist_terminal( let record = gw_state::ContentRecord { created_at_epoch_secs: now, request_id: subject.request_id.clone(), - ak: subject.ak.clone(), + ak: subject.ak.ak.clone(), user_id: subject.user_id.clone(), - tenant: subject.tenant.clone(), + tenant: subject.ak.tenant.clone(), kind: "terminal".to_owned(), content: body.to_string(), sealed: false, @@ -817,10 +815,11 @@ async fn persist_content( #[cfg(test)] mod tests { - use super::*; use gw_consts::Protocol; use gw_models::{ChatMsg, ModelParamV2}; + use super::*; + fn handler() -> OnlineHandler { let cfg = Arc::new(GatewayConfig::embedded_default().unwrap()); let state = Arc::new(GatewayState::from_config(&cfg)); @@ -1141,6 +1140,26 @@ mod tests { ); } + #[tokio::test] + async fn a_served_fallback_samples_one_success_for_the_requested_model() { + let (endpoint, _) = vendor_by_model().await; + let h = fallback_handler( + &endpoint, + "tenants: [{name: t1, models: [broken, healthy]}]", + ) + .await; + let ak = h.state().auth.authenticate("k1").await.unwrap(); + h.run(chat_req("broken", "hi"), ak).await.unwrap(); + let avail = &h.state().avail; + avail.flush().await; + let minute = gw_state::epoch_secs() / 60; + assert_eq!( + avail.window("broken", minute - 5, minute).await, + (1, 0), + "the client saw one success; the failed first attempt is not a sample" + ); + } + #[tokio::test] async fn an_exhausted_chain_reports_the_last_upstream_error() { let (endpoint, hits) = vendor_by_model().await; @@ -2637,6 +2656,7 @@ mod tests { let mut cfg = GatewayConfig::embedded_default().unwrap(); cfg.security.dlp_redact = true; cfg.security.blocklist = vec!["example.com".into()]; + cfg.security = std::mem::take(&mut cfg.security).compiled(); let cfg = Arc::new(cfg); let state = Arc::new(GatewayState::from_config(&cfg)); let h = OnlineHandler::new( @@ -2932,7 +2952,7 @@ mod tests { let submitter = OfflineHandler::new(online.clone()); let ak = state.auth.authenticate("ak-demo-123").await.unwrap(); - assert!(state.store.distributed_batches()); + assert!(state.store.distributes_batches()); let job = submitter .submit( ak, diff --git a/crates/handler/src/moderation.rs b/crates/handler/src/moderation.rs index 7f1c100..1a09fa4 100644 --- a/crates/handler/src/moderation.rs +++ b/crates/handler/src/moderation.rs @@ -109,7 +109,7 @@ impl Moderator for BedrockGuardrail { /// The moderator the config names, else the allow-all default. pub fn from_config(conf: Option<&ModerationConf>) -> Arc { let Some(conf) = conf else { - return default_moderator(); + return Arc::new(AllowModerator); }; let api_key = conf.api_key().unwrap_or_else(|| { tracing::warn!(var = %conf.api_key_env, "moderation api key env is unset; reviews will fail"); @@ -118,10 +118,6 @@ pub fn from_config(conf: Option<&ModerationConf>) -> Arc { Arc::new(BedrockGuardrail::new(conf, api_key)) } -pub fn default_moderator() -> Arc { - Arc::new(AllowModerator) -} - /// Map an `ApplyGuardrail` reply onto a verdict over the reviewed `text`. fn guardrail_verdict(text: &str, reply: &Value) -> Verdict { if reply["action"] != "GUARDRAIL_INTERVENED" { diff --git a/crates/handler/src/offline.rs b/crates/handler/src/offline.rs index 10c774f..864ffb6 100644 --- a/crates/handler/src/offline.rs +++ b/crates/handler/src/offline.rs @@ -2,6 +2,7 @@ //! instance, a distributed store persists items for any instance's drain loop. use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering::Relaxed}; use gw_models::{BatchItem, GatewayRequest, ModelParamV2}; use gw_state::{AkInfo, BatchItemResult, BatchJob, BatchStatus}; @@ -11,7 +12,7 @@ use crate::OnlineHandler; /// Batch orchestration built on top of the online handler. #[derive(Clone)] pub struct OfflineHandler { - pub online: OnlineHandler, + online: OnlineHandler, } impl OfflineHandler { @@ -24,20 +25,16 @@ impl OfflineHandler { &self, ak: Arc, model: String, - items: Vec, + mut items: Vec, ) -> gw_models::GResult { let store = self.online.state().store.clone(); - if store.distributed_batches() { + if store.distributes_batches() { // persist the EFFECTIVE user: execution, billing and erasure key on one identity - let items: Vec = items - .into_iter() - .map(|mut i| { - if let Some(owner) = ak.owner_override() { - i.user = owner.to_owned(); - } - i - }) - .collect(); + if let Some(owner) = ak.owner_override() { + for item in &mut items { + item.user = owner.to_owned(); + } + } // atomic: the job becomes claimable only once all items are saved store .batch_enqueue(&ak.ak, &ak.tenant, &model, &items) @@ -90,9 +87,8 @@ impl OfflineHandler { }; let done_indices: std::collections::HashSet = prior.iter().map(|r| r.index).collect(); - use std::sync::atomic::Ordering::Relaxed; // heartbeat: keeps a slow item from being judged stale, flips `lost` when the fence moves - let lost = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); + let lost = Arc::new(AtomicBool::new(false)); let hb = { let store = store.clone(); let id = id.to_owned(); @@ -122,7 +118,7 @@ impl OfflineHandler { break; } // re-read before dispatch (fail CLOSED): an erasure while queued blanks the stored item - if store.distributed_batches() { + if store.distributes_batches() { match store.batch_item_snapshot(id, index).await { Ok(Some(fresh)) => item = fresh, Ok(None) | Err(_) => { @@ -140,7 +136,7 @@ impl OfflineHandler { } let user = ak.attributed_user(&item.user).to_owned(); // local backends keep no item rows: the erasure marker stops the rest (fail closed) - let erased_mid_batch = !store.distributed_batches() + let erased_mid_batch = !store.distributes_batches() && store .user_erased_since(&ak.tenant, &user, captured_at) .await @@ -180,7 +176,7 @@ impl OfflineHandler { Ok(Err(e)) => failed_item(index, e.to_string(), user), Err(join_err) => failed_item(index, format!("item task failed: {join_err}"), user), }; - // if we lost the claim mid-run, don't persist — the new owner is authoritative + // a claim lost mid-run is not persisted: the new owner is authoritative if lost.load(Relaxed) { break; } @@ -213,7 +209,7 @@ impl OfflineHandler { let claimed = tokio::select! { biased; changed = shutdown.changed() => { - if stopping(changed, &shutdown) { + if is_stopping(changed, &shutdown) { return; } continue; @@ -276,7 +272,7 @@ impl OfflineHandler { } } -fn stopping( +fn is_stopping( changed: Result<(), tokio::sync::watch::error::RecvError>, shutdown: &tokio::sync::watch::Receiver, ) -> bool { @@ -289,7 +285,7 @@ async fn pause_or_stop( ) -> bool { tokio::select! { biased; - changed = shutdown.changed() => stopping(changed, shutdown), + changed = shutdown.changed() => is_stopping(changed, shutdown), _ = tokio::time::sleep(poll) => false, } } diff --git a/crates/handler/src/plugins.rs b/crates/handler/src/plugins.rs index 2942a79..587f83d 100644 --- a/crates/handler/src/plugins.rs +++ b/crates/handler/src/plugins.rs @@ -1,6 +1,8 @@ //! Rule-based request/response plugins from `config.security`: the pre-stage //! blocks and DLP-redacts inbound text, the post-stage redacts the response. +use std::fmt::Write as _; + use gw_config::{Action, SecurityConf}; use gw_models::{Block, ChatMsg, GatewayRequest, GatewayResponse, ModelParamV2}; @@ -38,12 +40,11 @@ impl<'a> ScanCounts<'a> { } } - fn visit(&mut self, s: &str) -> usize { - self.blocklist += i64::from(blocklist_hit(self.sec, s)); + fn visit(&mut self, s: &str) { + self.blocklist += i64::from(is_blocklisted(self.sec, s)); for (i, r) in self.sec.regexes.iter().enumerate() { self.regex[i] += r.re.find_iter(s).count() as i64; } - 0 } /// Fold the counts into a [`ScanOutcome`]; any block-action hit denies. @@ -198,7 +199,10 @@ pub fn security_check(sec: &SecurityConf, request: &mut GatewayRequest) -> ScanO return ScanOutcome::default(); } let mut counts = ScanCounts::new(sec); - for_each_request_text(request, SignedThinking::Visit, &mut |s, _| counts.visit(s)); + for_each_request_text(request, SignedThinking::Visit, &mut |s, _| { + counts.visit(s); + 0 + }); counts.outcome() } @@ -281,23 +285,15 @@ pub fn apply_mask_spans_frame( } /// Case-insensitive blocklist test; ASCII matches without allocating, non-ASCII copies once. -fn blocklist_hit(sec: &SecurityConf, text: &str) -> bool { - if text.is_ascii() { - return sec - .blocklist - .iter() - .any(|w| contains_ignore_ascii_case(text, w)); - } - let lower = text.to_lowercase(); - sec.blocklist.iter().any(|w| lower.contains(w)) -} - -fn contains_ignore_ascii_case(haystack: &str, needle: &str) -> bool { - let (h, n) = (haystack.as_bytes(), needle.as_bytes()); - if n.is_empty() || n.len() > h.len() { +fn is_blocklisted(sec: &SecurityConf, text: &str) -> bool { + let Some(matcher) = &sec.blocklist_matcher else { return false; + }; + if text.is_ascii() { + matcher.is_match(text) + } else { + matcher.is_match(&text.to_lowercase()) } - h.windows(n.len()).any(|w| w.eq_ignore_ascii_case(n)) } /// Walk every string leaf of a JSON value with a rewriting visitor; returns summed hits. @@ -673,26 +669,29 @@ fn walk_native_event( fragments: &mut EventFragments, f: &mut impl FnMut(&mut String) -> usize, ) -> usize { - if event["type"] == "message_start" { - let mut hits = walk_object_excluding(event, &["message"], f); - if let Some(message) = event.get_mut("message") { - hits += walk_object_excluding(message, &["content"], f); - if let Some(content) = message.get_mut("content") { - hits += walk_part_text(content, SignedThinking::Prose, &mut |s, _| f(s)); + match event["type"].as_str() { + Some("message_start") => { + let mut hits = walk_object_excluding(event, &["message"], f); + if let Some(message) = event.get_mut("message") { + hits += walk_object_excluding(message, &["content"], f); + if let Some(content) = message.get_mut("content") { + hits += walk_part_text(content, SignedThinking::Prose, &mut |s, _| f(s)); + } } + return hits; } - return hits; - } - if event["type"] == "content_block_start" { - let mut hits = walk_object_excluding(event, &["content_block"], f); - if let Some(block) = event.get_mut("content_block") { - hits += if block.as_object().is_some_and(is_signed_thinking_block) { - walk_signed_prose(block, &mut |s, _| f(s)) - } else { - walk_part_value(block, f) - }; + Some("content_block_start") => { + let mut hits = walk_object_excluding(event, &["content_block"], f); + if let Some(block) = event.get_mut("content_block") { + hits += if block.as_object().is_some_and(is_signed_thinking_block) { + walk_signed_prose(block, &mut |s, _| f(s)) + } else { + walk_part_value(block, f) + }; + } + return hits; } - return hits; + _ => {} } if event["type"] == "content_block_delta" && event.get("delta").is_some() { let opaque_key: Option<&'static str> = match event["delta"]["type"].as_str() { @@ -715,7 +714,7 @@ fn walk_native_event( } return hits; } - // Responses deltas join per output item so a pattern split across frames still matches + // deltas join per output item so a pattern split across frames still matches if let Some(index) = event["output_index"].as_u64() && let Some(text) = event["delta"].as_str() { @@ -746,13 +745,17 @@ fn collect_delta_payload( opaque_key: Option<&str>, fragments: &mut EventFragments, ) { + let mut path = String::from("delta"); let Some(object) = delta.as_object() else { - collect_delta_fragments(delta, index, "delta", fragments); + collect_delta_fragments(delta, index, &mut path, fragments); return; }; for (key, value) in object { if key != "type" && Some(key.as_str()) != opaque_key { - collect_delta_fragments(value, index, &format!("delta/{key}"), fragments); + path.push('/'); + path.push_str(key); + collect_delta_fragments(value, index, &mut path, fragments); + path.truncate("delta".len()); } if fragments.overflowed { break; @@ -763,19 +766,22 @@ fn collect_delta_payload( fn collect_delta_fragments( value: &serde_json::Value, index: u64, - path: &str, + path: &mut String, fragments: &mut EventFragments, ) { if fragments.overflowed { return; } + let len = path.len(); match value { serde_json::Value::String(text) => { fragments.append(index, path, text); } serde_json::Value::Array(values) => { for (position, value) in values.iter().enumerate() { - collect_delta_fragments(value, index, &format!("{path}/{position}"), fragments); + let _ = write!(path, "/{position}"); + collect_delta_fragments(value, index, path, fragments); + path.truncate(len); if fragments.overflowed { break; } @@ -783,7 +789,10 @@ fn collect_delta_fragments( } serde_json::Value::Object(object) => { for (key, value) in object { - collect_delta_fragments(value, index, &format!("{path}/{key}"), fragments); + path.push('/'); + path.push_str(key); + collect_delta_fragments(value, index, path, fragments); + path.truncate(len); if fragments.overflowed { break; } @@ -982,9 +991,10 @@ fn redact(text: &str) -> Option<(String, usize)> { #[cfg(test)] mod tests { - use super::*; use gw_models::ChatMsg; + use super::*; + #[test] fn mask_spans_map_across_slots_like_inbound_text() { let mut req = GatewayRequest { @@ -1020,6 +1030,7 @@ mod tests { dlp_redact: true, ..Default::default() } + .compiled() } #[test] @@ -1044,7 +1055,8 @@ mod tests { blocklist: vec!["forbiddenword".into(), "禁词".into()], dlp_redact: false, ..Default::default() - }; + } + .compiled(); let mut req = GatewayRequest { message: vec![ChatMsg::text("user", "前文 FORBIDDENWORD 后文")], ..Default::default() @@ -1112,7 +1124,8 @@ mod tests { blocklist: vec!["watch".into()], blocklist_action: Action::Flag, ..Default::default() - }; + } + .compiled(); let mut frame = serde_json::json!({"type":"input_text","text":"please watch this"}); let (out, text, _) = realtime_frame_scan(&s2, &mut frame, false); assert!(out.block.is_none(), "flag does not block realtime"); @@ -1127,7 +1140,8 @@ mod tests { blocklist_action: Action::Flag, detect_secrets: true, ..Default::default() - }; + } + .compiled(); let mut frame = serde_json::json!({ "type":"input_text","text":"key sk-abcdefghijklmnopqrstuvwxyz012345" }); @@ -1400,7 +1414,8 @@ mod tests { blocklist: vec!["watchword".into()], blocklist_action: Action::Flag, ..Default::default() - }; + } + .compiled(); let mut req = GatewayRequest { message: vec![ChatMsg::text("user", "contains watchword here")], ..Default::default() diff --git a/crates/models/src/cost.rs b/crates/models/src/cost.rs index 9d97300..934fcc5 100644 --- a/crates/models/src/cost.rs +++ b/crates/models/src/cost.rs @@ -104,20 +104,6 @@ pub fn long_context_scale( ) } -/// Weighted (prompt, completion) for the paths carrying no cache/reasoning -/// components (estimates, realtime turns). -pub fn weighted_pair(prompt: i64, completion: i64, rate: &TokenRate) -> (i64, i64) { - let input = TokenInput { - prompt, - completion, - ..Default::default() - }; - ( - weighted_prompt(&input, rate), - weighted_completion(&input, rate), - ) -} - /// Cache-normalized prompt (clamped at 0). fn normalize_prompt(input: &TokenInput, rate: &TokenRate) -> i64 { let mut prompt = input.prompt; @@ -231,15 +217,6 @@ mod tests { assert_eq!(weighted_completion(&input, &rate), 60); } - #[test] - fn weighted_pair_carries_flat_counts() { - let rate = TokenRate { - prompt_weight: 0.5, - ..Default::default() - }; - assert_eq!(weighted_pair(100, 50, &rate), (50, 50)); - } - #[test] fn negative_prompt_clamped_to_zero() { let rate = TokenRate { diff --git a/crates/models/src/lib.rs b/crates/models/src/lib.rs index aed58ca..58d344f 100644 --- a/crates/models/src/lib.rs +++ b/crates/models/src/lib.rs @@ -13,8 +13,7 @@ pub mod usage; pub use block::Block; pub use cost::{ - TokenInput, TokenRate, cost_micros, long_context_scale, weighted_completion, weighted_pair, - weighted_prompt, + TokenInput, TokenRate, cost_micros, long_context_scale, weighted_completion, weighted_prompt, }; pub use error::{GResult, GatewayError}; pub use params::{ diff --git a/crates/models/src/request.rs b/crates/models/src/request.rs index 964f3c0..421609b 100644 --- a/crates/models/src/request.rs +++ b/crates/models/src/request.rs @@ -84,12 +84,7 @@ impl GatewayRequest { .and_then(serde_json::Value::as_bool) == Some(true) }); - let replay = param - .raw - .get("input") - .and_then(serde_json::Value::as_array) - .is_some_and(|items| items.iter().any(|item| item["type"] == "reasoning")); - return reasoning || replay; + return reasoning || replays_responses_reasoning(param); } matches!( param @@ -113,14 +108,10 @@ impl GatewayRequest { /// protected block, or a Responses `input` reasoning item. A bare reasoning /// request that has produced nothing yet does not, so it may still fall back. pub fn replays_reasoning_output(&self) -> bool { - let responses_input_replay = self - .model_param_v2 + self.model_param_v2 .as_ref() .filter(|p| p.protocol == gw_consts::Protocol::Responses) - .and_then(|p| p.raw.get("input")) - .and_then(serde_json::Value::as_array) - .is_some_and(|items| items.iter().any(|item| item["type"] == "reasoning")); - responses_input_replay + .is_some_and(replays_responses_reasoning) || self.message.iter().any(|m| { m.reasoning_details.is_some() || m.parts @@ -131,6 +122,15 @@ impl GatewayRequest { } } +/// Whether a Responses `input` carries a reasoning item from an earlier turn. +fn replays_responses_reasoning(param: &ModelParamV2) -> bool { + param + .raw + .get("input") + .and_then(serde_json::Value::as_array) + .is_some_and(|items| items.iter().any(|item| item["type"] == "reasoning")) +} + /// One queued batch item: messages plus the end-user attribution, persisted so /// a distributed drainer still attributes and budgets it (owner overrides at billing). #[derive(Debug, Clone)] @@ -311,9 +311,10 @@ pub(crate) fn is_protected_anthropic_block(block: &serde_json::Value) -> bool { #[cfg(test)] mod tests { - use super::*; use gw_consts::Protocol; + use super::*; + #[test] fn dispatch_protocol() { let empty = GatewayRequest::default(); diff --git a/crates/protocol/src/anthropic.rs b/crates/protocol/src/anthropic.rs index 06efdab..9fa76a7 100644 --- a/crates/protocol/src/anthropic.rs +++ b/crates/protocol/src/anthropic.rs @@ -81,7 +81,11 @@ pub fn image_url_to_image(mut part: Value) -> Value { if part["type"] != "image_url" { return part; } - let url = match part["image_url"].get_mut("url").map(Value::take) { + let url = match part + .get_mut("image_url") + .and_then(|v| v.get_mut("url")) + .map(Value::take) + { Some(Value::String(url)) => url, _ => return part, }; @@ -105,7 +109,7 @@ pub fn image_to_image_url(mut block: Value) -> Value { if block["type"] != "image" { return block; } - let mut source = block["source"].take(); + let mut source = block.get_mut("source").map(Value::take).unwrap_or_default(); let url = match source["type"].as_str() { Some("base64") => format!( "data:{};base64,{}", diff --git a/crates/protocol/src/reasoning.rs b/crates/protocol/src/reasoning.rs index dc64e10..c5907b2 100644 --- a/crates/protocol/src/reasoning.rs +++ b/crates/protocol/src/reasoning.rs @@ -51,7 +51,7 @@ pub fn budget_effort(budget: i64) -> &'static str { /// The thinking dialect of a model on the Anthropic wire, by name. Vendors /// speaking that wire (MiniMax, GLM, Kimi) cloned the budget dialect. pub fn anthropic_thinking_dialect(model: &str) -> ThinkingDialect { - // Bedrock ids carry a vendor (and region) prefix: `us.anthropic.claude-…` + // ids on Bedrock carry a vendor (and region) prefix: `us.anthropic.claude-…` let model = model.find("claude").map_or(model, |i| &model[i..]); if !model.starts_with("claude") || model.starts_with("claude-3") { return ThinkingDialect::Budget; diff --git a/crates/server/src/main.rs b/crates/server/src/main.rs index c07aabb..6b1b4d4 100644 --- a/crates/server/src/main.rs +++ b/crates/server/src/main.rs @@ -63,7 +63,7 @@ async fn main() -> anyhow::Result<()> { None => cfg, }; - // GW_HOST / GW_PORT win over the config file (GW_HOST=0.0.0.0 for containers). + // the GW_HOST / GW_PORT env vars win over the config file (GW_HOST=0.0.0.0 for containers) let host = env::var("GW_HOST").unwrap_or_else(|_| cfg.listen.host.clone()); let port = env::var("GW_PORT") .ok() @@ -92,7 +92,7 @@ async fn main() -> anyhow::Result<()> { let purge_task = gw_task::spawn_content_purge(state.clone(), gw_task::PURGE_PERIOD); let rollup_task = gw_task::spawn_usage_rollup(state.clone(), gw_task::ROLLUP_PERIOD); let avail_task = gw_task::spawn_avail_flush(state.clone(), gw_task::AVAIL_FLUSH_PERIOD); - let distributed_batches = state.store.distributed_batches(); + let distributes_batches = state.store.distributes_batches(); let transport = select_transport()?; let postgres_url = cfg.storage.postgres_url.clone(); @@ -129,7 +129,7 @@ async fn main() -> anyhow::Result<()> { // fleet batch drain: on a distributed store any instance claims submitted batches let (batch_shutdown_tx, batch_shutdown_rx) = tokio::sync::watch::channel(false); - let batch_task = if distributed_batches { + let batch_task = if distributes_batches { let offline = app_state.offline.clone(); tracing::info!("batch drain loop started (distributed store)"); Some(tokio::spawn(async move { @@ -170,7 +170,7 @@ async fn main() -> anyhow::Result<()> { }); } - // SIGHUP → live reload (storage-backend changes still need a restart) + // a SIGHUP triggers a live reload (storage-backend changes still need a restart) #[cfg(unix)] { let app = app_state.clone(); @@ -286,7 +286,7 @@ async fn read_source_text(src: Option<&str>) -> Result, String } } -// GW_TRANSPORT: mock = zero egress, http = real HTTP, unset = mock:// in-process and real URLs over HTTP +// the GW_TRANSPORT env var: mock = zero egress, http = real HTTP, unset = mock:// in-process and real URLs over HTTP fn select_transport() -> anyhow::Result { Ok(match env::var("GW_TRANSPORT").as_deref() { Ok("mock") => { @@ -322,7 +322,7 @@ async fn shutdown_signal() { Ok(mut sig) => { sig.recv().await; } - Err(e) => tracing::error!("install SIGTERM handler: {e}"), + Err(e) => tracing::error!(error = %e, "install SIGTERM handler failed"), } }; #[cfg(not(unix))] diff --git a/crates/server/tests/compat.rs b/crates/server/tests/compat.rs index 8e4150e..a9668de 100644 --- a/crates/server/tests/compat.rs +++ b/crates/server/tests/compat.rs @@ -115,7 +115,7 @@ async fn live_chat_response_matches_canonical_key_sets() { .unwrap(); assert_eq!(resp.status(), StatusCode::OK); let ours = body_json(resp).await; - let canon: Value = serde_json::from_str(OPENAI_CHAT_CANONICAL).unwrap(); + let canon: Value = serde_json::from_str(OPENAI_CHAT_CANONICAL).expect("canonical parses"); assert_eq!(keys(&ours), keys(&canon), "top-level keys diverge"); assert_eq!( @@ -146,7 +146,7 @@ async fn live_messages_response_matches_canonical_key_sets() { .unwrap(); assert_eq!(resp.status(), StatusCode::OK); let ours = body_json(resp).await; - let canon: Value = serde_json::from_str(ANTHROPIC_MSG_CANONICAL).unwrap(); + let canon: Value = serde_json::from_str(ANTHROPIC_MSG_CANONICAL).expect("canonical parses"); assert_eq!(keys(&ours), keys(&canon), "top-level keys diverge"); assert_eq!( diff --git a/crates/server/tests/e2e.rs b/crates/server/tests/e2e.rs index 2432399..6a635ae 100644 --- a/crates/server/tests/e2e.rs +++ b/crates/server/tests/e2e.rs @@ -1347,6 +1347,7 @@ async fn pricing_dimensions_batch_discount_long_context_tier_and_per_image() { .unwrap(); assert_eq!(resp.status(), StatusCode::ACCEPTED); let id = body_json(resp).await["id"].as_str().unwrap().to_owned(); + let mut done = None; for _ in 0..100 { let j = body_json( app.clone() @@ -1356,10 +1357,12 @@ async fn pricing_dimensions_batch_discount_long_context_tier_and_per_image() { ) .await; if j["status"] == "completed" { + done = Some(j); break; } tokio::time::sleep(std::time::Duration::from_millis(10)).await; } + done.expect("batch finished"); let resp = app .clone() .oneshot(post( @@ -3624,13 +3627,13 @@ accounts: [{name: mock-openai-1, provider: openai, protocols: ["openai-chat"]}] async fn model_qpm_limit_third_call_429() { let app = app(); let body = r#"{"model":"qpm-mini","messages":[{"role":"user","content":"q"}]}"#; - for _ in 0..2 { + for i in 0..2 { let r = app .clone() .oneshot(post("/v1/chat/completions", Some("ak-demo-123"), body)) .await .unwrap(); - assert_eq!(r.status(), StatusCode::OK); + assert_eq!(r.status(), StatusCode::OK, "warm-up call {i}"); } let r = app .oneshot(post("/v1/chat/completions", Some("ak-demo-123"), body)) @@ -4364,13 +4367,13 @@ async fn bespoke_dashscope_native_wire() { async fn product_qpm_limit_third_call_429() { let app = app(); let body = r#"{"model":"gpt-4o","messages":[{"role":"user","content":"p"}]}"#; - for _ in 0..2 { + for i in 0..2 { let r = app .clone() .oneshot(post("/v1/chat/completions", Some("ak-prod-limited"), body)) .await .unwrap(); - assert_eq!(r.status(), StatusCode::OK); + assert_eq!(r.status(), StatusCode::OK, "warm-up call {i}"); } let r = app .oneshot(post("/v1/chat/completions", Some("ak-prod-limited"), body)) diff --git a/crates/server/tests/live_path.rs b/crates/server/tests/live_path.rs index dd2b402..ca7620e 100644 --- a/crates/server/tests/live_path.rs +++ b/crates/server/tests/live_path.rs @@ -40,8 +40,7 @@ async fn spawn_vendor() -> String { ); axum::response::Response::builder() .header("content-type", "text/event-stream") - .body(axum::body::Body::from(sse)) - .unwrap() + .body(axum::body::Body::from(sse)).expect("response") } else { let json = json!({ "id":"vendor-1","object":"chat.completion","created":1,"model":body["model"], @@ -77,10 +76,12 @@ async fn spawn_vendor() -> String { .unwrap() }), ); - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind"); + let addr = listener.local_addr().expect("addr"); tokio::spawn(async move { - axum::serve(listener, app).await.unwrap(); + axum::serve(listener, app).await.expect("serve"); }); format!("http://{addr}") } @@ -129,7 +130,7 @@ async fn full_pipeline_over_real_http() { .body(Body::from( r#"{"model":"gpt-live","messages":[{"role":"user","content":"is this live?"}]}"#, )) - .unwrap(); + .expect("response"); let resp = app.clone().oneshot(req).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK); let j = body_json(resp).await; @@ -150,7 +151,7 @@ async fn full_pipeline_over_real_http() { .uri("/internal/ledger") .header("authorization", "Bearer live-operator") .body(Body::empty()) - .unwrap(), + .expect("response"), ) .await .unwrap(); @@ -172,8 +173,7 @@ async fn streaming_pipeline_over_real_http() { .header("authorization", "Bearer ak-live") .body(Body::from( r#"{"model":"gpt-live","stream":true,"messages":[{"role":"user","content":"stream live"}]}"#, - )) - .unwrap(); + )).expect("response"); let resp = app.oneshot(req).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK); let ct = resp @@ -219,8 +219,7 @@ async fn claude_messages_over_real_http() { .header("authorization", "Bearer ak-live") .body(Body::from( r#"{"model":"claude-live","max_tokens":64,"messages":[{"role":"user","content":"is claude live?"}]}"#, - )) - .unwrap(); + )).expect("response"); let resp = app.clone().oneshot(req).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK); let j = body_json(resp).await; @@ -237,7 +236,7 @@ async fn claude_messages_over_real_http() { .uri("/internal/ledger") .header("authorization", "Bearer live-operator") .body(Body::empty()) - .unwrap(), + .expect("response"), ) .await .unwrap(); @@ -259,7 +258,7 @@ async fn auth_and_limits_still_apply_over_real_http() { .body(Body::from( r#"{"model":"gpt-live","messages":[{"role":"user","content":"x"}]}"#, )) - .unwrap() + .expect("response") }; let resp = app.clone().oneshot(chat("wrong")).await.unwrap(); assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); diff --git a/crates/server/tests/otel_trace.rs b/crates/server/tests/otel_trace.rs index ae2b37e..3577ff3 100644 --- a/crates/server/tests/otel_trace.rs +++ b/crates/server/tests/otel_trace.rs @@ -34,7 +34,7 @@ async fn request_span_exports_route_pipeline_fields_and_the_caller_context() { tracing::subscriber::set_global_default(subscriber).expect("first subscriber in this process"); opentelemetry::global::set_text_map_propagator(TraceContextPropagator::new()); - let cfg = Arc::new(GatewayConfig::embedded_default().unwrap()); + let cfg = Arc::new(GatewayConfig::embedded_default().expect("config")); let state = Arc::new(GatewayState::from_config(&cfg)); let app = gw_views::app(AppState::new( cfg, @@ -54,14 +54,14 @@ async fn request_span_exports_route_pipeline_fields_and_the_caller_context() { .body(Body::from( r#"{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}"#, )) - .unwrap(); + .expect("request"); let resp = app.clone().oneshot(chat).await.unwrap(); assert_eq!(resp.status(), StatusCode::OK); let miss = Request::builder() .uri("/v1/nothing") .header("authorization", "Bearer ak-demo-123") .body(Body::empty()) - .unwrap(); + .expect("request"); assert_eq!( app.oneshot(miss).await.unwrap().status(), StatusCode::NOT_FOUND diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index 7341245..d8c8dc1 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -95,7 +95,7 @@ impl BudgetScope { match self { Self::UserTokens => format!("{prefix}ub:{}:{user}", ak.tenant), Self::TenantCost => format!("{prefix}cb:tenant:{}", ak.tenant), - Self::KeyCost => format!("{prefix}cb:ak:{}", ak.ak), + Self::KeyCost => format!("{prefix}cb:ak:{}", ak.ak_id), Self::UserCost => format!("{prefix}cb:user:{}:{user}", ak.tenant), } } @@ -128,9 +128,9 @@ impl BillingLedger { pub(crate) fn repairing(store: Arc) -> Self { let (queue, mut pending) = mpsc::channel::(LEDGER_QUEUE_CAPACITY); - let deferred = store.deferred_ledger_writes(); + let deferred = store.defers_ledger_writes(); let worker_store = store.clone(); - // The bounded worker owns accepted rows through caller cancellation. + // the bounded worker owns accepted rows through caller cancellation tokio::spawn(async move { let mut batch = Vec::with_capacity(LEDGER_BATCH_MAX); let mut row_acks = Vec::with_capacity(LEDGER_BATCH_MAX); @@ -178,22 +178,27 @@ impl BillingLedger { rx.await.unwrap_or(0) } - async fn write(&self, record: &BillingRecord) { - if self.deferred + async fn write(&self, record: BillingRecord) { + let record = if self.deferred && let Some(queue) = &self.queue { - let queued = if record.user_id.is_empty() { - queue - .try_send(LedgerWrite::Row(record.clone(), None)) - .is_ok() - } else { - Self::queue_attributed(queue, record).await - }; - if queued { + if record.user_id.is_empty() { + match queue.try_send(LedgerWrite::Row(record, None)) { + Ok(()) => return, + Err(refused) => match refused.into_inner() { + LedgerWrite::Row(record, _) => record, + LedgerWrite::Flush(_) => return, + }, + } + } else if Self::queue_attributed(queue, &record).await { return; + } else { + record } - } - let Err(e) = self.store.ledger_add(record).await else { + } else { + record + }; + let Err(e) = self.store.ledger_add(&record).await else { return; }; metrics::counter!("gateway_ledger_write_failures_total").increment(1); @@ -202,11 +207,7 @@ impl BillingLedger { return; }; tracing::error!(error = %e, "billing ledger write failed; queued for repair"); - if queue - .send(LedgerWrite::Row(record.clone(), None)) - .await - .is_err() - { + if queue.send(LedgerWrite::Row(record, None)).await.is_err() { tracing::error!("billing ledger repair worker stopped"); } } @@ -469,16 +470,26 @@ pub async fn reserve_tpm( } } +/// What a settled request cost, for the budgets and the decision trail. +pub struct Settled { + pub total_tokens: i64, + pub cost_micros: i64, +} + /// Settle reserves to actuals, accrue the per-(AK, model) counter and write the /// ledger concurrently; a transient ledger failure goes to the bounded repair queue. pub async fn settle_and_bill( state: &GatewayState, cfg: &GatewayConfig, s: SettleInput<'_>, -) -> BillingRecord { +) -> Settled { let gov = state.governance.as_ref(); let total = clamp_tokens(s.billing.total); let record = billing_record(cfg, &s.billing); + let settled = Settled { + total_tokens: record.total_tokens, + cost_micros: record.cost_micros, + }; let settle_daily = gov.quota_settle(s.billing.ak, total - s.reserved, s.reserved_at); let consume_model = async { if let Some(key) = &s.model_quota_key { @@ -500,9 +511,9 @@ pub async fn settle_and_bill( None => {} } }; - let write_ledger = state.billing.write(&record); + let write_ledger = state.billing.write(record); tokio::join!(settle_daily, consume_model, settle_tpm, write_ledger); - record + settled } fn admit(ok: bool, deny: impl FnOnce() -> String) -> Result<(), String> { @@ -658,7 +669,7 @@ mod tests { store.fail_next_ledger_writes(2); let ledger = BillingLedger::repairing(store.clone()); let record = record("req-repair"); - ledger.write(&record).await; + ledger.write(record.clone()).await; let (count, rows) = tokio::time::timeout(Duration::from_secs(2), async { loop { let snapshot = store.ledger_snapshot(usize::MAX).await.unwrap(); @@ -684,7 +695,7 @@ mod tests { let mut row = record("req-attributed"); row.user_id = "user-42".into(); - ledger.write(&row).await; + ledger.write(row.clone()).await; let (count, rows) = store.ledger_snapshot(usize::MAX).await.unwrap(); assert_eq!(count, 1); @@ -702,7 +713,7 @@ mod tests { let mut row = record("req-fallback"); row.user_id = "user-42".into(); - ledger.write(&row).await; + ledger.write(row.clone()).await; let (count, rows) = store.ledger_snapshot(usize::MAX).await.unwrap(); assert_eq!(count, 1); @@ -735,7 +746,7 @@ mod tests { let ledger = BillingLedger::repairing(store.clone()); let queue = ledger.queue.as_ref().unwrap(); - ledger.write(&record("req-0")).await; + ledger.write(record("req-0")).await; tokio::time::timeout(Duration::from_secs(1), async { while queue.capacity() != LEDGER_QUEUE_CAPACITY { tokio::task::yield_now().await; @@ -745,12 +756,12 @@ mod tests { .expect("repair worker did not take the first row"); for i in 1..=LEDGER_QUEUE_CAPACITY { - ledger.write(&record(format!("req-{i}"))).await; + ledger.write(record(format!("req-{i}"))).await; } assert_eq!(queue.capacity(), 0); let blocked_record = record(format!("req-{}", LEDGER_QUEUE_CAPACITY + 1)); - let mut blocked = Box::pin(ledger.write(&blocked_record)); + let mut blocked = Box::pin(ledger.write(blocked_record.clone())); assert!( tokio::time::timeout(Duration::from_millis(20), blocked.as_mut()) .await @@ -781,7 +792,7 @@ mod tests { let store = Arc::new(crate::MemoryStore::default()); store.fail_next_ledger_writes(1); let ledger = BillingLedger::repairing(store.clone()); - ledger.write(&record("req-flush")).await; + ledger.write(record("req-flush")).await; ledger.flush().await; let (count, rows) = store.ledger_snapshot(usize::MAX).await.unwrap(); assert_eq!(count, 1); @@ -796,7 +807,7 @@ mod tests { tokio::time::timeout(Duration::from_secs(3_600), async { for i in 0..=LEDGER_QUEUE_CAPACITY { - ledger.write(&record(format!("req-{i}"))).await; + ledger.write(record(format!("req-{i}"))).await; } }) .await @@ -851,9 +862,12 @@ mod tests { let prev = BudgetScope::KeyCost.key(Some(previous_month(month)), &ak, ""); assert_eq!( BudgetScope::KeyCost.key(Some(month), &ak, ""), - format!("m:{y}{m:02}:cb:ak:k1") + format!("m:{y}{m:02}:cb:ak:{}", ak.ak_id) + ); + assert_eq!( + BudgetScope::KeyCost.key(None, &ak, ""), + format!("cb:ak:{}", ak.ak_id) ); - assert_eq!(BudgetScope::KeyCost.key(None, &ak, ""), "cb:ak:k1"); let untouched = budgets(gov, &cfg, &ak, "").await; assert_eq!(untouched.len(), 1); @@ -890,7 +904,7 @@ mod tests { err.starts_with("monthly cost budget exhausted for key:"), "{err}" ); - assert_eq!(gov.quota_used("cb:ak:k1").await, 0); + assert_eq!(gov.quota_used(&format!("cb:ak:{}", ak.ak_id)).await, 0); gov.quota_reset_all().await; assert!( check_budgets(gov, &cfg, &ak, "").await.is_err(), diff --git a/crates/state/src/configstore.rs b/crates/state/src/configstore.rs index 99c2597..a10048d 100644 --- a/crates/state/src/configstore.rs +++ b/crates/state/src/configstore.rs @@ -124,7 +124,7 @@ impl PostgresConfigStore { Ok(tx) } - // NOTIFY is transactional: peers hear the id only after the commit + // a NOTIFY is transactional: peers hear the id only after the commit async fn insert_notify( tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, yaml: &str, diff --git a/crates/state/src/content.rs b/crates/state/src/content.rs index 16f4ac8..a9ed87e 100644 --- a/crates/state/src/content.rs +++ b/crates/state/src/content.rs @@ -40,7 +40,7 @@ pub struct ContentRecord { } /// Whether a deployment key is configured (so `full` retention may store raw). -pub fn sealing_available() -> bool { +pub fn can_seal() -> bool { CIPHER.is_some() } diff --git a/crates/state/src/governance.rs b/crates/state/src/governance.rs index 874aa73..f917879 100644 --- a/crates/state/src/governance.rs +++ b/crates/state/src/governance.rs @@ -140,7 +140,6 @@ impl Governance for MemoryGovernance { /// Redis-backed governance for multi-replica deployments. Keys are namespaced /// under `gw:`; windows use INCR + EXPIRE so they self-expire. -#[derive(Clone)] pub struct RedisGovernance { conn: redis::aio::ConnectionManager, } diff --git a/crates/state/src/keystore.rs b/crates/state/src/keystore.rs index be06420..c90bde6 100644 --- a/crates/state/src/keystore.rs +++ b/crates/state/src/keystore.rs @@ -79,9 +79,13 @@ impl PostgresKeyStore { banned BOOLEAN NOT NULL DEFAULT FALSE, model_quotas TEXT NOT NULL DEFAULT '{}', owner TEXT, - source TEXT NOT NULL DEFAULT 'admin')", + source TEXT NOT NULL DEFAULT 'admin', + mcp_servers TEXT NOT NULL DEFAULT '[]', + mcp_tools TEXT NOT NULL DEFAULT '{}')", "ALTER TABLE access_keys ADD COLUMN IF NOT EXISTS owner TEXT", "ALTER TABLE access_keys ADD COLUMN IF NOT EXISTS suspended_until_epoch_secs BIGINT", + "ALTER TABLE access_keys ADD COLUMN IF NOT EXISTS mcp_servers TEXT NOT NULL DEFAULT '[]'", + "ALTER TABLE access_keys ADD COLUMN IF NOT EXISTS mcp_tools TEXT NOT NULL DEFAULT '{}'", ], ) .await?; @@ -104,7 +108,7 @@ impl PostgresKeyStore { async fn fetch(&self, ak: &str) -> Result>, sqlx::Error> { let row = sqlx::query( "SELECT ak, product, tenant, qps, daily_token_quota, tokens_per_minute, - expires_at_epoch_secs, banned, model_quotas, owner, suspended_until_epoch_secs FROM access_keys WHERE ak = $1", + expires_at_epoch_secs, banned, model_quotas, owner, suspended_until_epoch_secs, mcp_servers, mcp_tools FROM access_keys WHERE ak = $1", ) .bind(ak) .fetch_optional(&self.pool) @@ -153,7 +157,7 @@ impl KeyStore for PostgresKeyStore { } async fn patch(&self, ak: &str, patch: &KeyPatch) -> GResult> { - // FOR UPDATE: concurrent patches serialize instead of clobbering fields + // concurrent patches serialize under FOR UPDATE instead of clobbering fields let mut tx = self .pool .begin() @@ -161,7 +165,7 @@ impl KeyStore for PostgresKeyStore { .map_err(|e| crate::sqlx_err("begin patch", e))?; let row = sqlx::query( "SELECT ak, product, tenant, qps, daily_token_quota, tokens_per_minute, - expires_at_epoch_secs, banned, model_quotas, owner, suspended_until_epoch_secs FROM access_keys + expires_at_epoch_secs, banned, model_quotas, owner, suspended_until_epoch_secs, mcp_servers, mcp_tools FROM access_keys WHERE ak = $1 FOR UPDATE", ) .bind(ak) @@ -213,7 +217,7 @@ impl KeyStore for PostgresKeyStore { ) -> GResult> { let rows = sqlx::query( "SELECT ak, product, tenant, qps, daily_token_quota, tokens_per_minute, - expires_at_epoch_secs, banned, model_quotas, owner, suspended_until_epoch_secs FROM access_keys + expires_at_epoch_secs, banned, model_quotas, owner, suspended_until_epoch_secs, mcp_servers, mcp_tools FROM access_keys WHERE ($1::text IS NULL OR tenant = $1) ORDER BY ak LIMIT $2 OFFSET $3", ) .bind(tenant) @@ -260,10 +264,13 @@ async fn upsert( source: KeySource, ) -> Result<(), sqlx::Error> { let quotas = serde_json::to_string(&*info.model_quotas).unwrap_or_else(|_| "{}".into()); + let servers = serde_json::to_string(&info.mcp.servers).unwrap_or_else(|_| "[]".into()); + let tools = serde_json::to_string(&info.mcp.tools).unwrap_or_else(|_| "{}".into()); sqlx::query( "INSERT INTO access_keys (ak, product, tenant, qps, daily_token_quota, - tokens_per_minute, expires_at_epoch_secs, banned, model_quotas, owner, source) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) + tokens_per_minute, expires_at_epoch_secs, banned, model_quotas, owner, source, + mcp_servers, mcp_tools) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13) ON CONFLICT (ak) DO UPDATE SET product = EXCLUDED.product, tenant = EXCLUDED.tenant, qps = EXCLUDED.qps, daily_token_quota = EXCLUDED.daily_token_quota, @@ -271,6 +278,7 @@ async fn upsert( expires_at_epoch_secs = EXCLUDED.expires_at_epoch_secs, banned = EXCLUDED.banned, model_quotas = EXCLUDED.model_quotas, owner = EXCLUDED.owner, + mcp_servers = EXCLUDED.mcp_servers, mcp_tools = EXCLUDED.mcp_tools, source = CASE WHEN access_keys.source = 'config' AND EXCLUDED.source = 'admin' THEN 'config' ELSE EXCLUDED.source END", ) @@ -285,6 +293,8 @@ async fn upsert( .bind("as) .bind(&info.owner) .bind(source_str(source)) + .bind(&servers) + .bind(&tools) .execute(exec) .await .map(|_| ()) @@ -307,7 +317,10 @@ fn row_to_info(row: &sqlx::postgres::PgRow) -> AkInfo { ), owner: row.get(9), suspended_until_epoch_secs: row.get(10), - mcp: Default::default(), + mcp: std::sync::Arc::new(crate::McpAccess { + servers: serde_json::from_str(row.get::<&str, _>(11)).unwrap_or_default(), + tools: serde_json::from_str(row.get::<&str, _>(12)).unwrap_or_default(), + }), } } @@ -413,5 +426,30 @@ mod tests { "config ownership is sticky against admin overwrite" ); assert!(ks.authenticate("pk-b").await.is_some()); + + let cfg = gw_config::GatewayConfig::from_yaml( + "listen: {host: h, port: 1}\nmcp_servers: [{name: srv, endpoint: http://h/mcp}]\naccess_keys: [{ak: pk-mcp, product: p, qps: 1, daily_token_quota: 5, mcp_servers: [srv], mcp_tools: {srv: [t1]}}]", + ) + .unwrap(); + ks.reload_config_keys(&cfg.access_keys).await.unwrap(); + let k = ks + .authenticate("pk-mcp") + .await + .expect("config key persisted"); + assert_eq!(k.mcp.servers, vec!["srv".to_owned()]); + assert_eq!(k.mcp.tools["srv"], vec!["t1".to_owned()]); + + let mut admin = info("pk-plain", 1.0); + admin.mcp = std::sync::Arc::new(crate::McpAccess { + servers: vec!["srv".into()], + tools: Default::default(), + }); + ks.put(admin, KeySource::Admin).await.unwrap(); + let k = ks + .authenticate("pk-plain") + .await + .expect("admin key persisted"); + assert_eq!(k.mcp.servers, vec!["srv".to_owned()]); + assert!(k.mcp.tools.is_empty()); } } diff --git a/crates/state/src/lib.rs b/crates/state/src/lib.rs index 01cb198..d31acca 100644 --- a/crates/state/src/lib.rs +++ b/crates/state/src/lib.rs @@ -30,7 +30,7 @@ pub mod thinking_signature; pub use alerts::{AlertBus, AlertEvent}; pub use avail::{AvailState, AvailStore, classify}; pub use configstore::{CONFIG_CHANNEL, PostgresConfigStore}; -pub use content::{ContentRecord, sealing_available}; +pub use content::{ContentRecord, can_seal}; pub use governance::{Governance, MemoryGovernance, RedisGovernance}; pub use health::{HealthStore, RedisHealth}; pub use keystore::{KeyStore, PostgresKeyStore}; @@ -461,15 +461,9 @@ impl AccountPool { health: &dyn HealthStore, latency: Option<&latency::Latency>, ) -> Option> { - let candidates: Vec<&Arc> = self - .accounts - .iter() - .filter(|a| serves(a, p, provider)) - .collect(); - let checks = - futures::future::join_all(candidates.iter().map(|a| health.available(&a.name))).await; - let unhealthy: Vec<&str> = candidates - .iter() + let serving = || self.accounts.iter().filter(|a| serves(a, p, provider)); + let checks = futures::future::join_all(serving().map(|a| health.available(&a.name))).await; + let unhealthy: Vec<&str> = serving() .zip(checks) .filter(|(_, ok)| !ok) .map(|(a, _)| a.name.as_str()) @@ -1380,7 +1374,7 @@ mod tests { .await .unwrap(); assert!(st.store.file_get(&f.id).await.unwrap().is_some()); - assert!(!st.store.distributed_batches()); + assert!(!st.store.distributes_batches()); } #[test] diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index 6422f17..83ce92d 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -39,6 +39,9 @@ const ROLLUP_BACKFILL_SECS: i64 = 20 * 60; /// trailing replica lands a row in a rolled minute — a rolled minute's source /// set can only shrink, which keeps the max-upsert sound. const ROLLUP_SETTLE_SECS: i64 = ROLLUP_BUCKET_SECS; +/// In-process rollup retention: the monthly-budget window, then a bucket cap. +const ROLLUP_RETENTION_SECS: i64 = 62 * 86_400; +const ROLLUP_MAX_BUCKETS: usize = 1_000_000; /// Postgres advisory-lock key serializing the fleet's rollup: one replica /// advances per tick (the upsert is idempotent; the lock only avoids repeated scans). @@ -50,6 +53,8 @@ const ROLLUP_WATERMARK_SQL: &str = "SELECT COALESCE(MAX(minute_epoch), -60) + 60 /// A put prunes async video jobs older than this (vendor results expire far sooner). const VIDEO_JOB_RETENTION_SECS: i64 = 30 * 24 * 3600; +/// Rows per batch_items INSERT: four binds each under the 65535-parameter limit. +const BATCH_ITEM_CHUNK: usize = 16_000; // store pool size when storage.postgres_max_connections is unset const PG_MAX_CONNECTIONS: u32 = 10; @@ -275,9 +280,7 @@ pub fn billing_record(cfg: &gw_config::GatewayConfig, b: &BillingInput) -> Billi } }; let (vendor, vendor_unit) = cfg - .accounts - .iter() - .find(|a| a.name == b.account) + .find_account(b.account) .map(|a| { ( ( @@ -525,7 +528,7 @@ pub trait Store: Send + Sync + std::fmt::Debug { Ok(()) } /// Whether rows may batch off the request path; in-process backends stay synchronous. - fn deferred_ledger_writes(&self) -> bool { + fn defers_ledger_writes(&self) -> bool { false } /// Total count plus the most recent `limit` records in chronological order; @@ -667,7 +670,7 @@ pub trait Store: Send + Sync + std::fmt::Debug { /// Whether this backend runs a fleet work queue; local backends execute on /// the submitting instance. - fn distributed_batches(&self) -> bool { + fn distributes_batches(&self) -> bool { false } /// Atomically enqueue a batch and its items so a partial save never leaves @@ -769,6 +772,7 @@ pub struct MemoryStore { /// Minute buckets keyed by (minute, tenant, user, model); see /// [`Store::usage_rollup_advance`]. rollup: Mutex>, + rollup_max_buckets: Option, sec_events: Mutex>, audit: Mutex>, content: Mutex, @@ -792,6 +796,13 @@ impl MemoryStore { } } + pub fn with_rollup_cap(max_buckets: usize) -> Self { + Self { + rollup_max_buckets: Some(max_buckets), + ..Self::default() + } + } + #[cfg(test)] pub(crate) fn fail_next_ledger_writes(&self, count: usize) { self.ledger_failures.store(count, Ordering::Relaxed); @@ -814,13 +825,13 @@ impl Store for MemoryStore { )); } // watermark first: rollup-then-records is the lock order advance uses - let watermark = rollup_watermark(&lock(&self.rollup)); + let watermark = (self.ledger_max_rows > 0).then(|| rollup_watermark(&lock(&self.rollup))); let mut ledger = lock(&self.ledger); if !ledger.request_ids.insert(r.request_id.clone()) { return Ok(()); } ledger.rows.push(r.clone()); - if self.ledger_max_rows > 0 + if let Some(watermark) = watermark && ledger.rows.len() > self.ledger_max_rows && self .prune_seq @@ -991,6 +1002,13 @@ impl Store for MemoryStore { for (k, v) in fresh { rollup.entry(k).and_modify(|e| e.keep_max(&v)).or_insert(v); } + let floor = bucket_floor(now - ROLLUP_RETENTION_SECS); + let kept = rollup.split_off(&(floor, String::new(), String::new(), String::new())); + *rollup = kept; + let cap = self.rollup_max_buckets.unwrap_or(ROLLUP_MAX_BUCKETS); + while rollup.len() > cap { + rollup.pop_first(); + } Ok(written) } @@ -1367,6 +1385,7 @@ impl SqliteStore { served_model TEXT NOT NULL, account TEXT NOT NULL, unit_price_micros INTEGER NOT NULL, created_at_epoch_secs INTEGER NOT NULL, billed INTEGER NOT NULL DEFAULT 0)", + "CREATE INDEX IF NOT EXISTS video_jobs_created_idx ON video_jobs (created_at_epoch_secs)", "CREATE TABLE IF NOT EXISTS batches ( n INTEGER PRIMARY KEY AUTOINCREMENT, id TEXT UNIQUE NOT NULL, ak TEXT NOT NULL, tenant TEXT NOT NULL DEFAULT 'default', model TEXT NOT NULL, @@ -1804,7 +1823,7 @@ macro_rules! sql_store_impl { if rows.is_empty() { return Ok(()); } - // QueryBuilder numbers the placeholders per dialect, so no dialect_sql! here + // the QueryBuilder numbers placeholders per dialect, so no dialect_sql! here let mut qb = sqlx::QueryBuilder::new( "INSERT INTO billing (ak, product, tenant, model, served_model, protocol, account, prompt_tokens, completion_tokens, total_tokens, cost_micros, @@ -1839,7 +1858,7 @@ macro_rules! sql_store_impl { Ok(()) } - fn deferred_ledger_writes(&self) -> bool { + fn defers_ledger_writes(&self) -> bool { true } @@ -2410,6 +2429,7 @@ impl PostgresStore { served_model TEXT NOT NULL, account TEXT NOT NULL, unit_price_micros BIGINT NOT NULL, created_at_epoch_secs BIGINT NOT NULL, billed INTEGER NOT NULL DEFAULT 0)", + "CREATE INDEX IF NOT EXISTS video_jobs_created_idx ON video_jobs (created_at_epoch_secs)", "CREATE TABLE IF NOT EXISTS batches ( n BIGSERIAL PRIMARY KEY, id TEXT UNIQUE NOT NULL, ak TEXT NOT NULL, tenant TEXT NOT NULL DEFAULT 'default', model TEXT NOT NULL, @@ -2717,7 +2737,7 @@ sql_store_impl!(PostgresStore, postgres, { Ok(Some(done)) } - fn distributed_batches(&self) -> bool { + fn distributes_batches(&self) -> bool { true } @@ -2743,18 +2763,21 @@ sql_store_impl!(PostgresStore, postgres, { .fetch_one(&mut *tx) .await .map_err(|e| crate::sqlx_err("insert batch", e))?; - for (idx, item) in items.iter().enumerate() { - let json = serde_json::to_string(&item.messages).unwrap_or_else(|_| "[]".into()); - sqlx::query( - "INSERT INTO batch_items (batch_id, idx, messages, user_id) VALUES ($1, $2, $3, $4)", - ) - .bind(&id) - .bind(idx as i64) - .bind(json) - .bind(&item.user) - .execute(&mut *tx) - .await - .map_err(|e| crate::sqlx_err("save batch item", e))?; + for (chunk, items) in items.chunks(BATCH_ITEM_CHUNK).enumerate() { + let mut qb = sqlx::QueryBuilder::new( + "INSERT INTO batch_items (batch_id, idx, messages, user_id) ", + ); + qb.push_values(items.iter().enumerate(), |mut v, (i, item)| { + let json = serde_json::to_string(&item.messages).unwrap_or_else(|_| "[]".into()); + v.push_bind(&id) + .push_bind((chunk * BATCH_ITEM_CHUNK + i) as i64) + .push_bind(json) + .push_bind(&item.user); + }); + qb.build() + .execute(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("save batch items", e))?; } tx.commit() .await @@ -4071,6 +4094,38 @@ mod tests { assert_eq!(total, 2, "the cap holds after the prune"); } + #[tokio::test] + async fn memory_rollup_keeps_the_retention_window_and_caps_buckets() { + let now = 100 * 86_400; + let store = MemoryStore::default(); + let mut old = record("m1"); + old.created_at_epoch_secs = now - ROLLUP_RETENTION_SECS - 3_600; + let mut fresh = record("m1"); + fresh.created_at_epoch_secs = now - 3_600; + for r in [&old, &fresh] { + store.ledger_add(r).await.unwrap(); + } + store.usage_rollup_advance(now).await.unwrap(); + let minutes: Vec = lock(&store.rollup).keys().map(|k| k.0).collect(); + assert_eq!(minutes, vec![bucket_floor(fresh.created_at_epoch_secs)]); + + let store = MemoryStore::with_rollup_cap(2); + for i in 1..=3 { + let mut r = record("m1"); + r.request_id = format!("req-{i}"); + r.created_at_epoch_secs = now - i * 3_600; + store.ledger_add(&r).await.unwrap(); + } + store.usage_rollup_advance(now).await.unwrap(); + let minutes: Vec = lock(&store.rollup).keys().map(|k| k.0).collect(); + assert_eq!(minutes.len(), 2, "the cap holds"); + assert_eq!( + minutes[0], + bucket_floor(now - 2 * 3_600), + "the oldest minute is the one dropped" + ); + } + #[tokio::test] async fn ledger_batch_insert_is_multi_row_and_idempotent() { let dir = tempfile::tempdir().unwrap(); @@ -4268,7 +4323,7 @@ mod tests { ); assert_eq!(got.status, BatchStatus::Failed); - assert!(store.distributed_batches()); + assert!(store.distributes_batches()); let qmsgs = vec![ gw_models::BatchItem { messages: vec![gw_models::ChatMsg::text("user", "one")], diff --git a/crates/state/src/thinking_signature.rs b/crates/state/src/thinking_signature.rs index cc185b7..8f872b1 100644 --- a/crates/state/src/thinking_signature.rs +++ b/crates/state/src/thinking_signature.rs @@ -140,8 +140,8 @@ impl ThinkingSignatureAudit { }; let messages = &request.message; - // Anthropic validates only the latest logical assistant turn (older thinking may be - // omitted); consecutive same-role messages form one turn + // only the latest logical assistant turn is validated upstream, so older thinking may be omitted + // consecutive same-role messages form one turn let mut trailing_user_start = messages.len(); while trailing_user_start > 0 && messages[trailing_user_start - 1].role == gw_consts::role::USER @@ -543,7 +543,7 @@ impl CapturedBlock { } } - fn complete(&self) -> bool { + fn is_complete(&self) -> bool { match self { Self::Thinking { complete, .. } | Self::RedactedThinking { complete, .. } @@ -709,7 +709,7 @@ impl ThinkingStreamCapture { } fn register(&mut self) { - if self.disabled || self.blocks.values().any(|block| !block.complete()) { + if self.disabled || self.blocks.values().any(|block| !block.is_complete()) { return; } let mut sequence = ProtectedSequence::default(); @@ -739,7 +739,6 @@ impl ThinkingStreamCapture { } } self.audit.remember_sequence(&self.context, &sequence); - drop(sequence); self.blocks.clear(); self.captured_bytes = 0; self.registered = true; @@ -779,11 +778,12 @@ fn replace_bounded(target: &mut String, value: Option<&Value>) -> Option<()> { #[cfg(test)] mod tests { - use super::*; use gw_consts::Protocol; use gw_models::ModelParamV2; use serde_json::json; + use super::*; + fn message(role: &str, content: Value) -> ChatMsg { let mut message = ChatMsg::text(role, String::new()); message.parts = Some(content); diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index e047784..5f32ea5 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -73,8 +73,20 @@ pub type ConfigFuture = /// Reloads config from its source (file or the Postgres config store). pub type ConfigLoader = Arc ConfigFuture + Send + Sync>; +/// Handler state behind one Arc: axum clones it twice per request. #[derive(Clone)] -pub struct AppState { +pub struct AppState(Arc); + +impl std::ops::Deref for AppState { + type Target = AppInner; + + fn deref(&self) -> &AppInner { + &self.0 + } +} + +#[derive(Clone)] +pub struct AppInner { pub handler: OnlineHandler, pub offline: OfflineHandler, /// Client for the `/mcp/{server}` proxy; per-server timeouts apply per request. @@ -105,7 +117,7 @@ impl AppState { ) -> Self { let handler = OnlineHandler::new(config, transport); let offline = OfflineHandler::new(handler.clone()); - Self { + Self(Arc::new(AppInner { handler, offline, mcp: mcp_client(), @@ -113,12 +125,12 @@ impl AppState { mcp_sessions: mcp_sessions(), loader, config_store: None, - } + })) } /// Attach the fleet config store (enables `PUT /admin/config`). pub fn with_config_store(mut self, store: Arc) -> Self { - self.config_store = Some(store); + Arc::make_mut(&mut self.0).config_store = Some(store); self } @@ -135,13 +147,13 @@ impl AppState { /// The MCP proxy's client: no redirects, so a server or token endpoint cannot /// steer a credentialed request elsewhere. fn mcp_client() -> reqwest::Client { - reqwest::Client::builder() + #[allow(clippy::expect_used)] + // build fails only when TLS cannot initialize, where Client::new panics too + let client = reqwest::Client::builder() .redirect(reqwest::redirect::Policy::none()) .build() - .unwrap_or_else(|e| { - tracing::error!(error = %e, "mcp client fell back to the default client"); - reqwest::Client::new() - }) + .expect("mcp client builds"); + client } fn mcp_sessions() -> moka::sync::Cache> { @@ -339,7 +351,7 @@ fn rt_error_frame( } /// The AK carried as `gw-api-key.` in `Sec-WebSocket-Protocol`, the one header a browser can set. -fn ws_subprotocol_ak(headers: &HeaderMap) -> Option { +fn ws_subprotocol_ak(headers: &HeaderMap) -> Option<&str> { headers .get("sec-websocket-protocol")? .to_str() @@ -347,7 +359,6 @@ fn ws_subprotocol_ak(headers: &HeaderMap) -> Option { .split(',') .map(str::trim) .find_map(|p| p.strip_prefix("gw-api-key.")) - .map(str::to_owned) } /// GET /v1/realtime (WebSocket upgrade): bridge to the vendor's socket, or the mock for an endpoint-less account. @@ -363,7 +374,7 @@ async fn realtime_ws( Ok(ak) => ak, Err((st, msg)) => { let sub = match ws_subprotocol_ak(&headers) { - Some(k) => snap.state.auth.authenticate(&k).await, + Some(k) => snap.state.auth.authenticate(k).await, None => None, }; match sub { @@ -639,7 +650,7 @@ async fn bill_realtime_turn( let total = gw_state::clamp_tokens(bp.saturating_add(bc)); let model_quota_key = admission::model_quota_limit(cfg, ak, &m.requested) .map(|_| admission::model_quota_key(&ak.ak, &m.requested)); - let record = admission::settle_and_bill( + let settled = admission::settle_and_bill( state, cfg, admission::SettleInput { @@ -678,7 +689,7 @@ async fn bill_realtime_turn( ak, admit.user.as_str(), total, - record.cost_micros, + settled.cost_micros, ) .await; if !estimated { @@ -788,18 +799,16 @@ async fn realtime_session( } turn.record_text(delta); } - if socket + let sent = socket .send(send(json!({"type":"response.done", "usage":{"input_tokens": it, "output_tokens": ot}}))) .await - .is_err() - { - bill_realtime_turn(&turn.admit, &rtm, mt, &account, turn_tokens(it, ot), false) - .await; - return; - } + .is_ok(); bill_realtime_turn(&turn.admit, &rtm, mt, &account, turn_tokens(it, ot), false) .await; + if !sent { + return; + } } "session.close" => { let _ = socket.send(send(json!({"type":"session.closed"}))).await; @@ -862,7 +871,7 @@ async fn realtime_bridge( .replacen("http://", "ws://", 1); let key = account.api_key().unwrap_or_else(|| "mock".to_owned()); let gemini = gw_engines::realtime::is_gemini_realtime(account.wire_kind()); - // Gemini's Live socket is one bidi RPC authed by key; the model rides the setup frame + // a Gemini Live socket is one bidi RPC authed by key; the model rides the setup frame let url = if gemini { format!( "{ws_base}/ws/google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent?key={key}" @@ -909,7 +918,7 @@ async fn realtime_bridge( let mut pending: Option = None; // denied server-VAD turn: swallow its upstream frames until its terminal frame let mut suppress = false; - // Gemini sends cumulative usage on any server frame; the latest settles a bare turnComplete + // cumulative usage may arrive on any server frame; the latest settles a bare turnComplete let mut usage_snapshot: Option = None; // outbound DLP redactions summed within a turn, recorded once at its boundary let mut out_redacted = 0i64; @@ -1041,24 +1050,11 @@ async fn realtime_bridge( v["usageMetadata"] = u; } // turn boundary: settle the admitted turn; an ungated one bills unreserved - match pending.take() { - Some(turn) if it.saturating_add(ot) > 0 => { - bill_realtime_turn( - &turn.admit, - &rtm, - mt, - &account.name, - gw_models::TokenInput { - prompt: it, - completion: ot, - ..turn_audio(account.wire_kind(), &v) - }, - false, - ) - .await - } + let admit = match pending.take() { + Some(turn) if it.saturating_add(ot) > 0 => Some(turn.admit), Some(turn) => { - settle_realtime_abort(turn, &rtm, mt, &account.name).await + settle_realtime_abort(turn, &rtm, mt, &account.name).await; + None } None if it.saturating_add(ot) > 0 => { // re-authenticate so billing uses the key's current identity @@ -1070,7 +1066,7 @@ async fn realtime_bridge( .await .unwrap_or_else(|| ak.clone()); let user = billed.attributed_user(&hint).to_owned(); - let unreserved = RealtimeAdmit { + Some(RealtimeAdmit { ak: billed, user, reserved: 0, @@ -1078,22 +1074,24 @@ async fn realtime_bridge( at: gw_state::epoch_secs(), request_id: gw_handler::new_request_id(), snap, - }; - bill_realtime_turn( - &unreserved, - &rtm, - mt, - &account.name, - gw_models::TokenInput { - prompt: it, - completion: ot, - ..turn_audio(account.wire_kind(), &v) - }, - false, - ) - .await + }) } - None => {} + None => None, + }; + if let Some(admit) = admit { + bill_realtime_turn( + &admit, + &rtm, + mt, + &account.name, + gw_models::TokenInput { + prompt: it, + completion: ot, + ..turn_audio(account.wire_kind(), &v) + }, + false, + ) + .await; } recognized += 1; turn_ended = true; @@ -1284,12 +1282,9 @@ async fn list_models(State(s): State, Authed(ak): Authed) -> Response /// tenant and carry the operator's vendor-cost margin basis. async fn ledger( State(s): State, - headers: HeaderMap, + _: GlobalAdmin, Query(q): Query>, ) -> Response { - if let Err(r) = require_global_admin(&s, &headers) { - return r; - } let limit = q_num(&q, "limit", LEDGER_PAGE_DEFAULT).min(ADMIN_PAGE_MAX); match s.handler.state().store.ledger_snapshot(limit).await { Ok((count, records)) => Json(json!({ "count": count, "records": records })).into_response(), @@ -1299,10 +1294,7 @@ async fn ledger( /// Account pool view (name/provider/tier/priority/served model family). /// Global-token only: account names and health are operator internals. -async fn accounts(State(s): State, headers: HeaderMap) -> Response { - if let Err(r) = require_global_admin(&s, &headers) { - return r; - } +async fn accounts(State(s): State, _: GlobalAdmin) -> Response { let cfg = s.handler.cfg(); let health = &s.handler.state().health; let mut data: Vec = Vec::with_capacity(cfg.accounts.len()); @@ -1333,8 +1325,8 @@ fn user_header(headers: &HeaderMap) -> Option { /// The REST attribution precedence: `x-gw-user` header, else the dialect's own /// user field (batch items invert it — per-item `user` first). -fn user_hint(headers: &HeaderMap, field: &Value) -> Option { - user_header(headers).or_else(|| field.as_str().map(cap_user_hint)) +fn user_hint(hint: Option, field: &Value) -> Option { + hint.or_else(|| field.as_str().map(cap_user_hint)) } /// Bound an attribution hint at `USER_HINT_MAX_BYTES`: it keys governance @@ -1412,6 +1404,34 @@ impl axum::extract::FromRequestParts for Authed { } } +/// The `x-gw-user` attribution hint, read without cloning the header map. +pub struct UserHint(pub Option); + +impl axum::extract::FromRequestParts for UserHint { + type Rejection = Infallible; + + async fn from_request_parts( + parts: &mut axum::http::request::Parts, + _: &AppState, + ) -> Result { + Ok(UserHint(user_header(&parts.headers))) + } +} + +/// Proof of the global admin token; runs before any body extractor. +pub struct GlobalAdmin; + +impl axum::extract::FromRequestParts for GlobalAdmin { + type Rejection = Response; + + async fn from_request_parts( + parts: &mut axum::http::request::Parts, + s: &AppState, + ) -> Result { + require_global_admin(s, &parts.headers).map(|()| GlobalAdmin) + } +} + /// The shared body behind [`ApiJson`]/[`AnthJson`]: `axum::Json` with the /// rejection rendered through the surface's own envelope. #[allow(clippy::result_large_err)] // once per request; boxing would noise every call site @@ -1855,12 +1875,9 @@ fn ct_eq(a: &str, b: &str) -> bool { /// governance, store, health, and cache are preserved. async fn admin_reload( State(s): State, - headers: HeaderMap, + _: GlobalAdmin, AuditSourceIp(source): AuditSourceIp, ) -> Response { - if let Err(r) = require_global_admin(&s, &headers) { - return r; - } match s.reload().await { Ok(()) => { let cfg = s.handler.cfg(); @@ -2045,10 +2062,7 @@ async fn admin_key_delete( } /// GET /admin/config — the current fleet config document. Global admin only. -async fn admin_config_get(State(s): State, headers: HeaderMap) -> Response { - if let Err(r) = require_global_admin(&s, &headers) { - return r; - } +async fn admin_config_get(State(s): State, _: GlobalAdmin) -> Response { let store = match require_config_store(&s) { Ok(v) => v, Err(r) => return r, @@ -2065,14 +2079,7 @@ async fn admin_config_get(State(s): State, headers: HeaderMap) -> Resp } /// POST /admin/config/validate — parse and validate without publishing. -async fn admin_config_validate( - State(s): State, - headers: HeaderMap, - body: String, -) -> Response { - if let Err(r) = require_global_admin(&s, &headers) { - return r; - } +async fn admin_config_validate(_: GlobalAdmin, body: String) -> Response { match GatewayConfig::from_yaml(&body) { Ok(cfg) => Json(json!({ "valid": true, @@ -2090,12 +2097,9 @@ async fn admin_config_validate( /// GET /admin/config/versions — retained config heads, newest first. async fn admin_config_versions( State(s): State, - headers: HeaderMap, + _: GlobalAdmin, Query(q): Query>, ) -> Response { - if let Err(r) = require_global_admin(&s, &headers) { - return r; - } let store = match require_config_store(&s) { Ok(v) => v, Err(r) => return r, @@ -2111,14 +2115,11 @@ async fn admin_config_versions( /// this instance; peers converge via the store's change feed. Global admin only. async fn admin_config_put( State(s): State, - headers: HeaderMap, + _: GlobalAdmin, AuditSourceIp(source): AuditSourceIp, Query(q): Query>, body: String, ) -> Response { - if let Err(r) = require_global_admin(&s, &headers) { - return r; - } let store = match require_config_store(&s) { Ok(v) => v, Err(r) => return r, @@ -2177,13 +2178,10 @@ async fn admin_config_put( /// as a new head and reload this instance. Global admin only. async fn admin_config_rollback( State(s): State, - headers: HeaderMap, + _: GlobalAdmin, AuditSourceIp(source): AuditSourceIp, Path(source_id): Path, ) -> Response { - if let Err(r) = require_global_admin(&s, &headers) { - return r; - } let store = match require_config_store(&s) { Ok(v) => v, Err(r) => return r, @@ -2487,12 +2485,9 @@ async fn admin_security_events( /// Global admin only (the trail spans all tenants). async fn admin_audit_ops( State(s): State, - headers: HeaderMap, + _: GlobalAdmin, Query(q): Query>, ) -> Response { - if let Err(r) = require_global_admin(&s, &headers) { - return r; - } let limit = q_num(&q, "limit", LEDGER_PAGE_DEFAULT).min(ADMIN_PAGE_MAX); match s.handler.state().store.admin_audit_list(limit).await { Ok(entries) => Json(json!({ "entries": entries })).into_response(), @@ -2540,7 +2535,9 @@ async fn admin_content_get( row }) .collect(); - Json(json!({ "request_id": request_id, "entries": entries })).into_response() + let mut out = json!({ "request_id": request_id }); + out["entries"] = Value::Array(entries); + Json(out).into_response() } /// GET /admin/audit/content?user=&limit=&include= — one end user's retained rows, @@ -2832,7 +2829,7 @@ fn chat_reasoning(effort: Option, reasoning: Option) -> Option, - headers: HeaderMap, + UserHint(hint): UserHint, Authed(ak): Authed, ApiJson(body): ApiJson, ) -> Response { @@ -2883,7 +2880,7 @@ async fn chat_completions( ); param.typed = Some(typed); param.raw = Value::Object(body.extra); - let user_id = user_hint(&headers, ¶m.raw["user"]); + let user_id = user_hint(hint, ¶m.raw["user"]); let request = GatewayRequest { is_online: true, @@ -3281,23 +3278,12 @@ fn synth_chunks(outcome: &mut gw_engines::EngineOutcome) -> Vec Vec { let resp = &mut outcome.response; - if resp.anthropic_content.is_some() { - let chunks = gw_engines::anthropic_native_chunks(resp, None); - resp.anthropic_content = None; - return chunks; + if let Some(content) = resp.anthropic_content.take() { + return gw_engines::anthropic_native_chunks(resp, content, None); } - let mut chunks = text_chunks(resp); - if let Some(tc) = resp.tool_calls.take() { - chunks.push(gw_engines::StreamChunk { - tool_calls: Some(tc), - ..Default::default() - }); - } - chunks.push(gw_engines::StreamChunk { - finish_reason: Some(take_finish(resp)), - ..Default::default() - }); - chunks + // the raw pre-redaction deltas are never replayed: synth_chunks rebuilds from the redacted text + outcome.chunks.clear(); + synth_chunks(outcome) } fn stream_chunk_output_tokens(chunk: &gw_engines::StreamChunk) -> i64 { @@ -3312,7 +3298,7 @@ fn stream_chunk_output_tokens(chunk: &gw_engines::StreamChunk) -> i64 { tokens = tokens.saturating_add(encoder.encode_len(&tool_calls.to_string()) as i64); } if let Some(event) = &chunk.native_event { - // Anthropic deltas are objects keyed by kind; Responses deltas are strings + // deltas are objects keyed by kind on Anthropic and strings on Responses if let Some(value) = event["delta"].as_str() { tokens = tokens.saturating_add(encoder.encode_len(value) as i64); } @@ -3380,7 +3366,7 @@ async fn messages( body.extra.insert("system".into(), blocks); } param.raw = Value::Object(body.extra); - let user_id = user_hint(&headers, ¶m.raw["metadata"]["user_id"]); + let user_id = user_hint(user_header(&headers), ¶m.raw["metadata"]["user_id"]); let request = GatewayRequest { is_online: true, @@ -3597,15 +3583,15 @@ fn messages_stream_response( /// The wire pattern clients expect for a tool_use block: empty `input` /// in the start frame, the arguments as one input_json_delta, stop. - fn emit_tool_block(&mut self, block: &Value) { + fn emit_tool_block(&mut self, mut block: Value) { self.close_block(BlockKind::Text); let idx = self.next_idx; self.next_idx += 1; - self.queue.push_back(Self::ev( - "content_block_start", - json!({"type":"content_block_start","index":idx, - "content_block":{"type":"tool_use","id":block["id"],"name":block["name"],"input":{}}}), - )); + let mut start = json!({"type":"content_block_start","index":idx, + "content_block":{"type":"tool_use","input":{}}}); + start["content_block"]["id"] = block["id"].take(); + start["content_block"]["name"] = block["name"].take(); + self.queue.push_back(Self::ev("content_block_start", start)); self.queue.push_back(Self::ev( "content_block_delta", json!({"type":"content_block_delta","index":idx, @@ -3626,7 +3612,7 @@ fn messages_stream_response( self.ensure_message_start(); if let Some(frags) = self.tool_frags.take() { for block in anthropic_tool_blocks(Some(frags)) { - self.emit_tool_block(&block); + self.emit_tool_block(block); } } self.close_block(BlockKind::Text); @@ -3707,7 +3693,7 @@ fn messages_stream_response( .unwrap_or(false); if native { for block in anthropic_tool_blocks(Some(tc)) { - self.emit_tool_block(&block); + self.emit_tool_block(block); } } else { gw_engines::merge_tool_call_fragments(&mut self.tool_frags, &tc); @@ -3831,7 +3817,7 @@ fn response_v2_or_500(outcome: Option, engine: &str) /// as a single user message to CompletionsEngine. async fn completions( State(s): State, - headers: HeaderMap, + UserHint(hint): UserHint, Authed(ak): Authed, ApiJson(mut body): ApiJson, ) -> Response { @@ -3862,7 +3848,7 @@ async fn completions( gw_consts::Protocol::Completions, typed, vec![ChatMsg::text("user", prompt)], - user_hint(&headers, &body["user"]), + user_hint(hint, &body["user"]), ) .await { @@ -3902,7 +3888,7 @@ async fn completions( /// through ResponsesEngine and its native response is returned as-is. async fn responses( State(s): State, - headers: HeaderMap, + UserHint(hint): UserHint, Authed(ak): Authed, ApiJson(body): ApiJson, ) -> Response { @@ -3915,7 +3901,7 @@ async fn responses( return error_response(400, "input is required"); } let stream = body["stream"].as_bool().unwrap_or(false); - let user_id = user_hint(&headers, &body["user"]); + let user_id = user_hint(hint, &body["user"]); let stream_model = stream.then(|| model.clone()); let mut param = ModelParamV2::with_name(gw_consts::Protocol::Responses, model); param.raw = body; @@ -4070,7 +4056,7 @@ fn responses_stream_response( /// POST /v1/embeddings (OpenAI-compatible surface) async fn embeddings( State(s): State, - headers: HeaderMap, + UserHint(hint): UserHint, Authed(ak): Authed, ApiJson(mut body): ApiJson, ) -> Response { @@ -4090,7 +4076,7 @@ async fn embeddings( model, gw_consts::Protocol::Embeddings, typed, - user_hint(&headers, &body["user"]), + user_hint(hint, &body["user"]), "embeddings", "embeddings", started, @@ -4101,7 +4087,7 @@ async fn embeddings( /// POST /v1/images/generations (OpenAI-compatible image generation surface) async fn images_generations( State(s): State, - headers: HeaderMap, + UserHint(hint): UserHint, Authed(ak): Authed, ApiJson(mut body): ApiJson, ) -> Response { @@ -4123,7 +4109,7 @@ async fn images_generations( model, gw_consts::Protocol::Image, typed, - user_hint(&headers, &body["user"]), + user_hint(hint, &body["user"]), "images", "image", started, @@ -4135,7 +4121,7 @@ async fn images_generations( /// routes to the edit endpoint; the image arrives as base64 JSON. async fn images_edits( State(s): State, - headers: HeaderMap, + UserHint(hint): UserHint, Authed(ak): Authed, ApiJson(mut body): ApiJson, ) -> Response { @@ -4159,7 +4145,7 @@ async fn images_edits( model, gw_consts::Protocol::Image, typed, - user_hint(&headers, &body["user"]), + user_hint(hint, &body["user"]), "images_edits", "image", started, @@ -4170,7 +4156,7 @@ async fn images_edits( /// POST /v1/videos/generations — a `request_id` reply is remembered for the poll to bill. async fn videos_generations( State(s): State, - headers: HeaderMap, + UserHint(hint): UserHint, Authed(ak): Authed, ApiJson(mut body): ApiJson, ) -> Response { @@ -4187,7 +4173,7 @@ async fn videos_generations( aspect_ratio: gw_engines::engine::take_string(&mut body, "/aspect_ratio"), image: body.get_mut("image").map(Value::take), }); - let user = user_hint(&headers, &body["user"]); + let user = user_hint(hint, &body["user"]); let mut ctx = match run_family( &s, ak, @@ -4385,7 +4371,7 @@ async fn videos_content( /// POST /v1/audio/speech (TTS, returns audio bytes; OpenAI-compatible surface) async fn audio_speech( State(s): State, - headers: HeaderMap, + UserHint(hint): UserHint, Authed(ak): Authed, ApiJson(mut body): ApiJson, ) -> Response { @@ -4416,7 +4402,7 @@ async fn audio_speech( gw_consts::Protocol::Tts, typed, vec![], - user_hint(&headers, &body["user"]), + user_hint(hint, &body["user"]), ) .await { @@ -4488,7 +4474,7 @@ async fn audio_transcribe( gw_consts::Protocol::Stt, typed, vec![], - user_hint(&headers, &body["user"]), + user_hint(user_header(&headers), &body["user"]), ) .await { @@ -4521,7 +4507,7 @@ async fn audio_transcribe( /// an array of strings. async fn moderations( State(s): State, - headers: HeaderMap, + UserHint(hint): UserHint, Authed(ak): Authed, ApiJson(mut body): ApiJson, ) -> Response { @@ -4538,7 +4524,7 @@ async fn moderations( model, gw_consts::Protocol::Moderations, typed, - user_hint(&headers, &body["user"]), + user_hint(hint, &body["user"]), "moderations", "moderations", started, @@ -4549,7 +4535,7 @@ async fn moderations( /// POST /v1/search — web search as a routed backend: `{model, query, count?}`. async fn search( State(s): State, - headers: HeaderMap, + UserHint(hint): UserHint, Authed(ak): Authed, ApiJson(mut body): ApiJson, ) -> Response { @@ -4569,7 +4555,7 @@ async fn search( model, gw_consts::Protocol::Search, typed, - user_hint(&headers, &body["user"]), + user_hint(hint, &body["user"]), "search", "search", started, @@ -4580,7 +4566,7 @@ async fn search( /// POST /v1/rerank — Cohere/Jina-compatible: `{model, query, documents, top_n?}`. async fn rerank( State(s): State, - headers: HeaderMap, + UserHint(hint): UserHint, Authed(ak): Authed, ApiJson(mut body): ApiJson, ) -> Response { @@ -4602,7 +4588,7 @@ async fn rerank( model, gw_consts::Protocol::Rerank, typed, - user_hint(&headers, &body["user"]), + user_hint(hint, &body["user"]), "rerank", "rerank", started, @@ -4781,7 +4767,7 @@ async fn batches_get( ) -> Response { let found = s.handler.state().store.batch_get(&id).await; match tenant_owned(found, |j| &j.tenant, &ak.tenant, "batch", &id) { - Ok(job) => (StatusCode::OK, Json(json!(job))).into_response(), + Ok(job) => (StatusCode::OK, Json(job)).into_response(), Err(resp) => resp, } } @@ -5672,10 +5658,7 @@ mod tests { let yaml = "listen: {host: h, port: 1}\nadmin: {token_env: GW_TEST_CONTENT_ADMIN}\nmodels: [{name: gpt-4o, protocol: openai-chat}]\naccounts: [{name: a1, provider: openai, protocols: ['openai-chat']}]\ntenants: [{name: t1, retention: {content: full, days: 1}, security: {dlp_redact: false, detect_secrets: false}}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 100, daily_token_quota: 100000}]"; // SAFETY: unique var name for this test; no concurrent reader of it. unsafe { std::env::set_var("GW_TEST_CONTENT_ADMIN", "s3cret") }; - assert!( - !gw_state::sealing_available(), - "test env has no content key" - ); + assert!(!gw_state::can_seal(), "test env has no content key"); let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); let state = Arc::new(GatewayState::from_config(&cfg)); let app_state = AppState::new(cfg, state, Arc::new(gw_engines::MockTransport)); @@ -5774,7 +5757,7 @@ mod tests { ) .with_moderator(Arc::new(DenyModerator)); let offline = OfflineHandler::new(handler.clone()); - let app = AppState { + let app = AppState(Arc::new(AppInner { handler, offline, mcp: mcp_client(), @@ -5782,7 +5765,7 @@ mod tests { mcp_sessions: mcp_sessions(), loader: None, config_store: None, - }; + })); let ak = app.handler.state().auth.authenticate("k1").await.unwrap(); let cfg = app.handler.cfg(); let sec = cfg.security_for(&ak.tenant); @@ -5840,7 +5823,7 @@ mod tests { ) .with_moderator(Arc::new(FrameMaskModerator)); let offline = OfflineHandler::new(handler.clone()); - let app = AppState { + let app = AppState(Arc::new(AppInner { handler, offline, mcp: mcp_client(), @@ -5848,7 +5831,7 @@ mod tests { mcp_sessions: mcp_sessions(), loader: None, config_store: None, - }; + })); let ak = app.handler.state().auth.authenticate("k1").await.unwrap(); let mut frame = json!({"type":"input_text","text":"tell secret now"}); assert_eq!(rt_inbound_policy(&app, &ak, "", &mut frame).await, Ok(0)); diff --git a/crates/views/src/mcp.rs b/crates/views/src/mcp.rs index 2a123e2..9fab611 100644 --- a/crates/views/src/mcp.rs +++ b/crates/views/src/mcp.rs @@ -29,8 +29,8 @@ const FORWARDED_HEADERS: [&str; 5] = [ const RETURNED_HEADERS: [&str; 2] = ["content-type", "mcp-session-id"]; /// Methods whose results carry prose an agent reads; reviewed under `security.moderate`. const REVIEWED_METHODS: [&str; 3] = ["tools/call", "resources/read", "prompts/get"]; -/// Result fields that carry base64 binary, never prose: skipped so the review -/// neither reads nor rewrites an image, audio clip or blob resource. +/// Base64 payload fields, skipped only inside an image/audio block or a blob +/// resource so the review neither reads nor rewrites binary. const OPAQUE_KEYS: [&str; 2] = ["blob", "data"]; const JSONRPC_TOOL_DENIED: i64 = -32000; const JSONRPC_RESULT_BLOCKED: i64 = -32001; @@ -48,7 +48,7 @@ struct Call { /// `data` payload it could not parse, or framing it passes through. enum Segment { Message(Value), - Opaque(String), + Opaque, Raw(String), } @@ -267,13 +267,12 @@ fn parse_call(body: &[u8]) -> Result { let Value::Object(mut obj) = v else { return Err("JSON-RPC batches are not supported; send one message per request".to_owned()); }; - let method = obj - .get("method") - .and_then(Value::as_str) - .unwrap_or_default() - .to_owned(); - let tool = match obj.get("params").and_then(|p| p.get("name")) { - Some(Value::String(name)) if method == "tools/call" => Some(name.clone()), + let method = match obj.remove("method") { + Some(Value::String(method)) => method, + _ => String::new(), + }; + let tool = match obj.get_mut("params").and_then(|p| p.get_mut("name")) { + Some(Value::String(name)) if method == "tools/call" => Some(std::mem::take(name)), _ if method == "tools/call" => { return Err("tools/call needs a string params.name".to_owned()); } @@ -289,7 +288,7 @@ fn parse_call(body: &[u8]) -> Result { /// Keep only the allowlisted tools in every `tools/list` result; `None` when a message could not be parsed. fn filter_tool_list(bytes: &[u8], sse: bool, allowed: &[String]) -> Option> { let mut segments = parse_segments(bytes, sse); - if segments.iter().any(|seg| matches!(seg, Segment::Opaque(_))) { + if segments.iter().any(|seg| matches!(seg, Segment::Opaque)) { return None; } for tools in segments.iter_mut().filter_map(|seg| match seg { @@ -322,7 +321,7 @@ async fn moderate_result( sse: bool, ) -> Vec { let mut segments = parse_segments(bytes, sse); - if segments.iter().any(|seg| matches!(seg, Segment::Opaque(_))) { + if segments.iter().any(|seg| matches!(seg, Segment::Opaque)) { return blocked(snap, ak, server, label, id, UNREVIEWABLE, sse).await; } let texts: Vec<&mut String> = segments.iter_mut().flat_map(review_slots).collect(); @@ -389,14 +388,24 @@ fn collect_prose<'a>(v: &'a mut Value, out: &mut Vec<&'a mut String>) { match v { Value::String(s) => out.push(s), Value::Array(items) => items.iter_mut().for_each(|x| collect_prose(x, out)), - Value::Object(map) => map - .iter_mut() - .filter(|(k, _)| !OPAQUE_KEYS.contains(&k.as_str())) - .for_each(|(_, x)| collect_prose(x, out)), + Value::Object(map) => { + let binary = is_binary_node(map); + map.iter_mut() + .filter(|(k, _)| !(binary && OPAQUE_KEYS.contains(&k.as_str()))) + .for_each(|(_, x)| collect_prose(x, out)); + } _ => {} } } +/// An image/audio content block or a blob resource: its payload is base64, not prose. +fn is_binary_node(map: &serde_json::Map) -> bool { + matches!( + map.get("type").and_then(Value::as_str), + Some("image" | "audio") + ) || (map.contains_key("blob") && (map.contains_key("uri") || map.contains_key("mimeType"))) +} + /// A bare JSON body is one message; an event stream is its events, each event's `data` lines joined by newlines, framing kept verbatim. fn parse_segments(bytes: &[u8], sse: bool) -> Vec { let text = String::from_utf8_lossy(bytes); @@ -404,7 +413,7 @@ fn parse_segments(bytes: &[u8], sse: bool) -> Vec { if !sse { return vec![match serde_json::from_str(text) { Ok(msg) => Segment::Message(msg), - Err(_) => Segment::Opaque(text.to_owned()), + Err(_) => Segment::Opaque, }]; } let mut segments = Vec::new(); @@ -413,7 +422,7 @@ fn parse_segments(bytes: &[u8], sse: bool) -> Vec { if let Some(payload) = data.take() { segments.push(match serde_json::from_str(&payload) { Ok(msg) => Segment::Message(msg), - Err(_) => Segment::Opaque(payload), + Err(_) => Segment::Opaque, }); } }; @@ -444,14 +453,7 @@ fn serialize_segments(segments: Vec, sse: bool, hint: usize) -> Vec for seg in segments { match seg { Segment::Raw(s) => out.extend_from_slice(s.as_bytes()), - Segment::Opaque(payload) if sse => { - for line in payload.split('\n') { - out.extend_from_slice(b"data: "); - out.extend_from_slice(line.as_bytes()); - out.push(b'\n'); - } - } - Segment::Opaque(payload) => out.extend_from_slice(payload.as_bytes()), + Segment::Opaque => {} Segment::Message(msg) => { if sse { out.extend_from_slice(b"data: "); @@ -967,10 +969,9 @@ mod tests { let cfg = Arc::new(GatewayConfig::from_yaml(&app_yaml(base)).unwrap()); let state = Arc::new(GatewayState::from_config(&cfg)); let app_state = AppState::new(cfg, state.clone(), Arc::new(gw_engines::MockTransport)); - let app_state = AppState { - handler: app_state.handler.with_moderator(Arc::new(EmailMasker)), - ..app_state - }; + let mut inner = (*app_state.0).clone(); + inner.handler = inner.handler.with_moderator(Arc::new(EmailMasker)); + let app_state = AppState(Arc::new(inner)); (crate::app(app_state), state) } @@ -1072,6 +1073,10 @@ mod tests { Some("err_body") => { json!({"jsonrpc":"2.0","id":id,"error":{"code":-1,"message":"see bob@example.com"}}) } + Some("structured_data") => { + json!({"jsonrpc":"2.0","id":id,"result":{"structuredContent":{"data":"note bob@example.com"}, + "content":[{"type":"image","data":"Ym9iQGV4YW1wbGUuY29t"}]}}) + } _ => { json!({"jsonrpc":"2.0","id":id,"result":{"structuredContent":{"name":"contact bob@example.com"}}}) } @@ -1126,6 +1131,16 @@ mod tests { v["result"]["structuredContent"]["name"], "contact [MASKED]", "structured prose under any key is reviewed: {v}" ); + + let v = call_tool(&app, "structured_data").await; + assert_eq!( + v["result"]["structuredContent"]["data"], "note [MASKED]", + "a data field outside a binary block is prose: {v}" + ); + assert_eq!( + v["result"]["content"][0]["data"], "Ym9iQGV4YW1wbGUuY29t", + "an image block's payload is never read or rewritten" + ); } #[tokio::test] @@ -1173,11 +1188,7 @@ mod tests { matches!(segments[1], Segment::Message(_)), "joined data lines parse" ); - assert!( - segments - .iter() - .any(|s| matches!(s, Segment::Opaque(p) if p == "not json")) - ); + assert!(segments.iter().any(|s| matches!(s, Segment::Opaque))); assert!(filter_tool_list(sse, true, &["add".to_owned()]).is_none()); let bom = "\u{feff}{\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{\"tools\":[{\"name\":\"add\"},{\"name\":\"x\"}]}}"; let filtered = filter_tool_list(bom.as_bytes(), false, &["add".to_owned()]).unwrap(); diff --git a/docs/configuration.md b/docs/configuration.md index a343c8a..f820dec 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -191,7 +191,7 @@ wins. ```yaml accounts: - - name: openai-main + - name: openai-main # unique across accounts; a duplicate is rejected at load provider: openai priority: 1 # lower = preferred tier: ptu # ptu (provisioned, preferred) | paygo (default) diff --git a/docs/governance.md b/docs/governance.md index 63a8800..b357921 100644 --- a/docs/governance.md +++ b/docs/governance.md @@ -258,7 +258,9 @@ watermark or the trailing 20-minute window, whichever reaches further back — the first run rolls a pre-existing ledger whole, and a stalled task catches up on its own. Buckets only ever grow (a recompute over a partially pruned ledger keeps the more complete aggregate), and on Postgres a fleet elects one replica -per pass via an advisory lock. Usage queries are served from those buckets plus +per pass via an advisory lock. The in-process store keeps buckets for 62 days +and at most one million of them, dropping the oldest minutes first; Postgres +keeps them until an operator prunes the table. Usage queries are served from those buckets plus the raw ledger tail, so per-user cost stays correct after `storage.ledger_max_rows` prunes old billing rows. `since`/`until` bounds are minute-aligned, so a repeated query returns the same result whether a minute is diff --git a/docs/multi-instance.md b/docs/multi-instance.md index 489639b..51fc76e 100644 --- a/docs/multi-instance.md +++ b/docs/multi-instance.md @@ -10,7 +10,7 @@ shared, what stays local, and what the LB needs to do. | Rate limits / quotas / TPM (`Governance`) | Redis (`storage.redis_url`) | ✅ when Redis is set (includes pooled tenant QPS) | | Account health / cooldown (`HealthStore`) | Redis (`storage.redis_url`) | ✅ when Redis is set — one instance's cooldown benches the account for all | | Config: keys/models/providers/tenants (`ConfigStore`) | Postgres (`storage.postgres_url`) | ✅ when Postgres is set — versioned documents + a change feed | -| Access-key table (`KeyStore`) | Postgres (`storage.postgres_url`) | ✅ when Postgres is set — admin key CRUD is fleet-wide within ~2s and survives restarts | +| Access-key table (`KeyStore`) | Postgres (`storage.postgres_url`) | ✅ when Postgres is set — admin key CRUD is fleet-wide within ~2s and survives restarts; a key's MCP servers and tool allowlists are stored with it | | Billing ledger / files / batches / video jobs (`Store`) | Postgres (`storage.postgres_url`), else SQLite | ✅ with Postgres (a video poll may land on any instance; the settle claim is one atomic row update); SQLite stays per-node | | Request cache | in-process (moka), or Redis with `shared_cache: true` | ⚠️ per-instance by default; fleet-shared when `shared_cache` is set | | Thinking-signature audit | in-process only | ⚠️ per-instance; a continuation landing on another instance finds no anchor and fails open (forwarded, not rejected) |