From 974cee60db6221ebcd7a47b94ce8942a8249fb12 Mon Sep 17 00:00:00 2001 From: Anton Panasenko Date: Fri, 29 May 2026 04:47:09 -0700 Subject: [PATCH] fix(app-server-transport): refresh pairing auth in place --- .../src/transport/remote_control/enroll.rs | 23 + .../src/transport/remote_control/mod.rs | 14 + .../pairing_integration_tests.rs | 183 +++++++ .../src/transport/remote_control/tests.rs | 18 +- .../src/transport/remote_control/websocket.rs | 445 +++++++++++++++++- 5 files changed, 663 insertions(+), 20 deletions(-) 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 ec76158586..7f87ee308c 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 @@ -46,6 +46,13 @@ impl RemoteControlEnrollment { }) } + pub(super) fn server_token_refresh_delay(&self) -> Option { + let refresh_at = self.expires_at? + - time::Duration::seconds(REMOTE_CONTROL_SERVER_TOKEN_REFRESH_SKEW_SECS); + let refresh_delay = refresh_at - OffsetDateTime::now_utc(); + Some(std::time::Duration::try_from(refresh_delay).unwrap_or_default()) + } + pub(super) fn clear_server_token(&mut self) { self.remote_control_token = None; self.expires_at = None; @@ -435,6 +442,22 @@ mod tests { assert!(!expires_later.should_refresh_server_token()); } + #[test] + fn remote_control_enrollment_schedules_server_token_refresh_before_expiry() { + let refresh_delay = RemoteControlEnrollment { + account_id: "account_id".to_string(), + environment_id: "environment_id".to_string(), + server_id: "server_id".to_string(), + server_name: "server_name".to_string(), + remote_control_token: Some("remote-control-token".to_string()), + expires_at: Some(OffsetDateTime::now_utc() + time::Duration::seconds(31)), + } + .server_token_refresh_delay() + .expect("server token refresh should be scheduled"); + + assert!(refresh_delay <= std::time::Duration::from_secs(1)); + } + #[test] fn preview_remote_control_response_body_redacts_server_token() { assert_eq!( 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 76e55679bd..5f6dfda15f 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 @@ -60,6 +60,8 @@ pub struct RemoteControlHandle { status_tx: Arc>, state_db_available: bool, pairing: PairingClientState, + #[cfg(test)] + pairing_refresh_tx: Arc>, auth_change_rx: Arc>>, } @@ -230,6 +232,12 @@ impl RemoteControlHandle { }) } + #[cfg(test)] + fn request_pairing_auth_refresh(&self) { + self.pairing_refresh_tx + .send_modify(|revision| *revision = revision.wrapping_add(1)); + } + fn publish_status( &self, connection_status: RemoteControlConnectionStatus, @@ -309,6 +317,8 @@ pub async fn start_remote_control( let (enabled_tx, enabled_rx) = watch::channel(initial_enabled); let pairing = PairingClientState::new(); let websocket_pairing = pairing.clone(); + let (pairing_refresh_tx, pairing_refresh_rx) = watch::channel(0u64); + let websocket_pairing_refresh_tx = pairing_refresh_tx.clone(); let auth_change_rx = Arc::new(StdMutex::new(auth_manager.auth_change_receiver())); let server_name = gethostname().to_string_lossy().trim().to_string(); let remote_control_url = config.remote_control_url; @@ -358,6 +368,8 @@ pub async fn start_remote_control( transport_event_tx, status_publisher, pairing: websocket_pairing, + pairing_refresh_tx: websocket_pairing_refresh_tx, + pairing_refresh_rx, }, shutdown_token, enabled_rx, @@ -403,6 +415,8 @@ pub async fn start_remote_control( status_tx: Arc::new(status_tx), state_db_available, pairing, + #[cfg(test)] + pairing_refresh_tx: Arc::new(pairing_refresh_tx), auth_change_rx, }, )) diff --git a/codex-rs/app-server-transport/src/transport/remote_control/pairing_integration_tests.rs b/codex-rs/app-server-transport/src/transport/remote_control/pairing_integration_tests.rs index beb76408a5..3db576d8ef 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/pairing_integration_tests.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/pairing_integration_tests.rs @@ -394,3 +394,186 @@ async fn remote_control_handle_clears_pairing_client_after_auth_change() { shutdown_token.cancel(); let _ = remote_task.await; } + +#[tokio::test] +async fn remote_control_refreshes_server_token_while_connected() { + 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 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 (remote_task, remote_handle) = start_remote_control( + RemoteControlStartConfig { + remote_control_url, + installation_id: TEST_INSTALLATION_ID.to_string(), + }, + Some(remote_control_state_runtime(&codex_home).await), + remote_control_auth_manager(), + transport_event_tx, + shutdown_token.clone(), + /*app_server_client_name_rx*/ None, + /*initial_enabled*/ true, + ) + .await + .expect("remote control should start"); + + let enroll_request = accept_http_request(&listener).await; + respond_with_json( + enroll_request.stream, + remote_control_server_token_response( + "srv_e_test", + "env_test", + TEST_REMOTE_CONTROL_SERVER_TOKEN, + ), + ) + .await; + let mut first_websocket = accept_remote_control_connection(&listener).await; + + remote_handle.request_pairing_auth_refresh(); + let refresh_request = accept_http_request(&listener).await; + assert_eq!( + refresh_request.request_line, + "POST /backend-api/wham/remote/control/server/refresh HTTP/1.1" + ); + respond_with_json( + refresh_request.stream, + remote_control_server_token_response( + "srv_e_test", + "env_test", + TEST_REFRESHED_REMOTE_CONTROL_SERVER_TOKEN, + ), + ) + .await; + assert!( + timeout(Duration::from_millis(100), first_websocket.next()) + .await + .is_err(), + "server token refresh should keep the websocket open" + ); + + let pairing_task = tokio::spawn({ + let remote_handle = remote_handle.clone(); + async move { + remote_handle + .start_pairing(RemoteControlPairingStartParams::default()) + .await + } + }); + let pairing_request = accept_http_request(&listener).await; + assert_eq!( + pairing_request.request_line, + "POST /backend-api/wham/remote/control/server/pair HTTP/1.1" + ); + assert_eq!( + pairing_request.headers.get("authorization"), + Some(&format!( + "Bearer {TEST_REFRESHED_REMOTE_CONTROL_SERVER_TOKEN}" + )) + ); + respond_with_json( + pairing_request.stream, + json!({ + "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; + assert_eq!( + pairing_task + .await + .expect("pairing task should join") + .expect("pairing should use refreshed server token"), + codex_app_server_protocol::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, + } + ); + + first_websocket + .close(None) + .await + .expect("first websocket should close"); + + shutdown_token.cancel(); + let _ = remote_task.await; +} + +#[tokio::test] +async fn remote_control_schedules_server_token_refresh_while_connected() { + 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 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 (remote_task, _remote_handle) = start_remote_control( + RemoteControlStartConfig { + remote_control_url, + installation_id: TEST_INSTALLATION_ID.to_string(), + }, + Some(remote_control_state_runtime(&codex_home).await), + remote_control_auth_manager(), + transport_event_tx, + shutdown_token.clone(), + /*app_server_client_name_rx*/ None, + /*initial_enabled*/ true, + ) + .await + .expect("remote control should start"); + + let scheduled_refresh_expires_at = (OffsetDateTime::now_utc() + time::Duration::seconds(35)) + .format(&time::format_description::well_known::Rfc3339) + .expect("scheduled refresh expiry should format"); + let enroll_request = accept_http_request(&listener).await; + respond_with_json( + enroll_request.stream, + remote_control_server_token_response_with_expires_at( + "srv_e_test", + "env_test", + TEST_REMOTE_CONTROL_SERVER_TOKEN, + &scheduled_refresh_expires_at, + ), + ) + .await; + let mut first_websocket = accept_remote_control_connection(&listener).await; + + let refresh_request = timeout(Duration::from_secs(10), accept_http_request(&listener)) + .await + .expect("scheduled server token refresh should arrive"); + assert_eq!( + refresh_request.request_line, + "POST /backend-api/wham/remote/control/server/refresh HTTP/1.1" + ); + respond_with_json( + refresh_request.stream, + remote_control_server_token_response( + "srv_e_test", + "env_test", + TEST_REFRESHED_REMOTE_CONTROL_SERVER_TOKEN, + ), + ) + .await; + assert!( + timeout(Duration::from_millis(100), first_websocket.next()) + .await + .is_err(), + "scheduled server token refresh should keep the websocket open" + ); + + first_websocket + .close(None) + .await + .expect("first websocket should close"); + shutdown_token.cancel(); + let _ = remote_task.await; +} 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 2512aaede3..0e43e6d92e 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 @@ -144,6 +144,7 @@ fn remote_control_handle_with_pairing_client( installation_id: TEST_INSTALLATION_ID.to_string(), environment_id: Some("env_test".to_string()), }); + let (pairing_refresh_tx, _pairing_refresh_rx) = watch::channel(0u64); let pairing_client = Arc::new(StdMutex::new(Some(RemoteControlPairingClient::new( &normalize_remote_control_url(remote_control_url) .expect("remote control target should normalize"), @@ -162,6 +163,7 @@ fn remote_control_handle_with_pairing_client( client: pairing_client, generation: Arc::new(std::sync::atomic::AtomicU64::new(0)), }, + pairing_refresh_tx: Arc::new(pairing_refresh_tx), auth_change_rx: Arc::new(StdMutex::new(auth_change_rx)), } } @@ -170,12 +172,26 @@ fn remote_control_server_token_response( server_id: &str, environment_id: &str, remote_control_token: &str, +) -> serde_json::Value { + remote_control_server_token_response_with_expires_at( + server_id, + environment_id, + remote_control_token, + TEST_REMOTE_CONTROL_SERVER_TOKEN_EXPIRES_AT, + ) +} + +fn remote_control_server_token_response_with_expires_at( + server_id: &str, + environment_id: &str, + remote_control_token: &str, + expires_at: &str, ) -> serde_json::Value { json!({ "server_id": server_id, "environment_id": environment_id, "remote_control_token": remote_control_token, - "expires_at": TEST_REMOTE_CONTROL_SERVER_TOKEN_EXPIRES_AT, + "expires_at": expires_at, }) } 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 d6bf24f791..2ad1abc246 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 @@ -256,6 +256,8 @@ pub(crate) struct RemoteControlWebsocket { auth_recovery: UnauthorizedRecovery, auth_change_rx: watch::Receiver, pairing: PairingClientState, + _pairing_refresh_tx: watch::Sender, + pairing_refresh_rx: watch::Receiver, client_tracker: Arc>, state: Arc>, server_event_rx: Arc>>, @@ -289,13 +291,50 @@ enum ConnectionEndReason { EnabledWatchClosed, AuthChanged, AuthWatchClosed, + PairingRefreshWatchClosed, + ServerTokenRefreshRejected, + StaleEnrollment, ConnectionWorkerStopped, } +enum ConnectionLoopAction { + End(ConnectionEndReason), + RefreshServerToken, +} + +enum ConnectedServerTokenRefreshAction { + Continue, + Retry(io::Error), + End(ConnectionEndReason), +} + +struct ConnectedServerTokenRefreshRequest { + remote_control_target: RemoteControlTarget, + auth: RemoteControlConnectionAuth, + installation_id: String, + enrollment: RemoteControlEnrollment, + auth_change_revision: u64, +} + +struct PendingConnectedServerTokenRefreshRequest { + remote_control_target: RemoteControlTarget, + auth_manager: Arc, + auth_change_rx: watch::Receiver, + installation_id: String, + enrollment: RemoteControlEnrollment, +} + +struct ConnectedServerTokenRefreshResponse { + enrollment: RemoteControlEnrollment, + auth_change_revision: u64, +} + pub(super) struct RemoteControlChannels { pub(super) transport_event_tx: mpsc::Sender, pub(super) status_publisher: RemoteControlStatusPublisher, pub(super) pairing: PairingClientState, + pub(super) pairing_refresh_tx: watch::Sender, + pub(super) pairing_refresh_rx: watch::Receiver, } #[derive(Clone)] @@ -413,6 +452,8 @@ impl RemoteControlWebsocket { auth_recovery, auth_change_rx, pairing: channels.pairing, + _pairing_refresh_tx: channels.pairing_refresh_tx, + pairing_refresh_rx: channels.pairing_refresh_rx, client_tracker: Arc::new(Mutex::new(client_tracker)), state: Arc::new(Mutex::new(WebsocketState { outbound_buffer, @@ -498,7 +539,11 @@ impl RemoteControlWebsocket { }; let connection_end_reason = self - .run_connection(websocket_connection, shutdown_token) + .run_connection( + websocket_connection, + shutdown_token, + app_server_client_name.as_deref(), + ) .await; let status = self.status_publisher.status(); info!( @@ -710,6 +755,7 @@ impl RemoteControlWebsocket { &mut self, websocket_connection: WebSocketStream>, shutdown_token: CancellationToken, + app_server_client_name: Option<&str>, ) -> ConnectionEndReason { let (websocket_writer, websocket_reader) = websocket_connection.split(); let mut join_set = tokio::task::JoinSet::new(); @@ -731,26 +777,79 @@ impl RemoteControlWebsocket { )); let mut enabled_rx = self.enabled_rx.clone(); - let connection_end_reason = tokio::select! { - _ = shutdown_token.cancelled() => ConnectionEndReason::Shutdown, - changed = enabled_rx.wait_for(|enabled| !*enabled) => { - if changed.is_ok() { - self.status_publisher - .publish_status(RemoteControlConnectionStatus::Disabled); - ConnectionEndReason::Disabled - } else { - ConnectionEndReason::EnabledWatchClosed + let mut server_token_refresh_retry_delay = None; + let connection_end_reason = loop { + let server_token_refresh_delay = + server_token_refresh_retry_delay.take().or_else(|| { + self.enrollment + .as_ref() + .and_then(RemoteControlEnrollment::server_token_refresh_delay) + }); + let server_token_refresh = async move { + match server_token_refresh_delay { + Some(delay) => tokio::time::sleep(delay).await, + None => std::future::pending().await, + } + }; + tokio::pin!(server_token_refresh); + let connection_loop_action = tokio::select! { + _ = shutdown_token.cancelled() => ConnectionLoopAction::End(ConnectionEndReason::Shutdown), + changed = enabled_rx.wait_for(|enabled| !*enabled) => { + if changed.is_ok() { + self.status_publisher + .publish_status(RemoteControlConnectionStatus::Disabled); + ConnectionLoopAction::End(ConnectionEndReason::Disabled) + } else { + ConnectionLoopAction::End(ConnectionEndReason::EnabledWatchClosed) + } + } + changed = self.auth_change_rx.changed() => { + if changed.is_ok() { + self.auth_recovery = self.auth_manager.unauthorized_recovery(); + ConnectionLoopAction::End(ConnectionEndReason::AuthChanged) + } else { + ConnectionLoopAction::End(ConnectionEndReason::AuthWatchClosed) + } + } + changed = self.pairing_refresh_rx.changed() => { + if changed.is_ok() { + ConnectionLoopAction::RefreshServerToken + } else { + ConnectionLoopAction::End(ConnectionEndReason::PairingRefreshWatchClosed) + } + } + _ = &mut server_token_refresh => ConnectionLoopAction::RefreshServerToken, + _ = join_set.join_next() => ConnectionLoopAction::End(ConnectionEndReason::ConnectionWorkerStopped), + }; + match connection_loop_action { + ConnectionLoopAction::End(connection_end_reason) => break connection_end_reason, + ConnectionLoopAction::RefreshServerToken => { + match self + .refresh_connected_server_token_while_connected( + &mut enabled_rx, + &mut join_set, + &shutdown_token, + app_server_client_name, + ) + .await + { + ConnectedServerTokenRefreshAction::Continue => {} + ConnectedServerTokenRefreshAction::Retry(err) => { + warn!( + error = %err, + error_kind = ?err.kind(), + retry_delay = ?REMOTE_CONTROL_ACCOUNT_ID_RETRY_INTERVAL, + "failed to refresh connected app-server remote control server token" + ); + server_token_refresh_retry_delay = + Some(REMOTE_CONTROL_ACCOUNT_ID_RETRY_INTERVAL); + } + ConnectedServerTokenRefreshAction::End(connection_end_reason) => { + break connection_end_reason; + } + } } } - changed = self.auth_change_rx.changed() => { - if changed.is_ok() { - self.auth_recovery = self.auth_manager.unauthorized_recovery(); - ConnectionEndReason::AuthChanged - } else { - ConnectionEndReason::AuthWatchClosed - } - } - _ = join_set.join_next() => ConnectionEndReason::ConnectionWorkerStopped, }; clear_pairing_client(&self.pairing); shutdown_token.cancel(); @@ -760,6 +859,209 @@ impl RemoteControlWebsocket { connection_end_reason } + async fn refresh_connected_server_token_while_connected( + &mut self, + enabled_rx: &mut watch::Receiver, + join_set: &mut tokio::task::JoinSet<()>, + shutdown_token: &CancellationToken, + app_server_client_name: Option<&str>, + ) -> ConnectedServerTokenRefreshAction { + let refresh_request = match self.connected_server_token_refresh_request() { + Ok(refresh_request) => refresh_request, + Err(err) => return ConnectedServerTokenRefreshAction::Retry(err), + }; + let refresh_request = tokio::select! { + _ = shutdown_token.cancelled() => { + return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::Shutdown); + } + changed = enabled_rx.wait_for(|enabled| !*enabled) => { + if changed.is_ok() { + self.status_publisher + .publish_status(RemoteControlConnectionStatus::Disabled); + return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::Disabled); + } + return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::EnabledWatchClosed); + } + _ = join_set.join_next() => { + return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::ConnectionWorkerStopped); + } + refresh_request = prepare_connected_server_token_refresh(refresh_request) => refresh_request, + }; + let refresh_request = match refresh_request { + Ok(refresh_request) => refresh_request, + Err(err) => { + return self + .apply_connected_server_token_refresh(Err(err), app_server_client_name) + .await; + } + }; + if !mark_connected_refresh_auth_change_seen( + &mut self.auth_change_rx, + refresh_request.auth_change_revision, + ) { + clear_pairing_client(&self.pairing); + self.auth_recovery = self.auth_manager.unauthorized_recovery(); + return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::AuthChanged); + } + let refresh_result = tokio::select! { + _ = shutdown_token.cancelled() => { + return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::Shutdown); + } + changed = enabled_rx.wait_for(|enabled| !*enabled) => { + if changed.is_ok() { + self.status_publisher + .publish_status(RemoteControlConnectionStatus::Disabled); + return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::Disabled); + } + return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::EnabledWatchClosed); + } + changed = self.auth_change_rx.changed() => { + if changed.is_ok() { + self.auth_recovery = self.auth_manager.unauthorized_recovery(); + return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::AuthChanged); + } + return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::AuthWatchClosed); + } + _ = join_set.join_next() => { + return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::ConnectionWorkerStopped); + } + refresh_result = refresh_connected_server_token(refresh_request) => refresh_result, + }; + self.apply_connected_server_token_refresh(refresh_result, app_server_client_name) + .await + } + + fn connected_server_token_refresh_request( + &self, + ) -> io::Result { + let remote_control_target = self.remote_control_target.clone().ok_or_else(|| { + io::Error::other("missing remote control target while refreshing server token") + })?; + let enrollment = self.enrollment.clone().ok_or_else(|| { + io::Error::other("missing remote control enrollment while refreshing server token") + })?; + Ok(PendingConnectedServerTokenRefreshRequest { + remote_control_target, + auth_manager: self.auth_manager.clone(), + auth_change_rx: self.auth_change_rx.clone(), + installation_id: self.installation_id.clone(), + enrollment, + }) + } + + async fn apply_connected_server_token_refresh( + &mut self, + refresh_result: io::Result, + app_server_client_name: Option<&str>, + ) -> ConnectedServerTokenRefreshAction { + let refresh_response = match refresh_result { + Ok(refresh_response) => refresh_response, + Err(err) if err.kind() == ErrorKind::NotFound => { + self.clear_stale_connected_enrollment(app_server_client_name) + .await; + return ConnectedServerTokenRefreshAction::End( + ConnectionEndReason::StaleEnrollment, + ); + } + Err(err) if err.kind() == ErrorKind::PermissionDenied => { + if recover_remote_control_auth(&mut self.auth_recovery, &mut self.auth_change_rx) + .await + { + return ConnectedServerTokenRefreshAction::Retry(io::Error::other(format!( + "{err}; retrying after auth recovery" + ))); + } + if let Some(enrollment) = self.enrollment.as_mut() { + enrollment.clear_server_token(); + } + clear_pairing_client(&self.pairing); + return ConnectedServerTokenRefreshAction::End( + ConnectionEndReason::ServerTokenRefreshRejected, + ); + } + Err(err) if err.kind() == ErrorKind::InvalidInput => { + clear_pairing_client(&self.pairing); + self.auth_recovery = self.auth_manager.unauthorized_recovery(); + return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::AuthChanged); + } + Err(err) => return ConnectedServerTokenRefreshAction::Retry(err), + }; + if *self.auth_change_rx.borrow() != refresh_response.auth_change_revision { + clear_pairing_client(&self.pairing); + self.auth_recovery = self.auth_manager.unauthorized_recovery(); + return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::AuthChanged); + } + let current_enrollment = match self.enrollment.as_ref() { + Some(current_enrollment) => current_enrollment, + None => return ConnectedServerTokenRefreshAction::Continue, + }; + if !same_remote_control_enrollment_identity( + current_enrollment, + &refresh_response.enrollment, + ) { + return ConnectedServerTokenRefreshAction::Continue; + } + self.enrollment = Some(refresh_response.enrollment); + let Some(remote_control_target) = self.remote_control_target.as_ref() else { + return ConnectedServerTokenRefreshAction::Retry(io::Error::other( + "missing remote control target after refreshing server token", + )); + }; + let Some(enrollment) = self.enrollment.as_ref() else { + return ConnectedServerTokenRefreshAction::Retry(io::Error::other( + "missing remote control enrollment after refreshing server token", + )); + }; + match set_pairing_client( + &self.pairing, + remote_control_target, + enrollment, + refresh_response.auth_change_revision, + ) { + Ok(()) => ConnectedServerTokenRefreshAction::Continue, + Err(err) => ConnectedServerTokenRefreshAction::Retry(err), + } + } + + async fn clear_stale_connected_enrollment(&mut self, app_server_client_name: Option<&str>) { + let Some(remote_control_target) = self.remote_control_target.as_ref() else { + clear_pairing_client(&self.pairing); + self.enrollment = None; + self.status_publisher + .publish_environment_id(/*environment_id*/ None); + return; + }; + let Some(account_id) = self + .enrollment + .as_ref() + .map(|enrollment| enrollment.account_id.clone()) + else { + clear_pairing_client(&self.pairing); + return; + }; + info!( + "connected remote control server refresh returned HTTP 404; clearing stale enrollment before re-enrolling: websocket_url={}, account_id={}", + remote_control_target.websocket_url, account_id + ); + if let Some(state_db) = self.state_db.as_deref() { + clear_remote_control_enrollment( + state_db, + remote_control_target, + &account_id, + app_server_client_name, + &mut self.enrollment, + &self.status_publisher, + &self.pairing, + ) + .await; + return; + } + self.enrollment = None; + self.status_publisher + .publish_environment_id(/*environment_id*/ None); + clear_pairing_client(&self.pairing); + } + async fn join_connection_workers( join_set: &mut tokio::task::JoinSet<()>, shutdown_timeout: std::time::Duration, @@ -1186,6 +1488,73 @@ fn build_remote_control_websocket_request( Ok(request) } +async fn refresh_connected_server_token( + mut refresh_request: ConnectedServerTokenRefreshRequest, +) -> io::Result { + info!( + "refreshing connected remote control server token: websocket_url={}, refresh_url={}, account_id={}, server_id={}, environment_id={}", + refresh_request.remote_control_target.websocket_url, + refresh_request.remote_control_target.refresh_url, + refresh_request.auth.account_id, + refresh_request.enrollment.server_id, + refresh_request.enrollment.environment_id + ); + refresh_remote_control_server( + &refresh_request.remote_control_target, + &refresh_request.auth, + &refresh_request.installation_id, + &mut refresh_request.enrollment, + ) + .await?; + Ok(ConnectedServerTokenRefreshResponse { + enrollment: refresh_request.enrollment, + auth_change_revision: refresh_request.auth_change_revision, + }) +} + +async fn prepare_connected_server_token_refresh( + mut refresh_request: PendingConnectedServerTokenRefreshRequest, +) -> io::Result { + let auth = load_remote_control_auth(&refresh_request.auth_manager).await?; + // Loading auth may reload or proactively refresh through the same watch + // receiver. Treat the auth used for this refresh as the current revision + // before waiting for later auth changes to cancel the HTTP request. + let auth_change_revision = *refresh_request.auth_change_rx.borrow_and_update(); + if refresh_request.enrollment.account_id != auth.account_id { + return Err(io::Error::new( + ErrorKind::InvalidInput, + "remote control auth changed while refreshing server token", + )); + } + Ok(ConnectedServerTokenRefreshRequest { + remote_control_target: refresh_request.remote_control_target, + auth, + installation_id: refresh_request.installation_id, + enrollment: refresh_request.enrollment, + auth_change_revision, + }) +} + +fn mark_connected_refresh_auth_change_seen( + auth_change_rx: &mut watch::Receiver, + auth_change_revision: u64, +) -> bool { + if *auth_change_rx.borrow() != auth_change_revision { + return false; + } + auth_change_rx.borrow_and_update(); + true +} + +fn same_remote_control_enrollment_identity( + left: &RemoteControlEnrollment, + right: &RemoteControlEnrollment, +) -> bool { + left.account_id == right.account_id + && left.server_id == right.server_id + && left.environment_id == right.environment_id +} + pub(crate) async fn load_remote_control_auth( auth_manager: &Arc, ) -> io::Result { @@ -1802,6 +2171,41 @@ mod tests { ); } + #[test] + fn mark_connected_refresh_auth_change_seen_marks_loaded_auth_revision_seen() { + let (auth_change_tx, mut auth_change_rx) = watch::channel(0u64); + auth_change_tx.send_modify(|revision| *revision += 1); + let auth_change_revision = *auth_change_rx.borrow(); + + assert!(mark_connected_refresh_auth_change_seen( + &mut auth_change_rx, + auth_change_revision + )); + assert!( + !auth_change_rx + .has_changed() + .expect("auth change watch should remain open") + ); + } + + #[test] + fn mark_connected_refresh_auth_change_seen_preserves_racing_auth_change() { + let (auth_change_tx, mut auth_change_rx) = watch::channel(0u64); + auth_change_tx.send_modify(|revision| *revision += 1); + let auth_change_revision = *auth_change_rx.borrow(); + auth_change_tx.send_modify(|revision| *revision += 1); + + assert!(!mark_connected_refresh_auth_change_seen( + &mut auth_change_rx, + auth_change_revision + )); + assert!( + auth_change_rx + .has_changed() + .expect("auth change watch should remain open") + ); + } + async fn remote_control_state_runtime(codex_home: &TempDir) -> Arc { StateRuntime::init(codex_home.path().to_path_buf(), "test-provider".to_string()) .await @@ -2325,6 +2729,7 @@ mod tests { let (status_publisher, _status_rx) = remote_control_status_channel(); let shutdown_token = CancellationToken::new(); let (_enabled_tx, enabled_rx) = watch::channel(true); + let (pairing_refresh_tx, pairing_refresh_rx) = watch::channel(0u64); let websocket_task = tokio::spawn({ let shutdown_token = shutdown_token.clone(); async move { @@ -2341,6 +2746,8 @@ mod tests { transport_event_tx, status_publisher, pairing: test_pairing(), + pairing_refresh_tx, + pairing_refresh_rx, }, shutdown_token, enabled_rx,