From f016a0dd168591b5c912f2b9eacfae5325a537c7 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Tue, 3 Nov 2026 13:53:54 +0800 Subject: [PATCH 1/2] fix: automatically schedule CPU work and bound gateway memory Size execution from CPU and cgroup memory availability, retain buffer admission through AI/MCP responses and stream tails, bound audit backlog bytes and password hashing, and adopt Core 0.89.0. Add cancellation, elastic execution, bounded HTTP/Redis and audit ownership coverage. --- Cargo.lock | 255 +++++++++- Cargo.toml | 8 +- crates/auth/src/password.rs | 68 ++- crates/common/Cargo.toml | 8 + crates/common/src/audit/logger.rs | 18 +- crates/common/src/audit/types.rs | 96 ++++ crates/common/src/lib.rs | 1 + crates/common/src/resources.rs | 465 ++++++++++++++++++ crates/common/src/resources/cgroup.rs | 89 ++++ crates/common/src/resources/http.rs | 227 +++++++++ crates/gateway/src/bedrock/sigv4.rs | 35 +- crates/gateway/src/cache.rs | 10 +- crates/gateway/src/lifecycle/mod.rs | 28 +- crates/gateway/src/proxy/body_capture.rs | 2 +- crates/gateway/src/proxy/generate.rs | 279 +++++++---- crates/gateway/src/proxy/transport.rs | 5 +- crates/mcp-gateway/src/cache.rs | 28 +- crates/mcp-gateway/src/lifecycle/mod.rs | 19 + crates/mcp-gateway/src/pool.rs | 57 ++- crates/mcp-gateway/src/proxy.rs | 41 +- crates/mcp-gateway/src/proxy/jsonrpc.rs | 6 + crates/mcp-gateway/src/proxy/streaming.rs | 111 ++++- .../src/transport/streamable_http.rs | 30 +- crates/mcp-gateway/src/user_token.rs | 5 +- crates/server/Cargo.toml | 2 + crates/server/src/app.rs | 4 + crates/server/src/main.rs | 38 +- crates/test-support/tests/resources.rs | 73 +++ docs/operations/resources.md | 64 +++ 29 files changed, 1828 insertions(+), 244 deletions(-) create mode 100644 crates/common/src/resources.rs create mode 100644 crates/common/src/resources/cgroup.rs create mode 100644 crates/common/src/resources/http.rs create mode 100644 crates/test-support/tests/resources.rs create mode 100644 docs/operations/resources.md diff --git a/Cargo.lock b/Cargo.lock index 9c607902..32c07700 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -443,6 +443,12 @@ version = "1.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" +[[package]] +name = "bech32" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d86b93f97252c47b41663388e6d155714a9d0c398b99f1005cbc5f978b29f445" + [[package]] name = "bitflags" version = "1.3.2" @@ -610,7 +616,7 @@ dependencies = [ "num-traits", "serde", "wasm-bindgen", - "windows-link", + "windows-link 0.2.1", ] [[package]] @@ -1335,6 +1341,12 @@ dependencies = [ "smallvec", ] +[[package]] +name = "fastrand" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" + [[package]] name = "ff" version = "0.13.1" @@ -1575,8 +1587,8 @@ dependencies = [ "libc", "log", "rustversion", - "windows-link", - "windows-result", + "windows-link 0.1.3", + "windows-result 0.3.4", ] [[package]] @@ -1919,7 +1931,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core", + "windows-core 0.62.2", ] [[package]] @@ -2108,7 +2120,7 @@ dependencies = [ "portable-atomic", "portable-atomic-util", "serde_core", - "windows-link", + "windows-link 0.2.1", ] [[package]] @@ -2162,7 +2174,7 @@ dependencies = [ "simd_cesu8", "thiserror 2.0.21", "walkdir", - "windows-link", + "windows-link 0.2.1", ] [[package]] @@ -2274,6 +2286,15 @@ version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" +[[package]] +name = "libmimalloc-sys" +version = "0.1.49" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a45a52f43e1c16f667ccfe4dd8c85b7f7c204fd5e3bf46c5b0db9a5c3c0b8e9" +dependencies = [ + "cc", +] + [[package]] name = "libsqlite3-sys" version = "0.30.1" @@ -2284,6 +2305,12 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + [[package]] name = "litemap" version = "0.8.3" @@ -2409,6 +2436,15 @@ dependencies = [ "sketches-ddsketch", ] +[[package]] +name = "mimalloc" +version = "0.1.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d4139bb28d14ad1facf21d5eb8825051b326e172d216b39f6d31df53cc97862" +dependencies = [ + "libmimalloc-sys", +] + [[package]] name = "mime" version = "0.3.17" @@ -2468,6 +2504,15 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "ntapi" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3b335231dfd352ffb0f8017f3b6027a4917f7df785ea2143d8af2adc66980ae" +dependencies = [ + "winapi", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -2568,6 +2613,25 @@ dependencies = [ "url", ] +[[package]] +name = "objc2-core-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" +dependencies = [ + "bitflags 2.13.2", +] + +[[package]] +name = "objc2-io-kit" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33fafba39597d6dc1fb709123dfa8289d39406734be322956a69f0931c73bb15" +dependencies = [ + "libc", + "objc2-core-foundation", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -2689,7 +2753,7 @@ dependencies = [ "libc", "redox_syscall", "smallvec", - "windows-link", + "windows-link 0.2.1", ] [[package]] @@ -2944,7 +3008,7 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls", - "socket2 0.6.5", + "socket2 0.5.10", "thiserror 2.0.21", "tokio", "tracing", @@ -2983,7 +3047,7 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2 0.6.5", + "socket2 0.5.10", "tracing", "windows-sys 0.61.2", ] @@ -3362,6 +3426,19 @@ dependencies = [ "semver", ] +[[package]] +name = "rustix" +version = "1.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5e76742e5294153d4dde1207dbec5bdd81eef6fabdc0d2118babadabfd858aa" +dependencies = [ + "bitflags 2.13.2", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + [[package]] name = "rustls" version = "0.23.46" @@ -4166,6 +4243,20 @@ dependencies = [ "syn 3.0.6", ] +[[package]] +name = "sysinfo" +version = "0.36.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "252800745060e7b9ffb7b2badbd8b31cfa4aa2e61af879d0a3bf2a317c20217d" +dependencies = [ + "libc", + "memchr", + "ntapi", + "objc2-core-foundation", + "objc2-io-kit", + "windows", +] + [[package]] name = "system-configuration" version = "0.7.0" @@ -4187,6 +4278,19 @@ dependencies = [ "libc", ] +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.3.4", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + [[package]] name = "think-watch-auth" version = "3.10.0" @@ -4223,6 +4327,7 @@ version = "3.10.0" dependencies = [ "aes-gcm", "anyhow", + "async-stream", "async-trait", "aws-credential-types", "aws-sigv4", @@ -4236,6 +4341,8 @@ dependencies = [ "hex", "hmac 0.13.0", "http", + "http-body", + "http-body-util", "hyper-util", "metrics", "rand 0.10.3", @@ -4248,8 +4355,11 @@ dependencies = [ "sha1 0.11.0", "sha2 0.11.0", "sqlx", + "sysinfo", + "tempfile", "thiserror 2.0.21", "tokio", + "tower", "tracing", "tw-bedrock", "tw-breaker", @@ -4344,8 +4454,10 @@ dependencies = [ "futures", "hex", "hmac 0.13.0", + "libc", "metrics", "metrics-exporter-prometheus", + "mimalloc", "openidconnect", "p256 0.14.0", "rand 0.10.3", @@ -4838,8 +4950,8 @@ dependencies = [ [[package]] name = "tw-bedrock" -version = "0.86.0" -source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.86.0#62aec779f28b8f2f087670cc5096af829bd5ad20" +version = "0.89.0" +source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.89.0#c66a3b04ba297f608f2fc21501c73f21a78ab8ed" dependencies = [ "aws-credential-types", "aws-sigv4", @@ -4853,16 +4965,16 @@ dependencies = [ [[package]] name = "tw-breaker" -version = "0.86.0" -source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.86.0#62aec779f28b8f2f087670cc5096af829bd5ad20" +version = "0.89.0" +source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.89.0#c66a3b04ba297f608f2fc21501c73f21a78ab8ed" dependencies = [ "serde", ] [[package]] name = "tw-dialect" -version = "0.86.0" -source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.86.0#62aec779f28b8f2f087670cc5096af829bd5ad20" +version = "0.89.0" +source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.89.0#c66a3b04ba297f608f2fc21501c73f21a78ab8ed" dependencies = [ "memchr", "serde", @@ -4871,12 +4983,14 @@ dependencies = [ [[package]] name = "tw-guard" -version = "0.86.0" -source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.86.0#62aec779f28b8f2f087670cc5096af829bd5ad20" +version = "0.89.0" +source = "git+https://github.com/ThinkWatchProject/ThinkWatch-Core.git?tag=v0.89.0#c66a3b04ba297f608f2fc21501c73f21a78ab8ed" dependencies = [ "aho-corasick", "base64 0.22.1", + "bech32", "bytes", + "curve25519-dalek", "memchr", "regex", "regex-syntax", @@ -5238,6 +5352,41 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "windows" +version = "0.61.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9babd3a767a4c1aef6900409f85f5d53ce2544ccdfaa86dad48c91782c6d6893" +dependencies = [ + "windows-collections", + "windows-core 0.61.2", + "windows-future", + "windows-link 0.1.3", + "windows-numerics", +] + +[[package]] +name = "windows-collections" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3beeceb5e5cfd9eb1d76b381630e82c4241ccd0d27f1a39ed41b2760b255c5e8" +dependencies = [ + "windows-core 0.61.2", +] + +[[package]] +name = "windows-core" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0fdd3ddb90610c7638aa2b3a3ab2904fb9e5cdbecc643ddb3647212781c4ae3" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link 0.1.3", + "windows-result 0.3.4", + "windows-strings 0.4.2", +] + [[package]] name = "windows-core" version = "0.62.2" @@ -5246,9 +5395,20 @@ checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" dependencies = [ "windows-implement", "windows-interface", - "windows-link", - "windows-result", - "windows-strings", + "windows-link 0.2.1", + "windows-result 0.4.1", + "windows-strings 0.5.1", +] + +[[package]] +name = "windows-future" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc6a41e98427b19fe4b73c550f060b59fa592d7d686537eebf9385621bfbad8e" +dependencies = [ + "windows-core 0.61.2", + "windows-link 0.1.3", + "windows-threading", ] [[package]] @@ -5273,21 +5433,46 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "windows-link" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e6ad25900d524eaabdbbb96d20b4311e1e7ae1699af4fb28c17ae66c80d798a" + [[package]] name = "windows-link" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-numerics" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9150af68066c4c5c07ddc0ce30421554771e528bde427614c61038bc2c92c2b1" +dependencies = [ + "windows-core 0.61.2", + "windows-link 0.1.3", +] + [[package]] name = "windows-registry" version = "0.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "02752bf7fbdcce7f2a27a742f798510f3e5ad88dbe84871e5168e2120c3d5720" dependencies = [ - "windows-link", - "windows-result", - "windows-strings", + "windows-link 0.2.1", + "windows-result 0.4.1", + "windows-strings 0.5.1", +] + +[[package]] +name = "windows-result" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56f42bd332cc6c8eac5af113fc0c1fd6a8fd2aa08a0119358686e5160d0586c6" +dependencies = [ + "windows-link 0.1.3", ] [[package]] @@ -5296,7 +5481,16 @@ version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" dependencies = [ - "windows-link", + "windows-link 0.2.1", +] + +[[package]] +name = "windows-strings" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56e6c93f3a0c3b36176cb1327a4958a0353d5d166c2a35cb268ace15e91d3b57" +dependencies = [ + "windows-link 0.1.3", ] [[package]] @@ -5305,7 +5499,7 @@ version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" dependencies = [ - "windows-link", + "windows-link 0.2.1", ] [[package]] @@ -5323,7 +5517,7 @@ version = "0.61.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" dependencies = [ - "windows-link", + "windows-link 0.2.1", ] [[package]] @@ -5342,6 +5536,15 @@ dependencies = [ "windows_x86_64_msvc", ] +[[package]] +name = "windows-threading" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b66463ad2e0ea3bbf808b7f1d371311c80e115c0b71d60efc142cafbcfb057a6" +dependencies = [ + "windows-link 0.1.3", +] + [[package]] name = "windows_aarch64_gnullvm" version = "0.52.6" diff --git a/Cargo.toml b/Cargo.toml index 79fa1aa8..8c6fc137 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -70,10 +70,10 @@ opt-level = 3 # never re-exported through a local shim. And the reverse: something only # this side uses (the at-rest crypto, IMDSv2 credentials, the gateway error) # lives here, not in core. -tw-bedrock = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.86.0" } -tw-breaker = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.86.0" } -tw-dialect = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.86.0" } -tw-guard = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.86.0" } +tw-bedrock = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.89.0" } +tw-breaker = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.89.0" } +tw-dialect = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.89.0" } +tw-guard = { git = "https://github.com/ThinkWatchProject/ThinkWatch-Core.git", tag = "v0.89.0" } # Web framework axum = { version = "0.8", features = ["macros", "ws"] } diff --git a/crates/auth/src/password.rs b/crates/auth/src/password.rs index 4a2c3dae..4cb9446c 100644 --- a/crates/auth/src/password.rs +++ b/crates/auth/src/password.rs @@ -1,10 +1,7 @@ -use std::sync::{Arc, LazyLock}; - use argon2::{ Argon2, password_hash::{PasswordHasher, PasswordVerifier, phc::PasswordHash}, }; -use tokio::sync::Semaphore; const RANDOM_PASSWORD_LEN: usize = 16; const CHARSET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789!@#$%&*"; @@ -60,24 +57,13 @@ pub fn verify_password_or_dummy(password: &str, hash: Option<&str>) -> bool { // every other request scheduled on that thread, so the handlers use // these versions, which run it on the blocking pool. -/// How many Argon2 runs go at once: one per CPU, as many as the async -/// workers allowed when they ran it. Each run holds 19 MiB (`m=19456`) -/// while it lasts and the blocking pool grows to hundreds of threads, so -/// without this a burst of logins would hold gigabytes at once. -static ARGON2_PERMITS: LazyLock> = LazyLock::new(|| { - Arc::new(Semaphore::new( - std::thread::available_parallelism().map_or(1, |n| n.get()), - )) -}); - -/// `f` on the blocking pool once an [`ARGON2_PERMITS`] place is free. The -/// task keeps the place until it ends, even when the caller has gone. +/// Share CPU admission with the data plane and limit Argon2 memory using +/// the container allowance. The running task owns both places even when +/// its caller is cancelled. async fn on_blocking_pool( f: impl FnOnce() -> T + Send + 'static, ) -> Result { - // Never closed, so the acquire cannot fail; the permit inside the - // result is held all the same. - let permit = Arc::clone(&ARGON2_PERMITS).acquire_owned().await; + let permit = think_watch_common::resources::password_slot().await; tokio::task::spawn_blocking(move || { let _permit = permit; f() @@ -111,6 +97,48 @@ pub async fn verify_password_or_dummy_off_runtime(password: String, hash: Option mod tests { use super::*; + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn cancelling_callers_does_not_admit_more_running_password_jobs() { + use std::time::Duration; + let policy = think_watch_common::resources::policy(); + let jobs = policy.argon_jobs.min(policy.cpu_jobs); + let (started, mut beginnings) = tokio::sync::mpsc::unbounded_channel(); + let mut release = Vec::new(); + let mut callers = Vec::new(); + for _ in 0..jobs { + let (tx, rx) = std::sync::mpsc::channel::<()>(); + release.push(tx); + let started = started.clone(); + callers.push(tokio::spawn(on_blocking_pool(move || { + started.send(()).unwrap(); + let _ = rx.recv_timeout(Duration::from_secs(10)); + }))); + } + for _ in 0..jobs { + tokio::time::timeout(Duration::from_secs(5), beginnings.recv()) + .await + .unwrap() + .unwrap(); + } + for caller in callers { + caller.abort(); + assert!(caller.await.unwrap_err().is_cancelled()); + } + let next = tokio::spawn(on_blocking_pool(move || started.send(()).unwrap())); + assert!( + tokio::time::timeout(Duration::from_millis(50), beginnings.recv()) + .await + .is_err() + ); + drop(release); + tokio::time::timeout(Duration::from_secs(5), next) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(beginnings.recv().await, Some(())); + } + #[tokio::test] async fn the_off_runtime_versions_agree_with_the_blocking_ones() { let hash = hash_password_off_runtime("pw".into()).await.unwrap(); @@ -133,7 +161,9 @@ mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; static RUNNING: AtomicUsize = AtomicUsize::new(0); static PEAK: AtomicUsize = AtomicUsize::new(0); - let permits = std::thread::available_parallelism().map_or(1, |n| n.get()); + let permits = think_watch_common::resources::policy() + .argon_jobs + .min(think_watch_common::resources::policy().cpu_jobs); let tasks: Vec<_> = (0..permits * 4) .map(|_| { tokio::spawn(on_blocking_pool(|| { diff --git a/crates/common/Cargo.toml b/crates/common/Cargo.toml index b196e8bc..ca37e675 100644 --- a/crates/common/Cargo.toml +++ b/crates/common/Cargo.toml @@ -48,3 +48,11 @@ aws-sigv4 = { workspace = true } aws-credential-types = { workspace = true } http_1x = { package = "http", version = "1" } async-trait = "0.1" +http-body-util = { workspace = true } +http-body = "1" +sysinfo = { version = "0.36", default-features = false, features = ["system"] } + +[dev-dependencies] +async-stream = { workspace = true } +tempfile = "3" +tower = { workspace = true, features = ["util"] } diff --git a/crates/common/src/audit/logger.rs b/crates/common/src/audit/logger.rs index a0a194cd..df17df10 100644 --- a/crates/common/src/audit/logger.rs +++ b/crates/common/src/audit/logger.rs @@ -59,17 +59,15 @@ pub struct AuditLogger { /// Bounded audit channel capacity. /// -/// At ~80 bytes per entry that's about 8 MB of in-memory backlog, -/// which is fine for any reasonable host. 100k entries means a -/// 30-second ClickHouse outage at 3k req/s still survives; drops -/// are surfaced via a metric. +/// Secondary count limit. Entries include captured bodies, so a separate +/// process-wide byte budget is the primary memory bound, including retries. const AUDIT_CHANNEL_CAPACITY: usize = 100_000; /// Throttle window for the structured "audit drop" log line. The /// metric still increments on every drop, so dashboards see the /// real rate; the log is just for human inspection and one line /// per second is plenty. -fn log_audit_drop_throttled(err: &tokio::sync::mpsc::error::TrySendError) { +fn log_audit_drop_throttled(err: &impl std::fmt::Display) { use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; static LAST_LOG_SECS: AtomicU64 = AtomicU64::new(0); @@ -272,6 +270,8 @@ impl AuditLogger { // the REMAINING capacity, so depth = total - capacity. let depth = AUDIT_CHANNEL_CAPACITY.saturating_sub(tx_for_gauge.capacity()); metrics::gauge!("audit_log_queue_depth").set(depth as f64); + metrics::gauge!("audit_log_buffered_bytes") + .set(crate::resources::audit_budget().used() as f64); } } }); @@ -297,7 +297,7 @@ impl AuditLogger { .store(bps, std::sync::atomic::Ordering::Relaxed); } - pub fn log(&self, entry: AuditEntry) { + pub fn log(&self, mut entry: AuditEntry) { // Sampling: consult the atomic rate and skip the entry when a // uniform draw falls outside the keep window. A full keep // (10000) short-circuits the RNG so the hot path pays nothing @@ -309,6 +309,12 @@ impl AuditLogger { metrics::counter!("audit_log_sampled_out_total").increment(1); return; } + if !entry.reserve_memory(crate::resources::audit_budget()) { + metrics::counter!("audit_log_dropped_total").increment(1); + metrics::counter!("audit_log_memory_dropped_total").increment(1); + log_audit_drop_throttled(&"audit byte budget occupied (queue and retained batches)"); + return; + } if let Err(e) = self.tx.try_send(entry) { // Compliance signal: dropped audit entries are a real // operational incident, not a debug warning. Bump the diff --git a/crates/common/src/audit/types.rs b/crates/common/src/audit/types.rs index 61202580..212a6094 100644 --- a/crates/common/src/audit/types.rs +++ b/crates/common/src/audit/types.rs @@ -132,6 +132,10 @@ impl std::fmt::Display for BodyCaptureStatus { /// back out and feeds them into `send_webhook` again. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct AuditEntry { + /// Account for queued and retained data through its final clone. Never part + /// of ClickHouse rows or the durable webhook JSON contract. + #[serde(skip)] + memory: Option>, pub id: String, #[serde(skip, default = "default_log_type")] pub log_type: LogType, @@ -407,6 +411,7 @@ impl AuditEntry { )] pub fn new(action: impl Into) -> Self { Self { + memory: None, id: Uuid::new_v4().to_string(), log_type: LogType::Audit, user_id: None, @@ -430,6 +435,61 @@ impl AuditEntry { } } + pub(super) fn reserve_memory( + &mut self, + budget: &std::sync::Arc, + ) -> bool { + fn value_bytes(value: &serde_json::Value) -> usize { + use serde_json::Value; + match value { + Value::String(s) => s.capacity(), + Value::Array(values) => values + .capacity() + .saturating_mul(std::mem::size_of::()) + .saturating_add(values.iter().map(value_bytes).sum::()), + Value::Object(values) => values.iter().fold(0usize, |total, (key, value)| { + total + .saturating_add(key.capacity()) + .saturating_add(std::mem::size_of::<(String, Value)>() * 2) + .saturating_add(value_bytes(value)) + }), + _ => 0, + } + } + let strings = [ + Some(&self.id), + Some(&self.action), + Some(&self.created_at), + self.user_id.as_ref(), + self.user_email.as_ref(), + self.api_key_id.as_ref(), + self.api_key_lineage_id.as_ref(), + self.resource.as_ref(), + self.resource_id.as_ref(), + self.ip_address.as_ref(), + self.user_agent.as_ref(), + self.trace_id.as_ref(), + self.session_id.as_ref(), + self.request_body.as_ref(), + self.response_body.as_ref(), + self.body_capture_status.as_ref(), + ]; + let bytes = strings + .into_iter() + .flatten() + .fold(std::mem::size_of::(), |n, s| { + n.saturating_add(s.capacity()) + }) + .saturating_add(self.detail.as_ref().map_or(0, value_bytes)); + // Flush mappers and forwarders temporarily clone/serialize one entry. + let mut lease = budget.reserve(); + if !lease.grow(bytes.saturating_mul(2)) { + return false; + } + self.memory = Some(std::sync::Arc::new(lease)); + true + } + /// Create entry for gateway request logs. /// /// Bare constructor — same deprecation reasoning as `::new`. Use @@ -680,3 +740,39 @@ mod tests { assert_eq!(decoded.detail, original.detail); } } + +#[cfg(test)] +#[allow(deprecated)] +mod memory_tests { + use super::*; + use crate::resources::ByteBudget; + + #[test] + fn queue_and_retry_clones_retain_body_admission_without_changing_wire_json() { + let budget = ByteBudget::new(32 * 1024); + let mut entry = AuditEntry::new("tools.call").request_body("x".repeat(4096)); + let before = serde_json::to_value(&entry).unwrap(); + assert!(entry.reserve_memory(&budget)); + assert!(budget.used() >= 8192); + assert_eq!(before, serde_json::to_value(&entry).unwrap()); + let retained = entry.clone(); + drop(entry); + assert!(budget.used() >= 8192); + drop(retained); + assert_eq!(budget.used(), 0); + } + + #[test] + fn a_few_large_entries_hit_the_byte_limit_before_the_count_limit() { + let budget = ByteBudget::new(16 * 1024); + let mut first = AuditEntry::new("gateway").response_body("x".repeat(4096)); + let mut second = first.clone(); + assert!(first.reserve_memory(&budget)); + let used = budget.used(); + assert!(!second.reserve_memory(&budget)); + assert_eq!(used, budget.used()); + drop(first); + assert_eq!(budget.used(), 0); + assert!(second.reserve_memory(&budget)); + } +} diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs index 334950bd..495963c5 100644 --- a/crates/common/src/lib.rs +++ b/crates/common/src/lib.rs @@ -43,6 +43,7 @@ pub mod clickhouse_client; pub mod cost_decimal; // Decimal ↔ raw i64/i128 helpers for CH Decimal(18, 10) columns pub mod lifecycle; // Surface-agnostic request pipeline (see lifecycle::mod docs) pub mod limits; // rate-limit & budget evaluation +pub mod resources; // automatic CPU admission and process-wide byte budgets // --- Utilities --- pub mod crypto; // AES-256-GCM envelope for secrets at rest diff --git a/crates/common/src/resources.rs b/crates/common/src/resources.rs new file mode 100644 index 00000000..8f69b915 --- /dev/null +++ b/crates/common/src/resources.rs @@ -0,0 +1,465 @@ +//! Resource admission shared by AI, MCP, authentication and audit dispatch. +//! Limits describe this process (including cgroups), never the host alone. +//! Byte reservations survive cancellation, shared buffers and audit retries. + +use std::sync::{Arc, Mutex, OnceLock}; +use std::time::Duration; + +use bytes::Bytes; +use tokio::sync::{OwnedSemaphorePermit, Semaphore}; + +#[cfg(any(target_os = "linux", test))] +mod cgroup; +pub mod http; + +const MIB: usize = 1024 * 1024; +// Compact JSON arrays can use many times their wire size as Value nodes. This +// also covers conversion/capture copies rather than budgeting raw bytes alone. +pub const BODY_HEADROOM: usize = 32; +pub const IDLE_TIMEOUT: Duration = Duration::from_secs(10); + +#[derive(Debug)] +pub struct Policy { + pub io_workers: usize, + pub cpu_jobs: usize, + pub argon_jobs: usize, + pub body_bytes: usize, + pub audit_bytes: usize, + pub max_body_bytes: usize, +} + +impl Policy { + fn for_machine(cpus: usize, memory: u64) -> Self { + let memory = usize::try_from(memory).unwrap_or(usize::MAX); + let cpus = cpus.max(1); + let body_bytes = (memory / 4).clamp(8 * MIB, 8 * 1024 * MIB); + Self { + io_workers: cpus.min(8), + cpu_jobs: cpus.min((memory / (128 * MIB)).max(1)), + // Argon2's 19 MiB workspace, allocator headroom, and at most 1/8 + // of the process allowance. CPU admission is shared with scanning. + argon_jobs: cpus.min((memory / 8 / (32 * MIB)).max(1)), + body_bytes, + audit_bytes: (memory / 16).clamp(2 * MIB, 128 * MIB), + max_body_bytes: (body_bytes / BODY_HEADROOM).min(64 * MIB), + } + } +} + +pub fn policy() -> &'static Policy { + static POLICY: OnceLock = OnceLock::new(); + POLICY.get_or_init(|| { + let mut system = sysinfo::System::new(); + system.refresh_memory(); + let physical = match system.total_memory() { + 0 => 1 << 30, + memory => memory, + }; + let limit = system.cgroup_limits().map(|l| l.total_memory); + #[cfg(target_os = "linux")] + let limit = cgroup::memory_limit().or(limit); + let memory = limit + .filter(|&n| n > 0) + .map_or(physical, |n| n.min(physical)); + Policy::for_machine( + std::thread::available_parallelism().map_or(1, usize::from), + memory, + ) + }) +} + +struct Scheduler { + slots: Arc, + inline: Semaphore, +} + +impl Scheduler { + fn new(policy: &Policy) -> Self { + Self { + slots: Arc::new(Semaphore::new(policy.cpu_jobs)), + inline: Semaphore::new(policy.io_workers.div_ceil(2)), + } + } + + async fn run(&self, bytes: usize, f: impl FnOnce() -> T) -> T { + let _slot = self.slots.clone().acquire_owned().await.unwrap(); + let inline = (bytes < MIB) + .then(|| self.inline.try_acquire().ok()) + .flatten(); + let multi = tokio::runtime::Handle::current().runtime_flavor() + == tokio::runtime::RuntimeFlavor::MultiThread; + if inline.is_some() || !multi { + f() + } else { + tokio::task::block_in_place(f) + } + } +} + +fn scheduler() -> &'static Scheduler { + static SCHEDULER: OnceLock = OnceLock::new(); + SCHEDULER.get_or_init(|| Scheduler::new(policy())) +} + +pub async fn cpu(bytes: usize, f: impl FnOnce() -> T) -> T { + scheduler().run(bytes, f).await +} + +/// Acquire before spawning; cancellation of the caller must not release a +/// running password operation's CPU or memory admission. +pub async fn password_slot() -> (OwnedSemaphorePermit, OwnedSemaphorePermit) { + static PASSWORDS: OnceLock> = OnceLock::new(); + let memory = PASSWORDS.get_or_init(|| Arc::new(Semaphore::new(policy().argon_jobs))); + let memory = memory.clone().acquire_owned().await.unwrap(); + let cpu = scheduler().slots.clone().acquire_owned().await.unwrap(); + (memory, cpu) +} + +#[derive(Debug)] +pub struct ByteBudget { + limit: usize, + used: Mutex, +} + +impl ByteBudget { + pub fn new(limit: usize) -> Arc { + Arc::new(Self { + limit, + used: Mutex::new(0), + }) + } + + pub fn reserve(self: &Arc) -> Reservation { + Reservation { + budget: self.clone(), + charged: 0, + } + } + + pub fn used(&self) -> usize { + *self.used.lock().unwrap_or_else(|e| e.into_inner()) + } +} + +#[derive(Debug)] +pub struct Reservation { + budget: Arc, + charged: usize, +} + +impl Reservation { + pub fn grow(&mut self, bytes: usize) -> bool { + let extra = bytes.saturating_sub(self.charged); + let mut used = self.budget.used.lock().unwrap_or_else(|e| e.into_inner()); + if extra > self.budget.limit.saturating_sub(*used) { + return false; + } + *used += extra; + self.charged += extra; + true + } + + pub fn wrap(self, data: Bytes) -> Bytes { + Bytes::from_owner(Buffered { + data, + reservation: Arc::new(self), + }) + } +} + +impl Drop for Reservation { + fn drop(&mut self) { + *self.budget.used.lock().unwrap_or_else(|e| e.into_inner()) -= self.charged; + } +} + +struct Buffered { + data: Bytes, + reservation: Arc, +} + +impl AsRef<[u8]> for Buffered { + fn as_ref(&self) -> &[u8] { + let _ = &self.reservation; + &self.data + } +} + +pub fn body_budget() -> &'static Arc { + static BUDGET: OnceLock> = OnceLock::new(); + BUDGET.get_or_init(|| ByteBudget::new(policy().body_bytes)) +} + +pub fn audit_budget() -> &'static Arc { + static BUDGET: OnceLock> = OnceLock::new(); + BUDGET.get_or_init(|| ByteBudget::new(policy().audit_bytes)) +} + +/// A stream's retained timeline, shared with its detached post-invoke tail. +pub struct StreamBudget { + reservation: Reservation, + received: usize, + max: usize, +} + +impl StreamBudget { + pub fn new(budget: &Arc, max: usize) -> Self { + Self { + reservation: budget.reserve(), + received: 0, + max, + } + } + + pub fn append(&mut self, bytes: usize) -> bool { + let received = self.received.saturating_add(bytes); + if received > self.max + || !self + .reservation + .grow(received.saturating_mul(BODY_HEADROOM)) + { + return false; + } + self.received = received; + true + } +} + +pub type SharedStreamBudget = Arc>; + +pub fn stream_budget() -> SharedStreamBudget { + Arc::new(Mutex::new(StreamBudget::new( + body_budget(), + policy().max_body_bytes, + ))) +} + +#[derive(Debug, thiserror::Error)] +pub enum ReadError { + #[error("The response exceeds the gateway buffering limit")] + TooLarge, + #[error("The gateway buffering budget is occupied; retry shortly")] + Busy, + #[error(transparent)] + Transport(#[from] reqwest::Error), +} + +/// Charge before appending, including chunked and decompressed responses. Keep +/// room for JSON, conversion and allocator capacity; never wait holding a partial +/// body, since concurrent readers could otherwise deadlock each other. +pub async fn read_response(response: reqwest::Response, max: usize) -> Result { + let max = max.min(policy().max_body_bytes); + read_response_with_budget(response, max, body_budget()).await +} + +async fn read_response_with_budget( + mut response: reqwest::Response, + max: usize, + budget: &Arc, +) -> Result { + if response.content_length().is_some_and(|n| n > max as u64) { + return Err(ReadError::TooLarge); + } + let mut reservation = budget.reserve(); + let mut data = Vec::new(); + while let Some(chunk) = response.chunk().await? { + let size = data + .len() + .checked_add(chunk.len()) + .ok_or(ReadError::TooLarge)?; + if size > max { + return Err(ReadError::TooLarge); + } + if !reservation.grow(size.saturating_mul(BODY_HEADROOM)) { + metrics::counter!("gateway_memory_busy_total").increment(1); + return Err(ReadError::Busy); + } + data.extend_from_slice(&chunk); + } + Ok(reservation.wrap(data.into())) +} + +/// A transformed/captured body still owns the original read reservation. +pub fn retain(data: Bytes, source: Bytes) -> Bytes { + struct Retained { + data: Bytes, + _source: Bytes, + } + impl AsRef<[u8]> for Retained { + fn as_ref(&self) -> &[u8] { + &self.data + } + } + Bytes::from_owner(Retained { + data, + _source: source, + }) +} + +/// Reserve before fetching a Redis cache entry. STRLEN followed by bounded +/// GETRANGE also stays bounded if another writer replaces the key in between. +pub async fn cached_body(redis: &fred::clients::Client, key: &str) -> Option { + use fred::interfaces::KeysInterface; + let length: usize = redis.strlen(key).await.ok()?; + if length == 0 || length > policy().max_body_bytes { + return None; + } + let mut lease = body_budget().reserve(); + if !lease.grow(length.saturating_mul(BODY_HEADROOM)) { + return None; + } + let data: Vec = redis.getrange(key, 0, length).await.ok()?; + if data.len() > length { + return None; + } + Some(lease.wrap(data.into())) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn memory_limits_allocation_without_capping_a_well_provisioned_eighty_core_server() { + let large = Policy::for_machine(80, 128 << 30); + assert_eq!( + (large.io_workers, large.cpu_jobs, large.argon_jobs), + (8, 80, 80) + ); + let small = Policy::for_machine(80, 256 << 20); + assert_eq!((small.cpu_jobs, small.argon_jobs), (2, 1)); + assert!(small.body_bytes + small.audit_bytes <= 256 * MIB / 2); + } + + #[test] + fn reservations_follow_shared_slices_and_failed_growth_is_atomic() { + let budget = ByteBudget::new(10); + let mut lease = budget.reserve(); + assert!(lease.grow(8)); + assert!(!lease.grow(11)); + let data = lease.wrap(Bytes::from_static(b"hello")); + let slice = data.slice(1..3); + drop(data); + assert_eq!(budget.used(), 8); + drop(slice); + assert_eq!(budget.used(), 0); + } + + #[tokio::test] + async fn chunked_responses_obey_size_and_shared_memory_limits() { + async fn response() -> (reqwest::Response, tokio::task::JoinHandle<()>) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}/", listener.local_addr().unwrap()); + let router = axum::Router::new().route( + "/", + axum::routing::get(|| async { + axum::body::Body::from_stream(futures::stream::iter([ + Ok::<_, std::io::Error>(Bytes::from_static(b"12345678")), + Ok(Bytes::from_static(b"abcdefgh")), + ])) + }), + ); + let server = tokio::spawn(async move { axum::serve(listener, router).await.unwrap() }); + (reqwest::get(url).await.unwrap(), server) + } + for (limit, max, expected) in [ + (10 * BODY_HEADROOM, 32, "busy"), + (32 * BODY_HEADROOM, 12, "large"), + (32 * BODY_HEADROOM, 32, "ok"), + ] { + let budget = ByteBudget::new(limit); + let (response, server) = response().await; + assert!(response.content_length().is_none()); + match read_response_with_budget(response, max, &budget).await { + Err(ReadError::Busy) => assert_eq!(expected, "busy"), + Err(ReadError::TooLarge) => assert_eq!(expected, "large"), + Ok(bytes) => { + assert_eq!(expected, "ok"); + assert_eq!(&bytes[..], b"12345678abcdefgh"); + assert_eq!(budget.used(), 16 * BODY_HEADROOM); + let copy = bytes.clone(); + drop(bytes); + assert_eq!(budget.used(), 16 * BODY_HEADROOM); + drop(copy); + } + Err(error) => panic!("{error}"), + } + assert_eq!(budget.used(), 0); + server.abort(); + } + } + + #[test] + fn stream_timeline_reservations_cover_retained_history_and_reject_growth() { + let budget = ByteBudget::new(25 * BODY_HEADROOM); + let mut stream = StreamBudget::new(&budget, 64); + assert!(stream.append(20)); + assert!(!stream.append(10)); + assert_eq!(budget.used(), 20 * BODY_HEADROOM); + drop(stream); + assert_eq!(budget.used(), 0); + } + + #[test] + fn eighty_cpu_jobs_expand_beyond_io_workers_then_threads_retire() { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Instant; + let alive = Arc::new(AtomicUsize::new(0)); + let (started, stopped) = (alive.clone(), alive.clone()); + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .thread_keep_alive(Duration::from_millis(50)) + .on_thread_start(move || { + started.fetch_add(1, Ordering::SeqCst); + }) + .on_thread_stop(move || { + stopped.fetch_sub(1, Ordering::SeqCst); + }) + .enable_all() + .build() + .unwrap(); + let scheduler = Arc::new(Scheduler::new(&Policy { + io_workers: 2, + ..Policy::for_machine(80, 128 << 30) + })); + runtime.block_on(async { + let release = Arc::new((Mutex::new(false), std::sync::Condvar::new())); + let (began, mut beginnings) = tokio::sync::mpsc::unbounded_channel(); + let mut jobs = Vec::new(); + for _ in 0..80 { + let (scheduler, release, began) = + (scheduler.clone(), release.clone(), began.clone()); + jobs.push(tokio::spawn(async move { + scheduler + .run(4096, || { + began.send(()).unwrap(); + let guard = release.0.lock().unwrap(); + let _guard = release + .1 + .wait_timeout_while(guard, Duration::from_secs(10), |done| !*done) + .unwrap(); + }) + .await; + })); + } + for _ in 0..80 { + tokio::time::timeout(Duration::from_secs(5), beginnings.recv()) + .await + .expect("all 80 calls must start") + .unwrap(); + } + assert_eq!(scheduler.slots.available_permits(), 0); + tokio::time::sleep(Duration::from_millis(1)).await; + *release.0.lock().unwrap() = true; + release.1.notify_all(); + for job in jobs { + job.await.unwrap(); + } + }); + let deadline = Instant::now() + Duration::from_secs(5); + while alive.load(Ordering::SeqCst) != 2 && Instant::now() < deadline { + std::thread::sleep(Duration::from_millis(10)); + } + assert_eq!(alive.load(Ordering::SeqCst), 2); + } +} diff --git a/crates/common/src/resources/cgroup.rs b/crates/common/src/resources/cgroup.rs new file mode 100644 index 00000000..72ae6aa1 --- /dev/null +++ b/crates/common/src/resources/cgroup.rs @@ -0,0 +1,89 @@ +//! 同时看当前 cgroup 和它的祖先，systemd 的子组也不能越过父组的内存限额。 + +use std::path::{Component, Path}; + +#[cfg(target_os = "linux")] +pub(super) fn memory_limit() -> Option { + let membership = std::fs::read_to_string("/proc/self/cgroup").ok()?; + membership + .lines() + .filter_map(|line| { + let mut fields = line.splitn(3, ':'); + let hierarchy = fields.next()?; + let controllers = fields.next()?; + let path = fields.next()?; + if hierarchy == "0" && controllers.is_empty() { + lowest_limit(Path::new("/sys/fs/cgroup"), Path::new(path), "memory.max") + } else if controllers.split(',').any(|c| c == "memory") { + lowest_limit( + Path::new("/sys/fs/cgroup/memory"), + Path::new(path), + "memory.limit_in_bytes", + ) + } else { + None + } + }) + .min() +} + +fn lowest_limit(root: &Path, membership: &Path, file: &str) -> Option { + // cgroup namespace 可以把当前组挂成根目录。深路径不存在时，仍然检查挂载根的限额。 + if membership + .components() + .any(|c| matches!(c, Component::ParentDir | Component::Prefix(_))) + { + return None; + } + let relative = membership.strip_prefix("/").unwrap_or(membership); + let current = root.join(relative); + current + .ancestors() + .take_while(|p| p.starts_with(root)) + .filter_map(|p| { + std::fs::read_to_string(p.join(file)) + .ok()? + .trim() + .parse::() + .ok() + .filter(|&n| n > 0) + }) + .min() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn nested_groups_obey_the_smallest_ancestor_limit() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path(); + std::fs::create_dir_all(root.join("service/child")).unwrap(); + std::fs::write(root.join("memory.max"), "max\n").unwrap(); + std::fs::write(root.join("service/memory.max"), "536870912\n").unwrap(); + std::fs::write(root.join("service/child/memory.max"), "1073741824\n").unwrap(); + assert_eq!( + lowest_limit(root, Path::new("/service/child"), "memory.max"), + Some(512 << 20) + ); + } + + #[test] + fn a_namespaced_mount_root_and_v1_limits_are_supported() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write(dir.path().join("memory.limit_in_bytes"), "268435456").unwrap(); + assert_eq!( + lowest_limit( + dir.path(), + Path::new("/unavailable/host/path"), + "memory.limit_in_bytes" + ), + Some(256 << 20) + ); + assert_eq!( + lowest_limit(dir.path(), Path::new("/../escape"), "memory.limit_in_bytes"), + None + ); + } +} diff --git a/crates/common/src/resources/http.rs b/crates/common/src/resources/http.rs new file mode 100644 index 00000000..e07f7eb0 --- /dev/null +++ b/crates/common/src/resources/http.rs @@ -0,0 +1,227 @@ +//! Body admission precedes extraction; checking in the handler is too late. + +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use axum::body::{Body, Bytes, HttpBody}; +use axum::extract::Request; +use axum::middleware::Next; +use axum::response::{IntoResponse, Response}; +use futures::StreamExt; + +use super::{ByteBudget, policy}; + +/// Preserve Axum's JSON errors while moving large/concurrent decoding off I/O. +pub struct ScheduledJson(pub T); + +impl axum::extract::FromRequest for ScheduledJson +where + S: Send + Sync, + T: serde::de::DeserializeOwned, +{ + type Rejection = Response; + + async fn from_request(request: Request, state: &S) -> Result { + let json = request + .headers() + .get(axum::http::header::CONTENT_TYPE) + .and_then(|h| h.to_str().ok()) + .and_then(|h| h.split(';').next()) + .is_some_and(|h| { + let h = h.trim().to_ascii_lowercase(); + h == "application/json" || (h.starts_with("application/") && h.ends_with("+json")) + }); + if !json { + return Err(( + axum::http::StatusCode::UNSUPPORTED_MEDIA_TYPE, + "Expected request with Content-Type: application/json", + ) + .into_response()); + } + let bytes = Bytes::from_request(request, state) + .await + .map_err(IntoResponse::into_response)?; + super::cpu(bytes.len(), || axum::Json::::from_bytes(&bytes)) + .await + .map(|axum::Json(value)| Self(value)) + .map_err(IntoResponse::into_response) + } +} + +pub async fn admit(request: Request, next: Next) -> Response { + admit_with_budget(request, next, super::body_budget(), policy().max_body_bytes).await +} + +async fn admit_with_budget( + request: Request, + next: Next, + budget: &Arc, + max: usize, +) -> Response { + let (parts, body) = request.into_parts(); + if body.size_hint().upper().is_some_and(|n| n > max as u64) { + return rejection(axum::http::StatusCode::PAYLOAD_TOO_LARGE); + } + let mut data = Vec::new(); + let mut lease = budget.reserve(); + let mut chunks = body.into_data_stream(); + while let Some(chunk) = chunks.next().await { + let chunk = match chunk { + Ok(chunk) => chunk, + Err(error) => { + use std::error::Error; + let status = if error + .source() + .is_some_and(|e| e.is::()) + { + axum::http::StatusCode::PAYLOAD_TOO_LARGE + } else { + axum::http::StatusCode::BAD_REQUEST + }; + return rejection(status); + } + }; + let size = data.len().saturating_add(chunk.len()); + if size > max { + return rejection(axum::http::StatusCode::PAYLOAD_TOO_LARGE); + } + if !lease.grow(size.saturating_mul(super::BODY_HEADROOM)) { + return rejection(axum::http::StatusCode::TOO_MANY_REQUESTS); + } + data.extend_from_slice(&chunk); + } + let data = lease.wrap(data.into()); + let mut response = next + .run(Request::from_parts(parts, Body::from(data.clone()))) + .await; + let body = std::mem::replace(response.body_mut(), Body::empty()); + *response.body_mut() = Body::new(HeldBody { + inner: body, + _request: data, + }); + response +} + +fn rejection(status: axum::http::StatusCode) -> Response { + let busy = status == axum::http::StatusCode::TOO_MANY_REQUESTS; + let message = if busy { + "The gateway is buffering other requests. Retry shortly." + } else if status == axum::http::StatusCode::PAYLOAD_TOO_LARGE { + "The request body exceeds the gateway buffering limit." + } else { + "The request body could not be read." + }; + let mut response = (status, axum::Json(serde_json::json!({ + "error": {"type": if busy { "rate_limit_error" } else { "invalid_request_error" }, "message": message} + }))).into_response(); + if busy { + response + .headers_mut() + .insert(axum::http::header::RETRY_AFTER, "1".parse().unwrap()); + metrics::counter!("gateway_memory_busy_total").increment(1); + } + response +} + +/// A JSON extractor discards the input Bytes while its decoded object still +/// exists. Retain admission through the response, including streaming and drop. +struct HeldBody { + inner: Body, + _request: Bytes, +} + +impl HttpBody for HeldBody { + type Data = Bytes; + type Error = axum::Error; + + fn poll_frame( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + Pin::new(&mut self.inner).poll_frame(cx) + } + + fn is_end_stream(&self) -> bool { + self.inner.is_end_stream() + } + fn size_hint(&self) -> http_body::SizeHint { + self.inner.size_hint() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use axum::{Router, routing::post}; + use tower::ServiceExt; + + fn router(budget: Arc, max: usize) -> Router { + Router::new() + .route("/", post(|body: Bytes| async move { body })) + .layer(axum::middleware::from_fn( + move |request: Request, next: Next| { + let budget = budget.clone(); + async move { admit_with_budget(request, next, &budget, max).await } + }, + )) + } + + #[tokio::test] + async fn chunked_pressure_returns_retryable_429_and_frees_partial_buffers() { + let budget = ByteBudget::new(25 * super::super::BODY_HEADROOM); + let app = router(budget.clone(), 64); + let chunks = futures::stream::iter([ + Ok::<_, std::io::Error>(Bytes::from(vec![b'a'; 20])), + Ok(Bytes::from(vec![b'b'; 10])), + ]); + let request = Request::post("/").body(Body::from_stream(chunks)).unwrap(); + let response = app.clone().oneshot(request).await.unwrap(); + assert_eq!(response.status(), 429); + assert_eq!(response.headers()["retry-after"], "1"); + assert_eq!(budget.used(), 0); + + let request = Request::post("/").body(Body::from(vec![b'x'; 16])).unwrap(); + let response = app.oneshot(request).await.unwrap(); + assert_eq!(response.status(), 200); + assert_eq!( + budget.used(), + 16 * super::super::BODY_HEADROOM, + "input admission outlives the extractor" + ); + let body = axum::body::to_bytes(response.into_body(), 100) + .await + .unwrap(); + assert_eq!(body.len(), 16); + drop(body); + assert_eq!(budget.used(), 0); + } + + #[tokio::test] + async fn an_oversize_body_is_413_and_cancellation_releases_partial_admission() { + let budget = ByteBudget::new(25 * super::super::BODY_HEADROOM); + let app = router(budget.clone(), 20); + let response = app + .clone() + .oneshot(Request::post("/").body(Body::from(vec![0; 21])).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), 413); + assert_eq!(budget.used(), 0); + + let sent = Arc::new(tokio::sync::Notify::new()); + let started = sent.clone(); + let chunks = async_stream::stream! { + yield Ok::<_, std::io::Error>(Bytes::from_static(b"12345678")); + started.notify_one(); + std::future::pending::<()>().await; + }; + let task = + tokio::spawn(app.oneshot(Request::post("/").body(Body::from_stream(chunks)).unwrap())); + sent.notified().await; + assert_eq!(budget.used(), 8 * super::super::BODY_HEADROOM); + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + assert_eq!(budget.used(), 0); + } +} diff --git a/crates/gateway/src/bedrock/sigv4.rs b/crates/gateway/src/bedrock/sigv4.rs index 632d9940..f6dc4f31 100644 --- a/crates/gateway/src/bedrock/sigv4.rs +++ b/crates/gateway/src/bedrock/sigv4.rs @@ -195,27 +195,29 @@ impl Signer { let fail = |what: &str, e: reqwest::Error| SignError::Credentials(format!("{what}: {e}")); let imds = &self.imds; - let token = client + let token_response = client .put(format!("{imds}/latest/api/token")) .header("X-aws-ec2-metadata-token-ttl-seconds", "300") .send() .await .and_then(reqwest::Response::error_for_status) - .map_err(|e| fail("IMDSv2 token request", e))? - .text() + .map_err(|e| fail("IMDSv2 token request", e))?; + let token_bytes = think_watch_common::resources::read_response(token_response, 64 * 1024) .await - .map_err(|e| fail("IMDSv2 token read", e))?; + .map_err(|e| SignError::Credentials(format!("IMDSv2 token read: {e}")))?; + let token = String::from_utf8_lossy(&token_bytes); - let role = client + let role_response = client .get(format!("{imds}/latest/meta-data/iam/security-credentials/")) - .header("X-aws-ec2-metadata-token", &token) + .header("X-aws-ec2-metadata-token", token.as_ref()) .send() .await .and_then(reqwest::Response::error_for_status) - .map_err(|e| fail("IMDSv2 role lookup", e))? - .text() + .map_err(|e| fail("IMDSv2 role lookup", e))?; + let role_bytes = think_watch_common::resources::read_response(role_response, 64 * 1024) .await - .map_err(|e| fail("IMDSv2 role read", e))?; + .map_err(|e| SignError::Credentials(format!("IMDSv2 role read: {e}")))?; + let role = String::from_utf8_lossy(&role_bytes); // One role per instance profile; the listing is one name per line let role = role.lines().next().unwrap_or_default().trim(); if role.is_empty() { @@ -224,18 +226,21 @@ impl Signer { )); } - let creds: serde_json::Value = client + let credentials_response = client .get(format!( "{imds}/latest/meta-data/iam/security-credentials/{role}" )) - .header("X-aws-ec2-metadata-token", &token) + .header("X-aws-ec2-metadata-token", token.as_ref()) .send() .await .and_then(reqwest::Response::error_for_status) - .map_err(|e| fail("IMDSv2 credentials fetch", e))? - .json() - .await - .map_err(|e| fail("IMDSv2 credentials parse", e))?; + .map_err(|e| fail("IMDSv2 credentials fetch", e))?; + let credentials_bytes = + think_watch_common::resources::read_response(credentials_response, 64 * 1024) + .await + .map_err(|e| SignError::Credentials(format!("IMDSv2 credentials read: {e}")))?; + let creds: serde_json::Value = serde_json::from_slice(&credentials_bytes) + .map_err(|e| SignError::Credentials(format!("IMDSv2 credentials parse: {e}")))?; let field = |k: &str| { creds[k] diff --git a/crates/gateway/src/cache.rs b/crates/gateway/src/cache.rs index e8e3b933..42eef3c9 100644 --- a/crates/gateway/src/cache.rs +++ b/crates/gateway/src/cache.rs @@ -7,6 +7,7 @@ use xxhash_rust::xxh3::xxh3_128; /// A cached answer, in the caller's format with redaction placeholders intact. pub struct Cached { + pub memory: bytes::Bytes, pub body: Vec, /// Kept so a hit counts toward token limits and budgets the way the /// original call did. @@ -135,17 +136,20 @@ impl ResponseCache { /// Look up a cached answer. pub async fn get(&self, fingerprint: &[u8]) -> Option { let key = Self::cache_key_for(fingerprint); - let stored: Option = self.redis.get(&key).await.ok().flatten(); - stored.and_then(|json| { - serde_json::from_str::(&json) + let json = think_watch_common::resources::cached_body(&self.redis, &key).await?; + let memory = json.clone(); + think_watch_common::resources::cpu(json.len(), || { + serde_json::from_slice::(&json) .map_err(|e| tracing::warn!("Failed to read a cached response: {e}")) .ok() .map(|s| Cached { + memory, body: serde_json::to_vec(&s.body).unwrap_or_default(), prompt_tokens: s.prompt_tokens, completion_tokens: s.completion_tokens, }) }) + .await } /// Invalidate all cached responses by deleting keys matching the diff --git a/crates/gateway/src/lifecycle/mod.rs b/crates/gateway/src/lifecycle/mod.rs index 0d574b43..a844499d 100644 --- a/crates/gateway/src/lifecycle/mod.rs +++ b/crates/gateway/src/lifecycle/mod.rs @@ -51,6 +51,8 @@ pub(crate) struct ChatCompletionSurface; /// A whole answer, in the caller's format. pub struct Completed { + /// Retain the upstream byte reservation through conversion and audit. + pub memory: Option, /// As the upstream answered, in the caller's format, with redaction /// placeholders still in place and the answer's text not yet through /// the content rules — the form the cache stores, so a later caller @@ -96,6 +98,7 @@ pub struct ChatStreamCaptured { /// attempt and by the streaming tail task. Owned rather than borrowed so /// the detached tail needs no `'static` borrow. pub(crate) struct ChatRequestSnapshot { + pub memory: think_watch_common::resources::SharedStreamBudget, pub identity: GatewayRequestIdentity, /// The one id every audit row and gateway log carries for this request. pub trace_id: String, @@ -108,7 +111,7 @@ pub(crate) struct ChatRequestSnapshot { /// stripped anything, before outbound redaction: the audit row is the /// record of what the user wrote. Body capture applies its own /// redaction toggle on top. - pub request_for_audit: Vec, + pub request_for_audit: Bytes, /// Where the cache keeps this request's answer. `None` when the /// request must not be cached. pub cache_fingerprint: Option>, @@ -260,6 +263,7 @@ pub(crate) fn build_chat_pump( provider: &str, min_speed: Option, ) -> (axum::response::Response, StreamTail) { + let memory = request.memory.clone(); let readers = Arc::new(Mutex::new(Readers::default())); let readers_for_tail = Arc::clone(&readers); @@ -460,7 +464,7 @@ pub(crate) fn build_chat_pump( backoff = Some(wait); continue 'attempt; } - match deliver(&h.release(), &mut convert, &mut shaper, inspector.as_mut(), client) { + match think_watch_common::resources::cpu(0, || deliver(&h.release(), &mut convert, &mut shaper, inspector.as_mut(), client)).await { Ok(out) => { if !out.is_empty() { if last_frame.is_in(&out) { @@ -490,16 +494,27 @@ pub(crate) fn build_chat_pump( }; match item { Ok(chunk) => { + if !memory.lock().unwrap_or_else(|e| e.into_inner()).append(chunk.len()) { + metrics::counter!("gateway_memory_busy_total").increment(1); + let err = GatewayError::rate_limited("gateway_memory", 1); + let out = cut(&shaper, convert.as_mut(), client, &[], &err); + tell(&mut done_tx, StreamOutcome::UpstreamError { + error_type: err.error_tag().into(), + message: err.to_string(), status_code: err.status_code(), + }); + yield Ok(Bytes::from(out)); + return; + } if let Some(s) = speed.as_mut() { s.feed(&chunk, std::time::Instant::now()); } - let (sealed, said) = match readers.lock() { + let (sealed, said) = think_watch_common::resources::cpu(chunk.len(), || match readers.lock() { Ok(mut r) => { r.feed(&chunk); (r.sealed_refusal(), r.reported(false)) } Err(_) => (None, None), - }; + }).await; // The upstream cannot read the conversation's encrypted // history: the client is told why in plain words, in // place of the upstream's own error. @@ -538,7 +553,7 @@ pub(crate) fn build_chat_pump( hold.take().map(Opening::release).unwrap_or_default() } }; - match deliver(&chunk, &mut convert, &mut shaper, inspector.as_mut(), client) { + match think_watch_common::resources::cpu(chunk.len(), || deliver(&chunk, &mut convert, &mut shaper, inspector.as_mut(), client)).await { Ok(out) => { if !out.is_empty() { if last_frame.is_in(&out) { @@ -569,7 +584,7 @@ pub(crate) fn build_chat_pump( continue 'attempt; } // The opening it sent goes out before the error. - match deliver(&h.release(), &mut convert, &mut shaper, inspector.as_mut(), client) { + match think_watch_common::resources::cpu(0, || deliver(&h.release(), &mut convert, &mut shaper, inspector.as_mut(), client)).await { Ok(out) => { if !out.is_empty() { yield Ok(Bytes::from(out)); @@ -1513,6 +1528,7 @@ impl Surface for ChatCompletionSurface { }; if let Some(body) = body { let cached = crate::cache::Cached { + memory: Bytes::new(), body: body.clone(), prompt_tokens, completion_tokens, diff --git a/crates/gateway/src/proxy/body_capture.rs b/crates/gateway/src/proxy/body_capture.rs index c2889f8d..dc024ee0 100644 --- a/crates/gateway/src/proxy/body_capture.rs +++ b/crates/gateway/src/proxy/body_capture.rs @@ -198,7 +198,7 @@ async fn process_body( offloaded_flag: &mut bool, ) -> String { if redact_pii { - s = redaction.redact_blob(&s); + s = think_watch_common::resources::cpu(s.len(), || redaction.redact_blob(&s)).await; } if s.len() <= max_bytes { return s; diff --git a/crates/gateway/src/proxy/generate.rs b/crates/gateway/src/proxy/generate.rs index eb85bc3d..8321c059 100644 --- a/crates/gateway/src/proxy/generate.rs +++ b/crates/gateway/src/proxy/generate.rs @@ -32,6 +32,7 @@ //! are thinkwatch-core's, shared with the desktop gateway. use std::convert::Infallible; +use think_watch_common::resources; use crate::call_ctx::CallCtx; use crate::error::GatewayError; @@ -628,9 +629,12 @@ pub(crate) async fn send( ) -> Result<(reqwest::Response, Wire), GatewayError> { let upstream = &entry.upstream; let official = upstream.is_official(); - let first = outbound - .address(entry.protocol, model, official)? - .for_upstream(upstream, model); + let first = resources::cpu(usize::MAX, || { + outbound + .address(entry.protocol, model, official) + .map(|wire| wire.for_upstream(upstream, model)) + }) + .await?; // Only a request that can carry an effort asks the memory. let memory = if outbound.reasoning.any { effort::recall(redis, entry.provider_id, first.dialect, model).await @@ -724,9 +728,12 @@ pub(crate) async fn send( to = %protocol, "Upstream rejected the configured protocol — retrying with an alternate" ); - let wire = outbound - .address(*protocol, model, official)? - .for_upstream(upstream, model); + let wire = resources::cpu(usize::MAX, || { + outbound + .address(*protocol, model, official) + .map(|wire| wire.for_upstream(upstream, model)) + }) + .await?; let wire = if outbound.reasoning.any { let memory = effort::recall(redis, entry.provider_id, wire.dialect, model).await; wire.remembered(&memory) @@ -770,33 +777,46 @@ pub(crate) async fn read_whole( caller_model: &str, input_estimate: u64, ) -> Result { - let upstream = resp - .bytes() + let upstream = resources::read_response(resp, resources::policy().max_body_bytes) .await - .map_err(super::transport::transport_error)?; - - let mut sniffer = tw_dialect::usage::Sniffer::new(); - sniffer.feed(&upstream); - let (usage, usage_estimated) = - crate::usage_estimate::complete(sniffer.finish(), true, input_estimate, Some(&upstream)); - if usage_estimated { - metrics::counter!("gateway_usage_estimated_total").increment(1); - } + .map_err(|error| match error { + resources::ReadError::Busy => GatewayError::rate_limited("gateway_memory", 1), + resources::ReadError::Transport(error) => super::transport::transport_error(error), + resources::ReadError::TooLarge => { + GatewayError::ProviderInvalidResponse(error.to_string()) + } + })?; + let memory = upstream.clone(); + resources::cpu(upstream.len(), || { + let mut sniffer = tw_dialect::usage::Sniffer::new(); + sniffer.feed(&upstream); + let (usage, usage_estimated) = crate::usage_estimate::complete( + sniffer.finish(), + true, + input_estimate, + Some(&upstream), + ); + if usage_estimated { + metrics::counter!("gateway_usage_estimated_total").increment(1); + } - let body = match &wire.convert { - Some(session) => session.response(&upstream).ok_or_else(|| { - GatewayError::ProviderInvalidResponse( - "The upstream answered with something that is not JSON.".into(), - ) - })?, - None => upstream.to_vec(), - }; - Ok(Completed { - body: rewrite_model(&body, caller_model), - logged: None, - usage, - usage_estimated, + let body = match &wire.convert { + Some(session) => session.response(&upstream).ok_or_else(|| { + GatewayError::ProviderInvalidResponse( + "The upstream answered with something that is not JSON.".into(), + ) + })?, + None => upstream.to_vec(), + }; + Ok(Completed { + memory: Some(memory), + body: rewrite_model(&body, caller_model), + logged: None, + usage, + usage_estimated, + }) }) + .await } // ───────────────────────────────────────────── the pipeline @@ -860,11 +880,14 @@ async fn run( "(unknown)", request_started_at, ); - let raw: Value = serde_json::from_slice(&body).map_err(|_| { - early_ctx.emit(GatewayError::TransformError( - "The request body is not valid JSON.".into(), - )) - })?; + let original_body = body.clone(); + let raw: Value = resources::cpu(body.len(), || serde_json::from_slice(&body)) + .await + .map_err(|_| { + early_ctx.emit(GatewayError::TransformError( + "The request body is not valid JSON.".into(), + )) + })?; let (model, is_stream) = if surface.dialect == Dialect::Gemini { gemini_target(path).ok_or_else(|| { early_ctx.emit(GatewayError::TransformError(format!( @@ -927,7 +950,8 @@ async fn run( // the caller's 403, quoting their words (masked). The rules that // check answers run on the way back (the shaper for a stream, below // for a whole answer). - let screening = guards.content.screen(surface.dialect, &body); + let screening = + resources::cpu(body.len(), || guards.content.screen(surface.dialect, &body)).await; crate::content_filter::record(&state.audit, &caller, guards.content.mode, &screening); if let Some(hit) = screening.refusal() { return Err(ctx @@ -938,12 +962,14 @@ async fn run( // every hop, the audit row — the request is the changed one. let (body, raw) = match screening.body { Some(stripped) => { - let raw = serde_json::from_slice(&stripped).map_err(|_| { - ctx.emit(GatewayError::TransformError( + let raw = resources::cpu(stripped.len(), || serde_json::from_slice(&stripped)) + .await + .map_err(|_| { + ctx.emit(GatewayError::TransformError( "The request could not be read after the content filter changed text in it." .into(), )) - })?; + })?; (stripped, raw) } None => (body, raw), @@ -954,8 +980,10 @@ async fn run( // conversation kept on OpenAI's servers (`previous_response_id`), or // carries a compaction only OpenAI can read — can still be forwarded // to an upstream of its own format, so that is where it goes (step 8). - let decoded = - tw_dialect::convert::decode(surface.dialect, &raw, path, internal_query(surface.dialect)); + let decoded = resources::cpu(body.len(), || { + tw_dialect::convert::decode(surface.dialect, &raw, path, internal_query(surface.dialect)) + }) + .await; // 5. Outbound redaction: the text written in the request — located by // the structure of the caller's format, never its ids, types or @@ -968,19 +996,25 @@ async fn run( // The audit row keeps what the caller wrote (their at-rest // redaction setting applies on top). let shape = Shape::Generate(surface.dialect); - let (findings, ledger) = guards.redaction.look(&body, shape); + let (findings, ledger) = + resources::cpu(body.len(), || guards.redaction.look(&body, shape)).await; crate::redaction::record(&state.audit, &caller, &guards.redaction, &findings); - let request_for_audit = body.to_vec(); - let outbound_body = if ledger.is_empty() { - raw - } else { - let (replaced, _) = guards.redaction.replace(body.to_vec(), &ledger, shape); - serde_json::from_slice(&replaced).map_err(|_| { - ctx.emit(GatewayError::TransformError( - "The request could not be read after redaction.".into(), - )) - })? - }; + let request_for_audit = resources::retain(body.clone(), original_body); + let outbound_body = resources::cpu(body.len(), || { + if ledger.is_empty() { + Ok::<_, GatewayErrorResponse>(raw) + } else { + let (replaced, _) = guards.redaction.replace(body.to_vec(), &ledger, shape); + serde_json::from_slice(&replaced) + .map_err(|_| { + ctx.emit(GatewayError::TransformError( + "The request could not be read after redaction.".into(), + )) + }) + .map_err(Into::into) + } + }) + .await?; // 6. The model's settings. The output cap is applied to each hop as it // is addressed (`Outbound::cap_output`): whether to fill one in @@ -1005,22 +1039,27 @@ async fn run( // budgets like a real call would — the caller received them, and // otherwise repeating a cached prompt would get round every token // limit. - let cache_fingerprint = if surface.caches { - ResponseCache::fingerprint(&outbound_body).map(|mut fp| { - // The cap and the reasoning setting are applied after this, per - // hop: an answer made under one must not be served under - // another. - if let Some(cap) = max_output_tokens { - fp.extend_from_slice(format!("\nmax_output_tokens={cap}").as_bytes()); - } - if let Some(setting) = reasoning.setting { - fp.extend_from_slice(format!("\nreasoning_effort={}", setting.slug()).as_bytes()); - } - fp - }) - } else { - None - }; + let cache_fingerprint = resources::cpu(body.len(), || { + if surface.caches { + ResponseCache::fingerprint(&outbound_body).map(|mut fp| { + // The cap and the reasoning setting are applied after this, per + // hop: an answer made under one must not be served under + // another. + if let Some(cap) = max_output_tokens { + fp.extend_from_slice(format!("\nmax_output_tokens={cap}").as_bytes()); + } + if let Some(setting) = reasoning.setting { + fp.extend_from_slice( + format!("\nreasoning_effort={}", setting.slug()).as_bytes(), + ); + } + fp + }) + } else { + None + } + }) + .await; if let Some(fp) = &cache_fingerprint && let Some(cached) = state.cache.get(fp).await { @@ -1028,23 +1067,33 @@ async fn run( // With this caller's values in it: what they would run. A stored // answer passed the content rules and the inspection in force when // it was stored, not necessarily the ones in force now. - let restored = crate::redaction::restore_body(&ledger, &cached.body); - let restored = crate::content_filter::whole( - &guards.content, - &state.audit, - &caller, - "cache", - surface.dialect, - restored, - ); - if let Some(e) = crate::tool_inspection::check_whole( - &guards.tools, - &guards.redaction, - &state.audit, - &caller, - "cache", - &restored, - ) { + let restored = resources::cpu(cached.body.len(), || { + crate::redaction::restore_body(&ledger, &cached.body) + }) + .await; + let restored = resources::cpu(restored.len(), || { + crate::content_filter::whole( + &guards.content, + &state.audit, + &caller, + "cache", + surface.dialect, + restored, + ) + }) + .await; + if let Some(e) = resources::cpu(restored.len(), || { + crate::tool_inspection::check_whole( + &guards.tools, + &guards.redaction, + &state.audit, + &caller, + "cache", + &restored, + ) + }) + .await + { return Err(ctx.emit(e).into()); } // The tokens the stored answer records, weighted for this model, @@ -1110,12 +1159,13 @@ async fn run( // client gets the framing it asked for. let body = String::from_utf8_lossy(&restored).into_owned(); let events = async_stream::stream! { + let _memory = cached.memory; yield Ok::<_, Infallible>(axum::response::sse::Event::default().data(body)); yield Ok::<_, Infallible>(axum::response::sse::Event::default().data("[DONE]")); }; axum::response::sse::Sse::new(events).into_response() } else { - json_response(restored) + json_response(resources::retain(restored.into(), cached.memory)) }; response .headers_mut() @@ -1213,6 +1263,7 @@ async fn run( state: state.clone(), guards: guards.clone(), request: crate::lifecycle::ChatRequestSnapshot { + memory: resources::stream_budget(), identity: identity.clone(), trace_id: metadata.request_id.clone(), session_id: session_id.clone(), @@ -1361,28 +1412,39 @@ async fn run( // The answer as the caller will receive it: their values restored, // then its text through the content rules that check answers. - let restored = crate::redaction::restore_body(&answer_ledger, &completed.body); - let restored = crate::content_filter::whole( - &guards.content, - &state.audit, - &caller, - &entry.provider_name, - surface.dialect, - restored, - ); + let response_memory = completed.memory.clone(); + let restored = resources::cpu(completed.body.len(), || { + crate::redaction::restore_body(&answer_ledger, &completed.body) + }) + .await; + let restored = resources::cpu(restored.len(), || { + crate::content_filter::whole( + &guards.content, + &state.audit, + &caller, + &entry.provider_name, + surface.dialect, + restored, + ) + }) + .await; // Tool calls, on the whole answer before any of it has gone out — in // the form the caller would run them. A refusal is the gateway's // policy, not the upstream failing, so the route's health counts it as // a success; the answer is neither cached nor billed. - if let Some(e) = crate::tool_inspection::check_whole( - &guards.tools, - &guards.redaction, - &state.audit, - &caller, - &entry.provider_name, - &restored, - ) { + if let Some(e) = resources::cpu(restored.len(), || { + crate::tool_inspection::check_whole( + &guards.tools, + &guards.redaction, + &state.audit, + &caller, + &entry.provider_name, + &restored, + ) + }) + .await + { finalize_health(&state, &sel_record, true).await; return Err(ctx.emit(e).into()); } @@ -1408,7 +1470,10 @@ async fn run( "Audit log: request completed" ); - let mut response = json_response(restored); + let mut response = json_response(resources::retain( + restored.into(), + response_memory.unwrap_or_default(), + )); response .headers_mut() .insert("X-Cache", HeaderValue::from_static("MISS")); @@ -1475,13 +1540,13 @@ fn dialect_headers(headers: &HeaderMap) -> Vec<(String, String)> { .collect() } -fn json_response(body: Vec) -> axum::response::Response { +fn json_response(body: impl Into) -> axum::response::Response { ( [( header::CONTENT_TYPE, HeaderValue::from_static("application/json"), )], - body, + body.into(), ) .into_response() } diff --git a/crates/gateway/src/proxy/transport.rs b/crates/gateway/src/proxy/transport.rs index 84431279..8d301b4f 100644 --- a/crates/gateway/src/proxy/transport.rs +++ b/crates/gateway/src/proxy/transport.rs @@ -304,7 +304,10 @@ async fn check_status( } if !status.is_success() { let retry_after = retry_after(resp.headers()); - let body = resp.text().await.unwrap_or_default(); + let body = think_watch_common::resources::read_response(resp, 64 * 1024) + .await + .map(|bytes| String::from_utf8_lossy(&bytes).into_owned()) + .unwrap_or_else(|error| error.to_string()); tracing::warn!(provider = label, status = %status, body = %body, "upstream returned non-2xx"); const CLIENT_MAX: usize = 512; let shown = if body.len() > CLIENT_MAX { diff --git a/crates/mcp-gateway/src/cache.rs b/crates/mcp-gateway/src/cache.rs index a34ad8aa..50bb818d 100644 --- a/crates/mcp-gateway/src/cache.rs +++ b/crates/mcp-gateway/src/cache.rs @@ -105,16 +105,26 @@ impl McpResponseCache { ) -> Option { let scope_label = scope_metric_label(caller); let key = Self::cache_key(server_id, caller, request); - let cached: Option = self.redis.get(&key).await.ok().flatten(); + let cached = think_watch_common::resources::cached_body(&self.redis, &key).await; let hit = cached.is_some(); - let parsed = cached.and_then(|json| { - serde_json::from_str::(&json) - .map_err(|e| { - tracing::warn!("Failed to deserialize cached MCP response: {e}"); - e - }) - .ok() - }); + let parsed = if let Some(json) = cached { + let memory = json.clone(); + think_watch_common::resources::cpu(json.len(), || { + serde_json::from_slice::(&json) + .map_err(|e| { + tracing::warn!("Failed to deserialize cached MCP response: {e}"); + e + }) + .ok() + .map(|mut response| { + response.memory = Some(memory); + response + }) + }) + .await + } else { + None + }; // Three counters per lookup, each labelled by scope so the // dashboard can break hit-rate down by global vs per-caller: // * mcp_cache_total — denominator diff --git a/crates/mcp-gateway/src/lifecycle/mod.rs b/crates/mcp-gateway/src/lifecycle/mod.rs index 095f343f..435f0e7e 100644 --- a/crates/mcp-gateway/src/lifecycle/mod.rs +++ b/crates/mcp-gateway/src/lifecycle/mod.rs @@ -150,6 +150,10 @@ impl Surface for McpSurface { // shared response-code classifier so a single user's bad // INVALID_PARAMS doesn't trip the breaker for everyone. match &invoked.view { + CapturedView::Streaming { + outcome: StreamOutcome::UpstreamError { error_type, .. }, + .. + } if error_type == "gateway_memory" => {} CapturedView::Streaming { outcome: StreamOutcome::UpstreamError { .. }, .. @@ -240,6 +244,8 @@ impl Surface for McpSurface { /// shaped); the rest are surface-specific identifiers the pipeline /// is generic over, so they cannot live on `Invoked` itself. pub struct McpPostInvokeDeps { + /// Retains stream capture admission until cache/audit hooks finish. + pub stream_memory: think_watch_common::resources::SharedStreamBudget, pub proxy: McpProxy, pub server_id: Uuid, pub server_name: String, @@ -278,6 +284,19 @@ fn response_for_hooks(invoked: &Invoked, deps: &McpPostInvokeDeps) - match &invoked.view { CapturedView::Buffered(r) => r.clone(), CapturedView::Streaming { outcome, captured } => { + if let StreamOutcome::UpstreamError { + error_type, + message, + .. + } = outcome + && error_type == "gateway_memory" + { + return err_response( + deps.original_request_id.clone(), + crate::proxy::RESOURCE_BUSY, + message.clone(), + ); + } pick_response_envelope(captured, deps.original_request_id.as_ref()).unwrap_or_else( || { let msg = match outcome { diff --git a/crates/mcp-gateway/src/pool.rs b/crates/mcp-gateway/src/pool.rs index 9c5b74e2..d95892e4 100644 --- a/crates/mcp-gateway/src/pool.rs +++ b/crates/mcp-gateway/src/pool.rs @@ -2,6 +2,7 @@ use std::collections::HashMap; use std::sync::Arc; use reqwest::Client; +use think_watch_common::resources; use tokio::sync::RwLock; use uuid::Uuid; @@ -55,6 +56,9 @@ pub enum PoolError { #[error("Failed to parse upstream JSON-RPC response: {0}")] ParseError(String), + + #[error(transparent)] + Buffering(#[from] resources::ReadError), } /// Manages a pool of `McpConnection`s keyed by server ID. @@ -219,7 +223,10 @@ impl ConnectionPool { let status = resp.status(); if !status.is_success() { - let body = resp.text().await.unwrap_or_default(); + let body = resources::read_response(resp, 64 * 1024) + .await + .map(|bytes| String::from_utf8_lossy(&bytes).into_owned()) + .unwrap_or_else(|error| error.to_string()); return Err(PoolError::UpstreamError { status: status.as_u16(), body, @@ -237,28 +244,30 @@ impl ConnectionPool { .unwrap_or("") .to_lowercase(); - let (json_resp, stream_audit_body): (JsonRpcResponse, Option) = + let data = resources::read_response(resp, resources::policy().max_body_bytes).await?; + let (mut json_resp, stream_audit_body): (JsonRpcResponse, Option) = if content_type.contains("text/event-stream") { - let text = resp - .text() - .await - .map_err(|e| PoolError::ParseError(e.to_string()))?; + let text = String::from_utf8_lossy(&data); // Pick the response envelope by matching id against the // request id — without this, a tool that emits // `notifications/progress` before the final response // had its first progress notification returned AS the // response (notifications have id=None which deserializes // into JsonRpcResponse just fine, masking the real reply). - let (matched, full_audit) = parse_sse_json_rpc(&text, request.id.as_ref())?; + let (matched, full_audit) = resources::cpu(data.len(), || { + parse_sse_json_rpc(&text, request.id.as_ref()) + }) + .await?; (matched, Some(full_audit)) } else { - let envelope: JsonRpcResponse = resp - .json() - .await - .map_err(|e| PoolError::ParseError(e.to_string()))?; + let envelope: JsonRpcResponse = + resources::cpu(data.len(), || serde_json::from_slice(&data)) + .await + .map_err(|e| PoolError::ParseError(e.to_string()))?; (envelope, None) }; + json_resp.memory = Some(data); // Validate JSON-RPC version if json_resp.jsonrpc != "2.0" { return Err(PoolError::ParseError("Invalid JSON-RPC version".into())); @@ -309,7 +318,10 @@ impl ConnectionPool { // Same error-surfacing contract as send_request — buffer // the small error body so the caller can see the upstream // diagnostic instead of an opaque "non-2xx". - let body = resp.text().await.unwrap_or_default(); + let body = resources::read_response(resp, 64 * 1024) + .await + .map(|bytes| String::from_utf8_lossy(&bytes).into_owned()) + .unwrap_or_else(|error| error.to_string()); return Err(PoolError::UpstreamError { status: status.as_u16(), body, @@ -477,6 +489,27 @@ fn parse_sse_json_rpc( Ok((json_resp, audit_body)) } +pub(crate) async fn read_json_response( + response: reqwest::Response, +) -> Result { + let bytes = resources::read_response(response, resources::policy().max_body_bytes).await?; + let mut result: JsonRpcResponse = + resources::cpu(bytes.len(), || serde_json::from_slice(&bytes)) + .await + .map_err(|e| PoolError::ParseError(e.to_string()))?; + result.memory = Some(bytes); + Ok(result) +} + +pub(crate) fn error_response(id: Option, error: &PoolError) -> JsonRpcResponse { + let code = if matches!(error, PoolError::Buffering(resources::ReadError::Busy)) { + crate::proxy::RESOURCE_BUSY + } else { + crate::proxy::INTERNAL_ERROR + }; + crate::proxy::err_response(id, code, format!("Upstream server error: {error}")) +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/mcp-gateway/src/proxy.rs b/crates/mcp-gateway/src/proxy.rs index 762c358f..1210d6ad 100644 --- a/crates/mcp-gateway/src/proxy.rs +++ b/crates/mcp-gateway/src/proxy.rs @@ -23,7 +23,7 @@ mod streaming; use jsonrpc::ok_response; pub use jsonrpc::{ INTERNAL_ERROR, INVALID_PARAMS, INVALID_REQUEST, JsonRpcError, JsonRpcRequest, JsonRpcResponse, - METHOD_NOT_FOUND, NEEDS_USER_CREDENTIALS, err_response, + METHOD_NOT_FOUND, NEEDS_USER_CREDENTIALS, RESOURCE_BUSY, err_response, }; use streaming::build_replay_payload; pub(crate) use streaming::pick_response_envelope; @@ -794,6 +794,7 @@ impl McpProxy { ), }; return HandleOutcome::Buffered(JsonRpcResponse { + memory: None, jsonrpc: "2.0".to_owned(), id: request.id, result: None, @@ -824,6 +825,7 @@ impl McpProxy { CredentialOwner::AdminShared => "/admin/mcp/servers", }; return HandleOutcome::Buffered(JsonRpcResponse { + memory: None, jsonrpc: "2.0".to_owned(), id: request.id, result: None, @@ -882,6 +884,7 @@ impl McpProxy { // to `run_post_invoke` per branch. let logged_arguments = params.get("arguments").cloned(); let mut deps = crate::lifecycle::McpPostInvokeDeps { + stream_memory: think_watch_common::resources::stream_budget(), proxy: self.clone(), server_id, server_name: server_name.clone(), @@ -958,6 +961,8 @@ impl McpProxy { self.build_mcp_pump( resp, crate::proxy::streaming::PumpContext { + memory: deps.stream_memory.clone(), + request_id: request.id.clone(), identity: auth_identity, trace_id: auth_trace_id, started_at: auth_started_at, @@ -968,13 +973,9 @@ impl McpProxy { ) } else { // Upstream chose application/json — buffer it. - let response = match resp.json::().await { + let response = match crate::pool::read_json_response(resp).await { Ok(r) => r, - Err(e) => err_response( - request.id.clone(), - INTERNAL_ERROR, - format!("Upstream server error: parse failed: {e}"), - ), + Err(e) => crate::pool::error_response(request.id.clone(), &e), }; Invocation::Buffered(Invoked { identity: auth_identity, @@ -993,11 +994,7 @@ impl McpProxy { error = %e, "streaming upstream tools/call failed" ); - let response = err_response( - request.id.clone(), - INTERNAL_ERROR, - format!("Upstream server error: {e}"), - ); + let response = crate::pool::error_response(request.id.clone(), &e); Invocation::Buffered(Invoked { identity: auth_identity, trace_id: auth_trace_id, @@ -1049,11 +1046,7 @@ impl McpProxy { error = %e, "upstream tools/call failed" ); - let response = err_response( - request.id.clone(), - INTERNAL_ERROR, - format!("Upstream server error: {e}"), - ); + let response = crate::pool::error_response(request.id.clone(), &e); Invocation::Buffered(Invoked { identity: auth_identity, trace_id: auth_trace_id, @@ -1077,7 +1070,12 @@ impl McpProxy { let response = emitted .response .expect("Invocation::Buffered always yields Emitted.response"); - if ctx.wants_streaming { + if ctx.wants_streaming + && !response + .error + .as_ref() + .is_some_and(|e| e.code == RESOURCE_BUSY) + { // Client signalled SSE but the upstream replied // buffered. Wrap the response (and any captured // upstream timeline) as a sequence of SSE events. @@ -1120,6 +1118,13 @@ impl McpProxy { server_name: &str, response: &JsonRpcResponse, ) { + if response + .error + .as_ref() + .is_some_and(|error| error.code == RESOURCE_BUSY) + { + return; + } let is_server_failure = response .error .as_ref() diff --git a/crates/mcp-gateway/src/proxy/jsonrpc.rs b/crates/mcp-gateway/src/proxy/jsonrpc.rs index be117652..b9da25f0 100644 --- a/crates/mcp-gateway/src/proxy/jsonrpc.rs +++ b/crates/mcp-gateway/src/proxy/jsonrpc.rs @@ -15,6 +15,9 @@ pub struct JsonRpcRequest { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct JsonRpcResponse { + /// Internal ownership only; excluded from the JSON-RPC wire format. + #[serde(skip)] + pub memory: Option, pub jsonrpc: String, pub id: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -36,6 +39,7 @@ pub const INVALID_REQUEST: i32 = -32600; pub const METHOD_NOT_FOUND: i32 = -32601; pub const INVALID_PARAMS: i32 = -32602; pub const INTERNAL_ERROR: i32 = -32603; +pub const RESOURCE_BUSY: i32 = -32005; /// ThinkWatch-specific: the calling user (or the calling API key's /// chosen account label) has no upstream credential for the requested /// MCP server. The `data` field carries `server_id` and an optional @@ -48,6 +52,7 @@ pub(crate) fn ok_response( result: serde_json::Value, ) -> JsonRpcResponse { JsonRpcResponse { + memory: None, jsonrpc: "2.0".to_owned(), id, result: Some(result), @@ -61,6 +66,7 @@ pub fn err_response( message: impl Into, ) -> JsonRpcResponse { JsonRpcResponse { + memory: None, jsonrpc: "2.0".to_owned(), id, result: None, diff --git a/crates/mcp-gateway/src/proxy/streaming.rs b/crates/mcp-gateway/src/proxy/streaming.rs index d69834c9..2312e252 100644 --- a/crates/mcp-gateway/src/proxy/streaming.rs +++ b/crates/mcp-gateway/src/proxy/streaming.rs @@ -108,6 +108,8 @@ pub(crate) fn pick_response_envelope( /// /// [`run_post_invoke`]: think_watch_common::lifecycle::stages::run_post_invoke pub(crate) struct PumpContext { + pub memory: think_watch_common::resources::SharedStreamBudget, + pub request_id: Option, pub identity: crate::lifecycle::McpIdentity, pub trace_id: String, pub started_at: std::time::Instant, @@ -151,7 +153,13 @@ impl McpProxy { let bytes_source = upstream_resp .bytes_stream() .map(|r| r.map_err(|e| e.to_string())); - let body = build_passthrough_body(bytes_source, events_buf, done_tx); + let body = build_passthrough_body( + bytes_source, + events_buf, + done_tx, + ctx.memory.clone(), + ctx.request_id, + ); let payload = StreamingPayload { body, @@ -205,6 +213,8 @@ pub(crate) fn build_passthrough_body( source: S, events_buf: std::sync::Arc>>, done_tx: tokio::sync::oneshot::Sender, + memory: think_watch_common::resources::SharedStreamBudget, + request_id: Option, ) -> std::pin::Pin< Box< dyn futures::stream::Stream< @@ -232,6 +242,23 @@ where while let Some(chunk) = source.next().await { match chunk { Ok(bytes) => { + let admitted = memory.lock().unwrap_or_else(|e| e.into_inner()).append(bytes.len()); + if !admitted { + metrics::counter!("gateway_memory_busy_total").increment(1); + metrics::counter!("mcp_memory_busy_total").increment(1); + if let Some(tx) = done_tx.take() { + let _ = tx.send(StreamOutcome::UpstreamError { + error_type: "gateway_memory".into(), + message: "Gateway buffering capacity is occupied; retry shortly".into(), + status_code: 429, + }); + } + let error = crate::proxy::err_response(request_id.clone(), crate::proxy::RESOURCE_BUSY, + "Gateway buffering capacity is occupied; retry shortly"); + yield Ok::(Event::default().data(serde_json::to_string(&error).unwrap())); + text_buf.clear(); + break; + } // Reject the chunk BEFORE the allocation when its // post-append size would push us past the cap. // Tearing down the stream is preferable to letting @@ -264,7 +291,7 @@ where // Best-effort JSON parse so non-JSON // `data:` payloads (rare) still land in // the timeline verbatim. - let parsed = serde_json::from_str::(&payload) + let parsed = think_watch_common::resources::cpu(payload.len(), || serde_json::from_str::(&payload)).await .unwrap_or_else(|_| { serde_json::Value::String(payload.clone()) }); @@ -303,7 +330,7 @@ where && let Some(payload) = extract_sse_data_payload(trailing) && !payload.is_empty() { - let parsed = serde_json::from_str::(&payload) + let parsed = think_watch_common::resources::cpu(payload.len(), || serde_json::from_str::(&payload)).await .unwrap_or_else(|_| serde_json::Value::String(payload.clone())); if let Ok(mut g) = events_buf.lock() { g.push(parsed); @@ -352,11 +379,12 @@ pub(crate) fn build_replay_payload( { events.push(s); } - let stream = futures::stream::iter( - events - .into_iter() - .map(|s| Ok::<_, Infallible>(Event::default().data(s))), - ); + let stream = async_stream::stream! { + let _memory = response.memory; + for event in events { + yield Ok::<_, Infallible>(Event::default().data(event)); + } + }; StreamingPayload { body: Box::pin(stream), new_session_id: None, @@ -368,8 +396,24 @@ mod replay_tests { use super::*; use futures::StreamExt; + #[tokio::test] + async fn replay_retains_buffer_admission_until_the_slow_consumer_drops() { + let budget = think_watch_common::resources::ByteBudget::new(100); + let mut lease = budget.reserve(); + assert!(lease.grow(100)); + let mut response = make_response(1, serde_json::json!({"ok":true})); + response.memory = Some(lease.wrap(bytes::Bytes::from_static(b"upstream"))); + let mut payload = build_replay_payload(response, None); + assert_eq!(budget.used(), 100); + let _ = payload.body.next().await.unwrap().unwrap(); + assert_eq!(budget.used(), 100); + drop(payload); + assert_eq!(budget.used(), 0); + } + fn make_response(id: u64, result: serde_json::Value) -> JsonRpcResponse { JsonRpcResponse { + memory: None, jsonrpc: "2.0".to_owned(), id: Some(serde_json::json!(id)), result: Some(result), @@ -537,13 +581,54 @@ mod passthrough_tests { use futures::stream::{self, StreamExt}; use std::sync::{Arc, Mutex}; + #[tokio::test] + async fn small_events_share_admission_and_a_detached_tail_keeps_it_after_cancel() { + use think_watch_common::resources::{BODY_HEADROOM, ByteBudget, StreamBudget}; + let chunk = Bytes::from_static(b"data: {\"result\":1}\n\n"); + let budget = ByteBudget::new(chunk.len() * BODY_HEADROOM); + let memory = Arc::new(Mutex::new(StreamBudget::new(&budget, 1024))); + let tail_memory = memory.clone(); + let events = Arc::new(Mutex::new(Vec::new())); + let (tx, rx) = tokio::sync::oneshot::channel(); + let mut body = build_passthrough_body( + stream::iter([Ok(chunk.clone()), Ok(chunk.clone())]), + events.clone(), + tx, + memory, + Some(serde_json::json!(42)), + ); + let _ = body.next().await.unwrap().unwrap(); + let busy = body.next().await.unwrap().unwrap(); + let busy = format!("{busy:?}"); + assert!(busy.contains("-32005"), "{busy}"); + assert!( + busy.contains("42"), + "the busy response must identify the request: {busy}" + ); + assert!( + matches!(rx.await.unwrap(), StreamOutcome::UpstreamError { error_type, status_code: 429, .. } if error_type == "gateway_memory") + ); + assert_eq!(events.lock().unwrap().len(), 1); + drop(body); + assert_eq!(budget.used(), chunk.len() * BODY_HEADROOM); + drop(events); + drop(tail_memory); + assert_eq!(budget.used(), 0); + } + async fn run_body( chunks: Vec>, ) -> (usize, Vec, Option<&'static str>) { let events_buf: Arc>> = Arc::new(Mutex::new(Vec::new())); let (done_tx, done_rx) = tokio::sync::oneshot::channel::(); let source = stream::iter(chunks); - let body = build_passthrough_body(source, events_buf.clone(), done_tx); + let body = build_passthrough_body( + source, + events_buf.clone(), + done_tx, + think_watch_common::resources::stream_budget(), + None, + ); let mut count = 0usize; let mut body = body; while let Some(_ev) = body.next().await { @@ -629,7 +714,13 @@ mod passthrough_tests { "data: {\"jsonrpc\":\"2.0\",\"method\":\"x\"}\n\n", ))]) .chain(stream::pending()); - let body = build_passthrough_body(source, events_buf.clone(), done_tx); + let body = build_passthrough_body( + source, + events_buf.clone(), + done_tx, + think_watch_common::resources::stream_budget(), + None, + ); let mut body = body; let _ = body.next().await; drop(body); diff --git a/crates/mcp-gateway/src/transport/streamable_http.rs b/crates/mcp-gateway/src/transport/streamable_http.rs index a0f86554..b2455ac7 100644 --- a/crates/mcp-gateway/src/transport/streamable_http.rs +++ b/crates/mcp-gateway/src/transport/streamable_http.rs @@ -108,7 +108,7 @@ pub async fn handle_post( State(state): State>, headers: HeaderMap, axum::Extension(identity): axum::Extension, - Json(request): Json, + think_watch_common::resources::http::ScheduledJson(request): think_watch_common::resources::http::ScheduledJson, ) -> impl IntoResponse { // --- Session management ------------------------------------------------ let sessions = &state.proxy.sessions; @@ -162,6 +162,16 @@ pub async fn handle_post( match outcome { HandleOutcome::Buffered(response) => { + let memory = response.memory.clone().unwrap_or_default(); + if response + .error + .as_ref() + .is_some_and(|e| e.code == crate::proxy::RESOURCE_BUSY) + { + resp_headers.insert(axum::http::header::RETRY_AFTER, "1".parse().unwrap()); + return (StatusCode::TOO_MANY_REQUESTS, resp_headers, Json(response)) + .into_response(); + } // Buffered shape happens for initialize / tools/list, for // tools/call cache hits / early rejections, AND for // tools/call when the upstream replied with plain @@ -172,13 +182,25 @@ pub async fn handle_post( // plain JSON-RPC clients. if sse_capable { let payload = serde_json::to_string(&response).unwrap_or_default(); - let event = Ok::<_, std::convert::Infallible>(Event::default().data(payload)); - let stream = futures::stream::iter(std::iter::once(event)); + let stream = async_stream::stream! { + let _memory = memory; + yield Ok::<_, std::convert::Infallible>(Event::default().data(payload)); + }; let mut sse_resp = Sse::new(stream).into_response(); sse_resp.headers_mut().extend(resp_headers); sse_resp } else { - (StatusCode::OK, resp_headers, Json(response)).into_response() + let bytes = serde_json::to_vec(&response).unwrap_or_default(); + resp_headers.insert( + axum::http::header::CONTENT_TYPE, + "application/json".parse().unwrap(), + ); + ( + StatusCode::OK, + resp_headers, + think_watch_common::resources::retain(bytes.into(), memory), + ) + .into_response() } } HandleOutcome::Streaming(payload) => { diff --git a/crates/mcp-gateway/src/user_token.rs b/crates/mcp-gateway/src/user_token.rs index 24e1df1e..c92f008a 100644 --- a/crates/mcp-gateway/src/user_token.rs +++ b/crates/mcp-gateway/src/user_token.rs @@ -825,7 +825,10 @@ impl UserTokenResolver { }; let status = resp.status(); - let text = resp.text().await.unwrap_or_default(); + let bytes = think_watch_common::resources::read_response(resp, 64 * 1024) + .await + .map_err(|e| OAuthRefreshFailure::Transient(format!("token response read: {e}")))?; + let text = String::from_utf8_lossy(&bytes); if !status.is_success() { // Only treat actually-irrecoverable grant rejections as // Permanent — the original "any 4xx → Permanent" rule diff --git a/crates/server/Cargo.toml b/crates/server/Cargo.toml index c14e4140..49796768 100644 --- a/crates/server/Cargo.toml +++ b/crates/server/Cargo.toml @@ -11,6 +11,8 @@ name = "think-watch-server" path = "src/main.rs" [dependencies] +mimalloc = { version = "0.1.52", features = ["no_thp"] } +libc = "0.2" tw-breaker = { workspace = true } think-watch-common = { workspace = true } think-watch-auth = { workspace = true } diff --git a/crates/server/src/app.rs b/crates/server/src/app.rs index e4e551dc..91349b97 100644 --- a/crates/server/src/app.rs +++ b/crates/server/src/app.rs @@ -440,6 +440,9 @@ pub async fn create_gateway_app(_config: &AppConfig, state: AppState) -> anyhow: .merge(ai_routes) .merge(key_routes) .merge(mcp_routes) + .layer(axum::middleware::from_fn( + think_watch_common::resources::http::admit, + )) .layer(RequestBodyLimitLayer::new(10 * 1024 * 1024)) // 10MB for large prompts // When the timeout below starts counting: a whole answer's retries // stop in time for its error to be written (see `gateway_proxy::retry`). @@ -1186,6 +1189,7 @@ pub fn create_console_app(config: &AppConfig, state: AppState) -> anyhow::Result // upgrade path bypasses tower-http body transforms); permessage- // deflate negotiation belongs at the reverse-proxy layer. .layer(CompressionLayer::new()) + .layer(axum::middleware::from_fn(think_watch_common::resources::http::admit)) .layer(RequestBodyLimitLayer::new(1024 * 1024)) // 1MB for console API .layer(TimeoutLayer::with_status_code( axum::http::StatusCode::REQUEST_TIMEOUT, diff --git a/crates/server/src/main.rs b/crates/server/src/main.rs index 864b2658..2d66e50c 100644 --- a/crates/server/src/main.rs +++ b/crates/server/src/main.rs @@ -5,8 +5,40 @@ use think_watch_common::db; use think_watch_server::{app, handlers, init, tracing_ch}; use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt}; -#[tokio::main] -async fn main() -> anyhow::Result<()> { +#[global_allocator] +static ALLOCATOR: mimalloc::MiMalloc = mimalloc::MiMalloc; + +fn main() -> anyhow::Result<()> { + #[cfg(target_os = "linux")] + { + // A process-local policy, inherited by subsequently created threads. + // Sparse small-object heaps should not pin entire 2 MiB huge pages. + // SAFETY: PR_SET_THP_DISABLE only consumes integer arguments. + unsafe { + libc::prctl( + libc::PR_SET_THP_DISABLE, + 1 as libc::c_ulong, + 0 as libc::c_ulong, + 0 as libc::c_ulong, + 0 as libc::c_ulong, + ) + }; + } + let policy = think_watch_common::resources::policy(); + let workers = match std::env::var("TOKIO_WORKER_THREADS") { + Ok(value) => value.parse::()?.get(), + Err(std::env::VarError::NotPresent) => policy.io_workers, + Err(error) => return Err(error.into()), + }; + tokio::runtime::Builder::new_multi_thread() + .worker_threads(workers) + .thread_keep_alive(think_watch_common::resources::IDLE_TIMEOUT) + .enable_all() + .build()? + .block_on(serve()) +} + +async fn serve() -> anyhow::Result<()> { // Container healthcheck. The runtime image is distroless/static — // no shell, no curl — so the only executable available to Docker's // healthcheck is this binary. `think-watch-server healthcheck` @@ -29,6 +61,8 @@ async fn main() -> anyhow::Result<()> { .with(ch_layer) .init(); + tracing::info!(policy = ?think_watch_common::resources::policy(), allocator = "mimalloc", "Automatic resource admission enabled"); + let config = AppConfig::from_env()?; if let Err(e) = config.validate() { tracing::error!("Configuration validation failed: {e}"); diff --git a/crates/test-support/tests/resources.rs b/crates/test-support/tests/resources.rs new file mode 100644 index 00000000..9fe69211 --- /dev/null +++ b/crates/test-support/tests/resources.rs @@ -0,0 +1,73 @@ +//! Real Redis decoding and HTTP admission/recovery through the assembled app. + +use fred::prelude::*; +use think_watch_common::resources; +use think_watch_test_support::prelude::*; + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn redis_body_reads_preserve_bytes_and_reservations_until_the_last_clone() { + let app = TestApp::spawn().await; + let key = unique_name("resource-body"); + let data = b"\0\xff\"cached JSON\"\xe4\xb8\xad"; + let _: () = app + .state + .redis + .set(&key, data.as_slice(), None, None, false) + .await + .unwrap(); + let budget = resources::body_budget(); + let baseline = budget.used(); + let body = resources::cached_body(&app.state.redis, &key) + .await + .unwrap(); + assert_eq!(&body[..], data); + assert_eq!( + budget.used(), + baseline + data.len() * resources::BODY_HEADROOM + ); + let copy = body.clone(); + drop(body); + assert!(budget.used() > baseline); + drop(copy); + assert_eq!(budget.used(), baseline); + + let mut occupied = budget.reserve(); + assert!(occupied.grow(resources::policy().body_bytes - baseline)); + assert!( + resources::cached_body(&app.state.redis, &key) + .await + .is_none() + ); + drop(occupied); + assert!( + resources::cached_body(&app.state.redis, &key) + .await + .is_some() + ); + let _: () = app.state.redis.del(&key).await.unwrap(); +} + +#[ignore = "integration test — run via `make test-it`"] +#[tokio::test] +async fn http_body_pressure_is_retryable_and_releases_without_affecting_health() { + let app = TestApp::spawn().await; + let budget = resources::body_budget(); + let mut occupied = budget.reserve(); + assert!(occupied.grow(resources::policy().body_bytes - budget.used())); + let gateway = app.gateway_client(); + let response = gateway + .post("/v1/chat/completions", json!({"model":"test"})) + .await + .unwrap(); + response.assert_status(429); + assert_eq!(response.headers["retry-after"], "1"); + drop(occupied); + let response = gateway + .post("/v1/chat/completions", json!({"model":"test"})) + .await + .unwrap(); + response.assert_status(401); + let health = gateway.get("/health/live").await.unwrap(); + health.assert_ok(); +} diff --git a/docs/operations/resources.md b/docs/operations/resources.md new file mode 100644 index 00000000..2646bcb9 --- /dev/null +++ b/docs/operations/resources.md @@ -0,0 +1,64 @@ +# Automatic process resources + +The server sizes its execution and buffering policy at startup. No configuration +change is required. CPU availability includes the process affinity and container +CPU quota. Memory uses the smaller of physical RAM and the cgroup v1/v2 memory +limit, including ancestor limits. + +## Execution + +The I/O runtime keeps at most eight workers. Large or concurrent JSON decoding, +protocol conversion, redaction and content inspection use elastic blocking +execution. CPU admission is shared with password hashing: at most one job per +available CPU, further limited to one job per 128 MiB of memory allowance. +An 80 CPU server with sufficient memory can execute 80 jobs concurrently. +Extra execution threads retire after ten idle seconds. + +Argon2 also has its own memory admission: one concurrent hash per 32 MiB within +one eighth of the memory allowance, with a minimum of one. Cancelling an HTTP +request does not release a hash's admission until that hash actually finishes. +The existing `TOKIO_WORKER_THREADS` override remains supported. + +The server uses mimalloc with transparent huge pages disabled. On Linux it also +applies a process-local THP policy before creating runtime threads. This needs +no host sysctl or environment changes. + +## Buffered bodies and streams + +Requests, upstream AI/MCP responses, Redis cache reads and retained streaming +timelines share one byte budget: one quarter of the memory allowance, between +8 MiB and 8 GiB. Reservations allow 32 times the wire size for decoded JSON, +conversion, inspection and capture copies. A single response or stream timeline +is limited to the smaller of 64 MiB and one thirty-second of that budget. +The existing gateway 10 MiB and console 1 MiB request limits still apply. +Error bodies and credential endpoint responses have a separate per-read 64 KiB +limit and use the same shared budget. + +Admission is incremental, including chunked responses. Readers reject growth +before appending rather than waiting while holding partial bodies. Reservations +follow shared buffers through transformations, slow clients, cancellation and +detached stream audit work. Cache reads reserve before fetching and use bounded +Redis reads even if a key changes between the size check and fetch. + +Memory pressure produces HTTP 429 with `Retry-After: 1` before streaming starts. +Already-started streams terminate with an error event (MCP code `-32005`). +Local pressure does not count as an upstream failure for circuit breakers. +Oversized requests return 413; oversized buffered upstream responses fail rather +than allocating without a bound. Cache reads under pressure become misses. + +## Audit backlog + +Queued audit entries and retained ClickHouse retry batches share a separate byte +budget: one sixteenth of the memory allowance, between 2 MiB and 128 MiB. The +estimate includes owned strings, JSON allocations and room for temporary flush +copies. The existing 100,000 entry queue limit is a secondary bound. +When full, the logger drops new entries and emits a throttled error. + +Monitor `audit_log_buffered_bytes`, `audit_log_memory_dropped_total`, +`audit_log_dropped_total` and `gateway_memory_busy_total` alongside the existing +queue-depth metrics. A drop requires operational attention, especially when +ClickHouse is unavailable for an extended period. + +These budgets estimate application allocations; they are not a hard RSS cap. +Connection state, database pools, configuration, stacks and allocator overhead +also consume memory. Deployment memory and CPU limits remain authoritative. From 3ae4062f9632aff4b2def89621430a3dfa52d5c3 Mon Sep 17 00:00:00 2001 From: fylorn <249551762+fylorn@users.noreply.github.com> Date: Tue, 3 Nov 2026 14:01:12 +0800 Subject: [PATCH 2/2] fix: bound background discovery and diagnostic response reads Apply shared response admission to MCP tool catalogs, provider and registry probes, OAuth discovery and credential exchange, audit forwarder errors and S3 reads. Retain decoded catalog admission while persisting results and bound protocol-probe draining time. --- crates/auth/src/oauth/client.rs | 10 +++- crates/common/src/audit/forwarders.rs | 10 +++- crates/common/src/blob_store.rs | 8 ++- crates/common/src/resources.rs | 10 ++++ crates/server/src/handlers/mcp_oauth.rs | 6 ++- .../src/handlers/mcp_oauth/discovery.rs | 45 ++++++++++------- crates/server/src/handlers/mcp_shared.rs | 50 +++++++++---------- crates/server/src/handlers/mcp_store.rs | 10 ++-- crates/server/src/handlers/providers.rs | 14 +++++- crates/server/src/mcp_runtime.rs | 38 +++++++++----- crates/server/src/protocol_probe.rs | 6 ++- docs/operations/resources.md | 4 +- 12 files changed, 142 insertions(+), 69 deletions(-) diff --git a/crates/auth/src/oauth/client.rs b/crates/auth/src/oauth/client.rs index d4f62426..553914c7 100644 --- a/crates/auth/src/oauth/client.rs +++ b/crates/auth/src/oauth/client.rs @@ -98,7 +98,15 @@ pub async fn exchange_authorization_code( .map_err(TokenExchangeError::Unreachable)?; let status = resp.status(); - let resp_text = resp.text().await.map_err(TokenExchangeError::Unreachable)?; + let bytes = think_watch_common::resources::read_response(resp, 64 * 1024) + .await + .map_err(|e| match e { + think_watch_common::resources::ReadError::Transport(e) => { + TokenExchangeError::Unreachable(e) + } + other => TokenExchangeError::InvalidResponse(other.to_string()), + })?; + let resp_text = String::from_utf8_lossy(&bytes); // Two cases land here: // * non-2xx status (always treated as failure) diff --git a/crates/common/src/audit/forwarders.rs b/crates/common/src/audit/forwarders.rs index 90023b4e..31a2dd14 100644 --- a/crates/common/src/audit/forwarders.rs +++ b/crates/common/src/audit/forwarders.rs @@ -202,7 +202,10 @@ pub(super) async fn send_kafka( if resp.status().is_success() { Ok(()) } else { - let body = resp.text().await.unwrap_or_default(); + let bytes = crate::resources::read_response(resp, 64 * 1024) + .await + .map_err(|e| e.to_string())?; + let body = String::from_utf8_lossy(&bytes); Err(format!("Kafka REST proxy returned error: {body}")) } } @@ -300,7 +303,10 @@ pub(super) async fn send_webhook( Ok(resp) if resp.status().is_success() => Ok(()), Ok(resp) => { let status = resp.status(); - let rtext = resp.text().await.unwrap_or_default(); + let bytes = crate::resources::read_response(resp, 64 * 1024) + .await + .map_err(|e| e.to_string())?; + let rtext = String::from_utf8_lossy(&bytes); Err(format!("Webhook returned {status}: {rtext}")) } Err(e) => Err(format!("Webhook request failed: {e}")), diff --git a/crates/common/src/blob_store.rs b/crates/common/src/blob_store.rs index 5211d7a8..8a057190 100644 --- a/crates/common/src/blob_store.rs +++ b/crates/common/src/blob_store.rs @@ -378,8 +378,12 @@ impl S3Store { .await .map_err(|e| BlobError::Transport(e.to_string()))?; let status = resp.status(); - let bytes = resp - .bytes() + let max = if status.is_success() { + crate::resources::policy().max_body_bytes + } else { + 64 * 1024 + }; + let bytes = crate::resources::read_response(resp, max) .await .map_err(|e| BlobError::Transport(e.to_string()))?; if !status.is_success() { diff --git a/crates/common/src/resources.rs b/crates/common/src/resources.rs index 8f69b915..5f8b3987 100644 --- a/crates/common/src/resources.rs +++ b/crates/common/src/resources.rs @@ -252,6 +252,16 @@ pub async fn read_response(response: reqwest::Response, max: usize) -> Result( + response: reqwest::Response, + max: usize, +) -> anyhow::Result<(T, Bytes)> { + let bytes = read_response(response, max).await?; + let value = cpu(bytes.len(), || serde_json::from_slice(&bytes)).await?; + Ok((value, bytes)) +} + async fn read_response_with_budget( mut response: reqwest::Response, max: usize, diff --git a/crates/server/src/handlers/mcp_oauth.rs b/crates/server/src/handlers/mcp_oauth.rs index 88311499..f9297494 100644 --- a/crates/server/src/handlers/mcp_oauth.rs +++ b/crates/server/src/handlers/mcp_oauth.rs @@ -1340,14 +1340,16 @@ pub(crate) async fn resolve_upstream_subject( return Some(s); } let endpoint = userinfo_endpoint?; - let body = match http + let (body, _memory): (serde_json::Value, _) = match http .get(endpoint) .bearer_auth(access_token) .header("Accept", "application/json") .send() .await { - Ok(r) if r.status().is_success() => r.json::().await.ok()?, + Ok(r) if r.status().is_success() => think_watch_common::resources::read_json(r, 64 * 1024) + .await + .ok()?, Ok(r) => { tracing::debug!( endpoint = %endpoint, diff --git a/crates/server/src/handlers/mcp_oauth/discovery.rs b/crates/server/src/handlers/mcp_oauth/discovery.rs index 997a80d5..629bb533 100644 --- a/crates/server/src/handlers/mcp_oauth/discovery.rs +++ b/crates/server/src/handlers/mcp_oauth/discovery.rs @@ -77,13 +77,14 @@ pub async fn oauth_discover( for url in &candidates { match http.get(url).send().await { Ok(resp) if resp.status().is_success() => { - let body: serde_json::Value = match resp.json().await { - Ok(v) => v, - Err(e) => { - last_err = Some(format!("{url}: {e}")); - continue; - } - }; + let (body, _memory): (serde_json::Value, _) = + match think_watch_common::resources::read_json(resp, 64 * 1024).await { + Ok(v) => v, + Err(e) => { + last_err = Some(format!("{url}: {e}")); + continue; + } + }; return Ok(Json(parse_oauth_metadata(&body))); } Ok(resp) => { @@ -455,13 +456,14 @@ async fn fetch_protected_resource( return None; } }; - let body: serde_json::Value = match resp.json().await { - Ok(v) => v, - Err(e) => { - diag.push(format!("{url}: parse: {e}")); - return None; - } - }; + let (body, _memory): (serde_json::Value, _) = + match think_watch_common::resources::read_json(resp, 64 * 1024).await { + Ok(v) => v, + Err(e) => { + diag.push(format!("{url}: parse: {e}")); + return None; + } + }; let issuer = body .get("authorization_servers") .and_then(|v| v.as_array()) @@ -501,8 +503,10 @@ async fn fetch_authz_server_metadata( } match http.get(url).send().await { Ok(resp) if resp.status().is_success() => { - match resp.json::().await { - Ok(body) => { + match think_watch_common::resources::read_json::(resp, 64 * 1024) + .await + { + Ok((body, _memory)) => { diag.push(format!("found authz-server metadata at {url}")); return parse_authz_server_metadata(&body); } @@ -655,7 +659,14 @@ async fn register_dynamic_client( } }; let status = resp.status(); - let text = resp.text().await.unwrap_or_default(); + let bytes = match think_watch_common::resources::read_response(resp, 64 * 1024).await { + Ok(bytes) => bytes, + Err(error) => { + diag.push(format!("dynamic-registration response read: {error}")); + return None; + } + }; + let text = String::from_utf8_lossy(&bytes); if !status.is_success() { diag.push(format!( "dynamic-registration HTTP {status}: {}", diff --git a/crates/server/src/handlers/mcp_shared.rs b/crates/server/src/handlers/mcp_shared.rs index e8875ee0..6e044d21 100644 --- a/crates/server/src/handlers/mcp_shared.rs +++ b/crates/server/src/handlers/mcp_shared.rs @@ -111,35 +111,33 @@ pub async fn probe_mcp_endpoint( .unwrap_or("") .to_lowercase(); - let json: serde_json::Value = if content_type.contains("text/event-stream") { - match resp.text().await { - Ok(text) => crate::mcp_runtime::parse_sse_json(&text).unwrap_or_else(|e| { - tracing::warn!( - endpoint = %endpoint_url, - error = %e, - "MCP probe: failed to parse SSE JSON response" - ); - serde_json::Value::Null - }), - Err(e) => { - tracing::warn!( - endpoint = %endpoint_url, - error = %e, - "MCP probe: failed to read SSE response body" - ); - serde_json::Value::Null - } + let memory = match think_watch_common::resources::read_response( + resp, + think_watch_common::resources::policy().max_body_bytes, + ) + .await + { + Ok(bytes) => bytes, + Err(error) => { + return McpProbeOutcome { + success: false, + requires_auth: false, + message: format!("MCP response could not be read: {error}"), + latency_ms, + tools: vec![], + }; } - } else { - resp.json().await.unwrap_or_else(|e| { - tracing::warn!( - endpoint = %endpoint_url, - error = %e, - "MCP probe: failed to parse JSON response" - ); + }; + let json: serde_json::Value = think_watch_common::resources::cpu(memory.len(), || { + if content_type.contains("text/event-stream") { + crate::mcp_runtime::parse_sse_json(&String::from_utf8_lossy(&memory)) + } else { + serde_json::from_slice(&memory).map_err(Into::into) + }.unwrap_or_else(|e| { + tracing::warn!(endpoint = %endpoint_url, error = %e, "MCP probe: failed to parse response"); serde_json::Value::Null }) - }; + }).await; let result_field = json.get("result"); let tools: Vec = result_field diff --git a/crates/server/src/handlers/mcp_store.rs b/crates/server/src/handlers/mcp_store.rs index b57efbeb..cb79b150 100644 --- a/crates/server/src/handlers/mcp_store.rs +++ b/crates/server/src/handlers/mcp_store.rs @@ -274,10 +274,12 @@ pub async fn sync_registry( ))); } - let registry: RegistryResponse = resp - .json() - .await - .map_err(|e| AppError::BadRequest(format!("Invalid registry JSON: {e}")))?; + let (registry, _memory): (RegistryResponse, _) = think_watch_common::resources::read_json( + resp, + think_watch_common::resources::policy().max_body_bytes, + ) + .await + .map_err(|e| AppError::BadRequest(format!("Invalid registry JSON: {e}")))?; // Wrap the whole sync (upserts + post-loop cleanup DELETE) in a // single transaction. Previously each `INSERT ... ON CONFLICT diff --git a/crates/server/src/handlers/providers.rs b/crates/server/src/handlers/providers.rs index 72ad74e3..38b4dd19 100644 --- a/crates/server/src/handlers/providers.rs +++ b/crates/server/src/handlers/providers.rs @@ -806,7 +806,19 @@ async fn list_models(client: &reqwest::Client, url: &str, headers: &[(String, St }; let status = resp.status(); - let body: serde_json::Value = resp.json().await.unwrap_or(serde_json::Value::Null); + let memory = match think_watch_common::resources::read_response( + resp, + think_watch_common::resources::policy().max_body_bytes, + ) + .await + { + Ok(value) => value, + Err(error) => return Listing::Failed(error.to_string()), + }; + let body: serde_json::Value = think_watch_common::resources::cpu(memory.len(), || { + serde_json::from_slice(&memory).unwrap_or(serde_json::Value::Null) + }) + .await; if !status.is_success() { let message = body .get("error") diff --git a/crates/server/src/mcp_runtime.rs b/crates/server/src/mcp_runtime.rs index 1274c69e..a8084073 100644 --- a/crates/server/src/mcp_runtime.rs +++ b/crates/server/src/mcp_runtime.rs @@ -441,12 +441,19 @@ async fn try_discover_and_persist_tools_with_auth( .unwrap_or("") .to_lowercase(); - let json: serde_json::Value = if content_type.contains("text/event-stream") { - let text = resp.text().await?; - parse_sse_json(&text)? - } else { - resp.json().await? - }; + let memory = think_watch_common::resources::read_response( + resp, + think_watch_common::resources::policy().max_body_bytes, + ) + .await?; + let json: serde_json::Value = think_watch_common::resources::cpu(memory.len(), || { + if content_type.contains("text/event-stream") { + parse_sse_json(&String::from_utf8_lossy(&memory)) + } else { + serde_json::from_slice(&memory).map_err(Into::into) + } + }) + .await?; // Strict: require a `result` field so malformed responses are // caught instead of silently treated as success. let result = json @@ -569,12 +576,19 @@ pub async fn discover_user_tools( .and_then(|v| v.to_str().ok()) .unwrap_or("") .to_lowercase(); - let json: serde_json::Value = if content_type.contains("text/event-stream") { - let text = resp.text().await?; - parse_sse_json(&text)? - } else { - resp.json().await? - }; + let memory = think_watch_common::resources::read_response( + resp, + think_watch_common::resources::policy().max_body_bytes, + ) + .await?; + let json: serde_json::Value = think_watch_common::resources::cpu(memory.len(), || { + if content_type.contains("text/event-stream") { + parse_sse_json(&String::from_utf8_lossy(&memory)) + } else { + serde_json::from_slice(&memory).map_err(Into::into) + } + }) + .await?; let result = json .get("result") .ok_or_else(|| anyhow::anyhow!("tools/list response missing `result` field"))? diff --git a/crates/server/src/protocol_probe.rs b/crates/server/src/protocol_probe.rs index 1d42d3ea..c02d5675 100644 --- a/crates/server/src/protocol_probe.rs +++ b/crates/server/src/protocol_probe.rs @@ -247,7 +247,11 @@ async fn probe_one(materials: &ProviderMaterials, upstream_model: &str) -> Probe match attempt { Ok(Ok(resp)) => { // Drain it so the connection goes back to the pool. - let _ = resp.bytes().await; + let _ = tokio::time::timeout( + PROBE_TIMEOUT, + think_watch_common::resources::read_response(resp, 64 * 1024), + ) + .await; return Probe::Verdict(Verdict::Ok(candidate)); } Ok(Err(GatewayError::UpstreamAuthError { status, message })) => { diff --git a/docs/operations/resources.md b/docs/operations/resources.md index 2646bcb9..57a3e49d 100644 --- a/docs/operations/resources.md +++ b/docs/operations/resources.md @@ -32,7 +32,9 @@ conversion, inspection and capture copies. A single response or stream timeline is limited to the smaller of 64 MiB and one thirty-second of that budget. The existing gateway 10 MiB and console 1 MiB request limits still apply. Error bodies and credential endpoint responses have a separate per-read 64 KiB -limit and use the same shared budget. +limit and use the same shared budget. MCP tool discovery, provider and registry +catalogs, OAuth discovery, audit forwarder errors and S3 reads also use bounded +reads, so administrative/background requests cannot bypass admission on ingress. Admission is incremental, including chunked responses. Readers reject growth before appending rather than waiting while holding partial bodies. Reservations