openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-27 18:48:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Honor shared Retry-After deadlines for remote control (#44311)

Browse Source 
## Why

Remote-control requests could bypass a server's `Retry-After` deadline through pairing, auth changes, or reconnects. Proactive token refresh could also continue using a valid token despite an explicit server delay.

## What changed

- Share retry deadlines across enrollment, token refresh, pairing, pairing status, and WebSocket handshakes, preserving them across enrollment replacement and disable/enable transitions.
- Honor `Retry-After` on HTTP `429` and `503` pairing and handshake responses, including incomplete pairing response bodies.
- Add up to 30 seconds of jitter once per response without shortening the server's deadline, including for `Retry-After: 0`.
- Defer new requests while retaining valid tokens, and prevent auth changes from shortening the server delay while keeping shutdown responsive.

## Testing

Add regression coverage for shared deadlines, concurrent pairing and handshake failures, auth recovery, disable/enable transitions, retries after expiry, prompt shutdown, and jitter bounds.

GitOrigin-RevId: 69a5a46559031709a4558190c3142eb65e2ae634
This commit is contained in:
cooper-oai
2026-09-09 21:50:16 +00:00
committed by copyberry
parent 45eec73b11
commit 434efa95e6
9 changed files with 869 additions and 140 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server-transport/src/transport/remote_control/enroll.rs
+52 -12
View File
@@ -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<OffsetDateTime>,
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,
codex-rs/app-server-transport/src/transport/remote_control/mod.rs
+64 -5
View File
@@ -132,6 +132,8 @@ enum RemoteControlEnrollmentSelection {
struct RemoteControlEnrollmentState {
enrollment: StdMutex<Option<RemoteControlEnrollment>>,
// Keep an observed server deadline across enrollment replacement and enable transitions.
retry_at: StdMutex<Option<time::OffsetDateTime>>,
lock: Semaphore,
}
@@ -139,10 +141,48 @@ impl RemoteControlEnrollmentState {
fn new(enrollment: Option<RemoteControlEnrollment>) -> Self {
Self {
enrollment: StdMutex::new(enrollment),
retry_at: StdMutex::new(None),
lock: Semaphore::new(1),
}
}
async fn lock_for_request(&self) -> io::Result<RemoteControlEnrollmentLease<'_>> {
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<T>(&self, result: io::Result<T>) -> io::Result<T> {
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(&params)?;
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<AuthManager>,
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<RemoteControlEnrollment>,
current_enrollment: &mut RemoteControlEnrollmentLease<'_>,
auth_manager: &Arc<AuthManager>,
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)
}
}
codex-rs/app-server-transport/src/transport/remote_control/server_api.rs
+52 -9
View File
@@ -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<StatusCode>,
retry_at: Option<OffsetDateTime>,
}
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<StatusCode>,
retry_at: Option<OffsetDateTime>,
@@ -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<OffsetDateTime> {
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> {
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<OffsetDateTime> {
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<OffsetDateTime> {
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(
codex-rs/app-server-transport/src/transport/remote_control/server_api_tests.rs
+37
View File
@@ -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));
}
codex-rs/app-server-transport/src/transport/remote_control/tests.rs
+2
View File
@@ -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";
codex-rs/app-server-transport/src/transport/remote_control/tests/pairing_tests.rs
+5 -4
View File
@@ -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)
);
}
codex-rs/app-server-transport/src/transport/remote_control/tests/retry_tests.rs
+187
View File
@@ -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");
}
codex-rs/app-server-transport/src/transport/remote_control/websocket.rs
+74 -13
View File
@@ -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<u64>,
server_retry_delay: Option<std::time::Duration>,
) -> 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,
codex-rs/app-server-transport/src/transport/remote_control/websocket_refresh_tests.rs
+396 -97
View File
@@ -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,
&current_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,
&current_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,
&current_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,
&current_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,
&current_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)
);
}