From 1185c2331c2cfb3a056af0541a0e66c224a88926 Mon Sep 17 00:00:00 2001 From: Jared Lunde Date: Sat, 20 Jun 2026 09:26:25 -0700 Subject: [PATCH] Fix audit findings: CB race, silent billing zeros, config/SSE robustness MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses issues surfaced by the technical audit: - circuit_breaker: fix windowed failure-count race. The old two-atomic design (state + window_start) computed the window-reset decision outside the state CAS, so concurrent failures at a window boundary could each reset the count to 1 and drop one another — pinning the count and keeping a genuinely-broken provider's breaker stuck closed. Anchor the window in the packed timestamp so reset-and-increment commit in a single CAS; delete the now-redundant window_start atomic. Add a regression test. - circuit_breaker: add a record_success fast path so high-throughput success traffic doesn't bounce the breaker cache line on every response. - proxy: emit a warn + ai_usage_parse_errors_total when a managed 2xx response carries no parseable usage, instead of silently emitting a zero-token billing row indistinguishable from a real zero. - config: surface a malformed-TOML parse error at load with the file named, rather than swallowing it and letting it resurface as an opaque Figment error. half_open_permits(1) so the breaker matches the documented "a probe". - state: warn at boot when upstream_verify_cert is disabled (MitM opening valid only for the local test mock). - usage: tolerate CRLF SSE line endings (RFC 8895) — a trailing \r would otherwise fail the JSON parse and silently zero usage. - peek: correct an inaccurate comment about the "stream" pre-filter needle. - docs: keep ARCHITECTURE.md Metrics table in sync. Co-Authored-By: Claude Opus 4.8 (1M context) --- ARCHITECTURE.md | 1 + src/circuit_breaker.rs | 131 ++++++++++++++++++++++++----------------- src/config.rs | 54 +++++++++++++---- src/metrics.rs | 11 ++++ src/peek.rs | 10 ++-- src/proxy.rs | 25 +++++++- src/state.rs | 12 ++++ src/usage.rs | 22 ++++++- 8 files changed, 194 insertions(+), 72 deletions(-) diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 8f9b40d..c05ccb2 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -369,6 +369,7 @@ Prometheus on the default registry, exposed at `/metrics` on `metrics_listen`. | `ai_requests_in_flight` | Gauge | — | All in-flight requests (streaming + non-streaming) | | `ai_deny_set_size` | Gauge | — | Current number of denied tenants | | `ai_nats_connected` | Gauge | — | 1 if NATS watcher is connected, 0 otherwise | +| `ai_usage_parse_errors_total` | Counter | — | Managed 2xx responses with no parseable usage (emitted as a zero-token billing row) | --- diff --git a/src/circuit_breaker.rs b/src/circuit_breaker.rs index af8c1e1..fe70cb5 100644 --- a/src/circuit_breaker.rs +++ b/src/circuit_breaker.rs @@ -159,18 +159,16 @@ impl CircuitBreakerConfig { /// - Bits 62-63: State (0=closed, 1=open, 2=half-open) /// - Bits 48-61: Failure count (14 bits, max 16383) /// - Bits 32-47: Half-open permits remaining (16 bits) -/// - Bits 0-31: Timestamp of last state change (seconds since epoch, wraps every 136 years) -/// -/// For windowed mode, a second atomic tracks the window start timestamp. +/// - Bits 0-31: Timestamp (seconds since epoch, wraps every 136 years). In OPEN it's when the +/// circuit opened (drives the reset timeout); in CLOSED windowed mode it doubles as the current +/// failure window's start. Because the timestamp lives in the same word, a windowed failure can +/// reset-the-window and increment-the-count in a **single** CAS — so concurrent failures at a +/// window boundary can never each independently reset to 1 and drop one another. /// /// This packing ensures all state transitions are atomic via single CAS operations. pub struct CircuitBreaker { /// Packed state word. state: AtomicU64, - /// Window start timestamp (only used in windowed mode). - /// Stores seconds since epoch when the first failure in the current window occurred. - /// 0 means no active window. - window_start: AtomicU64, /// Configuration (immutable after construction). config: CircuitBreakerConfig, /// Clock function for getting current time in seconds. @@ -181,7 +179,6 @@ impl std::fmt::Debug for CircuitBreaker { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("CircuitBreaker") .field("state", &self.state) - .field("window_start", &self.window_start) .field("config", &self.config) .finish_non_exhaustive() } @@ -214,7 +211,6 @@ impl CircuitBreaker { let initial = Self::pack(STATE_CLOSED, 0, 0, clock()); Self { state: AtomicU64::new(initial), - window_start: AtomicU64::new(0), config, clock, } @@ -333,22 +329,25 @@ impl CircuitBreaker { /// In closed state, resets the failure counter (and window for windowed mode). /// In half-open state, closes the circuit (service is healthy again). pub fn record_success(&self) { - // Reset window start for windowed mode - self.window_start.store(0, Ordering::Release); + // Fast path: a healthy CLOSED breaker with no accrued failures is the overwhelmingly common + // case — every successful response calls this. Bail before any write so high-throughput + // success traffic to one provider doesn't bounce its breaker cache line across every worker + // thread on each response. The CLOSED timestamp is only read as a window anchor by the *next* + // failure (which re-anchors a fresh window when failures==0 regardless), so not refreshing it + // on an already-clean breaker changes no observable behavior. + let (state, failures, _, _) = Self::unpack(self.state.load(Ordering::Acquire)); + if state == STATE_CLOSED && failures == 0 { + return; + } loop { let packed = self.state.load(Ordering::Acquire); let (state, _, _, _) = Self::unpack(packed); let new_packed = match state { - STATE_CLOSED => { - // Reset failure count, keep closed - Self::pack(STATE_CLOSED, 0, 0, self.now_secs()) - } - STATE_HALF_OPEN => { - // Success in half-open: close the circuit - Self::pack(STATE_CLOSED, 0, 0, self.now_secs()) - } + // Reset the failure count (and re-anchor the window via the timestamp). Covers both + // CLOSED-with-accrued-failures and HALF_OPEN (a probe succeeded → close the circuit). + STATE_CLOSED | STATE_HALF_OPEN => Self::pack(STATE_CLOSED, 0, 0, self.now_secs()), STATE_OPEN => return, // Shouldn't record success while open _ => return, }; @@ -415,52 +414,40 @@ impl CircuitBreaker { } /// Record failure with windowed failure tracking. + /// + /// The current window's start lives in the packed timestamp (CLOSED state), so the + /// reset-the-window-and-set-to-1 decision and the within-window increment are made from the + /// **same** `packed` value and committed in the **same** CAS. That single-CAS atomicity is what + /// fixes the old two-atomic race: when many failures land at a window boundary, they retry the + /// CAS and each sees the latest count, so they linearize into a correct running total instead of + /// each independently resetting to 1 and dropping the others (which could pin the count at 1 and + /// keep a genuinely-broken provider's breaker stuck closed). fn record_failure_windowed(&self, threshold: u32, window_secs: u64) { let now = self.now_secs(); - // Handle window timing - let window_start = self.window_start.load(Ordering::Acquire); - let (new_window_start, reset_count) = if window_start == 0 { - // First failure, start new window - (now, true) - } else if now.wrapping_sub(window_start) >= window_secs { - // Window expired, start new window - (now, true) - } else { - // Within window, continue counting - (window_start, false) - }; - - // Update window start if needed (best-effort, races are acceptable) - if new_window_start != window_start { - let _ = self.window_start.compare_exchange( - window_start, - new_window_start, - Ordering::Release, - Ordering::Relaxed, - ); - } - - // Now update the main state loop { let packed = self.state.load(Ordering::Acquire); - let (state, failures, _, _) = Self::unpack(packed); + let (state, failures, _, ts) = Self::unpack(packed); let new_packed = match state { STATE_CLOSED => { - let new_failures = if reset_count { 1 } else { failures + 1 }; + // `failures == 0` ⇒ fresh window (cold start, or just after a success reset); + // `now - ts >= window` ⇒ the window this failure falls into has expired. Either + // way this failure starts a new window at count 1; otherwise it accrues into the + // current one. The anchor (window start) is carried in the timestamp field. + let window_expired = failures == 0 || now.wrapping_sub(ts) >= window_secs; + let (new_failures, anchor) = if window_expired { + (1, now) + } else { + (failures + 1, ts) + }; if new_failures >= u64::from(threshold) { - // Reset window when opening circuit - self.window_start.store(0, Ordering::Release); Self::pack(STATE_OPEN, 0, 0, now) } else { - Self::pack(STATE_CLOSED, new_failures, 0, now) + Self::pack(STATE_CLOSED, new_failures, 0, anchor) } } - STATE_HALF_OPEN => { - self.window_start.store(0, Ordering::Release); - Self::pack(STATE_OPEN, 0, 0, now) - } + STATE_HALF_OPEN => Self::pack(STATE_OPEN, 0, 0, now), STATE_OPEN => return, _ => return, }; @@ -496,7 +483,6 @@ impl CircuitBreaker { /// Reset the circuit breaker to closed state. pub fn reset(&self) { - self.window_start.store(0, Ordering::Release); let packed = Self::pack(STATE_CLOSED, 0, 0, self.now_secs()); self.state.store(packed, Ordering::Release); } @@ -850,6 +836,45 @@ mod tests { } } + /// A fixed clock pins every failure into the same window, so this isolates concurrent *counting* + /// from window timing — the exact condition the old two-atomic design dropped increments under. + fn fixed_clock() -> u64 { + 1000 + } + + #[test] + fn test_concurrent_windowed_counts_every_failure() { + // Regression for the window-boundary race: with all 20 failures inside one window and the + // threshold well above 20, every concurrent failure must be counted — none silently dropped. + // The old code computed the window-reset decision outside the state CAS, so losers of the + // window_start race could reset a peer's increment back to 1, pinning the count and keeping a + // genuinely-failing provider's breaker stuck closed. The single-CAS anchor makes that + // impossible: the count must land at exactly 20. + for _ in 0..200 { + let cb = Arc::new(CircuitBreaker::with_clock( + CircuitBreakerConfig::windowed(1000, Duration::from_secs(60)), + fixed_clock, + )); + + let handles: Vec<_> = (0..20) + .map(|_| { + let cb = Arc::clone(&cb); + thread::spawn(move || cb.record_failure()) + }) + .collect(); + + for h in handles { + h.join().unwrap(); + } + + assert_eq!( + cb.state(), + CircuitState::Closed { failure_count: 20 }, + "every concurrent in-window failure must be counted (no lost increments)" + ); + } + } + #[test] fn test_concurrent_windowed_failures() { for _ in 0..50 { diff --git a/src/config.rs b/src/config.rs index c454ad9..704d91b 100644 --- a/src/config.rs +++ b/src/config.rs @@ -274,7 +274,13 @@ impl AiConfig { ) .reset_timeout(std::time::Duration::from_secs( self.circuit_breaker_reset_secs, - )), + )) + // Admit exactly **one** half-open probe (matching ARCHITECTURE.md: "a probe is + // admitted"). The library default is 3, but for an egress breaker a single probe is the + // conservative choice — one request tests a possibly-still-broken provider, and a success + // closes it immediately for full traffic; 3 concurrent probes would send 3× the load at a + // provider we believe is down. Deliberately not a config knob: minimum effective surface. + .half_open_permits(1), ) } @@ -328,16 +334,26 @@ fn known_config_keys() -> std::collections::BTreeSet { fn reject_unknown_toml_keys(path: &Path) -> Result<()> { use figment::Provider as _; let known = known_config_keys(); - let unknown: std::collections::BTreeSet = Toml::file(path) - .data() - .map(|profiles| { - profiles - .into_values() - .flat_map(|dict| dict.into_keys()) - .filter(|k| !known.contains(k)) - .collect() - }) - .unwrap_or_default(); + // `data()` errors two ways we must distinguish: a *missing* file is benign (the gateway runs on + // defaults + env), but a *malformed* file is a hard error we must surface here with the file + // named — otherwise the syntax error only reappears later as an opaque Figment `extract()` + // failure with no path attribution. A missing file yields no keys and passes; any other error + // (parse failure, permission denied) fails the load loudly. + let unknown: std::collections::BTreeSet = match Toml::file(path).data() { + Ok(profiles) => profiles + .into_values() + .flat_map(|dict| dict.into_keys()) + .filter(|k| !known.contains(k)) + .collect(), + Err(e) if path.exists() => { + return Err(GatewayError::Config(format!( + "failed to parse {}: {e}", + path.display() + ))); + } + // No file (or it vanished between checks): nothing to validate — defaults + env apply. + Err(_) => return Ok(()), + }; if unknown.is_empty() { return Ok(()); } @@ -421,6 +437,22 @@ mod tests { } } + #[test] + fn rejects_malformed_toml_with_path_attribution() { + // A syntax error in the operator's TOML must fail the load *here*, naming the file — not + // silently pass this check and resurface later as an opaque Figment extract() error. + let path = temp_toml("malformed", "listen = \"unterminated\nrate_limit_rps = 7\n"); + let err = AiConfig::load_with_path(Some(&path)).unwrap_err(); + let _ = std::fs::remove_file(&path); + match err { + GatewayError::Config(msg) => assert!( + msg.contains("malformed") || msg.contains("parse"), + "error must indicate a parse failure and name the file, got: {msg}" + ), + other => panic!("expected Config error, got {other:?}"), + } + } + #[test] fn accepts_known_toml_keys() { // Every key here is a real `AiConfig` field (including the `[signing_keys]` table) — load diff --git a/src/metrics.rs b/src/metrics.rs index b3ac19e..754df34 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -56,6 +56,11 @@ pub struct Metrics { /// fail-open — it serves on the last-known set when NATS is down — so staleness is otherwise /// silent; this is the metric to alert "deny-set has been stale for >N minutes" on. pub nats_connected: IntGauge, + /// Managed `2xx` responses that carried no parseable usage block. Such a response still emits a + /// zero-token `ai.usage` row, which is indistinguishable downstream from a (non-existent) + /// legitimate zero-token generation — so a provider changing its usage wire shape would silently + /// zero out billing. This counter (paired with a `warn!`) is the alerting surface for that. + pub usage_parse_errors_total: IntCounter, } /// TTFT buckets (seconds). Tuned for LLM latency: sub-second prompts up through the multi-second @@ -133,6 +138,10 @@ impl Metrics { "ai_nats_connected", "Deny-set watcher NATS connectivity (1=connected, 0=disconnected)", ))?; + let usage_parse_errors_total = IntCounter::with_opts(Opts::new( + "ai_usage_parse_errors_total", + "Managed 2xx responses with no parseable usage (emitted as a zero-token billing row)", + ))?; r.register(Box::new(requests_total.clone()))?; r.register(Box::new(rejections_total.clone()))?; @@ -145,6 +154,7 @@ impl Metrics { r.register(Box::new(requests_in_flight.clone()))?; r.register(Box::new(deny_set_size.clone()))?; r.register(Box::new(nats_connected.clone()))?; + r.register(Box::new(usage_parse_errors_total.clone()))?; Ok(Arc::new(Self { requests_total, @@ -162,6 +172,7 @@ impl Metrics { requests_in_flight, deny_set_size, nats_connected, + usage_parse_errors_total, })) } } diff --git a/src/peek.rs b/src/peek.rs index c8ac81b..b39dd8d 100644 --- a/src/peek.rs +++ b/src/peek.rs @@ -165,11 +165,11 @@ impl ModelScanner { pub fn plan_stream_usage_injection(body: &[u8]) -> Option { let n = body.len(); // Cheap pre-filter: injection is only ever needed when a root-level `"stream"` key is present. - // If the substring `"stream"` doesn't occur *anywhere*, the structural answer is unconditionally - // `None`, so a single SIMD `memmem` pass lets us skip the whole walk — the common case, since - // most requests aren't streaming. (The needle is a substring of `"stream_options"` too, so a - // body carrying only stream_options still passes the filter and is correctly resolved to `None` - // by the walk below.) + // If the quoted token `"stream"` doesn't occur *anywhere*, the structural answer is + // unconditionally `None`, so a single SIMD `memmem` pass lets us skip the whole walk — the + // common case, since most requests aren't streaming. (Note `"stream_options"` does NOT contain + // the needle: the byte after `stream` is `_`, not a closing quote — so a body carrying only + // `stream_options` fails this pre-filter and returns `None` here, which is the correct answer.) memchr::memmem::find(body, b"\"stream\"")?; let mut i = 0; while i < n && body[i].is_ascii_whitespace() { diff --git a/src/proxy.rs b/src/proxy.rs index 9b77e18..5975c92 100644 --- a/src/proxy.rs +++ b/src/proxy.rs @@ -837,13 +837,34 @@ impl ProxyHttp for AiProxy { let tail = &rc.resp_tail[tail_start..]; // Extract usage facts from the tail (shape depends on dialect + streaming). - let usage = match (rc.dialect, rc.streaming) { + let parsed = match (rc.dialect, rc.streaming) { (Dialect::OpenAI, true) => usage::openai_stream(tail), (Dialect::OpenAI, false) => usage::openai_body(tail), (Dialect::Anthropic, true) => usage::anthropic_stream(tail), (Dialect::Anthropic, false) => usage::anthropic_body(tail), + }; + // A managed 2xx response is *expected* to carry usage; `None` there means the provider's + // usage block changed shape (a new API version, a wire change) and we're about to emit a + // zero-token billing row that looks exactly like a (non-existent) legitimate zero-token + // generation — silently zeroing that tenant's bill. Surface it on a counter + a warn so it + // can be alerted on. A `None` on a 4xx/5xx (error body has no usage) is normal, not logged. + if parsed.is_none() + && rc.managed + && let Some(s) = rc.upstream_status + && (200..300).contains(&s) + { + self.state.metrics.usage_parse_errors_total.inc(); + warn!( + request_id = %rc.request_id, + tenant_id = rc.tenant_id, + provider = rc.provider.name.as_str(), + dialect = ?rc.dialect, + stream = rc.streaming, + status = s, + "managed 2xx response carried no parseable usage; emitting a zero-token billing row", + ); } - .unwrap_or_default(); + let usage = parsed.unwrap_or_default(); let m = &self.state.metrics; // Pre-resolved fixed-label children (see `Metrics`) — no per-call `with_label_values` lookup. diff --git a/src/state.rs b/src/state.rs index 26a7446..7f58255 100644 --- a/src/state.rs +++ b/src/state.rs @@ -143,6 +143,18 @@ impl GatewayState { (no key swap, no deny-set, no billing). Expected only for a BYO-only deployment." ); } + // Disabling upstream cert verification turns every provider connection into an unverified + // channel — a transparent-MitM opening. It's legitimate *only* for the local self-signed TLS + // mock (e2e/bench). Warn loudly at boot so it can never be a silent production misconfig (an + // `AI_UPSTREAM_VERIFY_CERT=false` copied out of a bench env) that looks healthy. + if config.upstream_tls && !config.upstream_verify_cert { + warn!( + "upstream TLS certificate verification is DISABLED — connections to providers are \ + unauthenticated and vulnerable to interception. This is valid ONLY for a local \ + test/bench mock; never set upstream_verify_cert=false against a real provider." + ); + } + let providers = build_providers(&config, &metrics); let rate_limit = RateLimit::new(config.rate_limit_rps, config.byo_rate_limit_rps); diff --git a/src/usage.rs b/src/usage.rs index 7c97a37..fc4f922 100644 --- a/src/usage.rs +++ b/src/usage.rs @@ -98,7 +98,9 @@ fn sse_data_lines(sse: &[u8]) -> impl Iterator + '_ { // SSE strips *all* leading spaces after the field colon (not exactly one) — OpenAI/Anthropic // emit `data: ` (one space), but a config-added OpenAI-wire provider that pads with more // would otherwise leave whitespace in the payload and fail the JSON parse → silent zero usage. - let line = line.trim_ascii_start(); + // Trim the trailing end too: SSE permits CRLF line endings (RFC 8895), so splitting on `\n` + // leaves a trailing `\r` that would otherwise fail the JSON parse → another silent zero. + let line = line.trim_ascii(); (line != b"[DONE]").then_some(line) }) } @@ -261,6 +263,24 @@ mod tests { ); } + #[test] + fn tolerates_crlf_line_endings() { + // SSE permits CRLF (RFC 8895). Splitting on `\n` leaves a trailing `\r` on each line; the + // parser must strip it or the JSON parse silently fails → a phantom zero-token billing row. + let sse = + b"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":11,\"completion_tokens\":22}}\r\n\r\n\ + data: [DONE]\r\n\r\n"; + assert_eq!( + openai_stream(sse).unwrap(), + Usage { + input_tokens: 11, + output_tokens: 22, + cache_read_tokens: 0, + cache_write_tokens: 0 + } + ); + } + #[test] fn no_usage_returns_none() { // Absent `usage` and unparseable bodies must both meter as `None` — never a silent zero-token