diff --git a/codex-rs/codex-api/src/files.rs b/codex-rs/codex-api/src/files.rs index fab92ed28a..84564fb7ad 100644 --- a/codex-rs/codex-api/src/files.rs +++ b/codex-rs/codex-api/src/files.rs @@ -1,3 +1,4 @@ +use std::future::Future; use std::time::Duration; use crate::AuthProvider; @@ -18,6 +19,8 @@ pub const OPENAI_FILE_URI_PREFIX: &str = "sediment://"; pub const OPENAI_FILE_UPLOAD_LIMIT_BYTES: u64 = 512 * 1024 * 1024; const OPENAI_FILE_REQUEST_TIMEOUT: Duration = Duration::from_secs(60); +const OPENAI_FILE_BLOB_UPLOAD_TIMEOUT: Duration = Duration::from_secs(5 * 60); +const MAX_BLOB_UPLOAD_ATTEMPTS: u64 = 5; const OPENAI_FILE_FINALIZE_TIMEOUT: Duration = Duration::from_secs(30); const OPENAI_FILE_FINALIZE_RETRY_DELAY: Duration = Duration::from_millis(250); const OPENAI_FILE_USE_CASE: &str = "codex"; @@ -41,6 +44,8 @@ pub struct UploadedOpenAiFile { #[derive(Debug, thiserror::Error)] pub enum OpenAiFileError { + #[error("failed to open OpenAI file upload contents: {0}")] + ReadContents(#[source] std::io::Error), #[error( "file `{file_name}` is too large: {size_bytes} bytes exceeds the limit of {limit_bytes} bytes" )] @@ -118,15 +123,21 @@ pub fn openai_file_uri(file_id: &str) -> String { format!("{OPENAI_FILE_URI_PREFIX}{file_id}") } -pub async fn upload_openai_file( +/// Creates and finalizes one file record, reopening the same contents for blob PUT retries. +pub async fn upload_openai_file( base_url: &str, auth: &dyn AuthProvider, client_pool: &RouteAwareClientPool, file_name: String, file_size_bytes: u64, - contents: impl Stream> + Send + 'static, + mut open_contents: Open, hosted_upload: Option<&HostedFileUploadContext>, -) -> Result { +) -> Result +where + Open: FnMut() -> OpenFuture, + OpenFuture: Future>, + Contents: Stream> + Send + 'static, +{ if file_size_bytes > OPENAI_FILE_UPLOAD_LIMIT_BYTES { return Err(OpenAiFileError::FileTooLarge { file_name, @@ -186,19 +197,12 @@ pub async fn upload_openai_file( .ok() .and_then(|url| url.host_str().map(str::to_owned)) .unwrap_or_else(|| "unknown-host".to_string()); - let azure_client_request_id = Uuid::new_v4().to_string(); let upload_started_at = Instant::now(); - let upload_response = client_pool - .put(&create_payload.upload_url) - .timeout(Duration::from_secs(5 * 60)) - .header("x-ms-blob-type", "BlockBlob") - .header("x-ms-client-request-id", &azure_client_request_id) - .header(CONTENT_LENGTH, file_size_bytes) - .body_stream(contents) - .send() - .await - .map_err(|source| { - let elapsed_ms = upload_started_at.elapsed().as_millis(); + let deadline = upload_started_at + OPENAI_FILE_BLOB_UPLOAD_TIMEOUT; + let mut retry_delay = Duration::ZERO; + for attempt in 1..=MAX_BLOB_UPLOAD_ATTEMPTS { + let azure_client_request_id = Uuid::new_v4().to_string(); + let request_error = |source: RouteAwareRequestError| { let error_kind = if source.is_timeout() { "timeout" } else if source.is_connect() { @@ -210,53 +214,135 @@ pub async fn upload_openai_file( } else { "other" }; - tracing::event!( - target: "codex_otel.log_only", - tracing::Level::WARN, - event.name = "codex.openai_file_blob_upload_failed", - file_id = %create_payload.file_id, - host = %upload_host, - file_size_bytes, - elapsed_ms, - error_kind, - azure_client_request_id, - "OpenAI file blob upload transport failed" - ); OpenAiFileError::BlobUploadRequest { host: upload_host.clone(), - elapsed_ms, + elapsed_ms: upload_started_at.elapsed().as_millis(), error_kind, azure_client_request_id: azure_client_request_id.clone(), source: source.without_url(), } - })?; - let upload_status = upload_response.status(); - let cloudflare_ray_id = upload_response_header(&upload_response, "cf-ray"); - let azure_request_id = upload_response_header(&upload_response, "x-ms-request-id"); - let azure_error_code = upload_response_header(&upload_response, "x-ms-error-code"); - if !upload_status.is_success() { + }; + // Opening the stream, sending every attempt, and backoff share one deadline. + let result = tokio::time::timeout_at(deadline, async { + if Instant::now() >= deadline { + return Err(request_error(RouteAwareRequestError::Timeout)); + } + let contents = open_contents() + .await + .map_err(OpenAiFileError::ReadContents)?; + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + return Err(request_error(RouteAwareRequestError::Timeout)); + } + client_pool + .put(&create_payload.upload_url) + .timeout(remaining) + .header("x-ms-blob-type", "BlockBlob") + .header("x-ms-client-request-id", &azure_client_request_id) + .header(CONTENT_LENGTH, file_size_bytes) + .body_stream(contents) + .send() + .await + .map_err(request_error) + }) + .await + .unwrap_or_else(|_| Err(request_error(RouteAwareRequestError::Timeout))); + + let mut status = None; + let mut cloudflare_ray_id = "missing".to_string(); + let mut azure_request_id = "missing".to_string(); + let mut azure_error_code = "missing".to_string(); + let mut server_delay = None; + let (error, retryable, error_kind) = match result { + Ok(response) => { + let upload_status = response.status(); + status = Some(upload_status.to_string()); + cloudflare_ray_id = upload_response_header(&response, "cf-ray"); + azure_request_id = upload_response_header(&response, "x-ms-request-id"); + azure_error_code = upload_response_header(&response, "x-ms-error-code"); + if upload_status.is_success() { + if attempt > 1 { + tracing::event!( + target: "codex_otel.log_only", + tracing::Level::INFO, + event.name = "codex.openai_file_blob_upload_recovered", + file_id = %create_payload.file_id, + host = %upload_host, + file_size_bytes, + elapsed_ms = upload_started_at.elapsed().as_millis(), + attempt, + delay_ms = retry_delay.as_millis(), + status = status.as_deref(), + cloudflare_ray_id, + azure_client_request_id, + azure_request_id, + azure_error_code, + "OpenAI file blob upload recovered" + ); + } + break; + } + server_delay = blob_retry_after(response.headers()); + ( + OpenAiFileError::BlobUploadStatus { + host: upload_host.clone(), + status: upload_status, + azure_client_request_id: azure_client_request_id.clone(), + azure_request_id: azure_request_id.clone(), + azure_error_code: azure_error_code.clone(), + }, + upload_status == StatusCode::SERVICE_UNAVAILABLE, + "http_status", + ) + } + Err(error) => { + let (retryable, error_kind) = match &error { + OpenAiFileError::BlobUploadRequest { + source, error_kind, .. + } => ( + source.is_timeout() + || source.is_connect() + || source.is_body() + || source.is_request(), + *error_kind, + ), + _ => (false, "read"), + }; + (error, retryable, error_kind) + } + }; + retry_delay = server_delay + .unwrap_or_else(|| codex_client::backoff(Duration::from_millis(125), attempt)); + let will_retry = retryable + && attempt < MAX_BLOB_UPLOAD_ATTEMPTS + && retry_delay < deadline.saturating_duration_since(Instant::now()); + let event_name = if will_retry { + "codex.openai_file_blob_upload_retry" + } else { + "codex.openai_file_blob_upload_failed" + }; tracing::event!( target: "codex_otel.log_only", tracing::Level::WARN, - event.name = "codex.openai_file_blob_upload_failed", + event.name = event_name, file_id = %create_payload.file_id, host = %upload_host, file_size_bytes, elapsed_ms = upload_started_at.elapsed().as_millis(), - status = %upload_status, + attempt, + delay_ms = if will_retry { retry_delay.as_millis() } else { 0 }, + status = status.as_deref(), + error_kind, cloudflare_ray_id, azure_client_request_id, azure_request_id, azure_error_code, - "OpenAI file blob upload failed" + "OpenAI file blob upload attempt failed" ); - return Err(OpenAiFileError::BlobUploadStatus { - host: upload_host, - status: upload_status, - azure_client_request_id, - azure_request_id, - azure_error_code, - }); + if !will_retry { + return Err(error); + } + tokio::time::sleep(retry_delay).await; } let finalize_url = format!( @@ -363,6 +449,30 @@ fn upload_response_header(response: &HttpResponse, header: &str) -> String { .to_string() } +fn blob_retry_after(headers: &http::HeaderMap) -> Option { + if let Some(delay) = headers + .get("x-ms-retry-after-ms") + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse().ok()) + { + return Some(Duration::from_millis(delay)); + } + let value = headers.get(http::header::RETRY_AFTER)?.to_str().ok()?; + if let Ok(seconds) = value.parse() { + return Some(Duration::from_secs(seconds)); + } + let date = chrono::DateTime::parse_from_rfc2822(value).ok()?; + Some( + date.signed_duration_since(chrono::Utc::now()) + .to_std() + .unwrap_or_default(), + ) +} + +#[cfg(test)] +#[path = "files_retry_tests.rs"] +mod retry_tests; + #[cfg(test)] mod tests { use super::*; @@ -385,9 +495,9 @@ mod tests { use wiremock::matchers::path; #[derive(Clone, Copy)] - struct ChatGptTestAuth; + pub(super) struct ChatGptTestAuth; - fn default_http_client_pool() -> RouteAwareClientPool { + pub(super) fn default_http_client_pool() -> RouteAwareClientPool { RouteAwareClientPool::new_without_request_logging( HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), ClientRouteClass::Api, @@ -405,11 +515,11 @@ mod tests { } } - fn chatgpt_auth() -> ChatGptTestAuth { + pub(super) fn chatgpt_auth() -> ChatGptTestAuth { ChatGptTestAuth } - fn base_url_for(server: &MockServer) -> String { + pub(super) fn base_url_for(server: &MockServer) -> String { format!("{}/backend-api", server.uri()) } @@ -461,15 +571,15 @@ mod tests { .await; let base_url = base_url_for(&server); - let contents = - futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static(b"hello"))]); + let open_contents = + || async { Ok(futures::stream::iter([Ok(Bytes::from_static(b"hello"))])) }; let uploaded = upload_openai_file( &base_url, &chatgpt_auth(), &default_http_client_pool(), "hello.txt".to_string(), /*file_size_bytes*/ 5, - contents, + open_contents, /*hosted_upload*/ None, ) .await @@ -542,7 +652,7 @@ mod tests { &default_http_client_pool(), "report.pdf".to_string(), /*file_size_bytes*/ 8, - futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static(b"%PDF-1.4"))]), + || async { Ok(futures::stream::iter([Ok(Bytes::from_static(b"%PDF-1.4"))])) }, Some(&hosted_upload), ) .await @@ -595,7 +705,7 @@ mod tests { &default_http_client_pool(), "report.pdf".to_string(), /*file_size_bytes*/ 8, - futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static(b"%PDF-1.4"))]), + || async { Ok(futures::stream::iter([Ok(Bytes::from_static(b"%PDF-1.4"))])) }, Some(&hosted_upload), ) .await @@ -653,15 +763,18 @@ mod tests { let client_pool = default_http_client_pool(); let mut uploaded_files = Vec::new(); for (file_name, _, contents) in files { - let contents_stream = - futures::stream::iter([Ok::<_, std::io::Error>(Bytes::copy_from_slice(contents))]); + let open_contents = || async { + Ok(futures::stream::iter([Ok(Bytes::copy_from_slice( + contents, + ))])) + }; let uploaded = upload_openai_file( &base_url_for(&server), &chatgpt_auth(), &client_pool, file_name.to_string(), u64::try_from(contents.len()).expect("file size should fit in a u64"), - contents_stream, + open_contents, /*hosted_upload*/ None, ) .await @@ -721,7 +834,7 @@ mod tests { &default_http_client_pool(), "hello.txt".to_string(), /*file_size_bytes*/ 5, - futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static(b"hello"))]), + || async { Ok(futures::stream::iter([Ok(Bytes::from_static(b"hello"))])) }, /*hosted_upload*/ None, ) .await @@ -759,7 +872,7 @@ mod tests { &default_http_client_pool(), "hello.txt".to_string(), /*file_size_bytes*/ 5, - futures::stream::iter([Ok::<_, std::io::Error>(Bytes::from_static(b"hello"))]), + || async { Ok(futures::stream::iter([Ok(Bytes::from_static(b"hello"))])) }, /*hosted_upload*/ None, ) .await diff --git a/codex-rs/codex-api/src/files_retry_tests.rs b/codex-rs/codex-api/src/files_retry_tests.rs new file mode 100644 index 0000000000..037ae3615f --- /dev/null +++ b/codex-rs/codex-api/src/files_retry_tests.rs @@ -0,0 +1,222 @@ +//! Upload retries replay bytes without allocating or finalizing another file record. + +use super::tests::base_url_for; +use super::tests::chatgpt_auth; +use super::tests::default_http_client_pool; +use super::*; +use pretty_assertions::assert_eq; +use std::collections::HashSet; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::Request; +use wiremock::Respond; +use wiremock::ResponseTemplate; +use wiremock::matchers::body_bytes; +use wiremock::matchers::method; +use wiremock::matchers::path; + +async fn mock_upload( + server: &MockServer, + response: impl Respond + 'static, + puts: u64, + finalizes: u64, +) { + Mock::given(method("POST")) + .and(path("/backend-api/files")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "file_id": "file_123", + "upload_url": format!("{}/upload?sig=secret", server.uri()), + }))) + .expect(1) + .mount(server) + .await; + Mock::given(method("PUT")) + .and(path("/upload")) + .and(body_bytes(b"hello")) + .respond_with(response) + .expect(puts) + .mount(server) + .await; + Mock::given(method("POST")) + .and(path("/backend-api/files/file_123/uploaded")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "status": "success", "download_url": "https://example.com/file" + }))) + .expect(finalizes) + .mount(server) + .await; +} + +#[tokio::test] +async fn retries_503_within_budget_but_not_4xx() { + for (statuses, azure_code, retry_after) in [ + (vec![503, 200], "ServerBusy", None), + (vec![503; 5], "ServerBusy", None), + (vec![403], "AuthenticationFailed", None), + (vec![503], "ServerBusy", Some("301")), + ] { + let server = MockServer::start().await; + let attempts = statuses.len(); + let succeeds = statuses.last() == Some(&200); + let put_count = AtomicUsize::new(0); + mock_upload( + &server, + move |_request: &Request| { + let index = put_count.fetch_add(1, Ordering::SeqCst); + let response = ResponseTemplate::new(statuses[index]) + .insert_header("x-ms-request-id", format!("request-{index}")); + let response = match retry_after { + Some(delay) => response.insert_header("retry-after", delay), + None => response.insert_header("x-ms-retry-after-ms", "0"), + }; + response.insert_header("x-ms-error-code", azure_code) + }, + attempts as u64, + u64::from(succeeds), + ) + .await; + let mut opens = 0; + + let result = upload_openai_file( + &base_url_for(&server), + &chatgpt_auth(), + &default_http_client_pool(), + "hello.txt".into(), + /*file_size_bytes*/ 5, + || { + opens += 1; + async { Ok(futures::stream::iter([Ok(Bytes::from_static(b"hello"))])) } + }, + /*hosted_upload*/ None, + ) + .await; + + assert_eq!((result.is_ok(), opens), (succeeds, attempts)); + if let Err(error) = result { + let message = error.to_string(); + assert!(message.contains(&format!("azure_request_id=request-{}", attempts - 1))); + assert!(message.contains(azure_code)); + assert!(!format!("{error:?}").contains("sig=secret")); + } + let requests = server.received_requests().await.expect("requests"); + let mut request_ids = HashSet::new(); + for request in requests.iter().filter(|r| r.method == Method::PUT) { + let id = request.headers["x-ms-client-request-id"] + .to_str() + .expect("id"); + assert!(request_ids.insert(id.to_owned())); + } + assert_eq!(request_ids.len(), attempts); + server.verify().await; + } +} + +#[tokio::test] +async fn reopens_after_a_stream_closure_and_sends_the_complete_body() { + let server = MockServer::start().await; + mock_upload( + &server, + ResponseTemplate::new(200), + /*puts*/ 1, + /*finalizes*/ 1, + ) + .await; + let mut opens = 0; + + let result = upload_openai_file( + &base_url_for(&server), + &chatgpt_auth(), + &default_http_client_pool(), + "hello.txt".into(), + /*file_size_bytes*/ 5, + || { + opens += 1; + let chunks = if opens == 1 { + vec![ + Ok(Bytes::from_static(b"he")), + Err(std::io::Error::new( + std::io::ErrorKind::ConnectionReset, + "stream closed", + )), + ] + } else { + vec![Ok(Bytes::from_static(b"hello"))] + }; + async move { Ok(futures::stream::iter(chunks)) } + }, + /*hosted_upload*/ None, + ) + .await + .expect("replayed upload"); + + assert_eq!((result.file_id.as_str(), opens), ("file_123", 2)); + server.verify().await; +} + +#[tokio::test] +async fn reopening_cannot_extend_the_shared_upload_deadline() { + let server = MockServer::start().await; + mock_upload( + &server, + ResponseTemplate::new(503).insert_header("x-ms-retry-after-ms", "0"), + /*puts*/ 1, + /*finalizes*/ 0, + ) + .await; + let mut opens = 0; + let result = upload_openai_file( + &base_url_for(&server), + &chatgpt_auth(), + &default_http_client_pool(), + "hello.txt".into(), + /*file_size_bytes*/ 5, + || { + opens += 1; + let attempt = opens; + async move { + if attempt == 2 { + tokio::time::pause(); + tokio::time::sleep(OPENAI_FILE_BLOB_UPLOAD_TIMEOUT).await; + } + Ok(futures::stream::iter([Ok(Bytes::from_static(b"hello"))])) + } + }, + /*hosted_upload*/ None, + ) + .await; + + assert!(matches!( + result, + Err(OpenAiFileError::BlobUploadRequest { + error_kind: "timeout", + .. + }) + )); + assert_eq!(opens, 2); + server.verify().await; +} + +#[test] +fn server_retry_delays_accept_milliseconds_seconds_and_http_dates() { + for (milliseconds, retry_after, expected) in [ + (Some("250"), Some("5"), Some(Duration::from_millis(250))), + (Some("invalid"), Some("5"), Some(Duration::from_secs(5))), + ( + None, + Some("Wed, 21 Oct 2015 07:28:00 GMT"), + Some(Duration::ZERO), + ), + (None, Some("invalid"), None), + ] { + let mut headers = http::HeaderMap::new(); + if let Some(value) = milliseconds { + headers.insert("x-ms-retry-after-ms", value.parse().expect("header")); + } + if let Some(value) = retry_after { + headers.insert("retry-after", value.parse().expect("header")); + } + assert_eq!(blob_retry_after(&headers), expected); + } +} diff --git a/codex-rs/core/src/mcp_openai_file.rs b/codex-rs/core/src/mcp_openai_file.rs index 1e33c74215..04f4e9f9bf 100644 --- a/codex-rs/core/src/mcp_openai_file.rs +++ b/codex-rs/core/src/mcp_openai_file.rs @@ -23,6 +23,7 @@ use codex_sandboxing::policy_transforms::effective_file_system_sandbox_policy; use codex_sandboxing::policy_transforms::merge_permission_profiles; use serde_json::Value as JsonValue; use std::collections::HashMap; +use std::sync::Arc; struct FileArgumentLocation<'a> { field_name: &'a str, @@ -213,10 +214,12 @@ async fn build_uploaded_argument_value( OPENAI_FILE_UPLOAD_LIMIT_BYTES, ))); } - let contents = fs - .read_file_stream(&path_uri, sandbox.as_ref()) - .await - .map_err(|error| contextualize_error(error.to_string()))?; + // Keep the initial read permission check before allocating a remote file record. + let mut first_contents = Some( + fs.read_file_stream(&path_uri, sandbox.as_ref()) + .await + .map_err(|error| contextualize_error(error.to_string()))?, + ); let file_name = path_uri .basename() .or_else(|| { @@ -228,6 +231,18 @@ async fn build_uploaded_argument_value( }) }) .unwrap_or_else(|| "file".to_string()); + let open_contents = move || { + let first_contents = first_contents.take(); + let fs = Arc::clone(&fs); + let path_uri = path_uri.clone(); + let sandbox = sandbox.clone(); + async move { + match first_contents { + Some(contents) => Ok(contents), + None => fs.read_file_stream(&path_uri, sandbox.as_ref()).await, + } + } + }; let upload_auth = codex_model_provider::auth_provider_from_auth(auth); let uploaded = upload_openai_file( turn_context.config.chatgpt_base_url.trim_end_matches('/'), @@ -235,7 +250,7 @@ async fn build_uploaded_argument_value( &sess.services.openai_file_upload_client_pool, file_name, metadata.size, - contents, + open_contents, hosted_upload, ) .await diff --git a/codex-rs/core/tests/suite/openai_file_mcp.rs b/codex-rs/core/tests/suite/openai_file_mcp.rs index 9ae4b31214..db3d52b5bf 100644 --- a/codex-rs/core/tests/suite/openai_file_mcp.rs +++ b/codex-rs/core/tests/suite/openai_file_mcp.rs @@ -329,7 +329,7 @@ async fn codex_apps_file_params_pass_uploaded_file_to_post_tool_use_hook() -> Re } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn codex_apps_file_params_stream_allowed_file_under_restricted_read_policy() -> Result<()> { +async fn codex_apps_file_params_retry_stream_under_restricted_read_policy() -> Result<()> { skip_if_target_windows!( Ok(()), "Windows restricted-token sandbox cannot enforce deny-read policies" @@ -339,6 +339,14 @@ async fn codex_apps_file_params_stream_allowed_file_under_restricted_read_policy let server = start_mock_server().await; let apps_server = AppsTestServer::mount(&server).await?; mount_file_upload_mocks(&server, STREAMED_FILE_SIZE as u64).await; + Mock::given(method("PUT")) + .and(path("/upload/file_123")) + .respond_with(ResponseTemplate::new(503).insert_header("x-ms-retry-after-ms", "0")) + .with_priority(1) + .up_to_n_times(1) + .expect(1) + .mount(&server) + .await; let mut builder = apps_enabled_builder(apps_server.chatgpt_base_url) .with_config(|config| restrict_apps_upload_reads(config, "private.txt")) @@ -358,6 +366,15 @@ async fn codex_apps_file_params_stream_allowed_file_under_restricted_read_policy run_extract_turn(&test, &server, permission_profile).await?; + let requests = server.received_requests().await.expect("capture requests"); + let bodies: Vec<_> = requests + .iter() + .filter(|request| request.url.path() == "/upload/file_123") + .map(|request| request.body.as_slice()) + .collect(); + let contents = vec![b'x'; STREAMED_FILE_SIZE]; + assert_eq!(bodies, vec![contents.as_slice(); 2]); + let apps_tool_call = recorded_apps_tool_call_by_name(&server, CALENDAR_EXTRACT_TEXT_TOOL_NAME).await; assert_eq!(