From 64569aedd464ae04cfe4af95cb239184c3870946 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Fri, 3 Jul 2026 16:01:48 -0700 Subject: [PATCH] fix(core): keep capacity retries single-budget --- codex-rs/core/src/client.rs | 21 ++- codex-rs/core/src/compact_remote_v2.rs | 27 ++- codex-rs/core/src/responses_retry.rs | 16 +- codex-rs/core/src/session/turn.rs | 28 ++- codex-rs/core/tests/suite/client.rs | 166 +++++++++++++++--- .../core/tests/suite/client_websockets.rs | 36 ++++ codex-rs/core/tests/suite/compact.rs | 57 ++++++ codex-rs/core/tests/suite/compact_remote.rs | 2 +- 8 files changed, 306 insertions(+), 47 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 251b7c1dcf..edcd5dfad1 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -272,6 +272,7 @@ pub struct ModelClient { pub struct ModelClientSession { client: ModelClient, websocket_session: WebsocketSession, + caller_manages_http_request_retries: bool, /// Turn state for sticky routing. /// /// This is an `OnceLock` that stores the turn state value received from the server @@ -480,6 +481,7 @@ impl ModelClient { ModelClientSession { client: self.clone(), websocket_session: self.take_cached_websocket_session(), + caller_manages_http_request_retries: false, turn_state: Arc::new(OnceLock::new()), } } @@ -1118,6 +1120,10 @@ impl ModelClientSession { Arc::clone(&self.turn_state) } + pub(crate) fn use_caller_managed_http_request_retries(&mut self) { + self.caller_manages_http_request_retries = true; + } + fn reset_websocket_session(&mut self) { self.websocket_session.connection = None; self.websocket_session.last_request = None; @@ -1453,12 +1459,15 @@ impl ModelClientSession { let inference_trace_attempt = inference_trace.start_attempt(); inference_trace_attempt.add_request_headers(&mut options.extra_headers); inference_trace_attempt.record_started(&request); - let client = ApiResponsesClient::new( - transport, - client_setup.api_provider, - client_setup.api_auth, - ) - .with_telemetry(Some(request_telemetry), Some(sse_telemetry)); + let mut responses_api_provider = client_setup.api_provider; + if self.caller_manages_http_request_retries { + // Keep ordinary HTTP open failures in one outer request loop. + responses_api_provider.retry.retry_5xx = false; + responses_api_provider.retry.retry_transport = false; + } + let client = + ApiResponsesClient::new(transport, responses_api_provider, client_setup.api_auth) + .with_telemetry(Some(request_telemetry), Some(sse_telemetry)); let stream_result = client.stream_request(request, options).await; match stream_result { diff --git a/codex-rs/core/src/compact_remote_v2.rs b/codex-rs/core/src/compact_remote_v2.rs index 10e2c74e02..fadf5934d1 100644 --- a/codex-rs/core/src/compact_remote_v2.rs +++ b/codex-rs/core/src/compact_remote_v2.rs @@ -17,6 +17,7 @@ use crate::hook_runtime::run_post_compact_hooks; use crate::hook_runtime::run_pre_compact_hooks; use crate::responses_metadata::CodexResponsesMetadata; use crate::responses_metadata::CompactionTurnMetadata; +use crate::responses_retry::ResponsesRetryOutcome; use crate::responses_retry::ResponsesStreamRequest; use crate::responses_retry::handle_retryable_response_stream_error; use crate::session::session::Session; @@ -347,14 +348,18 @@ async fn run_remote_compaction_request_v2( responses_metadata: &CodexResponsesMetadata, cancellation_token: &CancellationToken, ) -> CodexResult { - let max_retries = turn_context + let max_stream_retries = turn_context .provider .info() .stream_max_retries() .min(MAX_REMOTE_COMPACTION_V2_STREAM_RETRIES); - let mut retries = 0; + let max_request_retries = turn_context.provider.info().request_max_retries(); + client_session.use_caller_managed_http_request_retries(); + let mut stream_retries = 0; + let mut request_retries = 0; let mut server_overloaded_retries = 0; loop { + let mut response_stream_opened = false; let result = async { let stream = client_session .stream( @@ -368,6 +373,7 @@ async fn run_remote_compaction_request_v2( &InferenceTraceContext::disabled(), ) .await?; + response_stream_opened = true; collect_compaction_output(stream).await } .or_cancel(cancellation_token) @@ -382,12 +388,17 @@ async fn run_remote_compaction_request_v2( if !err.is_retryable() && !should_retry_server_overload { return Err(err); } - let retry_counter = if should_retry_server_overload { - &mut server_overloaded_retries + let http_request_failed_to_open = + !sess.services.model_client.responses_websocket_enabled() + && !response_stream_opened; + let (retry_counter, max_retries) = if should_retry_server_overload { + (&mut server_overloaded_retries, max_stream_retries) + } else if http_request_failed_to_open { + (&mut request_retries, max_request_retries) } else { - &mut retries + (&mut stream_retries, max_stream_retries) }; - handle_retryable_response_stream_error( + let retry_outcome = handle_retryable_response_stream_error( retry_counter, max_retries, err, @@ -398,6 +409,10 @@ async fn run_remote_compaction_request_v2( cancellation_token.child_token(), ) .await?; + if retry_outcome == ResponsesRetryOutcome::TransportFallback { + request_retries = 0; + stream_retries = 0; + } } } } diff --git a/codex-rs/core/src/responses_retry.rs b/codex-rs/core/src/responses_retry.rs index eb55b47f94..f9b8a05b84 100644 --- a/codex-rs/core/src/responses_retry.rs +++ b/codex-rs/core/src/responses_retry.rs @@ -26,8 +26,14 @@ pub(crate) enum ResponsesStreamRequest { RemoteCompactionV2, } -/// Handles a retryable stream error and returns `Ok(())` when the caller should -/// retry the request loop. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ResponsesRetryOutcome { + Retrying, + TransportFallback, +} + +/// Handles a retryable stream error and reports whether the retry stays on the +/// current transport or falls back to another one. #[allow(clippy::too_many_arguments)] pub(crate) async fn handle_retryable_response_stream_error( retries: &mut u64, @@ -38,7 +44,7 @@ pub(crate) async fn handle_retryable_response_stream_error( turn_context: &TurnContext, request: ResponsesStreamRequest, cancellation_token: CancellationToken, -) -> Result<(), CodexErr> { +) -> Result { let is_server_overloaded = matches!(&err, CodexErr::ServerOverloaded); let max_retries = if is_server_overloaded { SERVER_OVERLOADED_MAX_RETRIES @@ -60,7 +66,7 @@ pub(crate) async fn handle_retryable_response_stream_error( ) .await; *retries = 0; - return Ok(()); + return Ok(ResponsesRetryOutcome::TransportFallback); } if *retries < max_retries { @@ -94,7 +100,7 @@ pub(crate) async fn handle_retryable_response_stream_error( .await; } return tokio::select! { - () = tokio::time::sleep(delay) => Ok(()), + () = tokio::time::sleep(delay) => Ok(ResponsesRetryOutcome::Retrying), () = cancellation_token.cancelled() => Err(CodexErr::TurnAborted), }; } diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index cda983359e..ba224497c3 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -37,6 +37,7 @@ use crate::mentions::collect_tool_mentions_from_messages; use crate::plugins::build_plugin_injections; use crate::responses_metadata::CodexResponsesMetadata; use crate::responses_metadata::CodexResponsesRequestKind; +use crate::responses_retry::ResponsesRetryOutcome; use crate::responses_retry::ResponsesStreamRequest; use crate::responses_retry::handle_retryable_response_stream_error; use crate::session::PreviousTurnSettings; @@ -1152,8 +1153,11 @@ async fn run_sampling_request( Arc::clone(&router), Arc::clone(&turn_diff_tracker), ); - let max_retries = turn_context.provider.info().stream_max_retries(); - let mut retries = 0; + let max_stream_retries = turn_context.provider.info().stream_max_retries(); + let max_request_retries = turn_context.provider.info().request_max_retries(); + client_session.use_caller_managed_http_request_retries(); + let mut stream_retries = 0; + let mut request_retries = 0; let mut server_overloaded_retries = 0; let mut initial_input = Some(input); let mut original_input = None; @@ -1171,12 +1175,14 @@ async fn run_sampling_request( turn_context.as_ref(), base_instructions.clone(), ); + let mut response_stream_opened = false; let err = match try_run_sampling_request( tool_runtime.clone(), Arc::clone(&sess), Arc::clone(&turn_context), Arc::clone(&turn_store), client_session, + &mut response_stream_opened, responses_metadata, Arc::clone(&turn_diff_tracker), &prompt, @@ -1211,12 +1217,16 @@ async fn run_sampling_request( return Err(err); } - let retry_counter = if should_retry_server_overload { - &mut server_overloaded_retries + let http_request_failed_to_open = + !sess.services.model_client.responses_websocket_enabled() && !response_stream_opened; + let (retry_counter, max_retries) = if should_retry_server_overload { + (&mut server_overloaded_retries, max_stream_retries) + } else if http_request_failed_to_open { + (&mut request_retries, max_request_retries) } else { - &mut retries + (&mut stream_retries, max_stream_retries) }; - handle_retryable_response_stream_error( + let retry_outcome = handle_retryable_response_stream_error( retry_counter, max_retries, err, @@ -1227,6 +1237,10 @@ async fn run_sampling_request( cancellation_token.child_token(), ) .await?; + if retry_outcome == ResponsesRetryOutcome::TransportFallback { + request_retries = 0; + stream_retries = 0; + } turn_context.turn_timing_state.record_sampling_retry(); } } @@ -1969,6 +1983,7 @@ async fn try_run_sampling_request( turn_context: Arc, turn_store: Arc, client_session: &mut ModelClientSession, + response_stream_opened: &mut bool, responses_metadata: &CodexResponsesMetadata, turn_diff_tracker: SharedTurnDiffTracker, prompt: &Prompt, @@ -2007,6 +2022,7 @@ async fn try_run_sampling_request( .instrument(trace_span!("stream_request")) .or_cancel(&cancellation_token) .await??; + *response_stream_opened = true; let mut in_flight: FuturesOrdered>> = FuturesOrdered::new(); let mut needs_follow_up = false; diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 1a7e8d0d9d..35b74f15de 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -63,11 +63,13 @@ use core_test_support::responses::ev_function_call; use core_test_support::responses::ev_message_item_added; use core_test_support::responses::ev_output_text_delta; use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_response_sequence; use core_test_support::responses::mount_sse_once; use core_test_support::responses::mount_sse_once_match; use core_test_support::responses::mount_sse_sequence; use core_test_support::responses::sse; use core_test_support::responses::sse_failed; +use core_test_support::responses::sse_response; use core_test_support::responses::strip_metadata_from_json; use core_test_support::responses_metadata as test_responses_metadata; use core_test_support::skip_if_no_network; @@ -3527,41 +3529,38 @@ async fn incomplete_response_emits_content_filter_error_message() -> anyhow::Res async fn server_overload_retries_use_a_separate_budget_on_the_same_turn() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); let server = MockServer::start().await; - let responses = mount_sse_sequence( + let overloaded = || { + ResponseTemplate::new(503).set_body_json(json!({ + "error": { + "code": "server_is_overloaded", + "message": "selected model is at capacity" + } + })) + }; + let responses = mount_response_sequence( &server, vec![ - sse_failed( + sse_response(sse_failed( "resp_stream_error", "server_error", "temporary stream failure", - ), - sse_failed( - "resp_overloaded_1", - "server_is_overloaded", - "selected model is at capacity", - ), - sse_failed( - "resp_overloaded_2", - "server_is_overloaded", - "selected model is at capacity", - ), - sse_failed( - "resp_overloaded_3", - "server_is_overloaded", - "selected model is at capacity", - ), - sse(vec![ + )), + ResponseTemplate::new(500).set_body_string("temporary gateway failure"), + overloaded(), + overloaded(), + overloaded(), + sse_response(sse(vec![ ev_response_created("resp_retried"), ev_completed("resp_retried"), - ]), + ])), ], ) .await; let test = test_codex() .with_config(|config| { - config.model_provider.request_max_retries = Some(0); - config.model_provider.stream_max_retries = Some(5); + config.model_provider.request_max_retries = Some(4); + config.model_provider.stream_max_retries = Some(1); }) .build(&server) .await?; @@ -3604,7 +3603,7 @@ async fn server_overload_retries_use_a_separate_budget_on_the_same_turn() -> any "Reconnecting... 3/3", ] ); - assert_eq!(requests.len(), 5); + assert_eq!(requests.len(), 6); let turn_id = requests[0].body_json()["client_metadata"]["turn_id"].clone(); assert!( requests @@ -3615,6 +3614,127 @@ async fn server_overload_retries_use_a_separate_budget_on_the_same_turn() -> any Ok(()) } +#[tokio::test] +async fn http_5xx_retries_respect_request_retry_budget() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + let server = MockServer::start().await; + let responses = mount_response_sequence( + &server, + vec![ + ResponseTemplate::new(503).set_body_string("temporary gateway failure"), + ResponseTemplate::new(503).set_body_string("temporary gateway failure"), + ], + ) + .await; + + let test = test_codex() + .with_config(|config| { + config.model_provider.request_max_retries = Some(1); + config.model_provider.stream_max_retries = Some(5); + }) + .build(&server) + .await?; + tokio::time::pause(); + + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "respect request retry budget".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + + let mut saw_terminal_error = false; + loop { + match test.codex.next_event().await?.msg { + EventMsg::StreamError(_) => tokio::time::advance(Duration::from_secs(1)).await, + EventMsg::Error(_) => saw_terminal_error = true, + EventMsg::TurnComplete(_) => break, + _ => {} + } + } + + assert!(saw_terminal_error); + assert_eq!(responses.requests().len(), 2); + Ok(()) +} + +#[tokio::test] +async fn mixed_network_and_5xx_respect_single_request_retry_budget() -> anyhow::Result<()> { + fn connection_reset(_: &wiremock::Request) -> std::io::Error { + std::io::Error::new( + std::io::ErrorKind::ConnectionReset, + "synthetic connection reset", + ) + } + + skip_if_no_network!(Ok(())); + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/responses")) + .respond_with_err(connection_reset) + .up_to_n_times(1) + .expect(1) + .with_priority(1) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/v1/responses")) + .respond_with(ResponseTemplate::new(500).set_body_string("temporary gateway failure")) + .with_priority(2) + .mount(&server) + .await; + + let test = test_codex() + .with_config(|config| { + config.model_provider.request_max_retries = Some(1); + config.model_provider.stream_max_retries = Some(5); + }) + .build(&server) + .await?; + tokio::time::pause(); + + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "respect one shared request retry budget".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + + let mut saw_terminal_error = false; + loop { + match test.codex.next_event().await?.msg { + EventMsg::StreamError(_) => tokio::time::advance(Duration::from_secs(1)).await, + EventMsg::Error(_) => saw_terminal_error = true, + EventMsg::TurnComplete(_) => break, + _ => {} + } + } + + assert!(saw_terminal_error); + let response_requests = server + .received_requests() + .await + .unwrap_or_default() + .into_iter() + .filter(|request| request.method == wiremock::http::Method::POST) + .filter(|request| request.url.path() == "/v1/responses") + .count(); + assert_eq!(response_requests, 2); + Ok(()) +} + /// We try to avoid setting env vars in tests because std::env::set_var() is /// process-wide and unsafe. Though for this test, we want to simulate the /// presence of an environment variable that the provider will read for auth, so diff --git a/codex-rs/core/tests/suite/client_websockets.rs b/codex-rs/core/tests/suite/client_websockets.rs index 7ce2cc81c6..7c1e23db34 100755 --- a/codex-rs/core/tests/suite/client_websockets.rs +++ b/codex-rs/core/tests/suite/client_websockets.rs @@ -1622,6 +1622,42 @@ async fn responses_websocket_connection_limit_error_reconnects_and_completes() { server.shutdown().await; } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn responses_websocket_5xx_uses_stream_retry_budget() { + skip_if_no_network!(); + + let websocket_server_error = json!({ + "type": "error", + "status": 500, + "error": { + "type": "server_error", + "message": "temporary websocket failure" + } + }); + let server = start_websocket_server(vec![ + vec![vec![websocket_server_error]], + vec![vec![ev_response_created("resp-1"), ev_completed("resp-1")]], + ]) + .await; + let mut builder = test_codex().with_config(|config| { + config.model_provider.request_max_retries = Some(0); + config.model_provider.stream_max_retries = Some(1); + }); + let test = builder + .build_with_websocket_server(&server) + .await + .expect("build websocket codex"); + + test.submit_turn("hello") + .await + .expect("submission should retry websocket 5xx as a stream failure"); + + let total_websocket_requests: usize = server.connections().iter().map(Vec::len).sum(); + assert_eq!(total_websocket_requests, 2); + + server.shutdown().await; +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn responses_websocket_uses_incremental_create_on_prefix() { skip_if_no_network!(); diff --git a/codex-rs/core/tests/suite/compact.rs b/codex-rs/core/tests/suite/compact.rs index 93cabd3383..4bf7a9d791 100644 --- a/codex-rs/core/tests/suite/compact.rs +++ b/codex-rs/core/tests/suite/compact.rs @@ -65,6 +65,7 @@ use std::path::Path; use std::sync::Arc; use tempfile::TempDir; use wiremock::MockServer; +use wiremock::ResponseTemplate; // --- Test helpers ----------------------------------------------------------- pub(super) const FIRST_REPLY: &str = "FIRST_REPLY"; @@ -684,6 +685,62 @@ async fn summarize_context_three_requests_and_instructions() { ); } +#[tokio::test] +async fn local_compact_keeps_provider_http_5xx_retries() -> Result<()> { + skip_if_no_network!(Ok(())); + let server = start_mock_server().await; + let responses = mount_response_sequence( + &server, + vec![ + sse_response(sse(vec![ + ev_assistant_message("m1", FIRST_REPLY), + ev_completed("r1"), + ])), + ResponseTemplate::new(500).set_body_string("temporary compact failure"), + sse_response(sse(vec![ + ev_assistant_message("m2", SUMMARY_TEXT), + ev_completed("r2"), + ])), + ], + ) + .await; + + let mut model_provider = non_openai_model_provider(&server); + model_provider.request_max_retries = Some(1); + model_provider.stream_max_retries = Some(0); + let mut builder = test_codex().with_config(move |config| { + config.model_provider = model_provider; + set_test_compact_prompt(config); + }); + let test = builder.build(&server).await?; + + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "hello before compact retry".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + test.codex.submit(Op::Compact).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + assert_eq!(responses.requests().len(), 3); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn manual_pre_compact_block_decision_does_not_block_compaction() { skip_if_no_network!(); diff --git a/codex-rs/core/tests/suite/compact_remote.rs b/codex-rs/core/tests/suite/compact_remote.rs index 9582409750..a1b3a59b19 100644 --- a/codex-rs/core/tests/suite/compact_remote.rs +++ b/codex-rs/core/tests/suite/compact_remote.rs @@ -1030,7 +1030,7 @@ async fn remote_compact_v2_retries_capacity_after_stream_retry_budget() -> Resul .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) .with_config(|config| { let _ = config.features.enable(Feature::RemoteCompactionV2); - config.model_provider.request_max_retries = Some(0); + config.model_provider.request_max_retries = Some(1); config.model_provider.stream_max_retries = Some(2); }), )