openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-06 15:29:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(app-server-transport): refresh pairing auth in place

Browse Source
This commit is contained in:
Anton Panasenko
2026-05-29 04:47:09 -07:00
parent 5b7077c8bb
commit 974cee60db
5 changed files with 663 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
codex-rs/app-server-transport/src/transport/remote_control/enroll.rs
View File

@@ -46,6 +46,13 @@ impl RemoteControlEnrollment {
})
}
pub(super) fn server_token_refresh_delay(&self) -> Option<std::time::Duration> {
let refresh_at = self.expires_at?
- time::Duration::seconds(REMOTE_CONTROL_SERVER_TOKEN_REFRESH_SKEW_SECS);
let refresh_delay = refresh_at - OffsetDateTime::now_utc();
Some(std::time::Duration::try_from(refresh_delay).unwrap_or_default())
}
pub(super) fn clear_server_token(&mut self) {
self.remote_control_token = None;
self.expires_at = None;
@@ -435,6 +442,22 @@ mod tests {
assert!(!expires_later.should_refresh_server_token());
}
#[test]
fn remote_control_enrollment_schedules_server_token_refresh_before_expiry() {
let refresh_delay = RemoteControlEnrollment {
account_id: "account_id".to_string(),
environment_id: "environment_id".to_string(),
server_id: "server_id".to_string(),
server_name: "server_name".to_string(),
remote_control_token: Some("remote-control-token".to_string()),
expires_at: Some(OffsetDateTime::now_utc() + time::Duration::seconds(31)),
}
.server_token_refresh_delay()
.expect("server token refresh should be scheduled");
assert!(refresh_delay <= std::time::Duration::from_secs(1));
}
#[test]
fn preview_remote_control_response_body_redacts_server_token() {
assert_eq!(

14
codex-rs/app-server-transport/src/transport/remote_control/mod.rs
View File

@@ -60,6 +60,8 @@ pub struct RemoteControlHandle {
status_tx: Arc<watch::Sender<RemoteControlStatusChangedNotification>>,
state_db_available: bool,
pairing: PairingClientState,
#[cfg(test)]
pairing_refresh_tx: Arc<watch::Sender<u64>>,
auth_change_rx: Arc<StdMutex<watch::Receiver<u64>>>,
}
@@ -230,6 +232,12 @@ impl RemoteControlHandle {
})
}
#[cfg(test)]
fn request_pairing_auth_refresh(&self) {
self.pairing_refresh_tx
.send_modify(|revision| *revision = revision.wrapping_add(1));
}
fn publish_status(
&self,
connection_status: RemoteControlConnectionStatus,
@@ -309,6 +317,8 @@ pub async fn start_remote_control(
let (enabled_tx, enabled_rx) = watch::channel(initial_enabled);
let pairing = PairingClientState::new();
let websocket_pairing = pairing.clone();
let (pairing_refresh_tx, pairing_refresh_rx) = watch::channel(0u64);
let websocket_pairing_refresh_tx = pairing_refresh_tx.clone();
let auth_change_rx = Arc::new(StdMutex::new(auth_manager.auth_change_receiver()));
let server_name = gethostname().to_string_lossy().trim().to_string();
let remote_control_url = config.remote_control_url;
@@ -358,6 +368,8 @@ pub async fn start_remote_control(
transport_event_tx,
status_publisher,
pairing: websocket_pairing,
pairing_refresh_tx: websocket_pairing_refresh_tx,
pairing_refresh_rx,
},
shutdown_token,
enabled_rx,
@@ -403,6 +415,8 @@ pub async fn start_remote_control(
status_tx: Arc::new(status_tx),
state_db_available,
pairing,
#[cfg(test)]
pairing_refresh_tx: Arc::new(pairing_refresh_tx),
auth_change_rx,
},
))

183
codex-rs/app-server-transport/src/transport/remote_control/pairing_integration_tests.rs
View File

@@ -394,3 +394,186 @@ async fn remote_control_handle_clears_pairing_client_after_auth_change() {
shutdown_token.cancel();
let _ = remote_task.await;
}
#[tokio::test]
async fn remote_control_refreshes_server_token_while_connected() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("listener should bind");
let remote_control_url = remote_control_url_for_listener(&listener);
let codex_home = TempDir::new().expect("temp dir should create");
let (transport_event_tx, _transport_event_rx) =
mpsc::channel::<TransportEvent>(CHANNEL_CAPACITY);
let shutdown_token = CancellationToken::new();
let (remote_task, remote_handle) = start_remote_control(
RemoteControlStartConfig {
remote_control_url,
installation_id: TEST_INSTALLATION_ID.to_string(),
},
Some(remote_control_state_runtime(&codex_home).await),
remote_control_auth_manager(),
transport_event_tx,
shutdown_token.clone(),
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
)
.await
.expect("remote control should start");
let enroll_request = accept_http_request(&listener).await;
respond_with_json(
enroll_request.stream,
remote_control_server_token_response(
"srv_e_test",
"env_test",
TEST_REMOTE_CONTROL_SERVER_TOKEN,
),
)
.await;
let mut first_websocket = accept_remote_control_connection(&listener).await;
remote_handle.request_pairing_auth_refresh();
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(
"srv_e_test",
"env_test",
TEST_REFRESHED_REMOTE_CONTROL_SERVER_TOKEN,
),
)
.await;
assert!(
timeout(Duration::from_millis(100), first_websocket.next())
.await
.is_err(),
"server token refresh should keep the websocket open"
);
let pairing_task = tokio::spawn({
let remote_handle = remote_handle.clone();
async move {
remote_handle
.start_pairing(RemoteControlPairingStartParams::default())
.await
}
});
let pairing_request = accept_http_request(&listener).await;
assert_eq!(
pairing_request.request_line,
"POST /backend-api/wham/remote/control/server/pair HTTP/1.1"
);
assert_eq!(
pairing_request.headers.get("authorization"),
Some(&format!(
"Bearer {TEST_REFRESHED_REMOTE_CONTROL_SERVER_TOKEN}"
))
);
respond_with_json(
pairing_request.stream,
json!({
"pairing_code": "pairing-code",
"manual_pairing_code": "ABCD-EFGH",
"server_id": "srv_e_test",
"environment_id": "env_test",
"expires_at": "3026-05-22T12:34:56Z",
}),
)
.await;
assert_eq!(
pairing_task
.await
.expect("pairing task should join")
.expect("pairing should use refreshed server token"),
codex_app_server_protocol::RemoteControlPairingStartResponse {
pairing_code: "pairing-code".to_string(),
manual_pairing_code: Some("ABCD-EFGH".to_string()),
environment_id: "env_test".to_string(),
expires_at: 33_336_362_096,
}
);
first_websocket
.close(None)
.await
.expect("first websocket should close");
shutdown_token.cancel();
let _ = remote_task.await;
}
#[tokio::test]
async fn remote_control_schedules_server_token_refresh_while_connected() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("listener should bind");
let remote_control_url = remote_control_url_for_listener(&listener);
let codex_home = TempDir::new().expect("temp dir should create");
let (transport_event_tx, _transport_event_rx) =
mpsc::channel::<TransportEvent>(CHANNEL_CAPACITY);
let shutdown_token = CancellationToken::new();
let (remote_task, _remote_handle) = start_remote_control(
RemoteControlStartConfig {
remote_control_url,
installation_id: TEST_INSTALLATION_ID.to_string(),
},
Some(remote_control_state_runtime(&codex_home).await),
remote_control_auth_manager(),
transport_event_tx,
shutdown_token.clone(),
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
)
.await
.expect("remote control should start");
let scheduled_refresh_expires_at = (OffsetDateTime::now_utc() + time::Duration::seconds(35))
.format(&time::format_description::well_known::Rfc3339)
.expect("scheduled refresh expiry should format");
let enroll_request = accept_http_request(&listener).await;
respond_with_json(
enroll_request.stream,
remote_control_server_token_response_with_expires_at(
"srv_e_test",
"env_test",
TEST_REMOTE_CONTROL_SERVER_TOKEN,
&scheduled_refresh_expires_at,
),
)
.await;
let mut first_websocket = accept_remote_control_connection(&listener).await;
let refresh_request = timeout(Duration::from_secs(10), accept_http_request(&listener))
.await
.expect("scheduled server token refresh should arrive");
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(
"srv_e_test",
"env_test",
TEST_REFRESHED_REMOTE_CONTROL_SERVER_TOKEN,
),
)
.await;
assert!(
timeout(Duration::from_millis(100), first_websocket.next())
.await
.is_err(),
"scheduled server token refresh should keep the websocket open"
);
first_websocket
.close(None)
.await
.expect("first websocket should close");
shutdown_token.cancel();
let _ = remote_task.await;
}

