diff --git a/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs b/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs index be10fffb84..638d797cf3 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs @@ -4,7 +4,10 @@ use super::protocol::RemoteControlPairingStatusResponse as BackendRemoteControlP use super::protocol::RemoteControlTarget; use super::protocol::StartRemoteControlPairingRequest; use super::protocol::StartRemoteControlPairingResponse; +use super::server_api::RemoteControlServerRequestError; +use super::server_api::retry_after_with_jitter; use axum::http::HeaderMap; +use axum::http::StatusCode; use codex_app_server_protocol::RemoteControlPairingStartResponse; use codex_app_server_protocol::RemoteControlPairingStatusResponse; use codex_login::default_client::create_client_without_request_logging; @@ -74,11 +77,17 @@ impl RemoteControlEnrollment { })?; let headers = response.headers().clone(); let status = response.status(); + let retry_at = retry_after_with_jitter(&headers, OffsetDateTime::now_utc()); let body = response.bytes().await.map_err(|err| { - io::Error::other(format!( - "failed to read remote control pairing response from `{}`: {err}", - self.remote_control_target.pair_url - )) + pairing_response_error( + format!( + "failed to read remote control pairing response from `{}`: {err}", + self.remote_control_target.pair_url + ), + status, + retry_at, + ErrorKind::Other, + ) })?; let body_preview = preview_remote_control_response_body(&body); if !status.is_success() { @@ -87,13 +96,15 @@ impl RemoteControlEnrollment { 404 => ErrorKind::NotFound, _ => ErrorKind::Other, }; - return Err(io::Error::new( - error_kind, + return Err(pairing_response_error( format!( "remote control pairing failed at `{}`: HTTP {status}, {}, body: {body_preview}", self.remote_control_target.pair_url, format_headers(&headers) ), + status, + retry_at, + error_kind, )); } @@ -169,11 +180,17 @@ impl RemoteControlEnrollment { })?; let headers = response.headers().clone(); let status = response.status(); + let retry_at = retry_after_with_jitter(&headers, OffsetDateTime::now_utc()); let body = response.bytes().await.map_err(|err| { - io::Error::other(format!( - "failed to read remote control pairing status response from `{}`: {err}", - self.remote_control_target.pair_status_url - )) + pairing_response_error( + format!( + "failed to read remote control pairing status response from `{}`: {err}", + self.remote_control_target.pair_status_url + ), + status, + retry_at, + ErrorKind::Other, + ) })?; let body_preview = preview_remote_control_response_body(&body); if !status.is_success() { @@ -182,13 +199,15 @@ impl RemoteControlEnrollment { 404 | 410 => ErrorKind::InvalidInput, _ => ErrorKind::Other, }; - return Err(io::Error::new( - error_kind, + return Err(pairing_response_error( format!( "remote control pairing status failed at `{}`: HTTP {status}, {}, body: {body_preview}", self.remote_control_target.pair_status_url, format_headers(&headers) ), + status, + retry_at, + error_kind, )); } @@ -242,6 +261,27 @@ impl RemoteControlEnrollment { } } +fn pairing_response_error( + message: String, + status: StatusCode, + retry_at: Option, + fallback_kind: ErrorKind, +) -> io::Error { + if matches!( + status, + StatusCode::TOO_MANY_REQUESTS | StatusCode::SERVICE_UNAVAILABLE + ) { + RemoteControlServerRequestError::io_error( + message, + Some(status), + retry_at, + /*timed_out*/ false, + ) + } else { + io::Error::new(fallback_kind, message) + } +} + pub(super) async fn load_persisted_remote_control_enrollment( state_db: Option<&StateRuntime>, remote_control_target: &RemoteControlTarget, diff --git a/codex-rs/app-server-transport/src/transport/remote_control/mod.rs b/codex-rs/app-server-transport/src/transport/remote_control/mod.rs index 32d8bf2ccf..7d8e9e2d86 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/mod.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/mod.rs @@ -132,6 +132,8 @@ enum RemoteControlEnrollmentSelection { struct RemoteControlEnrollmentState { enrollment: StdMutex>, + // Keep an observed server deadline across enrollment replacement and enable transitions. + retry_at: StdMutex>, lock: Semaphore, } @@ -139,10 +141,48 @@ impl RemoteControlEnrollmentState { fn new(enrollment: Option) -> Self { Self { enrollment: StdMutex::new(enrollment), + retry_at: StdMutex::new(None), lock: Semaphore::new(1), } } + async fn lock_for_request(&self) -> io::Result> { + self.check_retry_after()?; + let lease = self.lock().await; + self.check_retry_after()?; + Ok(lease) + } + + // Check immediately before admitting network work. Requests admitted before + // an overload response publishes its deadline may finish concurrently. + fn check_retry_after(&self) -> io::Result<()> { + let retry_at = *self + .retry_at + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(retry_at) = retry_at + && retry_at > time::OffsetDateTime::now_utc() + { + return Err(server_api::RemoteControlServerRequestError::retry_deferred( + retry_at, + )); + } + Ok(()) + } + + fn record_retry_after(&self, result: io::Result) -> io::Result { + if let Err(err) = &result + && let Some(retry_at) = server_api::remote_control_retry_at(err) + { + let mut current_retry_at = self + .retry_at + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + *current_retry_at = Some(current_retry_at.map_or(retry_at, |at| at.max(retry_at))); + } + result + } + async fn lock(&self) -> RemoteControlEnrollmentLease<'_> { let permit = match self.lock.acquire().await { Ok(permit) => permit, @@ -395,6 +435,7 @@ impl RemoteControlHandle { if !self.desired_state_tx.borrow().is_enabled() { return Err(Self::pairing_disabled_error()); } + let mut current_enrollment = self.current_enrollment.lock_for_request().await?; let mut auth = load_remote_control_auth(&self.auth_manager) .await .map_err(|_| pairing_unavailable_error())?; @@ -402,7 +443,6 @@ impl RemoteControlHandle { let installation_id = status.installation_id; let app_server_client_name = self.pairing_persistence_key(app_server_client_name)?; let app_server_client_name = app_server_client_name.as_deref(); - let mut current_enrollment = self.current_enrollment.lock().await; let mut enrollment = self .load_or_enroll_pairing_server( &mut current_enrollment, @@ -443,6 +483,7 @@ impl RemoteControlHandle { let pairing_request = || protocol::StartRemoteControlPairingRequest { manual_code: params.manual_code, }; + self.current_enrollment.check_retry_after()?; let pairing_response = match enrollment.start_pairing(pairing_request()).await { Err(err) if err.kind() == io::ErrorKind::PermissionDenied => { clear_pairing_server_token(&mut current_enrollment, &mut enrollment)?; @@ -454,6 +495,7 @@ impl RemoteControlHandle { &mut enrollment, ) .await?; + self.current_enrollment.check_retry_after()?; enrollment.start_pairing(pairing_request()).await } Err(err) if err.kind() == io::ErrorKind::NotFound => { @@ -467,11 +509,15 @@ impl RemoteControlHandle { RemoteControlEnrollmentSelection::ReplaceExisting, ) .await?; + self.current_enrollment.check_retry_after()?; enrollment.start_pairing(pairing_request()).await } pairing_response => pairing_response, }; if let Err(err) = &pairing_response { + if server_api::remote_control_retry_at(err).is_some() { + return self.current_enrollment.record_retry_after(pairing_response); + } match err.kind() { io::ErrorKind::NotFound => { self.load_or_enroll_pairing_server( @@ -593,6 +639,7 @@ impl RemoteControlHandle { RemoteControlEnrollmentSelection::ReplaceExisting => {} } + self.current_enrollment.check_retry_after()?; // Reused enrollments must still reach durable persistence during shutdown. let enrollment = tokio::select! { biased; @@ -603,12 +650,13 @@ impl RemoteControlHandle { )); } result = enroll_pairing_server( + &self.current_enrollment, &self.auth_manager, auth, &remote_control_target, installation_id, server_name, - ) => result?, + ) => self.current_enrollment.record_retry_after(result)?, }; Ok((enrollment, true)) } @@ -635,12 +683,12 @@ impl RemoteControlHandle { if !self.desired_state_tx.borrow().is_enabled() { return Err(Self::pairing_disabled_error()); } + let mut current_enrollment = self.current_enrollment.lock_for_request().await?; let mut auth = load_remote_control_auth(&self.auth_manager) .await .map_err(|_| pairing_unavailable_error())?; let app_server_client_name = self.pairing_persistence_key.borrow().clone(); let app_server_client_name = app_server_client_name.as_deref(); - let mut current_enrollment = self.current_enrollment.lock().await; let mut enrollment = current_enrollment .as_ref() .filter(|enrollment| enrollment.account_id == auth.account_id) @@ -678,6 +726,7 @@ impl RemoteControlHandle { let status_code = remote_control_pairing_status_code(¶ms)?; let pairing_status_request = || protocol::RemoteControlPairingStatusRequest::from(status_code.clone()); + self.current_enrollment.check_retry_after()?; let pairing_status_response = match enrollment.pairing_status(pairing_status_request()).await { Err(err) if err.kind() == io::ErrorKind::PermissionDenied => { @@ -690,11 +739,17 @@ impl RemoteControlHandle { &mut enrollment, ) .await?; + self.current_enrollment.check_retry_after()?; enrollment.pairing_status(pairing_status_request()).await } pairing_status_response => pairing_status_response, }; if let Err(err) = &pairing_status_response { + if server_api::remote_control_retry_at(err).is_some() { + return self + .current_enrollment + .record_retry_after(pairing_status_response); + } match err.kind() { io::ErrorKind::NotFound => { self.load_or_enroll_pairing_server( @@ -784,6 +839,7 @@ impl RemoteControlHandle { } async fn enroll_pairing_server( + current_enrollment: &RemoteControlEnrollmentState, auth_manager: &Arc, auth: &mut auth::RemoteControlConnectionAuth, remote_control_target: &protocol::RemoteControlTarget, @@ -806,6 +862,7 @@ async fn enroll_pairing_server( } Err(err) => return Err(err), } + current_enrollment.check_retry_after()?; enroll_remote_control_server(remote_control_target, auth, installation_id, server_name).await } @@ -831,12 +888,13 @@ fn remote_control_pairing_status_code( } async fn refresh_pairing_enrollment( - current_enrollment: &mut Option, + current_enrollment: &mut RemoteControlEnrollmentLease<'_>, auth_manager: &Arc, auth: &mut auth::RemoteControlConnectionAuth, installation_id: &str, enrollment: &mut RemoteControlEnrollment, ) -> io::Result<()> { + current_enrollment.state.check_retry_after()?; let mut refresh_result = refresh_remote_control_server(auth, installation_id, enrollment).await; if refresh_result .as_ref() @@ -848,6 +906,7 @@ async fn refresh_pairing_enrollment( match load_remote_control_auth(auth_manager).await { Ok(recovered_auth) if recovered_auth.account_id == enrollment.account_id => { *auth = recovered_auth; + current_enrollment.state.check_retry_after()?; refresh_result = refresh_remote_control_server(auth, installation_id, enrollment).await; } @@ -869,7 +928,7 @@ async fn refresh_pairing_enrollment( if !replace_current_enrollment(current_enrollment, enrollment) { Err(pairing_unavailable_error()) } else { - refresh_result + current_enrollment.state.record_retry_after(refresh_result) } } diff --git a/codex-rs/app-server-transport/src/transport/remote_control/server_api.rs b/codex-rs/app-server-transport/src/transport/remote_control/server_api.rs index fcd4ae9a42..ab7d7fd8a0 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/server_api.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/server_api.rs @@ -21,6 +21,8 @@ use time::OffsetDateTime; use time::format_description::well_known::Rfc3339; use tracing::warn; +const REMOTE_CONTROL_RETRY_AFTER_JITTER_MAX_MILLIS: u64 = 30_000; + const REMOTE_CONTROL_ENROLL_TIMEOUT: Duration = Duration::from_secs(30); const REMOTE_CONTROL_SERVER_TOKEN_REFRESH_BACKOFF_MIN_SECS: u64 = 24; const REMOTE_CONTROL_SERVER_TOKEN_REFRESH_BACKOFF_MAX_SECS: u64 = 36; @@ -28,14 +30,25 @@ const REMOTE_CONTROL_SERVER_TOKEN_REFRESH_BACKOFF_MAX_SECS: u64 = 36; pub(super) const REMOTE_CONTROL_INSTALLATION_ID_HEADER: &str = "x-codex-installation-id"; #[derive(Debug)] -struct RemoteControlServerRequestError { +pub(super) struct RemoteControlServerRequestError { message: String, status: Option, retry_at: Option, } impl RemoteControlServerRequestError { - fn io_error( + pub(super) fn retry_deferred(retry_at: OffsetDateTime) -> io::Error { + io::Error::new( + ErrorKind::WouldBlock, + Self { + message: format!("remote control retry deferred until {retry_at}"), + status: None, + retry_at: Some(retry_at), + }, + ) + } + + pub(super) fn io_error( message: String, status: Option, retry_at: Option, @@ -132,9 +145,8 @@ pub(super) async fn refresh_remote_control_server( && let Some(next_refresh_at) = enrollment.next_refresh_at && next_refresh_at > now { - return Err(io::Error::new( - ErrorKind::WouldBlock, - format!("remote control server token refresh deferred until {next_refresh_at}"), + return Err(RemoteControlServerRequestError::retry_deferred( + next_refresh_at, )); } let refresh_url = enrollment.remote_control_target.refresh_url.clone(); @@ -166,7 +178,13 @@ pub(super) async fn refresh_remote_control_server( == RemoteControlServerTokenRefreshRequirement::Required; let (refresh_delay, next_refresh_at) = refresh_deferral(refresh_error.retry_at, now); enrollment.next_refresh_at = Some(next_refresh_at); - if refresh_is_required { + // An explicit server deadline takes precedence over valid-token fallback. + // Keep the token, but let callers defer new requests until this deadline. + if refresh_is_required + || refresh_error + .retry_at + .is_some_and(|retry_at| retry_at > now) + { warn!( refresh_url, server_id = %enrollment.server_id, @@ -174,7 +192,7 @@ pub(super) async fn refresh_remote_control_server( error = %err, ?refresh_delay, %next_refresh_at, - "required remote control server token refresh failed; deferring next attempt" + "remote control server token refresh failed; deferring new requests" ); return Err(err); } @@ -244,7 +262,7 @@ where })?; let headers = response.headers().clone(); let status = response.status(); - let retry_at = parse_retry_after(&headers, OffsetDateTime::now_utc()); + let retry_at = retry_after_with_jitter(&headers, OffsetDateTime::now_utc()); let body = response.bytes().await.map_err(|err| { let timed_out = err.is_timeout(); RemoteControlServerRequestError::io_error( @@ -298,6 +316,31 @@ fn remote_control_server_request_error( err.get_ref()?.downcast_ref() } +pub(super) fn remote_control_retry_at(err: &io::Error) -> Option { + let request_error = remote_control_server_request_error(err)?; + if !request_error.is_transient(err.kind()) { + return None; + } + request_error.retry_at +} + +pub(super) fn remote_control_retry_delay(err: &io::Error) -> Option { + Duration::try_from(remote_control_retry_at(err)? - OffsetDateTime::now_utc()).ok() +} + +pub(super) fn retry_after_with_jitter( + headers: &HeaderMap, + received_at: OffsetDateTime, +) -> Option { + let retry_at = parse_retry_after(headers, received_at)?; + // Keep the server's deadline as a lower bound, then spread clients across + // the next 30 seconds. Sample once per response so retries share a deadline. + let jitter = time::Duration::milliseconds( + rand::rng().random_range(0..=REMOTE_CONTROL_RETRY_AFTER_JITTER_MAX_MILLIS) as i64, + ); + Some(retry_at.checked_add(jitter).unwrap_or(retry_at)) +} + fn parse_retry_after(headers: &HeaderMap, received_at: OffsetDateTime) -> Option { let retry_after = headers .get(axum::http::header::RETRY_AFTER)? @@ -309,7 +352,7 @@ fn parse_retry_after(headers: &HeaderMap, received_at: OffsetDateTime) -> Option } else { OffsetDateTime::from(httpdate::parse_http_date(retry_after).ok()?) }; - (retry_at > received_at).then_some(retry_at) + (retry_at >= received_at).then_some(retry_at) } fn refresh_deferral( diff --git a/codex-rs/app-server-transport/src/transport/remote_control/server_api_tests.rs b/codex-rs/app-server-transport/src/transport/remote_control/server_api_tests.rs index 786dc65a90..0c636d3c0e 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/server_api_tests.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/server_api_tests.rs @@ -282,3 +282,40 @@ fn http_date_retry_after_preserves_absolute_deadline() { (Duration::from_secs(90), retry_at) ); } + +#[test] +fn retry_after_jitter_never_shortens_the_server_deadline() { + let now = + OffsetDateTime::from_unix_timestamp(1_700_000_000).expect("test timestamp should parse"); + for value in ["120", "Tue, 14 Nov 2023 22:15:20 GMT"] { + let mut headers = HeaderMap::new(); + headers.insert( + axum::http::header::RETRY_AFTER, + value.parse().expect("retry hint should be a valid header"), + ); + let deadline = retry_after_with_jitter(&headers, now) + .expect("a valid retry hint should produce a deadline"); + assert!( + (now + time::Duration::seconds(120)..=now + time::Duration::seconds(150)) + .contains(&deadline) + ); + // Reading the response body must not start a new relative wait or resample jitter. + let (delay, deferred_until) = + refresh_deferral(Some(deadline), now + time::Duration::seconds(30)); + assert_eq!(deferred_until, deadline); + assert!((Duration::from_secs(90)..=Duration::from_secs(120)).contains(&delay)); + } +} + +#[test] +fn zero_retry_after_preserves_bounded_jitter() { + let now = OffsetDateTime::UNIX_EPOCH; + let mut headers = HeaderMap::new(); + headers.insert( + axum::http::header::RETRY_AFTER, + axum::http::HeaderValue::from_static("0"), + ); + let deadline = retry_after_with_jitter(&headers, now) + .expect("a zero-second retry hint should produce a deadline"); + assert!((now..=now + time::Duration::seconds(30)).contains(&deadline)); +} diff --git a/codex-rs/app-server-transport/src/transport/remote_control/tests.rs b/codex-rs/app-server-transport/src/transport/remote_control/tests.rs index e1bfaa9aca..3878c32410 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/tests.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/tests.rs @@ -68,6 +68,8 @@ use tokio_util::sync::CancellationToken; mod clients_tests; mod pairing_tests; +#[path = "tests/retry_tests.rs"] +mod retry_tests; const TEST_INSTALLATION_ID: &str = "11111111-1111-4111-8111-111111111111"; const TEST_REMOTE_CONTROL_URL: &str = "http://127.0.0.1:1/backend-api/wham/remote/control"; diff --git a/codex-rs/app-server-transport/src/transport/remote_control/tests/pairing_tests.rs b/codex-rs/app-server-transport/src/transport/remote_control/tests/pairing_tests.rs index a7567f26e8..b343242c8a 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/tests/pairing_tests.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/tests/pairing_tests.rs @@ -770,6 +770,7 @@ async fn pairing_publishes_refresh_deferral_after_auth_recovery() { recovered_refresh_request.headers.get("authorization"), Some(&"Bearer fresh-token".to_string()) ); + let response_started_at = OffsetDateTime::now_utc(); respond_with_status_and_headers( recovered_refresh_request.stream, "502 Bad Gateway", @@ -777,6 +778,7 @@ async fn pairing_publishes_refresh_deferral_after_auth_recovery() { "upstream unavailable", ) .await; + response_started_at }); let codex_home = TempDir::new().expect("temp dir should create"); let auth_manager = auth_manager_with_replacement(&codex_home, "account_id").await; @@ -790,7 +792,6 @@ async fn pairing_publishes_refresh_deferral_after_auth_recovery() { .expect("current enrollment should exist") .expires_at = Some(OffsetDateTime::now_utc() - time::Duration::seconds(1)); - let refresh_started_at = OffsetDateTime::now_utc(); let refresh_err = remote_handle .start_pairing( RemoteControlPairingStartParams::default(), @@ -806,7 +807,7 @@ async fn pairing_publishes_refresh_deferral_after_auth_recovery() { ) .await .expect_err("published deadline should throttle the next pairing refresh"); - server_task.await.expect("server task should finish"); + let response_started_at = server_task.await.expect("server task should finish"); assert!(refresh_err.to_string().contains("HTTP 502 Bad Gateway")); assert_eq!(deferred_err.kind(), io::ErrorKind::WouldBlock); @@ -816,8 +817,8 @@ async fn pairing_publishes_refresh_deferral_after_auth_recovery() { .and_then(|enrollment| enrollment.next_refresh_at) .expect("required refresh failure should publish its retry deadline"); assert!( - (refresh_started_at + time::Duration::seconds(120) - ..=refresh_completed_at + time::Duration::seconds(120)) + (response_started_at + time::Duration::seconds(120) + ..=refresh_completed_at + time::Duration::seconds(150)) .contains(&next_refresh_at) ); } diff --git a/codex-rs/app-server-transport/src/transport/remote_control/tests/retry_tests.rs b/codex-rs/app-server-transport/src/transport/remote_control/tests/retry_tests.rs new file mode 100644 index 0000000000..2ab8a88291 --- /dev/null +++ b/codex-rs/app-server-transport/src/transport/remote_control/tests/retry_tests.rs @@ -0,0 +1,187 @@ +//! Exercises overload retry deadlines over HTTP and WebSocket connections. +//! Auth reloads must respect the deadline; shutdown must remain prompt. + +use super::*; +use pretty_assertions::assert_eq; + +#[tokio::test] +async fn rate_limited_enrollment_respects_retry_after() { + assert_overload_retry_after("429 Too Many Requests", /*reject_enrollment*/ true).await; +} + +#[tokio::test] +async fn rate_limited_websocket_respects_retry_after() { + assert_overload_retry_after("429 Too Many Requests", /*reject_enrollment*/ false).await; +} + +#[tokio::test] +async fn enrollment_resumes_after_retry_after() { + assert_overload_retry_after("503 Service Unavailable", /*reject_enrollment*/ true).await; +} + +#[tokio::test] +async fn websocket_resumes_after_retry_after() { + assert_overload_retry_after("503 Service Unavailable", /*reject_enrollment*/ false).await; +} + +async fn assert_overload_retry_after(status: &str, reject_enrollment: bool) { + let verify_recovery = status == "503 Service Unavailable"; + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let codex_home = TempDir::new().expect("temp dir should create"); + let (transport_event_tx, _transport_event_rx) = mpsc::channel(CHANNEL_CAPACITY); + let shutdown_token = CancellationToken::new(); + let auth_manager = remote_control_auth_manager_with_home(&codex_home); + let (remote_task, remote_handle) = start_remote_control( + RemoteControlStartConfig { + remote_control_url: remote_control_url_for_listener(&listener), + installation_id: TEST_INSTALLATION_ID.to_string(), + policy: RemoteControlPolicy::Allowed, + }, + Some(remote_control_state_runtime(&codex_home).await), + auth_manager.clone(), + transport_event_tx, + shutdown_token.clone(), + /*app_server_client_name_rx*/ None, + RemoteControlStartupMode::EnabledEphemeral, + ) + .await + .expect("remote control should start"); + let mut status_rx = remote_handle.status_receiver(); + let mut rejected_request = accept_http_request(&listener).await; + assert_eq!( + rejected_request.request_line, + "POST /backend-api/wham/remote/control/server/enroll HTTP/1.1" + ); + if !reject_enrollment { + respond_with_json( + rejected_request.stream, + remote_control_server_token_response( + "srv_e_test", + "env_test", + TEST_REMOTE_CONTROL_SERVER_TOKEN, + ), + ) + .await; + rejected_request = accept_http_request(&listener).await; + assert_eq!( + rejected_request.request_line, + "GET /backend-api/wham/remote/control/server HTTP/1.1" + ); + } + let response_started_at = std::time::Instant::now(); + respond_with_status_and_headers( + rejected_request.stream, + status, + &[("Retry-After", if verify_recovery { "3" } else { "120" })], + "overloaded", + ) + .await; + timeout( + Duration::from_secs(5), + status_rx.wait_for(|status| status.status == RemoteControlConnectionStatus::Errored), + ) + .await + .expect("the overload response should be processed") + .expect("the status channel should remain open"); + + let auth_changes = auth_manager.auth_change_receiver(); + save_auth( + codex_home.path(), + &remote_control_auth_dot_json(Some("account_id")), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("updated credentials should save"); + auth_manager.reload().await; + assert!( + auth_changes + .has_changed() + .expect("auth watch should remain open") + ); + + if !verify_recovery { + let pairing_error = timeout( + Duration::from_secs(1), + remote_handle.start_pairing( + RemoteControlPairingStartParams::default(), + /*app_server_client_name*/ None, + ), + ) + .await + .expect("pairing should report the active retry delay promptly") + .expect_err("pairing must respect the shared server retry delay"); + let retry_delay = server_api::remote_control_retry_delay(&pairing_error) + .expect("pairing should preserve the overload retry deadline"); + assert!( + retry_delay >= Duration::from_secs(120).saturating_sub(response_started_at.elapsed()), + "pairing must preserve the full remaining Retry-After interval" + ); + assert!( + timeout(Duration::from_secs(2), listener.accept()) + .await + .is_err(), + "no enrollment, refresh or handshake should retry before Retry-After ({status})" + ); + remote_handle.disable_ephemeral().await; + assert!( + timeout(Duration::from_secs(2), listener.accept()) + .await + .is_err(), + "disabling remote control must stop connection attempts" + ); + remote_handle + .enable_ephemeral() + .expect("remote control should enable again"); + assert!( + timeout(Duration::from_secs(2), listener.accept()) + .await + .is_err(), + "re-enabling remote control must preserve the server retry delay" + ); + } + if !reject_enrollment { + assert_eq!( + remote_handle + .current_enrollment + .snapshot() + .and_then(|enrollment| enrollment.remote_control_token), + Some(TEST_REMOTE_CONTROL_SERVER_TOKEN.to_string()), + "overload must preserve the enrolled token" + ); + } + if verify_recovery { + remote_handle.disable_ephemeral().await; + remote_handle + .enable_ephemeral() + .expect("remote control should enable again"); + let (stream, _) = timeout(Duration::from_secs(45), listener.accept()) + .await + .expect("retry should resume after the server delay and bounded jitter") + .expect("retry connection should succeed"); + assert!( + response_started_at.elapsed() >= Duration::from_secs(3), + "the retry must wait for the full Retry-After interval" + ); + let mut reader = BufReader::new(stream); + let mut request_line = String::new(); + timeout(Duration::from_secs(5), reader.read_line(&mut request_line)) + .await + .expect("retry should send the request promptly") + .expect("retry request should read"); + assert_eq!( + request_line.trim_end(), + if reject_enrollment { + "POST /backend-api/wham/remote/control/server/enroll HTTP/1.1" + } else { + "GET /backend-api/wham/remote/control/server HTTP/1.1" + } + ); + } + shutdown_token.cancel(); + timeout(Duration::from_secs(1), remote_task) + .await + .expect("shutdown must interrupt the server retry delay") + .expect("remote task should finish"); +} diff --git a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs index 6803b99b04..cabac25f56 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs @@ -29,8 +29,11 @@ use crate::transport::remote_control::enroll::preview_remote_control_response_bo use crate::transport::remote_control::enroll::update_persisted_remote_control_enrollment; use crate::transport::remote_control::host_device::REMOTE_CONTROL_HOST_DEVICE_KIND_HEADER; use crate::transport::remote_control::host_device::host_device_kind; +use crate::transport::remote_control::server_api::RemoteControlServerRequestError; use crate::transport::remote_control::server_api::enroll_remote_control_server; use crate::transport::remote_control::server_api::refresh_remote_control_server; +use crate::transport::remote_control::server_api::remote_control_retry_delay; +use crate::transport::remote_control::server_api::retry_after_with_jitter; use axum::http::HeaderValue; use base64::Engine; use codex_app_server_protocol::RemoteControlConnectionStatus; @@ -776,14 +779,20 @@ impl RemoteControlWebsocket { if !self.desired_state_rx.borrow().is_enabled() { return ConnectOutcome::Disabled; } + let server_retry_delay = remote_control_retry_delay(&err); let reconnect_delay = if err.kind() == ErrorKind::WouldBlock { - REMOTE_CONTROL_ACCOUNT_ID_RETRY_INTERVAL + server_retry_delay + .map_or(REMOTE_CONTROL_ACCOUNT_ID_RETRY_INTERVAL, |delay| { + REMOTE_CONTROL_ACCOUNT_ID_RETRY_INTERVAL.max(delay) + }) } else { self.status_publisher .publish_status(RemoteControlConnectionStatus::Errored); let reconnect_attempt = self.reconnect_attempt.saturating_add(1); let (reconnect_delay, reconnect_backoff_reset) = next_reconnect_delay(&mut self.reconnect_attempt); + let reconnect_delay = server_retry_delay + .map_or(reconnect_delay, |delay| reconnect_delay.max(delay)); let enrollment = self.current_enrollment.snapshot(); warn!( websocket_url = %remote_control_target.websocket_url, @@ -816,7 +825,8 @@ impl RemoteControlWebsocket { } return ConnectOutcome::Disabled; } - changed = self.auth_change_rx.changed() => { + // Auth changes may shorten local backoff after the server deadline. + changed = wait_for_auth_change(&mut self.auth_change_rx, server_retry_delay) => { if changed.is_err() { return ConnectOutcome::Shutdown; } @@ -1309,6 +1319,16 @@ async fn build_remote_control_websocket_request( Ok(request) } +async fn wait_for_auth_change( + auth_change_rx: &mut watch::Receiver, + server_retry_delay: Option, +) -> Result<(), watch::error::RecvError> { + if let Some(delay) = server_retry_delay { + tokio::time::sleep(delay).await; + } + auth_change_rx.changed().await +} + fn next_reconnect_delay(reconnect_attempt: &mut u64) -> (std::time::Duration, bool) { let reconnect_delay = backoff(*reconnect_attempt).min(REMOTE_CONTROL_RECONNECT_BACKOFF_CAP); let reconnect_backoff_reset = reconnect_delay == REMOTE_CONTROL_RECONNECT_BACKOFF_CAP; @@ -1334,17 +1354,18 @@ pub(super) async fn connect_remote_control_websocket( ensure_rustls_crypto_provider(); let (auth, enrollment) = { - let mut current_enrollment = current_enrollment.lock().await; - let auth = prepare_remote_control_enrollment( + let mut lease = current_enrollment.lock_for_request().await?; + let auth_result = prepare_remote_control_enrollment( remote_control_target, state_db, &mut auth_context, - &mut current_enrollment, + &mut lease, connect_options, status_publisher, ) - .await?; - let enrollment = current_enrollment.as_ref().cloned().ok_or_else(|| { + .await; + let auth = current_enrollment.record_retry_after(auth_result)?; + let enrollment = lease.as_ref().cloned().ok_or_else(|| { io::Error::other("missing remote control enrollment after enrollment step") })?; (auth, enrollment) @@ -1357,6 +1378,7 @@ pub(super) async fn connect_remote_control_websocket( ) .await?; + current_enrollment.check_retry_after()?; let websocket_connect_result = tokio::time::timeout( REMOTE_CONTROL_WEBSOCKET_CONNECT_TIMEOUT, connect_async(request), @@ -1417,6 +1439,24 @@ pub(super) async fn connect_remote_control_websocket( "remote control websocket returned unrecognized HTTP 404; preserving enrollment before retry" ); } + tungstenite::Error::Http(response) + if matches!(response.status().as_u16(), 429 | 503) => + { + return current_enrollment.record_retry_after(Err( + RemoteControlServerRequestError::io_error( + format_remote_control_websocket_connect_error( + &remote_control_target.websocket_url, + &err, + ), + Some(response.status()), + retry_after_with_jitter( + response.headers(), + time::OffsetDateTime::now_utc(), + ), + /*timed_out*/ false, + ), + )); + } tungstenite::Error::Http(response) if matches!(response.status().as_u16(), 401 | 403) => { @@ -1669,23 +1709,24 @@ async fn replace_remote_control_enrollment_if_matches( "remote control requires sqlite state db", )); }; - let mut current_enrollment = current_enrollment.lock().await; - if !current_enrollment + let mut lease = current_enrollment.lock_for_request().await?; + if !lease .as_ref() .is_some_and(|current| same_remote_control_enrollment(current, enrollment)) { return Ok(()); } - enroll_and_persist_remote_control_server( + let result = enroll_and_persist_remote_control_server( remote_control_target, state_db, auth_context, - &mut current_enrollment, + &mut lease, connect_options, status_publisher, RemoteControlEnrollmentSelection::ReplaceExisting, ) - .await + .await; + current_enrollment.record_retry_after(result) } async fn clear_remote_control_server_token_if_matches( @@ -1893,6 +1934,26 @@ mod tests { Arc::new(RemoteControlEnrollmentState::new(enrollment)) } + #[tokio::test(start_paused = true)] + async fn queued_auth_change_waits_for_server_delay() { + let (auth_change_tx, mut auth_change_rx) = watch::channel(0); + auth_change_tx.send_replace(1); + let started = tokio::time::Instant::now(); + let auth_change = wait_for_auth_change(&mut auth_change_rx, Some(Duration::from_secs(5))); + tokio::pin!(auth_change); + + assert!( + timeout(Duration::from_secs(4), auth_change.as_mut()) + .await + .is_err() + ); + timeout(Duration::from_secs(2), auth_change) + .await + .expect("queued credentials should be usable as soon as the server delay expires") + .expect("auth watch should remain open"); + assert_eq!(started.elapsed(), Duration::from_secs(5)); + } + #[test] fn next_reconnect_delay_resets_after_cap() { let mut reconnect_attempt = 9; @@ -2024,7 +2085,7 @@ mod tests { format!("http://{addr}/backend-api/") } - fn remote_control_auth_dot_json(access_token: &str) -> AuthDotJson { + pub(super) fn remote_control_auth_dot_json(access_token: &str) -> AuthDotJson { #[derive(serde::Serialize)] struct Header { alg: &'static str, diff --git a/codex-rs/app-server-transport/src/transport/remote_control/websocket_refresh_tests.rs b/codex-rs/app-server-transport/src/transport/remote_control/websocket_refresh_tests.rs index ee4ea15f83..042a22e11d 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/websocket_refresh_tests.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/websocket_refresh_tests.rs @@ -3,6 +3,7 @@ use super::tests::TEST_INSTALLATION_ID; use super::tests::TEST_REMOTE_CONTROL_SERVER_TOKEN; use super::tests::accept_http_request; use super::tests::enabled_desired_state_sender; +use super::tests::remote_control_auth_dot_json; use super::tests::remote_control_auth_manager; use super::tests::remote_control_enrollment; use super::tests::remote_control_state_runtime; @@ -12,11 +13,16 @@ use super::tests::respond_with_status_and_headers; use super::tests::test_current_enrollment; use super::*; use crate::transport::remote_control::protocol::normalize_remote_control_url; +use crate::transport::remote_control::server_api::remote_control_retry_at; use crate::transport::remote_control::tests::remote_control_handle_with_current_enrollment; use codex_app_server_protocol::RemoteControlPairingStartParams; -use codex_app_server_protocol::RemoteControlPairingStartResponse; +use codex_app_server_protocol::RemoteControlPairingStatusParams; +use codex_config::types::AuthCredentialsStoreMode; +use codex_login::AuthKeyringBackendKind; +use codex_login::save_auth; use pretty_assertions::assert_eq; use tempfile::TempDir; +use tokio::io::AsyncWriteExt; use tokio::net::TcpListener; use tokio::net::TcpStream; use tokio::time::Duration; @@ -180,20 +186,7 @@ async fn websocket_retry_after_throttles_pairing_refresh() { "upstream unavailable", ) .await; - let first_websocket = accept_test_websocket(&listener).await; - let (pairing_stream, request_line) = accept_http_request(&listener).await; - assert_eq!( - request_line, - "POST /backend-api/wham/remote/control/server/pair HTTP/1.1" - ); - respond_with_status_and_headers( - pairing_stream, - "200 OK", - &[], - r#"{"pairing_code":"pairing-code","manual_pairing_code":"ABCD-EFGH","server_id":"srv_e_test","environment_id":"env_test","expires_at":"3026-05-22T12:34:56Z"}"#, - ) - .await; - first_websocket + listener }); let codex_home = TempDir::new().expect("temp dir should create"); let state_db = remote_control_state_runtime(&codex_home).await; @@ -210,14 +203,14 @@ async fn websocket_retry_after_throttles_pairing_refresh() { .expires_at = Some(time::OffsetDateTime::now_utc() + time::Duration::minutes(4)); let current_enrollment = remote_handle.current_enrollment.clone(); let refresh_started_at = time::OffsetDateTime::now_utc(); - connect_test_websocket( + let refresh_error = connect_test_websocket( &remote_control_target, state_db.as_ref(), &auth_manager, ¤t_enrollment, ) .await - .expect("first websocket should connect after deferred refresh"); + .expect_err("an explicit server deadline must defer the handshake even with a valid token"); let refresh_completed_at = time::OffsetDateTime::now_utc(); let next_refresh_at = current_enrollment .snapshot() @@ -225,119 +218,423 @@ async fn websocket_retry_after_throttles_pairing_refresh() { .expect("Retry-After should set a retry deadline"); assert!( (refresh_started_at + time::Duration::seconds(120) - ..=refresh_completed_at + time::Duration::seconds(120)) + ..=refresh_completed_at + time::Duration::seconds(150)) .contains(&next_refresh_at) ); - let pairing_response = remote_handle + let pairing_error = remote_handle .start_pairing( RemoteControlPairingStartParams::default(), /*app_server_client_name*/ None, ) .await - .expect("websocket Retry-After should throttle pairing refresh"); - let first_server_websocket = server_task.await.expect("server task should succeed"); - + .expect_err("the refresh deadline must also defer pairing"); + let listener = server_task.await.expect("server task should succeed"); assert_eq!( - pairing_response, - RemoteControlPairingStartResponse { - pairing_code: "pairing-code".to_string(), - manual_pairing_code: Some("ABCD-EFGH".to_string()), - environment_id: "env_test".to_string(), - expires_at: 33_336_362_096, - } + remote_control_retry_at(&refresh_error), + Some(next_refresh_at) ); - drop(first_server_websocket); + assert_eq!( + remote_control_retry_at(&pairing_error), + Some(next_refresh_at) + ); + assert_eq!( + current_enrollment + .snapshot() + .and_then(|enrollment| enrollment.remote_control_token), + Some(TEST_REMOTE_CONTROL_SERVER_TOKEN.to_string()) + ); + timeout(Duration::from_millis(100), listener.accept()) + .await + .expect_err("no handshake or pairing should bypass the proactive refresh deadline"); } #[tokio::test] async fn pairing_http_date_retry_after_throttles_websocket_refresh() { + for status in [ + "429 Too Many Requests", + "503 Service Unavailable", + "502 Bad Gateway", + ] { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let remote_control_url = remote_control_url_for_listener(&listener); + let remote_control_target = + normalize_remote_control_url(&remote_control_url).expect("target should parse"); + let retry_after = + httpdate::fmt_http_date(std::time::SystemTime::now() + Duration::from_secs(120)); + let expected_next_refresh_at = time::OffsetDateTime::from( + httpdate::parse_http_date(&retry_after).expect("Retry-After date should parse"), + ); + let server_task = tokio::spawn(async move { + let (refresh_stream, request_line) = accept_http_request(&listener).await; + assert_eq!( + request_line, + "POST /backend-api/wham/remote/control/server/refresh HTTP/1.1" + ); + respond_with_status_and_headers( + refresh_stream, + status, + &[("retry-after", &retry_after)], + "upstream unavailable", + ) + .await; + listener + }); + let codex_home = TempDir::new().expect("temp dir should create"); + let state_db = remote_control_state_runtime(&codex_home).await; + let auth_manager = remote_control_auth_manager(); + let mut remote_handle = remote_control_handle_with_current_enrollment( + &remote_control_url, + auth_manager.clone(), + ); + remote_handle.state_db = Some(state_db.clone()); + remote_handle + .current_enrollment + .lock() + .await + .as_mut() + .expect("current enrollment should exist") + .expires_at = Some(time::OffsetDateTime::now_utc() + time::Duration::minutes(4)); + let current_enrollment = remote_handle.current_enrollment.clone(); + + let refresh_error = remote_handle + .start_pairing( + RemoteControlPairingStartParams::default(), + /*app_server_client_name*/ None, + ) + .await + .expect_err("an explicit server deadline must defer pairing even with a valid token"); + let listener = server_task.await.expect("server task should succeed"); + let retry_at = remote_control_retry_at(&refresh_error) + .expect("the proactive refresh response should preserve its deadline"); + let enrollment = current_enrollment + .snapshot() + .expect("enrollment should remain"); + assert_eq!(enrollment.next_refresh_at, Some(retry_at)); + assert_eq!( + enrollment.remote_control_token.as_deref(), + Some(TEST_REMOTE_CONTROL_SERVER_TOKEN) + ); + assert!( + (expected_next_refresh_at..=expected_next_refresh_at + time::Duration::seconds(30)) + .contains(&retry_at) + ); + let connect_error = connect_test_websocket( + &remote_control_target, + state_db.as_ref(), + &auth_manager, + ¤t_enrollment, + ) + .await + .expect_err("the refresh deadline must also defer the handshake"); + assert_eq!(remote_control_retry_at(&connect_error), Some(retry_at)); + let pairing_error = remote_handle + .start_pairing( + RemoteControlPairingStartParams::default(), + /*app_server_client_name*/ None, + ) + .await + .expect_err("another pairing request must retain the same deadline"); + assert_eq!(remote_control_retry_at(&pairing_error), Some(retry_at)); + timeout(Duration::from_millis(100), listener.accept()) + .await + .expect_err("no pairing, refresh or handshake should bypass the deadline"); + } +} + +#[tokio::test] +async fn pairing_during_pending_handshake_respects_later_overload() { let listener = TcpListener::bind("127.0.0.1:0") .await .expect("listener should bind"); let remote_control_url = remote_control_url_for_listener(&listener); let remote_control_target = normalize_remote_control_url(&remote_control_url).expect("target should parse"); - let retry_after = - httpdate::fmt_http_date(std::time::SystemTime::now() + Duration::from_secs(120)); - let expected_next_refresh_at = time::OffsetDateTime::from( - httpdate::parse_http_date(&retry_after).expect("Retry-After date should parse"), - ); - let server_task = tokio::spawn(async move { - let (refresh_stream, request_line) = accept_http_request(&listener).await; - assert_eq!( - request_line, - "POST /backend-api/wham/remote/control/server/refresh HTTP/1.1" - ); - respond_with_status_and_headers( - refresh_stream, - "502 Bad Gateway", - &[("retry-after", &retry_after)], - "upstream unavailable", - ) - .await; - let (pairing_stream, request_line) = accept_http_request(&listener).await; - assert_eq!( - request_line, - "POST /backend-api/wham/remote/control/server/pair HTTP/1.1" - ); - respond_with_status_and_headers( - pairing_stream, - "200 OK", - &[], - r#"{"pairing_code":"pairing-code","manual_pairing_code":"ABCD-EFGH","server_id":"srv_e_test","environment_id":"env_test","expires_at":"3026-05-22T12:34:56Z"}"#, - ) - .await; - accept_test_websocket(&listener).await - }); let codex_home = TempDir::new().expect("temp dir should create"); let state_db = remote_control_state_runtime(&codex_home).await; let auth_manager = remote_control_auth_manager(); let mut remote_handle = remote_control_handle_with_current_enrollment(&remote_control_url, auth_manager.clone()); remote_handle.state_db = Some(state_db.clone()); - remote_handle - .current_enrollment - .lock() - .await - .as_mut() - .expect("current enrollment should exist") - .expires_at = Some(time::OffsetDateTime::now_utc() + time::Duration::minutes(4)); let current_enrollment = remote_handle.current_enrollment.clone(); - - let pairing_response = remote_handle - .start_pairing( - RemoteControlPairingStartParams::default(), - /*app_server_client_name*/ None, - ) - .await - .expect("pairing should continue after proactive refresh failure"); - assert_eq!( - current_enrollment - .snapshot() - .and_then(|enrollment| enrollment.next_refresh_at), - Some(expected_next_refresh_at) - ); - connect_test_websocket( + let connect = connect_test_websocket( &remote_control_target, state_db.as_ref(), &auth_manager, ¤t_enrollment, + ); + tokio::pin!(connect); + let (handshake_stream, request_line) = tokio::select! { + result = &mut connect => panic!("handshake should wait for its response: {result:?}"), + request = accept_http_request(&listener) => request, + }; + assert_eq!( + request_line, + "GET /backend-api/wham/remote/control/server HTTP/1.1" + ); + + let pairing = remote_handle.start_pairing( + RemoteControlPairingStartParams::default(), + /*app_server_client_name*/ None, + ); + tokio::pin!(pairing); + let (pairing_stream, request_line) = tokio::select! { + result = &mut pairing => panic!("pairing should wait for its response: {result:?}"), + request = accept_http_request(&listener) => request, + }; + assert_eq!( + request_line, + "POST /backend-api/wham/remote/control/server/pair HTTP/1.1" + ); + respond_with_status_and_headers( + pairing_stream, + "200 OK", + &[], + r#"{"pairing_code":"pairing-code","manual_pairing_code":"ABCD-EFGH","server_id":"srv_e_test","environment_id":"env_test","expires_at":"3026-05-22T12:34:56Z"}"#, + ) + .await; + timeout(Duration::from_secs(1), pairing) + .await + .expect("pairing must stay responsive while the handshake is pending") + .expect("pairing should succeed before any overload is observed"); + respond_with_status_and_headers( + handshake_stream, + "503 Service Unavailable", + &[("Retry-After", "120")], + "overloaded", + ) + .await; + let connect_error = connect + .await + .expect_err("the handshake should report overload"); + let retry_at = remote_control_retry_at(&connect_error) + .expect("the handshake should preserve its retry deadline"); + let pairing_error = timeout( + Duration::from_secs(1), + remote_handle.start_pairing( + RemoteControlPairingStartParams::default(), + /*app_server_client_name*/ None, + ), ) .await - .expect("pairing Retry-After should throttle websocket refresh"); - let server_websocket = server_task.await.expect("server task should succeed"); + .expect("new pairing should report the deadline promptly") + .expect_err("new pairing must honor the handshake deadline"); + assert_eq!(remote_control_retry_at(&pairing_error), Some(retry_at)); + timeout(Duration::from_millis(100), listener.accept()) + .await + .expect_err("new pairing must not bypass the newly recorded deadline"); +} - assert_eq!( - pairing_response, - RemoteControlPairingStartResponse { - pairing_code: "pairing-code".to_string(), - manual_pairing_code: Some("ABCD-EFGH".to_string()), - environment_id: "env_test".to_string(), - expires_at: 33_336_362_096, +#[tokio::test] +async fn pairing_auth_recovery_respects_concurrent_handshake_overload() { + for (pairing_status, auth_endpoint) in + [("401 Unauthorized", "refresh"), ("404 Not Found", "enroll")] + { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let remote_control_url = remote_control_url_for_listener(&listener); + let remote_control_target = + normalize_remote_control_url(&remote_control_url).expect("target should parse"); + let codex_home = TempDir::new().expect("temp dir should create"); + save_auth( + codex_home.path(), + &remote_control_auth_dot_json("stale-token"), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("stale auth should save"); + let auth_manager = AuthManager::shared( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*forced_chatgpt_workspace_id*/ None, + /*chatgpt_base_url*/ None, + AuthKeyringBackendKind::default(), + codex_login::test_support::transport_default_auth_route_config(), + ) + .await; + let state_db = remote_control_state_runtime(&codex_home).await; + let mut remote_handle = remote_control_handle_with_current_enrollment( + &remote_control_url, + auth_manager.clone(), + ); + remote_handle.state_db = Some(state_db.clone()); + let current_enrollment = remote_handle.current_enrollment.clone(); + let connect = connect_test_websocket( + &remote_control_target, + state_db.as_ref(), + &auth_manager, + ¤t_enrollment, + ); + tokio::pin!(connect); + let (handshake_stream, request_line) = tokio::select! { + result = &mut connect => panic!("handshake should wait for its response: {result:?}"), + request = accept_http_request(&listener) => request, + }; + assert_eq!( + request_line, + "GET /backend-api/wham/remote/control/server HTTP/1.1" + ); + + let pairing = remote_handle.start_pairing( + RemoteControlPairingStartParams::default(), + /*app_server_client_name*/ None, + ); + tokio::pin!(pairing); + let (pairing_stream, request_line) = tokio::select! { + result = &mut pairing => panic!("pairing should wait for its response: {result:?}"), + request = accept_http_request(&listener) => request, + }; + assert_eq!( + request_line, + "POST /backend-api/wham/remote/control/server/pair HTTP/1.1" + ); + respond_with_status_and_headers(pairing_stream, pairing_status, &[], "retry auth").await; + let (auth_stream, request_line) = tokio::select! { + result = &mut pairing => panic!("auth should wait for its response: {result:?}"), + request = accept_http_request(&listener) => request, + }; + assert_eq!( + request_line, + format!("POST /backend-api/wham/remote/control/server/{auth_endpoint} HTTP/1.1") + ); + + respond_with_status_and_headers( + handshake_stream, + "503 Service Unavailable", + &[("Retry-After", "120")], + "overloaded", + ) + .await; + let connect_error = connect + .await + .expect_err("the handshake should report overload"); + let retry_at = remote_control_retry_at(&connect_error) + .expect("the handshake should preserve its retry deadline"); + save_auth( + codex_home.path(), + &remote_control_auth_dot_json("fresh-token"), + AuthCredentialsStoreMode::File, + AuthKeyringBackendKind::default(), + ) + .expect("replacement auth should save"); + respond_with_status_and_headers(auth_stream, "401 Unauthorized", &[], "stale auth").await; + let pairing_error = timeout(TEST_HTTP_ACCEPT_TIMEOUT, pairing) + .await + .expect("auth recovery should report the deadline promptly") + .expect_err("auth recovery must honor the handshake deadline"); + assert_eq!(remote_control_retry_at(&pairing_error), Some(retry_at)); + timeout(Duration::from_millis(100), listener.accept()) + .await + .expect_err("auth recovery must not send another request during overload"); + } +} + +#[tokio::test] +async fn pairing_overload_defers_pairing_and_websocket_requests() { + for status in ["429 Too Many Requests", "503 Service Unavailable"] { + for check_status in [false, true] { + for incomplete_body in [false, true] { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let remote_control_url = remote_control_url_for_listener(&listener); + let remote_control_target = + normalize_remote_control_url(&remote_control_url).expect("target should parse"); + let server_task = tokio::spawn(async move { + let (mut stream, request_line) = accept_http_request(&listener).await; + assert_eq!( + request_line, + if check_status { + "POST /backend-api/wham/remote/control/server/pair/status HTTP/1.1" + } else { + "POST /backend-api/wham/remote/control/server/pair HTTP/1.1" + } + ); + if incomplete_body { + stream.write_all(format!( + "HTTP/1.1 {status}\r\nContent-Length: 100\r\nRetry-After: 120\r\nConnection: close\r\n\r\npartial" + ).as_bytes()).await.expect("partial response should send"); + stream.shutdown().await.expect("response should close"); + } else { + respond_with_status_and_headers( + stream, + status, + &[("Retry-After", "120")], + "overloaded", + ) + .await; + } + listener + }); + let codex_home = TempDir::new().expect("temp dir should create"); + let state_db = remote_control_state_runtime(&codex_home).await; + let auth_manager = remote_control_auth_manager(); + let mut remote_handle = remote_control_handle_with_current_enrollment( + &remote_control_url, + auth_manager.clone(), + ); + remote_handle.state_db = Some(state_db.clone()); + let current_enrollment = remote_handle.current_enrollment.clone(); + let status_params = || RemoteControlPairingStatusParams { + pairing_code: Some("pairing-code".to_string()), + manual_pairing_code: None, + }; + let response = if check_status { + remote_handle + .pairing_status(status_params()) + .await + .map(|_| ()) + } else { + remote_handle + .start_pairing( + RemoteControlPairingStartParams::default(), + /*app_server_client_name*/ None, + ) + .await + .map(|_| ()) + }; + let error = response.expect_err("pairing should report overload"); + let listener = server_task.await.expect("server task should succeed"); + let retry_at = remote_control_retry_at(&error) + .expect("even a partial response must preserve the retry deadline"); + let pairing_error = remote_handle + .start_pairing( + RemoteControlPairingStartParams::default(), + /*app_server_client_name*/ None, + ) + .await + .expect_err("pairing must wait for the server deadline"); + let status_error = remote_handle + .pairing_status(status_params()) + .await + .expect_err("pairing status must wait for the server deadline"); + let connect_error = connect_test_websocket( + &remote_control_target, + state_db.as_ref(), + &auth_manager, + ¤t_enrollment, + ) + .await + .expect_err("the handshake must wait for the same deadline"); + for error in [pairing_error, status_error, connect_error] { + assert_eq!(remote_control_retry_at(&error), Some(retry_at)); + } + assert_eq!( + current_enrollment + .snapshot() + .and_then(|enrollment| enrollment.remote_control_token), + Some(TEST_REMOTE_CONTROL_SERVER_TOKEN.to_string()), + ); + timeout(Duration::from_millis(100), listener.accept()) + .await + .expect_err("no remote-control request should bypass the pairing deadline"); + } } - ); - drop(server_websocket); + } } async fn assert_refresh_failure_blocks_websocket( @@ -399,14 +696,16 @@ async fn assert_refresh_failure_blocks_websocket( server_task.await.expect("server task should succeed"); assert!(refresh_err.to_string().contains("HTTP 502 Bad Gateway")); assert_eq!(deferred_err.kind(), io::ErrorKind::WouldBlock); - assert!(deferred_err.to_string().contains("refresh deferred until")); + let retry_at = remote_control_retry_at(&refresh_err) + .expect("the overload response should retain its retry deadline"); + assert_eq!(remote_control_retry_at(&deferred_err), Some(retry_at)); let next_refresh_at = current_enrollment .snapshot() .and_then(|enrollment| enrollment.next_refresh_at) .expect("required refresh failure should set a retry deadline"); assert!( (refresh_started_at + time::Duration::seconds(120) - ..=refresh_completed_at + time::Duration::seconds(120)) + ..=refresh_completed_at + time::Duration::seconds(150)) .contains(&next_refresh_at) ); }