gateway: close the full mutation run's last survivors; client cut mid-upload is not billed (D259)

.cargo/mutants.toml

@@ -59,8 +59,10 @@ exclude_re = [
   # `map_response` clones a same-wire body as is: the guard only skips that clone.
   'replace match guard client == Endpoint::Responses with false in SseBridge::assembled',
   # `str_is`'s raw compare is a fast path: `escaped_key_is` decodes a string with no escape to
-  # itself, so skipping the raw compare (`false`) gives the same answer, only slower.
-  '''replace match guard !inner\.contains\(&b'\\\\'\) with false in str_is''',
+  # itself, so skipping the raw compare (`false`) gives the same answer, only slower. `str_value`'s
+  # raw borrow is the same fast path: serde decodes a valid string with no escape to its own bytes.
+  # (A raw control byte, which serde refuses, is not valid JSON.)
+  '''replace match guard !inner\.contains\(&b'\\\\'\) with false in (str_is|str_value)''',
   # `unserved` ORs `1 << i` for distinct indices `i`, where `|` and `^` are the same operation.
   'replace \| with \^ in unserved',
   # The const `TALLY` / `TALLY4` tables pack `state << 3 | class` and `next | add << 4` from
@@ -70,6 +72,9 @@ exclude_re = [
   # `de_service_tier`'s visitors for a non-string tier answer `Ok(None)`, which is `Option`'s
   # default.
   'replace de_service_tier::<impl serde::de::Visitor<.de> for V>::visit_(bool|i64|u64|f64|none|unit) -> Result<Self::Value, E> with Ok\(Default::default\(\)\)',
+  # Its `expecting` only words an "invalid type" error, and the visitor accepts every JSON type, so
+  # serde never builds one.
+  'replace de_service_tier::<impl serde::de::Visitor<.de> for V>::expecting -> std::fmt::Result with Ok\(Default::default\(\)\)',
   # A push that leaves the buffer at exactly the cap: `>=` "compacts" it with a self-copy and a
   # no-op truncate, and `ring()` reads it as a ring either way.
   'replace > with >= in UsageTail::push',
@@ -92,4 +97,30 @@ exclude_re = [
   # `Ctx` implements `Deref` as `DerefMut`'s supertrait: every hook holds the context mutably, so
   # nothing calls `deref` itself (a panicking body passes the integration suites).
   'replace <impl std::ops::Deref for Ctx>::deref -> &Self::Target',
+  # `MAX_CATALOG_OUTPUT_TOKENS` folds the catalog's largest `max_output_tokens`: replacing the
+  # running maximum with an equal value (`>=`) changes nothing. A const initializer's mutant names
+  # no function, only the file and position.
+  'translate\.rs:\d+:\d+: replace > with >=$',
+  # `KEY_NODES` sizes the key trie as the root plus one node per key byte. The keys share prefixes
+  # (`in`put / `in`structions), so the trie uses fewer: one node less (`1 * n`) still
+  # holds it, and every index stays below the bound `feed` takes it modulo.
+  'usage\.rs:\d+:\d+: replace \+ with \*$',
+  # `run` cuts a run at a space so the halves walk as two chains; a space resets the state, so where
+  # (or whether) it cuts changes no count, only the speed. The cut's arithmetic is tested; its
+  # length gate (the only `>=` in `run`) is that choice.
+  'replace >= with < in InputTally::run',
+  # `tracks_terminal` only skips work: the terminal verdict is read on a managed request's billing
+  # row alone, and a non-stream or assembled body never ends in a terminal SSE event, so watching
+  # more of them (`||`, `true`) reaches the same verdicts. Watching none (`false`) is tested.
+  'replace && with \|\| in RequestCtx::tracks_terminal',
+  'replace RequestCtx::tracks_terminal -> bool with true',
+  # A large body's re-run takes its parent's `unserved` mask only when it walks the row's
+  # `candidates`, which the mask was computed against; walking the Responses arm it takes 0. No row
+  # with a Responses arm has a candidate the mask can flag (pinned by
+  # `route::tests::no_row_with_a_responses_arm_has_an_unserved_candidate`), so the parent's mask
+  # is 0 there too and the guard (`true`) changes nothing today.
+  'replace match guard std::ptr::eq\(arms, row\.candidates\) with true in <impl ProxyHttp for AiProxy>::request_filter',
+  # `keep_usage_head` keeps a head only where it is read (an Anthropic stream); widening the gate
+  # (`||`) keeps a bounded head nothing reads.
+  'replace && with \|\| in RequestCtx::keep_usage_head',
 ]

crates/gateway/ARCHITECTURE.md

@@ -1757,8 +1757,11 @@ Two more endings the provider bills and we used to write as zero get the same tr
   (D252), so a non-stream pro call is answered and billed in full rather than cut at 600s and billed
   an estimate for nothing delivered.
   "Had the whole request" is `body_delivered`: the connection was up (`upstream_request_filter`
-  ran), the client's body was read to its end, and nothing failed writing it upstream. A request
-  that never reached a provider (every breaker open, a connect failure, a reset mid-upload) stays
+  ran), the client's body was read to its end, and nothing failed writing it upstream. A client that
+  closes before its body's framing-defined end (bytes of a `Content-Length` still to come, a chunked
+  body's last chunk missing) did not get it there, though pingora marks such a body done
+  (`client_body_cut_short`, D259). A request that never reached a provider (every breaker open, a
+  connect failure, a reset mid-upload, the client closing mid-upload) stays
   unbilled, and so does one the upstream **closed or reset** before any head (the gateway's 502): the
   peer declined to answer, and nothing says it processed the request. The real edges in front of
   Anthropic, OpenAI and OpenRouter answer a request they forwarded and then lost with an HTTP error

crates/gateway/src/peek.rs

@@ -609,7 +609,9 @@ pub fn scan_buffered(body: &[u8]) -> BufferedScan {
             }
             b'[' => depth += 1,
             b'}' | b']' => depth = depth.saturating_sub(1),
-            b':' if depth == 1 => expect_key = false,
+            // Unguarded by depth, as in both walks this one replaces: a nested `:` is always
+            // followed by the root `,` or `}` that closes its value before the next root key.
+            b':' => expect_key = false,
             b',' if depth == 1 => {
                 expect_key = true;
                 last_key_is_stream = false;
@@ -630,8 +632,8 @@ pub fn scan_buffered(body: &[u8]) -> BufferedScan {
                 continue;
             }
             // A negative limit is no limit to cap. (A string or structured value never puts a digit
-            // at depth 1 before the `,` that ends it.)
-            b'-' if depth == 1 => limit_key = None,
+            // at depth 1 before the `,` that ends it, so a nested `-` may clear it too.)
+            b'-' => limit_key = None,
             b't' if depth == 1 && last_key_is_stream => {
                 if body[j..].starts_with(b"true") {
                     stream_true = true;
@@ -1229,6 +1231,10 @@ mod tests {
         let n = last_member(b, 0, "n").unwrap().unwrap();
         assert!(!str_is(b, n.value, "null"));
         assert_eq!(str_value(b, n.value), None);
+        // An escaped value is decoded, as the provider's parser decodes it.
+        let e = br#"{"type":"im\u0061ge"}"#;
+        let m = last_member(e, 0, "type").unwrap().unwrap();
+        assert_eq!(str_value(e, m.value).as_deref(), Some("image"));
         // A string far longer than the name is not decoded to be compared.
         let long = format!(r#"{{"k":"{}"}}"#, "\\u0041".repeat(10_000));
         let m = last_member(long.as_bytes(), 0, "k").unwrap().unwrap();
@@ -1533,6 +1539,10 @@ mod tests {
             b"not json".to_vec(),
             b"".to_vec(),
             b"{}".to_vec(),
+            // Not JSON: a key with no `:` before its nested value. Both walks read the nested `:`
+            // as the root's, so neither takes the next root string for a key.
+            br#"{"a" {"b":1} "model":"x"}"#.to_vec(),
+            br#"{"a" {"b":1} "stream":true}"#.to_vec(),
             format!(r#"{{"messages":[{{"content":"{big}"}}],"stream":true,"model":"gpt-4o"}}"#)
                 .into_bytes(),
             format!(r#"{{"system":"{big} \"stream\":true","model":"x"}}"#).into_bytes(),

crates/gateway/src/proxy.rs

@@ -650,11 +650,39 @@ fn body_delivered(session: &mut Session, rc: &RequestCtx, e: Option<&pingora_cor
     rc.upstream_phase == UpstreamPhase::Connected
         && session.as_mut().is_body_done()
         && !e.is_some_and(|e| {
-            e.esource() == &pingora_core::ErrorSource::Upstream
-                && (matches!(e.etype(), WriteError | WriteTimedout) || h2_body_unsent(e))
+            (e.esource() == &pingora_core::ErrorSource::Upstream
+                && (matches!(e.etype(), WriteError | WriteTimedout) || h2_body_unsent(e)))
+                || client_body_cut_short(e)
         })
 }
 
+/// Whether the client closed its HTTP/1.1 connection before its body's framing-defined end
+/// (D259): a `Content-Length` body with bytes still to come, or a chunked one without its
+/// terminating chunk. Pingora then marks the body *done* (`is_body_done`), though never complete,
+/// so without this the half a provider was sent read as the whole request, and `logging` billed
+/// it an estimate. Like [`h2_body_unsent`], the context string is pingora's only handle on it.
+fn client_body_cut_short(e: &pingora_core::Error) -> bool {
+    if e.esource() != &pingora_core::ErrorSource::Downstream {
+        return false;
+    }
+    let mut at = Some(e);
+    while let Some(x) = at {
+        if x.context.as_ref().is_some_and(|c| {
+            let c = c.as_str();
+            c.starts_with("Peer prematurely closed connection with")
+                || c.starts_with("Connection prematurely closed without the termination chunk")
+        }) {
+            return true;
+        }
+        at = x
+            .cause
+            .as_deref()
+            .and_then(|c| c.downcast_ref::<Box<pingora_core::Error>>())
+            .map(|b| &**b);
+    }
+    false
+}
+
 /// Whether writing the request body to an HTTP/2 stream failed because the stream had already
 /// closed with body bytes still unsent: pingora's `reserve_and_send` found no capacity to get
 /// (`cannot reserve capacity`, `while waiting for capacity`), a write failure it labels `H2Error`
@@ -1048,6 +1076,27 @@ impl RequestCtx {
             || self.signed.is_some()
     }
 
+    /// Whether the bytes sent to the client are watched for the stream's terminal event
+    /// (`closed_after_terminal`): a managed stream's, unless it is assembled into one JSON body.
+    /// The verdict is read only on a managed request's billing row, and a JSON body has no
+    /// terminal event, so watching anything else would scan every chunk for nothing.
+    fn tracks_terminal(&self) -> bool {
+        self.managed && self.streaming && !self.auto.as_ref().is_some_and(|a| catalog_assembling(a))
+    }
+
+    /// Anthropic SSE only: keep a bounded head of the upstream response so `message_start`'s input
+    /// and cache token counts survive the tail's compaction. At most `USAGE_HEAD_CAP` bytes in all,
+    /// satisfied within the first chunk or two, after which this copies nothing: one small
+    /// allocation on the one path that needs it, and nothing anywhere else. See `USAGE_HEAD_CAP`
+    /// for why only this dialect needs it.
+    fn keep_usage_head(&mut self, chunk: &[u8]) {
+        if self.streaming && self.dialect == Dialect::Anthropic {
+            let want = USAGE_HEAD_CAP.saturating_sub(self.resp_head.len());
+            self.resp_head
+                .extend_from_slice(&chunk[..want.min(chunk.len())]);
+        }
+    }
+
     /// Point the forwarded path at the candidate about to be attempted.
     ///
     /// Taken wholesale from the catalog rather than composed from a mount and the client's suffix:
@@ -4603,26 +4652,8 @@ impl ProxyHttp for AiProxy {
             streaming: false,
             inject_eligible,
             background_check,
-            // Only the inject-eligible path ever buffers the request body (to splice
-            // `stream_options` after the root `{`; the `stream` key can appear anywhere in the root
-            // object, so the decision needs the whole body — buffering is inherent here, not
-            // incidental). When it does, pre-size from the declared Content-Length so accumulation is
-            // a single allocation instead of a geometric realloc chain; capped at `MAX_REQUEST_BODY`
-            // so a lying header can't pre-allocate unbounded memory. Every other request leaves this
-            // empty and never buffers.
-            //
-            // The `+ STREAM_OPTIONS_FRAG.len()` is headroom for the splice, which
-            // `apply_stream_usage_injection` performs *in place*: with it the injection never
-            // reallocates, so a body arrives, is spliced, and goes upstream on one allocation.
-            req_buf: match (
-                inject_eligible || background_check || model_route.is_some() || signed.is_some(),
-                declared_len,
-            ) {
-                (true, Some(len)) => {
-                    Vec::with_capacity(len.min(MAX_REQUEST_BODY) + STREAM_OPTIONS_FRAG.len())
-                }
-                _ => Vec::new(),
-            },
+            // Pre-sized below, once the context says whether this request rewrites its body.
+            req_buf: Vec::new(),
             // Grown lazily by the response tap (`response_body_filter`), not pre-reserved: a
             // non-streaming response — the common case — is a few hundred bytes, so reserving the
             // full 64KB cap up front would waste an allocation on every request to hold ~200B. A
@@ -4672,6 +4703,20 @@ impl ProxyHttp for AiProxy {
             terminal: TerminalTracker::default(),
             signed,
         });
+        // A body buffered for a rewrite (`RequestCtx::rewrites_body`) is pre-sized from the
+        // declared Content-Length, so accumulation is a single allocation instead of a geometric
+        // realloc chain; capped at `MAX_REQUEST_BODY` so a lying header can't pre-allocate
+        // unbounded memory. Every other request leaves it empty and never buffers.
+        //
+        // The `+ STREAM_OPTIONS_FRAG.len()` is headroom for the splice, which
+        // `apply_stream_usage_injection` performs *in place*: with it the injection never
+        // reallocates, so a body arrives, is spliced, and goes upstream on one allocation.
+        if let Some(rc) = ctx.rc.as_mut()
+            && rc.rewrites_body()
+            && let Some(len) = declared_len
+        {
+            rc.req_buf = Vec::with_capacity(len.min(MAX_REQUEST_BODY) + STREAM_OPTIONS_FRAG.len());
+        }
         // Admitted: count it in-flight. Released in `logging`, or by `Ctx`'s drop if a panic
         // skipped `logging`, so the gauge cannot leak. `active_streams` only covers SSE; this
         // covers every request.
@@ -5873,19 +5918,8 @@ impl ProxyHttp for AiProxy {
                 rc.resp_model_scanner.feed(chunk);
             }
 
-            // Anthropic SSE only: keep a bounded head so `message_start`'s input + cache token
-            // counts survive the tail's compaction. Bounded by `USAGE_HEAD_CAP` and satisfied
-            // within the first chunk or two, after which the length check makes this a no-op — so
-            // it costs one small allocation on the one path that needs it, and nothing anywhere
-            // else. See `USAGE_HEAD_CAP` for why only this dialect needs it.
-            if rc.streaming
-                && rc.dialect == Dialect::Anthropic
-                && rc.resp_head.len() < USAGE_HEAD_CAP
-            {
-                let want = USAGE_HEAD_CAP - rc.resp_head.len();
-                rc.resp_head
-                    .extend_from_slice(&chunk[..want.min(chunk.len())]);
-            }
+            // Anthropic SSE only: the head that keeps `message_start` (see `keep_usage_head`).
+            rc.keep_usage_head(chunk);
 
             rc.resp_tail.push(chunk);
             // Managed only. Counts *upstream* bytes (pre-translate), like the tail.
@@ -5972,15 +6006,12 @@ impl ProxyHttp for AiProxy {
                     tap.push(&out);
                 }
             }
-            if rc.managed
-                && rc.streaming
-                && !rc.auto.as_ref().is_some_and(|a| catalog_assembling(a))
-            {
+            if rc.tracks_terminal() {
                 rc.terminal.feed(&out);
             }
             *body = Some(Bytes::from(out));
         } else if !chunk.is_empty() {
-            if rc.managed && rc.streaming {
+            if rc.tracks_terminal() {
                 rc.terminal.feed(chunk);
             }
             // Capture tap — same passive-tap contract as the usage tail above (copy, never withhold),
@@ -6966,6 +6997,25 @@ mod tests {
         }
     }
 
+    /// The Anthropic stream head is bounded at `USAGE_HEAD_CAP` however the stream is chunked: the
+    /// chunk that crosses the cap is cut there, and later chunks add nothing. Other responses keep
+    /// no head.
+    #[test]
+    fn the_usage_head_holds_at_most_its_cap() {
+        let mut rc = test_ctx(false);
+        rc.streaming = true;
+        rc.keep_usage_head(&[b'a'; 1000]);
+        assert!(rc.resp_head.is_empty(), "an OpenAI stream keeps no head");
+        rc.dialect = Dialect::Anthropic;
+        rc.keep_usage_head(&[b'a'; 5 * 1024]);
+        assert_eq!(rc.resp_head.len(), 5 * 1024);
+        rc.keep_usage_head(&[b'b'; 10 * 1024]);
+        assert_eq!(rc.resp_head.len(), USAGE_HEAD_CAP);
+        assert_eq!(rc.resp_head.last(), Some(&b'b'));
+        rc.keep_usage_head(b"more");
+        assert_eq!(rc.resp_head.len(), USAGE_HEAD_CAP);
+    }
+
     /// The replay contract, exercised directly: `upstream_peer` resets the body phase before pingora
     /// re-feeds the buffered prefix through `request_body_filter`, so attempt 2 sees the body once.
     ///