agent/codex
1
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(remote-control): allow pairing while disabled (#26215)

Browse Source 
## Why

`remoteControl/pairing/start` creates authorization for future
remote-control connections, so it should not require the live websocket
to already be enabled. Requiring enable first made pairing depend on
presence instead of the persisted server enrollment that pairing
actually uses.

Pairing also needs to recover when that persisted server row is stale.
If `/server/pair` returns `404`, making the first pairing attempt fail
forces a manual retry even though the client can clear the stale row and
create a replacement enrollment immediately.

## What Changed

- Allow `remoteControl/pairing/start` to reuse or create the persisted
remote-control server enrollment while remote control is disabled.
- Keep the selected in-memory enrollment across disable and share it
with websocket connect so a later enable uses the same selected server.
- Thread the app-server client name through pairing so stdio persistence
keeps using the websocket-owned enrollment key.
- Recover pairing server-token auth failures through the existing
refresh/auth-recovery path.
- Recover stale pairing enrollment on `/server/pair` `404` by clearing
the stale selected enrollment, re-enrolling once, and retrying pairing
once.
- Add focused disabled-pairing and stale-pairing recovery coverage.

## Verification

-
`remote_control_pairing_start_returns_pairing_artifacts_while_disabled`
exercises pairing before enable.
- `remote_control_handle_reenrolls_after_stale_pairing_enrollment`
exercises stale `/server/pair` `404` recovery without a manual retry.

Related: N/A
This commit is contained in:
Anton Panasenko
2026-06-04 22:12:23 -07:00
committed by GitHub
Unverified
parent a2f5874b7a
commit 64e0829cab
9 changed files with 702 additions and 244 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server-transport/src/transport/remote_control/mod.rs
+293 -60
View File
@@ -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<watch::Sender<bool>>,
status_tx: Arc<watch::Sender<RemoteControlStatusChangedNotification>>,
state_db_available: bool,
state_db: Option<Arc<StateRuntime>>,
remote_control_url: String,
current_enrollment: CurrentRemoteControlEnrollment,
pairing_persistence_key: RemoteControlPairingPersistenceKey,
pairing_persistence_key_required: bool,
auth_manager: Arc<AuthManager>,
}
type CurrentRemoteControlEnrollment = Arc<StdMutex<Option<RemoteControlEnrollment>>>;
// 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<RemoteControlEnrollmentState>;
type RemoteControlPairingPersistenceKey = watch::Sender<Option<String>>;
struct RemoteControlEnrollmentState {
enrollment: StdMutex<Option<RemoteControlEnrollment>>,
lock: Semaphore,
}
impl RemoteControlEnrollmentState {
fn new(enrollment: Option<RemoteControlEnrollment>) -> 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<RemoteControlEnrollment> {
self.enrollment
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone()
}
}
struct RemoteControlEnrollmentLease<'a> {
state: &'a RemoteControlEnrollmentState,
enrollment: Option<RemoteControlEnrollment>,
_permit: SemaphorePermit<'a>,
}
impl Deref for RemoteControlEnrollmentLease<'_> {
type Target = Option<RemoteControlEnrollment>;
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<RemoteControlPairingStartResponse> {
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<RemoteControlEnrollment>,
auth: &mut auth::RemoteControlConnectionAuth,
installation_id: &str,
server_name: &str,
app_server_client_name: Option<&str>,
) -> io::Result<RemoteControlEnrollment> {
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<Option<String>> {
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<AuthManager>,
auth: &mut auth::RemoteControlConnectionAuth,
remote_control_target: &protocol::RemoteControlTarget,
installation_id: &str,
server_name: &str,
) -> io::Result<RemoteControlEnrollment> {
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<RemoteControlEnrollment>,
state_db: Option<&StateRuntime>,
app_server_client_name: Option<&str>,
auth_manager: &Arc<AuthManager>,
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<RemoteControlEnrollment>,
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<RemoteControlEnrollment>,
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<RemoteControlEnrollment>,
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<RemoteControlEnrollment>,
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<RemoteControlEnrollment>,
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<RemoteControlEnrollment>,
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,
},
))
codex-rs/app-server-transport/src/transport/remote_control/tests.rs
+17 -13
View File
@@ -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,
}
}
codex-rs/app-server-transport/src/transport/remote_control/tests/clients_tests.rs
+4 -1
View File
@@ -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,
}
}
codex-rs/app-server-transport/src/transport/remote_control/tests/pairing_tests.rs
+142 -14
View File
@@ -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
}
});
codex-rs/app-server-transport/src/transport/remote_control/websocket.rs
+197 -151
View File
@@ -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<RemoteControlEnrollment>,
auth_recovery: UnauthorizedRecovery,
auth_change_rx: watch::Receiver<u64>,
current_enrollment: CurrentRemoteControlEnrollment,
pairing_persistence_key: RemoteControlPairingPersistenceKey,
client_tracker: Arc<Mutex<ClientTracker>>,
state: Arc<Mutex<WebsocketState>>,
server_event_rx: Arc<Mutex<mpsc::Receiver<super::QueuedServerEnvelope>>>,
@@ -294,6 +297,7 @@ pub(super) struct RemoteControlChannels {
pub(super) transport_event_tx: mpsc::Sender<TransportEvent>,
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<RemoteControlEnrollment>,
current_enrollment: &CurrentRemoteControlEnrollment,
connect_options: RemoteControlConnectOptions<'_>,
status_publisher: &RemoteControlStatusPublisher,
current_enrollment: &CurrentRemoteControlEnrollment,
) -> io::Result<(
WebSocketStream<MaybeTlsStream<TcpStream>>,
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<RemoteControlEnrollment>,
connect_options: RemoteControlConnectOptions<'_>,
status_publisher: &RemoteControlStatusPublisher,
) -> io::Result<RemoteControlConnectionAuth> {
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<RemoteControlEnrollment>,
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<RemoteControlEnrollment>,
) -> 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,
&current_enrollment,
RemoteControlConnectOptions {
installation_id: TEST_INSTALLATION_ID,
server_name: "test-server",
@@ -1761,7 +1824,6 @@ mod tests {
app_server_client_name: None,
},
&status_publisher,
&current_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,
&current_enrollment,
RemoteControlConnectOptions {
installation_id: TEST_INSTALLATION_ID,
server_name: "test-server",
@@ -1832,7 +1888,6 @@ mod tests {
app_server_client_name: None,
},
&status_publisher,
&current_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,
&current_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,
&current_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,
&current_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,
&current_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,
codex-rs/app-server/src/message_processor.rs
+1 -1
View File
@@ -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
codex-rs/app-server/src/request_processors/remote_control_processor.rs
+2 -1
View File
@@ -52,9 +52,10 @@ impl RemoteControlRequestProcessor {
pub(crate) async fn pairing_start(
&self,
params: RemoteControlPairingStartParams,
app_server_client_name: Option<&str>,
) -> Result<RemoteControlPairingStartResponse, JSONRPCErrorError> {
self.handle()?
.start_pairing(params)
.start_pairing(params, app_server_client_name)
.await
.map_err(map_pairing_start_error)
}
codex-rs/app-server/src/request_processors/remote_control_processor/remote_control_processor_tests.rs
+4 -1
View File
@@ -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");
codex-rs/app-server/tests/suite/v2/remote_control.rs
+42 -2
View File
@@ -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!({