18
codex-rs/app-server-transport/src/transport/remote_control/tests.rs
View File

@@ -144,6 +144,7 @@ fn remote_control_handle_with_pairing_client(
installation_id: TEST_INSTALLATION_ID.to_string(),
environment_id: Some("env_test".to_string()),
});
let (pairing_refresh_tx, _pairing_refresh_rx) = watch::channel(0u64);
let pairing_client = Arc::new(StdMutex::new(Some(RemoteControlPairingClient::new(
&normalize_remote_control_url(remote_control_url)
.expect("remote control target should normalize"),
@@ -162,6 +163,7 @@ fn remote_control_handle_with_pairing_client(
client: pairing_client,
generation: Arc::new(std::sync::atomic::AtomicU64::new(0)),
},
pairing_refresh_tx: Arc::new(pairing_refresh_tx),
auth_change_rx: Arc::new(StdMutex::new(auth_change_rx)),
}
}
@@ -170,12 +172,26 @@ fn remote_control_server_token_response(
server_id: &str,
environment_id: &str,
remote_control_token: &str,
) -> serde_json::Value {
remote_control_server_token_response_with_expires_at(
server_id,
environment_id,
remote_control_token,
TEST_REMOTE_CONTROL_SERVER_TOKEN_EXPIRES_AT,
)
}
fn remote_control_server_token_response_with_expires_at(
server_id: &str,
environment_id: &str,
remote_control_token: &str,
expires_at: &str,
) -> serde_json::Value {
json!({
"server_id": server_id,
"environment_id": environment_id,
"remote_control_token": remote_control_token,
"expires_at": TEST_REMOTE_CONTROL_SERVER_TOKEN_EXPIRES_AT,
"expires_at": expires_at,
})
}

445
codex-rs/app-server-transport/src/transport/remote_control/websocket.rs
View File

@@ -256,6 +256,8 @@ pub(crate) struct RemoteControlWebsocket {
auth_recovery: UnauthorizedRecovery,
auth_change_rx: watch::Receiver<u64>,
pairing: PairingClientState,
_pairing_refresh_tx: watch::Sender<u64>,
pairing_refresh_rx: watch::Receiver<u64>,
client_tracker: Arc<Mutex<ClientTracker>>,
state: Arc<Mutex<WebsocketState>>,
server_event_rx: Arc<Mutex<mpsc::Receiver<super::QueuedServerEnvelope>>>,
@@ -289,13 +291,50 @@ enum ConnectionEndReason {
EnabledWatchClosed,
AuthChanged,
AuthWatchClosed,
PairingRefreshWatchClosed,
ServerTokenRefreshRejected,
StaleEnrollment,
ConnectionWorkerStopped,
}
enum ConnectionLoopAction {
End(ConnectionEndReason),
RefreshServerToken,
}
enum ConnectedServerTokenRefreshAction {
Continue,
Retry(io::Error),
End(ConnectionEndReason),
}
struct ConnectedServerTokenRefreshRequest {
remote_control_target: RemoteControlTarget,
auth: RemoteControlConnectionAuth,
installation_id: String,
enrollment: RemoteControlEnrollment,
auth_change_revision: u64,
}
struct PendingConnectedServerTokenRefreshRequest {
remote_control_target: RemoteControlTarget,
auth_manager: Arc<AuthManager>,
auth_change_rx: watch::Receiver<u64>,
installation_id: String,
enrollment: RemoteControlEnrollment,
}
struct ConnectedServerTokenRefreshResponse {
enrollment: RemoteControlEnrollment,
auth_change_revision: u64,
}
pub(super) struct RemoteControlChannels {
pub(super) transport_event_tx: mpsc::Sender<TransportEvent>,
pub(super) status_publisher: RemoteControlStatusPublisher,
pub(super) pairing: PairingClientState,
pub(super) pairing_refresh_tx: watch::Sender<u64>,
pub(super) pairing_refresh_rx: watch::Receiver<u64>,
}
#[derive(Clone)]
@@ -413,6 +452,8 @@ impl RemoteControlWebsocket {
auth_recovery,
auth_change_rx,
pairing: channels.pairing,
_pairing_refresh_tx: channels.pairing_refresh_tx,
pairing_refresh_rx: channels.pairing_refresh_rx,
client_tracker: Arc::new(Mutex::new(client_tracker)),
state: Arc::new(Mutex::new(WebsocketState {
outbound_buffer,
@@ -498,7 +539,11 @@ impl RemoteControlWebsocket {
};
let connection_end_reason = self
.run_connection(websocket_connection, shutdown_token)
.run_connection(
websocket_connection,
shutdown_token,
app_server_client_name.as_deref(),
)
.await;
let status = self.status_publisher.status();
info!(
@@ -710,6 +755,7 @@ impl RemoteControlWebsocket {
&mut self,
websocket_connection: WebSocketStream<MaybeTlsStream<TcpStream>>,
shutdown_token: CancellationToken,
app_server_client_name: Option<&str>,
) -> ConnectionEndReason {
let (websocket_writer, websocket_reader) = websocket_connection.split();
let mut join_set = tokio::task::JoinSet::new();
@@ -731,26 +777,79 @@ impl RemoteControlWebsocket {
));
let mut enabled_rx = self.enabled_rx.clone();
let connection_end_reason = tokio::select! {
_ = shutdown_token.cancelled() => ConnectionEndReason::Shutdown,
changed = enabled_rx.wait_for(|enabled| !*enabled) => {
if changed.is_ok() {
self.status_publisher
.publish_status(RemoteControlConnectionStatus::Disabled);
ConnectionEndReason::Disabled
} else {
ConnectionEndReason::EnabledWatchClosed
let mut server_token_refresh_retry_delay = None;
let connection_end_reason = loop {
let server_token_refresh_delay =
server_token_refresh_retry_delay.take().or_else(|| {
self.enrollment
.as_ref()
.and_then(RemoteControlEnrollment::server_token_refresh_delay)
});
let server_token_refresh = async move {
match server_token_refresh_delay {
Some(delay) => tokio::time::sleep(delay).await,
None => std::future::pending().await,
}
};
tokio::pin!(server_token_refresh);
let connection_loop_action = tokio::select! {
_ = shutdown_token.cancelled() => ConnectionLoopAction::End(ConnectionEndReason::Shutdown),
changed = enabled_rx.wait_for(|enabled| !*enabled) => {
if changed.is_ok() {
self.status_publisher
.publish_status(RemoteControlConnectionStatus::Disabled);
ConnectionLoopAction::End(ConnectionEndReason::Disabled)
} else {
ConnectionLoopAction::End(ConnectionEndReason::EnabledWatchClosed)
}
}
changed = self.auth_change_rx.changed() => {
if changed.is_ok() {
self.auth_recovery = self.auth_manager.unauthorized_recovery();
ConnectionLoopAction::End(ConnectionEndReason::AuthChanged)
} else {
ConnectionLoopAction::End(ConnectionEndReason::AuthWatchClosed)
}
}
changed = self.pairing_refresh_rx.changed() => {
if changed.is_ok() {
ConnectionLoopAction::RefreshServerToken
} else {
ConnectionLoopAction::End(ConnectionEndReason::PairingRefreshWatchClosed)
}
}
_ = &mut server_token_refresh => ConnectionLoopAction::RefreshServerToken,
_ = join_set.join_next() => ConnectionLoopAction::End(ConnectionEndReason::ConnectionWorkerStopped),
};
match connection_loop_action {
ConnectionLoopAction::End(connection_end_reason) => break connection_end_reason,
ConnectionLoopAction::RefreshServerToken => {
match self
.refresh_connected_server_token_while_connected(
&mut enabled_rx,
&mut join_set,
&shutdown_token,
app_server_client_name,
)
.await
{
ConnectedServerTokenRefreshAction::Continue => {}
ConnectedServerTokenRefreshAction::Retry(err) => {
warn!(
error = %err,
error_kind = ?err.kind(),
retry_delay = ?REMOTE_CONTROL_ACCOUNT_ID_RETRY_INTERVAL,
"failed to refresh connected app-server remote control server token"
);
server_token_refresh_retry_delay =
Some(REMOTE_CONTROL_ACCOUNT_ID_RETRY_INTERVAL);
}
ConnectedServerTokenRefreshAction::End(connection_end_reason) => {
break connection_end_reason;
}
}
}
}
changed = self.auth_change_rx.changed() => {
if changed.is_ok() {
self.auth_recovery = self.auth_manager.unauthorized_recovery();
ConnectionEndReason::AuthChanged
} else {
ConnectionEndReason::AuthWatchClosed
}
}
_ = join_set.join_next() => ConnectionEndReason::ConnectionWorkerStopped,
};
clear_pairing_client(&self.pairing);
shutdown_token.cancel();
@@ -760,6 +859,209 @@ impl RemoteControlWebsocket {
connection_end_reason
}
async fn refresh_connected_server_token_while_connected(
&mut self,
enabled_rx: &mut watch::Receiver<bool>,
join_set: &mut tokio::task::JoinSet<()>,
shutdown_token: &CancellationToken,
app_server_client_name: Option<&str>,
) -> ConnectedServerTokenRefreshAction {
let refresh_request = match self.connected_server_token_refresh_request() {
Ok(refresh_request) => refresh_request,
Err(err) => return ConnectedServerTokenRefreshAction::Retry(err),
};
let refresh_request = tokio::select! {
_ = shutdown_token.cancelled() => {
return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::Shutdown);
}
changed = enabled_rx.wait_for(|enabled| !*enabled) => {
if changed.is_ok() {
self.status_publisher
.publish_status(RemoteControlConnectionStatus::Disabled);
return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::Disabled);
}
return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::EnabledWatchClosed);
}
_ = join_set.join_next() => {
return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::ConnectionWorkerStopped);
}
refresh_request = prepare_connected_server_token_refresh(refresh_request) => refresh_request,
};
let refresh_request = match refresh_request {
Ok(refresh_request) => refresh_request,
Err(err) => {
return self
.apply_connected_server_token_refresh(Err(err), app_server_client_name)
.await;
}
};
if !mark_connected_refresh_auth_change_seen(
&mut self.auth_change_rx,
refresh_request.auth_change_revision,
) {
clear_pairing_client(&self.pairing);
self.auth_recovery = self.auth_manager.unauthorized_recovery();
return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::AuthChanged);
}
let refresh_result = tokio::select! {
_ = shutdown_token.cancelled() => {
return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::Shutdown);
}
changed = enabled_rx.wait_for(|enabled| !*enabled) => {
if changed.is_ok() {
self.status_publisher
.publish_status(RemoteControlConnectionStatus::Disabled);
return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::Disabled);
}
return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::EnabledWatchClosed);
}
changed = self.auth_change_rx.changed() => {
if changed.is_ok() {
self.auth_recovery = self.auth_manager.unauthorized_recovery();
return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::AuthChanged);
}
return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::AuthWatchClosed);
}
_ = join_set.join_next() => {
return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::ConnectionWorkerStopped);
}
refresh_result = refresh_connected_server_token(refresh_request) => refresh_result,
};
self.apply_connected_server_token_refresh(refresh_result, app_server_client_name)
.await
}
fn connected_server_token_refresh_request(
&self,
) -> io::Result<PendingConnectedServerTokenRefreshRequest> {
let remote_control_target = self.remote_control_target.clone().ok_or_else(|| {
io::Error::other("missing remote control target while refreshing server token")
})?;
let enrollment = self.enrollment.clone().ok_or_else(|| {
io::Error::other("missing remote control enrollment while refreshing server token")
})?;
Ok(PendingConnectedServerTokenRefreshRequest {
remote_control_target,
auth_manager: self.auth_manager.clone(),
auth_change_rx: self.auth_change_rx.clone(),
installation_id: self.installation_id.clone(),
enrollment,
})
}
async fn apply_connected_server_token_refresh(
&mut self,
refresh_result: io::Result<ConnectedServerTokenRefreshResponse>,
app_server_client_name: Option<&str>,
) -> ConnectedServerTokenRefreshAction {
let refresh_response = match refresh_result {
Ok(refresh_response) => refresh_response,
Err(err) if err.kind() == ErrorKind::NotFound => {
self.clear_stale_connected_enrollment(app_server_client_name)
.await;
return ConnectedServerTokenRefreshAction::End(
ConnectionEndReason::StaleEnrollment,
);
}
Err(err) if err.kind() == ErrorKind::PermissionDenied => {
if recover_remote_control_auth(&mut self.auth_recovery, &mut self.auth_change_rx)
.await
{
return ConnectedServerTokenRefreshAction::Retry(io::Error::other(format!(
"{err}; retrying after auth recovery"
)));
}
if let Some(enrollment) = self.enrollment.as_mut() {
enrollment.clear_server_token();
}
clear_pairing_client(&self.pairing);
return ConnectedServerTokenRefreshAction::End(
ConnectionEndReason::ServerTokenRefreshRejected,
);
}
Err(err) if err.kind() == ErrorKind::InvalidInput => {
clear_pairing_client(&self.pairing);
self.auth_recovery = self.auth_manager.unauthorized_recovery();
return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::AuthChanged);
}
Err(err) => return ConnectedServerTokenRefreshAction::Retry(err),
};
if *self.auth_change_rx.borrow() != refresh_response.auth_change_revision {
clear_pairing_client(&self.pairing);
self.auth_recovery = self.auth_manager.unauthorized_recovery();
return ConnectedServerTokenRefreshAction::End(ConnectionEndReason::AuthChanged);
}
let current_enrollment = match self.enrollment.as_ref() {
Some(current_enrollment) => current_enrollment,
None => return ConnectedServerTokenRefreshAction::Continue,
};
if !same_remote_control_enrollment_identity(
current_enrollment,
&refresh_response.enrollment,
) {
return ConnectedServerTokenRefreshAction::Continue;
}
self.enrollment = Some(refresh_response.enrollment);
let Some(remote_control_target) = self.remote_control_target.as_ref() else {
return ConnectedServerTokenRefreshAction::Retry(io::Error::other(
"missing remote control target after refreshing server token",
));
};
let Some(enrollment) = self.enrollment.as_ref() else {
return ConnectedServerTokenRefreshAction::Retry(io::Error::other(
"missing remote control enrollment after refreshing server token",
));
};
match set_pairing_client(
&self.pairing,
remote_control_target,
enrollment,
refresh_response.auth_change_revision,
) {
Ok(()) => ConnectedServerTokenRefreshAction::Continue,
Err(err) => ConnectedServerTokenRefreshAction::Retry(err),
}
}
async fn clear_stale_connected_enrollment(&mut self, app_server_client_name: Option<&str>) {
let Some(remote_control_target) = self.remote_control_target.as_ref() else {
clear_pairing_client(&self.pairing);
self.enrollment = None;
self.status_publisher
.publish_environment_id(/*environment_id*/ None);
return;
};
let Some(account_id) = self
.enrollment
.as_ref()
.map(|enrollment| enrollment.account_id.clone())
else {
clear_pairing_client(&self.pairing);
return;
};
info!(
"connected remote control server refresh returned HTTP 404; clearing stale enrollment before re-enrolling: websocket_url={}, account_id={}",
remote_control_target.websocket_url, account_id
);
if let Some(state_db) = self.state_db.as_deref() {
clear_remote_control_enrollment(
state_db,
remote_control_target,
&account_id,
app_server_client_name,
&mut self.enrollment,
&self.status_publisher,
&self.pairing,
)
.await;
return;
}
self.enrollment = None;
self.status_publisher
.publish_environment_id(/*environment_id*/ None);
clear_pairing_client(&self.pairing);
}
async fn join_connection_workers(
join_set: &mut tokio::task::JoinSet<()>,
shutdown_timeout: std::time::Duration,
@@ -1186,6 +1488,73 @@ fn build_remote_control_websocket_request(
Ok(request)
}
async fn refresh_connected_server_token(
mut refresh_request: ConnectedServerTokenRefreshRequest,
) -> io::Result<ConnectedServerTokenRefreshResponse> {
info!(
"refreshing connected remote control server token: websocket_url={}, refresh_url={}, account_id={}, server_id={}, environment_id={}",
refresh_request.remote_control_target.websocket_url,
refresh_request.remote_control_target.refresh_url,
refresh_request.auth.account_id,
refresh_request.enrollment.server_id,
refresh_request.enrollment.environment_id
);
refresh_remote_control_server(
&refresh_request.remote_control_target,
&refresh_request.auth,
&refresh_request.installation_id,
&mut refresh_request.enrollment,
)
.await?;
Ok(ConnectedServerTokenRefreshResponse {
enrollment: refresh_request.enrollment,
auth_change_revision: refresh_request.auth_change_revision,
})
}
async fn prepare_connected_server_token_refresh(
mut refresh_request: PendingConnectedServerTokenRefreshRequest,
) -> io::Result<ConnectedServerTokenRefreshRequest> {
let auth = load_remote_control_auth(&refresh_request.auth_manager).await?;
// Loading auth may reload or proactively refresh through the same watch
// receiver. Treat the auth used for this refresh as the current revision
// before waiting for later auth changes to cancel the HTTP request.
let auth_change_revision = *refresh_request.auth_change_rx.borrow_and_update();
if refresh_request.enrollment.account_id != auth.account_id {
return Err(io::Error::new(
ErrorKind::InvalidInput,
"remote control auth changed while refreshing server token",
));
}
Ok(ConnectedServerTokenRefreshRequest {
remote_control_target: refresh_request.remote_control_target,
auth,
installation_id: refresh_request.installation_id,
enrollment: refresh_request.enrollment,
auth_change_revision,
})
}
fn mark_connected_refresh_auth_change_seen(
auth_change_rx: &mut watch::Receiver<u64>,
auth_change_revision: u64,
) -> bool {
if *auth_change_rx.borrow() != auth_change_revision {
return false;
}
auth_change_rx.borrow_and_update();
true
}
fn same_remote_control_enrollment_identity(
left: &RemoteControlEnrollment,
right: &RemoteControlEnrollment,
) -> bool {
left.account_id == right.account_id
&& left.server_id == right.server_id
&& left.environment_id == right.environment_id
}
pub(crate) async fn load_remote_control_auth(
auth_manager: &Arc<AuthManager>,
) -> io::Result<RemoteControlConnectionAuth> {
@@ -1802,6 +2171,41 @@ mod tests {
);
}
#[test]
fn mark_connected_refresh_auth_change_seen_marks_loaded_auth_revision_seen() {
let (auth_change_tx, mut auth_change_rx) = watch::channel(0u64);
auth_change_tx.send_modify(|revision| *revision += 1);
let auth_change_revision = *auth_change_rx.borrow();
assert!(mark_connected_refresh_auth_change_seen(
&mut auth_change_rx,
auth_change_revision
));
assert!(
!auth_change_rx
.has_changed()
.expect("auth change watch should remain open")
);
}
#[test]
fn mark_connected_refresh_auth_change_seen_preserves_racing_auth_change() {
let (auth_change_tx, mut auth_change_rx) = watch::channel(0u64);
auth_change_tx.send_modify(|revision| *revision += 1);
let auth_change_revision = *auth_change_rx.borrow();
auth_change_tx.send_modify(|revision| *revision += 1);
assert!(!mark_connected_refresh_auth_change_seen(
&mut auth_change_rx,
auth_change_revision
));
assert!(
auth_change_rx
.has_changed()
.expect("auth change watch should remain open")
);
}
async fn remote_control_state_runtime(codex_home: &TempDir) -> Arc<StateRuntime> {
StateRuntime::init(codex_home.path().to_path_buf(), "test-provider".to_string())
.await
@@ -2325,6 +2729,7 @@ mod tests {
let (status_publisher, _status_rx) = remote_control_status_channel();
let shutdown_token = CancellationToken::new();
let (_enabled_tx, enabled_rx) = watch::channel(true);
let (pairing_refresh_tx, pairing_refresh_rx) = watch::channel(0u64);
let websocket_task = tokio::spawn({
let shutdown_token = shutdown_token.clone();
async move {
@@ -2341,6 +2746,8 @@ mod tests {
transport_event_tx,
status_publisher,
pairing: test_pairing(),
pairing_refresh_tx,
pairing_refresh_rx,
},
shutdown_token,
enabled_rx,