diff --git a/codex-rs/codex-api/src/api_bridge.rs b/codex-rs/codex-api/src/api_bridge.rs index 70da94e3fb..c68a7071a4 100644 --- a/codex-rs/codex-api/src/api_bridge.rs +++ b/codex-rs/codex-api/src/api_bridge.rs @@ -210,6 +210,9 @@ pub fn map_api_error(err: ApiError) -> CodexErr { CodexErr::ConnectionFailed(ConnectionFailedError { source }) } TransportError::Network(msg) | TransportError::Build(msg) => CodexErr::Stream(msg), + error @ TransportError::ResponseTooLarge { .. } => { + CodexErr::InvalidRequest(error.to_string()) + } }, ApiError::RateLimit(msg) => CodexErr::Stream(msg), } diff --git a/codex-rs/codex-api/src/endpoint/models.rs b/codex-rs/codex-api/src/endpoint/models.rs index ecd50a0f1d..9e504f2af1 100644 --- a/codex-rs/codex-api/src/endpoint/models.rs +++ b/codex-rs/codex-api/src/endpoint/models.rs @@ -2,6 +2,7 @@ use crate::auth::SharedAuthProvider; use crate::endpoint::session::EndpointSession; use crate::error::ApiError; use crate::provider::Provider; +use bytes::Bytes; use codex_client::HttpTransport; use codex_client::RequestTelemetry; use codex_protocol::openai_models::ModelInfo; @@ -48,6 +49,34 @@ impl ModelsClient { request_url: String, extra_headers: HeaderMap, ) -> Result<(Vec, Option), ApiError> { + let (body, header_etag) = self + .list_models_raw( + request_url, + extra_headers, + /*response_body_limit_bytes*/ None, + ) + .await?; + let ModelsResponse { models } = + serde_json::from_slice::(&body).map_err(|e| { + ApiError::Stream(format!( + "failed to decode models response: {e}; body: {}", + String::from_utf8_lossy(&body) + )) + })?; + + Ok((models, header_etag)) + } + + /// Fetches a catalog using the provider's auth and retry policy without decoding it. + /// + /// Callers accepting provider-controlled catalogs must set a response-body limit + /// and validate the native response without including raw values in diagnostics. + pub async fn list_models_raw( + &self, + request_url: String, + extra_headers: HeaderMap, + response_body_limit_bytes: Option, + ) -> Result<(Bytes, Option), ApiError> { let resp = self .session .execute_with( @@ -57,6 +86,7 @@ impl ModelsClient { /*body*/ None, move |req| { req.url.clone_from(&request_url); + req.response_body_limit_bytes = response_body_limit_bytes; }, ) .await?; @@ -67,22 +97,16 @@ impl ModelsClient { .and_then(|value| value.to_str().ok()) .map(ToString::to_string); - let ModelsResponse { models } = serde_json::from_slice::(&resp.body) - .map_err(|e| { - ApiError::Stream(format!( - "failed to decode models response: {e}; body: {}", - String::from_utf8_lossy(&resp.body) - )) - })?; - - Ok((models, header_etag)) + Ok((resp.body, header_etag)) } } #[cfg(test)] mod tests { use super::*; + use crate::auth::AuthError; use crate::auth::AuthProvider; + use crate::auth::AuthProviderFuture; use crate::provider::RetryConfig; use codex_client::Request; use codex_client::Response; @@ -140,6 +164,29 @@ mod tests { fn add_auth_headers(&self, _headers: &mut HeaderMap) {} } + #[derive(Default)] + struct RetryAuth { + request_limits: Mutex>>, + } + + impl AuthProvider for RetryAuth { + fn add_auth_headers(&self, headers: &mut HeaderMap) { + headers.insert(http::header::AUTHORIZATION, "Bearer test".parse().unwrap()); + } + + fn apply_auth(&self, mut request: Request) -> AuthProviderFuture<'_> { + Box::pin(async move { + let mut limits = self.request_limits.lock().unwrap(); + limits.push(request.response_body_limit_bytes); + if limits.len() == 1 { + return Err(AuthError::Transient("retry authentication".to_string())); + } + self.add_auth_headers(&mut request.headers); + Ok(request) + }) + } + } + fn provider(base_url: &str) -> Provider { Provider { name: "test".to_string(), @@ -157,6 +204,52 @@ mod tests { } } + #[tokio::test] + async fn response_body_limit_survives_auth_retry_without_limiting_other_models_requests() { + let transport = CapturingTransport::default(); + let auth = Arc::new(RetryAuth::default()); + let mut provider = provider("https://example.com/api/codex"); + provider.retry.max_attempts = 2; + let request_url = ModelsClient::::request_url(&provider, "0.99.0"); + let client = ModelsClient::new(transport.clone(), provider, auth.clone()); + + client + .list_models_raw( + request_url.clone(), + HeaderMap::new(), + /*response_body_limit_bytes*/ Some(64), + ) + .await + .expect("limited request should succeed after auth retry"); + { + let request = transport.last_request.lock().unwrap(); + let request = request.as_ref().unwrap(); + assert_eq!(request.response_body_limit_bytes, Some(64)); + assert_eq!(request.url, request_url); + assert_eq!(request.headers[http::header::AUTHORIZATION], "Bearer test"); + } + + let (models, _) = client + .list_models(request_url, HeaderMap::new()) + .await + .expect("ordinary request on the same client should remain unbounded"); + assert!(models.is_empty()); + assert_eq!( + transport + .last_request + .lock() + .unwrap() + .as_ref() + .unwrap() + .response_body_limit_bytes, + None + ); + assert_eq!( + *auth.request_limits.lock().unwrap(), + vec![Some(64), Some(64), None] + ); + } + #[tokio::test] async fn appends_client_version_query() { let response = ModelsResponse { models: Vec::new() }; diff --git a/codex-rs/codex-api/src/provider.rs b/codex-rs/codex-api/src/provider.rs index 849513af55..ea4363e5c6 100644 --- a/codex-rs/codex-api/src/provider.rs +++ b/codex-rs/codex-api/src/provider.rs @@ -82,6 +82,7 @@ impl Provider { body: None, compression: RequestCompression::None, timeout: None, + response_body_limit_bytes: None, } } diff --git a/codex-rs/codex-client/src/retry.rs b/codex-rs/codex-client/src/retry.rs index f78aa2e1e3..86071851f6 100644 --- a/codex-rs/codex-client/src/retry.rs +++ b/codex-rs/codex-client/src/retry.rs @@ -31,7 +31,9 @@ impl RetryOn { TransportError::Timeout | TransportError::Connection(_) | TransportError::Network(_) => self.retry_transport, - _ => false, + TransportError::Build(_) + | TransportError::RetryLimit + | TransportError::ResponseTooLarge { .. } => false, } } } diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/metrics.rs b/codex-rs/ext/guardian-v2/src/async_scorer/metrics.rs index 138dbb08ee..4ecc681ffe 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/metrics.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/metrics.rs @@ -41,6 +41,7 @@ pub(super) fn sampler_failure_reason(error: &LunaSamplerError) -> &'static str { ApiError::Transport(TransportError::Network(_)) => "network_error", ApiError::Transport(TransportError::RetryLimit) => "retry_limit", ApiError::Transport(TransportError::Build(_)) => "request_build_error", + ApiError::Transport(TransportError::ResponseTooLarge { .. }) => "response_too_large", ApiError::Stream(_) => "stream_error", ApiError::ContextWindowExceeded => "context_window_exceeded", ApiError::QuotaExceeded => "quota_exceeded", diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs index cd9d5451d4..c9c69bb95c 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs @@ -82,7 +82,9 @@ impl SamplingExecution { | LunaSamplerError::IncompatibleCompaction | LunaSamplerError::InputTooLarge | LunaSamplerError::Api( - ApiError::Transport(TransportError::Build(_)) + ApiError::Transport( + TransportError::Build(_) | TransportError::ResponseTooLarge { .. }, + ) | ApiError::ContextWindowExceeded | ApiError::QuotaExceeded | ApiError::UsageNotIncluded diff --git a/codex-rs/http-client/src/error.rs b/codex-rs/http-client/src/error.rs index e3a4f8aca6..43ace4f7a7 100644 --- a/codex-rs/http-client/src/error.rs +++ b/codex-rs/http-client/src/error.rs @@ -24,6 +24,8 @@ pub enum TransportError { Network(String), #[error("request build error: {0}")] Build(String), + #[error("response body exceeds the {max_bytes} byte limit")] + ResponseTooLarge { max_bytes: usize }, } #[derive(Debug, Error)] diff --git a/codex-rs/http-client/src/request.rs b/codex-rs/http-client/src/request.rs index 7608860c3d..b21ce6280e 100644 --- a/codex-rs/http-client/src/request.rs +++ b/codex-rs/http-client/src/request.rs @@ -81,6 +81,9 @@ pub struct Request { pub body: Option, pub compression: RequestCompression, pub timeout: Option, + /// Maximum accepted response-body bytes, including streaming and error bodies. + /// `None` leaves the response unbounded. + pub response_body_limit_bytes: Option, } impl Request { @@ -92,6 +95,7 @@ impl Request { body: None, compression: RequestCompression::None, timeout: None, + response_body_limit_bytes: None, } } diff --git a/codex-rs/http-client/src/transport.rs b/codex-rs/http-client/src/transport.rs index 337b1e8675..ca351c1f00 100644 --- a/codex-rs/http-client/src/transport.rs +++ b/codex-rs/http-client/src/transport.rs @@ -1,3 +1,8 @@ +//! Shared HTTP transport with optional limits on complete response bodies. +//! +//! Each request's limit applies to observed body bytes, including unsuccessful responses; +//! declared content lengths are only an additional early-rejection check. + use crate::client::HttpClient; use crate::client::RequestBuilder; use crate::error::TransportError; @@ -5,6 +10,7 @@ use crate::request::Request; use crate::request::RequestBody; use crate::request::Response; use bytes::Bytes; +use bytes::BytesMut; use futures::StreamExt; use futures::stream::BoxStream; use http::HeaderMap; @@ -59,6 +65,7 @@ impl ReqwestTransport { body: _, compression: _, timeout, + response_body_limit_bytes: _, } = req; let mut builder = self.client.request( @@ -115,13 +122,23 @@ impl HttpTransport for ReqwestTransport { self.trace_request(&req); let url = req.url.clone(); + let response_body_limit_bytes = req.response_body_limit_bytes; let builder = self.build(req)?; let resp = builder.send().await.map_err(Self::map_error)?; let status = resp.status(); let headers = resp.headers().clone(); - let bytes = resp.bytes().await.map_err(Self::map_error)?; + let bytes = match response_body_limit_bytes { + Some(max_bytes) => bounded_response_bytes(resp, max_bytes).await, + None => resp.bytes().await.map_err(Self::map_error), + }; if !status.is_success() { - let body = String::from_utf8(bytes.to_vec()).ok(); + let body = match bytes { + Ok(bytes) => String::from_utf8(bytes.to_vec()).ok(), + Err(error @ TransportError::ResponseTooLarge { .. }) => return Err(error), + // Keep bounded diagnostic-body failures from hiding HTTP auth/retry status. + Err(_) if response_body_limit_bytes.is_some() => None, + Err(error) => return Err(error), + }; return Err(TransportError::Http { status, url: Some(url), @@ -132,7 +149,7 @@ impl HttpTransport for ReqwestTransport { Ok(Response { status, headers, - body: bytes, + body: bytes?, }) } @@ -140,12 +157,27 @@ impl HttpTransport for ReqwestTransport { self.trace_request(&req); let url = req.url.clone(); + let response_body_limit_bytes = req.response_body_limit_bytes; let builder = self.build(req)?; let resp = builder.send().await.map_err(Self::map_error)?; let status = resp.status(); let headers = resp.headers().clone(); if !status.is_success() { - let body = resp.text().await.ok(); + let body = match response_body_limit_bytes { + Some(max_bytes) => match bounded_response_bytes(resp, max_bytes).await { + Ok(bytes) => { + // Reuse the unbounded path's charset, BOM and replacement decoding, + // but only after the body has passed the byte limit. + let mut buffered = http::Response::new(bytes); + *buffered.headers_mut() = headers.clone(); + reqwest::Response::from(buffered).text().await.ok() + } + Err(error @ TransportError::ResponseTooLarge { .. }) => return Err(error), + // A failed diagnostic body must not hide HTTP auth or retry semantics. + Err(_) => None, + }, + None => resp.text().await.ok(), + }; return Err(TransportError::Http { status, url: Some(url), @@ -153,17 +185,68 @@ impl HttpTransport for ReqwestTransport { body, }); } - let stream = resp - .bytes_stream() - .map(|result| result.map_err(Self::map_error)); + let bytes = match response_body_limit_bytes { + Some(max_bytes) => bounded_response_stream(resp, max_bytes)?, + None => Box::pin( + resp.bytes_stream() + .map(|result| result.map_err(Self::map_error)), + ), + }; Ok(StreamResponse { status, headers, - bytes: Box::pin(stream), + bytes, }) } } +fn bounded_response_stream( + response: reqwest::Response, + max_bytes: usize, +) -> Result { + if response + .content_length() + .is_some_and(|length| length > max_bytes as u64) + { + return Err(TransportError::ResponseTooLarge { max_bytes }); + } + + let stream = futures::stream::try_unfold( + (response, max_bytes), + move |(mut response, remaining)| async move { + match response + .chunk() + .await + .map_err(|error| ReqwestTransport::map_error(error.without_url()))? + { + Some(chunk) if chunk.len() <= remaining => { + let remaining = remaining - chunk.len(); + Ok(Some((chunk, (response, remaining)))) + } + Some(_) => Err(TransportError::ResponseTooLarge { max_bytes }), + None => Ok(None), + } + }, + ); + Ok(Box::pin(stream)) +} + +async fn bounded_response_bytes( + response: reqwest::Response, + max_bytes: usize, +) -> Result { + let mut stream = bounded_response_stream(response, max_bytes)?; + let mut body = BytesMut::new(); + while let Some(chunk) = stream.next().await { + body.extend_from_slice(&chunk?); + } + Ok(body.freeze()) +} + #[cfg(test)] #[path = "transport_tests.rs"] mod tests; + +#[cfg(test)] +#[path = "transport_limit_tests.rs"] +mod limit_tests; diff --git a/codex-rs/http-client/src/transport_limit_tests.rs b/codex-rs/http-client/src/transport_limit_tests.rs new file mode 100644 index 0000000000..1c977854a1 --- /dev/null +++ b/codex-rs/http-client/src/transport_limit_tests.rs @@ -0,0 +1,452 @@ +//! HTTP fixtures exercise opt-in body limits for both buffered and streaming callers. + +use super::*; +use pretty_assertions::assert_eq; +use std::io::Read; +use std::io::Write; +use std::net::TcpListener; +use std::net::TcpStream; +use std::time::Duration; +use std::time::Instant; + +#[derive(Clone, Copy, Debug)] +enum Delivery { + Buffered, + Streamed, +} + +#[tokio::test] +async fn accepts_complete_bodies_at_the_limit() { + for delivery in [Delivery::Buffered, Delivery::Streamed] { + for (framing, wire_body, max_bytes, expected_body) in [ + ("Content-Length: 4\r\n", "test", 4, "test"), + ("", "test", 4, "test"), + ( + "Transfer-Encoding: chunked\r\n", + "2\r\nte\r\n2\r\nst\r\n0\r\n\r\n", + 4, + "test", + ), + ("Content-Length: 0\r\n", "", 0, ""), + ] { + let response = receive_response( + format!("HTTP/1.1 200 OK\r\n{framing}Connection: close\r\n\r\n{wire_body}"), + &transport(), + Some(max_bytes), + delivery, + ) + .await + .expect("a complete body within the limit should succeed"); + let mut expected_headers = HeaderMap::new(); + expected_headers.insert("connection", "close".parse().unwrap()); + for line in framing.lines() { + let (name, value) = line.split_once(':').unwrap(); + expected_headers.insert( + http::HeaderName::from_bytes(name.as_bytes()).unwrap(), + value.trim().parse().unwrap(), + ); + } + assert_eq!( + (response.status, response.headers, response.body), + ( + StatusCode::OK, + expected_headers, + Bytes::copy_from_slice(expected_body.as_bytes()), + ), + "{delivery:?} with framing {framing:?}", + ); + } + } +} + +#[tokio::test] +async fn rejects_oversized_success_and_error_bodies_without_sensitive_diagnostics() { + for delivery in [Delivery::Buffered, Delivery::Streamed] { + for status in ["200 OK", "500 Internal Server Error"] { + for (framing, wire_body) in [ + ("Content-Length: 11\r\n", "body-secret"), + ("", "body-secret"), + ( + "Transfer-Encoding: chunked\r\n", + "4\r\nbody\r\n7\r\n-secret\r\n0\r\n\r\n", + ), + ( + "Transfer-Encoding: chunked\r\nContent-Length: 1\r\n", + "4\r\nbody\r\n7\r\n-secret\r\n0\r\n\r\n", + ), + ] { + let error = receive_response( + format!( + "HTTP/1.1 {status}\r\n{framing}X-Secret: header-secret\r\nConnection: close\r\n\r\n{wire_body}" + ), + &transport(), + Some(4), + delivery, + ) + .await + .expect_err("the observed body must not exceed the limit"); + assert_eq!( + (error.to_string(), format!("{error:?}")), + ( + "response body exceeds the 4 byte limit".to_string(), + "ResponseTooLarge { max_bytes: 4 }".to_string(), + ), + "{delivery:?}, {status}, framing {framing:?}", + ); + } + } + } +} + +#[tokio::test] +async fn rejects_oversized_declared_lengths_before_reading_the_body() { + for delivery in [Delivery::Buffered, Delivery::Streamed] { + for status in [ + StatusCode::OK, + StatusCode::UNAUTHORIZED, + StatusCode::TOO_MANY_REQUESTS, + ] { + let error = receive_response( + format!( + "HTTP/1.1 {status}\r\nContent-Length: 1000000\r\nConnection: close\r\n\r\n" + ), + &transport(), + Some(4), + delivery, + ) + .await + .expect_err("an oversized declaration should fail before the incomplete body"); + assert!(matches!( + error, + TransportError::ResponseTooLarge { max_bytes: 4 } + )); + } + } +} + +#[tokio::test] +async fn streaming_limit_counts_across_chunks_and_terminates_on_error() { + let (continue_tx, continue_rx) = std::sync::mpsc::channel(); + let (url, server) = start_server(move |connection| { + connection + .write_all(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n4\r\npart\r\n") + .expect("write the first chunk"); + continue_rx + .recv_timeout(Duration::from_secs(5)) + .expect("client should consume the first chunk"); + connection + .write_all(b"4\r\nnext\r\n0\r\n\r\n") + .expect("write the second chunk"); + }); + let mut request = Request::new(Method::GET, url); + request.response_body_limit_bytes = Some(6); + let mut response = transport() + .stream(request) + .await + .expect("stream should open before the body limit is reached"); + assert_eq!( + response.bytes.next().await.unwrap().unwrap(), + Bytes::from_static(b"part"), + ); + continue_tx.send(()).expect("release the second chunk"); + let error = response + .bytes + .next() + .await + .expect("second chunk should produce the limit error") + .expect_err("two individually small chunks should exceed the total limit"); + assert!(matches!( + error, + TransportError::ResponseTooLarge { max_bytes: 6 } + )); + assert!(response.bytes.next().await.is_none()); + server.join().expect("HTTP fixture should finish"); +} + +#[tokio::test] +async fn preserves_http_failures_within_the_limit() { + for delivery in [Delivery::Buffered, Delivery::Streamed] { + let error = receive_response( + "HTTP/1.1 403 Forbidden\r\nContent-Length: 6\r\nConnection: close\r\n\r\ndenied" + .to_string(), + &transport(), + Some(6), + delivery, + ) + .await + .expect_err("a small error body should retain HTTP error semantics"); + let TransportError::Http { + status, + headers, + body, + .. + } = error + else { + panic!("expected HTTP error, got {error:?}"); + }; + assert_eq!( + (status, headers, body), + ( + StatusCode::FORBIDDEN, + Some(HeaderMap::from_iter([ + (http::header::CONTENT_LENGTH, "6".parse().unwrap()), + (http::header::CONNECTION, "close".parse().unwrap()), + ])), + Some("denied".to_string()), + ), + ); + } +} + +#[tokio::test] +async fn bounded_errors_preserve_status_when_the_body_is_interrupted() { + for delivery in [Delivery::Buffered, Delivery::Streamed] { + for status in [StatusCode::UNAUTHORIZED, StatusCode::TOO_MANY_REQUESTS] { + for max_bytes in [None, Some(5), Some(4), Some(3)] { + // The first chunk is complete, but the peer closes before the next chunk. + let error = receive_response( + format!( + "HTTP/1.1 {status}\r\nTransfer-Encoding: chunked\r\nRetry-After: 7\r\nWWW-Authenticate: Bearer\r\nConnection: close\r\n\r\n4\r\npart\r\n" + ), + &transport(), + max_bytes, + delivery, + ) + .await + .expect_err("an interrupted error body must still fail"); + if matches!(delivery, Delivery::Buffered) && max_bytes.is_none() { + assert!(matches!(error, TransportError::Network(_))); + continue; + } + if max_bytes == Some(3) { + assert!(matches!( + error, + TransportError::ResponseTooLarge { max_bytes: 3 } + )); + continue; + } + let TransportError::Http { + status: actual_status, + headers, + body, + .. + } = error + else { + panic!("expected preserved HTTP error, got {error:?}"); + }; + assert_eq!( + (actual_status, headers, body), + ( + status, + Some(HeaderMap::from_iter([ + (http::header::TRANSFER_ENCODING, "chunked".parse().unwrap()), + (http::header::RETRY_AFTER, "7".parse().unwrap()), + (http::header::WWW_AUTHENTICATE, "Bearer".parse().unwrap()), + (http::header::CONNECTION, "close".parse().unwrap()), + ])), + None, + ), + "{delivery:?}, max_bytes={max_bytes:?}", + ); + } + } + } +} + +#[tokio::test] +async fn bounded_successes_report_interrupted_bodies_as_network_errors() { + for delivery in [Delivery::Buffered, Delivery::Streamed] { + let error = receive_response( + "HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n4\r\npart\r\n" + .to_string(), + &transport(), + Some(4), + delivery, + ) + .await + .expect_err("a successful status cannot hide an incomplete response body"); + assert!(matches!(error, TransportError::Network(_))); + } +} + +#[tokio::test] +async fn bounded_streamed_errors_preserve_text_decoding() { + for (content_type, bytes, expected) in [ + ( + "text/plain; charset=windows-1252", + [0x63, 0x61, 0x66, 0xe9].as_slice(), + "café", + ), + ( + "text/plain; charset=utf-8", + b"error: \xff".as_slice(), + "error: �", + ), + ( + "text/plain; charset=unknown", + b"error: \xff".as_slice(), + "error: �", + ), + ("text/plain", b"\xef\xbb\xbfdenied".as_slice(), "denied"), + ] { + for bounded in [false, true] { + let mut wire = format!( + "HTTP/1.1 403 Forbidden\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + bytes.len() + ) + .into_bytes(); + wire.extend_from_slice(bytes); + let error = receive_response( + wire, + &transport(), + bounded.then_some(bytes.len()), + Delivery::Streamed, + ) + .await + .expect_err("retain the HTTP error and decoded diagnostics"); + let TransportError::Http { + status, + headers, + body, + .. + } = error + else { + panic!("expected HTTP error, got {error:?}"); + }; + assert_eq!( + (status, headers, body), + ( + StatusCode::FORBIDDEN, + Some(HeaderMap::from_iter([ + (http::header::CONTENT_TYPE, content_type.parse().unwrap()), + ( + http::header::CONTENT_LENGTH, + bytes.len().to_string().parse().unwrap() + ), + (http::header::CONNECTION, "close".parse().unwrap()), + ])), + Some(expected.to_string()), + ), + "{content_type}, bounded={bounded}", + ); + } + } +} + +#[tokio::test] +async fn shared_transport_keeps_each_requests_response_limit_independent() { + let transport = transport(); + for delivery in [Delivery::Buffered, Delivery::Streamed] { + for max_bytes in [Some(4), None, Some(11), Some(0), None] { + let response = receive_response( + "HTTP/1.1 200 OK\r\nContent-Length: 11\r\nConnection: close\r\n\r\nbody-secret" + .to_string(), + &transport, + max_bytes, + delivery, + ) + .await; + match max_bytes { + Some(max_bytes) if max_bytes < 11 => assert!(matches!( + response, + Err(TransportError::ResponseTooLarge { max_bytes: actual }) if actual == max_bytes + )), + _ => assert_eq!( + response.expect("only this request's limit applies").body, + Bytes::from_static(b"body-secret"), + ), + } + } + } +} + +fn transport() -> ReqwestTransport { + ReqwestTransport::new( + reqwest::Client::builder() + .no_proxy() + .build() + .expect("test client should build"), + ) +} + +async fn receive_response( + wire_response: impl Into>, + transport: &ReqwestTransport, + response_body_limit_bytes: Option, + delivery: Delivery, +) -> Result { + let wire_response = wire_response.into(); + let (url, server) = start_server(move |connection| { + connection + .write_all(&wire_response) + .expect("write HTTP response"); + }); + let mut request = Request::new(Method::GET, url); + request.headers.insert( + http::header::AUTHORIZATION, + "Bearer auth-secret".parse().unwrap(), + ); + request.timeout = Some(Duration::from_secs(5)); + request.response_body_limit_bytes = response_body_limit_bytes; + let result = match delivery { + Delivery::Buffered => transport.execute(request).await, + Delivery::Streamed => { + async { + let mut response = transport.stream(request).await?; + let mut body = BytesMut::new(); + while let Some(chunk) = response.bytes.next().await { + body.extend_from_slice(&chunk?); + } + Ok(Response { + status: response.status, + headers: response.headers, + body: body.freeze(), + }) + } + .await + } + }; + server.join().expect("HTTP fixture should finish"); + result +} + +fn start_server( + respond: impl FnOnce(&mut TcpStream) + Send + 'static, +) -> (String, std::thread::JoinHandle<()>) { + let listener = TcpListener::bind(("127.0.0.1", 0)).expect("HTTP fixture should bind"); + let address = listener.local_addr().expect("HTTP fixture address"); + listener + .set_nonblocking(true) + .expect("nonblocking listener"); + let server = std::thread::spawn(move || { + let deadline = Instant::now() + Duration::from_secs(5); + let (mut connection, _) = loop { + match listener.accept() { + Ok(connection) => break connection, + Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => { + assert!(Instant::now() < deadline, "HTTP fixture accept timed out"); + std::thread::sleep(Duration::from_millis(1)); + } + Err(error) => panic!("HTTP fixture should accept: {error}"), + } + }; + connection + .set_nonblocking(false) + .expect("HTTP fixture connection should use blocking reads"); + connection + .set_read_timeout(Some(Duration::from_secs(5))) + .expect("HTTP fixture read timeout"); + let mut request = Vec::new(); + let mut buffer = [0; 1024]; + while !request.windows(4).any(|window| window == b"\r\n\r\n") { + let count = connection.read(&mut buffer).expect("read HTTP request"); + assert_ne!(count, 0, "HTTP request should contain complete headers"); + request.extend_from_slice(&buffer[..count]); + } + respond(&mut connection); + }); + ( + format!("http://{address}/models?credential=url-secret"), + server, + ) +} diff --git a/codex-rs/response-debug-context/src/lib.rs b/codex-rs/response-debug-context/src/lib.rs index 42825a7c98..4754705490 100644 --- a/codex-rs/response-debug-context/src/lib.rs +++ b/codex-rs/response-debug-context/src/lib.rs @@ -64,6 +64,7 @@ pub fn telemetry_transport_error_message(error: &TransportError) -> String { match error { TransportError::Http { status, .. } => format!("http {}", status.as_u16()), TransportError::RetryLimit => "retry limit reached".to_string(), + TransportError::ResponseTooLarge { .. } => "response body too large".to_string(), TransportError::Timeout => "timeout".to_string(), TransportError::Connection(err) => err.to_string(), TransportError::Network(err) => err.to_string(),