fix(app-server-transport): keep pairing auth current

codex-rs/app-server-transport/src/transport/remote_control/mod.rs

@@ -58,6 +58,7 @@ pub struct RemoteControlHandle {
     status_tx: Arc<watch::Sender<RemoteControlStatusChangedNotification>>,
     state_db_available: bool,
     pairing_client: Arc<StdMutex<Option<RemoteControlPairingClient>>>,
+    auth_change_rx: Arc<StdMutex<watch::Receiver<u64>>>,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -147,22 +148,48 @@ impl RemoteControlHandle {
             ));
         }
 
-        let pairing_client = self
-            .pairing_client
-            .lock()
-            .unwrap_or_else(std::sync::PoisonError::into_inner)
-            .clone()
-            .ok_or_else(|| {
-                io::Error::new(
-                    io::ErrorKind::InvalidInput,
-                    "remote control pairing is unavailable until enrollment completes",
-                )
-            })?;
-        pairing_client
+        let auth_change_revision = self.auth_change_revision();
+        let pairing_client = {
+            let mut pairing_client = self
+                .pairing_client
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner);
+            let current_pairing_client = pairing_client
+                .as_ref()
+                .filter(|pairing_client| {
+                    pairing_client.matches_auth_change_revision(auth_change_revision)
+                })
+                .cloned();
+            if current_pairing_client.is_none() {
+                *pairing_client = None;
+            }
+            current_pairing_client
+        }
+        .ok_or_else(Self::pairing_unavailable_error)?;
+        let pairing_response = pairing_client
             .start(protocol::StartRemoteControlPairingRequest {
                 manual_code: params.manual_code,
             })
-            .await
+            .await;
+        if self.auth_change_revision() != auth_change_revision {
+            return Err(Self::pairing_unavailable_error());
+        }
+        pairing_response
+    }
+
+    fn auth_change_revision(&self) -> u64 {
+        *self
+            .auth_change_rx
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .borrow()
+    }
+
+    fn pairing_unavailable_error() -> io::Error {
+        io::Error::new(
+            io::ErrorKind::InvalidInput,
+            "remote control pairing is unavailable until enrollment completes",
+        )
     }
 
     fn publish_status(
@@ -242,6 +269,7 @@ pub async fn start_remote_control(
     let (enabled_tx, enabled_rx) = watch::channel(initial_enabled);
     let pairing_client = Arc::new(StdMutex::new(None));
     let websocket_pairing_client = Arc::clone(&pairing_client);
+    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;
     let installation_id = config.installation_id;
@@ -335,6 +363,7 @@ pub async fn start_remote_control(
             status_tx: Arc::new(status_tx),
             state_db_available,
             pairing_client,
+            auth_change_rx,
         },
     ))
 }

codex-rs/app-server-transport/src/transport/remote_control/pairing.rs

