Skip to content

Add remote control instance naming #21430

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ebrevdo wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add remote control instance naming #21430

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions codex-rs/app-server-transport/src/transport/remote_control/enroll.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,9 +193,16 @@ pub(crate) fn format_headers(headers: &HeaderMap) -> String {
pub(super) async fn enroll_remote_control_server(
remote_control_target: &RemoteControlTarget,
auth: &RemoteControlConnectionAuth,
remote_control_instance_name: Option<&str>,
) -> io::Result<RemoteControlEnrollment> {
let enroll_url = &remote_control_target.enroll_url;
let server_name = gethostname().to_string_lossy().trim().to_string();
let server_name = match remote_control_instance_name {
Some(remote_control_instance_name) => {
format!("{server_name} - {remote_control_instance_name}")
}
None => server_name,
};
let request = EnrollRemoteServerRequest {
name: server_name.clone(),
os: std::env::consts::OS,
Expand Down Expand Up @@ -459,6 +466,7 @@ mod tests {
auth_provider: codex_model_provider::unauthenticated_auth_provider(),
account_id: "account_id".to_string(),
},
/*remote_control_instance_name*/ None,
)
.await
.expect_err("invalid response should fail to parse");
Expand Down
2 changes: 2 additions & 0 deletions codex-rs/app-server-transport/src/transport/remote_control/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@ pub async fn start_remote_control(
auth_manager: Arc<AuthManager>,
transport_event_tx: mpsc::Sender<TransportEvent>,
shutdown_token: CancellationToken,
remote_control_instance_name: Option<String>,
app_server_client_name_rx: Option<oneshot::Receiver<String>>,
initial_enabled: bool,
) -> io::Result<(JoinHandle<()>, RemoteControlHandle)> {
Expand Down Expand Up @@ -105,6 +106,7 @@ pub async fn start_remote_control(
},
shutdown_token,
enabled_rx,
remote_control_instance_name,
)
.run(app_server_client_name_rx)
.await;
Expand Down
115 changes: 115 additions & 0 deletions codex-rs/app-server-transport/src/transport/remote_control/tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,7 @@ async fn remote_control_transport_manages_virtual_clients_and_routes_messages()
remote_control_auth_manager(),
transport_event_tx,
shutdown_token.clone(),
/*remote_control_instance_name*/ None,
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
)
Expand Down Expand Up @@ -454,6 +455,7 @@ async fn remote_control_transport_reconnects_after_disconnect() {
remote_control_auth_manager(),
transport_event_tx,
shutdown_token.clone(),
/*remote_control_instance_name*/ None,
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
)
Expand Down Expand Up @@ -533,6 +535,7 @@ async fn remote_control_start_allows_remote_control_invalid_url_when_disabled()
remote_control_auth_manager(),
transport_event_tx,
shutdown_token.clone(),
/*remote_control_instance_name*/ None,
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ false,
)
Expand Down Expand Up @@ -569,6 +572,7 @@ async fn remote_control_start_allows_missing_auth_when_enabled() {
auth_manager,
transport_event_tx,
shutdown_token.clone(),
/*remote_control_instance_name*/ None,
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
)
Expand Down Expand Up @@ -601,6 +605,7 @@ async fn remote_control_start_reports_missing_state_db_as_disabled_when_enabled(
remote_control_auth_manager(),
transport_event_tx,
shutdown_token.clone(),
/*remote_control_instance_name*/ None,
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
)
Expand Down Expand Up @@ -650,6 +655,7 @@ async fn remote_control_handle_set_enabled_stops_and_restarts_connections() {
remote_control_auth_manager(),
transport_event_tx,
shutdown_token.clone(),
/*remote_control_instance_name*/ None,
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
)
Expand Down Expand Up @@ -734,6 +740,7 @@ async fn remote_control_transport_clears_outgoing_buffer_when_backend_acks() {
remote_control_auth_manager(),
transport_event_tx,
shutdown_token.clone(),
/*remote_control_instance_name*/ None,
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
)
Expand Down Expand Up @@ -909,6 +916,7 @@ async fn remote_control_http_mode_enrolls_before_connecting() {
remote_control_auth_manager(),
transport_event_tx,
shutdown_token.clone(),
/*remote_control_instance_name*/ None,
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
)
Expand Down Expand Up @@ -1098,6 +1106,109 @@ async fn remote_control_http_mode_enrolls_before_connecting() {
let _ = remote_task.await;
}

