Fix audit findings: CB race, silent billing zeros, config/SSE robustness

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)    |
 
 ---
 

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 {

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<String> {
 fn reject_unknown_toml_keys(path: &Path) -> Result<()> {
     use figment::Provider as _;
     let known = known_config_keys();
-    let unknown: std::collections::BTreeSet<String> = 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<String> = 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

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,
         }))
     }
 }

src/peek.rs

@@ -165,11 +165,11 @@ impl ModelScanner {
 pub fn plan_stream_usage_injection(body: &[u8]) -> Option<usize> {
     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() {

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.

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);