@@ -19,6 +19,7 @@ pub(super) struct RemoteControlPairingClient {
     server_id: String,
     environment_id: String,
     expires_at: OffsetDateTime,
+    auth_change_revision: u64,
 }
 
 impl RemoteControlPairingClient {
@@ -28,6 +29,7 @@ impl RemoteControlPairingClient {
         server_id: String,
         environment_id: String,
         expires_at: OffsetDateTime,
+        auth_change_revision: u64,
     ) -> Self {
         Self {
             pairing_url: remote_control_target.pair_url.clone(),
@@ -35,9 +37,14 @@ impl RemoteControlPairingClient {
             server_id,
             environment_id,
             expires_at,
+            auth_change_revision,
         }
     }
 
+    pub(super) fn matches_auth_change_revision(&self, auth_change_revision: u64) -> bool {
+        self.auth_change_revision == auth_change_revision
+    }
+
     pub(super) async fn start(
         &self,
         request: StartRemoteControlPairingRequest,

codex-rs/app-server-transport/src/transport/remote_control/pairing_tests.rs

@@ -105,6 +105,7 @@ async fn start_remote_control_pairing_uses_server_token_and_maps_response() {
         "server-id".to_string(),
         "environment-id".to_string(),
         OffsetDateTime::from_unix_timestamp(33_336_362_096).expect("future timestamp should parse"),
+        /*auth_change_revision*/ 0,
     );
 
     let response = client
@@ -173,6 +174,7 @@ async fn start_remote_control_pairing_preserves_backend_error_context() {
         "server-id".to_string(),
         "environment-id".to_string(),
         OffsetDateTime::from_unix_timestamp(33_336_362_096).expect("future timestamp should parse"),
+        /*auth_change_revision*/ 0,
     );
 
     let err = client

codex-rs/app-server-transport/src/transport/remote_control/tests.rs

@@ -61,6 +61,7 @@ use tokio_tungstenite::tungstenite;
 use tokio_util::sync::CancellationToken;
 
 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";
 const TEST_REMOTE_CONTROL_SERVER_TOKEN: &str = "Remote Control Token";
 const TEST_REFRESHED_REMOTE_CONTROL_SERVER_TOKEN: &str = "Refreshed Remote Control Token";
 const TEST_REMOTE_CONTROL_SERVER_TOKEN_EXPIRES_AT: &str = "2999-01-01T00:00:00Z";
@@ -131,6 +132,35 @@ fn test_server_name() -> String {
     gethostname().to_string_lossy().trim().to_string()
 }
 
+fn remote_control_handle_with_pairing_client(
+    remote_control_url: &str,
+    auth_change_rx: watch::Receiver<u64>,
+) -> RemoteControlHandle {
+    let (enabled_tx, _enabled_rx) = watch::channel(/*init*/ true);
+    let (status_tx, _status_rx) = watch::channel(RemoteControlStatusChangedNotification {
+        status: RemoteControlConnectionStatus::Connected,
+        server_name: test_server_name(),
+        installation_id: TEST_INSTALLATION_ID.to_string(),
+        environment_id: Some("env_test".to_string()),
+    });
+    let pairing_client = Arc::new(StdMutex::new(Some(RemoteControlPairingClient::new(
+        &normalize_remote_control_url(remote_control_url)
+            .expect("remote control target should normalize"),
+        TEST_REMOTE_CONTROL_SERVER_TOKEN.to_string(),
+        "srv_e_test".to_string(),
+        "env_test".to_string(),
+        OffsetDateTime::from_unix_timestamp(33_336_362_096).expect("future timestamp should parse"),
+        /*auth_change_revision*/ 0,
+    ))));
+    RemoteControlHandle {
+        enabled_tx: Arc::new(enabled_tx),
+        status_tx: Arc::new(status_tx),
+        state_db_available: true,
+        pairing_client,
+        auth_change_rx: Arc::new(StdMutex::new(auth_change_rx)),
+    }
+}
+
 fn remote_control_server_token_response(
     server_id: &str,
     environment_id: &str,
@@ -900,25 +930,10 @@ async fn remote_control_handle_enable_disable_stops_and_restarts_connections() {
 
 #[tokio::test]
 async fn remote_control_handle_disable_clears_stale_pairing_client() {
-    let (enabled_tx, _enabled_rx) = watch::channel(/*init*/ true);
-    let (status_tx, _status_rx) = watch::channel(RemoteControlStatusChangedNotification {
-        status: RemoteControlConnectionStatus::Connected,
-        server_name: test_server_name(),
-        installation_id: TEST_INSTALLATION_ID.to_string(),
-        environment_id: Some("env_test".to_string()),
-    });
-    let pairing_client = Arc::new(StdMutex::new(Some(RemoteControlPairingClient::new(
-        &normalize_remote_control_url("http://127.0.0.1:1/backend-api/wham/remote/control")
-            .expect("remote control target should normalize"),
-        TEST_REMOTE_CONTROL_SERVER_TOKEN.to_string(),
-        OffsetDateTime::from_unix_timestamp(33_336_362_096).expect("future timestamp should parse"),
-    ))));
-    let remote_handle = RemoteControlHandle {
-        enabled_tx: Arc::new(enabled_tx),
-        status_tx: Arc::new(status_tx),
-        state_db_available: true,
-        pairing_client,
-    };
+    let remote_handle = remote_control_handle_with_pairing_client(
+        TEST_REMOTE_CONTROL_URL,
+        watch::channel(/*init*/ 0u64).1,
+    );
 
     assert_eq!(
         remote_handle.disable(),
@@ -940,6 +955,195 @@ async fn remote_control_handle_disable_clears_stale_pairing_client() {
     );
 }
 
+#[tokio::test]
+async fn remote_control_handle_rejects_pairing_client_after_auth_change() {
+    let (auth_change_tx, auth_change_rx) = watch::channel(/*init*/ 0u64);
+    let remote_handle =
+        remote_control_handle_with_pairing_client(TEST_REMOTE_CONTROL_URL, auth_change_rx);
+    auth_change_tx.send_modify(|revision| *revision += 1);
+
+    assert_eq!(
+        remote_handle
+            .start_pairing(RemoteControlPairingStartParams::default())
+            .await
+            .expect_err("pairing should wait for current-account enrollment")
+            .to_string(),
+        "remote control pairing is unavailable until enrollment completes"
+    );
+}
+
+#[tokio::test]
+async fn remote_control_handle_discards_pairing_response_after_auth_change() {
+    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");
+    save_auth(
+        codex_home.path(),
+        &remote_control_auth_dot_json(Some("account_id")),
+        AuthCredentialsStoreMode::File,
+    )
+    .expect("initial auth should save");
+    let auth_manager = AuthManager::shared(
+        codex_home.path().to_path_buf(),
+        /*enable_codex_api_key_env*/ false,
+        AuthCredentialsStoreMode::File,
+        /*chatgpt_base_url*/ None,
+    )
+    .await;
+    let remote_handle = remote_control_handle_with_pairing_client(
+        &remote_control_url,
+        auth_manager.auth_change_receiver(),
+    );
+    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_REMOTE_CONTROL_SERVER_TOKEN}"))
+    );
+    assert_eq!(
+        serde_json::from_str::<serde_json::Value>(&pairing_request.body)
+            .expect("pairing request body should deserialize"),
+        json!({ "manual_code": false })
+    );
+
+    save_auth(
+        codex_home.path(),
+        &remote_control_auth_dot_json(Some("next_account_id")),
+        AuthCredentialsStoreMode::File,
+    )
+    .expect("next auth should save");
+    auth_manager.reload().await;
+    respond_with_json(
+        pairing_request.stream,
+        json!({
+            "pairing_code": "stale-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_err("stale pairing response should be discarded")
+            .to_string(),
+        "remote control pairing is unavailable until enrollment completes"
+    );
+}
+
+#[tokio::test]
+async fn remote_control_handle_clears_pairing_client_after_auth_change() {
+    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");
+    save_auth(
+        codex_home.path(),
+        &remote_control_auth_dot_json(Some("account_id")),
+        AuthCredentialsStoreMode::File,
+    )
+    .expect("initial auth should save");
+    let auth_manager = AuthManager::shared(
+        codex_home.path().to_path_buf(),
+        /*enable_codex_api_key_env*/ false,
+        AuthCredentialsStoreMode::File,
+        /*chatgpt_base_url*/ None,
+    )
+    .await;
+    let (transport_event_tx, _transport_event_rx) =
+        mpsc::channel::<TransportEvent>(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),
+        auth_manager.clone(),
+        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_initial",
+            "env_initial",
+            TEST_REMOTE_CONTROL_SERVER_TOKEN,
+        ),
+    )
+    .await;
+    let mut first_websocket = accept_remote_control_connection(&listener).await;
+
+    save_auth(
+        codex_home.path(),
+        &remote_control_auth_dot_json(Some("next_account_id")),
+        AuthCredentialsStoreMode::File,
+    )
+    .expect("next auth should save");
+    auth_manager.reload().await;
+    expect_remote_control_connection_closed(
+        &mut first_websocket,
+        "auth change should close the stale websocket",
+    )
+    .await;
+    assert_eq!(
+        remote_handle
+            .start_pairing(RemoteControlPairingStartParams::default())
+            .await
+            .expect_err("pairing should wait for current-account enrollment")
+            .to_string(),
+        "remote control pairing is unavailable until enrollment completes"
+    );
+
+    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(
+            "srv_e_next",
+            "env_next",
+            TEST_REFRESHED_REMOTE_CONTROL_SERVER_TOKEN,
+        ),
+    )
+    .await;
+    let mut second_websocket = accept_remote_control_connection(&listener).await;
+    second_websocket
+        .close(None)
+        .await
+        .expect("second websocket should close");
+
+    shutdown_token.cancel();
+    let _ = remote_task.await;
+}
+
 #[tokio::test]
 async fn remote_control_transport_clears_outgoing_buffer_when_backend_acks() {
     let listener = TcpListener::bind("127.0.0.1:0")
@@ -1893,6 +2097,34 @@ async fn accept_remote_control_connection(listener: &TcpListener) -> WebSocketSt
         .expect("websocket handshake should succeed")
 }
 
+async fn expect_remote_control_connection_closed(
+    websocket: &mut WebSocketStream<TcpStream>,
+    timeout_message: &str,
+) {
+    loop {
+        let frame = timeout(Duration::from_secs(5), websocket.next())
+            .await
+            .expect(timeout_message);
+        let Some(frame) = frame else {
+            return;
+        };
+        let Ok(frame) = frame else {
+            return;
+        };
+        match frame {
+            tungstenite::Message::Close(_) => return,
+            tungstenite::Message::Ping(payload) => {
+                websocket
+                    .send(tungstenite::Message::Pong(payload))
+                    .await
+                    .expect("websocket pong should send");
+            }
+            tungstenite::Message::Pong(_) | tungstenite::Message::Frame(_) => {}
+            frame => panic!("unexpected websocket frame while waiting for close: {frame:?}"),
+        }
+    }
+}
+
 async fn accept_http_request(listener: &TcpListener) -> CapturedHttpRequest {
     let (stream, _) = timeout(Duration::from_secs(5), listener.accept())
         .await

codex-rs/app-server-transport/src/transport/remote_control/websocket.rs

@@ -286,6 +286,8 @@ enum ConnectionEndReason {
     Shutdown,
     Disabled,
     EnabledWatchClosed,
+    AuthChanged,
+    AuthWatchClosed,
     ConnectionWorkerStopped,
 }
 
@@ -703,7 +705,7 @@ impl RemoteControlWebsocket {
     }
 
     async fn run_connection(
-        &self,
+        &mut self,
         websocket_connection: WebSocketStream<MaybeTlsStream<TcpStream>>,
         shutdown_token: CancellationToken,
     ) -> ConnectionEndReason {
@@ -738,8 +740,17 @@ impl RemoteControlWebsocket {
                     ConnectionEndReason::EnabledWatchClosed
                 }
             }
+            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_client);
         shutdown_token.cancel();
 
         Self::join_connection_workers(&mut join_set, REMOTE_CONTROL_CONNECTION_SHUTDOWN_TIMEOUT)
@@ -1252,6 +1263,7 @@ pub(super) async fn connect_remote_control_websocket(
         ));
     };
 
+    let auth_change_revision = *auth_context.auth_change_rx.borrow();
     let auth = match load_remote_control_auth(auth_context.auth_manager).await {
         Ok(auth) => auth,
         Err(err) => {
@@ -1386,7 +1398,6 @@ 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")
     })?;
-    set_pairing_client(pairing_client, remote_control_target, enrollment_ref)?;
     let request = build_remote_control_websocket_request(
         &remote_control_target.websocket_url,
         enrollment_ref,
@@ -1410,8 +1421,17 @@ pub(super) async fn connect_remote_control_websocket(
     })?;
 
     match websocket_connect_result {
-        Ok((websocket_stream, response)) => Ok((websocket_stream, response.map(|_| ()))),
+        Ok((websocket_stream, response)) => {
+            set_pairing_client(
+                pairing_client,
+                remote_control_target,
+                enrollment_ref,
+                auth_change_revision,
+            )?;
+            Ok((websocket_stream, response.map(|_| ())))
+        }
         Err(err) => {
+            clear_pairing_client(pairing_client);
             match &err {
                 tungstenite::Error::Http(response) if response.status().as_u16() == 404 => {
                     info!(
@@ -1490,6 +1510,7 @@ fn set_pairing_client(
     pairing_client: &Arc<StdMutex<Option<RemoteControlPairingClient>>>,
     remote_control_target: &RemoteControlTarget,
     enrollment: &RemoteControlEnrollment,
+    auth_change_revision: u64,
 ) -> io::Result<()> {
     let remote_control_token = enrollment.remote_control_token.clone().ok_or_else(|| {
         io::Error::new(
@@ -1512,6 +1533,7 @@ fn set_pairing_client(
             enrollment.server_id.clone(),
             enrollment.environment_id.clone(),
             expires_at,
+            auth_change_revision,
         ));
     Ok(())
 }
@@ -1890,6 +1912,12 @@ mod tests {
 
         server_task.await.expect("server task should succeed");
         assert_eq!(err.to_string(), expected_error);
+        assert!(
+            pairing_client
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner)
+                .is_none()
+        );
         assert_eq!(
             status_rx.borrow().clone(),
             RemoteControlStatusChangedNotification {