feat(app-server): persist remote-control desired state (#27445)

## Why

Remote-control runtime enablement and persisted enrollment preference
were represented by separate flags. That made startup rehydration, RPC
persistence, and new-enrollment seeding race with one another, and it
did not cleanly distinguish runtime-only CLI or daemon starts from
durable app-server RPC changes.

## What Changed

- Replace the parallel enablement, seed, and rehydration flags with one
transport-owned `RemoteControlDesiredState`.
- Add nullable enrollment-scoped persistence and preserve existing
preferences during enrollment upserts.
- Rehydrate plain startup only after auth and client scope resolve,
without overwriting a concurrent RPC transition.
- Make ordinary `remoteControl/enable` and `remoteControl/disable`
durable while retaining `ephemeral: true` for runtime-only callers.
- Have the daemon explicitly request ephemeral enablement and regenerate
the app-server schemas.

## Verification

- Covered migration and `NULL`/`0`/`1` persistence round trips.
- Covered plain-start rehydration and runtime-only versus durable
enrollment seeding.
- Covered durable enable, durable disable, and ephemeral enable through
app-server RPC.
- Covered the daemon's exact `{ "ephemeral": true }` request payload.

Related issue: N/A (internal remote-control persistence architecture
change).
This commit is contained in:
Anton Panasenko authored and GitHub committed 2026-06-11 21:28:52 -07:00
1 parent be338ee9a2
commit d61dfeb23a
33 files changed
+2157 -412

No files matched your search

+3
View File
@@ -11,8 +11,10 @@ pub use transport::AppServerTransport;
pub use transport::AppServerTransportParseError;
pub use transport::CHANNEL_CAPACITY;
pub use transport::ConnectionOrigin;
pub use transport::REMOTE_CONTROL_DISABLED_ENV_VAR;
pub use transport::RemoteControlHandle;
pub use transport::RemoteControlStartConfig;
pub use transport::RemoteControlStartupMode;
pub use transport::RemoteControlUnavailable;
pub use transport::TransportEvent;
pub use transport::acquire_app_server_startup_lock;
@@ -24,3 +26,4 @@ pub use transport::start_control_socket_acceptor;
pub use transport::start_remote_control;
pub use transport::start_stdio_connection;
pub use transport::start_websocket_acceptor;
pub use transport::take_remote_control_disabled_env;
@@ -30,10 +30,13 @@ mod unix_socket;
mod unix_socket_tests;
mod websocket;
pub use remote_control::REMOTE_CONTROL_DISABLED_ENV_VAR;
pub use remote_control::RemoteControlHandle;
pub use remote_control::RemoteControlStartConfig;
pub use remote_control::RemoteControlStartupMode;
pub use remote_control::RemoteControlUnavailable;
pub use remote_control::start_remote_control;
pub use remote_control::take_remote_control_disabled_env;
pub use stdio::start_stdio_connection;
pub use unix_socket::AppServerStartupLock;
pub use unix_socket::acquire_app_server_startup_lock;
@@ -0,0 +1,158 @@
use super::RemoteControlHandle;
use super::RemoteControlUnavailable;
use super::enroll::update_persisted_remote_control_enrollment;
use super::protocol::normalize_remote_control_url;
use super::publish_current_enrollment;
use super::websocket::RemoteControlStatusPublisher;
use codex_app_server_protocol::RemoteControlStatusChangedNotification;
use codex_state::RemoteControlEnrollmentRecord;
use std::io;
use tokio::sync::Semaphore;
use tokio::sync::SemaphorePermit;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum RemoteControlDesiredState {
// `Unknown` exists only on plain startup before auth and enrollment scope resolve. Persisted
// `1` is `Enabled { persistence_preference: Some(true) }`; `0`, `NULL`, or no row are
// `Disabled`. Runtime-only enable is `Enabled { persistence_preference: None }`, so new rows
// keep `NULL`; durable RPC enable uses `Some(true)`, so new rows get `1`. Durable disable writes
// `0` before entering `Disabled`; runtime-only disable does not write. `Disabled` carries no
// preference because disabled sessions do not create enrollments.
Unknown,
Disabled,
Enabled {
persistence_preference: Option<bool>,
},
}
impl RemoteControlDesiredState {
pub(super) fn is_enabled(self) -> bool {
matches!(self, Self::Enabled { .. })
}
}
pub(super) async fn acquire_persistence_lock(lock: &Semaphore) -> SemaphorePermit<'_> {
lock.acquire().await.unwrap_or_else(|_| unreachable!())
}
pub(super) fn desired_state_from_persisted_enrollment(
enrollment: Option<RemoteControlEnrollmentRecord>,
) -> RemoteControlDesiredState {
if enrollment.and_then(|enrollment| enrollment.remote_control_enabled) == Some(true) {
RemoteControlDesiredState::Enabled {
persistence_preference: Some(true),
}
} else {
RemoteControlDesiredState::Disabled
}
}
impl RemoteControlHandle {
pub async fn resolve_persisted_preference(
&self,
app_server_client_name: Option<&str>,
) -> io::Result<bool> {
let _transition = self
.desired_state_rpc_lock
.acquire()
.await
.unwrap_or_else(|_| unreachable!());
if !matches!(
*self.desired_state_tx.borrow(),
RemoteControlDesiredState::Unknown
) {
return Ok(self.desired_state_tx.borrow().is_enabled());
}
let state_db = self
.state_db
.as_deref()
.ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, RemoteControlUnavailable))?;
let auth = super::auth::load_remote_control_auth(&self.auth_manager).await?;
let remote_control_target = normalize_remote_control_url(&self.remote_control_url)?;
let app_server_client_name = self.pairing_persistence_key(app_server_client_name)?;
let enrollment = state_db
.get_remote_control_enrollment(
&remote_control_target.websocket_url,
&auth.account_id,
app_server_client_name.as_deref(),
)
.await
.map_err(io::Error::other)?;
let desired_state = desired_state_from_persisted_enrollment(enrollment);
self.desired_state_tx.send_if_modified(|state| {
if !matches!(*state, RemoteControlDesiredState::Unknown) {
return false;
}
*state = desired_state;
true
});
Ok(self.desired_state_tx.borrow().is_enabled())
}
pub async fn enable(
&self,
app_server_client_name: Option<&str>,
) -> io::Result<RemoteControlStatusChangedNotification> {
let _transition = self
.desired_state_rpc_lock
.acquire()
.await
.unwrap_or_else(|_| unreachable!());
let state_db = self
.state_db
.as_deref()
.ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, RemoteControlUnavailable))?;
let mut auth = super::auth::load_remote_control_auth(&self.auth_manager).await?;
let remote_control_target = normalize_remote_control_url(&self.remote_control_url)?;
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 status = self.status();
let mut current_enrollment = self.current_enrollment.lock().await;
let (enrollment, _) = self
.load_or_enroll_server(
&current_enrollment,
&mut auth,
&status.installation_id,
&status.server_name,
app_server_client_name,
super::RemoteControlEnrollmentSelection::ReuseOrCreate,
)
.await?;
let current_auth = super::auth::load_remote_control_auth(&self.auth_manager).await?;
if current_auth.account_id != auth.account_id {
return Err(io::Error::new(
io::ErrorKind::Interrupted,
"remote control account changed during enrollment",
));
}
let _persistence = acquire_persistence_lock(&self.desired_state_persistence_lock).await;
let updated = state_db
.set_remote_control_enabled(
&remote_control_target.websocket_url,
&auth.account_id,
app_server_client_name,
/*remote_control_enabled*/ true,
)
.await
.map_err(io::Error::other)?;
if updated == 0 {
update_persisted_remote_control_enrollment(
Some(state_db),
&remote_control_target,
&auth.account_id,
app_server_client_name,
Some(&enrollment),
Some(true),
)
.await?;
}
publish_current_enrollment(&mut current_enrollment, &enrollment);
self.enable_with_preference(Some(true))
.map_err(|err| io::Error::new(io::ErrorKind::NotFound, err))?;
RemoteControlStatusPublisher::new(self.status_tx.as_ref().clone())
.publish_environment_id(Some(enrollment.environment_id));
Ok(self.status())
}
}
@@ -286,6 +286,7 @@ pub(super) async fn update_persisted_remote_control_enrollment(
account_id: &str,
app_server_client_name: Option<&str>,
enrollment: Option<&RemoteControlEnrollment>,
remote_control_enabled: Option<bool>,
) -> io::Result<()> {
let Some(state_db) = state_db else {
return Err(io::Error::new(
@@ -316,6 +317,7 @@ pub(super) async fn update_persisted_remote_control_enrollment(
server_id: enrollment.server_id.clone(),
environment_id: enrollment.environment_id.clone(),
server_name: enrollment.server_name.clone(),
remote_control_enabled,
})
.await
.map_err(io::Error::other)?;
@@ -648,6 +650,7 @@ mod tests {
"account-a",
Some("desktop-client"),
Some(&first_enrollment),
/*remote_control_enabled*/ None,
)
.await
.expect("first enrollment should persist");
@@ -657,6 +660,7 @@ mod tests {
"account-a",
Some("desktop-client"),
Some(&second_enrollment),
/*remote_control_enabled*/ None,
)
.await
.expect("second enrollment should persist");
@@ -730,6 +734,7 @@ mod tests {
"account-a",
/*app_server_client_name*/ None,
Some(&first_enrollment),
/*remote_control_enabled*/ None,
)
.await
.expect("first enrollment should persist");
@@ -739,6 +744,7 @@ mod tests {
"account-a",
/*app_server_client_name*/ None,
Some(&second_enrollment),
/*remote_control_enabled*/ None,
)
.await
.expect("second enrollment should persist");
@@ -749,6 +755,7 @@ mod tests {
"account-a",
/*app_server_client_name*/ None,
/*enrollment*/ None,
/*remote_control_enabled*/ None,
)
.await
.expect("matching enrollment should clear");
@@ -1,6 +1,7 @@
mod auth;
mod client_tracker;
mod clients;
mod desired_state;
mod enroll;
mod protocol;
mod segment;
@@ -8,6 +9,8 @@ mod websocket;
use self::auth::load_remote_control_auth;
use self::auth::recover_remote_control_auth;
use self::desired_state::RemoteControlDesiredState;
use self::desired_state::acquire_persistence_lock;
use self::enroll::RemoteControlEnrollment;
use self::enroll::enroll_remote_control_server;
use self::enroll::load_persisted_remote_control_enrollment;
@@ -63,6 +66,26 @@ pub struct RemoteControlStartConfig {
pub installation_id: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RemoteControlStartupMode {
ResolvePersisted,
DisabledEphemeral,
EnabledEphemeral,
}
/// Internal marker used by the daemon to disable remote control without requiring a new CLI flag.
pub const REMOTE_CONTROL_DISABLED_ENV_VAR: &str =
"CODEX_INTERNAL_APP_SERVER_REMOTE_CONTROL_DISABLED";
/// Reads and removes the daemon's internal disabled-start marker before worker threads start.
pub fn take_remote_control_disabled_env() -> bool {
let disabled =
std::env::var_os(REMOTE_CONTROL_DISABLED_ENV_VAR).is_some_and(|value| value == "1");
// SAFETY: app-server calls this synchronously at process startup, before spawning threads.
unsafe { std::env::remove_var(REMOTE_CONTROL_DISABLED_ENV_VAR) };
disabled
}
pub(super) struct QueuedServerEnvelope {
pub(super) event: ServerEvent,
pub(super) client_id: ClientId,
@@ -72,9 +95,10 @@ pub(super) struct QueuedServerEnvelope {
#[derive(Clone)]
pub struct RemoteControlHandle {
enabled_tx: Arc<watch::Sender<bool>>,
desired_state_tx: Arc<watch::Sender<RemoteControlDesiredState>>,
desired_state_rpc_lock: Arc<Semaphore>,
desired_state_persistence_lock: Arc<Semaphore>,
status_tx: Arc<watch::Sender<RemoteControlStatusChangedNotification>>,
state_db_available: bool,
state_db: Option<Arc<StateRuntime>>,
remote_control_url: String,
current_enrollment: CurrentRemoteControlEnrollment,
@@ -83,11 +107,17 @@ pub struct RemoteControlHandle {
auth_manager: Arc<AuthManager>,
}
// Pairing and websocket connect share one selected server so they cannot enroll or clear
// Pairing and websocket connect share one selected server so they cannot enroll or replace
// different persisted rows while either path is awaiting backend I/O.
type CurrentRemoteControlEnrollment = Arc<RemoteControlEnrollmentState>;
type RemoteControlPairingPersistenceKey = watch::Sender<Option<String>>;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum RemoteControlEnrollmentSelection {
ReuseOrCreate,
ReplaceExisting,
}
struct RemoteControlEnrollmentState {
enrollment: StdMutex<Option<RemoteControlEnrollment>>,
lock: Semaphore,
@@ -166,23 +196,45 @@ impl fmt::Display for RemoteControlUnavailable {
impl Error for RemoteControlUnavailable {}
impl RemoteControlHandle {
pub fn enable(
pub fn enable_ephemeral(
&self,
) -> Result<RemoteControlStatusChangedNotification, RemoteControlUnavailable> {
if !self.state_db_available {
self.enable_with_preference(/*persistence_preference*/ None)
}
fn enable_with_preference(
&self,
persistence_preference: Option<bool>,
) -> Result<RemoteControlStatusChangedNotification, RemoteControlUnavailable> {
if self.state_db.is_none() {
warn!("remote control cannot be enabled because sqlite state db is unavailable");
return Err(RemoteControlUnavailable);
}
let enabled_changed = self.enabled_tx.send_if_modified(|state| {
let changed = !*state;
*state = true;
let mut effective_persistence_preference = persistence_preference;
let desired_state_changed = self.desired_state_tx.send_if_modified(|state| {
if effective_persistence_preference.is_none()
&& matches!(
*state,
RemoteControlDesiredState::Enabled {
persistence_preference: Some(true)
}
)
{
effective_persistence_preference = Some(true);
}
let next_state = RemoteControlDesiredState::Enabled {
persistence_preference: effective_persistence_preference,
};
let changed = *state != next_state;
*state = next_state;
changed
});
let status = self.status();
info!(
enabled_changed,
desired_state_changed,
?effective_persistence_preference,
current_status = ?status.status,
environment_id = ?status.environment_id,
installation_id = %status.installation_id,
@@ -199,15 +251,43 @@ impl RemoteControlHandle {
Ok(self.publish_status(RemoteControlConnectionStatus::Connecting))
}
pub fn disable(&self) -> RemoteControlStatusChangedNotification {
let enabled_changed = self.enabled_tx.send_if_modified(|state| {
let changed = *state;
*state = false;
pub async fn disable(
&self,
app_server_client_name: Option<&str>,
) -> io::Result<RemoteControlStatusChangedNotification> {
let _transition = self
.desired_state_rpc_lock
.acquire()
.await
.unwrap_or_else(|_| unreachable!());
let _persistence = acquire_persistence_lock(&self.desired_state_persistence_lock).await;
self.persist_preference(
app_server_client_name,
/*remote_control_enabled*/ false,
)
.await?;
Ok(self.transition_disabled())
}
pub async fn disable_ephemeral(&self) -> RemoteControlStatusChangedNotification {
let _transition = self
.desired_state_rpc_lock
.acquire()
.await
.unwrap_or_else(|_| unreachable!());
let _persistence = acquire_persistence_lock(&self.desired_state_persistence_lock).await;
self.transition_disabled()
}
fn transition_disabled(&self) -> RemoteControlStatusChangedNotification {
let desired_state_changed = self.desired_state_tx.send_if_modified(|state| {
let changed = *state != RemoteControlDesiredState::Disabled;
*state = RemoteControlDesiredState::Disabled;
changed
});
let status = self.status();
info!(
enabled_changed,
desired_state_changed,
current_status = ?status.status,
environment_id = ?status.environment_id,
installation_id = %status.installation_id,
@@ -217,6 +297,30 @@ impl RemoteControlHandle {
self.publish_status(RemoteControlConnectionStatus::Disabled)
}
async fn persist_preference(
&self,
app_server_client_name: Option<&str>,
remote_control_enabled: bool,
) -> io::Result<()> {
let state_db = self
.state_db
.as_deref()
.ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, RemoteControlUnavailable))?;
let auth = load_remote_control_auth(&self.auth_manager).await?;
let remote_control_target = normalize_remote_control_url(&self.remote_control_url)?;
let app_server_client_name = self.pairing_persistence_key(app_server_client_name)?;
state_db
.set_remote_control_enabled(
&remote_control_target.websocket_url,
&auth.account_id,
app_server_client_name.as_deref(),
remote_control_enabled,
)
.await
.map_err(io::Error::other)?;
Ok(())
}
pub fn status(&self) -> RemoteControlStatusChangedNotification {
self.status_tx.borrow().clone()
}
@@ -230,6 +334,9 @@ impl RemoteControlHandle {
params: RemoteControlPairingStartParams,
app_server_client_name: Option<&str>,
) -> io::Result<RemoteControlPairingStartResponse> {
if !self.desired_state_tx.borrow().is_enabled() {
return Err(Self::pairing_disabled_error());
}
let mut auth = load_remote_control_auth(&self.auth_manager)
.await
.map_err(|_| pairing_unavailable_error())?;
@@ -245,19 +352,35 @@ impl RemoteControlHandle {
&installation_id,
&status.server_name,
app_server_client_name,
RemoteControlEnrollmentSelection::ReuseOrCreate,
)
.await?;
if enrollment.should_refresh_server_token() {
refresh_pairing_enrollment(
let refresh_result = refresh_pairing_enrollment(
&mut current_enrollment,
self.state_db.as_deref(),
app_server_client_name,
&self.auth_manager,
&mut auth,
&installation_id,
&mut enrollment,
)
.await?;
.await;
if refresh_result
.as_ref()
.is_err_and(|err| err.kind() == io::ErrorKind::NotFound)
{
enrollment = self
.load_or_enroll_pairing_server(
&mut current_enrollment,
&mut auth,
&installation_id,
&status.server_name,
app_server_client_name,
RemoteControlEnrollmentSelection::ReplaceExisting,
)
.await?;
} else {
refresh_result?;
}
}
let pairing_request = || protocol::StartRemoteControlPairingRequest {
manual_code: params.manual_code,
@@ -267,8 +390,6 @@ impl RemoteControlHandle {
clear_pairing_server_token(&mut current_enrollment, &mut enrollment)?;
refresh_pairing_enrollment(
&mut current_enrollment,
self.state_db.as_deref(),
app_server_client_name,
&self.auth_manager,
&mut auth,
&installation_id,
@@ -278,13 +399,6 @@ impl RemoteControlHandle {
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,
@@ -292,6 +406,7 @@ impl RemoteControlHandle {
&installation_id,
&status.server_name,
app_server_client_name,
RemoteControlEnrollmentSelection::ReplaceExisting,
)
.await?;
enrollment.start_pairing(pairing_request()).await
@@ -301,13 +416,15 @@ impl RemoteControlHandle {
if let Err(err) = &pairing_response {
match err.kind() {
io::ErrorKind::NotFound => {
clear_pairing_enrollment(
self.load_or_enroll_pairing_server(
&mut current_enrollment,
self.state_db.as_deref(),
&mut auth,
&installation_id,
&status.server_name,
app_server_client_name,
&enrollment,
RemoteControlEnrollmentSelection::ReplaceExisting,
)
.await;
.await?;
return Err(pairing_unavailable_error());
}
io::ErrorKind::PermissionDenied => {
@@ -323,6 +440,9 @@ impl RemoteControlHandle {
if current_auth.account_id != auth.account_id {
return Err(pairing_unavailable_error());
}
if !self.desired_state_tx.borrow().is_enabled() {
return Err(Self::pairing_disabled_error());
}
pairing_response
}
@@ -333,31 +453,86 @@ impl RemoteControlHandle {
installation_id: &str,
server_name: &str,
app_server_client_name: Option<&str>,
selection: RemoteControlEnrollmentSelection,
) -> io::Result<RemoteControlEnrollment> {
if let Some(enrollment) = current_enrollment
.as_ref()
.filter(|enrollment| enrollment.account_id == auth.account_id)
.cloned()
{
let (enrollment, created) = self
.load_or_enroll_server(
current_enrollment,
auth,
installation_id,
server_name,
app_server_client_name,
selection,
)
.await?;
if !created {
publish_current_enrollment(current_enrollment, &enrollment);
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(
let _persistence = acquire_persistence_lock(&self.desired_state_persistence_lock).await;
let persistence_preference = match *self.desired_state_tx.borrow() {
RemoteControlDesiredState::Enabled {
persistence_preference,
} => persistence_preference,
RemoteControlDesiredState::Unknown | RemoteControlDesiredState::Disabled => {
return Err(Self::pairing_disabled_error());
}
};
update_persisted_remote_control_enrollment(
Some(state_db),
&remote_control_target,
&enrollment.remote_control_target,
&auth.account_id,
app_server_client_name,
Some(&enrollment),
persistence_preference,
)
.await?
{
enrollment.server_name = server_name.to_string();
publish_current_enrollment(current_enrollment, &enrollment);
return Ok(enrollment);
.await?;
publish_current_enrollment(current_enrollment, &enrollment);
Ok(enrollment)
}
async fn load_or_enroll_server(
&self,
current_enrollment: &Option<RemoteControlEnrollment>,
auth: &mut auth::RemoteControlConnectionAuth,
installation_id: &str,
server_name: &str,
app_server_client_name: Option<&str>,
selection: RemoteControlEnrollmentSelection,
) -> io::Result<(RemoteControlEnrollment, bool)> {
let remote_control_target = normalize_remote_control_url(&self.remote_control_url)?;
match selection {
RemoteControlEnrollmentSelection::ReuseOrCreate => {
if let Some(enrollment) = current_enrollment
.as_ref()
.filter(|enrollment| enrollment.account_id == auth.account_id)
.cloned()
{
return Ok((enrollment, false));
}
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();
return Ok((enrollment, false));
}
}
RemoteControlEnrollmentSelection::ReplaceExisting => {}
}
let enrollment = enroll_pairing_server(
@@ -368,16 +543,7 @@ impl RemoteControlHandle {
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)
Ok((enrollment, true))
}
fn pairing_persistence_key(
@@ -398,7 +564,7 @@ impl RemoteControlHandle {
&self,
params: RemoteControlPairingStatusParams,
) -> io::Result<RemoteControlPairingStatusResponse> {
if !*self.enabled_tx.borrow() {
if !self.desired_state_tx.borrow().is_enabled() {
return Err(Self::pairing_disabled_error());
}
let mut auth = load_remote_control_auth(&self.auth_manager)
@@ -412,18 +578,34 @@ impl RemoteControlHandle {
.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 server_name = status.server_name;
if enrollment.should_refresh_server_token() {
refresh_pairing_enrollment(
let refresh_result = refresh_pairing_enrollment(
&mut current_enrollment,
self.state_db.as_deref(),
app_server_client_name,
&self.auth_manager,
&mut auth,
&installation_id,
&mut enrollment,
)
.await?;
.await;
if refresh_result
.as_ref()
.is_err_and(|err| err.kind() == io::ErrorKind::NotFound)
{
self.load_or_enroll_pairing_server(
&mut current_enrollment,
&mut auth,
&installation_id,
&server_name,
app_server_client_name,
RemoteControlEnrollmentSelection::ReplaceExisting,
)
.await?;
return Err(pairing_unavailable_error());
}
refresh_result?;
}
let status_code = remote_control_pairing_status_code(&params)?;
let pairing_status_request =
@@ -434,8 +616,6 @@ impl RemoteControlHandle {
clear_pairing_server_token(&mut current_enrollment, &mut enrollment)?;
refresh_pairing_enrollment(
&mut current_enrollment,
self.state_db.as_deref(),
app_server_client_name,
&self.auth_manager,
&mut auth,
&installation_id,
@@ -449,13 +629,15 @@ impl RemoteControlHandle {
if let Err(err) = &pairing_status_response {
match err.kind() {
io::ErrorKind::NotFound => {
clear_pairing_enrollment(
self.load_or_enroll_pairing_server(
&mut current_enrollment,
self.state_db.as_deref(),
&mut auth,
&installation_id,
&server_name,
app_server_client_name,
&enrollment,
RemoteControlEnrollmentSelection::ReplaceExisting,
)
.await;
.await?;
return Err(pairing_unavailable_error());
}
io::ErrorKind::PermissionDenied => {
@@ -465,7 +647,7 @@ impl RemoteControlHandle {
_ => {}
}
}
if !*self.enabled_tx.borrow() {
if !self.desired_state_tx.borrow().is_enabled() {
return Err(Self::pairing_disabled_error());
}
let current_auth = load_remote_control_auth(&self.auth_manager)
@@ -580,8 +762,6 @@ fn remote_control_pairing_status_code(
async fn refresh_pairing_enrollment(
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,
@@ -589,14 +769,7 @@ 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,
state_db,
app_server_client_name,
enrollment,
err,
)
.await;
return Err(err);
}
let mut auth_recovery = auth_manager.unauthorized_recovery();
let mut auth_change_rx = auth_manager.auth_change_receiver();
@@ -609,16 +782,7 @@ async fn refresh_pairing_enrollment(
if auth.account_id != enrollment.account_id {
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,
state_db,
app_server_client_name,
enrollment,
err,
)
.await;
}
refresh_remote_control_server(auth, installation_id, enrollment).await?
}
if replace_current_enrollment(current_enrollment, enrollment) {
Ok(())
@@ -627,52 +791,6 @@ async fn refresh_pairing_enrollment(
}
}
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_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: &mut Option<RemoteControlEnrollment>,
enrollment: &mut RemoteControlEnrollment,
@@ -729,21 +847,6 @@ fn replace_current_enrollment(
true
}
fn clear_current_enrollment_if_matches(
current_enrollment: &mut Option<RemoteControlEnrollment>,
enrollment: &RemoteControlEnrollment,
) -> bool {
if current_enrollment
.as_ref()
.is_some_and(|current| same_remote_control_enrollment(current, enrollment))
{
*current_enrollment = None;
true
} else {
false
}
}
fn same_remote_control_enrollment(
left: &RemoteControlEnrollment,
right: &RemoteControlEnrollment,
@@ -762,11 +865,22 @@ pub async fn start_remote_control(
transport_event_tx: mpsc::Sender<TransportEvent>,
shutdown_token: CancellationToken,
app_server_client_name_rx: Option<oneshot::Receiver<String>>,
initial_enabled: bool,
startup_mode: RemoteControlStartupMode,
) -> io::Result<(JoinHandle<()>, RemoteControlHandle)> {
let state_db_available = state_db.is_some();
let requested_initial_enabled = initial_enabled;
let initial_enabled = initial_enabled && state_db_available;
let requested_initial_enabled = startup_mode == RemoteControlStartupMode::EnabledEphemeral;
let desired_state = if !state_db_available {
RemoteControlDesiredState::Disabled
} else {
match startup_mode {
RemoteControlStartupMode::ResolvePersisted => RemoteControlDesiredState::Unknown,
RemoteControlStartupMode::DisabledEphemeral => RemoteControlDesiredState::Disabled,
RemoteControlStartupMode::EnabledEphemeral => RemoteControlDesiredState::Enabled {
persistence_preference: None,
},
}
};
let initial_enabled = desired_state.is_enabled();
if requested_initial_enabled && !state_db_available {
warn!("remote control disabled because sqlite state db is unavailable");
}
@@ -776,7 +890,12 @@ pub async fn start_remote_control(
None
};
let (enabled_tx, enabled_rx) = watch::channel(initial_enabled);
let (desired_state_tx, _desired_state_rx) = watch::channel(desired_state);
let desired_state_tx = Arc::new(desired_state_tx);
let desired_state_rpc_lock = Arc::new(Semaphore::new(1));
let desired_state_persistence_lock = Arc::new(Semaphore::new(1));
let websocket_desired_state_tx = desired_state_tx.clone();
let websocket_desired_state_persistence_lock = desired_state_persistence_lock.clone();
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();
@@ -804,7 +923,7 @@ pub async fn start_remote_control(
installation_id = %installation_id,
server_name = %server_name,
state_db_available,
initial_enabled,
?desired_state,
"starting app-server remote control websocket task"
);
let remote_control_url_for_log = remote_control_url.clone();
@@ -817,7 +936,7 @@ pub async fn start_remote_control(
remote_control_url = %remote_control_url_for_log,
installation_id = %installation_id_for_log,
server_name = %server_name_for_log,
initial_enabled,
?desired_state,
"app-server remote control websocket task started"
);
let websocket_task = RemoteControlWebsocket::new(
@@ -834,9 +953,10 @@ pub async fn start_remote_control(
status_publisher,
current_enrollment: websocket_current_enrollment,
pairing_persistence_key: websocket_pairing_persistence_key,
desired_state_persistence_lock: websocket_desired_state_persistence_lock,
},
shutdown_token,
enabled_rx,
websocket_desired_state_tx,
)
.run(app_server_client_name_rx);
match AssertUnwindSafe(websocket_task).catch_unwind().await {
@@ -875,9 +995,10 @@ pub async fn start_remote_control(
Ok((
join_handle,
RemoteControlHandle {
enabled_tx: Arc::new(enabled_tx),
desired_state_tx,
desired_state_rpc_lock,
desired_state_persistence_lock,
status_tx: Arc::new(status_tx),
state_db_available,
state_db: handle_state_db,
remote_control_url: handle_remote_control_url,
current_enrollment,
@@ -9,6 +9,8 @@ use super::protocol::ClientId;
use super::protocol::StreamId;
use super::protocol::normalize_remote_control_url;
use super::websocket::REMOTE_CONTROL_PROTOCOL_VERSION;
use super::websocket::RemoteControlWebsocket;
use super::websocket::RemoteControlWebsocketConfig;
use super::*;
use crate::outgoing_message::OutgoingMessage;
use crate::outgoing_message::QueuedOutgoingMessage;
@@ -33,6 +35,7 @@ use codex_login::CodexAuth;
use codex_login::save_auth;
use codex_login::token_data::TokenData;
use codex_login::token_data::parse_chatgpt_jwt_claims;
use codex_state::RemoteControlEnrollmentRecord;
use codex_state::StateRuntime;
use futures::SinkExt;
use futures::StreamExt;
@@ -126,6 +129,137 @@ async fn remote_control_state_runtime(codex_home: &TempDir) -> Arc<StateRuntime>
.expect("state runtime should initialize")
}
#[tokio::test]
async fn plain_start_resolves_persisted_remote_control_preference() {
let cases = [
("enabled", Some(Some(true))),
("disabled", Some(Some(false))),
("unset", Some(None)),
("missing", None),
];
let codex_home = TempDir::new().expect("temp dir should create");
let state_db = remote_control_state_runtime(&codex_home).await;
let remote_control_target = normalize_remote_control_url(TEST_REMOTE_CONTROL_URL)
.expect("remote control target should normalize");
for (name, stored_preference) in cases {
let Some(remote_control_enabled) = stored_preference else {
continue;
};
state_db
.upsert_remote_control_enrollment(&RemoteControlEnrollmentRecord {
websocket_url: remote_control_target.websocket_url.clone(),
account_id: "account_id".to_string(),
app_server_client_name: Some(name.to_string()),
server_id: format!("server-{name}"),
environment_id: format!("environment-{name}"),
server_name: format!("server-name-{name}"),
remote_control_enabled,
})
.await
.expect("enrollment should persist");
}
let (transport_event_tx, _transport_event_rx) = mpsc::channel(CHANNEL_CAPACITY);
let (status_tx, _status_rx) = watch::channel(RemoteControlStatusChangedNotification {
status: RemoteControlConnectionStatus::Disabled,
server_name: test_server_name(),
installation_id: TEST_INSTALLATION_ID.to_string(),
environment_id: None,
});
let (desired_state_tx, _desired_state_rx) = watch::channel(RemoteControlDesiredState::Unknown);
let desired_state_tx = Arc::new(desired_state_tx);
let mut websocket = RemoteControlWebsocket::new(
RemoteControlWebsocketConfig {
remote_control_url: TEST_REMOTE_CONTROL_URL.to_string(),
installation_id: TEST_INSTALLATION_ID.to_string(),
remote_control_target: None,
server_name: test_server_name(),
},
Some(state_db),
remote_control_auth_manager(),
RemoteControlChannels {
transport_event_tx,
status_publisher: RemoteControlStatusPublisher::new(status_tx),
current_enrollment: Arc::new(RemoteControlEnrollmentState::new(
/*enrollment*/ None,
)),
pairing_persistence_key: watch::channel(None).0,
desired_state_persistence_lock: Arc::new(Semaphore::new(1)),
},
CancellationToken::new(),
desired_state_tx.clone(),
);
for (name, stored_preference) in cases {
desired_state_tx.send_replace(RemoteControlDesiredState::Unknown);
assert!(websocket.resolve_unknown_desired_state(Some(name)).await);
let expected = if stored_preference == Some(Some(true)) {
RemoteControlDesiredState::Enabled {
persistence_preference: Some(true),
}
} else {
RemoteControlDesiredState::Disabled
};
assert_eq!(*desired_state_tx.borrow(), expected, "case {name}");
}
}
#[tokio::test]
async fn explicit_disabled_start_ignores_persisted_enable() {
let codex_home = TempDir::new().expect("temp dir should create");
let state_db = remote_control_state_runtime(&codex_home).await;
let remote_control_target = normalize_remote_control_url(TEST_REMOTE_CONTROL_URL)
.expect("remote control target should normalize");
let enrollment = RemoteControlEnrollmentRecord {
websocket_url: remote_control_target.websocket_url,
account_id: "account_id".to_string(),
app_server_client_name: None,
server_id: "server-id".to_string(),
environment_id: "environment-id".to_string(),
server_name: "server-name".to_string(),
remote_control_enabled: Some(true),
};
state_db
.upsert_remote_control_enrollment(&enrollment)
.await
.expect("enrollment should persist");
let (transport_event_tx, _transport_event_rx) = mpsc::channel(CHANNEL_CAPACITY);
let shutdown_token = CancellationToken::new();
let (remote_task, remote_handle) = start_remote_control(
RemoteControlStartConfig {
remote_control_url: TEST_REMOTE_CONTROL_URL.to_string(),
installation_id: TEST_INSTALLATION_ID.to_string(),
},
Some(state_db.clone()),
remote_control_auth_manager(),
transport_event_tx,
shutdown_token.clone(),
/*app_server_client_name_rx*/ None,
RemoteControlStartupMode::DisabledEphemeral,
)
.await
.expect("remote control should start disabled");
assert_eq!(
*remote_handle.desired_state_tx.borrow(),
RemoteControlDesiredState::Disabled
);
assert_eq!(
state_db
.get_remote_control_enrollment(
&enrollment.websocket_url,
&enrollment.account_id,
/*app_server_client_name*/ None,
)
.await
.expect("enrollment should load"),
Some(enrollment)
);
shutdown_token.cancel();
remote_task.await.expect("remote control task should join");
}
fn remote_control_url_for_listener(listener: &TcpListener) -> String {
let addr = listener
.local_addr()
@@ -141,7 +275,10 @@ fn remote_control_handle_with_current_enrollment(
remote_control_url: &str,
auth_manager: Arc<AuthManager>,
) -> RemoteControlHandle {
let (enabled_tx, _enabled_rx) = watch::channel(/*init*/ true);
let (desired_state_tx, _desired_state_rx) =
watch::channel(RemoteControlDesiredState::Enabled {
persistence_preference: None,
});
let (status_tx, _status_rx) = watch::channel(RemoteControlStatusChangedNotification {
status: RemoteControlConnectionStatus::Connecting,
server_name: test_server_name(),
@@ -165,9 +302,10 @@ fn remote_control_handle_with_current_enrollment(
},
)));
RemoteControlHandle {
enabled_tx: Arc::new(enabled_tx),
desired_state_tx: Arc::new(desired_state_tx),
desired_state_rpc_lock: Arc::new(Semaphore::new(1)),
desired_state_persistence_lock: Arc::new(Semaphore::new(1)),
status_tx: Arc::new(status_tx),
state_db_available: true,
state_db: None,
remote_control_url: remote_control_url.to_string(),
current_enrollment,
@@ -177,6 +315,44 @@ fn remote_control_handle_with_current_enrollment(
}
}
#[tokio::test]
async fn ephemeral_enable_preserves_durable_preference() {
let codex_home = TempDir::new().expect("temp dir should create");
let mut remote_handle = remote_control_handle_with_current_enrollment(
TEST_REMOTE_CONTROL_URL,
remote_control_auth_manager(),
);
remote_handle.state_db = Some(remote_control_state_runtime(&codex_home).await);
remote_handle
.desired_state_tx
.send_replace(RemoteControlDesiredState::Enabled {
persistence_preference: Some(true),
});
remote_handle
.enable_ephemeral()
.expect("ephemeral enable should succeed");
assert_eq!(
*remote_handle.desired_state_tx.borrow(),
RemoteControlDesiredState::Enabled {
persistence_preference: Some(true),
}
);
remote_handle
.desired_state_tx
.send_replace(RemoteControlDesiredState::Disabled);
remote_handle
.enable_ephemeral()
.expect("ephemeral enable should succeed");
assert_eq!(
*remote_handle.desired_state_tx.borrow(),
RemoteControlDesiredState::Enabled {
persistence_preference: None,
}
);
}
fn remote_control_server_token_response(
server_id: &str,
environment_id: &str,
@@ -242,7 +418,10 @@ async fn remote_control_transport_manages_virtual_clients_and_routes_messages()
.await
.expect("listener should bind");
let remote_control_url = remote_control_url_for_listener(&listener);
let remote_control_target = normalize_remote_control_url(&remote_control_url)
.expect("remote control target should normalize");
let codex_home = TempDir::new().expect("temp dir should create");
let state_db = remote_control_state_runtime(&codex_home).await;
let (transport_event_tx, mut transport_event_rx) =
mpsc::channel::<TransportEvent>(CHANNEL_CAPACITY);
let shutdown_token = CancellationToken::new();
@@ -251,12 +430,12 @@ async fn remote_control_transport_manages_virtual_clients_and_routes_messages()
remote_control_url,
installation_id: TEST_INSTALLATION_ID.to_string(),
},
Some(remote_control_state_runtime(&codex_home).await),
Some(state_db.clone()),
remote_control_auth_manager(),
transport_event_tx,
shutdown_token.clone(),
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
RemoteControlStartupMode::EnabledEphemeral,
)
.await
.expect("remote control should start");
@@ -276,6 +455,16 @@ async fn remote_control_transport_manages_virtual_clients_and_routes_messages()
)
.await;
let mut websocket = accept_remote_control_connection(&listener).await;
let enrollment = state_db
.get_remote_control_enrollment(
&remote_control_target.websocket_url,
"account_id",
/*app_server_client_name*/ None,
)
.await
.expect("new enrollment should load")
.expect("new enrollment should exist");
assert_eq!(enrollment.remote_control_enabled, None);
expect_remote_control_status(
&mut status_rx,
/*expected_status*/ None,
@@ -539,7 +728,7 @@ async fn remote_control_transport_reconnects_after_disconnect() {
transport_event_tx,
shutdown_token.clone(),
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
RemoteControlStartupMode::EnabledEphemeral,
)
.await
.expect("remote control should start");
@@ -640,7 +829,7 @@ async fn remote_control_transport_refreshes_server_token_after_websocket_unautho
transport_event_tx,
shutdown_token.clone(),
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
RemoteControlStartupMode::EnabledEphemeral,
)
.await
.expect("remote control should start");
@@ -720,7 +909,7 @@ async fn remote_control_start_allows_remote_control_invalid_url_when_disabled()
transport_event_tx,
shutdown_token.clone(),
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ false,
RemoteControlStartupMode::ResolvePersisted,
)
.await
.expect("disabled remote control should not validate the URL at startup");
@@ -759,7 +948,7 @@ async fn remote_control_start_allows_missing_auth_when_enabled() {
transport_event_tx,
shutdown_token.clone(),
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
RemoteControlStartupMode::EnabledEphemeral,
)
.await
.expect("remote control should start before ChatGPT auth is available");
@@ -794,7 +983,7 @@ async fn remote_control_start_reports_missing_state_db_as_disabled_when_enabled(
transport_event_tx,
shutdown_token.clone(),
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
RemoteControlStartupMode::EnabledEphemeral,
)
.await
.expect("remote control should start disabled without sqlite state db");
@@ -814,7 +1003,9 @@ async fn remote_control_start_reports_missing_state_db_as_disabled_when_enabled(
.expect_err("remote control should not connect without sqlite state db");
assert_eq!(
remote_handle.enable().expect_err("enable should fail"),
remote_handle
.enable_ephemeral()
.expect_err("enable should fail"),
super::RemoteControlUnavailable
);
timeout(Duration::from_millis(100), listener.accept())
@@ -851,7 +1042,7 @@ async fn remote_control_handle_enable_disable_stops_and_restarts_connections() {
transport_event_tx,
shutdown_token.clone(),
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
RemoteControlStartupMode::EnabledEphemeral,
)
.await
.expect("remote control should start");
@@ -884,7 +1075,10 @@ async fn remote_control_handle_enable_disable_stops_and_restarts_connections() {
.await;
assert_eq!(
remote_handle.disable(),
remote_handle
.disable(Some("rpc-client"))
.await
.expect("disable should succeed"),
RemoteControlStatusChangedNotification {
status: RemoteControlConnectionStatus::Disabled,
server_name: test_server_name(),
@@ -910,12 +1104,15 @@ async fn remote_control_handle_enable_disable_stops_and_restarts_connections() {
.expect_err("disabled remote control should not reconnect");
assert_eq!(
remote_handle.enable().expect("enable should succeed"),
remote_handle
.enable(Some("rpc-client"))
.await
.expect("enable should succeed"),
RemoteControlStatusChangedNotification {
status: RemoteControlConnectionStatus::Connecting,
server_name: test_server_name(),
installation_id: TEST_INSTALLATION_ID.to_string(),
environment_id: None,
environment_id: Some("env_test".to_string()),
}
);
expect_remote_control_status_snapshot(
@@ -924,7 +1121,7 @@ async fn remote_control_handle_enable_disable_stops_and_restarts_connections() {
status: RemoteControlConnectionStatus::Connecting,
server_name: test_server_name(),
installation_id: TEST_INSTALLATION_ID.to_string(),
environment_id: None,
environment_id: Some("env_test".to_string()),
},
)
.await;
@@ -964,7 +1161,7 @@ async fn remote_control_transport_clears_outgoing_buffer_when_backend_acks() {
transport_event_tx,
shutdown_token.clone(),
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
RemoteControlStartupMode::EnabledEphemeral,
)
.await
.expect("remote control should start");
@@ -1146,7 +1343,7 @@ async fn remote_control_http_mode_enrolls_before_connecting() {
transport_event_tx,
shutdown_token.clone(),
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
RemoteControlStartupMode::EnabledEphemeral,
)
.await
.expect("remote control should start");
@@ -1376,6 +1573,7 @@ async fn remote_control_http_mode_refreshes_persisted_enrollment_before_connecti
"account_id",
/*app_server_client_name*/ None,
Some(&persisted_enrollment),
/*remote_control_enabled*/ None,
)
.await
.expect("persisted enrollment should save");
@@ -1393,7 +1591,7 @@ async fn remote_control_http_mode_refreshes_persisted_enrollment_before_connecti
transport_event_tx,
shutdown_token.clone(),
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
RemoteControlStartupMode::EnabledEphemeral,
)
.await
.expect("remote control should start");
@@ -1482,6 +1680,7 @@ async fn remote_control_stdio_mode_waits_for_client_name_before_connecting() {
"account_id",
Some(app_server_client_name),
Some(&persisted_enrollment),
/*remote_control_enabled*/ None,
)
.await
.expect("persisted enrollment should save");
@@ -1500,7 +1699,7 @@ async fn remote_control_stdio_mode_waits_for_client_name_before_connecting() {
transport_event_tx,
shutdown_token.clone(),
Some(app_server_client_name_rx),
/*initial_enabled*/ true,
RemoteControlStartupMode::EnabledEphemeral,
)
.await
.expect("remote control should start");
@@ -1581,7 +1780,7 @@ async fn remote_control_waits_for_account_id_before_enrolling() {
transport_event_tx,
shutdown_token.clone(),
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
RemoteControlStartupMode::EnabledEphemeral,
)
.await
.expect("remote control should start before account id is available");
@@ -1625,6 +1824,123 @@ async fn remote_control_waits_for_account_id_before_enrolling() {
let _ = remote_task.await;
}
#[tokio::test]
async fn persisted_enable_does_not_follow_auth_to_an_account_without_a_preference() {
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_a")),
AuthCredentialsStoreMode::File,
)
.expect("account A auth should save");
let state_db = remote_control_state_runtime(&codex_home).await;
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_control_target =
normalize_remote_control_url(&remote_control_url).expect("target should parse");
let enrollment = RemoteControlEnrollment {
remote_control_target: remote_control_target.clone(),
account_id: "account_a".to_string(),
environment_id: "env_a".to_string(),
server_id: "srv_e_a".to_string(),
server_name: "server-a".to_string(),
remote_control_token: None,
expires_at: None,
};
update_persisted_remote_control_enrollment(
Some(state_db.as_ref()),
&remote_control_target,
"account_a",
/*app_server_client_name*/ None,
Some(&enrollment),
/*remote_control_enabled*/ Some(true),
)
.await
.expect("account A enrollment should save");
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(state_db.clone()),
auth_manager.clone(),
transport_event_tx,
shutdown_token.clone(),
/*app_server_client_name_rx*/ None,
RemoteControlStartupMode::ResolvePersisted,
)
.await
.expect("remote control should start");
let refresh_request = accept_http_request(&listener).await;
assert_eq!(
refresh_request.request_line,
"POST /backend-api/wham/remote/control/server/refresh HTTP/1.1"
);
respond_with_json(
refresh_request.stream,
remote_control_server_token_response(
&enrollment.server_id,
&enrollment.environment_id,
TEST_REFRESHED_REMOTE_CONTROL_SERVER_TOKEN,
),
)
.await;
let (_handshake_request, mut websocket) =
accept_remote_control_backend_connection(&listener).await;
save_auth(
codex_home.path(),
&remote_control_auth_dot_json(Some("account_b")),
AuthCredentialsStoreMode::File,
)
.expect("account B auth should save");
auth_manager.reload().await;
websocket
.close(None)
.await
.expect("backend websocket should close");
let mut desired_state_rx = remote_handle.desired_state_tx.subscribe();
timeout(
Duration::from_secs(1),
desired_state_rx.wait_for(|state| *state == RemoteControlDesiredState::Disabled),
)
.await
.expect("account B missing preference should disable remote control")
.expect("desired state channel should stay open");
timeout(Duration::from_millis(100), listener.accept())
.await
.expect_err("disabled account B should not enroll");
assert_eq!(
state_db
.get_remote_control_enrollment(
&remote_control_target.websocket_url,
"account_b",
/*app_server_client_name*/ None,
)
.await
.expect("account B enrollment should load"),
None
);
shutdown_token.cancel();
let _ = remote_task.await;
}
#[tokio::test]
async fn remote_control_http_mode_reenrolls_when_refresh_reports_stale_enrollment() {
let listener = TcpListener::bind("127.0.0.1:0")
@@ -1660,6 +1976,7 @@ async fn remote_control_http_mode_reenrolls_when_refresh_reports_stale_enrollmen
"account_id",
/*app_server_client_name*/ None,
Some(&stale_enrollment),
/*remote_control_enabled*/ Some(true),
)
.await
.expect("stale enrollment should save");
@@ -1677,7 +1994,7 @@ async fn remote_control_http_mode_reenrolls_when_refresh_reports_stale_enrollmen
transport_event_tx,
shutdown_token.clone(),
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
RemoteControlStartupMode::ResolvePersisted,
)
.await
.expect("remote control should start");
@@ -1695,12 +2012,6 @@ async fn remote_control_http_mode_reenrolls_when_refresh_reports_stale_enrollmen
)
.await;
respond_with_status(refresh_request.stream, "404 Not Found", "").await;
expect_remote_control_status(
&mut status_rx,
/*expected_status*/ None,
/*expected_environment_id*/ None,
)
.await;
let enroll_request = accept_http_request(&listener).await;
assert_eq!(
@@ -1729,15 +2040,23 @@ async fn remote_control_http_mode_reenrolls_when_refresh_reports_stale_enrollmen
Some(&refreshed_enrollment.server_id)
);
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)
state_db
.get_remote_control_enrollment(
&remote_control_target.websocket_url,
"account_id",
/*app_server_client_name*/ None,
)
.await
.expect("refreshed enrollment should load"),
Some(RemoteControlEnrollmentRecord {
websocket_url: remote_control_target.websocket_url.clone(),
account_id: "account_id".to_string(),
app_server_client_name: None,
server_id: refreshed_enrollment.server_id.clone(),
environment_id: refreshed_enrollment.environment_id.clone(),
server_name: refreshed_enrollment.server_name.clone(),
remote_control_enabled: Some(true),
})
);
shutdown_token.cancel();
@@ -1779,6 +2098,7 @@ async fn remote_control_http_mode_reenrolls_after_explicit_missing_server_404()
"account_id",
/*app_server_client_name*/ None,
Some(&stale_enrollment),
/*remote_control_enabled*/ Some(true),
)
.await
.expect("stale enrollment should save");
@@ -1796,7 +2116,7 @@ async fn remote_control_http_mode_reenrolls_after_explicit_missing_server_404()
transport_event_tx,
shutdown_token.clone(),
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
RemoteControlStartupMode::ResolvePersisted,
)
.await
.expect("remote control should start");
@@ -1838,12 +2158,6 @@ async fn remote_control_http_mode_reenrolls_after_explicit_missing_server_404()
&json!({"detail": "Remote app server not found"}).to_string(),
)
.await;
expect_remote_control_status(
&mut status_rx,
/*expected_status*/ None,
/*expected_environment_id*/ None,
)
.await;
let enroll_request = accept_http_request(&listener).await;
assert_eq!(
@@ -1872,15 +2186,125 @@ async fn remote_control_http_mode_reenrolls_after_explicit_missing_server_404()
Some(&refreshed_enrollment.server_id)
);
assert_eq!(
load_persisted_remote_control_enrollment(
Some(state_db.as_ref()),
&remote_control_target,
"account_id",
/*app_server_client_name*/ None,
)
state_db
.get_remote_control_enrollment(
&remote_control_target.websocket_url,
"account_id",
/*app_server_client_name*/ None,
)
.await
.expect("refreshed enrollment should load"),
Some(RemoteControlEnrollmentRecord {
websocket_url: remote_control_target.websocket_url.clone(),
account_id: "account_id".to_string(),
app_server_client_name: None,
server_id: refreshed_enrollment.server_id.clone(),
environment_id: refreshed_enrollment.environment_id.clone(),
server_name: refreshed_enrollment.server_name.clone(),
remote_control_enabled: Some(true),
})
);
shutdown_token.cancel();
let _ = remote_task.await;
}
#[tokio::test]
async fn remote_control_http_mode_preserves_stale_enrollment_when_reenrollment_fails() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("refreshed enrollment should load"),
Some(refreshed_enrollment)
.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 remote_control_target =
normalize_remote_control_url(&remote_control_url).expect("target should parse");
let stale_enrollment = RemoteControlEnrollment {
remote_control_target: remote_control_target.clone(),
account_id: "account_id".to_string(),
environment_id: "env_stale".to_string(),
server_id: "srv_e_stale".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),
/*remote_control_enabled*/ Some(true),
)
.await
.expect("stale enrollment should save");
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(state_db.clone()),
remote_control_auth_manager_with_home(&codex_home),
transport_event_tx,
shutdown_token.clone(),
/*app_server_client_name_rx*/ None,
RemoteControlStartupMode::ResolvePersisted,
)
.await
.expect("remote control should start");
let refresh_request = accept_http_request(&listener).await;
assert_eq!(
refresh_request.request_line,
"POST /backend-api/wham/remote/control/server/refresh HTTP/1.1"
);
respond_with_status(refresh_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_status(enroll_request.stream, "500 Internal Server Error", "failed").await;
let retry_refresh_request = accept_http_request(&listener).await;
assert_eq!(
retry_refresh_request.request_line,
"POST /backend-api/wham/remote/control/server/refresh HTTP/1.1"
);
respond_with_status(
retry_refresh_request.stream,
"500 Internal Server Error",
"failed",
)
.await;
assert_eq!(
*remote_handle.current_enrollment.lock().await,
Some(stale_enrollment.clone())
);
assert_eq!(
state_db
.get_remote_control_enrollment(
&remote_control_target.websocket_url,
"account_id",
/*app_server_client_name*/ None,
)
.await
.expect("stale enrollment should load"),
Some(RemoteControlEnrollmentRecord {
websocket_url: remote_control_target.websocket_url,
account_id: "account_id".to_string(),
app_server_client_name: None,
server_id: stale_enrollment.server_id,
environment_id: stale_enrollment.environment_id,
server_name: stale_enrollment.server_name,
remote_control_enabled: Some(true),
})
);
shutdown_token.cancel();
@@ -1912,6 +2336,7 @@ async fn remote_control_http_mode_preserves_enrollment_after_generic_websocket_4
"account_id",
/*app_server_client_name*/ None,
Some(&stale_enrollment),
/*remote_control_enabled*/ None,
)
.await
.expect("stale enrollment should save");
@@ -1929,7 +2354,7 @@ async fn remote_control_http_mode_preserves_enrollment_after_generic_websocket_4
transport_event_tx,
shutdown_token.clone(),
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
RemoteControlStartupMode::EnabledEphemeral,
)
.await
.expect("remote control should start");
@@ -13,7 +13,7 @@ fn client_management_handle(
remote_control_url: String,
auth_manager: Arc<AuthManager>,
) -> RemoteControlHandle {
let (enabled_tx, _enabled_rx) = watch::channel(/*init*/ false);
let desired_state_tx = watch::channel(RemoteControlDesiredState::Disabled).0;
let (status_tx, _status_rx) = watch::channel(RemoteControlStatusChangedNotification {
status: RemoteControlConnectionStatus::Disabled,
server_name: test_server_name(),
@@ -21,9 +21,10 @@ fn client_management_handle(
environment_id: None,
});
RemoteControlHandle {
enabled_tx: Arc::new(enabled_tx),
desired_state_tx: Arc::new(desired_state_tx),
desired_state_rpc_lock: Arc::new(Semaphore::new(1)),
desired_state_persistence_lock: Arc::new(Semaphore::new(1)),
status_tx: Arc::new(status_tx),
state_db_available: false,
state_db: None,
remote_control_url,
current_enrollment: Arc::new(RemoteControlEnrollmentState::new(/*enrollment*/ None)),
@@ -645,7 +645,9 @@ async fn remote_control_handle_disable_keeps_current_enrollment() {
remote_control_auth_manager(),
);
remote_handle.disable();
remote_handle
.desired_state_tx
.send_replace(RemoteControlDesiredState::Disabled);
assert!(
remote_handle.current_enrollment.lock().await.is_some(),
"disabled remote control should keep the selected pairing server"
@@ -665,7 +667,6 @@ async fn remote_control_handle_reenrolls_after_stale_pairing_enrollment() {
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()
@@ -688,9 +689,15 @@ async fn remote_control_handle_reenrolls_after_stale_pairing_enrollment() {
"account_id",
/*app_server_client_name*/ None,
Some(&stale_enrollment),
/*remote_control_enabled*/ Some(true),
)
.await
.expect("stale enrollment should save");
remote_handle
.desired_state_tx
.send_replace(RemoteControlDesiredState::Enabled {
persistence_preference: Some(true),
});
let server_refreshed_enrollment = refreshed_enrollment.clone();
let server_task = tokio::spawn(async move {
let stale_pairing_request = accept_http_request(&listener).await;
@@ -761,15 +768,17 @@ async fn remote_control_handle_reenrolls_after_stale_pairing_enrollment() {
}
);
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)
state_db
.get_remote_control_enrollment(
&remote_control_target.websocket_url,
"account_id",
/*app_server_client_name*/ None,
)
.await
.expect("refreshed enrollment should load")
.expect("refreshed enrollment should exist")
.remote_control_enabled,
Some(true)
);
}
@@ -1,5 +1,9 @@
use super::CurrentRemoteControlEnrollment;
use super::RemoteControlEnrollmentSelection;
use super::RemoteControlPairingPersistenceKey;
use super::desired_state::RemoteControlDesiredState;
use super::desired_state::acquire_persistence_lock;
use super::desired_state::desired_state_from_persisted_enrollment;
use super::protocol::ClientEnvelope;
use super::protocol::ClientEvent;
use super::protocol::ClientId;
@@ -45,6 +49,7 @@ use std::io::ErrorKind;
use std::sync::Arc;
use tokio::net::TcpStream;
use tokio::sync::Mutex;
use tokio::sync::Semaphore;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::sync::watch;
@@ -264,7 +269,9 @@ pub(crate) struct RemoteControlWebsocket {
state: Arc<Mutex<WebsocketState>>,
server_event_rx: Arc<Mutex<mpsc::Receiver<super::QueuedServerEnvelope>>>,
used_rx: watch::Receiver<usize>,
enabled_rx: watch::Receiver<bool>,
desired_state_tx: Arc<watch::Sender<RemoteControlDesiredState>>,
desired_state_rx: watch::Receiver<RemoteControlDesiredState>,
desired_state_persistence_lock: Arc<Semaphore>,
}
pub(crate) struct RemoteControlWebsocketConfig {
@@ -280,6 +287,11 @@ pub(super) struct RemoteControlAuthContext<'a> {
auth_change_rx: &'a mut watch::Receiver<u64>,
}
struct RemoteControlEnrollmentAuthContext<'a, 'b> {
auth: &'a RemoteControlConnectionAuth,
recovery: &'a mut RemoteControlAuthContext<'b>,
}
enum ConnectOutcome {
Connected(Box<WebSocketStream<MaybeTlsStream<TcpStream>>>),
Disabled,
@@ -299,6 +311,7 @@ pub(super) struct RemoteControlChannels {
pub(super) status_publisher: RemoteControlStatusPublisher,
pub(super) current_enrollment: CurrentRemoteControlEnrollment,
pub(super) pairing_persistence_key: RemoteControlPairingPersistenceKey,
pub(super) desired_state_persistence_lock: Arc<Semaphore>,
}
#[derive(Clone)]
@@ -341,7 +354,7 @@ impl RemoteControlStatusPublisher {
}
}
fn publish_environment_id(&self, environment_id: Option<String>) {
pub(super) fn publish_environment_id(&self, environment_id: Option<String>) {
let mut status_change = None;
self.tx.send_if_modified(|status| {
if status.status == RemoteControlConnectionStatus::Disabled {
@@ -380,6 +393,8 @@ pub(super) struct RemoteControlConnectOptions<'a> {
server_name: &'a str,
subscribe_cursor: Option<&'a str>,
app_server_client_name: Option<&'a str>,
desired_state_tx: &'a watch::Sender<RemoteControlDesiredState>,
desired_state_persistence_lock: &'a Semaphore,
}
impl RemoteControlWebsocket {
@@ -389,7 +404,7 @@ impl RemoteControlWebsocket {
auth_manager: Arc<AuthManager>,
channels: RemoteControlChannels,
shutdown_token: CancellationToken,
enabled_rx: watch::Receiver<bool>,
desired_state_tx: Arc<watch::Sender<RemoteControlDesiredState>>,
) -> Self {
let shutdown_token = shutdown_token.child_token();
let (server_event_tx, server_event_rx) = mpsc::channel(super::CHANNEL_CAPACITY);
@@ -402,6 +417,7 @@ impl RemoteControlWebsocket {
let auth_recovery = auth_manager.unauthorized_recovery();
let auth_change_rx = auth_manager.auth_change_receiver();
let desired_state_rx = desired_state_tx.subscribe();
Self {
remote_control_url: config.remote_control_url,
installation_id: config.installation_id,
@@ -426,7 +442,9 @@ impl RemoteControlWebsocket {
})),
server_event_rx: Arc::new(Mutex::new(server_event_rx)),
used_rx,
enabled_rx,
desired_state_tx,
desired_state_rx,
desired_state_persistence_lock: channels.desired_state_persistence_lock,
}
}
@@ -442,7 +460,7 @@ impl RemoteControlWebsocket {
remote_control_url = %self.remote_control_url,
installation_id = %self.installation_id,
server_name = %self.server_name,
initial_enabled = *self.enabled_rx.borrow(),
initial_desired_state = ?*self.desired_state_rx.borrow(),
"app-server remote control websocket loop started"
);
let app_server_client_name = match self
@@ -464,6 +482,16 @@ impl RemoteControlWebsocket {
};
self.pairing_persistence_key
.send_replace(app_server_client_name.clone());
if matches!(
*self.desired_state_rx.borrow(),
RemoteControlDesiredState::Unknown
) && !self
.resolve_unknown_desired_state(app_server_client_name.as_deref())
.await
{
self.client_tracker.lock().await.shutdown().await;
return;
}
loop {
if !self.wait_until_enabled().await {
@@ -513,7 +541,7 @@ impl RemoteControlWebsocket {
connection_end_reason = ?connection_end_reason,
current_status = ?status.status,
environment_id = ?status.environment_id,
enabled = *self.enabled_rx.borrow(),
desired_state = ?*self.desired_state_rx.borrow(),
"app-server remote control websocket connection cycle ended"
);
}
@@ -546,10 +574,96 @@ impl RemoteControlWebsocket {
}
}
pub(super) async fn resolve_unknown_desired_state(
&mut self,
app_server_client_name: Option<&str>,
) -> bool {
let remote_control_target = match super::protocol::normalize_remote_control_url(
&self.remote_control_url,
) {
Ok(remote_control_target) => remote_control_target,
Err(err) => {
warn!(
"remote control preference cannot be resolved because the URL is invalid: {err}"
);
self.transition_unknown_to(RemoteControlDesiredState::Disabled);
return true;
}
};
self.remote_control_target = Some(remote_control_target.clone());
let Some(state_db) = self.state_db.clone() else {
self.transition_unknown_to(RemoteControlDesiredState::Disabled);
return true;
};
loop {
if !matches!(
*self.desired_state_rx.borrow(),
RemoteControlDesiredState::Unknown
) {
return true;
}
let auth = match load_remote_control_auth(&self.auth_manager).await {
Ok(auth) => auth,
Err(err) => {
info!(
error = %err,
"waiting to resolve remote control preference until authentication is available"
);
if !self.wait_for_preference_resolution_retry().await {
return false;
}
continue;
}
};
let enrollment = match state_db
.get_remote_control_enrollment(
&remote_control_target.websocket_url,
&auth.account_id,
app_server_client_name,
)
.await
{
Ok(enrollment) => enrollment,
Err(err) => {
warn!(
error = %err,
"failed to resolve persisted remote control preference; retrying"
);
if !self.wait_for_preference_resolution_retry().await {
return false;
}
continue;
}
};
let desired_state = desired_state_from_persisted_enrollment(enrollment);
self.transition_unknown_to(desired_state);
return true;
}
}
fn transition_unknown_to(&self, desired_state: RemoteControlDesiredState) {
self.desired_state_tx.send_if_modified(|state| {
if !matches!(*state, RemoteControlDesiredState::Unknown) {
return false;
}
*state = desired_state;
true
});
}
async fn wait_for_preference_resolution_retry(&mut self) -> bool {
tokio::select! {
_ = self.shutdown_token.cancelled() => false,
changed = self.desired_state_rx.changed() => changed.is_ok(),
_ = tokio::time::sleep(REMOTE_CONTROL_ACCOUNT_ID_RETRY_INTERVAL) => true,
}
}
async fn wait_until_enabled(&mut self) -> bool {
tokio::select! {
_ = self.shutdown_token.cancelled() => false,
enabled = self.enabled_rx.wait_for(|enabled| *enabled) => enabled.is_ok(),
desired_state = self.desired_state_rx.wait_for(|state| state.is_enabled()) => desired_state.is_ok(),
}
}
@@ -573,7 +687,7 @@ impl RemoteControlWebsocket {
warn!("remote control is enabled but the URL is invalid: {err}");
tokio::select! {
_ = shutdown_token.cancelled() => return ConnectOutcome::Shutdown,
changed = self.enabled_rx.wait_for(|enabled| !*enabled) => {
changed = self.desired_state_rx.wait_for(|state| !state.is_enabled()) => {
if changed.is_err() {
return ConnectOutcome::Shutdown;
}
@@ -604,15 +718,18 @@ impl RemoteControlWebsocket {
server_name: &self.server_name,
subscribe_cursor: subscribe_cursor.as_deref(),
app_server_client_name,
desired_state_tx: &self.desired_state_tx,
desired_state_persistence_lock: &self.desired_state_persistence_lock,
};
let auth_context = RemoteControlAuthContext {
auth_manager: &self.auth_manager,
auth_recovery: &mut self.auth_recovery,
auth_change_rx: &mut self.auth_change_rx,
};
let mut disabled_rx = self.desired_state_rx.clone();
let connect_result = tokio::select! {
_ = shutdown_token.cancelled() => return ConnectOutcome::Shutdown,
changed = self.enabled_rx.wait_for(|enabled| !*enabled) => {
changed = disabled_rx.wait_for(|state| !state.is_enabled()) => {
if changed.is_err() {
return ConnectOutcome::Shutdown;
}
@@ -633,7 +750,7 @@ impl RemoteControlWebsocket {
match connect_result {
Ok((websocket_connection, response)) => {
if !*self.enabled_rx.borrow() {
if !self.desired_state_rx.borrow().is_enabled() {
return ConnectOutcome::Disabled;
}
self.reconnect_attempt = 0;
@@ -654,7 +771,7 @@ impl RemoteControlWebsocket {
return ConnectOutcome::Connected(Box::new(websocket_connection));
}
Err(err) => {
if !*self.enabled_rx.borrow() {
if !self.desired_state_rx.borrow().is_enabled() {
return ConnectOutcome::Disabled;
}
let reconnect_delay = if err.kind() == ErrorKind::WouldBlock {
@@ -691,7 +808,7 @@ impl RemoteControlWebsocket {
};
tokio::select! {
_ = shutdown_token.cancelled() => return ConnectOutcome::Shutdown,
changed = self.enabled_rx.wait_for(|enabled| !*enabled) => {
changed = self.desired_state_rx.wait_for(|state| !state.is_enabled()) => {
if changed.is_err() {
return ConnectOutcome::Shutdown;
}
@@ -736,10 +853,10 @@ impl RemoteControlWebsocket {
shutdown_token.clone(),
));
let mut enabled_rx = self.enabled_rx.clone();
let mut desired_state_rx = self.desired_state_rx.clone();
let connection_end_reason = tokio::select! {
_ = shutdown_token.cancelled() => ConnectionEndReason::Shutdown,
changed = enabled_rx.wait_for(|enabled| !*enabled) => {
changed = desired_state_rx.wait_for(|state| !state.is_enabled()) => {
if changed.is_ok() {
self.status_publisher
.publish_status(RemoteControlConnectionStatus::Disabled);
@@ -1253,22 +1370,25 @@ pub(super) async fn connect_remote_control_websocket(
if websocket_response_reports_missing_remote_app_server(response) =>
{
info!(
"remote control websocket returned HTTP 404; clearing stale enrollment before re-enrolling: websocket_url={}, account_id={}, server_id={}, environment_id={}",
"remote control websocket returned HTTP 404; replacing stale enrollment: 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(
replace_remote_control_enrollment_if_matches(
state_db,
remote_control_target,
&auth.account_id,
connect_options.app_server_client_name,
RemoteControlEnrollmentAuthContext {
auth: &auth,
recovery: &mut auth_context,
},
current_enrollment,
&enrollment,
connect_options,
status_publisher,
)
.await;
.await?;
}
tungstenite::Error::Http(response) if response.status().as_u16() == 404 => {
let response_body = response
@@ -1337,6 +1457,14 @@ async fn prepare_remote_control_enrollment(
};
let enrollment_account_id = enrollment.as_ref().map(|enrollment| &enrollment.account_id);
if enrollment_account_id.is_some_and(|account_id| account_id != &auth.account_id) {
resolve_desired_state_after_account_change(
state_db,
remote_control_target,
auth_context.auth_manager,
&auth.account_id,
connect_options,
)
.await?;
info!(
"clearing in-memory remote control enrollment because account id changed: websocket_url={}, previous_account_id={:?}, current_account_id={:?}",
remote_control_target.websocket_url,
@@ -1347,6 +1475,12 @@ async fn prepare_remote_control_enrollment(
);
*enrollment = None;
status_publisher.publish_environment_id(/*environment_id*/ None);
if !connect_options.desired_state_tx.borrow().is_enabled() {
return Err(io::Error::new(
ErrorKind::Interrupted,
"remote control disabled after account changed",
));
}
}
if let Some(enrollment) = enrollment.as_mut() {
enrollment.remote_control_target = remote_control_target.clone();
@@ -1373,14 +1507,17 @@ async fn prepare_remote_control_enrollment(
});
}
enroll_remote_control_server_if_missing(
enroll_and_persist_remote_control_server(
remote_control_target,
state_db,
&auth,
auth_context,
RemoteControlEnrollmentAuthContext {
auth: &auth,
recovery: auth_context,
},
enrollment,
connect_options,
status_publisher,
RemoteControlEnrollmentSelection::ReuseOrCreate,
)
.await?;
@@ -1412,26 +1549,20 @@ async fn prepare_remote_control_enrollment(
Ok(()) => {}
Err(err) if err.kind() == ErrorKind::NotFound => {
info!(
"remote control server refresh returned HTTP 404; clearing stale enrollment before re-enrolling: websocket_url={}, account_id={}, server_id={}, environment_id={}",
"remote control server refresh returned HTTP 404; replacing stale enrollment: websocket_url={}, account_id={}, server_id={}, environment_id={}",
remote_control_target.websocket_url, auth.account_id, server_id, environment_id
);
clear_remote_control_enrollment(
state_db,
remote_control_target,
&auth.account_id,
connect_options.app_server_client_name,
enrollment,
status_publisher,
)
.await;
enroll_remote_control_server_if_missing(
enroll_and_persist_remote_control_server(
remote_control_target,
state_db,
&auth,
auth_context,
RemoteControlEnrollmentAuthContext {
auth: &auth,
recovery: auth_context,
},
enrollment,
connect_options,
status_publisher,
RemoteControlEnrollmentSelection::ReplaceExisting,
)
.await?;
}
@@ -1454,6 +1585,51 @@ async fn prepare_remote_control_enrollment(
Ok(auth)
}
async fn resolve_desired_state_after_account_change(
state_db: &StateRuntime,
remote_control_target: &RemoteControlTarget,
auth_manager: &Arc<AuthManager>,
account_id: &str,
connect_options: RemoteControlConnectOptions<'_>,
) -> io::Result<()> {
let durable_enabled = RemoteControlDesiredState::Enabled {
persistence_preference: Some(true),
};
if *connect_options.desired_state_tx.borrow() != durable_enabled {
return Ok(());
}
let _persistence =
acquire_persistence_lock(connect_options.desired_state_persistence_lock).await;
if *connect_options.desired_state_tx.borrow() != durable_enabled {
return Ok(());
}
let enrollment = state_db
.get_remote_control_enrollment(
&remote_control_target.websocket_url,
account_id,
connect_options.app_server_client_name,
)
.await
.map_err(io::Error::other)?;
let current_auth = load_remote_control_auth(auth_manager).await?;
if current_auth.account_id != account_id {
return Err(io::Error::new(
ErrorKind::WouldBlock,
"remote control account changed while resolving persisted preference",
));
}
let resolved_state = desired_state_from_persisted_enrollment(enrollment);
connect_options.desired_state_tx.send_if_modified(|state| {
if *state != durable_enabled || *state == resolved_state {
return false;
}
*state = resolved_state;
true
});
Ok(())
}
fn websocket_response_reports_missing_remote_app_server(
response: &tungstenite::http::Response<Option<Vec<u8>>>,
) -> bool {
@@ -1466,57 +1642,38 @@ fn websocket_response_reports_missing_remote_app_server(
})
}
async fn clear_remote_control_enrollment(
state_db: &StateRuntime,
remote_control_target: &RemoteControlTarget,
account_id: &str,
app_server_client_name: Option<&str>,
enrollment: &mut Option<RemoteControlEnrollment>,
status_publisher: &RemoteControlStatusPublisher,
) {
if let Err(clear_err) = update_persisted_remote_control_enrollment(
Some(state_db),
remote_control_target,
account_id,
app_server_client_name,
/*enrollment*/ None,
)
.await
{
warn!("failed to clear stale remote control enrollment in sqlite state db: {clear_err}");
}
*enrollment = None;
status_publisher.publish_environment_id(/*environment_id*/ None);
}
async fn clear_remote_control_enrollment_if_matches(
async fn replace_remote_control_enrollment_if_matches(
state_db: Option<&StateRuntime>,
remote_control_target: &RemoteControlTarget,
account_id: &str,
app_server_client_name: Option<&str>,
auth_context: RemoteControlEnrollmentAuthContext<'_, '_>,
current_enrollment: &CurrentRemoteControlEnrollment,
enrollment: &RemoteControlEnrollment,
connect_options: RemoteControlConnectOptions<'_>,
status_publisher: &RemoteControlStatusPublisher,
) {
) -> io::Result<()> {
let Some(state_db) = state_db else {
return;
return Err(io::Error::new(
ErrorKind::NotFound,
"remote control requires sqlite state db",
));
};
let mut current_enrollment = current_enrollment.lock().await;
if !current_enrollment
.as_ref()
.is_some_and(|current| same_remote_control_enrollment(current, enrollment))
{
return;
return Ok(());
}
clear_remote_control_enrollment(
state_db,
enroll_and_persist_remote_control_server(
remote_control_target,
account_id,
app_server_client_name,
state_db,
auth_context,
&mut current_enrollment,
connect_options,
status_publisher,
RemoteControlEnrollmentSelection::ReplaceExisting,
)
.await;
.await
}
async fn clear_remote_control_server_token_if_matches(
@@ -1534,26 +1691,39 @@ async fn clear_remote_control_server_token_if_matches(
Ok(())
}
async fn enroll_remote_control_server_if_missing(
async fn enroll_and_persist_remote_control_server(
remote_control_target: &RemoteControlTarget,
state_db: &StateRuntime,
auth: &RemoteControlConnectionAuth,
auth_context: &mut RemoteControlAuthContext<'_>,
auth_context: RemoteControlEnrollmentAuthContext<'_, '_>,
enrollment: &mut Option<RemoteControlEnrollment>,
connect_options: RemoteControlConnectOptions<'_>,
status_publisher: &RemoteControlStatusPublisher,
selection: RemoteControlEnrollmentSelection,
) -> io::Result<()> {
if enrollment.is_some() {
return Ok(());
match selection {
RemoteControlEnrollmentSelection::ReuseOrCreate => {
if enrollment.is_some() {
return Ok(());
}
}
RemoteControlEnrollmentSelection::ReplaceExisting => {}
}
if !connect_options.desired_state_tx.borrow().is_enabled() {
return Err(io::Error::new(
ErrorKind::Interrupted,
"remote control disabled before enrollment",
));
}
info!(
"creating new remote control enrollment: websocket_url={}, enroll_url={}, account_id={}",
remote_control_target.websocket_url, remote_control_target.enroll_url, auth.account_id
remote_control_target.websocket_url,
remote_control_target.enroll_url,
auth_context.auth.account_id
);
let new_enrollment = match enroll_remote_control_server(
remote_control_target,
auth,
auth_context.auth,
connect_options.installation_id,
connect_options.server_name,
)
@@ -1563,8 +1733,8 @@ async fn enroll_remote_control_server_if_missing(
Err(err)
if err.kind() == ErrorKind::PermissionDenied
&& recover_remote_control_auth(
auth_context.auth_recovery,
auth_context.auth_change_rx,
auth_context.recovery.auth_recovery,
auth_context.recovery.auth_change_rx,
)
.await =>
{
@@ -1574,12 +1744,26 @@ async fn enroll_remote_control_server_if_missing(
}
Err(err) => return Err(err),
};
let _persistence =
acquire_persistence_lock(connect_options.desired_state_persistence_lock).await;
let persistence_preference = match *connect_options.desired_state_tx.borrow() {
RemoteControlDesiredState::Enabled {
persistence_preference,
} => persistence_preference,
RemoteControlDesiredState::Unknown | RemoteControlDesiredState::Disabled => {
return Err(io::Error::new(
ErrorKind::Interrupted,
"remote control disabled during enrollment",
));
}
};
if let Err(err) = update_persisted_remote_control_enrollment(
Some(state_db),
remote_control_target,
&auth.account_id,
&auth_context.auth.account_id,
connect_options.app_server_client_name,
Some(&new_enrollment),
persistence_preference,
)
.await
{
@@ -1758,6 +1942,13 @@ mod tests {
(RemoteControlStatusPublisher::new(status_tx), status_rx)
}
fn enabled_desired_state_sender() -> watch::Sender<RemoteControlDesiredState> {
watch::channel(RemoteControlDesiredState::Enabled {
persistence_preference: None,
})
.0
}
#[test]
fn mark_recovery_auth_change_seen_marks_only_recovery_revision_seen() {
let (auth_change_tx, mut auth_change_rx) = watch::channel(0u64);
@@ -1896,6 +2087,8 @@ mod tests {
server_name: "test-server",
subscribe_cursor: None,
app_server_client_name: None,
desired_state_tx: &enabled_desired_state_sender(),
desired_state_persistence_lock: &Semaphore::new(1),
},
&status_publisher,
)
@@ -1960,6 +2153,8 @@ mod tests {
server_name: "test-server",
subscribe_cursor: None,
app_server_client_name: None,
desired_state_tx: &enabled_desired_state_sender(),
desired_state_persistence_lock: &Semaphore::new(1),
},
&status_publisher,
)
@@ -2042,6 +2237,8 @@ mod tests {
server_name: "test-server",
subscribe_cursor: None,
app_server_client_name: None,
desired_state_tx: &enabled_desired_state_sender(),
desired_state_persistence_lock: &Semaphore::new(1),
},
&status_publisher,
)
@@ -2136,6 +2333,8 @@ mod tests {
server_name: "test-server",
subscribe_cursor: None,
app_server_client_name: None,
desired_state_tx: &enabled_desired_state_sender(),
desired_state_persistence_lock: &Semaphore::new(1),
},
&status_publisher,
)
@@ -2201,6 +2400,8 @@ mod tests {
server_name: "test-server",
subscribe_cursor: None,
app_server_client_name: None,
desired_state_tx: &enabled_desired_state_sender(),
desired_state_persistence_lock: &Semaphore::new(1),
},
&status_publisher,
)
@@ -2251,6 +2452,8 @@ mod tests {
server_name: "test-server",
subscribe_cursor: None,
app_server_client_name: None,
desired_state_tx: &enabled_desired_state_sender(),
desired_state_persistence_lock: &Semaphore::new(1),
},
&status_publisher,
)
@@ -2288,7 +2491,10 @@ mod tests {
drop(transport_event_rx);
let (status_publisher, _status_rx) = remote_control_status_channel();
let shutdown_token = CancellationToken::new();
let (_enabled_tx, enabled_rx) = watch::channel(true);
let (desired_state_tx, _desired_state_rx) =
watch::channel(RemoteControlDesiredState::Enabled {
persistence_preference: None,
});
let websocket_task = tokio::spawn({
let shutdown_token = shutdown_token.clone();
async move {
@@ -2306,9 +2512,10 @@ mod tests {
status_publisher,
current_enrollment: test_current_enrollment(/*enrollment*/ None),
pairing_persistence_key: watch::channel(None).0,
desired_state_persistence_lock: Arc::new(Semaphore::new(1)),
},
shutdown_token,
enabled_rx,
Arc::new(desired_state_tx),
)
.run(/*app_server_client_name_rx*/ None)
.await