diff --git a/README.md b/README.md index aa1a856..b554ba3 100644 --- a/README.md +++ b/README.md @@ -513,6 +513,13 @@ Sign in with your **ChatGPT Plus/Pro account**, not an OpenAI API account. The token file includes the extracted `chatgpt_account_id` so the proxy can set the `ChatGPT-Account-Id` header on every upstream call. +When no proxy-owned credential exists, the proxy can bootstrap from the native +Codex CLI's `~/.codex/auth.json`. Treat that as a migration convenience, not a +long-term shared credential: OAuth refresh tokens can rotate, so two concurrent +writers can strand one another. If native Codex and this proxy run on the same +machine, complete `codex auth login` or `codex auth device` once so the proxy has +its own OAuth session. + #### `codex auth device` Same OAuth flow, but for headless machines. Prints a short user code and a URL; @@ -727,7 +734,22 @@ affecting other keys. Codex uses the WebSocket Responses transport by default. Set `CCP_CODEX_TRANSPORT=http` to use the older HTTP SSE transport for debugging or compatibility, or `CCP_CODEX_TRANSPORT=auto` to try WebSocket with HTTP fallback -only when setup fails before a request is sent upstream. +after any retryable WebSocket transport failure. + +`auto` also acts as a completion-gated reliability mode: it buffers a complete +WebSocket response before sending Anthropic events to Claude Code and safely +falls back from WebSocket to HTTP after a retryable WebSocket failure, then +retries statusless transport failures up to three times with jittered exponential +backoff while nothing has reached the client. Retry and exhaustion events are +written as structured `buffered_transport_retry*` log records. This prevents +duplicate tool calls after a partial upstream stream, at the cost of delaying +visible reasoning, text, and tool events until the Codex response finishes. Use +`auto` when unattended completion matters more than live token streaming. A +WebSocket transport failure keeps that proxy instance on HTTP for 60 seconds, +and the complete buffered operation has a 10-minute wall-clock limit. Incoming +request bodies also have a 30-second read deadline so an abandoned upload cannot +leave a proxy request wedged indefinitely. + `CCP_CODEX_PREVIOUS_RESPONSE_ID=1` enables opt-in WebSocket continuation for append-only turns. Continuation keeps in-memory state keyed by Claude Code session id, reuses a session WebSocket while it remains open, and sends diff --git a/src/providers/codex/auth/manager.rs b/src/providers/codex/auth/manager.rs index 2226e14..c5dee38 100644 --- a/src/providers/codex/auth/manager.rs +++ b/src/providers/codex/auth/manager.rs @@ -1,5 +1,8 @@ -use std::sync::{Arc, Mutex}; -use std::time::{SystemTime, UNIX_EPOCH}; +use std::sync::Arc; +#[cfg(test)] +use std::sync::Mutex; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use tokio::sync::Mutex as AsyncMutex; use super::constants::{CLIENT_ID, ISSUER, REFRESH_MARGIN_MS}; use super::jwt::{TokenResponse, extract_account_id, validate_token_response}; @@ -8,14 +11,30 @@ use crate::auth::AuthStorage; pub struct CodexAuthManager> { pub store: CodexTokenStore, - cached: Arc>>, + #[cfg(test)] + test_auth: Arc>>, + refresh_lock: Arc>, + refresh_client: reqwest::Client, + token_endpoint: String, } impl> CodexAuthManager { pub fn new(store: CodexTokenStore) -> Self { + Self::new_with_token_endpoint(store, format!("{ISSUER}/oauth/token")) + } + + fn new_with_token_endpoint(store: CodexTokenStore, token_endpoint: String) -> Self { Self { store, - cached: Arc::new(Mutex::new(None)), + #[cfg(test)] + test_auth: Arc::new(Mutex::new(None)), + refresh_lock: Arc::new(AsyncMutex::new(())), + refresh_client: reqwest::Client::builder() + .connect_timeout(Duration::from_secs(15)) + .timeout(Duration::from_secs(30)) + .build() + .expect("failed to create Codex OAuth refresh client"), + token_endpoint, } } @@ -26,77 +45,89 @@ impl> CodexAuthManager { .as_millis() as u64 } - pub fn get_auth(&self) -> Result { - let cached = { - let guard = self.cached.lock().map_err(|e| anyhow::anyhow!("{e}"))?; - guard.clone() - }; - let stored = match cached { - Some(ref auth) => auth.clone(), - None => { - let loaded = self.store.load_auth()?; - match loaded { - Some(auth) => { - let mut guard = self.cached.lock().map_err(|e| anyhow::anyhow!("{e}"))?; - *guard = Some(auth.clone()); - auth - } - None => { - anyhow::bail!("Not authenticated. Run: claude-code-proxy codex auth login"); - } - } - } - }; + pub async fn get_auth(&self) -> Result { + let stored = self.load_auth()?.ok_or_else(|| { + anyhow::anyhow!("Not authenticated. Run: claude-code-proxy codex auth login") + })?; if stored.expires > Self::now_ms() + REFRESH_MARGIN_MS { return Ok(stored); } - self.refresh_now(&stored) + self.refresh(false, None).await } - pub fn force_refresh(&self) -> Result { - let stored = { - let guard = self.cached.lock().map_err(|e| anyhow::anyhow!("{e}"))?; - guard.clone() - }; - let stored = match stored { - Some(auth) => auth, - None => { - let loaded = self.store.load_auth()?; - loaded.ok_or_else(|| anyhow::anyhow!("Not authenticated"))? - } - }; - self.refresh_now(&stored) + pub async fn force_refresh(&self, rejected_access: &str) -> Result { + self.refresh(true, Some(rejected_access)).await } - fn refresh_now(&self, current: &StoredAuth) -> Result { + fn load_auth(&self) -> Result, anyhow::Error> { + #[cfg(test)] + if let Some(auth) = self + .test_auth + .lock() + .map_err(|e| anyhow::anyhow!("{e}"))? + .clone() + { + return Ok(Some(auth)); + } + + self.store.load_auth() + } + + async fn refresh( + &self, + force: bool, + rejected_access: Option<&str>, + ) -> Result { + let _refresh_guard = self.refresh_lock.lock().await; + + // Reload from durable storage after acquiring the single-flight lock. + // Another request may have rotated and persisted the token while this + // caller was waiting. + let current = self + .load_auth()? + .ok_or_else(|| anyhow::anyhow!("Not authenticated"))?; + + if (!force && current.expires > Self::now_ms() + REFRESH_MARGIN_MS) + || rejected_access.is_some_and(|access| current.access != access) + { + return Ok(current); + } + + self.refresh_now(¤t).await + } + + async fn refresh_now(&self, current: &StoredAuth) -> Result { if current.refresh.is_empty() { anyhow::bail!("No refresh token stored; re-authenticate"); } - let client = reqwest::blocking::Client::new(); let form = [ ("client_id", CLIENT_ID.to_string()), ("grant_type", "refresh_token".to_string()), ("refresh_token", current.refresh.clone()), ]; - let resp = client - .post(format!("{ISSUER}/oauth/token")) + let resp = self + .refresh_client + .post(&self.token_endpoint) .form(&form) .send() + .await .map_err(|e| anyhow::anyhow!("refresh network error: {e}"))?; let status = resp.status().as_u16(); if status == 401 || status == 403 { + if let Some(latest) = self.store.load_auth()? + && latest.access != current.access { - let mut guard = self.cached.lock().map_err(|e| anyhow::anyhow!("{e}"))?; - *guard = None; + return Ok(latest); } let _ = self.store.clear_auth(); let err_msg = resp .text() + .await .unwrap_or_else(|_| "Token refresh unauthorized".to_string()); anyhow::bail!("{err_msg}"); } @@ -107,6 +138,7 @@ impl> CodexAuthManager { let tokens: TokenResponse = resp .json() + .await .map_err(|e| anyhow::anyhow!("failed to parse token response: {e}"))?; validate_token_response(&tokens)?; let account_id = extract_account_id(&tokens).or_else(|| current.account_id.clone()); @@ -118,10 +150,6 @@ impl> CodexAuthManager { account_id, }; self.store.save_auth(next.clone())?; - { - let mut guard = self.cached.lock().map_err(|e| anyhow::anyhow!("{e}"))?; - *guard = Some(next.clone()); - } Ok(next) } @@ -139,37 +167,33 @@ impl> CodexAuthManager { account_id, }; self.store.save_auth(auth.clone())?; - { - let mut guard = self.cached.lock().map_err(|e| anyhow::anyhow!("{e}"))?; - *guard = Some(auth.clone()); - } Ok(auth) } - pub fn set_cached(&self, auth: StoredAuth) { - if let Ok(mut guard) = self.cached.lock() { + #[cfg(test)] + pub fn set_test_auth(&self, auth: StoredAuth) { + if let Ok(mut guard) = self.test_auth.lock() { *guard = Some(auth); } } - - pub fn reset_cache(&self) { - if let Ok(mut guard) = self.cached.lock() { - *guard = None; - } - } } #[cfg(test)] mod tests { use super::*; use crate::auth::InMemoryAuthStore; + use std::io::{Read, Write}; + use std::net::TcpListener; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::thread; fn test_store() -> CodexTokenStore> { CodexTokenStore::new(InMemoryAuthStore::new()) } - #[test] - fn get_auth_returns_stored() { + #[tokio::test] + async fn get_auth_returns_stored() { let store = test_store(); let auth = StoredAuth { access: "test_access".into(), @@ -179,22 +203,120 @@ mod tests { }; store.save_auth(auth.clone()).unwrap(); let manager = CodexAuthManager::new(store); - let result = manager.get_auth().unwrap(); + let result = manager.get_auth().await.unwrap(); assert_eq!(result.access, "test_access"); assert_eq!(result.account_id.as_deref(), Some("acct_1")); } - #[test] - fn get_auth_fails_when_no_auth() { + #[tokio::test] + async fn get_auth_fails_when_no_auth() { let store = test_store(); let manager = CodexAuthManager::new(store); - assert!(manager.get_auth().is_err()); + assert!(manager.get_auth().await.is_err()); assert!( manager .get_auth() + .await .unwrap_err() .to_string() .contains("Not authenticated") ); } + + #[tokio::test(flavor = "multi_thread")] + async fn concurrent_expired_auth_refreshes_once() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let addr = listener.local_addr().unwrap(); + let refreshes = Arc::new(AtomicUsize::new(0)); + let server_refreshes = refreshes.clone(); + let server = thread::spawn(move || { + let (mut stream, _) = listener.accept().unwrap(); + let mut request = [0u8; 4096]; + let read = stream.read(&mut request).unwrap(); + assert!(read > 0); + assert!(String::from_utf8_lossy(&request[..read]).contains("refresh_token=stale")); + server_refreshes.fetch_add(1, Ordering::SeqCst); + + let body = br#"{"access_token":"rotated","refresh_token":"rotated-refresh","expires_in":3600}"#; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n", + body.len() + ); + stream.write_all(response.as_bytes()).unwrap(); + stream.write_all(body).unwrap(); + }); + + let store = test_store(); + store + .save_auth(StoredAuth { + access: "expired".into(), + refresh: "stale".into(), + expires: 0, + account_id: Some("acct_1".into()), + }) + .unwrap(); + let manager = Arc::new(CodexAuthManager::new_with_token_endpoint( + store, + format!("http://{addr}/oauth/token"), + )); + let (first, second) = tokio::join!(manager.get_auth(), manager.get_auth()); + let results = [first.unwrap(), second.unwrap()]; + server.join().unwrap(); + + assert_eq!(refreshes.load(Ordering::SeqCst), 1); + assert!(results.iter().all(|auth| auth.access == "rotated")); + assert!(results.iter().all(|auth| auth.refresh == "rotated-refresh")); + } + + #[tokio::test] + async fn stale_401_reuses_already_rotated_auth() { + let store = test_store(); + store + .save_auth(StoredAuth { + access: "rotated".into(), + refresh: "rotated-refresh".into(), + expires: u64::MAX, + account_id: Some("acct_1".into()), + }) + .unwrap(); + let manager = CodexAuthManager::new_with_token_endpoint( + store, + "http://127.0.0.1:1/should-not-be-called".into(), + ); + + let auth = manager.force_refresh("rejected").await.unwrap(); + assert_eq!(auth.access, "rotated"); + assert_eq!(auth.refresh, "rotated-refresh"); + } + + #[tokio::test] + async fn durable_rotation_and_logout_are_observed_by_shared_manager() { + let store = test_store(); + store + .save_auth(StoredAuth { + access: "first".into(), + refresh: "first-refresh".into(), + expires: u64::MAX, + account_id: Some("acct_1".into()), + }) + .unwrap(); + let manager = CodexAuthManager::new(store); + assert_eq!(manager.get_auth().await.unwrap().access, "first"); + + manager + .store + .save_auth(StoredAuth { + access: "rotated".into(), + refresh: "rotated-refresh".into(), + expires: u64::MAX, + account_id: Some("acct_2".into()), + }) + .unwrap(); + let rotated = manager.get_auth().await.unwrap(); + assert_eq!(rotated.access, "rotated"); + assert_eq!(rotated.account_id.as_deref(), Some("acct_2")); + + manager.store.clear_auth().unwrap(); + assert!(manager.get_auth().await.is_err()); + } } diff --git a/src/providers/codex/client.rs b/src/providers/codex/client.rs index b67bd1a..c2f7d78 100644 --- a/src/providers/codex/client.rs +++ b/src/providers/codex/client.rs @@ -1,7 +1,9 @@ +use std::sync::Mutex; use std::time::{Duration, Instant}; use crate::anthropic::sse::parse_sse_events; use crate::config; +use crate::logging::create_logger; use crate::provider::RequestContext; use crate::retry::{compute_backoff_delay, should_retry_status, sleep}; use crate::traffic::TrafficCapture; @@ -11,6 +13,11 @@ use super::auth::manager::CodexAuthManager; use super::auth::token_store::{DefaultCodexAuthStore, StoredAuth, file_store}; use super::translate::request::ResponsesRequest; +const MAX_BUFFERED_TRANSPORT_RETRIES: u32 = 3; +const MAX_BUFFERED_TRANSPORT_ATTEMPTS: u32 = MAX_BUFFERED_TRANSPORT_RETRIES + 1; +const BUFFERED_REQUEST_TIMEOUT: Duration = Duration::from_secs(600); +const WEBSOCKET_FAILURE_COOLDOWN: Duration = Duration::from_secs(60); + // --------------------------------------------------------------------------- // Errors // --------------------------------------------------------------------------- @@ -206,6 +213,8 @@ pub struct CodexHttpClient { auth_manager: CodexAuthManager, base_url: String, header_timeout_ms: u64, + buffered_timeout: Duration, + websocket_cooldown_until: Mutex>, #[allow(dead_code)] header_timeout_retries: u32, } @@ -222,12 +231,13 @@ impl CodexHttpClient { Self { client: reqwest::Client::builder() .connect_timeout(Duration::from_secs(15)) - .timeout(Duration::from_millis(timeout_ms + 10_000)) .build() .expect("failed to create HTTP client"), auth_manager: CodexAuthManager::new(file_store()), base_url: config::codex_base_url(CODEX_API_ENDPOINT), header_timeout_ms: timeout_ms, + buffered_timeout: BUFFERED_REQUEST_TIMEOUT, + websocket_cooldown_until: Mutex::new(None), header_timeout_retries: 1, } } @@ -242,6 +252,8 @@ impl CodexHttpClient { auth_manager, base_url, header_timeout_ms: 60_000, + buffered_timeout: BUFFERED_REQUEST_TIMEOUT, + websocket_cooldown_until: Mutex::new(None), header_timeout_retries: 1, } } @@ -252,12 +264,15 @@ impl CodexHttpClient { base_url: String, header_timeout_ms: u64, header_timeout_retries: u32, + buffered_timeout: Duration, ) -> Self { Self { client, auth_manager: CodexAuthManager::new(file_store()), base_url, header_timeout_ms, + buffered_timeout, + websocket_cooldown_until: Mutex::new(None), header_timeout_retries, } } @@ -266,16 +281,54 @@ impl CodexHttpClient { &self.auth_manager } + fn resolve_transport( + &self, + transport: crate::config::CodexTransport, + ) -> crate::config::CodexTransport { + use crate::config::CodexTransport; + + if transport != CodexTransport::Auto { + return transport; + } + match *self + .websocket_cooldown_until + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + { + Some(until) if until > Instant::now() => CodexTransport::Http, + _ => CodexTransport::WebSocket, + } + } + + fn trip_websocket_cooldown(&self) { + *self + .websocket_cooldown_until + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = + Some(Instant::now() + WEBSOCKET_FAILURE_COOLDOWN); + } + pub async fn post_codex( &self, body: &ResponsesRequest, ctx: &RequestContext, continuation: Option<&super::continuation::ContinuationCandidate>, + ) -> Result { + self.post_codex_with_transport(body, ctx, continuation, crate::config::codex_transport()) + .await + } + + async fn post_codex_with_transport( + &self, + body: &ResponsesRequest, + ctx: &RequestContext, + continuation: Option<&super::continuation::ContinuationCandidate>, + transport: crate::config::CodexTransport, ) -> Result { use super::continuation::clear_continuation; use crate::config::CodexTransport; - let mut auth = self.auth_manager.get_auth().map_err(|e| CodexError { + let mut auth = self.auth_manager.get_auth().await.map_err(|e| CodexError { status: 401, message: "Auth error".to_string(), detail: Some(e.to_string()), @@ -283,98 +336,89 @@ impl CodexHttpClient { origin: CodexErrorOrigin::Auth, })?; - let transport = crate::config::codex_transport(); - let pool_key = websocket_pool_key(ctx, continuation); + let initial_pool_key = websocket_pool_key(ctx, continuation); if should_reset_websocket_pool(continuation) - && let Some(key) = pool_key + && let Some(key) = initial_pool_key { super::websocket::invalidate_codex_websocket_pool_key(key); } - for transport_attempt in 0..=3u32 { - let result = match transport { - CodexTransport::Http => { - let body_json = serde_json::to_string(body).map_err(|e| CodexError { - status: 500, - message: "Failed to serialize request".to_string(), - detail: Some(e.to_string()), - retry_after: None, - origin: CodexErrorOrigin::Http, - })?; - self.attempt_post_http(&auth, &body_json, ctx, body.client_metadata.is_some()) + let mut active_continuation = continuation; + let mut active_transport = self.resolve_transport(transport); + let mut auth_refresh_attempted = false; + let mut transport_failures = 0u32; + let deadline = Instant::now() + self.buffered_timeout; + loop { + let pool_key = websocket_pool_key(ctx, active_continuation); + let remaining = deadline.saturating_duration_since(Instant::now()); + let result = match tokio::time::timeout(remaining, async { + match active_transport { + CodexTransport::Http => { + let body_json = serde_json::to_string(body).map_err(|e| CodexError { + status: 500, + message: "Failed to serialize request".to_string(), + detail: Some(e.to_string()), + retry_after: None, + origin: CodexErrorOrigin::Http, + })?; + self.attempt_post_http( + &auth, + &body_json, + ctx, + body.client_metadata.is_some(), + ) .await - } - CodexTransport::WebSocket => { - let ws_headers = - build_codex_headers(&auth, ctx, body.client_metadata.is_some())?; - let ws_headers = super::websocket::codex_websocket_headers(&ws_headers); - let ws_body = build_websocket_request(body, continuation); - - super::websocket::codex_websocket_request( - &self.base_url, - &ws_headers, - &ws_body, - ctx, - ctx.traffic.as_deref(), - pool_key, - super::websocket::WEBSOCKET_CONNECT_TIMEOUT_MS, - super::websocket::WEBSOCKET_IDLE_TIMEOUT_MS, - continuation, - ) - .await - } - CodexTransport::Auto => { - let ws_headers = - build_codex_headers(&auth, ctx, body.client_metadata.is_some())?; - let ws_headers = super::websocket::codex_websocket_headers(&ws_headers); - let ws_body = build_websocket_request(body, continuation); - - // Try WebSocket first - let ws_result = super::websocket::codex_websocket_request( - &self.base_url, - &ws_headers, - &ws_body, - ctx, - ctx.traffic.as_deref(), - pool_key, - super::websocket::WEBSOCKET_CONNECT_TIMEOUT_MS, - super::websocket::WEBSOCKET_IDLE_TIMEOUT_MS, - continuation, - ) - .await; - - match ws_result { - Ok(response) => Ok(response), - Err(err) - if err.status == 0 - && err.detail.as_deref() == Some("websocket_pre_request") => - { - // Fall back to HTTP only if WebSocket failed before sending - let body_json = - serde_json::to_string(body).map_err(|e| CodexError { - status: 500, - message: "Failed to serialize request".to_string(), - detail: Some(e.to_string()), - retry_after: None, - origin: CodexErrorOrigin::Http, - })?; - self.attempt_post_http( - &auth, - &body_json, - ctx, - body.client_metadata.is_some(), - ) - .await - } - Err(err) => Err(err), } + CodexTransport::WebSocket => { + let ws_headers = + build_codex_headers(&auth, ctx, body.client_metadata.is_some())?; + let ws_headers = super::websocket::codex_websocket_headers(&ws_headers); + let ws_body = build_websocket_request(body, active_continuation); + + super::websocket::codex_websocket_request( + &self.base_url, + &ws_headers, + &ws_body, + ctx, + ctx.traffic.as_deref(), + pool_key, + super::websocket::WEBSOCKET_CONNECT_TIMEOUT_MS, + super::websocket::WEBSOCKET_IDLE_TIMEOUT_MS, + active_continuation, + ) + .await + } + CodexTransport::Auto => unreachable!("auto transport must be resolved"), + } + }) + .await + { + Ok(result) => result, + Err(_) => { + return Err(CodexError { + status: 504, + message: format!( + "Buffered Codex request exceeded {}ms", + self.buffered_timeout.as_millis() + ), + detail: Some("buffered_request_timeout".to_string()), + retry_after: None, + origin: match active_transport { + CodexTransport::WebSocket => CodexErrorOrigin::WebSocket, + _ => CodexErrorOrigin::Http, + }, + }); } }; - if should_refresh_after_unauthorized(&result, transport_attempt) { - match self.auth_manager.force_refresh() { + if should_refresh_after_unauthorized(&result, auth_refresh_attempted) { + auth_refresh_attempted = true; + match self.auth_manager.force_refresh(&auth.access).await { Ok(new_auth) => { auth = new_auth; + if let Some(key) = pool_key { + super::websocket::invalidate_codex_websocket_pool_key(key); + } continue; } Err(e) => { @@ -389,6 +433,44 @@ impl CodexHttpClient { } } + if let Ok(response) = &result + && let Some(failure) = retryable_buffered_upstream_failure(&response.body) + { + if failure.status != 429 && transport_failures < MAX_BUFFERED_TRANSPORT_RETRIES { + let delay = + compute_backoff_delay(transport_failures, failure.retry_after.as_deref()); + log_buffered_retry( + ctx, + active_transport, + transport_failures + 1, + delay.wait_ms, + failure.status, + "upstream_event", + &failure.message, + ); + transport_failures += 1; + sleep(delay.wait_ms).await; + continue; + } + + if failure.status != 429 { + log_buffered_retry_exhausted( + ctx, + active_transport, + failure.status, + "upstream_event", + &failure.message, + ); + } + return Err(CodexError { + status: failure.status, + message: failure.message, + detail: None, + retry_after: failure.retry_after, + origin: CodexErrorOrigin::Http, + }); + } + match result { Ok(response) if response.status == 401 => { let detail = String::from_utf8_lossy(&response.body).to_string(); @@ -416,12 +498,6 @@ impl CodexHttpClient { .iter() .find(|(k, _)| k.to_lowercase() == "retry-after") .map(|(_, v)| v.clone()); - if transport_attempt < 3 { - let delay = - compute_backoff_delay(transport_attempt, retry_after.as_deref()); - sleep(delay.wait_ms).await; - continue; - } let detail = String::from_utf8_lossy(&response.body).to_string(); return Err(CodexError { status: 429, @@ -431,74 +507,113 @@ impl CodexHttpClient { origin: CodexErrorOrigin::Http, }); } + Ok(response) if should_retry_status(response.status) => { + if transport_failures < MAX_BUFFERED_TRANSPORT_RETRIES { + let retry_after = response + .headers + .iter() + .find(|(key, _)| key.eq_ignore_ascii_case("retry-after")) + .map(|(_, value)| value.as_str()); + let delay = compute_backoff_delay(transport_failures, retry_after); + log_buffered_retry( + ctx, + active_transport, + transport_failures + 1, + delay.wait_ms, + response.status, + "upstream", + "retryable upstream status", + ); + transport_failures += 1; + sleep(delay.wait_ms).await; + continue; + } + log_buffered_retry_exhausted( + ctx, + active_transport, + response.status, + "upstream", + "retryable upstream status", + ); + let retry_after = response + .headers + .iter() + .find(|(key, _)| key.eq_ignore_ascii_case("retry-after")) + .map(|(_, value)| value.clone()); + return Err(CodexError { + status: response.status, + message: format!( + "Upstream returned retryable status {} after buffered retries", + response.status + ), + detail: None, + retry_after, + origin: CodexErrorOrigin::Http, + }); + } Ok(response) => return Ok(response), - Err(err) if should_retry_without_continuation(&err, continuation) => { + Err(err) if should_retry_without_continuation(&err, active_continuation) => { clear_continuation(ctx.session_id.as_deref()); if let Some(key) = pool_key { super::websocket::invalidate_codex_websocket_pool_key(key); } - - let retry_headers = - build_codex_headers(&auth, ctx, body.client_metadata.is_some())?; - let full_body_json = serde_json::to_string(body).map_err(|e| CodexError { - status: 500, - message: "Failed to serialize request".to_string(), - detail: Some(e.to_string()), - retry_after: None, - origin: CodexErrorOrigin::Http, - })?; - - match transport { - CodexTransport::Http => { - return self - .attempt_post_http( - &auth, - &full_body_json, - ctx, - body.client_metadata.is_some(), - ) - .await; - } - CodexTransport::WebSocket | CodexTransport::Auto => { - let ws_retry_headers = - super::websocket::codex_websocket_headers(&retry_headers); - let ws_retry_body = build_websocket_request(body, None); - return super::websocket::codex_websocket_request( - &self.base_url, - &ws_retry_headers, - &ws_retry_body, - ctx, - ctx.traffic.as_deref(), - pool_key, - super::websocket::WEBSOCKET_CONNECT_TIMEOUT_MS, - super::websocket::WEBSOCKET_IDLE_TIMEOUT_MS, - None, - ) - .await; - } - } + active_continuation = None; + continue; } Err(err) => { // Determine if retryable let retryable = is_retryable_transport_error(&err); - if retryable && transport_attempt < 3 { + if retryable + && transport == CodexTransport::Auto + && active_transport == CodexTransport::WebSocket + && transport_failures < MAX_BUFFERED_TRANSPORT_RETRIES + { + log_buffered_retry( + ctx, + active_transport, + transport_failures + 1, + 0, + err.status, + codex_error_origin_name(err.origin), + &err.message, + ); + transport_failures += 1; + if let Some(key) = pool_key { + super::websocket::invalidate_codex_websocket_pool_key(key); + } + self.trip_websocket_cooldown(); + active_transport = CodexTransport::Http; + continue; + } + if retryable && transport_failures < MAX_BUFFERED_TRANSPORT_RETRIES { let delay = - compute_backoff_delay(transport_attempt, err.retry_after.as_deref()); + compute_backoff_delay(transport_failures, err.retry_after.as_deref()); + log_buffered_retry( + ctx, + active_transport, + transport_failures + 1, + delay.wait_ms, + err.status, + codex_error_origin_name(err.origin), + &err.message, + ); + transport_failures += 1; sleep(delay.wait_ms).await; continue; } + if retryable { + log_buffered_retry_exhausted( + ctx, + active_transport, + err.status, + codex_error_origin_name(err.origin), + &err.message, + ); + } return Err(err); } } } - - Err(CodexError { - status: 0, - message: "Max retries exceeded".to_string(), - detail: None, - retry_after: None, - origin: CodexErrorOrigin::Http, - }) } pub async fn stream_codex_websocket_events( @@ -507,7 +622,7 @@ impl CodexHttpClient { ctx: &RequestContext, continuation: Option<&super::continuation::ContinuationCandidate>, ) -> Result { - let auth = self.auth_manager.get_auth().map_err(|e| CodexError { + let auth = self.auth_manager.get_auth().await.map_err(|e| CodexError { status: 401, message: "Auth error".to_string(), detail: Some(e.to_string()), @@ -624,7 +739,7 @@ impl CodexHttpClient { let send_fut = req_builder.body(body_json.to_string()).send(); let header_timeout_dur = Duration::from_millis(self.header_timeout_ms); - let resp = tokio::time::timeout(header_timeout_dur, send_fut) + let mut resp = tokio::time::timeout(header_timeout_dur, send_fut) .await .map_err(|_| CodexError { status: 0, @@ -663,7 +778,46 @@ impl CodexHttpClient { .map(|(k, v)| (k.to_string(), v.to_str().unwrap_or("").to_string())) .collect(); - let body_bytes = resp.bytes().await.unwrap_or_default().to_vec(); + let mut body_bytes = Vec::new(); + loop { + let chunk = match tokio::time::timeout( + Duration::from_millis(super::websocket::WEBSOCKET_IDLE_TIMEOUT_MS), + resp.chunk(), + ) + .await + { + Err(_) => { + return Err(CodexError { + status: 0, + message: format!( + "Timed out waiting {}ms for the next Codex response body chunk", + super::websocket::WEBSOCKET_IDLE_TIMEOUT_MS + ), + detail: Some("http_response_body".to_string()), + retry_after: None, + origin: CodexErrorOrigin::Http, + }); + } + Ok(Ok(chunk)) => chunk, + Ok(Err(e)) => { + return Err(CodexError { + status: 0, + message: format!("Transport error reading Codex response body: {e}"), + detail: Some("http_response_body".to_string()), + retry_after: None, + origin: CodexErrorOrigin::Http, + }); + } + }; + + let Some(chunk) = chunk else { + break; + }; + body_bytes.extend_from_slice(&chunk); + if has_terminal_event(&body_bytes) { + break; + } + } if let Some(traffic) = ctx.traffic.as_deref() { write_upstream_response_capture( @@ -683,6 +837,21 @@ impl CodexHttpClient { } } +fn has_terminal_event(body: &[u8]) -> bool { + parse_sse_events(body).into_iter().any(|event| { + let Ok(payload) = serde_json::from_str::(&event.data) else { + return false; + }; + super::is_codex_terminal_event(&payload) + || (payload.get("type").and_then(|value| value.as_str()) == Some("codex.rate_limits") + && payload + .get("rate_limits") + .and_then(|limits| limits.get("limit_reached")) + .and_then(|value| value.as_bool()) + == Some(true)) + }) +} + fn write_codex_http_request_capture( traffic: &TrafficCapture, url: &str, @@ -802,15 +971,162 @@ fn summarize_json_request_size(body: &serde_json::Value, body_json: &str) -> ser } fn is_retryable_transport_error(err: &CodexError) -> bool { - (err.detail.as_deref() == Some("websocket_pre_request") && should_retry_status(err.status)) - || (err.status == 0 - && (err.message.contains("Timed out waiting") - || err.message.contains("Transport error") - || err.message.contains("connection reset") - || err.message.contains("connection closed") - || err.message.contains("timed out") - || err.message.contains("econnreset") - || err.message.contains("etimedout"))) + if err.detail.as_deref() == Some("websocket_pre_request") { + return err.status == 0 || should_retry_status(err.status); + } + if err.status != 0 { + return false; + } + + let message = err.message.to_ascii_lowercase(); + [ + "timed out", + "timeout", + "transport error", + "websocket send error", + "websocket stream error", + "websocket connection closed", + "failed to ping pooled connection", + "connection reset", + "connection closed", + "econnreset", + "etimedout", + "broken pipe", + "epipe", + "unexpected eof", + ] + .iter() + .any(|needle| message.contains(needle)) +} + +struct BufferedUpstreamFailure { + status: u16, + message: String, + retry_after: Option, +} + +fn retryable_buffered_upstream_failure(body: &[u8]) -> Option { + for event in parse_sse_events(body) { + if event.data == "[DONE]" { + continue; + } + let Ok(payload) = serde_json::from_str::(&event.data) else { + continue; + }; + let event_type = payload.get("type").and_then(|value| value.as_str()); + if !matches!( + event_type, + Some("response.failed" | "response.error" | "error" | "codex.rate_limits") + ) { + continue; + } + + let message = payload + .get("response") + .and_then(|response| response.get("error")) + .and_then(|error| error.get("message")) + .and_then(|value| value.as_str()) + .or_else(|| { + payload + .get("error") + .and_then(|error| error.get("message")) + .and_then(|value| value.as_str()) + }) + .unwrap_or("Retryable upstream failure") + .to_string(); + if !super::retryable_live_start_payload(&payload, &message) { + continue; + } + + let status = payload + .get("status") + .or_else(|| payload.get("status_code")) + .and_then(|value| value.as_u64()) + .or_else(|| { + payload + .get("error") + .and_then(|error| error.get("status")) + .and_then(|value| value.as_u64()) + }) + .or_else(|| { + payload + .get("response") + .and_then(|response| response.get("error")) + .and_then(|error| error.get("status")) + .and_then(|value| value.as_u64()) + }) + .and_then(|status| u16::try_from(status).ok()) + .unwrap_or_else(|| { + if event_type == Some("codex.rate_limits") { + 429 + } else if message.to_ascii_lowercase().contains("overloaded") { + 529 + } else { + 503 + } + }); + + return Some(BufferedUpstreamFailure { + status, + message, + retry_after: super::retry_after_from_live_payload(&payload), + }); + } + + None +} + +fn codex_error_origin_name(origin: CodexErrorOrigin) -> &'static str { + match origin { + CodexErrorOrigin::Http => "http", + CodexErrorOrigin::WebSocket => "websocket", + CodexErrorOrigin::Auth => "auth", + } +} + +fn log_buffered_retry( + ctx: &RequestContext, + transport: crate::config::CodexTransport, + failed_attempt: u32, + delay_ms: u64, + status: u16, + origin: &str, + reason: &str, +) { + let mut fields = serde_json::Map::new(); + fields.insert("reqId".into(), serde_json::json!(ctx.req_id)); + fields.insert("transport".into(), serde_json::json!(transport.as_str())); + fields.insert("failedAttempt".into(), serde_json::json!(failed_attempt)); + fields.insert("nextAttempt".into(), serde_json::json!(failed_attempt + 1)); + fields.insert( + "maxAttempts".into(), + serde_json::json!(MAX_BUFFERED_TRANSPORT_ATTEMPTS), + ); + fields.insert("delayMs".into(), serde_json::json!(delay_ms)); + fields.insert("status".into(), serde_json::json!(status)); + fields.insert("origin".into(), serde_json::json!(origin)); + fields.insert("reason".into(), serde_json::json!(reason)); + create_logger("codex").warn("buffered_transport_retry", Some(fields)); +} + +fn log_buffered_retry_exhausted( + ctx: &RequestContext, + transport: crate::config::CodexTransport, + status: u16, + origin: &str, + reason: &str, +) { + let mut fields = serde_json::Map::new(); + fields.insert("reqId".into(), serde_json::json!(ctx.req_id)); + fields.insert("transport".into(), serde_json::json!(transport.as_str())); + fields.insert( + "attempts".into(), + serde_json::json!(MAX_BUFFERED_TRANSPORT_ATTEMPTS), + ); + fields.insert("status".into(), serde_json::json!(status)); + fields.insert("origin".into(), serde_json::json!(origin)); + fields.insert("reason".into(), serde_json::json!(reason)); + create_logger("codex").warn("buffered_transport_retry_exhausted", Some(fields)); } fn is_retryable_reqwest_error(err: &reqwest::Error) -> bool { @@ -827,9 +1143,9 @@ fn is_retryable_reqwest_error(err: &reqwest::Error) -> bool { fn should_refresh_after_unauthorized( result: &Result, - transport_attempt: u32, + auth_refresh_attempted: bool, ) -> bool { - if transport_attempt != 0 { + if auth_refresh_attempted { return false; } match result { @@ -885,6 +1201,62 @@ fn should_reset_websocket_pool( #[cfg(test)] mod tests { use super::*; + use futures_util::{SinkExt, StreamExt}; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::{Arc, Mutex}; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpListener; + use tokio_tungstenite::{accept_async, tungstenite::Message}; + + fn buffered_test_client(base_url: String) -> CodexHttpClient { + let client = CodexHttpClient::new_for_test( + reqwest::Client::new(), + base_url, + 1_000, + 1, + Duration::from_secs(5), + ); + client.auth_manager().set_test_auth(StoredAuth { + access: "test_access".into(), + refresh: "test_refresh".into(), + expires: u64::MAX, + account_id: Some("test_account".into()), + }); + client + } + + fn buffered_test_context(req_id: &str) -> RequestContext { + RequestContext { + req_id: req_id.into(), + session_id: None, + session_seq: None, + provider: "codex".into(), + traffic: None, + monitor: None, + } + } + + fn buffered_test_request() -> ResponsesRequest { + ResponsesRequest { + model: "gpt-5.6-sol".into(), + instructions: None, + input: vec![], + tools: None, + tool_choice: None, + store: false, + stream: true, + parallel_tool_calls: true, + include: None, + client_metadata: None, + service_tier: None, + prompt_cache_key: None, + text: super::super::translate::request::ResponsesText { + verbosity: None, + format: None, + }, + reasoning: None, + } + } #[test] fn codex_error_display() { @@ -926,6 +1298,636 @@ mod tests { assert!(!is_retryable_transport_error(&err)); } + #[test] + fn statusless_websocket_failures_are_retryable() { + for message in [ + "WebSocket stream error: IO error: Connection reset by peer", + "WebSocket connection closed before terminal Codex response event", + "WebSocket idle timeout after 300000ms", + "WebSocket response start timeout after 300000ms", + "WebSocket protocol error: Connection reset without closing handshake", + ] { + let err = CodexError { + status: 0, + message: message.to_string(), + detail: None, + retry_after: None, + origin: CodexErrorOrigin::WebSocket, + }; + + assert!( + is_retryable_transport_error(&err), + "expected retryable WebSocket error: {message}" + ); + } + } + + #[test] + fn statusless_http_transport_matching_is_case_insensitive() { + let err = CodexError { + status: 0, + message: "Transport error: Connection reset by peer".to_string(), + detail: None, + retry_after: None, + origin: CodexErrorOrigin::Http, + }; + + assert!(is_retryable_transport_error(&err)); + } + + #[test] + fn statusful_websocket_failures_are_not_replayed_as_transport_errors() { + let err = CodexError { + status: 400, + message: "Connection reset by peer".to_string(), + detail: None, + retry_after: None, + origin: CodexErrorOrigin::WebSocket, + }; + + assert!(!is_retryable_transport_error(&err)); + } + + #[test] + fn statusless_websocket_semantic_and_config_errors_are_not_retryable() { + for message in [ + "Previous response not found", + "WebSocket binary frames not supported", + "Unsupported WebSocket URL scheme: ftp", + "Failed to build WebSocket request: invalid header", + ] { + let err = CodexError { + status: 0, + message: message.to_string(), + detail: None, + retry_after: None, + origin: CodexErrorOrigin::WebSocket, + }; + + assert!( + !is_retryable_transport_error(&err), + "expected non-retryable WebSocket error: {message}" + ); + } + } + + #[test] + fn buffered_rate_limit_event_preserves_429_and_retry_after() { + let body = b"data: {\"type\":\"codex.rate_limits\",\"rate_limits\":{\"limit_reached\":true,\"primary\":{\"reset_after_seconds\":7}}}\n\n"; + let failure = retryable_buffered_upstream_failure(body).unwrap(); + assert_eq!(failure.status, 429); + assert_eq!(failure.retry_after.as_deref(), Some("7")); + } + + #[tokio::test] + async fn buffered_auto_falls_back_to_http_before_releasing_partial_output() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let attempts = Arc::new(AtomicUsize::new(0)); + let server_attempts = attempts.clone(); + + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + { + let mut websocket = accept_async(stream).await.unwrap(); + let request = websocket.next().await.unwrap().unwrap(); + assert!(request.is_text()); + server_attempts.fetch_add(1, Ordering::SeqCst); + websocket + .send(Message::Text( + r#"{"type":"response.output_text.delta","delta":"discard me"}"#.into(), + )) + .await + .unwrap(); + } + + let (mut stream, _) = listener.accept().await.unwrap(); + let mut request = vec![0u8; 16 * 1024]; + let read = stream.read(&mut request).await.unwrap(); + assert!(request[..read].starts_with(b"POST ")); + server_attempts.fetch_add(1, Ordering::SeqCst); + let body = b"data: {\"type\":\"response.output_text.delta\",\"delta\":\"keep me\"}\n\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_ok\",\"usage\":{}}}\n\n"; + let headers = format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n", + body.len() + ); + stream.write_all(headers.as_bytes()).await.unwrap(); + stream.write_all(body).await.unwrap(); + stream.shutdown().await.unwrap(); + }); + + let client = buffered_test_client(format!("http://{addr}/backend-api/codex/responses")); + let ctx = buffered_test_context("buffered-retry"); + let request = buffered_test_request(); + + let response = tokio::time::timeout( + Duration::from_secs(6), + client.post_codex_with_transport( + &request, + &ctx, + None, + crate::config::CodexTransport::Auto, + ), + ) + .await + .expect("buffered retry timed out") + .expect("buffered retry failed"); + server.await.unwrap(); + + let body = String::from_utf8(response.body).unwrap(); + assert_eq!(attempts.load(Ordering::SeqCst), 2); + assert_eq!( + client.resolve_transport(crate::config::CodexTransport::Auto), + crate::config::CodexTransport::Http + ); + assert!(body.contains("keep me")); + assert!(body.contains("response.completed")); + assert!(!body.contains("discard me")); + } + + #[tokio::test] + async fn buffered_http_retries_body_reset_after_headers() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let attempts = Arc::new(AtomicUsize::new(0)); + let server_attempts = attempts.clone(); + + let server = tokio::spawn(async move { + for attempt in 0..2 { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut request = vec![0u8; 16 * 1024]; + let read = stream.read(&mut request).await.unwrap(); + assert!(read > 0); + assert!(request[..read].starts_with(b"POST ")); + server_attempts.fetch_add(1, Ordering::SeqCst); + + if attempt == 0 { + let partial = b"data: discard me\n\n"; + let headers = format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n", + partial.len() + 128 + ); + stream.write_all(headers.as_bytes()).await.unwrap(); + stream.write_all(partial).await.unwrap(); + stream.shutdown().await.unwrap(); + continue; + } + + let body = b"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_http_ok\",\"usage\":{}}}\n\n"; + let headers = format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n", + body.len() + ); + stream.write_all(headers.as_bytes()).await.unwrap(); + stream.write_all(body).await.unwrap(); + stream.shutdown().await.unwrap(); + } + }); + + let client = buffered_test_client(format!("http://{addr}/responses")); + let response = tokio::time::timeout( + Duration::from_secs(8), + client.post_codex_with_transport( + &buffered_test_request(), + &buffered_test_context("http-body-reset"), + None, + crate::config::CodexTransport::Http, + ), + ) + .await + .expect("buffered HTTP retry timed out") + .expect("buffered HTTP retry failed"); + server.await.unwrap(); + + let body = String::from_utf8(response.body).unwrap(); + assert_eq!(attempts.load(Ordering::SeqCst), 2); + assert!(body.contains("resp_http_ok")); + assert!(!body.contains("discard me")); + } + + #[tokio::test] + async fn buffered_http_stops_at_terminal_before_invalid_framing() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + let server = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut request = vec![0u8; 16 * 1024]; + let read = stream.read(&mut request).await.unwrap(); + assert!(request[..read].starts_with(b"POST ")); + + let body = b"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_complete\",\"usage\":{}}}\n\n"; + let headers = format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n", + body.len() + 128 + ); + stream.write_all(headers.as_bytes()).await.unwrap(); + stream.write_all(body).await.unwrap(); + stream.shutdown().await.unwrap(); + }); + + let client = buffered_test_client(format!("http://{addr}/responses")); + let response = client + .post_codex_with_transport( + &buffered_test_request(), + &buffered_test_context("http-complete-framing-error"), + None, + crate::config::CodexTransport::Http, + ) + .await + .expect("complete response should survive invalid HTTP framing"); + server.await.unwrap(); + + assert!( + String::from_utf8(response.body) + .unwrap() + .contains("resp_complete") + ); + } + + #[tokio::test] + async fn buffered_http_preserves_terminal_rate_limit_before_invalid_framing() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + let server = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut request = vec![0u8; 16 * 1024]; + assert!(stream.read(&mut request).await.unwrap() > 0); + + let body = b"data: {\"type\":\"codex.rate_limits\",\"rate_limits\":{\"limit_reached\":true,\"primary\":{\"reset_after_seconds\":7}}}\n\n"; + let headers = format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n", + body.len() + 128 + ); + stream.write_all(headers.as_bytes()).await.unwrap(); + stream.write_all(body).await.unwrap(); + }); + + let client = buffered_test_client(format!("http://{addr}/responses")); + let result = client + .post_codex_with_transport( + &buffered_test_request(), + &buffered_test_context("http-rate-limit-framing-error"), + None, + crate::config::CodexTransport::Http, + ) + .await; + server.await.unwrap(); + + let Err(error) = result else { + panic!("terminal rate limit should fail fast"); + }; + assert_eq!(error.status, 429); + assert_eq!(error.retry_after.as_deref(), Some("7")); + } + + #[test] + fn terminal_event_detection_matches_codex_lifecycle() { + for event_type in [ + "response.completed", + "response.incomplete", + "response.failed", + "response.error", + "response.done", + "error", + ] { + let body = format!("data: {{\"type\":\"{event_type}\"}}\n\n"); + assert!(has_terminal_event(body.as_bytes()), "{event_type}"); + } + + for body in [ + b"data: {\"type\":\"response.output_text.delta\",\"delta\":\"response.completed\"}\n\n" + .as_slice(), + b"data: {\"nested\":{\"type\":\"response.completed\"}}\n\n".as_slice(), + b"data: {\"type\":\"codex.rate_limits\",\"rate_limits\":{\"limit_reached\":false}}\n\n" + .as_slice(), + b"data: not-json\n\n".as_slice(), + ] { + assert!(!has_terminal_event(body)); + } + + let rate_limit = + b"data: {\"type\":\"codex.rate_limits\",\"rate_limits\":{\"limit_reached\":true}}\n\n"; + assert!(has_terminal_event(rate_limit)); + } + + #[tokio::test] + async fn buffered_transport_retries_503_before_success() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let attempts = Arc::new(AtomicUsize::new(0)); + let server_attempts = attempts.clone(); + + let server = tokio::spawn(async move { + for attempt in 0..2 { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut request = vec![0u8; 16 * 1024]; + let read = stream.read(&mut request).await.unwrap(); + assert!(read > 0); + server_attempts.fetch_add(1, Ordering::SeqCst); + + let (status, retry_after, body): (&str, &str, &[u8]) = if attempt == 0 { + ("503 Service Unavailable", "retry-after: 0\r\n", b"retry me") + } else { + ("200 OK", "", b"data: keep me\n\n") + }; + let response = format!( + "HTTP/1.1 {status}\r\ncontent-length: {}\r\n{retry_after}connection: close\r\n\r\n", + body.len() + ); + stream.write_all(response.as_bytes()).await.unwrap(); + stream.write_all(body).await.unwrap(); + stream.shutdown().await.unwrap(); + } + }); + + let client = buffered_test_client(format!("http://{addr}/responses")); + let response = client + .post_codex_with_transport( + &buffered_test_request(), + &buffered_test_context("retry-503"), + None, + crate::config::CodexTransport::Http, + ) + .await + .expect("503 retry failed"); + server.await.unwrap(); + + assert_eq!(attempts.load(Ordering::SeqCst), 2); + assert_eq!(response.status, 200); + assert_eq!(response.body, b"data: keep me\n\n"); + } + + #[tokio::test] + async fn buffered_rate_limit_fails_fast() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + let server = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let mut request = vec![0u8; 16 * 1024]; + assert!(stream.read(&mut request).await.unwrap() > 0); + stream + .write_all( + b"HTTP/1.1 429 Too Many Requests\r\ncontent-length: 7\r\nretry-after: 60\r\nconnection: close\r\n\r\nlimited", + ) + .await + .unwrap(); + }); + + let client = buffered_test_client(format!("http://{addr}/responses")); + let result = tokio::time::timeout( + Duration::from_secs(1), + client.post_codex_with_transport( + &buffered_test_request(), + &buffered_test_context("rate-limit"), + None, + crate::config::CodexTransport::Http, + ), + ) + .await + .expect("rate limit was retried"); + let Err(error) = result else { + panic!("rate limit should fail"); + }; + server.await.unwrap(); + + assert_eq!(error.status, 429); + assert_eq!(error.retry_after.as_deref(), Some("60")); + } + + #[tokio::test] + async fn buffered_request_has_wall_clock_deadline() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_async(stream).await.unwrap(); + let _ = websocket.next().await; + tokio::time::sleep(Duration::from_secs(1)).await; + }); + let client = CodexHttpClient::new_for_test( + reqwest::Client::new(), + format!("http://{addr}/responses"), + 1_000, + 1, + Duration::from_millis(50), + ); + client.auth_manager().set_test_auth(StoredAuth { + access: "test_access".into(), + refresh: "test_refresh".into(), + expires: u64::MAX, + account_id: Some("test_account".into()), + }); + + let result = client + .post_codex_with_transport( + &buffered_test_request(), + &buffered_test_context("deadline"), + None, + crate::config::CodexTransport::Auto, + ) + .await; + let Err(error) = result else { + panic!("deadline should fail"); + }; + server.abort(); + + assert_eq!(error.status, 504); + assert_eq!(error.detail.as_deref(), Some("buffered_request_timeout")); + } + + #[tokio::test] + async fn exhausted_retry_budget_does_not_add_http_fallback() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let attempts = Arc::new(AtomicUsize::new(0)); + let server_attempts = attempts.clone(); + let server = tokio::spawn(async move { + for attempt in 0..=MAX_BUFFERED_TRANSPORT_RETRIES { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_async(stream).await.unwrap(); + let _ = websocket.next().await; + server_attempts.fetch_add(1, Ordering::SeqCst); + if attempt < MAX_BUFFERED_TRANSPORT_RETRIES { + websocket + .send(Message::Text( + r#"{"type":"response.failed","response":{"error":{"type":"overloaded_error","message":"overloaded","status":529,"retry_after_seconds":0}}}"#.into(), + )) + .await + .unwrap(); + } + } + }); + let client = buffered_test_client(format!("http://{addr}/responses")); + + let result = client + .post_codex_with_transport( + &buffered_test_request(), + &buffered_test_context("retry-boundary"), + None, + crate::config::CodexTransport::Auto, + ) + .await; + let Err(error) = result else { + panic!("exhausted retry budget should fail"); + }; + server.await.unwrap(); + + assert_eq!(error.status, 0); + assert_eq!( + attempts.load(Ordering::SeqCst), + MAX_BUFFERED_TRANSPORT_ATTEMPTS as usize + ); + } + + #[tokio::test] + async fn buffered_auto_retries_overloaded_response_failed_event() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let attempts = Arc::new(AtomicUsize::new(0)); + let server_attempts = attempts.clone(); + + let server = tokio::spawn(async move { + for attempt in 0..2 { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_async(stream).await.unwrap(); + let request = websocket.next().await.unwrap().unwrap(); + assert!(request.is_text()); + server_attempts.fetch_add(1, Ordering::SeqCst); + + if attempt == 0 { + websocket + .send(Message::Text( + r#"{"type":"response.failed","response":{"error":{"type":"overloaded_error","message":"overloaded","status":529,"retry_after_seconds":0}}}"# + .into(), + )) + .await + .unwrap(); + continue; + } + + websocket + .send(Message::Text( + r#"{"type":"response.output_text.delta","delta":"keep me"}"#.into(), + )) + .await + .unwrap(); + websocket + .send(Message::Text( + r#"{"type":"response.completed","response":{"id":"resp_after_overload","usage":{}}}"# + .into(), + )) + .await + .unwrap(); + } + }); + + let client = buffered_test_client(format!("http://{addr}/responses")); + let response = client + .post_codex_with_transport( + &buffered_test_request(), + &buffered_test_context("response-failed-overload"), + None, + crate::config::CodexTransport::Auto, + ) + .await + .expect("overloaded response.failed retry failed"); + server.await.unwrap(); + + assert_eq!(attempts.load(Ordering::SeqCst), 2); + assert!( + String::from_utf8(response.body) + .unwrap() + .contains("keep me") + ); + } + + #[tokio::test] + async fn buffered_continuation_fallback_keeps_transport_retry_budget() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let requests = Arc::new(Mutex::new(Vec::new())); + let server_requests = requests.clone(); + + let server = tokio::spawn(async move { + for attempt in 0..2 { + let (stream, _) = listener.accept().await.unwrap(); + let mut websocket = accept_async(stream).await.unwrap(); + let request = websocket.next().await.unwrap().unwrap(); + let request: serde_json::Value = + serde_json::from_str(request.to_text().unwrap()).unwrap(); + server_requests.lock().unwrap().push(request); + + if attempt == 0 { + websocket + .send(Message::Text( + r#"{"type":"error","error":{"code":"previous_response_not_found","message":"previous response not found","status":400}}"# + .into(), + )) + .await + .unwrap(); + continue; + } + websocket + .send(Message::Text( + r#"{"type":"response.output_text.delta","delta":"discard me"}"#.into(), + )) + .await + .unwrap(); + } + + let (mut stream, _) = listener.accept().await.unwrap(); + let mut request = vec![0u8; 16 * 1024]; + let read = stream.read(&mut request).await.unwrap(); + assert!(request[..read].starts_with(b"POST ")); + let body = b"data: {\"type\":\"response.output_text.delta\",\"delta\":\"keep me\"}\n\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_full_context\",\"usage\":{}}}\n\n"; + let headers = format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n", + body.len() + ); + stream.write_all(headers.as_bytes()).await.unwrap(); + stream.write_all(body).await.unwrap(); + }); + + let client = buffered_test_client(format!("http://{addr}/responses")); + let mut ctx = buffered_test_context("continuation-then-reset"); + ctx.session_id = Some("session-continuation".into()); + let continuation = super::super::continuation::ContinuationCandidate { + previous_response_id: Some("resp_previous".into()), + input_delta: None, + input_delta_count: 1, + disabled_reason: None, + }; + let response = tokio::time::timeout( + Duration::from_secs(10), + client.post_codex_with_transport( + &buffered_test_request(), + &ctx, + Some(&continuation), + crate::config::CodexTransport::Auto, + ), + ) + .await + .expect("continuation fallback retry timed out") + .expect("continuation fallback retry failed"); + server.await.unwrap(); + + let requests = requests.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert_eq!( + requests[0]["previous_response_id"], + serde_json::json!("resp_previous") + ); + assert!(requests[1].get("previous_response_id").is_none()); + assert!( + String::from_utf8(response.body) + .unwrap() + .contains("keep me") + ); + } + #[test] fn codex_headers_include_session_and_beta() { let auth = StoredAuth { @@ -1182,13 +2184,13 @@ mod tests { origin: CodexErrorOrigin::WebSocket, }); - assert!(should_refresh_after_unauthorized(&http_unauthorized, 0)); + assert!(should_refresh_after_unauthorized(&http_unauthorized, false)); assert!(should_refresh_after_unauthorized( &websocket_unauthorized, - 0 + false )); - assert!(!should_refresh_after_unauthorized(&forbidden, 0)); - assert!(!should_refresh_after_unauthorized(&http_unauthorized, 1)); + assert!(!should_refresh_after_unauthorized(&forbidden, false)); + assert!(!should_refresh_after_unauthorized(&http_unauthorized, true)); } #[test] diff --git a/src/providers/codex/mod.rs b/src/providers/codex/mod.rs index f775f8b..9a46f70 100644 --- a/src/providers/codex/mod.rs +++ b/src/providers/codex/mod.rs @@ -46,7 +46,9 @@ use self::translate::stream::translate_stream_bytes_with_traffic; // Provider // --------------------------------------------------------------------------- -pub struct CodexProvider; +pub struct CodexProvider { + client: Arc, +} impl Default for CodexProvider { fn default() -> Self { @@ -56,7 +58,9 @@ impl Default for CodexProvider { impl CodexProvider { pub fn new() -> Self { - Self + Self { + client: Arc::new(CodexHttpClient::new()), + } } } @@ -131,7 +135,7 @@ impl Provider for CodexProvider { ); // Post to upstream with continuation - let client = Arc::new(CodexHttpClient::new()); + let client = self.client.clone(); if let Some(monitor) = ctx.monitor.as_ref() { monitor.upstream_started(&ctx.req_id); } @@ -817,6 +821,22 @@ fn map_codex_error_to_response(err: &client::CodexError) -> Response { let headers = [(http::header::RETRY_AFTER, retry_after)]; (headers, resp).into_response() } + status @ (500 | 502 | 503 | 504 | 529) => { + let response = json_error( + StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_GATEWAY), + if status == 529 { + "overloaded_error" + } else { + "api_error" + }, + codex_error_message(err), + ); + if let Some(retry_after) = err.retry_after.as_deref() { + ([(http::header::RETRY_AFTER, retry_after)], response).into_response() + } else { + response + } + } _ => json_error( StatusCode::BAD_GATEWAY, "api_error", @@ -1095,6 +1115,32 @@ mod tests { ); } + #[tokio::test] + async fn exhausted_retryable_status_is_preserved() { + let err = client::CodexError { + status: 503, + message: "Upstream unavailable after buffered retries".to_string(), + detail: None, + retry_after: Some("7".to_string()), + origin: client::CodexErrorOrigin::Http, + }; + + let response = map_codex_error_to_response(&err); + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + assert_eq!( + response + .headers() + .get(http::header::RETRY_AFTER) + .and_then(|value| value.to_str().ok()), + Some("7") + ); + let body = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + let body: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(body["error"]["type"], "api_error"); + } + #[test] fn live_start_payload_retry_detection_covers_rate_limit_and_overload() { assert!(retryable_live_start_payload( diff --git a/src/retry.rs b/src/retry.rs index 48b6196..9b61873 100644 --- a/src/retry.rs +++ b/src/retry.rs @@ -1,5 +1,7 @@ use std::time::Duration; +use rand::Rng; + pub const RETRY_INITIAL_DELAY_MS: u64 = 2000; pub const RETRY_MAX_DELAY_MS: u64 = 30_000; pub const RETRY_BACKOFF_FACTOR: u64 = 2; @@ -31,8 +33,10 @@ pub fn compute_backoff_delay(attempt: u32, retry_after: Option<&str>) -> Backoff if exp > RETRY_MAX_DELAY_MS { exp = RETRY_MAX_DELAY_MS; } - let jitter = exp / 2; - let wait_ms = (exp / 2) + (jitter / 2); + // Concurrent agents should not wake and retry in lockstep. Pick uniformly + // from the latter half of the exponential window while preserving the + // existing lower and upper bounds. + let wait_ms = rand::thread_rng().gen_range((exp / 2)..=exp); BackoffOutcome { wait_ms, exceeds_budget: false, diff --git a/src/server.rs b/src/server.rs index dce16f7..45a7ee5 100644 --- a/src/server.rs +++ b/src/server.rs @@ -23,7 +23,7 @@ use std::future::Future; use std::io::Write; use std::path::{Path, PathBuf}; use std::sync::Arc; -use std::time::Instant; +use std::time::{Duration, Instant}; use tokio::net::TcpListener; use uuid::Uuid; @@ -154,9 +154,14 @@ async fn dispatch_request( } let request_guard = RequestMonitorGuard::new(state.monitor.clone(), req_id.clone()); let now = current_millis(); - let body_bytes = match axum::body::to_bytes(req.into_body(), usize::MAX).await { - Ok(bytes) => bytes, - Err(err) => { + let body_bytes = match tokio::time::timeout( + Duration::from_secs(30), + axum::body::to_bytes(req.into_body(), usize::MAX), + ) + .await + { + Ok(Ok(bytes)) => bytes, + Ok(Err(err)) => { let response = json_error( StatusCode::BAD_REQUEST, "invalid_request_error", @@ -196,6 +201,31 @@ async fn dispatch_request( ); return response; } + Err(_) => { + let response = json_error( + StatusCode::REQUEST_TIMEOUT, + "api_error", + "Timed out reading request body", + ); + log_request_completed( + &log, + RequestLogContext { + req_id: &req_id, + provider: None, + model: None, + count_tokens, + status: response.status(), + started_at, + }, + ); + monitor_failed( + state.monitor.as_ref(), + &req_id, + Some(response.status()), + "Timed out reading request body", + ); + return response; + } }; let body: crate::anthropic::schema::MessagesRequest = match parse_json_body(&body_bytes) { diff --git a/tests/smoke_cutover.rs b/tests/smoke_cutover.rs index 7bfa167..be839f7 100644 --- a/tests/smoke_cutover.rs +++ b/tests/smoke_cutover.rs @@ -14,6 +14,7 @@ use std::path::{Path, PathBuf}; use std::sync::{Arc, Mutex, OnceLock}; use std::time::Duration; use tempfile::TempDir; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpListener; use tokio_tungstenite::tungstenite::Message; use tower::util::ServiceExt; @@ -245,6 +246,49 @@ async fn spawn_websocket_delayed_terminal_upstream() -> String { addr_str } +async fn spawn_websocket_partial_close_then_http_success_upstream( + attempts: Arc>, +) -> String { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let addr_str = format!("http://{addr}"); + + tokio::spawn(async move { + let Ok((stream, _)) = listener.accept().await else { + return; + }; + let Ok(mut ws) = tokio_tungstenite::accept_async(stream).await else { + return; + }; + let _ = ws.next().await; + *attempts.lock().unwrap() += 1; + let _ = ws + .send(Message::Text( + r#"{"type":"response.output_text.delta","delta":"discard me"}"#.into(), + )) + .await; + drop(ws); + + let Ok((mut stream, _)) = listener.accept().await else { + return; + }; + let mut request = vec![0u8; 16 * 1024]; + if stream.read(&mut request).await.unwrap_or(0) == 0 { + return; + } + *attempts.lock().unwrap() += 1; + let body = b"data: {\"type\":\"response.output_item.added\",\"output_index\":0,\"item\":{\"type\":\"message\",\"id\":\"msg_keep\"}}\n\ndata: {\"type\":\"response.output_text.delta\",\"output_index\":0,\"delta\":\"keep me\"}\n\ndata: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"type\":\"message\"}}\n\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_keep\",\"usage\":{\"input_tokens\":5,\"output_tokens\":2}}}\n\n"; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n", + body.len() + ); + let _ = stream.write_all(response.as_bytes()).await; + let _ = stream.write_all(body).await; + }); + + addr_str +} + async fn spawn_websocket_error_upstream(message: &'static str) -> String { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); @@ -835,6 +879,42 @@ async fn smoke_codex_websocket_messages_uses_mock_upstream() { assert!(sent.get("stream").is_none()); } +#[allow(clippy::await_holding_lock)] +#[tokio::test(flavor = "multi_thread")] +async fn smoke_codex_auto_buffers_failed_attempt_before_downstream_stream() { + let _guard = env_lock(); + let config = TempDir::new().unwrap(); + write_auth(config.path(), "codex"); + clear_codex_websocket_pool_for_tests(); + + let attempts = Arc::new(Mutex::new(0usize)); + let upstream = spawn_websocket_partial_close_then_http_success_upstream(attempts.clone()).await; + + let _config_env = EnvGuard::set("CCP_CONFIG_DIR", config.path()); + let _base_url_env = EnvGuard::set("CCP_CODEX_BASE_URL", &upstream); + let _transport_env = EnvGuard::set("CCP_CODEX_TRANSPORT", "auto"); + let response = tokio::time::timeout( + Duration::from_secs(8), + call_messages_body(json!({ + "model": "gpt-5.5", + "max_tokens": 64, + "stream": true, + "messages": [{"role":"user","content":"hello"}] + })), + ) + .await + .expect("buffered Auto request timed out"); + + assert_eq!(response.status(), StatusCode::OK); + let body = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + let body = String::from_utf8(body.to_vec()).unwrap(); + assert_eq!(*attempts.lock().unwrap(), 2); + assert!(body.contains("keep me")); + assert!(!body.contains("discard me")); +} + #[allow(clippy::await_holding_lock)] #[tokio::test(flavor = "multi_thread")] async fn smoke_codex_websocket_stream_returns_delta_before_terminal() {