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 38e34e6d4..443dcc897 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 @@ -9,7 +9,10 @@ mod websocket; use self::auth::load_remote_control_auth; use self::auth::recover_remote_control_auth; use self::enroll::RemoteControlEnrollment; +use self::enroll::enroll_remote_control_server; +use self::enroll::load_persisted_remote_control_enrollment; use self::enroll::refresh_remote_control_server; +use self::enroll::update_persisted_remote_control_enrollment; use crate::transport::remote_control::websocket::RemoteControlChannels; use crate::transport::remote_control::websocket::RemoteControlStatusPublisher; use crate::transport::remote_control::websocket::RemoteControlWebsocket; @@ -36,9 +39,13 @@ use gethostname::gethostname; use std::error::Error; use std::fmt; use std::io; +use std::ops::Deref; +use std::ops::DerefMut; use std::panic::AssertUnwindSafe; use std::sync::Arc; use std::sync::Mutex as StdMutex; +use tokio::sync::Semaphore; +use tokio::sync::SemaphorePermit; use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio::sync::watch; @@ -65,12 +72,81 @@ pub struct RemoteControlHandle { enabled_tx: Arc>, status_tx: Arc>, state_db_available: bool, + state_db: Option>, remote_control_url: String, current_enrollment: CurrentRemoteControlEnrollment, + pairing_persistence_key: RemoteControlPairingPersistenceKey, + pairing_persistence_key_required: bool, auth_manager: Arc, } -type CurrentRemoteControlEnrollment = Arc>>; +// Pairing and websocket connect share one selected server so they cannot enroll or clear +// different persisted rows while either path is awaiting backend I/O. +type CurrentRemoteControlEnrollment = Arc; +type RemoteControlPairingPersistenceKey = watch::Sender>; + +struct RemoteControlEnrollmentState { + enrollment: StdMutex>, + lock: Semaphore, +} + +impl RemoteControlEnrollmentState { + fn new(enrollment: Option) -> Self { + Self { + enrollment: StdMutex::new(enrollment), + lock: Semaphore::new(1), + } + } + + async fn lock(&self) -> RemoteControlEnrollmentLease<'_> { + let permit = match self.lock.acquire().await { + Ok(permit) => permit, + Err(_) => unreachable!("remote control enrollment lock should stay open"), + }; + RemoteControlEnrollmentLease { + state: self, + enrollment: self.snapshot(), + _permit: permit, + } + } + + fn snapshot(&self) -> Option { + self.enrollment + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + } +} + +struct RemoteControlEnrollmentLease<'a> { + state: &'a RemoteControlEnrollmentState, + enrollment: Option, + _permit: SemaphorePermit<'a>, +} + +impl Deref for RemoteControlEnrollmentLease<'_> { + type Target = Option; + + fn deref(&self) -> &Self::Target { + &self.enrollment + } +} + +impl DerefMut for RemoteControlEnrollmentLease<'_> { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.enrollment + } +} + +impl Drop for RemoteControlEnrollmentLease<'_> { + fn drop(&mut self) { + *self + .state + .enrollment + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = self.enrollment.take(); + } +} #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct RemoteControlUnavailable; @@ -126,8 +202,6 @@ impl RemoteControlHandle { *state = false; changed }); - clear_current_enrollment(&self.current_enrollment); - let status = self.status(); info!( enabled_changed, @@ -151,28 +225,30 @@ impl RemoteControlHandle { pub async fn start_pairing( &self, params: RemoteControlPairingStartParams, + app_server_client_name: Option<&str>, ) -> io::Result { - if !*self.enabled_tx.borrow() { - return Err(Self::pairing_disabled_error()); - } let mut auth = load_remote_control_auth(&self.auth_manager) .await .map_err(|_| pairing_unavailable_error())?; - let mut enrollment = { - let current_enrollment = self - .current_enrollment - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - current_enrollment - .as_ref() - .filter(|enrollment| enrollment.account_id == auth.account_id) - .cloned() - } - .ok_or_else(pairing_unavailable_error)?; - let installation_id = self.status().installation_id; + let status = self.status(); + 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, + &mut auth, + &installation_id, + &status.server_name, + app_server_client_name, + ) + .await?; if enrollment.should_refresh_server_token() { refresh_pairing_enrollment( - &self.current_enrollment, + &mut current_enrollment, + self.state_db.as_deref(), + app_server_client_name, &self.auth_manager, &mut auth, &installation_id, @@ -185,9 +261,11 @@ impl RemoteControlHandle { }; let pairing_response = match enrollment.start_pairing(pairing_request()).await { Err(err) if err.kind() == io::ErrorKind::PermissionDenied => { - clear_pairing_server_token(&self.current_enrollment, &mut enrollment)?; + clear_pairing_server_token(&mut current_enrollment, &mut enrollment)?; refresh_pairing_enrollment( - &self.current_enrollment, + &mut current_enrollment, + self.state_db.as_deref(), + app_server_client_name, &self.auth_manager, &mut auth, &installation_id, @@ -196,24 +274,46 @@ impl RemoteControlHandle { .await?; enrollment.start_pairing(pairing_request()).await } + Err(err) if err.kind() == io::ErrorKind::NotFound => { + clear_pairing_enrollment( + &mut current_enrollment, + self.state_db.as_deref(), + app_server_client_name, + &enrollment, + ) + .await; + enrollment = self + .load_or_enroll_pairing_server( + &mut current_enrollment, + &mut auth, + &installation_id, + &status.server_name, + app_server_client_name, + ) + .await?; + enrollment.start_pairing(pairing_request()).await + } pairing_response => pairing_response, }; if let Err(err) = &pairing_response { match err.kind() { io::ErrorKind::NotFound => { - clear_current_enrollment_if_matches(&self.current_enrollment, &enrollment); + clear_pairing_enrollment( + &mut current_enrollment, + self.state_db.as_deref(), + app_server_client_name, + &enrollment, + ) + .await; return Err(pairing_unavailable_error()); } io::ErrorKind::PermissionDenied => { - clear_pairing_server_token(&self.current_enrollment, &mut enrollment)?; + clear_pairing_server_token(&mut current_enrollment, &mut enrollment)?; return Err(pairing_unavailable_error()); } _ => {} } } - if !*self.enabled_tx.borrow() { - return Err(Self::pairing_disabled_error()); - } let current_auth = load_remote_control_auth(&self.auth_manager) .await .map_err(|_| pairing_unavailable_error())?; @@ -223,6 +323,74 @@ impl RemoteControlHandle { pairing_response } + async fn load_or_enroll_pairing_server( + &self, + current_enrollment: &mut Option, + auth: &mut auth::RemoteControlConnectionAuth, + installation_id: &str, + server_name: &str, + app_server_client_name: Option<&str>, + ) -> io::Result { + if let Some(enrollment) = current_enrollment + .as_ref() + .filter(|enrollment| enrollment.account_id == auth.account_id) + .cloned() + { + return Ok(enrollment); + } + + let remote_control_target = normalize_remote_control_url(&self.remote_control_url)?; + let state_db = self + .state_db + .as_deref() + .ok_or_else(pairing_unavailable_error)?; + if let Some(mut enrollment) = load_persisted_remote_control_enrollment( + Some(state_db), + &remote_control_target, + &auth.account_id, + app_server_client_name, + ) + .await? + { + enrollment.server_name = server_name.to_string(); + publish_current_enrollment(current_enrollment, &enrollment); + return Ok(enrollment); + } + + let enrollment = enroll_pairing_server( + &self.auth_manager, + auth, + &remote_control_target, + installation_id, + server_name, + ) + .await?; + update_persisted_remote_control_enrollment( + Some(state_db), + &remote_control_target, + &auth.account_id, + app_server_client_name, + Some(&enrollment), + ) + .await?; + publish_current_enrollment(current_enrollment, &enrollment); + Ok(enrollment) + } + + fn pairing_persistence_key( + &self, + app_server_client_name: Option<&str>, + ) -> io::Result> { + if self.pairing_persistence_key_required && self.pairing_persistence_key.borrow().is_none() + { + let app_server_client_name = + app_server_client_name.ok_or_else(pairing_unavailable_error)?; + self.pairing_persistence_key + .send_replace(Some(app_server_client_name.to_string())); + } + Ok(self.pairing_persistence_key.borrow().clone()) + } + pub async fn list_clients( &self, params: RemoteControlClientsListParams, @@ -239,13 +407,6 @@ impl RemoteControlHandle { .await } - fn pairing_disabled_error() -> io::Error { - io::Error::new( - io::ErrorKind::InvalidInput, - "remote control pairing requires remote control to be enabled", - ) - } - fn publish_status( &self, connection_status: RemoteControlConnectionStatus, @@ -277,8 +438,36 @@ impl RemoteControlHandle { } } +async fn enroll_pairing_server( + auth_manager: &Arc, + auth: &mut auth::RemoteControlConnectionAuth, + remote_control_target: &protocol::RemoteControlTarget, + installation_id: &str, + server_name: &str, +) -> io::Result { + match enroll_remote_control_server(remote_control_target, auth, installation_id, server_name) + .await + { + Ok(enrollment) => return Ok(enrollment), + Err(err) if err.kind() == io::ErrorKind::PermissionDenied => { + let mut auth_recovery = auth_manager.unauthorized_recovery(); + let mut auth_change_rx = auth_manager.auth_change_receiver(); + if !recover_remote_control_auth(&mut auth_recovery, &mut auth_change_rx).await { + return Err(err); + } + *auth = load_remote_control_auth(auth_manager) + .await + .map_err(|_| pairing_unavailable_error())?; + } + Err(err) => return Err(err), + } + enroll_remote_control_server(remote_control_target, auth, installation_id, server_name).await +} + async fn refresh_pairing_enrollment( - current_enrollment: &CurrentRemoteControlEnrollment, + current_enrollment: &mut Option, + state_db: Option<&StateRuntime>, + app_server_client_name: Option<&str>, auth_manager: &Arc, auth: &mut auth::RemoteControlConnectionAuth, installation_id: &str, @@ -286,7 +475,14 @@ async fn refresh_pairing_enrollment( ) -> io::Result<()> { if let Err(err) = refresh_remote_control_server(auth, installation_id, enrollment).await { if err.kind() != io::ErrorKind::PermissionDenied { - return handle_pairing_refresh_error(current_enrollment, enrollment, err); + return handle_pairing_refresh_error( + current_enrollment, + state_db, + app_server_client_name, + enrollment, + err, + ) + .await; } let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut auth_change_rx = auth_manager.auth_change_receiver(); @@ -300,7 +496,14 @@ async fn refresh_pairing_enrollment( return Err(pairing_unavailable_error()); } if let Err(err) = refresh_remote_control_server(auth, installation_id, enrollment).await { - return handle_pairing_refresh_error(current_enrollment, enrollment, err); + return handle_pairing_refresh_error( + current_enrollment, + state_db, + app_server_client_name, + enrollment, + err, + ) + .await; } } if replace_current_enrollment(current_enrollment, enrollment) { @@ -310,21 +513,54 @@ async fn refresh_pairing_enrollment( } } -fn handle_pairing_refresh_error( - current_enrollment: &CurrentRemoteControlEnrollment, +async fn handle_pairing_refresh_error( + current_enrollment: &mut Option, + state_db: Option<&StateRuntime>, + app_server_client_name: Option<&str>, enrollment: &RemoteControlEnrollment, err: io::Error, ) -> io::Result<()> { if err.kind() == io::ErrorKind::NotFound { - clear_current_enrollment_if_matches(current_enrollment, enrollment); + clear_pairing_enrollment( + current_enrollment, + state_db, + app_server_client_name, + enrollment, + ) + .await; Err(pairing_unavailable_error()) } else { Err(err) } } +async fn clear_pairing_enrollment( + current_enrollment: &mut Option, + state_db: Option<&StateRuntime>, + app_server_client_name: Option<&str>, + enrollment: &RemoteControlEnrollment, +) { + if !clear_current_enrollment_if_matches(current_enrollment, enrollment) { + return; + } + let Some(state_db) = state_db else { + return; + }; + if let Err(err) = update_persisted_remote_control_enrollment( + Some(state_db), + &enrollment.remote_control_target, + &enrollment.account_id, + app_server_client_name, + /*enrollment*/ None, + ) + .await + { + warn!("failed to clear stale pairing enrollment: {err}"); + } +} + fn clear_pairing_server_token( - current_enrollment: &CurrentRemoteControlEnrollment, + current_enrollment: &mut Option, enrollment: &mut RemoteControlEnrollment, ) -> io::Result<()> { enrollment.clear_server_token(); @@ -359,27 +595,16 @@ fn remote_control_status_with_connection_status( } fn publish_current_enrollment( - current_enrollment: &CurrentRemoteControlEnrollment, + current_enrollment: &mut Option, enrollment: &RemoteControlEnrollment, ) { - *current_enrollment - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(enrollment.clone()); -} - -fn clear_current_enrollment(current_enrollment: &CurrentRemoteControlEnrollment) { - *current_enrollment - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) = None; + *current_enrollment = Some(enrollment.clone()); } fn replace_current_enrollment( - current_enrollment: &CurrentRemoteControlEnrollment, + current_enrollment: &mut Option, enrollment: &RemoteControlEnrollment, ) -> bool { - let mut current_enrollment = current_enrollment - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); if !current_enrollment .as_ref() .is_some_and(|current| same_remote_control_enrollment(current, enrollment)) @@ -391,17 +616,17 @@ fn replace_current_enrollment( } fn clear_current_enrollment_if_matches( - current_enrollment: &CurrentRemoteControlEnrollment, + current_enrollment: &mut Option, enrollment: &RemoteControlEnrollment, -) { - let mut current_enrollment = current_enrollment - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); +) -> bool { if current_enrollment .as_ref() .is_some_and(|current| same_remote_control_enrollment(current, enrollment)) { *current_enrollment = None; + true + } else { + false } } @@ -438,9 +663,13 @@ pub async fn start_remote_control( }; let (enabled_tx, enabled_rx) = watch::channel(initial_enabled); - let current_enrollment = Arc::new(StdMutex::new(None)); + let current_enrollment = Arc::new(RemoteControlEnrollmentState::new(/*enrollment*/ None)); let websocket_current_enrollment = current_enrollment.clone(); + let pairing_persistence_key_required = app_server_client_name_rx.is_some(); + let (pairing_persistence_key, _pairing_persistence_key_rx) = watch::channel(None); + let websocket_pairing_persistence_key = pairing_persistence_key.clone(); let handle_auth_manager = auth_manager.clone(); + let handle_state_db = state_db.clone(); let server_name = gethostname().to_string_lossy().trim().to_string(); let remote_control_url = config.remote_control_url; let installation_id = config.installation_id; @@ -490,6 +719,7 @@ pub async fn start_remote_control( transport_event_tx, status_publisher, current_enrollment: websocket_current_enrollment, + pairing_persistence_key: websocket_pairing_persistence_key, }, shutdown_token, enabled_rx, @@ -534,8 +764,11 @@ pub async fn start_remote_control( enabled_tx: Arc::new(enabled_tx), status_tx: Arc::new(status_tx), state_db_available, + state_db: handle_state_db, remote_control_url: handle_remote_control_url, current_enrollment, + pairing_persistence_key, + pairing_persistence_key_required, auth_manager: handle_auth_manager, }, )) 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 a1a284da9..3e9a911b2 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 @@ -40,7 +40,6 @@ use pretty_assertions::assert_eq; use serde_json::json; use std::collections::BTreeMap; use std::sync::Arc; -use std::sync::Mutex as StdMutex; use tempfile::TempDir; use time::OffsetDateTime; use tokio::io::AsyncBufReadExt; @@ -148,24 +147,29 @@ fn remote_control_handle_with_current_enrollment( }); let remote_control_target = normalize_remote_control_url(remote_control_url) .expect("remote control target should normalize"); - let current_enrollment = Arc::new(StdMutex::new(Some(RemoteControlEnrollment { - remote_control_target, - account_id: "account_id".to_string(), - environment_id: "env_test".to_string(), - server_id: "srv_e_test".to_string(), - server_name: test_server_name(), - remote_control_token: Some(TEST_REMOTE_CONTROL_SERVER_TOKEN.to_string()), - expires_at: Some( - OffsetDateTime::from_unix_timestamp(33_336_362_096) - .expect("future timestamp should parse"), - ), - }))); + let current_enrollment = Arc::new(RemoteControlEnrollmentState::new(Some( + RemoteControlEnrollment { + remote_control_target, + account_id: "account_id".to_string(), + environment_id: "env_test".to_string(), + server_id: "srv_e_test".to_string(), + server_name: test_server_name(), + remote_control_token: Some(TEST_REMOTE_CONTROL_SERVER_TOKEN.to_string()), + expires_at: Some( + OffsetDateTime::from_unix_timestamp(33_336_362_096) + .expect("future timestamp should parse"), + ), + }, + ))); RemoteControlHandle { enabled_tx: Arc::new(enabled_tx), status_tx: Arc::new(status_tx), state_db_available: true, + state_db: None, remote_control_url: remote_control_url.to_string(), current_enrollment, + pairing_persistence_key: watch::channel(None).0, + pairing_persistence_key_required: false, auth_manager, } } diff --git a/codex-rs/app-server-transport/src/transport/remote_control/tests/clients_tests.rs b/codex-rs/app-server-transport/src/transport/remote_control/tests/clients_tests.rs index 22aa29076..060063909 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/tests/clients_tests.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/tests/clients_tests.rs @@ -24,8 +24,11 @@ fn client_management_handle( enabled_tx: Arc::new(enabled_tx), status_tx: Arc::new(status_tx), state_db_available: false, + state_db: None, remote_control_url, - current_enrollment: Arc::new(StdMutex::new(None)), + current_enrollment: Arc::new(RemoteControlEnrollmentState::new(/*enrollment*/ None)), + pairing_persistence_key: watch::channel(None).0, + pairing_persistence_key_required: false, auth_manager, } } 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 176aeb4ed..7586f0b4d 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 @@ -131,13 +131,16 @@ async fn remote_control_handle_starts_pairing_before_websocket_connects() { remote_handle .current_enrollment .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) + .await .as_mut() .expect("current enrollment should exist") .expires_at = Some(OffsetDateTime::now_utc() + time::Duration::seconds(29)); let response = remote_handle - .start_pairing(RemoteControlPairingStartParams { manual_code: true }) + .start_pairing( + RemoteControlPairingStartParams { manual_code: true }, + /*app_server_client_name*/ None, + ) .await .expect("pairing should use the current server before websocket connect"); server_task.await.expect("server task should finish"); @@ -219,7 +222,10 @@ async fn remote_control_handle_refreshes_after_pairing_auth_failure() { ); let response = remote_handle - .start_pairing(RemoteControlPairingStartParams::default()) + .start_pairing( + RemoteControlPairingStartParams::default(), + /*app_server_client_name*/ None, + ) .await .expect("pairing should refresh after server token auth failure"); server_task.await.expect("server task should finish"); @@ -332,13 +338,16 @@ async fn remote_control_handle_recovers_auth_before_refreshing_pairing() { remote_handle .current_enrollment .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) + .await .as_mut() .expect("current enrollment should exist") .expires_at = Some(OffsetDateTime::now_utc() + time::Duration::seconds(29)); let response = remote_handle - .start_pairing(RemoteControlPairingStartParams::default()) + .start_pairing( + RemoteControlPairingStartParams::default(), + /*app_server_client_name*/ None, + ) .await .expect("pairing should refresh after auth recovery"); server_task.await.expect("server task should finish"); @@ -414,21 +423,137 @@ async fn start_remote_control_pairing_preserves_expiry_parse_error_context() { } #[tokio::test] -async fn remote_control_handle_disable_clears_current_enrollment() { +async fn remote_control_handle_disable_keeps_current_enrollment() { let remote_handle = remote_control_handle_with_current_enrollment( TEST_REMOTE_CONTROL_URL, remote_control_auth_manager(), ); remote_handle.disable(); - remote_handle.enable().expect("enable should succeed"); + assert!( + remote_handle.current_enrollment.lock().await.is_some(), + "disabled remote control should keep the selected pairing server" + ); +} + +#[tokio::test] +async fn remote_control_handle_reenrolls_after_stale_pairing_enrollment() { + 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 state_db = remote_control_state_runtime(&codex_home).await; + let mut remote_handle = remote_control_handle_with_current_enrollment( + &remote_control_url, + remote_control_auth_manager_with_home(&codex_home), + ); + remote_handle.state_db = Some(state_db.clone()); + remote_handle.disable(); + let stale_enrollment = remote_handle + .current_enrollment + .lock() + .await + .clone() + .expect("current enrollment should exist"); + let remote_control_target = stale_enrollment.remote_control_target.clone(); + let refreshed_enrollment = RemoteControlEnrollment { + remote_control_target: remote_control_target.clone(), + account_id: "account_id".to_string(), + environment_id: "env_refreshed".to_string(), + server_id: "srv_e_refreshed".to_string(), + server_name: test_server_name(), + remote_control_token: None, + expires_at: None, + }; + update_persisted_remote_control_enrollment( + Some(state_db.as_ref()), + &remote_control_target, + "account_id", + /*app_server_client_name*/ None, + Some(&stale_enrollment), + ) + .await + .expect("stale enrollment should save"); + let server_refreshed_enrollment = refreshed_enrollment.clone(); + let server_task = tokio::spawn(async move { + let stale_pairing_request = accept_http_request(&listener).await; + assert_eq!( + stale_pairing_request.request_line, + "POST /backend-api/wham/remote/control/server/pair HTTP/1.1" + ); + assert_eq!( + stale_pairing_request.headers.get("authorization"), + Some(&format!("Bearer {TEST_REMOTE_CONTROL_SERVER_TOKEN}")) + ); + respond_with_status(stale_pairing_request.stream, "404 Not Found", "").await; + + let enroll_request = accept_http_request(&listener).await; + assert_eq!( + enroll_request.request_line, + "POST /backend-api/wham/remote/control/server/enroll HTTP/1.1" + ); + respond_with_json( + enroll_request.stream, + remote_control_server_token_response( + &server_refreshed_enrollment.server_id, + &server_refreshed_enrollment.environment_id, + TEST_REFRESHED_REMOTE_CONTROL_SERVER_TOKEN, + ), + ) + .await; + + let refreshed_pairing_request = accept_http_request(&listener).await; + assert_eq!( + refreshed_pairing_request.request_line, + "POST /backend-api/wham/remote/control/server/pair HTTP/1.1" + ); + assert_eq!( + refreshed_pairing_request.headers.get("authorization"), + Some(&format!( + "Bearer {TEST_REFRESHED_REMOTE_CONTROL_SERVER_TOKEN}" + )) + ); + respond_with_json( + refreshed_pairing_request.stream, + json!({ + "pairing_code": "pairing-code", + "manual_pairing_code": "ABCD-EFGH", + "server_id": server_refreshed_enrollment.server_id, + "environment_id": server_refreshed_enrollment.environment_id, + "expires_at": "3026-05-22T12:34:56Z", + }), + ) + .await; + }); + let response = remote_handle + .start_pairing( + RemoteControlPairingStartParams::default(), + /*app_server_client_name*/ None, + ) + .await + .expect("pairing should re-enroll after stale enrollment"); + server_task.await.expect("server task should finish"); + assert_eq!( - remote_handle - .start_pairing(RemoteControlPairingStartParams::default()) - .await - .expect_err("re-enabled remote control should wait for a current server") - .to_string(), - "remote control pairing is unavailable until enrollment completes" + response, + RemoteControlPairingStartResponse { + pairing_code: "pairing-code".to_string(), + manual_pairing_code: Some("ABCD-EFGH".to_string()), + environment_id: "env_refreshed".to_string(), + expires_at: 33_336_362_096, + } + ); + assert_eq!( + load_persisted_remote_control_enrollment( + Some(state_db.as_ref()), + &remote_control_target, + "account_id", + /*app_server_client_name*/ None, + ) + .await + .expect("refreshed enrollment should load"), + Some(refreshed_enrollment) ); } @@ -458,7 +583,10 @@ async fn remote_control_handle_discards_pairing_response_after_auth_change() { let remote_handle = remote_handle.clone(); async move { remote_handle - .start_pairing(RemoteControlPairingStartParams::default()) + .start_pairing( + RemoteControlPairingStartParams::default(), + /*app_server_client_name*/ None, + ) .await } }); 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 60792df54..43fce8ed3 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 @@ -1,13 +1,13 @@ use super::CurrentRemoteControlEnrollment; -use super::clear_current_enrollment; +use super::RemoteControlPairingPersistenceKey; use super::protocol::ClientEnvelope; use super::protocol::ClientEvent; use super::protocol::ClientId; use super::protocol::RemoteControlTarget; use super::protocol::ServerEnvelope; use super::protocol::StreamId; -use super::publish_current_enrollment; use super::remote_control_status_with_connection_status; +use super::same_remote_control_enrollment; use super::segment::ClientSegmentObservation; use super::segment::ClientSegmentReassembler; use super::segment::REMOTE_CONTROL_SEGMENT_MAX_BYTES; @@ -55,6 +55,9 @@ use tokio_tungstenite::connect_async; use tokio_tungstenite::tungstenite; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tokio_util::sync::CancellationToken; + +#[cfg(test)] +use super::RemoteControlEnrollmentState; use tracing::error; use tracing::info; use tracing::warn; @@ -252,10 +255,10 @@ pub(crate) struct RemoteControlWebsocket { status_publisher: RemoteControlStatusPublisher, shutdown_token: CancellationToken, reconnect_attempt: u64, - enrollment: Option, auth_recovery: UnauthorizedRecovery, auth_change_rx: watch::Receiver, current_enrollment: CurrentRemoteControlEnrollment, + pairing_persistence_key: RemoteControlPairingPersistenceKey, client_tracker: Arc>, state: Arc>, server_event_rx: Arc>>, @@ -294,6 +297,7 @@ pub(super) struct RemoteControlChannels { pub(super) transport_event_tx: mpsc::Sender, pub(super) status_publisher: RemoteControlStatusPublisher, pub(super) current_enrollment: CurrentRemoteControlEnrollment, + pub(super) pairing_persistence_key: RemoteControlPairingPersistenceKey, } #[derive(Clone)] @@ -407,10 +411,10 @@ impl RemoteControlWebsocket { status_publisher: channels.status_publisher, shutdown_token, reconnect_attempt: 0, - enrollment: None, auth_recovery, auth_change_rx, current_enrollment: channels.current_enrollment, + pairing_persistence_key: channels.pairing_persistence_key, client_tracker: Arc::new(Mutex::new(client_tracker)), state: Arc::new(Mutex::new(WebsocketState { outbound_buffer, @@ -457,6 +461,8 @@ impl RemoteControlWebsocket { return; } }; + self.pairing_persistence_key + .send_replace(app_server_client_name.clone()); loop { if !self.wait_until_enabled().await { @@ -579,15 +585,15 @@ impl RemoteControlWebsocket { loop { let subscribe_cursor = self.state.lock().await.subscribe_cursor.clone(); - let enrollment = self.enrollment.as_ref(); + let enrollment = self.current_enrollment.snapshot(); info!( websocket_url = %remote_control_target.websocket_url, installation_id = %self.installation_id, server_name = %self.server_name, reconnect_attempt = self.reconnect_attempt.saturating_add(1), has_enrollment = enrollment.is_some(), - server_id = ?enrollment.map(|enrollment| enrollment.server_id.as_str()), - environment_id = ?enrollment.map(|enrollment| enrollment.environment_id.as_str()), + server_id = ?enrollment.as_ref().map(|enrollment| enrollment.server_id.as_str()), + environment_id = ?enrollment.as_ref().map(|enrollment| enrollment.environment_id.as_str()), subscribe_cursor_present = subscribe_cursor.is_some(), app_server_client_name = ?app_server_client_name, "connecting to app-server remote control websocket" @@ -611,15 +617,17 @@ impl RemoteControlWebsocket { } return ConnectOutcome::Disabled; } - connect_result = connect_remote_control_websocket( - &remote_control_target, - self.state_db.as_deref(), - auth_context, - &mut self.enrollment, - connect_options, - &self.status_publisher, - &self.current_enrollment, - ) => connect_result, + connect_result = async { + connect_remote_control_websocket( + &remote_control_target, + self.state_db.as_deref(), + auth_context, + &self.current_enrollment, + connect_options, + &self.status_publisher, + ) + .await + } => connect_result, }; match connect_result { @@ -631,13 +639,13 @@ impl RemoteControlWebsocket { self.auth_recovery = self.auth_manager.unauthorized_recovery(); self.status_publisher .publish_status(RemoteControlConnectionStatus::Connected); - let enrollment = self.enrollment.as_ref(); + let enrollment = self.current_enrollment.snapshot(); info!( websocket_url = %remote_control_target.websocket_url, installation_id = %self.installation_id, server_name = %self.server_name, - server_id = ?enrollment.map(|enrollment| enrollment.server_id.as_str()), - environment_id = ?enrollment.map(|enrollment| enrollment.environment_id.as_str()), + server_id = ?enrollment.as_ref().map(|enrollment| enrollment.server_id.as_str()), + environment_id = ?enrollment.as_ref().map(|enrollment| enrollment.environment_id.as_str()), subscribe_cursor_present = subscribe_cursor.is_some(), response_headers = %format_headers(response.headers()), "connected to app-server remote control websocket" @@ -656,7 +664,7 @@ impl RemoteControlWebsocket { let reconnect_attempt = self.reconnect_attempt.saturating_add(1); let (reconnect_delay, reconnect_backoff_reset) = next_reconnect_delay(&mut self.reconnect_attempt); - let enrollment = self.enrollment.as_ref(); + let enrollment = self.current_enrollment.snapshot(); warn!( websocket_url = %remote_control_target.websocket_url, installation_id = %self.installation_id, @@ -667,8 +675,8 @@ impl RemoteControlWebsocket { reconnect_delay = ?reconnect_delay, reconnect_backoff_reset, has_enrollment = enrollment.is_some(), - server_id = ?enrollment.map(|enrollment| enrollment.server_id.as_str()), - environment_id = ?enrollment.map(|enrollment| enrollment.environment_id.as_str()), + server_id = ?enrollment.as_ref().map(|enrollment| enrollment.server_id.as_str()), + environment_id = ?enrollment.as_ref().map(|enrollment| enrollment.environment_id.as_str()), subscribe_cursor_present = subscribe_cursor.is_some(), "failed to connect to app-server remote control websocket" ); @@ -1189,19 +1197,108 @@ pub(super) async fn connect_remote_control_websocket( remote_control_target: &RemoteControlTarget, state_db: Option<&StateRuntime>, mut auth_context: RemoteControlAuthContext<'_>, - enrollment: &mut Option, + current_enrollment: &CurrentRemoteControlEnrollment, connect_options: RemoteControlConnectOptions<'_>, status_publisher: &RemoteControlStatusPublisher, - current_enrollment: &CurrentRemoteControlEnrollment, ) -> io::Result<( WebSocketStream>, tungstenite::http::Response<()>, )> { ensure_rustls_crypto_provider(); + let (auth, enrollment) = { + let mut current_enrollment = current_enrollment.lock().await; + let auth = prepare_remote_control_enrollment( + remote_control_target, + state_db, + &mut auth_context, + &mut current_enrollment, + connect_options, + status_publisher, + ) + .await?; + let enrollment = current_enrollment.as_ref().cloned().ok_or_else(|| { + io::Error::other("missing remote control enrollment after enrollment step") + })?; + (auth, enrollment) + }; + let request = build_remote_control_websocket_request( + &remote_control_target.websocket_url, + &enrollment, + connect_options.installation_id, + connect_options.subscribe_cursor, + )?; + + let websocket_connect_result = tokio::time::timeout( + REMOTE_CONTROL_WEBSOCKET_CONNECT_TIMEOUT, + connect_async(request), + ) + .await + .map_err(|_| { + io::Error::new( + ErrorKind::TimedOut, + format!( + "timed out connecting to remote control websocket at `{}` after {:?}", + remote_control_target.websocket_url, REMOTE_CONTROL_WEBSOCKET_CONNECT_TIMEOUT + ), + ) + })?; + + match websocket_connect_result { + Ok((websocket_stream, response)) => Ok((websocket_stream, response.map(|_| ()))), + Err(err) => { + match &err { + tungstenite::Error::Http(response) if response.status().as_u16() == 404 => { + info!( + "remote control websocket returned HTTP 404; clearing stale enrollment before re-enrolling: websocket_url={}, account_id={}, server_id={}, environment_id={}", + remote_control_target.websocket_url, + auth.account_id, + enrollment.server_id, + enrollment.environment_id + ); + clear_remote_control_enrollment_if_matches( + state_db, + remote_control_target, + &auth.account_id, + connect_options.app_server_client_name, + current_enrollment, + &enrollment, + status_publisher, + ) + .await; + } + tungstenite::Error::Http(response) + if matches!(response.status().as_u16(), 401 | 403) => + { + clear_remote_control_server_token_if_matches(current_enrollment, &enrollment) + .await?; + return Err(io::Error::other(format!( + "remote control websocket auth failed with HTTP {}; refreshing server token before reconnect", + response.status() + ))); + } + _ => {} + } + Err(io::Error::other( + format_remote_control_websocket_connect_error( + &remote_control_target.websocket_url, + &err, + ), + )) + } + } +} + +async fn prepare_remote_control_enrollment( + remote_control_target: &RemoteControlTarget, + state_db: Option<&StateRuntime>, + auth_context: &mut RemoteControlAuthContext<'_>, + enrollment: &mut Option, + connect_options: RemoteControlConnectOptions<'_>, + status_publisher: &RemoteControlStatusPublisher, +) -> io::Result { let Some(state_db) = state_db else { *enrollment = None; - clear_current_enrollment(current_enrollment); return Err(io::Error::new( ErrorKind::NotFound, "remote control requires sqlite state db", @@ -1214,7 +1311,6 @@ pub(super) async fn connect_remote_control_websocket( if err.kind() == ErrorKind::PermissionDenied { *enrollment = None; status_publisher.publish_environment_id(/*environment_id*/ None); - clear_current_enrollment(current_enrollment); } return Err(err); } @@ -1231,7 +1327,6 @@ pub(super) async fn connect_remote_control_websocket( ); *enrollment = None; status_publisher.publish_environment_id(/*environment_id*/ None); - clear_current_enrollment(current_enrollment); } if let Some(enrollment) = enrollment.as_mut() { enrollment.remote_control_target = remote_control_target.clone(); @@ -1262,7 +1357,7 @@ pub(super) async fn connect_remote_control_websocket( remote_control_target, state_db, &auth, - &mut auth_context, + auth_context, enrollment, connect_options, status_publisher, @@ -1307,14 +1402,13 @@ pub(super) async fn connect_remote_control_websocket( connect_options.app_server_client_name, enrollment, status_publisher, - current_enrollment, ) .await; enroll_remote_control_server_if_missing( remote_control_target, state_db, &auth, - &mut auth_context, + auth_context, enrollment, connect_options, status_publisher, @@ -1337,82 +1431,7 @@ pub(super) async fn connect_remote_control_websocket( } } - let enrollment_ref = enrollment.as_ref().ok_or_else(|| { - io::Error::other("missing remote control enrollment after enrollment step") - })?; - publish_current_enrollment(current_enrollment, enrollment_ref); - let request = build_remote_control_websocket_request( - &remote_control_target.websocket_url, - enrollment_ref, - connect_options.installation_id, - connect_options.subscribe_cursor, - )?; - - let websocket_connect_result = tokio::time::timeout( - REMOTE_CONTROL_WEBSOCKET_CONNECT_TIMEOUT, - connect_async(request), - ) - .await - .map_err(|_| { - io::Error::new( - ErrorKind::TimedOut, - format!( - "timed out connecting to remote control websocket at `{}` after {:?}", - remote_control_target.websocket_url, REMOTE_CONTROL_WEBSOCKET_CONNECT_TIMEOUT - ), - ) - })?; - - match websocket_connect_result { - Ok((websocket_stream, response)) => Ok((websocket_stream, response.map(|_| ()))), - Err(err) => { - match &err { - tungstenite::Error::Http(response) if response.status().as_u16() == 404 => { - info!( - "remote control websocket returned HTTP 404; clearing stale enrollment before re-enrolling: websocket_url={}, account_id={}, server_id={}, environment_id={}", - remote_control_target.websocket_url, - auth.account_id, - enrollment_ref.server_id, - enrollment_ref.environment_id - ); - clear_remote_control_enrollment( - state_db, - remote_control_target, - &auth.account_id, - connect_options.app_server_client_name, - enrollment, - status_publisher, - current_enrollment, - ) - .await; - } - tungstenite::Error::Http(response) - if matches!(response.status().as_u16(), 401 | 403) => - { - enrollment - .as_mut() - .ok_or_else(|| { - io::Error::other( - "missing remote control enrollment after websocket auth failure", - ) - })? - .clear_server_token(); - clear_current_enrollment(current_enrollment); - return Err(io::Error::other(format!( - "remote control websocket auth failed with HTTP {}; refreshing server token before reconnect", - response.status() - ))); - } - _ => {} - } - Err(io::Error::other( - format_remote_control_websocket_connect_error( - &remote_control_target.websocket_url, - &err, - ), - )) - } - } + Ok(auth) } async fn clear_remote_control_enrollment( @@ -1422,7 +1441,6 @@ async fn clear_remote_control_enrollment( app_server_client_name: Option<&str>, enrollment: &mut Option, status_publisher: &RemoteControlStatusPublisher, - current_enrollment: &CurrentRemoteControlEnrollment, ) { if let Err(clear_err) = update_persisted_remote_control_enrollment( Some(state_db), @@ -1437,7 +1455,51 @@ async fn clear_remote_control_enrollment( } *enrollment = None; status_publisher.publish_environment_id(/*environment_id*/ None); - clear_current_enrollment(current_enrollment); +} + +async fn clear_remote_control_enrollment_if_matches( + state_db: Option<&StateRuntime>, + remote_control_target: &RemoteControlTarget, + account_id: &str, + app_server_client_name: Option<&str>, + current_enrollment: &CurrentRemoteControlEnrollment, + enrollment: &RemoteControlEnrollment, + status_publisher: &RemoteControlStatusPublisher, +) { + let Some(state_db) = state_db else { + return; + }; + let mut current_enrollment = current_enrollment.lock().await; + if !current_enrollment + .as_ref() + .is_some_and(|current| same_remote_control_enrollment(current, enrollment)) + { + return; + } + clear_remote_control_enrollment( + state_db, + remote_control_target, + account_id, + app_server_client_name, + &mut current_enrollment, + status_publisher, + ) + .await; +} + +async fn clear_remote_control_server_token_if_matches( + current_enrollment: &CurrentRemoteControlEnrollment, + enrollment: &RemoteControlEnrollment, +) -> io::Result<()> { + let mut current_enrollment = current_enrollment.lock().await; + current_enrollment + .as_mut() + .filter(|current| same_remote_control_enrollment(current, enrollment)) + .ok_or_else(|| { + io::Error::other("missing remote control enrollment after websocket auth failure") + })? + .clear_server_token(); + Ok(()) } async fn enroll_remote_control_server_if_missing( @@ -1585,8 +1647,10 @@ mod tests { } } - fn test_current_enrollment() -> CurrentRemoteControlEnrollment { - Arc::new(std::sync::Mutex::new(None)) + fn test_current_enrollment( + enrollment: Option, + ) -> CurrentRemoteControlEnrollment { + Arc::new(RemoteControlEnrollmentState::new(enrollment)) } #[test] @@ -1739,10 +1803,9 @@ mod tests { let auth_manager = remote_control_auth_manager(); let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut auth_change_rx = auth_manager.auth_change_receiver(); - let mut enrollment = Some(remote_control_enrollment(Some( + let current_enrollment = test_current_enrollment(Some(remote_control_enrollment(Some( TEST_REMOTE_CONTROL_SERVER_TOKEN, - ))); - let current_enrollment = test_current_enrollment(); + )))); let (status_publisher, status_rx) = remote_control_status_channel(); let err = match connect_remote_control_websocket( @@ -1753,7 +1816,7 @@ mod tests { auth_recovery: &mut auth_recovery, auth_change_rx: &mut auth_change_rx, }, - &mut enrollment, + ¤t_enrollment, RemoteControlConnectOptions { installation_id: TEST_INSTALLATION_ID, server_name: "test-server", @@ -1761,7 +1824,6 @@ mod tests { app_server_client_name: None, }, &status_publisher, - ¤t_enrollment, ) .await { @@ -1771,12 +1833,7 @@ mod tests { server_task.await.expect("server task should succeed"); assert_eq!(err.to_string(), expected_error); - assert!( - current_enrollment - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .is_some() - ); + assert!(current_enrollment.lock().await.is_some()); assert_eq!( status_rx.borrow().clone(), RemoteControlStatusChangedNotification { @@ -1801,10 +1858,9 @@ mod tests { let auth_manager = remote_control_auth_manager(); let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut auth_change_rx = auth_manager.auth_change_receiver(); - let mut enrollment = Some(remote_control_enrollment(Some( + let current_enrollment = test_current_enrollment(Some(remote_control_enrollment(Some( TEST_REMOTE_CONTROL_SERVER_TOKEN, - ))); - let current_enrollment = test_current_enrollment(); + )))); let (status_publisher, status_rx) = remote_control_status_channel(); let server_task = tokio::spawn(async move { @@ -1824,7 +1880,7 @@ mod tests { auth_recovery: &mut auth_recovery, auth_change_rx: &mut auth_change_rx, }, - &mut enrollment, + ¤t_enrollment, RemoteControlConnectOptions { installation_id: TEST_INSTALLATION_ID, server_name: "test-server", @@ -1832,7 +1888,6 @@ mod tests { app_server_client_name: None, }, &status_publisher, - ¤t_enrollment, ) .await .expect_err("unauthorized response should fail the websocket connect"); @@ -1853,13 +1908,7 @@ mod tests { ); let mut expected_enrollment = remote_control_enrollment(/*remote_control_token*/ None); expected_enrollment.remote_control_target = remote_control_target; - assert_eq!(enrollment, Some(expected_enrollment)); - assert!( - current_enrollment - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .is_none() - ); + assert_eq!(*current_enrollment.lock().await, Some(expected_enrollment)); } #[tokio::test] @@ -1896,7 +1945,7 @@ mod tests { .await; let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut auth_change_rx = auth_manager.auth_change_receiver(); - let mut enrollment = None; + let current_enrollment = test_current_enrollment(/*enrollment*/ None); let (status_publisher, status_rx) = remote_control_status_channel(); save_auth( codex_home.path(), @@ -1913,7 +1962,7 @@ mod tests { auth_recovery: &mut auth_recovery, auth_change_rx: &mut auth_change_rx, }, - &mut enrollment, + ¤t_enrollment, RemoteControlConnectOptions { installation_id: TEST_INSTALLATION_ID, server_name: "test-server", @@ -1921,7 +1970,6 @@ mod tests { app_server_client_name: None, }, &status_publisher, - &test_current_enrollment(), ) .await .expect_err("unauthorized enrollment should fail the websocket connect"); @@ -1989,9 +2037,9 @@ mod tests { .await; let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut auth_change_rx = auth_manager.auth_change_receiver(); - let mut enrollment = Some(remote_control_enrollment( + let current_enrollment = test_current_enrollment(Some(remote_control_enrollment( /*remote_control_token*/ None, - )); + ))); let (status_publisher, status_rx) = remote_control_status_channel(); save_auth( codex_home.path(), @@ -2008,7 +2056,7 @@ mod tests { auth_recovery: &mut auth_recovery, auth_change_rx: &mut auth_change_rx, }, - &mut enrollment, + ¤t_enrollment, RemoteControlConnectOptions { installation_id: TEST_INSTALLATION_ID, server_name: "test-server", @@ -2016,7 +2064,6 @@ mod tests { app_server_client_name: None, }, &status_publisher, - &test_current_enrollment(), ) .await .expect_err("unauthorized refresh should fail the websocket connect"); @@ -2061,9 +2108,9 @@ mod tests { let auth_manager = remote_control_auth_manager(); let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut auth_change_rx = auth_manager.auth_change_receiver(); - let mut enrollment = Some(remote_control_enrollment(Some( + let current_enrollment = test_current_enrollment(Some(remote_control_enrollment(Some( TEST_REMOTE_CONTROL_SERVER_TOKEN, - ))); + )))); let (status_publisher, _status_rx) = remote_control_status_channel(); let err = connect_remote_control_websocket( @@ -2074,7 +2121,7 @@ mod tests { auth_recovery: &mut auth_recovery, auth_change_rx: &mut auth_change_rx, }, - &mut enrollment, + ¤t_enrollment, RemoteControlConnectOptions { installation_id: TEST_INSTALLATION_ID, server_name: "test-server", @@ -2082,14 +2129,13 @@ mod tests { app_server_client_name: None, }, &status_publisher, - &test_current_enrollment(), ) .await .expect_err("missing sqlite state db should fail remote control"); assert_eq!(err.kind(), ErrorKind::NotFound); assert_eq!(err.to_string(), "remote control requires sqlite state db"); - assert_eq!(enrollment, None); + assert_eq!(*current_enrollment.lock().await, None); } #[tokio::test] @@ -2107,9 +2153,9 @@ mod tests { .await; let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut auth_change_rx = auth_manager.auth_change_receiver(); - let mut enrollment = Some(remote_control_enrollment(Some( + let current_enrollment = test_current_enrollment(Some(remote_control_enrollment(Some( TEST_REMOTE_CONTROL_SERVER_TOKEN, - ))); + )))); let (status_publisher, mut status_rx) = remote_control_status_channel(); status_publisher.publish_environment_id(Some("env_test".to_string())); status_rx @@ -2125,7 +2171,7 @@ mod tests { auth_recovery: &mut auth_recovery, auth_change_rx: &mut auth_change_rx, }, - &mut enrollment, + ¤t_enrollment, RemoteControlConnectOptions { installation_id: TEST_INSTALLATION_ID, server_name: "test-server", @@ -2133,7 +2179,6 @@ mod tests { app_server_client_name: None, }, &status_publisher, - &test_current_enrollment(), ) .await .expect_err("missing auth should fail remote control"); @@ -2143,7 +2188,7 @@ mod tests { err.to_string(), "remote control requires ChatGPT authentication" ); - assert_eq!(enrollment, None); + assert_eq!(*current_enrollment.lock().await, None); assert_eq!( status_rx.borrow().clone(), RemoteControlStatusChangedNotification { @@ -2185,7 +2230,8 @@ mod tests { RemoteControlChannels { transport_event_tx, status_publisher, - current_enrollment: test_current_enrollment(), + current_enrollment: test_current_enrollment(/*enrollment*/ None), + pairing_persistence_key: watch::channel(None).0, }, shutdown_token, enabled_rx, diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index d89f0b064..6c578ce8e 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -918,7 +918,7 @@ impl MessageProcessor { .map(|response| Some(response.into())), ClientRequest::RemoteControlPairingStart { params, .. } => self .remote_control_processor - .pairing_start(params) + .pairing_start(params, app_server_client_name.as_deref()) .await .map(|response| Some(response.into())), ClientRequest::RemoteControlClientsList { params, .. } => self diff --git a/codex-rs/app-server/src/request_processors/remote_control_processor.rs b/codex-rs/app-server/src/request_processors/remote_control_processor.rs index 804339c17..22223687b 100644 --- a/codex-rs/app-server/src/request_processors/remote_control_processor.rs +++ b/codex-rs/app-server/src/request_processors/remote_control_processor.rs @@ -52,9 +52,10 @@ impl RemoteControlRequestProcessor { pub(crate) async fn pairing_start( &self, params: RemoteControlPairingStartParams, + app_server_client_name: Option<&str>, ) -> Result { self.handle()? - .start_pairing(params) + .start_pairing(params, app_server_client_name) .await .map_err(map_pairing_start_error) } diff --git a/codex-rs/app-server/src/request_processors/remote_control_processor/remote_control_processor_tests.rs b/codex-rs/app-server/src/request_processors/remote_control_processor/remote_control_processor_tests.rs index 4381e4895..e336491cd 100644 --- a/codex-rs/app-server/src/request_processors/remote_control_processor/remote_control_processor_tests.rs +++ b/codex-rs/app-server/src/request_processors/remote_control_processor/remote_control_processor_tests.rs @@ -6,7 +6,10 @@ use pretty_assertions::assert_eq; #[tokio::test] async fn pairing_start_returns_internal_error_when_remote_control_is_unavailable() { let err = RemoteControlRequestProcessor::new(/*remote_control_handle*/ None) - .pairing_start(RemoteControlPairingStartParams::default()) + .pairing_start( + RemoteControlPairingStartParams::default(), + /*app_server_client_name*/ None, + ) .await .expect_err("missing remote control should fail pairing"); diff --git a/codex-rs/app-server/tests/suite/v2/remote_control.rs b/codex-rs/app-server/tests/suite/v2/remote_control.rs index b5384d3ba..c48ffda0f 100644 --- a/codex-rs/app-server/tests/suite/v2/remote_control.rs +++ b/codex-rs/app-server/tests/suite/v2/remote_control.rs @@ -193,6 +193,42 @@ async fn remote_control_pairing_start_returns_pairing_artifacts() -> Result<()> Ok(()) } +#[tokio::test] +async fn remote_control_pairing_start_returns_pairing_artifacts_while_disabled() -> Result<()> { + let codex_home = TempDir::new()?; + let mut backend = PairingRemoteControlBackend::start(codex_home.path()).await?; + let mut mcp = TestAppServer::new(codex_home.path()).await?; + timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; + + let request_id = mcp + .send_remote_control_pairing_start_request(RemoteControlPairingStartParams { + manual_code: true, + }) + .await?; + let response: JSONRPCResponse = timeout( + DEFAULT_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + assert_eq!( + timeout(DEFAULT_TIMEOUT, backend.wait_for_enroll_request()).await??, + "POST /backend-api/wham/remote/control/server/enroll HTTP/1.1" + ); + assert_eq!(response.result.get("serverId"), None); + let received: RemoteControlPairingStartResponse = to_response(response)?; + + assert_eq!( + received, + RemoteControlPairingStartResponse { + pairing_code: "pairing-code".to_string(), + manual_pairing_code: Some("ABCD-EFGH".to_string()), + environment_id: "environment-id".to_string(), + expires_at: 33_336_362_096, + } + ); + Ok(()) +} + #[tokio::test] async fn remote_control_client_management_works_while_disabled() -> Result<()> { let codex_home = TempDir::new()?; @@ -376,8 +412,12 @@ impl PairingRemoteControlBackend { ) .await?; - let _websocket_request = read_http_request(&listener).await?; - let pair_http_request = read_http_request(&listener).await?; + let request_after_enroll = read_http_request(&listener).await?; + let pair_http_request = if request_after_enroll.request_line.starts_with("GET ") { + read_http_request(&listener).await? + } else { + request_after_enroll + }; respond_with_json( pair_http_request.reader.into_inner(), serde_json::json!({