From c11a5cbe30921a93a4c99e7f30de709f16680d72 Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 9 Sep 2026 19:05:31 +0900 Subject: [PATCH 01/17] deps: grpc-go 1.83.2 for the xDS DoS advisory in the control plane An indirect dependency; the control plane runs no xDS server, so the advisory is not reachable here, but the pinned range is the flagged one. --- control-plane/go.mod | 2 +- control-plane/go.sum | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) 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= From af832e7a3034fd58e491b30cadc6ae8bc3a8d7cb Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 9 Sep 2026 19:06:10 +0900 Subject: [PATCH 02/17] deps: vitest 4.1.11 for the mocker path-traversal advisory (web dev dependency) --- control-plane/web/package-lock.json | 100 ++++++++++++++-------------- 1 file changed, 50 insertions(+), 50 deletions(-) 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" From 3003ae385ac46550676d3e0d9a3b836d071966e0 Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 9 Sep 2026 19:25:39 +0900 Subject: [PATCH 03/17] review: comment register and budget Inline comments lead with a lowercase word (identifiers keep their case) across 26 sites; four lines that restated the code or a neighbouring doc are cut; two narrative comments are rephrased to one fact per line. --- crates/config/src/lib.rs | 2 +- crates/consts/src/error_class.rs | 5 ++--- crates/dag/src/token_estimate.rs | 1 - crates/engines/src/bedrock.rs | 2 +- crates/engines/src/bespoke.rs | 2 +- crates/engines/src/claude_engine.rs | 6 +++--- crates/engines/src/families.rs | 6 +++--- crates/engines/src/mock_transport.rs | 3 +-- crates/engines/src/openai_engine.rs | 2 +- crates/engines/src/realtime.rs | 1 - crates/engines/src/sigv4.rs | 2 +- crates/handler/src/offline.rs | 2 +- crates/handler/src/plugins.rs | 2 +- crates/protocol/src/reasoning.rs | 2 +- crates/server/src/main.rs | 6 +++--- crates/state/src/admission.rs | 2 +- crates/state/src/configstore.rs | 2 +- crates/state/src/keystore.rs | 2 +- crates/state/src/store.rs | 2 +- crates/state/src/thinking_signature.rs | 4 ++-- crates/views/src/lib.rs | 6 +++--- 21 files changed, 29 insertions(+), 33 deletions(-) diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 1177a71..599bdf9 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -1344,7 +1344,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"], 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/token_estimate.rs b/crates/dag/src/token_estimate.rs index 8325063..15dc260 100644 --- a/crates/dag/src/token_estimate.rs +++ b/crates/dag/src/token_estimate.rs @@ -128,7 +128,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; diff --git a/crates/engines/src/bedrock.rs b/crates/engines/src/bedrock.rs index d7aba68..8575e66 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()), ] } diff --git a/crates/engines/src/bespoke.rs b/crates/engines/src/bespoke.rs index f1f2c85..f9ce710 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 diff --git a/crates/engines/src/claude_engine.rs b/crates/engines/src/claude_engine.rs index 7b9d1d1..0936d2f 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() { diff --git a/crates/engines/src/families.rs b/crates/engines/src/families.rs index cbb99a8..a12f61c 100644 --- a/crates/engines/src/families.rs +++ b/crates/engines/src/families.rs @@ -68,7 +68,7 @@ impl VertexEngine { } fn build_body(&self) -> Value { - // Gemini has no system role: system turns go to systemInstruction, never contents + // there is no system role on Gemini: system turns go to systemInstruction, never contents let contents: Vec = self .base .request @@ -673,7 +673,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 +742,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| { diff --git a/crates/engines/src/mock_transport.rs b/crates/engines/src/mock_transport.rs index fdea77f..ee0e51a 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; diff --git a/crates/engines/src/openai_engine.rs b/crates/engines/src/openai_engine.rs index fe9e856..6275e66 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, diff --git a/crates/engines/src/realtime.rs b/crates/engines/src/realtime.rs index f2d0f95..2dd3010 100644 --- a/crates/engines/src/realtime.rs +++ b/crates/engines/src/realtime.rs @@ -99,7 +99,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..9b6520a 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); diff --git a/crates/handler/src/offline.rs b/crates/handler/src/offline.rs index 10c774f..8d801d6 100644 --- a/crates/handler/src/offline.rs +++ b/crates/handler/src/offline.rs @@ -180,7 +180,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; } diff --git a/crates/handler/src/plugins.rs b/crates/handler/src/plugins.rs index 2942a79..eb6b58f 100644 --- a/crates/handler/src/plugins.rs +++ b/crates/handler/src/plugins.rs @@ -715,7 +715,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() { 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..6b6e079 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() @@ -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") => { diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index 7341245..3ffb037 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -130,7 +130,7 @@ impl BillingLedger { let (queue, mut pending) = mpsc::channel::(LEDGER_QUEUE_CAPACITY); let deferred = store.deferred_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); 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/keystore.rs b/crates/state/src/keystore.rs index be06420..a0209c5 100644 --- a/crates/state/src/keystore.rs +++ b/crates/state/src/keystore.rs @@ -153,7 +153,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() diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index 6422f17..72c2fd2 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -1804,7 +1804,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, diff --git a/crates/state/src/thinking_signature.rs b/crates/state/src/thinking_signature.rs index cc185b7..7ed83f9 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 diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index e047784..2f7af15 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -862,7 +862,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 +909,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; @@ -3312,7 +3312,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); } From 74d25b418abdcd00068ebbec19c94837f3e14329 Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 9 Sep 2026 19:29:27 +0900 Subject: [PATCH 04/17] review: import groups inside test modules The this-crate group (use super::*) comes last, blank-line separated, in the five test modules that had it first; a function-body use in the offline handler moves to the file's std group. --- crates/dag/src/token_estimate.rs | 3 ++- crates/handler/src/lib.rs | 3 ++- crates/handler/src/offline.rs | 4 ++-- crates/handler/src/plugins.rs | 3 ++- crates/models/src/request.rs | 3 ++- crates/state/src/thinking_signature.rs | 3 ++- 6 files changed, 12 insertions(+), 7 deletions(-) diff --git a/crates/dag/src/token_estimate.rs b/crates/dag/src/token_estimate.rs index 15dc260..1c3fa48 100644 --- a/crates/dag/src/token_estimate.rs +++ b/crates/dag/src/token_estimate.rs @@ -184,9 +184,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/handler/src/lib.rs b/crates/handler/src/lib.rs index 615126f..2855308 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -817,10 +817,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)); diff --git a/crates/handler/src/offline.rs b/crates/handler/src/offline.rs index 8d801d6..20b5799 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}; @@ -90,9 +91,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(); diff --git a/crates/handler/src/plugins.rs b/crates/handler/src/plugins.rs index eb6b58f..a29f3a9 100644 --- a/crates/handler/src/plugins.rs +++ b/crates/handler/src/plugins.rs @@ -982,9 +982,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 { diff --git a/crates/models/src/request.rs b/crates/models/src/request.rs index 964f3c0..f2c9cc0 100644 --- a/crates/models/src/request.rs +++ b/crates/models/src/request.rs @@ -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/state/src/thinking_signature.rs b/crates/state/src/thinking_signature.rs index 7ed83f9..d7bee3d 100644 --- a/crates/state/src/thinking_signature.rs +++ b/crates/state/src/thinking_signature.rs @@ -779,11 +779,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); From eabaf4127f6bafd46b70e5b3fab2a7dfab586157 Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 9 Sep 2026 19:31:54 +0900 Subject: [PATCH 05/17] review: predicate names, visibility, and two single-use wrappers Boolean fns read as predicates: is_stopping, is_complete, is_blocklisted, defers_ledger_writes, distributes_batches, can_seal. Items nothing outside their file names become private: whole_seconds, signing_key, the two token encoders, the offline handler's inner field, three engine re-exports; an unused Clone derive goes. default_moderator and is_response_create are inlined at their one caller. --- crates/dag/src/token_estimate.rs | 7 +++---- crates/engines/src/families.rs | 2 +- crates/engines/src/lib.rs | 4 ++-- crates/engines/src/realtime.rs | 7 +------ crates/engines/src/sigv4.rs | 2 +- crates/handler/src/lib.rs | 4 ++-- crates/handler/src/moderation.rs | 6 +----- crates/handler/src/offline.rs | 14 +++++++------- crates/handler/src/plugins.rs | 4 ++-- crates/server/src/main.rs | 4 ++-- crates/state/src/admission.rs | 2 +- crates/state/src/content.rs | 2 +- crates/state/src/governance.rs | 1 - crates/state/src/lib.rs | 4 ++-- crates/state/src/store.rs | 10 +++++----- crates/state/src/thinking_signature.rs | 4 ++-- crates/views/src/lib.rs | 5 +---- 17 files changed, 34 insertions(+), 48 deletions(-) diff --git a/crates/dag/src/token_estimate.rs b/crates/dag/src/token_estimate.rs index 1c3fa48..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; diff --git a/crates/engines/src/families.rs b/crates/engines/src/families.rs index a12f61c..3654b59 100644 --- a/crates/engines/src/families.rs +++ b/crates/engines/src/families.rs @@ -575,7 +575,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(|| { 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/realtime.rs b/crates/engines/src/realtime.rs index 2dd3010..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 diff --git a/crates/engines/src/sigv4.rs b/crates/engines/src/sigv4.rs index 9b6520a..7764c5d 100644 --- a/crates/engines/src/sigv4.rs +++ b/crates/engines/src/sigv4.rs @@ -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/handler/src/lib.rs b/crates/handler/src/lib.rs index 2855308..26b70a1 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -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(); @@ -2933,7 +2933,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 20b5799..577e8d7 100644 --- a/crates/handler/src/offline.rs +++ b/crates/handler/src/offline.rs @@ -12,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 { @@ -28,7 +28,7 @@ impl OfflineHandler { 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() @@ -122,7 +122,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 +140,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 @@ -213,7 +213,7 @@ impl OfflineHandler { let claimed = tokio::select! { biased; changed = shutdown.changed() => { - if stopping(changed, &shutdown) { + if is_stopping(changed, &shutdown) { return; } continue; @@ -276,7 +276,7 @@ impl OfflineHandler { } } -fn stopping( +fn is_stopping( changed: Result<(), tokio::sync::watch::error::RecvError>, shutdown: &tokio::sync::watch::Receiver, ) -> bool { @@ -289,7 +289,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 a29f3a9..91e5214 100644 --- a/crates/handler/src/plugins.rs +++ b/crates/handler/src/plugins.rs @@ -39,7 +39,7 @@ impl<'a> ScanCounts<'a> { } fn visit(&mut self, s: &str) -> usize { - self.blocklist += i64::from(blocklist_hit(self.sec, s)); + 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; } @@ -281,7 +281,7 @@ 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 { +fn is_blocklisted(sec: &SecurityConf, text: &str) -> bool { if text.is_ascii() { return sec .blocklist diff --git a/crates/server/src/main.rs b/crates/server/src/main.rs index 6b6e079..2593f9e 100644 --- a/crates/server/src/main.rs +++ b/crates/server/src/main.rs @@ -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 { diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index 3ffb037..ca5d0ac 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -128,7 +128,7 @@ 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 tokio::spawn(async move { 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/lib.rs b/crates/state/src/lib.rs index 01cb198..c17bfcd 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}; @@ -1380,7 +1380,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 72c2fd2..a32de28 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -525,7 +525,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 +667,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 @@ -1839,7 +1839,7 @@ macro_rules! sql_store_impl { Ok(()) } - fn deferred_ledger_writes(&self) -> bool { + fn defers_ledger_writes(&self) -> bool { true } @@ -2717,7 +2717,7 @@ sql_store_impl!(PostgresStore, postgres, { Ok(Some(done)) } - fn distributed_batches(&self) -> bool { + fn distributes_batches(&self) -> bool { true } @@ -4268,7 +4268,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 d7bee3d..bcf1de3 100644 --- a/crates/state/src/thinking_signature.rs +++ b/crates/state/src/thinking_signature.rs @@ -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(); diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index 2f7af15..1767416 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -5672,10 +5672,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)); From 21f0d951363f6e308eb8b4475c56e91cf2d4706b Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 9 Sep 2026 19:44:39 +0900 Subject: [PATCH 06/17] review: request-path ownership Engines move the bodies they own instead of copying them: the Vertex engine takes its messages like its siblings and gemini_parts consumes each turn; family_outcome owns the model name so four callers move it; anthropic_native_chunks takes the content blocks so the views tail moves them and the engine copies once instead of twice; the DashScope stream callback takes the frame; the OpenAI system prepend inserts in place. The MCP proxy takes method and tool name out of the owned envelope, and an unparsable segment carries no payload (the whole non-JSON reply was copied and then discarded). Smaller borrow-first fixes: the realtime subprotocol key, tool-block start events, the terminal-retention subject (one Arc instead of two Strings), the Bedrock invoke URI, a search summary, two Default restatements, one json! deep copy, and Json(job). Behaviour: a Converse content block or tool whose container is not an object is skipped instead of panicking the request task (serde_json's IndexMut panics on a non-object), and a malformed image part no longer gains a stray null key on its way upstream. --- crates/engines/src/bedrock.rs | 4 +- crates/engines/src/bespoke.rs | 14 ++--- crates/engines/src/claude_engine.rs | 97 ++++++++++++++++------------- crates/engines/src/converse.rs | 26 ++++++-- crates/engines/src/families.rs | 69 ++++++++++---------- crates/engines/src/openai_engine.rs | 8 +-- crates/handler/src/lib.rs | 12 ++-- crates/protocol/src/anthropic.rs | 8 ++- crates/views/src/lib.rs | 33 +++++----- crates/views/src/mcp.rs | 38 ++++------- 10 files changed, 157 insertions(+), 152 deletions(-) diff --git a/crates/engines/src/bedrock.rs b/crates/engines/src/bedrock.rs index 8575e66..171baf3 100644 --- a/crates/engines/src/bedrock.rs +++ b/crates/engines/src/bedrock.rs @@ -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 f9ce710..c53e9c6 100644 --- a/crates/engines/src/bespoke.rs +++ b/crates/engines/src/bespoke.rs @@ -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 0936d2f..4a0e2d4 100644 --- a/crates/engines/src/claude_engine.rs +++ b/crates/engines/src/claude_engine.rs @@ -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 3654b59..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 { + 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 = self - .base - .request - .message - .iter() + 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; @@ -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/openai_engine.rs b/crates/engines/src/openai_engine.rs index 6275e66..2a4d7b7 100644 --- a/crates/engines/src/openai_engine.rs +++ b/crates/engines/src/openai_engine.rs @@ -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/handler/src/lib.rs b/crates/handler/src/lib.rs index 26b70a1..cb0c388 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -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, 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/views/src/lib.rs b/crates/views/src/lib.rs index 1767416..76f1b03 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -339,7 +339,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 +347,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 +362,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 { @@ -2540,7 +2539,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, @@ -3281,10 +3282,8 @@ 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() { @@ -3597,15 +3596,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 +3625,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 +3706,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); @@ -4781,7 +4780,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, } } diff --git a/crates/views/src/mcp.rs b/crates/views/src/mcp.rs index 2a123e2..4d83b31 100644 --- a/crates/views/src/mcp.rs +++ b/crates/views/src/mcp.rs @@ -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(); @@ -404,7 +403,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 +412,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 +443,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: "); @@ -1173,11 +1165,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(); From d1fd371135579d94abedd45c89e028d62a67761f Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 9 Sep 2026 19:55:10 +0900 Subject: [PATCH 07/17] review: claim-path allocations The billing ledger takes its row by value: the unattributed deferred path moves it into the queue and recovers it from a refused send, so a billed request no longer clones nine Strings; settle_and_bill returns the two totals its callers read. select_healthy probes the serving accounts without materializing them first. The serving account's vendor prices come from a hash index instead of a scan of every account (account names must now be unique, which the pool, health and latency tables already assumed). The memory ledger takes the rollup lock only when a row cap is set. The blocklist scan returns before lower-casing when the list is empty, and the streaming DLP walk reuses one path buffer per delta instead of formatting a String per key. --- crates/config/src/lib.rs | 8 +++++ crates/dag/src/nodes.rs | 8 ++--- crates/handler/src/plugins.rs | 25 ++++++++++--- crates/state/src/admission.rs | 67 ++++++++++++++++++++--------------- crates/state/src/lib.rs | 12 ++----- crates/state/src/store.rs | 8 ++--- crates/views/src/lib.rs | 4 +-- 7 files changed, 79 insertions(+), 53 deletions(-) diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 599bdf9..56ba041 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -764,6 +764,8 @@ pub struct GatewayConfig { product_idx: HashMap, #[serde(skip)] tenant_idx: HashMap, + #[serde(skip)] + account_idx: HashMap, } impl GatewayConfig { @@ -781,6 +783,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 @@ -1101,6 +1104,7 @@ impl GatewayConfig { check_unique("product", self.products.iter().map(|p| p.name.as_str()))?; check_unique("provider", self.providers.iter().map(|p| p.name.as_str()))?; check_unique("tenant", self.tenants.iter().map(|t| t.name.as_str()))?; + check_unique("account", self.accounts.iter().map(|a| a.name.as_str()))?; check_unique( "mcp server", self.mcp_servers.iter().map(|m| m.name.as_str()), @@ -1278,6 +1282,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) diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 8aec472..4f71d02 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -837,7 +837,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 +875,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/handler/src/plugins.rs b/crates/handler/src/plugins.rs index 91e5214..092570d 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}; @@ -282,6 +284,9 @@ pub fn apply_mask_spans_frame( /// Case-insensitive blocklist test; ASCII matches without allocating, non-ASCII copies once. fn is_blocklisted(sec: &SecurityConf, text: &str) -> bool { + if sec.blocklist.is_empty() { + return false; + } if text.is_ascii() { return sec .blocklist @@ -746,13 +751,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 +772,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 +795,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; } diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index ca5d0ac..676d6e0 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -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 diff --git a/crates/state/src/lib.rs b/crates/state/src/lib.rs index c17bfcd..d31acca 100644 --- a/crates/state/src/lib.rs +++ b/crates/state/src/lib.rs @@ -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()) diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index a32de28..22e0c84 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -275,9 +275,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| { ( ( @@ -814,13 +812,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 diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index 76f1b03..dda182f 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -638,7 +638,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 { @@ -677,7 +677,7 @@ async fn bill_realtime_turn( ak, admit.user.as_str(), total, - record.cost_micros, + settled.cost_micros, ) .await; if !estimated { From 8b4b117e96d886476a763696d3fad7492282b9e3 Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 9 Sep 2026 20:06:02 +0900 Subject: [PATCH 08/17] review: request-entry allocations in the view layer AppState sits behind one Arc, so the two clones axum makes per request cost one refcount each instead of fifteen. A UserHint extractor reads x-gw-user from the request parts, so eleven REST handlers no longer take a cloned HeaderMap they read one header from; a GlobalAdmin extractor replaces the three-line token prologue on nine admin handlers and, as a FromRequestParts extractor, rejects before the config body is buffered. ModelConf parses its protocol kind once at load instead of scanning the wire table on every lookup. --- crates/config/src/lib.rs | 7 +- crates/views/src/lib.rs | 157 +++++++++++++++++++++------------------ crates/views/src/mcp.rs | 7 +- 3 files changed, 92 insertions(+), 79 deletions(-) diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 56ba041..797c8fe 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 } } @@ -873,6 +875,9 @@ impl GatewayConfig { compile_security(sec); } } + for m in &mut self.models { + m.kind = Protocol::from_wire(&m.protocol); + } Ok(()) } diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index dda182f..4173807 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 } @@ -1283,12 +1295,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(), @@ -1298,10 +1307,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()); @@ -1332,8 +1338,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 @@ -1411,6 +1417,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 @@ -1854,12 +1888,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(); @@ -2044,10 +2075,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, @@ -2064,14 +2092,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, @@ -2089,12 +2110,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, @@ -2110,14 +2128,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, @@ -2176,13 +2191,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, @@ -2486,12 +2498,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(), @@ -2833,7 +2842,7 @@ fn chat_reasoning(effort: Option, reasoning: Option) -> Option, - headers: HeaderMap, + UserHint(hint): UserHint, Authed(ak): Authed, ApiJson(body): ApiJson, ) -> Response { @@ -2884,7 +2893,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, @@ -3379,7 +3388,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, @@ -3830,7 +3839,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 { @@ -3861,7 +3870,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 { @@ -3901,7 +3910,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 { @@ -3914,7 +3923,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; @@ -4069,7 +4078,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 { @@ -4089,7 +4098,7 @@ async fn embeddings( model, gw_consts::Protocol::Embeddings, typed, - user_hint(&headers, &body["user"]), + user_hint(hint, &body["user"]), "embeddings", "embeddings", started, @@ -4100,7 +4109,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 { @@ -4122,7 +4131,7 @@ async fn images_generations( model, gw_consts::Protocol::Image, typed, - user_hint(&headers, &body["user"]), + user_hint(hint, &body["user"]), "images", "image", started, @@ -4134,7 +4143,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 { @@ -4158,7 +4167,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, @@ -4169,7 +4178,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 { @@ -4186,7 +4195,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, @@ -4384,7 +4393,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 { @@ -4415,7 +4424,7 @@ async fn audio_speech( gw_consts::Protocol::Tts, typed, vec![], - user_hint(&headers, &body["user"]), + user_hint(hint, &body["user"]), ) .await { @@ -4487,7 +4496,7 @@ async fn audio_transcribe( gw_consts::Protocol::Stt, typed, vec![], - user_hint(&headers, &body["user"]), + user_hint(user_header(&headers), &body["user"]), ) .await { @@ -4520,7 +4529,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 { @@ -4537,7 +4546,7 @@ async fn moderations( model, gw_consts::Protocol::Moderations, typed, - user_hint(&headers, &body["user"]), + user_hint(hint, &body["user"]), "moderations", "moderations", started, @@ -4548,7 +4557,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 { @@ -4568,7 +4577,7 @@ async fn search( model, gw_consts::Protocol::Search, typed, - user_hint(&headers, &body["user"]), + user_hint(hint, &body["user"]), "search", "search", started, @@ -4579,7 +4588,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 { @@ -4601,7 +4610,7 @@ async fn rerank( model, gw_consts::Protocol::Rerank, typed, - user_hint(&headers, &body["user"]), + user_hint(hint, &body["user"]), "rerank", "rerank", started, @@ -5770,7 +5779,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(), @@ -5778,7 +5787,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); @@ -5836,7 +5845,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(), @@ -5844,7 +5853,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 4d83b31..2d9d630 100644 --- a/crates/views/src/mcp.rs +++ b/crates/views/src/mcp.rs @@ -959,10 +959,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) } From 7ef480626f38d6dbafb4bbe03dc72fee6ef6a07b Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 9 Sep 2026 20:30:38 +0900 Subject: [PATCH 09/17] review: simplification pass One fail-loud DagContext::model_param accessor replaces four guards on a state resolve_model already rules out (the tenant-entitlement guard would have silently allowed an empty model name); weighted_pair is inlined at its one caller; the Responses reasoning-replay scan is written once; the security-scan rows are extended without a temporary Vec; batch items take the owner override in place; the realtime relay bills a turn from one call site and the mock session from one; the redacted stream tail delegates to synth_chunks after clearing the raw deltas; the Converse mock loses a dead Cohere branch and scans for the last user turn once; walk_native_event dispatches once on the event type; ScanCounts::visit returns nothing; a redundant drop and an always-true streamed_live go. The MCP client's unreachable fallback, which would have followed redirects, is an expected invariant instead; the SIGTERM install error logs its value. --- crates/dag/src/context.rs | 10 ++- crates/dag/src/nodes.rs | 34 ++++------ crates/engines/src/mock_transport.rs | 28 ++------ crates/engines/src/pump.rs | 2 +- crates/handler/src/lib.rs | 12 ++-- crates/handler/src/offline.rs | 16 ++--- crates/handler/src/plugins.rs | 45 +++++++------ crates/models/src/cost.rs | 23 ------- crates/models/src/lib.rs | 3 +- crates/models/src/request.rs | 24 +++---- crates/server/src/main.rs | 2 +- crates/state/src/thinking_signature.rs | 1 - crates/views/src/lib.rs | 92 ++++++++++---------------- 13 files changed, 116 insertions(+), 176 deletions(-) diff --git a/crates/dag/src/context.rs b/crates/dag/src/context.rs index 2dfb2b8..4434445 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 { @@ -78,6 +78,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 4f71d02..59f41d6 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(()); @@ -388,12 +379,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) @@ -702,8 +688,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 +752,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(()); } diff --git a/crates/engines/src/mock_transport.rs b/crates/engines/src/mock_transport.rs index ee0e51a..7ff0d2c 100644 --- a/crates/engines/src/mock_transport.rs +++ b/crates/engines/src/mock_transport.rs @@ -29,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| { @@ -64,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) @@ -320,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(); @@ -332,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/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/handler/src/lib.rs b/crates/handler/src/lib.rs index cb0c388..eaee4b1 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 = diff --git a/crates/handler/src/offline.rs b/crates/handler/src/offline.rs index 577e8d7..864ffb6 100644 --- a/crates/handler/src/offline.rs +++ b/crates/handler/src/offline.rs @@ -25,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.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) diff --git a/crates/handler/src/plugins.rs b/crates/handler/src/plugins.rs index 092570d..d183e3c 100644 --- a/crates/handler/src/plugins.rs +++ b/crates/handler/src/plugins.rs @@ -40,12 +40,11 @@ impl<'a> ScanCounts<'a> { } } - fn visit(&mut self, s: &str) -> usize { + 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. @@ -200,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() } @@ -678,26 +680,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() { 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 f2c9cc0..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)] diff --git a/crates/server/src/main.rs b/crates/server/src/main.rs index 2593f9e..6b1b4d4 100644 --- a/crates/server/src/main.rs +++ b/crates/server/src/main.rs @@ -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/state/src/thinking_signature.rs b/crates/state/src/thinking_signature.rs index bcf1de3..8f872b1 100644 --- a/crates/state/src/thinking_signature.rs +++ b/crates/state/src/thinking_signature.rs @@ -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; diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index 4173807..5f32ea5 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -147,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> { @@ -799,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; @@ -1052,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 @@ -1081,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, @@ -1089,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; @@ -3294,18 +3281,9 @@ fn redacted_stream_tail(outcome: &mut gw_engines::EngineOutcome) -> Vec i64 { From 4accf144a7a7ef9b2950ae305be857b512fa620b Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 9 Sep 2026 20:40:56 +0900 Subject: [PATCH 10/17] fix: persist MCP entitlements in the Postgres key store The access_keys table had no columns for a key's MCP servers and tool allowlists, so every key read back from Postgres carried an empty grant and the MCP proxy answered 404 for all of them whenever storage.postgres_url was set. Two JSON text columns, added with ADD COLUMN IF NOT EXISTS for existing tables, ride through the upsert and the three reads; the Postgres-gated test round-trips a config-declared grant and an admin put. --- crates/state/src/keystore.rs | 52 +++++++++++++++++++++++++++++++----- 1 file changed, 45 insertions(+), 7 deletions(-) diff --git a/crates/state/src/keystore.rs b/crates/state/src/keystore.rs index a0209c5..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) @@ -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()); } } From fc1ac1b5db5dfb77f4de9f556cf20b2d8ddcc57e Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 9 Sep 2026 20:44:56 +0900 Subject: [PATCH 11/17] review: age index on video jobs; batch items in one statement per chunk The retention DELETE that runs on every async video submit scanned a table with no index on created_at_epoch_secs; both dialects now carry one. Postgres batch_enqueue inserted items one round trip at a time inside the submit transaction; it now uses the same push_values shape as ledger_add_batch, chunked under the parameter limit. --- crates/state/src/store.rs | 31 +++++++++++++++++++------------ 1 file changed, 19 insertions(+), 12 deletions(-) diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index 22e0c84..11798c6 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -50,6 +50,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; @@ -1365,6 +1367,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, @@ -2408,6 +2411,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, @@ -2741,18 +2745,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 From 0e8056d0fba26acbacc708f87ca82dfac4411137 Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 9 Sep 2026 21:03:25 +0900 Subject: [PATCH 12/17] perf: compile the blocklist into one automaton at config load Every string leaf of every request and response ran a windowed case-insensitive compare per blocklist word (a 20-word list over a 100 KB prompt is millions of comparisons per request). The list is built into an Aho-Corasick automaton where it is lower-cased at load, so a leaf costs one pass; a list that fails to build is logged and ignored, as a failing regex rule already is. SecurityConf::compiled runs the same step for a policy built outside the loader. Semantics are unchanged: ASCII text matches case-insensitively, other text is lower-cased first. --- Cargo.lock | 1 + Cargo.toml | 1 + crates/config/Cargo.toml | 1 + crates/config/src/lib.rs | 22 ++++++++++++++++++++++ crates/handler/src/lib.rs | 1 + crates/handler/src/plugins.rs | 34 ++++++++++++++-------------------- 6 files changed, 40 insertions(+), 20 deletions(-) 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/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 797c8fe..07b343a 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -390,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, @@ -415,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. @@ -1482,6 +1491,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/handler/src/lib.rs b/crates/handler/src/lib.rs index eaee4b1..8bd056a 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -2634,6 +2634,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( diff --git a/crates/handler/src/plugins.rs b/crates/handler/src/plugins.rs index d183e3c..587f83d 100644 --- a/crates/handler/src/plugins.rs +++ b/crates/handler/src/plugins.rs @@ -286,25 +286,14 @@ pub fn apply_mask_spans_frame( /// Case-insensitive blocklist test; ASCII matches without allocating, non-ASCII copies once. fn is_blocklisted(sec: &SecurityConf, text: &str) -> bool { - if sec.blocklist.is_empty() { + let Some(matcher) = &sec.blocklist_matcher else { return false; - } + }; 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() { - return false; + 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. @@ -1041,6 +1030,7 @@ mod tests { dlp_redact: true, ..Default::default() } + .compiled() } #[test] @@ -1065,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() @@ -1133,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"); @@ -1148,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" }); @@ -1421,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() From 82501e58b9df5b523a43d16e8f7f7fa0b6b005dd Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 9 Sep 2026 21:07:50 +0900 Subject: [PATCH 13/17] review: name the fixture that failed in the integration tests Setup calls that bind, serve, build a transport, a mock response or a request carry a short expect message, as the suite's own helpers already did; the batch-pricing test asserts the batch actually completed instead of falling through its poll loop; the two QPM warm-up loops say which call failed. --- crates/engines/tests/http_transport_wire.rs | 45 +++++++++++---------- crates/server/tests/compat.rs | 4 +- crates/server/tests/e2e.rs | 11 +++-- crates/server/tests/live_path.rs | 25 ++++++------ crates/server/tests/otel_trace.rs | 6 +-- 5 files changed, 47 insertions(+), 44 deletions(-) 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/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 From 04ded6b1ec640ee1117ac6c9b3873a7d7a5bafd5 Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 9 Sep 2026 21:08:42 +0900 Subject: [PATCH 14/17] docs: account names are unique; MCP grants persist with Postgres keys --- docs/configuration.md | 2 +- docs/multi-instance.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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/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) | From 8c3dee30ff91c8733766bbbcb5ea20383bfcbabc Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 9 Sep 2026 21:14:04 +0900 Subject: [PATCH 15/17] review: drop a duplicated account-uniqueness check validate() already rejected duplicate account names further down; the second check added with the account index was redundant. --- crates/config/src/lib.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 07b343a..a0c4e99 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -1118,7 +1118,6 @@ impl GatewayConfig { check_unique("product", self.products.iter().map(|p| p.name.as_str()))?; check_unique("provider", self.providers.iter().map(|p| p.name.as_str()))?; check_unique("tenant", self.tenants.iter().map(|t| t.name.as_str()))?; - check_unique("account", self.accounts.iter().map(|a| a.name.as_str()))?; check_unique( "mcp server", self.mcp_servers.iter().map(|m| m.name.as_str()), From ba4d2ce4171e72af8e46c1db911205f0d3732f97 Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 9 Sep 2026 21:45:59 +0900 Subject: [PATCH 16/17] fix: key budgets, MCP review and availability samples (external review) The per-key cost counter was keyed by the access key itself, so the fail-open warning on a Redis read error logged the credential; it is now keyed by the key's fingerprint like every other log field (a budget in flight at deploy time starts over). The MCP result review skipped any field named data or blob at any depth, so structuredContent: {"data": ...} prose escaped it; the payload is skipped only inside an image/audio block or a blob resource. A request served by a fallback model recorded one failure for the first attempt and one success for the fallback, both against the requested name, so one successful call read as a 50% error rate. The handler now tells the pipeline when a fallback is still ahead and a failed attempt samples only when it is the client-visible outcome. --- crates/dag/src/context.rs | 4 ++++ crates/dag/src/nodes.rs | 21 +++++++++++--------- crates/handler/src/lib.rs | 22 +++++++++++++++++++++ crates/state/src/admission.rs | 11 +++++++---- crates/views/src/mcp.rs | 36 +++++++++++++++++++++++++++++------ 5 files changed, 75 insertions(+), 19 deletions(-) diff --git a/crates/dag/src/context.rs b/crates/dag/src/context.rs index 4434445..46f93dd 100644 --- a/crates/dag/src/context.rs +++ b/crates/dag/src/context.rs @@ -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, diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 59f41d6..e12d5f2 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -309,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, @@ -482,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(); @@ -506,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)) } @@ -527,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, diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 8bd056a..c7dcf84 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -256,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() @@ -1138,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; diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index 676d6e0..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), } } @@ -862,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); @@ -901,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/views/src/mcp.rs b/crates/views/src/mcp.rs index 2d9d630..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; @@ -388,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); @@ -1063,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"}}}) } @@ -1117,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] From 868d1fd448acfca9962f2e0bf11fda3738d5d2cf Mon Sep 17 00:00:00 2001 From: CMGS Date: Wed, 9 Sep 2026 21:45:59 +0900 Subject: [PATCH 17/17] review: bound the in-process usage rollup The memory store's minute buckets were insert-only and keyed by the client-supplied user id, so a caller rotating that id grew the map for the life of the process. Each rollup pass now drops buckets older than 62 days (the monthly-budget window) and, above one million buckets, the oldest minutes first. Postgres keeps its table until an operator prunes it, as before. --- crates/state/src/store.rs | 50 +++++++++++++++++++++++++++++++++++++++ docs/governance.md | 4 +++- 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index 11798c6..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). @@ -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); @@ -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) } @@ -4076,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(); 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