diff --git a/.cargo/mutants.toml b/.cargo/mutants.toml index 77132521..2e050dc4 100644 --- a/.cargo/mutants.toml +++ b/.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:: for V>::visit_(bool|i64|u64|f64|none|unit) -> Result 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:: 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 ::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 ::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', ] diff --git a/crates/gateway/ARCHITECTURE.md b/crates/gateway/ARCHITECTURE.md index 9f1de94e..14c06951 100644 --- a/crates/gateway/ARCHITECTURE.md +++ b/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 diff --git a/crates/gateway/src/peek.rs b/crates/gateway/src/peek.rs index 50808949..9891cc28 100644 --- a/crates/gateway/src/peek.rs +++ b/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(), diff --git a/crates/gateway/src/proxy.rs b/crates/gateway/src/proxy.rs index 26ff2479..d254724f 100644 --- a/crates/gateway/src/proxy.rs +++ b/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::>()) + .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. /// diff --git a/crates/gateway/src/route.rs b/crates/gateway/src/route.rs index 7afa9b61..51716acd 100644 --- a/crates/gateway/src/route.rs +++ b/crates/gateway/src/route.rs @@ -818,6 +818,31 @@ mod tests { } } + /// A large body's re-run reads [`unserved`] from its parent, which computed it against the + /// row's `candidates`; a re-run that walks the Responses arm instead takes 0 (`proxy`'s + /// `std::ptr::eq(arms, row.candidates)` guard). That is exact only while no row with a Responses + /// arm has a candidate `unserved` can flag. A row that breaks this needs the parent to compute + /// the mask for the arm the re-run walks (and the guard's exclusion in `.cargo/mutants.toml` + /// dropped). + #[test] + fn no_row_with_a_responses_arm_has_an_unserved_candidate() { + for route in providers::catalog::MODEL_ROUTES { + if route.responses.is_empty() { + continue; + } + for c in route.candidates.iter().chain(route.responses) { + assert!( + providers::catalog::serves_structured_outputs(c) + && providers::catalog::serves_file_input(c), + "{}: {} {}", + route.model, + providers::by_id(c.provider).name, + c.upstream_model, + ); + } + } + } + #[test] fn is_default_prefix_boundary_checks() { // The real default-prefix shape: exactly "/v1" or "/v1/…". @@ -1130,6 +1155,59 @@ mod tests { assert_eq!(p.pool_auth[1].value.expose(), "Bearer sk-b"); } + /// A refused key 0 starts the next request on key 1, in a pool of two (the smallest with a + /// key to skip to) or more, and the provider is not cooling while a key is fresh. + #[test] + fn a_refused_first_key_starts_requests_on_the_next() { + for keys in [&["sk-a", "sk-b"][..], &["sk-a", "sk-b", "sk-c"]] { + let p = Provider::resolve( + "openai", + "api.openai.com:443".to_string(), + Dialect::OpenAi, + AuthScheme::Bearer, + keys, + ProviderMetrics::disconnected(), + None, + ); + assert_eq!(p.first_key(), 0); + p.mark_key_bad(0); + assert_eq!( + p.first_key(), + 1, + "{} keys: key 0 is cooling off", + keys.len() + ); + assert!(!p.cooling(), "a fresh key can still take the request"); + for i in 1..keys.len() { + p.mark_key_bad(i as u8); + } + assert!(p.cooling()); + assert_eq!(p.first_key(), 0, "every key cooling: key 0 still answers"); + } + } + + /// The cooldown is measured on a running clock: a refusal whose cooldown has passed no longer + /// skips the key. + #[test] + fn a_key_stops_cooling_once_its_cooldown_has_passed() { + let p = Provider::resolve( + "openai", + "api.openai.com:443".to_string(), + Dialect::OpenAi, + AuthScheme::Bearer, + &["sk-a", "sk-b"], + ProviderMetrics::disconnected(), + None, + ); + // A cooldown that ends 200 ms from now, rather than `KEY_COOLDOWN`'s minute. + p.pool_auth[0] + .bad_until_ms + .store(clock_ms().saturating_add(200), Ordering::Relaxed); + assert_eq!(p.first_key(), 1); + std::thread::sleep(Duration::from_millis(250)); + assert_eq!(p.first_key(), 0, "the cooldown has passed"); + } + #[test] fn precomputed_header_values_match_the_string_form() { // The per-request insert now clones these instead of re-validating and re-copying the diff --git a/crates/gateway/src/translate.rs b/crates/gateway/src/translate.rs index 8c0bbb40..dc788b43 100644 --- a/crates/gateway/src/translate.rs +++ b/crates/gateway/src/translate.rs @@ -8021,7 +8021,9 @@ fn next_event(buf: &[u8], start: &mut usize, scanned: &mut usize) -> Option<(usi } } -/// End (exclusive) of the first `\n\n` or `\r\n\r\n`, judging only `\n`s at or after `from`. +/// End (exclusive) of the first blank line: a line ending (`\n`, or the `\n` of `\r\n`) followed by +/// an empty line ended by `\n` or `\r\n`, judging only `\n`s at or after `from`. That covers +/// `\n\n` and `\r\n\r\n`, and the mixed endings SSE allows too. /// /// One forward pass over newlines. The earlier two-`windows` form ran both searches to completion, /// so a stream using only one line ending walked the whole buffer for the other on every call. @@ -8031,11 +8033,7 @@ fn find_event_end(buf: &[u8], from: usize) -> Option { let nl = at + off; match buf.get(nl + 1) { Some(b'\n') => return Some(nl + 2), - Some(b'\r') - if nl > 0 && buf.get(nl - 1) == Some(&b'\r') && buf.get(nl + 2) == Some(&b'\n') => - { - return Some(nl + 3); - } + Some(b'\r') if buf.get(nl + 2) == Some(&b'\n') => return Some(nl + 3), _ => {} } at = nl + 1; @@ -8113,6 +8111,19 @@ mod tests { assert_eq!(b.pending_len(), big.len()); } + /// A blank line ends an event whatever line endings frame it; a lone `\r\n` (one line ending) + /// does not. + #[test] + fn an_event_ends_at_a_blank_line_of_either_ending() { + assert_eq!(find_event_end(b"data: a\n\nb", 0), Some(9)); + assert_eq!(find_event_end(b"data: a\r\n\r\nb", 0), Some(11)); + assert_eq!(find_event_end(b"data: a\n\r\nb", 0), Some(10)); + assert_eq!(find_event_end(b"data: a\r\n\nb", 0), Some(10)); + assert_eq!(find_event_end(b"\n\r\nb", 0), Some(3)); + assert_eq!(find_event_end(b"data: a\r\ndata: b\r\n", 0), None); + assert_eq!(find_event_end(b"data: a\n\r", 0), None, "undecided"); + } + /// The resumable search must find exactly what a from-scratch search finds, including a /// terminator split across chunks at every possible point, for both line endings. #[test] @@ -8121,6 +8132,8 @@ mod tests { &b"data: a\n\ndata: bb\n\n: c\n\n"[..], &b"data: a\r\n\r\ndata: bb\r\n\r\n"[..], &b"data: a\r\ndata: b\n\ndata: c\r\n\r\n"[..], + // A `\n` line ending then a `\r\n` blank line, and the reverse: SSE allows both. + &b"data: a\n\r\ndata: b\r\n\ndata: c\n\n"[..], ] { // As `SseBridge::feed` holds it: events taken by cursor, the buffer compacted per part. let events = |parts: &[&[u8]]| { diff --git a/crates/gateway/src/translate_response_tests.rs b/crates/gateway/src/translate_response_tests.rs index ee1acebd..9cdf3b2c 100644 --- a/crates/gateway/src/translate_response_tests.rs +++ b/crates/gateway/src/translate_response_tests.rs @@ -3303,3 +3303,75 @@ fn a_responses_refusal_streams_onto_chat() { .collect(); assert_eq!(refusal, "Nope.", "{evs:#?}"); } + +/// A Chat stream's reasoning reaches a Responses client as one reasoning item per block: a new +/// OpenRouter block index opens one (the same index continues it), so does text after a signature +/// (even text with no index), and a whole signed `thinking` block opens its own, its text intact. +/// Each item announces its summary part once, before its first delta. +#[test] +fn chat_reasoning_blocks_become_one_responses_item_each() { + let src = concat!( + "data: {\"id\":\"g\",\"model\":\"m\",\"choices\":[{\"index\":0,\"delta\":{\"reasoning_details\":[{\"type\":\"reasoning.text\",\"text\":\"a1 \",\"format\":\"anthropic-claude-v1\",\"index\":0}]}}]}\n\n", + "data: {\"id\":\"g\",\"model\":\"m\",\"choices\":[{\"index\":0,\"delta\":{\"reasoning_details\":[{\"type\":\"reasoning.text\",\"text\":\"a2\",\"format\":\"anthropic-claude-v1\",\"index\":0}]}}]}\n\n", + "data: {\"id\":\"g\",\"model\":\"m\",\"choices\":[{\"index\":0,\"delta\":{\"reasoning_details\":[{\"type\":\"reasoning.text\",\"text\":\"b\",\"format\":\"anthropic-claude-v1\",\"index\":1}]}}]}\n\n", + "data: {\"id\":\"g\",\"model\":\"m\",\"choices\":[{\"index\":0,\"delta\":{\"reasoning_details\":[{\"type\":\"reasoning.text\",\"signature\":\"s1\",\"format\":\"anthropic-claude-v1\",\"index\":1}]}}]}\n\n", + "data: {\"id\":\"g\",\"model\":\"m\",\"choices\":[{\"index\":0,\"delta\":{\"thinking\":[{\"type\":\"thinking\",\"thinking\":\"c\",\"signature\":\"s2\"}]}}]}\n\n", + "data: {\"id\":\"g\",\"model\":\"m\",\"choices\":[{\"index\":0,\"delta\":{\"reasoning\":\"d\"}}]}\n\n", + "data: {\"id\":\"g\",\"model\":\"m\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"done\"},\"finish_reason\":\"stop\"}]}\n\n", + "data: {\"id\":\"g\",\"model\":\"m\",\"choices\":[],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":1,\"total_tokens\":2}}\n\n", + "data: [DONE]\n\n", + ); + let evs = stream(Chat, Responses, src); + let resp = assert_responses_lifecycle(&evs, "response.completed"); + let out = resp["output"].as_array().unwrap(); + let summaries: Vec<(&str, &Value)> = out + .iter() + .map(|i| (i["type"].as_str().unwrap(), &i["summary"])) + .collect(); + let s = |t: &str| json!([{"type": "summary_text", "text": t}]); + assert_eq!( + summaries, + [ + ("reasoning", &s("a1 a2")), + ("reasoning", &s("b")), + ("reasoning", &s("c")), + ("reasoning", &s("d")), + ("message", &Value::Null), + ], + "{resp:#}" + ); + // One `summary_part.added` per item, each before that item's first delta. + let mut announced = Vec::new(); + for (name, v) in &evs { + let item = v["item_id"].as_str().unwrap_or_default().to_owned(); + match name.as_str() { + "response.reasoning_summary_part.added" => { + assert!(!announced.contains(&item), "{item} announced twice"); + announced.push(item); + } + "response.reasoning_summary_text.delta" => { + assert!(announced.contains(&item), "{item}: a delta before its part"); + } + _ => {} + } + } + assert_eq!(announced.len(), 4); +} + +/// A Chat tool call whose first chunk already carries its arguments reaches a Responses client +/// with them as an arguments delta, not only on the finished item. +#[test] +fn a_chat_tool_call_opened_with_its_arguments_streams_them_to_responses() { + let src = concat!( + "data: {\"id\":\"c\",\"model\":\"gpt\",\"choices\":[{\"index\":0,\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"type\":\"function\",\"function\":{\"name\":\"get_weather\",\"arguments\":\"{\\\"city\\\":\\\"Paris\\\"}\"}}]}}]}\n\n", + "data: {\"id\":\"c\",\"model\":\"gpt\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"tool_calls\"}]}\n\n", + "data: {\"id\":\"c\",\"model\":\"gpt\",\"choices\":[],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":1,\"total_tokens\":2}}\n\n", + "data: [DONE]\n\n", + ); + let evs = stream(Chat, Responses, src); + let deltas: String = named(&evs, "response.function_call_arguments.delta") + .iter() + .map(|v| v["delta"].as_str().unwrap()) + .collect(); + assert_eq!(deltas, "{\"city\":\"Paris\"}", "{evs:#?}"); +} diff --git a/crates/gateway/src/usage.rs b/crates/gateway/src/usage.rs index 132fc6de..7127e1cb 100644 --- a/crates/gateway/src/usage.rs +++ b/crates/gateway/src/usage.rs @@ -103,9 +103,6 @@ fn de_service_tier<'de, D: serde::Deserializer<'de>>( fn visit_unit(self) -> Result { Ok(None) } - fn visit_some>(self, d: D) -> Result { - d.deserialize_any(V) - } fn visit_bool(self, _: bool) -> Result { Ok(None) } @@ -1056,16 +1053,15 @@ static TALLY: [u8; 128] = { WIDE => (s, 0), LETTER => match from { LETTERS | LETTERS_APOSTROPHE => (LETTERS, 0), - // One punctuation byte joins the letters after it. - PUNCT1 => (LETTERS, 1), PUNCT2 => (LETTERS, 2), + // A new word, which one punctuation byte before it joins (`PUNCT1`). _ => (LETTERS, 1), }, DIGIT => match from { DIGITS1 => (DIGITS2, 0), DIGITS2 => (DIGITS3, 0), - DIGITS3 => (DIGITS1, 1), LETTERS_APOSTROPHE | PUNCT1 | PUNCT2 => (DIGITS1, 2), + // A new group: a fourth digit (`DIGITS3`) starts one too. _ => (DIGITS1, 1), }, _ => match from { @@ -1201,15 +1197,11 @@ impl InputTally { fn text(&mut self, rest: &[u8]) -> usize { let mut i = 0; loop { - if self.aux > 0 { - // The hex digits of a `\u` escape, already classed as one wide byte. - let n = usize::from(self.aux).min(rest.len() - i); - i += n; - self.aux -= n as u8; - if self.aux > 0 { - return i; - } - } + // The hex digits of a `\u` escape, already classed as one wide byte (none outside + // one). Digits still pending leave nothing of `rest`, which the next line returns on. + let n = usize::from(self.aux).min(rest.len() - i); + i += n; + self.aux -= n as u8; let Some(&b) = rest.get(i) else { return i }; if self.flags & ESCAPE != 0 { self.flags &= !ESCAPE; @@ -1700,6 +1692,33 @@ mod tests { } } + /// A long run of plain text (no escape to split it first) is walked as two chains cut at a + /// space past its middle: the count is the one a walk of short pieces, never cut, gives. + #[test] + fn a_long_plain_run_counts_as_its_short_pieces_do() { + let words = [ + "alpha", "it's", "x2", "(foo)", "1234567", "café", "a-b", "Mr.", "日本", "zz", + ]; + let mut seed = 0x9E37_79B9_7F4A_7C15_u64; + for len in [256, 257, 300, 511, 1024, 4096] { + let mut text = String::new(); + while text.len() < len { + seed ^= seed << 13; + seed ^= seed >> 7; + seed ^= seed << 17; + text.push_str(words[(seed % words.len() as u64) as usize]); + // Mostly no space between words, so the cut's space is rarely next to the middle. + if seed.is_multiple_of(5) { + text.push(' '); + } + } + let body = format!(r#"{{"content":"{text}"}}"#); + let whole = tally(&[body.as_bytes()]).estimate_tokens(); + let pieces: Vec<&[u8]> = body.as_bytes().chunks(64).collect(); + assert_eq!(tally(&pieces).estimate_tokens(), whole, "{text}"); + } + } + #[test] fn input_tally_is_small() { assert_eq!(std::mem::size_of::(), 12); diff --git a/crates/gateway/tests/cancellation.rs b/crates/gateway/tests/cancellation.rs index fc80efbd..e5e1c43c 100644 --- a/crates/gateway/tests/cancellation.rs +++ b/crates/gateway/tests/cancellation.rs @@ -408,3 +408,128 @@ async fn a_catalog_walk_resends_a_reused_connection_reset_before_reading() { "no pooled connection was reused, so the scenario did not fire" ); } + +/// A listener that accepts connections, reads whatever arrives and never answers; returns its port +/// and how many connections it has accepted. +async fn silent_upstream() -> (u16, std::sync::Arc) { + use tokio::io::AsyncReadExt; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let accepted = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let counter = accepted.clone(); + tokio::spawn(async move { + while let Ok((mut s, _)) = listener.accept().await { + counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + tokio::spawn(async move { + let mut buf = [0u8; 4096]; + while matches!(s.read(&mut buf).await, Ok(n) if n > 0) {} + }); + } + }); + (port, accepted) +} + +/// A client that resets its connection mid-upload ends its request: the walk does not fail over to +/// the next candidate (or resend to the same one) with a body nobody will finish sending. (A clean +/// close reads as the body's end, so the request counts as delivered and is not resent either.) +/// claim: REL-9 +#[tokio::test] +async fn a_client_gone_mid_upload_is_not_failed_over() { + use std::sync::atomic::Ordering::SeqCst; + use tokio::io::AsyncWriteExt; + let (primary, primary_conns) = silent_upstream().await; + let (fallback, fallback_conns) = silent_upstream().await; + let (pubkey, sk) = test_keypair(1); + let gw = Gateway::builder( + unused_nats_port(), + &format!("127.0.0.1:{primary}"), + &b64(&pubkey), + ) + .providers(&["openai", "openrouter"]) + .provider_authority("openrouter", &format!("127.0.0.1:{fallback}")) + .start() + .await; + let payload = body(); + let mut s = tokio::net::TcpStream::connect(("127.0.0.1", gw.port)) + .await + .unwrap(); + let head = format!( + "POST /auto/chat/completions HTTP/1.1\r\nhost: gw\r\nauthorization: Bearer {}\r\n\ + content-type: application/json\r\nx-beyond-model: gpt-4o-mini\r\n\ + x-beyond-order: openai,openrouter\r\ncontent-length: {}\r\n\r\n", + billing_vkey(&sk, 81), + payload.len() + ); + s.write_all(head.as_bytes()).await.unwrap(); + s.write_all(&payload.as_bytes()[..payload.len() / 2]) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(10), async { + while primary_conns.load(SeqCst) == 0 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("the request reached the primary"); + tokio::time::sleep(Duration::from_millis(200)).await; + s.set_zero_linger().unwrap(); + drop(s); + // The abort ends the request in `logging`; nothing is sent anywhere after it. + gw.wait_for_log_line(&["upstream request errored"]).await; + tokio::time::sleep(Duration::from_millis(300)).await; + assert_eq!( + (primary_conns.load(SeqCst), fallback_conns.load(SeqCst)), + (1, 0), + "log:\n{}", + gw.log() + ); +} + +/// A provider-routed request whose client stalls mid-upload on a pooled connection ends at the read +/// timeout, and is not resent: the timeout is a liveness verdict on an attempt that may still be +/// running, not a reused connection's failure before delivery (D248's `ReusedOnly` rule is for a +/// body write that met a closed stream only). +/// claim: REL-1 +#[tokio::test] +async fn a_read_timeout_on_a_reused_connection_is_not_resent() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let mock = ReplyUpstream::start(|_, _| Reply::ok()).await; + let (pubkey, _sk) = test_keypair(1); + let gw = Gateway::builder(unused_nats_port(), &mock.authority(), &b64(&pubkey)) + .config_line("read_timeout_secs = 1") + .start() + .await; + // Pools a connection to the upstream. + let first = test_client() + .post(format!("{}/openai/v1/chat/completions", gw.url())) + .header("authorization", "Bearer sk-byo-test") + .header("content-type", "application/json") + .body(body()) + .send() + .await + .unwrap(); + assert_eq!(first.status().as_u16(), 200); + let payload = body(); + let mut s = tokio::net::TcpStream::connect(("127.0.0.1", gw.port)) + .await + .unwrap(); + let head = format!( + "POST /openai/v1/chat/completions HTTP/1.1\r\nhost: gw\r\n\ + authorization: Bearer sk-byo-test\r\ncontent-type: application/json\r\n\ + content-length: {}\r\n\r\n", + payload.len() + ); + s.write_all(head.as_bytes()).await.unwrap(); + s.write_all(&payload.as_bytes()[..payload.len() / 2]) + .await + .unwrap(); + // The rest never comes: the gateway answers once the upstream read times out. + let mut buf = [0u8; 4096]; + let n = tokio::time::timeout(Duration::from_secs(10), s.read(&mut buf)) + .await + .expect("an answer") + .unwrap_or(0); + let text = String::from_utf8_lossy(&buf[..n]); + tokio::time::sleep(Duration::from_millis(300)).await; + assert_eq!(mock.hits(), 2, "{text}\nlog:\n{}", gw.log()); +} diff --git a/crates/gateway/tests/catalog_capabilities.rs b/crates/gateway/tests/catalog_capabilities.rs index cad4ebcc..4def5305 100644 --- a/crates/gateway/tests/catalog_capabilities.rs +++ b/crates/gateway/tests/catalog_capabilities.rs @@ -489,6 +489,18 @@ async fn a_pdf_request_skips_a_candidate_that_reads_none() { "{path}: the PDF reached xAI as input_file: {sent}" ); } + // A body past pingora's 64 KiB replay buffer is walked by a `FullBody` re-run, which takes the + // mask its parent computed: the same skip. + let big_pdf = format!( + "data:application/pdf;base64,JVBERi0xLjQK{}", + "A".repeat(96 * 1024) + ); + let large = json!({"model": "grok-build-0.1", "messages": [{"role": "user", "content": [ + {"type": "file", "file": {"filename": "a.pdf", "file_data": big_pdf}}, + {"type": "text", "text": "What is the code word?"}]}]}); + let resp = post(&gw, &key, "/v1/chat/completions", &order, &large).await; + assert_eq!(resp.status().as_u16(), 200, "large body"); + assert_eq!(provider_of(&resp).as_deref(), Some("xai"), "large body"); assert_eq!( openrouter.hits(), 0, diff --git a/crates/gateway/tests/cut_short.rs b/crates/gateway/tests/cut_short.rs index 2df98a81..2dbd12a7 100644 --- a/crates/gateway/tests/cut_short.rs +++ b/crates/gateway/tests/cut_short.rs @@ -211,6 +211,49 @@ async fn an_error_only_stream_is_not_billed_an_estimate() { assert_eq!(row["output_tokens"].as_u64(), Some(0), "{row}"); } +/// A stream whose error event comes *after* output (OpenRouter's mid-stream `"error":{…}`, the +/// provider failing part way through an answer) is work the provider did and bills: estimated from +/// what was relayed, unlike an error-only stream. +/// claim: BIL-3, BIL-12 +#[tokio::test] +async fn an_error_after_output_is_billed_an_estimate() { + let (pubkey, sk) = test_keypair(47); + let delta = "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"word \"}}]}\n\n"; + let sse = format!( + "{}data: {{\"error\":{{\"message\":\"upstream failed\",\"code\":502}}}}\n\n", + delta.repeat(3) + ); + let mock = ReplyUpstream::start(move |_, _| Reply::Full { + status: 200, + content_type: "text/event-stream", + body: bytes::Bytes::from(sse.clone()), + }) + .await; + let gw = Gateway::builder(unused_nats_port(), &mock.authority(), &b64(&pubkey)) + .start() + .await; + let resp = test_client() + .post(format!("{}/openai/v1/chat/completions", gw.url())) + .header("authorization", format!("Bearer {}", vkey(&sk, 47))) + .header("content-type", "application/json") + .body( + r#"{"model":"gpt-4o-mini","stream":true,"messages":[{"role":"user","content":"hi"}]}"#, + ) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + let _ = resp.text().await.unwrap(); + + let row = usage_row(&gw).await; + assert_eq!(row["usage_estimated"], true, "{row}"); + assert_eq!( + row["output_tokens"].as_u64(), + Some(3), + "one token per delta relayed: {row}" + ); +} + /// The over-count this guards: a Messages client's image is `{"type":"base64","data":"…"}` — not a /// data URI — and when an OpenAI-wire candidate serves it, input is estimated from the body. That /// payload must not be counted as text either. @@ -242,3 +285,56 @@ async fn anthropic_format_images_do_not_inflate_the_input_estimate() { "user (1) what is this? (4); an Anthropic base64 source is not text: {row}" ); } + +/// A client that closes before the end of its own body (a `Content-Length` body with bytes still +/// to come, a chunked one without its last chunk) never gave the provider the request: the half +/// it was sent is no prompt it bills. Not estimated (D259), unlike a client that sent the whole +/// body and then gave up waiting (D130). +/// claim: BIL-3 +/// defect: D259 +#[tokio::test] +async fn a_client_that_closes_mid_upload_is_not_billed_an_estimate() { + use tokio::io::AsyncWriteExt; + let payload = r#"{"model":"gpt-4o-mini","messages":[{"role":"user","content":"Explain TCP congestion control in detail, please."}]}"#; + let half = &payload[..payload.len() / 2]; + for (what, framing, sent) in [ + ( + "content-length", + format!("content-length: {}", payload.len()), + half.to_owned(), + ), + ( + "chunked", + "transfer-encoding: chunked".to_owned(), + format!("{:x}\r\n{half}\r\n", half.len()), + ), + ] { + let (pubkey, sk) = test_keypair(48); + let mock = ReplyUpstream::start(|_, _| Reply::HeaderStall).await; + let gw = Gateway::builder(unused_nats_port(), &mock.authority(), &b64(&pubkey)) + .start() + .await; + let mut s = tokio::net::TcpStream::connect(("127.0.0.1", gw.port)) + .await + .unwrap(); + let head = format!( + "POST /openai/v1/chat/completions HTTP/1.1\r\nhost: gw\r\nauthorization: Bearer {}\r\n\ + content-type: application/json\r\n{framing}\r\n\r\n", + vkey(&sk, 48), + ); + s.write_all(head.as_bytes()).await.unwrap(); + s.write_all(sent.as_bytes()).await.unwrap(); + tokio::time::timeout(Duration::from_secs(10), async { + while mock.hits() == 0 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("the request reached the upstream"); + tokio::time::sleep(Duration::from_millis(200)).await; + drop(s); + let row = usage_row(&gw).await; + assert_eq!(row["usage_estimated"], false, "{what}: {row}"); + assert_eq!(row["input_tokens"].as_u64(), Some(0), "{what}: {row}"); + } +} diff --git a/crates/gateway/tests/managed_headers.rs b/crates/gateway/tests/managed_headers.rs index 94c247c6..cad5f028 100644 --- a/crates/gateway/tests/managed_headers.rs +++ b/crates/gateway/tests/managed_headers.rs @@ -191,6 +191,38 @@ async fn a_translated_walk_merges_its_own_beta_with_the_allowed_client_tokens() ); } +/// The gateway supplies `anthropic-version` only where it wrote the Messages body itself: a Chat +/// Completions client translated onto Claude gets `2023-06-01`; a Messages client relayed to +/// Anthropic sends its own version or none, and the provider answers what it sent. +/// claim: SEC-6 +#[tokio::test] +async fn anthropic_version_is_added_only_to_a_translated_walk() { + let (pubkey, sk) = test_keypair(1); + let mock = MockUpstream::start(Mode::AnthropicJson).await; + let gw = Gateway::builder(unused_nats_port(), &mock.authority(), &b64(&pubkey)) + .providers(&["anthropic"]) + .start() + .await; + let auth = || ("authorization", format!("Bearer {}", managed_key(&sk))); + const CLAUDE_CHAT: &str = + r#"{"model":"claude-opus-4-8","messages":[{"role":"user","content":"hi"}]}"#; + assert_eq!( + post(&gw, "/v1/chat/completions", auth(), CLAUDE_CHAT, &[]).await, + 200 + ); + let cap = mock.captured().unwrap(); + assert_eq!(cap.path, "/v1/messages"); + assert_eq!(cap.anthropic_version.as_deref(), Some("2023-06-01")); + + assert_eq!(post(&gw, "/v1/messages", auth(), MESSAGES, &[]).await, 200); + let cap = mock.captured().unwrap(); + assert_eq!(cap.path, "/v1/messages"); + assert_eq!( + cap.anthropic_version, None, + "a relayed Messages request is sent as the client wrote it" + ); +} + /// Send one raw HTTP/1.1 request (so `Connection: Upgrade` goes out exactly as written) and read /// the response. async fn raw_request(port: u16, request: String) -> RawResponse { diff --git a/crates/gateway/tests/reliability_large_body.rs b/crates/gateway/tests/reliability_large_body.rs index 88bad8fa..10da8845 100644 --- a/crates/gateway/tests/reliability_large_body.rs +++ b/crates/gateway/tests/reliability_large_body.rs @@ -9,6 +9,7 @@ mod common; use beyond_ai::key::{VirtualKey, mint}; use common::*; +use std::time::Duration; const MODEL: &str = "gpt-4o-mini"; @@ -304,6 +305,90 @@ async fn a_reset_during_the_upload_fails_over() { ); } +/// The small-body form of the test above: a header-won walk streams a small body through, so a +/// provider that resets after reading only the request head, while the client is still sending +/// the body, cannot have the request. The walk fails over to the next candidate (pingora replays +/// what it buffered and reads the rest), rather than resending to the provider that reset it. +/// claim: REL-1, REL-21 +/// defect: D51 +#[tokio::test] +async fn a_reset_before_a_small_body_is_sent_fails_over() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let (pubkey, sk) = test_keypair(1); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let primary = listener.local_addr().unwrap(); + let resets = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let counter = resets.clone(); + let task = tokio::spawn(async move { + while let Ok((mut s, _)) = listener.accept().await { + let mut seen = Vec::new(); + let mut buf = [0u8; 4096]; + while !seen.windows(4).any(|w| w == b"\r\n\r\n") { + match s.read(&mut buf).await { + Ok(0) | Err(_) => break, + Ok(n) => seen.extend_from_slice(&buf[..n]), + } + } + counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let _ = s.set_zero_linger(); + drop(s); + } + }); + let fallback = MockUpstream::start(Mode::Json).await; + let gw = Gateway::builder(unused_nats_port(), &primary.to_string(), &b64(&pubkey)) + .providers(&["openai", "openrouter"]) + .provider_authority("openrouter", &fallback.authority()) + .start() + .await; + let payload = body(SMALL); + let (first, rest) = payload.as_bytes().split_at(payload.len() / 2); + let mut s = tokio::net::TcpStream::connect(("127.0.0.1", gw.port)) + .await + .unwrap(); + let head = format!( + "POST /auto/chat/completions HTTP/1.1\r\nhost: gw\r\nauthorization: Bearer {}\r\n\ + content-type: application/json\r\nx-beyond-model: {MODEL}\r\n\ + x-beyond-order: openai,openrouter\r\ncontent-length: {}\r\n\r\n", + vkey(&sk), + payload.len() + ); + s.write_all(head.as_bytes()).await.unwrap(); + s.write_all(first).await.unwrap(); + // The primary resets while the rest of the body is still with the client. + tokio::time::timeout(Duration::from_secs(10), async { + while resets.load(std::sync::atomic::Ordering::SeqCst) == 0 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("the request reached the primary"); + tokio::time::sleep(Duration::from_millis(300)).await; + s.write_all(rest).await.unwrap(); + let mut answer = Vec::new(); + let mut buf = [0u8; 4096]; + let _ = tokio::time::timeout(Duration::from_secs(10), async { + while !answer.windows(4).any(|w| w == b"\r\n\r\n") { + match s.read(&mut buf).await { + Ok(0) | Err(_) => break, + Ok(n) => answer.extend_from_slice(&buf[..n]), + } + } + }) + .await; + task.abort(); + let text = String::from_utf8_lossy(&answer); + assert_eq!( + ( + text.starts_with("HTTP/1.1 200"), + resets.load(std::sync::atomic::Ordering::SeqCst), + fallback.hits() + ), + (true, 1, 1), + "{text}\nlog:\n{}", + gw.log() + ); +} + /// A large body's failover to the next candidate still charges the candidate it left: the reset /// is that provider failing, so its breaker hears of it. Only a same-candidate retry (a reused /// connection the provider had closed) gives the permit back without an outcome; a failover that diff --git a/crates/gateway/tests/reliability_lifecycle.rs b/crates/gateway/tests/reliability_lifecycle.rs index b3694f18..ea4d2164 100644 --- a/crates/gateway/tests/reliability_lifecycle.rs +++ b/crates/gateway/tests/reliability_lifecycle.rs @@ -775,7 +775,29 @@ async fn a_large_body_reserves_every_copy_it_will_hold_before_it_is_read() { "{what}: {text}" ); } - // 1.5 MiB on `/openai/…`, buffered for `stream_options`: past the budget even once. + // 1.5 MiB declared on `/openai/…`, buffered for `stream_options`: past the budget even once, + // and refused from the declared length too, before a byte of it is read. + { + let addr = gw.url().trim_start_matches("http://").to_owned(); + let mut s = tokio::net::TcpStream::connect(&addr).await.unwrap(); + let head = format!( + "POST /openai/v1/chat/completions HTTP/1.1\r\nhost: gw\r\nauthorization: Bearer {}\r\n\ + content-type: application/json\r\ncontent-length: {}\r\n\r\n", + vkey(&sk, 35), + 1536 * 1024 + ); + s.write_all(head.as_bytes()).await.unwrap(); + let mut buf = vec![0u8; 4096]; + let n = tokio::time::timeout(Duration::from_secs(5), s.read(&mut buf)) + .await + .expect("provider route: no answer until the body arrives") + .unwrap(); + let text = String::from_utf8_lossy(&buf[..n]); + assert!( + text.starts_with("HTTP/1.1 503") && text.contains("too many large request bodies"), + "provider route, declared: {text}" + ); + } let resp = send("/openai/v1/chat/completions", None, chat(1536 * 1024)) .await .unwrap(); @@ -914,6 +936,33 @@ async fn a_body_declared_at_the_size_cap_is_not_refused_up_front() { } } +/// The size cap is inclusive on the bytes that stream through too: a body of exactly +/// `MAX_REQUEST_BODY` (100 MiB) relayed to a `/{provider}` route is answered, not aborted at its +/// last byte. +/// claim: SEC-19 +#[tokio::test] +async fn a_body_of_exactly_the_size_cap_streams_through() { + let (pubkey, _sk) = test_keypair(1); + let mock = MockUpstream::start(Mode::Json).await; + let gw = Gateway::builder(unused_nats_port(), &mock.authority(), &b64(&pubkey)) + .start() + .await; + const CAP: usize = 100 * 1024 * 1024; + let close = br#""}]}"#; + let mut body = br#"{"model":"gpt-4o-mini","messages":[{"role":"user","content":""#.to_vec(); + body.resize(CAP - close.len(), b'x'); + body.extend_from_slice(close); + let resp = test_client() + .post(format!("{}/openai/v1/chat/completions", gw.url())) + .header("authorization", "Bearer sk-byo-test") + .header("content-type", "application/json") + .body(body) + .send() + .await + .unwrap(); + assert_eq!(resp.status().as_u16(), 200, "{}", gw.log()); +} + /// `client_write_timeout_secs = 0` disables the downstream write timeout rather than making it /// zero: a client that is slow to start reading a large stream still gets every byte of it. /// claim: REL-10 diff --git a/crates/gateway/tests/signed_ids.rs b/crates/gateway/tests/signed_ids.rs index 368fb628..23bd9831 100644 --- a/crates/gateway/tests/signed_ids.rs +++ b/crates/gateway/tests/signed_ids.rs @@ -357,6 +357,33 @@ async fn a_json_response_has_every_id_signed() { ); } +/// Signing changes the answer's length, so the provider's `Content-Length` is dropped and an +/// HTTP/1.1 client gets the body chunked. +/// claim: SEC-25 +#[tokio::test] +async fn a_signed_json_response_is_chunked_to_an_http1_client() { + let (_mock, gw, sk) = gpt_gateway(Mode::Raw(200, "application/json", RESPONSE)).await; + let resp = test_client() + .post(format!("{}/v1/responses", gw.url())) + .header("authorization", format!("Bearer {}", billing_vkey(&sk, A))) + .header("content-type", "application/json") + .body(r#"{"model":"gpt-4o","input":"hi"}"#) + .send() + .await + .unwrap(); + assert_eq!(resp.status().as_u16(), 200); + assert_eq!(resp.version(), reqwest::Version::HTTP_11); + assert!(resp.headers().get("content-length").is_none()); + assert_eq!( + resp.headers() + .get("transfer-encoding") + .and_then(|v| v.to_str().ok()), + Some("chunked") + ); + let text = resp.text().await.unwrap(); + assert!(!text.contains(RESP), "{text}"); +} + /// A streamed answer signs every id on every event that carries one, and the same item is the same /// signed id on all of them (`item.id` on `output_item.*`, `item_id` on each content event, the /// output on `completed`), so a client that keys items by id still joins them. diff --git a/crates/gateway/tests/tenant_concurrency.rs b/crates/gateway/tests/tenant_concurrency.rs index ec32bbde..9344e3e1 100644 --- a/crates/gateway/tests/tenant_concurrency.rs +++ b/crates/gateway/tests/tenant_concurrency.rs @@ -157,3 +157,39 @@ async fn a_cancelled_request_releases_its_slot() { "three cancellations left the tenant with a slot" ); } + +/// The cap bounds managed spend only: a BYO request is the caller's own key and holds no slot, so +/// BYO requests in flight together are each served, however low the cap. +/// claim: A3 +#[tokio::test] +async fn byo_requests_take_no_tenant_slot() { + let (pubkey, _sk) = test_keypair(54); + let mock = MockUpstream::start(Mode::Slow(1_000)).await; + let gw = Gateway::builder(unused_nats_port(), &mock.authority(), &b64(&pubkey)) + .tenant_max_in_flight(1) + .start() + .await; + let client = test_client(); + let held = { + let (client, url) = (client.clone(), gw.url()); + tokio::spawn(async move { + client + .post(format!("{url}/openai/v1/chat/completions")) + .header("authorization", "Bearer sk-byo-test") + .header("content-type", "application/json") + .body(body()) + .send() + .await + .unwrap() + .status() + .as_u16() + }) + }; + wait_for_hits(&mock, 1).await; + assert_eq!( + post(&client, &gw, "sk-byo-test").await.status().as_u16(), + 200, + "a second BYO request in flight is served" + ); + assert_eq!(held.await.unwrap(), 200); +} diff --git a/crates/gateway/tests/translate.rs b/crates/gateway/tests/translate.rs index 20238406..9a7bd166 100644 --- a/crates/gateway/tests/translate.rs +++ b/crates/gateway/tests/translate.rs @@ -59,6 +59,16 @@ async fn openai_sdk_nonstream_claude_round_trips_json() { .await .unwrap(); assert_eq!(resp.status().as_u16(), 200); + // Translation changes the length: the provider's `Content-Length` is dropped and an HTTP/1.1 + // client gets the body chunked. + assert_eq!(resp.version(), reqwest::Version::HTTP_11); + assert!(resp.headers().get("content-length").is_none()); + assert_eq!( + resp.headers() + .get("transfer-encoding") + .and_then(|v| v.to_str().ok()), + Some("chunked") + ); let text = resp.text().await.unwrap(); let v: Value = serde_json::from_str(&text).expect(&text); assert_eq!(v["object"], "chat.completion"); diff --git a/verify/defects.toml b/verify/defects.toml index 54772a88..cf452788 100644 --- a/verify/defects.toml +++ b/verify/defects.toml @@ -2567,5 +2567,19 @@ evidence = "Hermetic: a_402_the_walk_fails_over_on_cools_an_unfunded_key. With x claims = ["REL-4"] state = "fixed" test = "a_402_the_walk_fails_over_on_cools_an_unfunded_key" -fix = "3bd64c4" +fix = "a7c753a" note = "Reading the refused 402's body before abandoning it, as first proposed, is not possible: pingora 0.9's upstream_response_filter gets the head only, the upstream session is not reachable from any hook, and an Err there drops the connection with the body unread (the same reason D180 could not fail over the request that saw the answer). So the status decides, where it can: remedy::unfunded_402 says a 402 is out of credit by its status alone on every provider but OpenRouter, whose 402 may be one request larger than the balance (\"This request requires more credits\", which must not cool the shared key: D180, D200). upstream_response_filter cools the key (AiProxy::cool_unfunded_key: same metric reason, same warn line as logging's D180 path) on both abandoning paths, the in-place retry and a recorded FullBody retry; a FullBody final attempt is relayed and cooled from its body as before. No new state (Provider::mark_key_bad, the existing key cooldown) and no work on any other status, the 2xx path included. logging no longer cools from the body of an abandoned FullBody attempt (relay_abandoned), which raced the parent's drop and would count a head-cooled 402 twice; this also makes a FullBody key walk's abandoned bodies deterministically unread, as ARCHITECTURE.md already said. Limit: an OpenRouter 402 the walk fails over on still cools nothing, so a row that leads with OpenRouter (or an x-beyond-order that puts it first) keeps a round trip to an unfunded OpenRouter account; by default the aggregator is last in the walk (D253), where its 402 is relayed and cooled from the body. Also tested: remedy::tests::only_openrouter_needs_the_body_of_a_402; the same test asserts the OpenRouter too-costly 402 fails over and cools nothing (OpenRouter is sent all three requests) for both body sizes." + +[[defect]] +id = "D259" +severity = "low" +area = "Billing" +title = "A client that closed mid-upload was billed an input estimate for a request the provider never had whole" +impact = "BIL-3 bills a wait given up on only once the provider has the whole request (D130). A client that closed its connection before the end of its own body (a Content-Length body with bytes still to come, a chunked body without its terminating chunk) left the provider half a request, which it never processes or bills, yet the ai.usage row was usage_estimated: true with the half body's input estimate (outcome client_cancelled)." +confidence = "high" +evidence = "pingora-core 0.9 http/v1/body.rs: on EOF mid-body the reader sets body_state to Done and returns ConnectionClosed ('Peer prematurely closed connection with N bytes of body remaining to read', or 'Connection prematurely closed without the termination chunk'); body_done() is true for Done as well as Complete. proxy.rs body_delivered read session.is_body_done(), so logging's no_head (managed, no head, gave_up_waiting: any Downstream error) estimated the prompt. Hermetic: half a body then FIN against a silent upstream -> row {usage_estimated: true, input_tokens: 1, outcome: client_cancelled}. Found triaging the error_while_proxy Downstream mutant (cargo-mutants full run)." +claims = ["BIL-3", "BIL-20"] +state = "fixed" +test = "a_client_that_closes_mid_upload_is_not_billed_an_estimate" +fix = "28bc659" +note = "Fixed: body_delivered also requires that the error is not a downstream premature close (client_body_cut_short, matched on pingora's context strings through the cause chain, as h2_body_unsent matches its own; pingora's session API exposes no body_complete). Covers both framings; error_while_proxy's Downstream early return already kept such a request from being resent. A client that sends its whole body and then gives up is still billed (D130). HTTP/2 downstream resets are a different error and were not affected."