mirror of
https://github.com/openai/codex.git
synced 2026-11-13 06:45:07 +00:00
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:
1 parent
be338ee9a2
commit
d61dfeb23a
33 files changed
+2157
-412
No files matched your search
@@ -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(
|
||||
¤t_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(¶ms)?;
|
||||
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
|
||||
|
||||
Reference in new issue
Block a user