openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-19 17:28:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Retry transient Guardian sampling failures (#38997)

Browse Source 
## 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
This commit is contained in:
jif
2026-08-17 12:45:50 +00:00
committed by copyberry
parent 3666a46ea3
commit 4885eb6c52
2 changed files with 229 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/ext/guardian-v2/src/sampler.rs
+103 -16
View File
@@ -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<UnauthorizedRecovery>,
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<String, LunaSamplerError> {
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) => {
codex-rs/ext/guardian-v2/src/sampler_tests.rs
+126
View File
@@ -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<String> {
@@ -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(())
}