#[tokio::test]
async fn remote_control_http_mode_uses_instance_name_for_enrollment() {
let listener = TcpListener::bind("127.0.0.1:0")
.await
.expect("listener should bind");
let remote_control_url = remote_control_url_for_listener(&listener);
let codex_home = TempDir::new().expect("temp dir should create");
let state_db = remote_control_state_runtime(&codex_home).await;
let remote_control_target =
normalize_remote_control_url(&remote_control_url).expect("target should parse");
let default_enrollment = RemoteControlEnrollment {
account_id: "account_id".to_string(),
environment_id: "env_default".to_string(),
server_id: "srv_e_default".to_string(),
server_name: "default-server".to_string(),
};
update_persisted_remote_control_enrollment(
Some(state_db.as_ref()),
&remote_control_target,
"account_id",
/*app_server_client_name*/ None,
Some(&default_enrollment),
)
.await
.expect("default persisted enrollment should save");

let instance_name = "next-build";
let hostname = gethostname().to_string_lossy().trim().to_string();
let expected_server_name = format!("{hostname} - {instance_name}");
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(
remote_control_url,
Some(state_db.clone()),
remote_control_auth_manager_with_home(&codex_home),
transport_event_tx,
shutdown_token.clone(),
Some(instance_name.to_string()),
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
)
.await
.expect("remote control should start");

let enroll_request = accept_http_request(&listener).await;
assert_eq!(
serde_json::from_str::<serde_json::Value>(&enroll_request.body)
.expect("enroll body should deserialize"),
json!({
"name": expected_server_name,
"os": std::env::consts::OS,
"arch": std::env::consts::ARCH,
"app_server_version": env!("CARGO_PKG_VERSION"),
})
);
respond_with_json(
enroll_request.stream,
json!({ "server_id": "srv_e_instance", "environment_id": "env_instance" }),
)
.await;

let (handshake_request, _websocket) = accept_remote_control_backend_connection(&listener).await;
assert_eq!(
handshake_request.headers.get("x-codex-server-id"),
Some(&"srv_e_instance".to_string())
);
assert_eq!(
handshake_request.headers.get("x-codex-name"),
Some(&base64::engine::general_purpose::STANDARD.encode(&expected_server_name))
);
assert_eq!(
load_persisted_remote_control_enrollment(
Some(state_db.as_ref()),
&remote_control_target,
"account_id",
Some(instance_name),
)
.await
.expect("instance persisted enrollment should load"),
Some(RemoteControlEnrollment {
account_id: "account_id".to_string(),
environment_id: "env_instance".to_string(),
server_id: "srv_e_instance".to_string(),
server_name: expected_server_name,
})
);
assert_eq!(
load_persisted_remote_control_enrollment(
Some(state_db.as_ref()),
&remote_control_target,
"account_id",
/*app_server_client_name*/ None,
)
.await
.expect("default persisted enrollment should load"),
Some(default_enrollment)
);

shutdown_token.cancel();
let _ = remote_task.await;
}

