From 4885eb6c52c4d9915698671395aaa327072afca9 Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 17 Aug 2026 12:33:24 +0000 Subject: [PATCH] Retry transient Guardian sampling failures (#38997) ## Why Guardian sampling can encounter temporary connection or service failures while opening a WebSocket, starting a request, or consuming its response stream. ## What changed - Retry transient transport, streaming, rate-limit, and server failures up to two times across the full sampling flow. - Recover from unauthorized responses through the configured auth manager and discard pooled connections that may contain stale credentials. - Allow sampler initialization to continue when its initial WebSocket prewarm fails, so the first sample can establish a fresh connection. ## Testing Added coverage for recovery after prewarm failures and `503` responses, and for stopping once the transient retry budget is exhausted. GitOrigin-RevId: af74c2df95c4241184ca641ef4418db6aa4ab695 --- codex-rs/ext/guardian-v2/src/sampler.rs | 119 ++++++++++++++--- codex-rs/ext/guardian-v2/src/sampler_tests.rs | 126 ++++++++++++++++++ 2 files changed, 229 insertions(+), 16 deletions(-) diff --git a/codex-rs/ext/guardian-v2/src/sampler.rs b/codex-rs/ext/guardian-v2/src/sampler.rs index 358d0a9724..97ccf41623 100644 --- a/codex-rs/ext/guardian-v2/src/sampler.rs +++ b/codex-rs/ext/guardian-v2/src/sampler.rs @@ -15,11 +15,13 @@ use codex_api::ResponsesApiRequest; use codex_api::ResponsesWebsocketClient; use codex_api::ResponsesWebsocketConnection; use codex_api::ResponsesWsRequest; +use codex_api::TransportError; use codex_api::build_session_headers; use codex_api::create_text_param_for_request; use codex_http_client::HttpClientFactory; use codex_login::AgentIdentityAuthPolicy; use codex_login::CodexAuth; +use codex_login::UnauthorizedRecovery; use codex_login::default_client::add_originator_header; use codex_login::default_client::default_headers; use codex_model_provider::AgentIdentitySessionFallback; @@ -31,6 +33,7 @@ use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::SessionSource; use http::HeaderValue; +use http::StatusCode; use serde_json::Value; use thiserror::Error; use tokio::sync::OwnedSemaphorePermit; @@ -41,6 +44,7 @@ pub(crate) const MODEL: &str = "gpt-5.6-luna"; const MAX_OUTPUT_BYTES: usize = 8 * 1024; const INITIAL_WEBSOCKET_CONNECTIONS: usize = 2; const MAX_WEBSOCKET_CONNECTIONS: usize = 16; +const MAX_SAMPLING_RETRIES: usize = 2; const MAX_WEBSOCKET_AGE: Duration = Duration::from_secs(55 * 60); const RESPONSES_WEBSOCKETS_BETA: &str = "responses_websockets=2026-02-06"; const RESPONSES_LITE_METADATA_KEY: &str = @@ -153,10 +157,9 @@ impl LunaSampler { capacity: Arc::new(Semaphore::new(MAX_WEBSOCKET_CONNECTIONS)), active_requests: Mutex::new(VecDeque::with_capacity(MAX_WEBSOCKET_CONNECTIONS)), }; - for index in 0..INITIAL_WEBSOCKET_CONNECTIONS { + for _ in 0..INITIAL_WEBSOCKET_CONNECTIONS { let connection = match sampler.open_connection().await { Ok(connection) => connection, - Err(error) if index == 0 => return Err(error), Err(_) => break, }; sampler @@ -277,6 +280,63 @@ impl LunaSampler { }) } + async fn retry_after_failure( + &self, + error: &LunaSamplerError, + auth_recovery: &mut Option, + retries: &mut usize, + ) -> bool { + let retryable = match error { + LunaSamplerError::ConnectionTimeout + | LunaSamplerError::Api( + ApiError::Retryable { .. } | ApiError::Stream(_) | ApiError::ServerOverloaded, + ) + | LunaSamplerError::Api(ApiError::Transport( + TransportError::RetryLimit + | TransportError::Timeout + | TransportError::Connection(_) + | TransportError::Network(_), + )) => true, + LunaSamplerError::Api(ApiError::Transport(TransportError::Http { status, .. })) + | LunaSamplerError::Api(ApiError::Api { status, .. }) => { + if *status == StatusCode::UNAUTHORIZED { + let Some(recovery) = auth_recovery.as_mut() else { + return false; + }; + if !recovery.has_next() || recovery.next().await.is_err() { + return false; + } + self.idle_connections + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clear(); + return true; + } else { + status.is_server_error() || *status == StatusCode::TOO_MANY_REQUESTS + } + } + LunaSamplerError::Provider(_) + | LunaSamplerError::MissingOutput + | LunaSamplerError::OutputTooLarge + | LunaSamplerError::Superseded + | LunaSamplerError::Api( + ApiError::Transport(TransportError::Build(_)) + | ApiError::ContextWindowExceeded + | ApiError::QuotaExceeded + | ApiError::UsageNotIncluded + | ApiError::RateLimit(_) + | ApiError::InvalidRequest { .. } + | ApiError::MisalignmentPolicyViolation { .. } + | ApiError::CyberPolicy { .. }, + ) => false, + }; + if retryable && *retries < MAX_SAMPLING_RETRIES { + *retries += 1; + return true; + } + false + } + /// Sends one structured, tool-less request on an exclusively leased WebSocket. pub async fn sample(&self, request: LunaSamplingRequest) -> Result { let metadata = HashMap::from([ @@ -377,13 +437,29 @@ impl LunaSampler { }); } let mut retries = 0; + let mut auth_recovery = self + .config + .provider + .auth_manager() + .map(|manager| manager.unauthorized_recovery()); 'retry: loop { - let lease = tokio::select! { + let lease = match tokio::select! { biased; _ = &mut superseded => return Err(LunaSamplerError::Superseded), - lease = self.lease_connection() => lease?, + lease = self.lease_connection() => lease, + } { + Ok(lease) => lease, + Err(error) => { + if self + .retry_after_failure(&error, &mut auth_recovery, &mut retries) + .await + { + continue; + } + return Err(error); + } }; - let mut stream = lease + let mut stream = match lease .connection .connection .stream_request( @@ -392,7 +468,19 @@ impl LunaSampler { /*turn_state*/ None, ) .await - .map_err(LunaSamplerError::Api)?; + { + Ok(stream) => stream, + Err(error) => { + let error = LunaSamplerError::Api(error); + if self + .retry_after_failure(&error, &mut auth_recovery, &mut retries) + .await + { + continue; + } + return Err(error); + } + }; let mut output = String::new(); let mut deltas = String::new(); @@ -409,17 +497,16 @@ impl LunaSampler { } { let event = match event { Ok(event) => event, - Err(error) - if retries < INITIAL_WEBSOCKET_CONNECTIONS - && matches!( - error, - ApiError::Retryable { .. } | ApiError::Stream(_) - ) => - { - retries += 1; - continue 'retry; + Err(error) => { + let error = LunaSamplerError::Api(error); + if self + .retry_after_failure(&error, &mut auth_recovery, &mut retries) + .await + { + continue 'retry; + } + return Err(error); } - Err(error) => return Err(LunaSamplerError::Api(error)), }; match event { ResponseEvent::OutputTextDelta(delta) => { diff --git a/codex-rs/ext/guardian-v2/src/sampler_tests.rs b/codex-rs/ext/guardian-v2/src/sampler_tests.rs index c4b55b6fef..3a73acdbc8 100644 --- a/codex-rs/ext/guardian-v2/src/sampler_tests.rs +++ b/codex-rs/ext/guardian-v2/src/sampler_tests.rs @@ -20,6 +20,8 @@ use core_test_support::skip_if_no_network; use pretty_assertions::assert_eq; use serde_json::json; use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::time::Duration; use tokio::net::TcpListener; use tokio::net::TcpStream; @@ -27,6 +29,7 @@ use tokio::net::TcpStream; use super::LunaSampler; use super::LunaSamplerConfig; use super::LunaSamplingRequest; +use super::MAX_SAMPLING_RETRIES; use super::MAX_WEBSOCKET_CONNECTIONS; async fn proxy_websocket_servers(servers: &[&responses::WebSocketTestServer]) -> Result { @@ -366,6 +369,51 @@ async fn sampler_returns_complete_json_before_terminal_response_events() -> Resu Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn sampler_recovers_after_initial_prewarm_failures() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_websocket_server(vec![vec![vec![ + ev_assistant_message("recovered", r#"{"score":0.25}"#), + ev_completed("recovered"), + ]]]) + .await; + let listener = TcpListener::bind("127.0.0.1:0").await?; + let address = listener.local_addr()?; + let target = server.uri().trim_start_matches("ws://").to_owned(); + let failed_connections = Arc::new(AtomicUsize::new(0)); + let observed_failures = Arc::clone(&failed_connections); + tokio::spawn(async move { + for _ in 0..=MAX_SAMPLING_RETRIES { + let Ok((connection, _)) = listener.accept().await else { + return; + }; + observed_failures.fetch_add(/*val*/ 1, Ordering::Relaxed); + drop(connection); + } + while let Ok((mut incoming, _)) = listener.accept().await { + let target = target.clone(); + tokio::spawn(async move { + let Ok(mut outgoing) = TcpStream::connect(target).await else { + return; + }; + let _ = tokio::io::copy_bidirectional(&mut incoming, &mut outgoing).await; + }); + } + }); + + let sampler = LunaSampler::connect(sampler_config(format!("http://{address}/v1"))).await?; + assert_eq!(failed_connections.load(Ordering::Relaxed), 1); + + assert_eq!( + sampler.sample(sample_request("turn-1")).await?, + r#"{"score":0.25}"# + ); + assert_eq!(server.handshakes().len(), 1); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn sampler_remains_available_when_second_prewarm_fails() -> Result<()> { skip_if_no_network!(Ok(())); @@ -560,3 +608,81 @@ async fn sampler_retries_expired_websockets_on_another_warm_connection() -> Resu ); Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn sampler_reconnects_after_transient_service_failures() -> Result<()> { + skip_if_no_network!(Ok(())); + + let unavailable = || { + vec![vec![vec![json!({ + "type": "error", + "status": 503, + "error": { + "type": "server_error", + "message": "temporarily unavailable" + } + })]]] + }; + let first = responses::start_websocket_server(unavailable()).await; + let second = responses::start_websocket_server(unavailable()).await; + let recovered = responses::start_websocket_server(vec![vec![vec![ + ev_assistant_message("recovered", r#"{"score":0.25}"#), + ev_completed("recovered"), + ]]]) + .await; + let sampler = LunaSampler::connect(sampler_config( + proxy_websocket_servers(&[&first, &second, &recovered]).await?, + )) + .await?; + + assert_eq!( + sampler.sample(sample_request("turn-1")).await?, + r#"{"score":0.25}"# + ); + assert_eq!(first.single_connection().len(), 1); + assert_eq!(second.single_connection().len(), 1); + assert_eq!(recovered.single_connection().len(), 1); + assert_eq!( + recovered.single_handshake().header("authorization"), + Some("Bearer test-api-key".to_owned()) + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn sampler_limits_transient_recovery_attempts() -> Result<()> { + skip_if_no_network!(Ok(())); + + let unavailable = || { + vec![vec![vec![json!({ + "type": "error", + "status": 503, + "error": { + "type": "server_error", + "message": "temporarily unavailable" + } + })]]] + }; + let first = responses::start_websocket_server(unavailable()).await; + let second = responses::start_websocket_server(unavailable()).await; + let third = responses::start_websocket_server(unavailable()).await; + let unused = responses::start_websocket_server(unavailable()).await; + let sampler = LunaSampler::connect(sampler_config( + proxy_websocket_servers(&[&first, &second, &third, &unused]).await?, + )) + .await?; + + let error = sampler + .sample(sample_request("turn-1")) + .await + .expect_err("sampling should stop after the bounded retries"); + + assert!(error.to_string().contains("503")); + assert_eq!(first.single_connection().len(), 1); + assert_eq!(second.single_connection().len(), 1); + assert_eq!(third.single_connection().len(), 1); + assert!(unused.handshakes().is_empty()); + + Ok(()) +}