diff --git a/crates/gateway/src/proxy.rs b/crates/gateway/src/proxy.rs index cc9f9cae..90b27fbf 100644 --- a/crates/gateway/src/proxy.rs +++ b/crates/gateway/src/proxy.rs @@ -1939,7 +1939,8 @@ impl AiProxy { /// Never exempt on a `FullBody` re-run, whose parent reserved two copies of the body, not the /// `Value`s built from it. A translation that needs more than the whole budget is a 413 naming /// translation; one that does not fit now is the budget's retryable 503. With the budget off - /// there is nothing to hold, and nothing is refused. + /// there is nothing to hold, and nothing is refused. On a response the status line is already + /// downstream, so either refusal aborts it. fn hold_translation_heap( &self, body: &[u8], @@ -5748,19 +5749,14 @@ impl ProxyHttp for AiProxy { a.route.card.output_cap().unwrap_or(0), ); let serving = catalog_serving_endpoint(a.as_ref()); - let upstream_model = - a.candidate_at(a.candidate).map_or("", |c| c.upstream_model); - let openai_host = a - .candidate_at(a.candidate) - .is_some_and(|c| c.provider == providers::ProviderId::OpenAi); - let stream_only = a - .candidate_at(a.candidate) - .is_some_and(providers::catalog::stream_only); - let reads_developer = a - .candidate_at(a.candidate) - .is_none_or(providers::catalog::reads_developer_role); - let tool_thinking = a - .candidate_at(a.candidate) + let candidate = a.candidate_at(a.candidate); + let upstream_model = candidate.map_or("", |c| c.upstream_model); + let openai_host = + candidate.is_some_and(|c| c.provider == providers::ProviderId::OpenAi); + let stream_only = candidate.is_some_and(providers::catalog::stream_only); + let reads_developer = + candidate.is_none_or(providers::catalog::reads_developer_role); + let tool_thinking = candidate .map_or(providers::catalog::ToolThinking::Free, |c| { providers::catalog::tool_thinking(c) }); @@ -5835,7 +5831,7 @@ impl ProxyHttp for AiProxy { } // Native OpenAI Chat takes `max_completion_tokens` on every model and 400s // `max_tokens` on its reasoning ones; a translated body already says the former. - if a.candidate_at(a.candidate).is_some_and(native_openai_chat) + if candidate.is_some_and(native_openai_chat) && rename_max_tokens(&mut buf, &scan.limit_keys) { scan = peek::scan_buffered(&buf); @@ -6241,6 +6237,20 @@ impl ProxyHttp for AiProxy { return Err(self.translate_overflow(&rc.request_id, "json_body")); } if end_of_stream { + // The response's `Value`s are held in the body budget like a request's + // (D216). Headers are already downstream, so a refusal aborts the response, + // as `translate_overflow` does. A same-wire 2xx is relayed as bytes. + let _heap = if upstream == t.client && (200..300).contains(&status) { + None + } else { + self.hold_translation_heap(&t.json_buf).inspect_err(|_| { + warn!( + request_id = %rc.request_id, + bytes = t.json_buf.len(), + "translated response exceeds the body budget; aborting", + ); + })? + }; translate::response_json_tools( upstream, t.client, diff --git a/crates/gateway/tests/body_memory.rs b/crates/gateway/tests/body_memory.rs index 1eca500a..8ee457b9 100644 --- a/crates/gateway/tests/body_memory.rs +++ b/crates/gateway/tests/body_memory.rs @@ -156,6 +156,48 @@ async fn a_body_too_costly_to_translate_is_a_413_naming_translation() { ); } +/// The response half of D216: a non-streaming Claude answer whose tool input is 4 MiB of tiny +/// objects, translated for a Chat client, takes the same ~1 GiB of `Value`s. Its heap is held in the +/// budget like a request's, and it does not fit: the status line is already downstream, so the +/// response is aborted rather than translated, and the gateway's memory stays where it was. +/// claim: SEC-19 +/// defect: D216 +#[tokio::test] +async fn a_response_too_costly_to_translate_is_aborted() { + let answer = format!( + r#"{{"id":"msg_1","type":"message","role":"assistant","model":"claude-opus-4-8","content":[{{"type":"tool_use","id":"toolu_1","name":"f","input":{{"x":[{}]}}}}],"stop_reason":"tool_use","usage":{{"input_tokens":1,"output_tokens":1}}}}"#, + tiny_objects(4 * MIB) + ); + let len = answer.len(); + let answer: &'static str = Box::leak(answer.into_boxed_str()); + let (pubkey, sk) = test_keypair(1); + let mock = MockUpstream::start(Mode::Raw(200, "application/json", answer)).await; + let gw = Gateway::builder(unused_nats_port(), &mock.authority(), &b64(&pubkey)) + .providers(&["anthropic"]) + .start() + .await; + let (_, text, rise) = send( + &gw, + &sk, + "/v1/chat/completions", + r#"{"model":"claude-opus-4-8","max_tokens":16,"messages":[{"role":"user","content":"go"}]}"# + .to_owned(), + ) + .await; + assert_eq!(mock.hits(), 1); + assert!( + !text.contains("tool_calls"), + "the response was translated ({} bytes)", + text.len() + ); + assert!( + rise <= copies_budget(len), + "peak RSS rose {rise} KiB for a {} KiB response (budget {} KiB)", + len / 1024, + copies_budget(len) + ); +} + /// A translated stream that outgrows every answer the catalog allows (the largest `max_output`, /// 384,000 tokens, at 128 bytes a token: `translate::MAX_STREAM_OUTPUT`, 46.9 MiB) is cut, rather /// than gathered whole: a Responses client's bridge keeps every text delta for the closing