#[tokio::test]
async fn remote_control_http_mode_reuses_persisted_enrollment_before_reenrolling() {
let listener = TcpListener::bind("127.0.0.1:0")
Expand Down Expand Up @@ -1133,6 +1244,7 @@ async fn remote_control_http_mode_reuses_persisted_enrollment_before_reenrolling
remote_control_auth_manager_with_home(&codex_home),
transport_event_tx,
shutdown_token.clone(),
/*remote_control_instance_name*/ None,
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
)
Expand Down Expand Up @@ -1201,6 +1313,7 @@ async fn remote_control_stdio_mode_waits_for_client_name_before_connecting() {
remote_control_auth_manager_with_home(&codex_home),
transport_event_tx,
shutdown_token.clone(),
/*remote_control_instance_name*/ None,
Some(app_server_client_name_rx),
/*initial_enabled*/ true,
)
Expand Down Expand Up @@ -1260,6 +1373,7 @@ async fn remote_control_waits_for_account_id_before_enrolling() {
auth_manager,
transport_event_tx,
shutdown_token.clone(),
/*remote_control_instance_name*/ None,
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
)
Expand Down Expand Up @@ -1343,6 +1457,7 @@ async fn remote_control_http_mode_clears_stale_persisted_enrollment_after_404()
remote_control_auth_manager_with_home(&codex_home),
transport_event_tx,
shutdown_token.clone(),
/*remote_control_instance_name*/ None,
/*app_server_client_name_rx*/ None,
/*initial_enabled*/ true,
)
Expand Down
27 changes: 23 additions & 4 deletions codex-rs/app-server-transport/src/transport/remote_control/websocket.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -217,6 +217,7 @@ pub(crate) struct RemoteControlWebsocket {
remote_control_target: Option<RemoteControlTarget>,
state_db: Option<Arc<StateRuntime>>,
auth_manager: Arc<AuthManager>,
remote_control_instance_name: Option<String>,
status_publisher: RemoteControlStatusPublisher,
shutdown_token: CancellationToken,
reconnect_attempt: u64,
Expand Down Expand Up @@ -292,6 +293,7 @@ impl RemoteControlStatusPublisher {
pub(super) struct RemoteControlConnectOptions<'a> {
subscribe_cursor: Option<&'a str>,
app_server_client_name: Option<&'a str>,
remote_control_instance_name: Option<&'a str>,
}

impl RemoteControlWebsocket {
Expand All @@ -303,6 +305,7 @@ impl RemoteControlWebsocket {
channels: RemoteControlChannels,
shutdown_token: CancellationToken,
enabled_rx: watch::Receiver<bool>,
remote_control_instance_name: Option<String>,
) -> Self {
let shutdown_token = shutdown_token.child_token();
let (server_event_tx, server_event_rx) = mpsc::channel(super::CHANNEL_CAPACITY);
Expand All @@ -319,6 +322,7 @@ impl RemoteControlWebsocket {
remote_control_target,
state_db,
auth_manager,
remote_control_instance_name,
status_publisher: channels.status_publisher,
shutdown_token,
reconnect_attempt: 0,
Expand Down Expand Up @@ -444,6 +448,7 @@ impl RemoteControlWebsocket {
let connect_options = RemoteControlConnectOptions {
subscribe_cursor: subscribe_cursor.as_deref(),
app_server_client_name,
remote_control_instance_name: self.remote_control_instance_name.as_deref(),
};
let connect_result = tokio::select! {
_ = shutdown_token.cancelled() => return ConnectOutcome::Shutdown,
Expand Down Expand Up @@ -1029,6 +1034,9 @@ pub(super) async fn connect_remote_control_websocket(
return Err(err);
}
};
let enrollment_app_server_client_name = connect_options
.remote_control_instance_name
.or(connect_options.app_server_client_name);
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) {
info!(
Expand All @@ -1052,7 +1060,7 @@ pub(super) async fn connect_remote_control_websocket(
Some(state_db),
remote_control_target,
&auth.account_id,
connect_options.app_server_client_name,
enrollment_app_server_client_name,
)
.await?;
if let Some(loaded_enrollment) = loaded_enrollment.as_ref() {
Expand All @@ -1066,7 +1074,12 @@ pub(super) async fn connect_remote_control_websocket(
"creating new remote control enrollment: websocket_url={}, enroll_url={}, account_id={}",
remote_control_target.websocket_url, remote_control_target.enroll_url, auth.account_id
);
let new_enrollment = match enroll_remote_control_server(remote_control_target, &auth).await
let new_enrollment = match enroll_remote_control_server(
remote_control_target,
&auth,
connect_options.remote_control_instance_name,
)
.await
{
Ok(new_enrollment) => new_enrollment,
Err(err)
Expand All @@ -1083,7 +1096,7 @@ pub(super) async fn connect_remote_control_websocket(
Some(state_db),
remote_control_target,
&auth.account_id,
connect_options.app_server_client_name,
enrollment_app_server_client_name,
Some(&new_enrollment),
)
.await
Expand Down Expand Up @@ -1129,7 +1142,7 @@ pub(super) async fn connect_remote_control_websocket(
Some(state_db),
remote_control_target,
&auth.account_id,
connect_options.app_server_client_name,
enrollment_app_server_client_name,
/*enrollment*/ None,
)
.await
Expand Down Expand Up @@ -1361,6 +1374,7 @@ mod tests {
RemoteControlConnectOptions {
subscribe_cursor: None,
app_server_client_name: None,
remote_control_instance_name: None,
},
&status_publisher,
)
Expand Down Expand Up @@ -1437,6 +1451,7 @@ mod tests {
RemoteControlConnectOptions {
subscribe_cursor: None,
app_server_client_name: None,
remote_control_instance_name: None,
},
&status_publisher,
)
Expand Down Expand Up @@ -1517,6 +1532,7 @@ mod tests {
RemoteControlConnectOptions {
subscribe_cursor: None,
app_server_client_name: None,
remote_control_instance_name: None,
},
&status_publisher,
)
Expand Down Expand Up @@ -1569,6 +1585,7 @@ mod tests {
RemoteControlConnectOptions {
subscribe_cursor: None,
app_server_client_name: None,
remote_control_instance_name: None,
},
&status_publisher,
)
Expand Down Expand Up @@ -1616,6 +1633,7 @@ mod tests {
RemoteControlConnectOptions {
subscribe_cursor: None,
app_server_client_name: None,
remote_control_instance_name: None,
},
&status_publisher,
)
Expand Down Expand Up @@ -1666,6 +1684,7 @@ mod tests {
},
shutdown_token,
enabled_rx,
/*remote_control_instance_name*/ None,
)
.run(/*app_server_client_name_rx*/ None)
.await
Expand Down
Loading
Loading