diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 792e11b90a..f2b19882f1 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -532,10 +532,19 @@ pub async fn run_main_with_transport( let shutdown_when_no_connections = single_client_mode; let graceful_signal_restart_enabled = !single_client_mode; + let auth_manager = AuthManager::shared( + config.codex_home.clone(), + false, + config.cli_auth_credentials_store_mode, + ); + auth_manager.set_forced_chatgpt_workspace_id(config.forced_chatgpt_workspace_id.clone()); + let remote_control_url = config.experimental_app_server_remote_control_url.clone(); if let Some(remote_control_url) = remote_control_url { let accept_handle = start_remote_control( remote_control_url, + config.codex_home.clone(), + auth_manager.clone(), transport_event_tx.clone(), transport_shutdown_token.clone(), ) diff --git a/codex-rs/app-server/src/transport.rs b/codex-rs/app-server/src/transport.rs index 1ad29254ac..19bf16357a 100644 --- a/codex-rs/app-server/src/transport.rs +++ b/codex-rs/app-server/src/transport.rs @@ -17,6 +17,8 @@ use axum::routing::get; use codex_app_server_protocol::JSONRPCErrorError; use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::ServerRequest; +use codex_core::AuthManager; +use codex_core::default_client::build_reqwest_client; use futures::SinkExt; use futures::StreamExt; use owo_colors::OwoColorize; @@ -29,6 +31,8 @@ use std::collections::HashSet; use std::io::ErrorKind; use std::io::Result as IoResult; use std::net::SocketAddr; +use std::path::Path; +use std::path::PathBuf; use std::str::FromStr; use std::sync::Arc; use std::sync::RwLock; @@ -40,13 +44,19 @@ use tokio::io::AsyncWriteExt; use tokio::io::BufReader; use tokio::io::{self}; use tokio::net::TcpListener; +use tokio::net::TcpStream; use tokio::sync::mpsc; use tokio::task::JoinHandle; use tokio::time::Duration; use tokio::time::Instant; use tokio::time::MissedTickBehavior; +use tokio_tungstenite::MaybeTlsStream; +use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::connect_async; use tokio_tungstenite::tungstenite::Message as TungsteniteMessage; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::http::HeaderValue; +use tokio_tungstenite::tungstenite::{self}; use tokio_util::sync::CancellationToken; use tracing::debug; use tracing::error; @@ -59,8 +69,13 @@ use tracing::warn; pub(crate) const CHANNEL_CAPACITY: usize = 128; const REMOTE_CONTROL_CLIENT_IDLE_TIMEOUT: Duration = Duration::from_secs(10 * 60); const REMOTE_CONTROL_IDLE_SWEEP_INTERVAL: Duration = Duration::from_secs(30); +const REMOTE_CONTROL_ENROLL_TIMEOUT: Duration = Duration::from_secs(30); const REMOTE_CONTROL_RECONNECT_INITIAL_BACKOFF: Duration = Duration::from_secs(1); const REMOTE_CONTROL_RECONNECT_MAX_BACKOFF: Duration = Duration::from_secs(30); +const REMOTE_CONTROL_PROTOCOL_VERSION: &str = "2"; +const REMOTE_CONTROL_SERVER_NAME: &str = "codex-app-server"; +const REMOTE_CONTROL_ACCOUNT_ID_HEADER: &str = "chatgpt-account-id"; +const REMOTE_CONTROL_STATE_FILE: &str = "remote_control.toml"; fn colorize(text: &str, style: Style) -> String { text.if_supports_color(Stream::Stderr, |value| value.style(style)) @@ -712,20 +727,29 @@ pub(crate) async fn route_outgoing_envelope( #[serde(transparent)] pub struct ClientId(pub String); +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ClientActivityState { + Foreground, + Background, +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "type", rename_all = "snake_case")] pub enum ClientEvent { ClientMessage { - #[serde(rename = "clientId")] + #[serde(rename = "client_id", alias = "clientId")] client_id: ClientId, message: JSONRPCMessage, }, Ping { - #[serde(rename = "clientId")] + #[serde(rename = "client_id", alias = "clientId")] client_id: ClientId, + #[serde(skip_serializing_if = "Option::is_none")] + state: Option, }, ClientClosed { - #[serde(rename = "clientId")] + #[serde(rename = "client_id", alias = "clientId")] client_id: ClientId, }, } @@ -734,12 +758,12 @@ pub enum ClientEvent { #[serde(tag = "type", rename_all = "snake_case")] pub enum ServerEvent { ServerMessage { - #[serde(rename = "clientId")] + #[serde(rename = "client_id")] client_id: ClientId, message: Box, }, Pong { - #[serde(rename = "clientId")] + #[serde(rename = "client_id")] client_id: ClientId, status: PongStatus, }, @@ -758,12 +782,168 @@ struct RemoteControlClientState { last_activity_at: Instant, } +#[derive(Debug, Clone, PartialEq, Eq)] +struct RemoteControlTarget { + websocket_url: String, + enroll_url: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct RemoteControlEnrollment { + server_id: String, + server_name: String, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct RemoteControlConnectionAuth { + bearer_token: String, + account_id: Option, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] +struct RemoteControlStateToml { + #[serde(default)] + enrollments: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +struct PersistedRemoteControlEnrollment { + websocket_url: String, + account_id: Option, + server_id: String, + server_name: String, +} + +#[derive(Debug, Serialize)] +struct EnrollRemoteServerRequest<'a> { + name: &'a str, + os: &'a str, + arch: &'a str, + app_server_version: &'a str, +} + +#[derive(Debug, Deserialize)] +struct EnrollRemoteServerResponse { + server_id: String, +} + +fn remote_control_state_path(codex_home: &Path) -> PathBuf { + codex_home.join(REMOTE_CONTROL_STATE_FILE) +} + +fn matches_persisted_remote_control_enrollment( + entry: &PersistedRemoteControlEnrollment, + remote_control_target: &RemoteControlTarget, + account_id: Option<&str>, +) -> bool { + entry.websocket_url == remote_control_target.websocket_url + && entry.account_id.as_deref() == account_id +} + +async fn load_remote_control_state(state_path: &Path) -> IoResult { + let contents = match tokio::fs::read_to_string(state_path).await { + Ok(contents) => contents, + Err(err) if err.kind() == ErrorKind::NotFound => { + return Ok(RemoteControlStateToml::default()); + } + Err(err) => return Err(err), + }; + + toml::from_str(&contents).map_err(|err| { + std::io::Error::new( + ErrorKind::InvalidData, + format!( + "failed to parse remote control state `{}`: {err}", + state_path.display() + ), + ) + }) +} + +async fn write_remote_control_state( + state_path: &Path, + state: &RemoteControlStateToml, +) -> IoResult<()> { + if let Some(parent) = state_path.parent() { + tokio::fs::create_dir_all(parent).await?; + } + let serialized = toml::to_string(state).map_err(std::io::Error::other)?; + tokio::fs::write(state_path, serialized).await +} + +async fn load_persisted_remote_control_enrollment( + state_path: &Path, + remote_control_target: &RemoteControlTarget, + account_id: Option<&str>, +) -> Option { + let state = match load_remote_control_state(state_path).await { + Ok(state) => state, + Err(err) => { + warn!("{err}"); + return None; + } + }; + + state + .enrollments + .into_iter() + .find(|entry| { + matches_persisted_remote_control_enrollment(entry, remote_control_target, account_id) + }) + .map(|entry| RemoteControlEnrollment { + server_id: entry.server_id, + server_name: entry.server_name, + }) +} + +async fn update_persisted_remote_control_enrollment( + state_path: &Path, + remote_control_target: &RemoteControlTarget, + account_id: Option<&str>, + enrollment: Option<&RemoteControlEnrollment>, +) -> IoResult<()> { + let mut state = match load_remote_control_state(state_path).await { + Ok(state) => state, + Err(err) if err.kind() == ErrorKind::InvalidData => { + warn!("{err}"); + RemoteControlStateToml::default() + } + Err(err) => return Err(err), + }; + + state.enrollments.retain(|entry| { + !matches_persisted_remote_control_enrollment(entry, remote_control_target, account_id) + }); + + if let Some(enrollment) = enrollment { + state.enrollments.push(PersistedRemoteControlEnrollment { + websocket_url: remote_control_target.websocket_url.clone(), + account_id: account_id.map(str::to_owned), + server_id: enrollment.server_id.clone(), + server_name: enrollment.server_name.clone(), + }); + } + + if state.enrollments.is_empty() { + match tokio::fs::remove_file(state_path).await { + Ok(()) => Ok(()), + Err(err) if err.kind() == ErrorKind::NotFound => Ok(()), + Err(err) => Err(err), + } + } else { + write_remote_control_state(state_path, &state).await + } +} + pub(crate) async fn start_remote_control( remote_control_url: String, + codex_home: PathBuf, + auth_manager: Arc, transport_event_tx: mpsc::Sender, shutdown_token: CancellationToken, ) -> IoResult> { let remote_control_url = normalize_remote_control_url(&remote_control_url)?; + let remote_control_state_path = remote_control_state_path(&codex_home); Ok(tokio::spawn(async move { let local_shutdown_token = shutdown_token.child_token(); let (client_event_tx, client_event_rx) = mpsc::channel(CHANNEL_CAPACITY); @@ -772,6 +952,8 @@ pub(crate) async fn start_remote_control( let mut websocket_task = tokio::spawn(run_remote_control_websocket_loop( remote_control_url, + remote_control_state_path, + auth_manager, client_event_tx, server_event_rx, local_shutdown_token.clone(), @@ -800,15 +982,23 @@ pub(crate) async fn start_remote_control( })) } -fn normalize_remote_control_url(remote_control_url: &str) -> IoResult { +fn normalize_remote_control_url(remote_control_url: &str) -> IoResult { + let remote_control_url = remote_control_url.trim_end_matches('/'); + if remote_control_url.starts_with("ws://") || remote_control_url.starts_with("wss://") { - return Ok(remote_control_url.to_string()); + return Ok(RemoteControlTarget { + websocket_url: remote_control_url.to_string(), + enroll_url: None, + }); } + if let Some(rest) = remote_control_url.strip_prefix("http://") { - return Ok(format!("ws://{rest}")); + return Ok(normalize_http_remote_control_url(rest, "http://", "ws://")); } if let Some(rest) = remote_control_url.strip_prefix("https://") { - return Ok(format!("wss://{rest}")); + return Ok(normalize_http_remote_control_url( + rest, "https://", "wss://", + )); } Err(std::io::Error::new( @@ -819,6 +1009,29 @@ fn normalize_remote_control_url(remote_control_url: &str) -> IoResult { )) } +fn normalize_http_remote_control_url( + rest: &str, + http_scheme: &str, + websocket_scheme: &str, +) -> RemoteControlTarget { + let rest = if let Some(rest) = rest.strip_suffix("/remote/control/server/enroll") { + format!("{rest}/remote/control/server") + } else if rest.ends_with("/remote/control/server") { + rest.to_string() + } else if let Some(rest) = rest.strip_suffix("/server/enroll") { + format!("{rest}/server") + } else if rest.ends_with("/server") { + rest.to_string() + } else { + format!("{rest}/remote/control/server") + }; + + RemoteControlTarget { + websocket_url: format!("{websocket_scheme}{rest}"), + enroll_url: Some(format!("{http_scheme}{rest}/enroll")), + } +} + async fn run_remote_control_manager( transport_event_tx: mpsc::Sender, mut client_event_rx: mpsc::Receiver, @@ -918,7 +1131,7 @@ async fn run_remote_control_manager( break; } } - ClientEvent::Ping { client_id } => { + ClientEvent::Ping { client_id, .. } => { let status = match clients.get_mut(&client_id) { Some(client) => { client.last_activity_at = Instant::now(); @@ -1033,14 +1246,235 @@ async fn run_remote_control_client_outbound( let _ = writer_exited_tx.send(client_id).await; } +async fn load_remote_control_auth( + auth_manager: &AuthManager, +) -> IoResult { + let auth = match auth_manager.auth().await { + Some(auth) => auth, + None => { + auth_manager.reload(); + auth_manager.auth().await.ok_or_else(|| { + std::io::Error::new( + ErrorKind::PermissionDenied, + "remote control requires ChatGPT authentication", + ) + })? + } + }; + + if !auth.is_chatgpt_auth() { + return Err(std::io::Error::new( + ErrorKind::PermissionDenied, + "remote control requires ChatGPT authentication; API key auth is not supported", + )); + } + + Ok(RemoteControlConnectionAuth { + bearer_token: auth.get_token().map_err(std::io::Error::other)?, + account_id: auth.get_account_id(), + }) +} + +async fn enroll_remote_control_server( + remote_control_target: &RemoteControlTarget, + auth: &RemoteControlConnectionAuth, +) -> IoResult { + let Some(enroll_url) = remote_control_target.enroll_url.as_deref() else { + return Err(std::io::Error::other( + "remote control enrollment requires an HTTP(S) URL", + )); + }; + + let request = EnrollRemoteServerRequest { + name: REMOTE_CONTROL_SERVER_NAME, + os: std::env::consts::OS, + arch: std::env::consts::ARCH, + app_server_version: env!("CARGO_PKG_VERSION"), + }; + let client = build_reqwest_client(); + let mut http_request = client + .post(enroll_url) + .timeout(REMOTE_CONTROL_ENROLL_TIMEOUT) + .bearer_auth(&auth.bearer_token) + .json(&request); + if let Some(account_id) = auth.account_id.as_deref() { + http_request = http_request.header(REMOTE_CONTROL_ACCOUNT_ID_HEADER, account_id); + } + + let response = http_request.send().await.map_err(|err| { + std::io::Error::other(format!( + "failed to enroll remote control server at `{enroll_url}`: {err}" + )) + })?; + let status = response.status(); + if !status.is_success() { + let body = response.text().await.unwrap_or_default(); + return Err(std::io::Error::other(format!( + "remote control server enrollment failed at `{enroll_url}`: HTTP {status} {body}" + ))); + } + + let enrollment = response + .json::() + .await + .map_err(|err| { + std::io::Error::other(format!( + "failed to parse remote control enrollment response from `{enroll_url}`: {err}" + )) + })?; + + Ok(RemoteControlEnrollment { + server_id: enrollment.server_id, + server_name: REMOTE_CONTROL_SERVER_NAME.to_string(), + }) +} + +fn set_remote_control_header( + headers: &mut tungstenite::http::HeaderMap, + name: &'static str, + value: &str, +) -> IoResult<()> { + let header_value = HeaderValue::from_str(value).map_err(|err| { + std::io::Error::new( + ErrorKind::InvalidInput, + format!("invalid remote control header `{name}`: {err}"), + ) + })?; + headers.insert(name, header_value); + Ok(()) +} + +fn build_remote_control_websocket_request( + websocket_url: &str, + enrollment: &RemoteControlEnrollment, + auth: &RemoteControlConnectionAuth, +) -> IoResult> { + let mut request = websocket_url.into_client_request().map_err(|err| { + std::io::Error::new( + ErrorKind::InvalidInput, + format!("invalid remote control websocket URL `{websocket_url}`: {err}"), + ) + })?; + let headers = request.headers_mut(); + set_remote_control_header(headers, "x-codex-server-id", &enrollment.server_id)?; + set_remote_control_header(headers, "x-codex-name", &enrollment.server_name)?; + set_remote_control_header( + headers, + "x-codex-protocol-version", + REMOTE_CONTROL_PROTOCOL_VERSION, + )?; + set_remote_control_header( + headers, + "authorization", + &format!("Bearer {}", auth.bearer_token), + )?; + if let Some(account_id) = auth.account_id.as_deref() { + set_remote_control_header(headers, REMOTE_CONTROL_ACCOUNT_ID_HEADER, account_id)?; + } + Ok(request) +} + +async fn connect_remote_control_websocket( + remote_control_target: &RemoteControlTarget, + remote_control_state_path: &Path, + auth_manager: &AuthManager, + enrollment: &mut Option, +) -> IoResult>> { + if remote_control_target.enroll_url.is_none() { + return connect_async(remote_control_target.websocket_url.as_str()) + .await + .map(|(websocket_stream, _response)| websocket_stream) + .map_err(|err| { + std::io::Error::other(format!( + "failed to connect app-server remote control websocket `{}`: {err}", + remote_control_target.websocket_url + )) + }); + } + + let auth = load_remote_control_auth(auth_manager).await?; + if enrollment.is_none() { + *enrollment = load_persisted_remote_control_enrollment( + remote_control_state_path, + remote_control_target, + auth.account_id.as_deref(), + ) + .await; + } + + if enrollment.is_none() { + let new_enrollment = enroll_remote_control_server(remote_control_target, &auth).await?; + if let Err(err) = update_persisted_remote_control_enrollment( + remote_control_state_path, + remote_control_target, + auth.account_id.as_deref(), + Some(&new_enrollment), + ) + .await + { + warn!( + "failed to persist remote control enrollment in `{}`: {err}", + remote_control_state_path.display() + ); + } + *enrollment = Some(new_enrollment); + } + + let enrollment_ref = match enrollment.as_ref() { + Some(enrollment) => enrollment, + None => { + return Err(std::io::Error::other( + "missing remote control enrollment after enrollment step", + )); + } + }; + let request = build_remote_control_websocket_request( + &remote_control_target.websocket_url, + enrollment_ref, + &auth, + )?; + + match connect_async(request).await { + Ok((websocket_stream, _response)) => Ok(websocket_stream), + Err(err) => { + if matches!( + &err, + tungstenite::Error::Http(response) if response.status().as_u16() == 404 + ) { + if let Err(clear_err) = update_persisted_remote_control_enrollment( + remote_control_state_path, + remote_control_target, + auth.account_id.as_deref(), + None, + ) + .await + { + warn!( + "failed to clear stale remote control enrollment in `{}`: {clear_err}", + remote_control_state_path.display() + ); + } + *enrollment = None; + } + Err(std::io::Error::other(format!( + "failed to connect app-server remote control websocket `{}`: {err}", + remote_control_target.websocket_url + ))) + } + } +} + async fn run_remote_control_websocket_loop( - remote_control_url: String, + remote_control_target: RemoteControlTarget, + remote_control_state_path: PathBuf, + auth_manager: Arc, client_event_tx: mpsc::Sender, mut server_event_rx: mpsc::Receiver, shutdown_token: CancellationToken, ) { let mut reconnect_backoff = REMOTE_CONTROL_RECONNECT_INITIAL_BACKOFF; let mut wait_before_connect = false; + let mut enrollment = None::; let mut pending_server_event = None::; loop { @@ -1061,15 +1495,23 @@ async fn run_remote_control_websocket_loop( _ = shutdown_token.cancelled() => { break; } - connect_result = connect_async(remote_control_url.as_str()) => { + connect_result = connect_remote_control_websocket( + &remote_control_target, + remote_control_state_path.as_path(), + auth_manager.as_ref(), + &mut enrollment, + ) => { match connect_result { - Ok((websocket_stream, _response)) => { + Ok(websocket_stream) => { reconnect_backoff = REMOTE_CONTROL_RECONNECT_INITIAL_BACKOFF; - info!("connected to app-server remote control websocket: {remote_control_url}"); + info!( + "connected to app-server remote control websocket: {}", + remote_control_target.websocket_url + ); websocket_stream } Err(err) => { - warn!("failed to connect app-server remote control websocket `{remote_control_url}`: {err}"); + warn!("{err}"); if slept_before_connect { reconnect_backoff = reconnect_backoff .saturating_mul(2) @@ -1115,8 +1557,8 @@ async fn run_remote_control_websocket_loop( return; } } - Err(err) => { - warn!("failed to deserialize remote-control client event: {err}"); + Err(_) => { + warn!("failed to deserialize remote-control client event"); } } } @@ -1160,19 +1602,40 @@ mod tests { use super::*; use crate::error_code::OVERLOADED_ERROR_CODE; use codex_app_server_protocol::CommandExecutionRequestApprovalSkillMetadata; + use codex_core::CodexAuth; + use codex_core::test_support::auth_manager_from_auth; + use codex_core::test_support::auth_manager_from_auth_with_home; use codex_utils_absolute_path::AbsolutePathBuf; use pretty_assertions::assert_eq; use serde_json::json; + use std::collections::BTreeMap; use std::path::PathBuf; + use tempfile::TempDir; + use tokio::io::AsyncBufReadExt; + use tokio::io::AsyncReadExt; + use tokio::io::AsyncWriteExt; + use tokio::io::BufReader; use tokio::net::TcpStream; use tokio::time::timeout; use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::accept_async; + use tokio_tungstenite::accept_hdr_async; fn absolute_path(path: &str) -> AbsolutePathBuf { AbsolutePathBuf::from_absolute_path(path).expect("absolute path") } + fn remote_control_auth_manager() -> Arc { + auth_manager_from_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + } + + fn remote_control_auth_manager_with_home(codex_home: &TempDir) -> Arc { + auth_manager_from_auth_with_home( + CodexAuth::create_dummy_chatgpt_auth_for_testing(), + codex_home.path().to_path_buf(), + ) + } + #[test] fn app_server_transport_parses_stdio_listen_url() { let transport = AppServerTransport::from_listen_url(AppServerTransport::DEFAULT_LISTEN_URL) @@ -1820,19 +2283,49 @@ mod tests { fn normalize_remote_control_url_rewrites_http_schemes() { assert_eq!( normalize_remote_control_url("ws://example.com/control").expect("valid ws url"), - "ws://example.com/control" + RemoteControlTarget { + websocket_url: "ws://example.com/control".to_string(), + enroll_url: None, + } ); assert_eq!( normalize_remote_control_url("wss://example.com/control").expect("valid wss url"), - "wss://example.com/control" + RemoteControlTarget { + websocket_url: "wss://example.com/control".to_string(), + enroll_url: None, + } ); assert_eq!( - normalize_remote_control_url("http://example.com/control").expect("valid http url"), - "ws://example.com/control" + normalize_remote_control_url("http://example.com/backend-api/wham") + .expect("valid http prefix"), + RemoteControlTarget { + websocket_url: "ws://example.com/backend-api/wham/remote/control/server" + .to_string(), + enroll_url: Some( + "http://example.com/backend-api/wham/remote/control/server/enroll".to_string(), + ), + } ); assert_eq!( - normalize_remote_control_url("https://example.com/control").expect("valid https url"), - "wss://example.com/control" + normalize_remote_control_url( + "https://example.com/backend-api/wham/remote/control/server" + ) + .expect("valid https full path"), + RemoteControlTarget { + websocket_url: "wss://example.com/backend-api/wham/remote/control/server" + .to_string(), + enroll_url: Some( + "https://example.com/backend-api/wham/remote/control/server/enroll".to_string(), + ), + } + ); + assert_eq!( + normalize_remote_control_url("http://example.com/legacy/server") + .expect("valid legacy http url"), + RemoteControlTarget { + websocket_url: "ws://example.com/legacy/server".to_string(), + enroll_url: Some("http://example.com/legacy/server/enroll".to_string()), + } ); } @@ -1846,6 +2339,132 @@ mod tests { ); } + #[tokio::test] + async fn persisted_remote_control_enrollment_round_trips_by_target_and_account() { + let codex_home = TempDir::new().expect("temp dir should create"); + let state_path = remote_control_state_path(codex_home.path()); + let first_target = normalize_remote_control_url("http://example.com/remote/control") + .expect("first target should parse"); + let second_target = normalize_remote_control_url("http://example.com/other/control") + .expect("second target should parse"); + let first_enrollment = RemoteControlEnrollment { + server_id: "srv_e_first".to_string(), + server_name: REMOTE_CONTROL_SERVER_NAME.to_string(), + }; + let second_enrollment = RemoteControlEnrollment { + server_id: "srv_e_second".to_string(), + server_name: REMOTE_CONTROL_SERVER_NAME.to_string(), + }; + + update_persisted_remote_control_enrollment( + state_path.as_path(), + &first_target, + Some("account-a"), + Some(&first_enrollment), + ) + .await + .expect("first enrollment should persist"); + update_persisted_remote_control_enrollment( + state_path.as_path(), + &second_target, + Some("account-a"), + Some(&second_enrollment), + ) + .await + .expect("second enrollment should persist"); + + assert_eq!( + load_persisted_remote_control_enrollment( + state_path.as_path(), + &first_target, + Some("account-a"), + ) + .await, + Some(first_enrollment.clone()) + ); + assert_eq!( + load_persisted_remote_control_enrollment( + state_path.as_path(), + &first_target, + Some("account-b"), + ) + .await, + None + ); + assert_eq!( + load_persisted_remote_control_enrollment( + state_path.as_path(), + &second_target, + Some("account-a"), + ) + .await, + Some(second_enrollment) + ); + } + + #[tokio::test] + async fn clearing_persisted_remote_control_enrollment_removes_only_matching_entry() { + let codex_home = TempDir::new().expect("temp dir should create"); + let state_path = remote_control_state_path(codex_home.path()); + let first_target = normalize_remote_control_url("http://example.com/remote/control") + .expect("first target should parse"); + let second_target = normalize_remote_control_url("http://example.com/other/control") + .expect("second target should parse"); + let first_enrollment = RemoteControlEnrollment { + server_id: "srv_e_first".to_string(), + server_name: REMOTE_CONTROL_SERVER_NAME.to_string(), + }; + let second_enrollment = RemoteControlEnrollment { + server_id: "srv_e_second".to_string(), + server_name: REMOTE_CONTROL_SERVER_NAME.to_string(), + }; + + update_persisted_remote_control_enrollment( + state_path.as_path(), + &first_target, + Some("account-a"), + Some(&first_enrollment), + ) + .await + .expect("first enrollment should persist"); + update_persisted_remote_control_enrollment( + state_path.as_path(), + &second_target, + Some("account-a"), + Some(&second_enrollment), + ) + .await + .expect("second enrollment should persist"); + + update_persisted_remote_control_enrollment( + state_path.as_path(), + &first_target, + Some("account-a"), + None, + ) + .await + .expect("matching enrollment should clear"); + + assert_eq!( + load_persisted_remote_control_enrollment( + state_path.as_path(), + &first_target, + Some("account-a"), + ) + .await, + None + ); + assert_eq!( + load_persisted_remote_control_enrollment( + state_path.as_path(), + &second_target, + Some("account-a"), + ) + .await, + Some(second_enrollment) + ); + } + #[test] fn remote_control_client_is_alive_respects_activity_timeout() { let base = Instant::now(); @@ -1876,11 +2495,14 @@ mod tests { .local_addr() .expect("listener should have a local addr") ); + let codex_home = TempDir::new().expect("temp dir should create"); let (transport_event_tx, mut transport_event_rx) = mpsc::channel::(CHANNEL_CAPACITY); let shutdown_token = CancellationToken::new(); let remote_handle = start_remote_control( remote_control_url, + codex_home.path().to_path_buf(), + remote_control_auth_manager(), transport_event_tx, shutdown_token.clone(), ) @@ -1893,6 +2515,7 @@ mod tests { &mut websocket, ClientEvent::Ping { client_id: client_id.clone(), + state: Some(ClientActivityState::Foreground), }, ) .await; @@ -1900,7 +2523,7 @@ mod tests { read_server_event(&mut websocket).await, json!({ "type": "pong", - "clientId": "client-1", + "client_id": "client-1", "status": "unknown", }) ); @@ -2007,6 +2630,7 @@ mod tests { &mut websocket, ClientEvent::Ping { client_id: client_id.clone(), + state: Some(ClientActivityState::Foreground), }, ) .await; @@ -2014,7 +2638,7 @@ mod tests { read_server_event(&mut websocket).await, json!({ "type": "pong", - "clientId": "client-1", + "client_id": "client-1", "status": "active", }) ); @@ -2032,7 +2656,7 @@ mod tests { read_server_event(&mut websocket).await, json!({ "type": "server_message", - "clientId": "client-1", + "client_id": "client-1", "message": { "method": "codex/event/test", "params": { @@ -2062,12 +2686,19 @@ mod tests { other => panic!("expected connection close event, got {other:?}"), } - send_client_event(&mut websocket, ClientEvent::Ping { client_id }).await; + send_client_event( + &mut websocket, + ClientEvent::Ping { + client_id, + state: Some(ClientActivityState::Foreground), + }, + ) + .await; assert_eq!( read_server_event(&mut websocket).await, json!({ "type": "pong", - "clientId": "client-1", + "client_id": "client-1", "status": "unknown", }) ); @@ -2087,11 +2718,14 @@ mod tests { .local_addr() .expect("listener should have a local addr") ); + let codex_home = TempDir::new().expect("temp dir should create"); let (transport_event_tx, mut transport_event_rx) = mpsc::channel::(CHANNEL_CAPACITY); let shutdown_token = CancellationToken::new(); let remote_handle = start_remote_control( remote_control_url, + codex_home.path().to_path_buf(), + remote_control_auth_manager(), transport_event_tx, shutdown_token.clone(), ) @@ -2138,6 +2772,356 @@ mod tests { let _ = remote_handle.await; } + #[tokio::test] + async fn remote_control_http_mode_enrolls_before_connecting() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let remote_control_url = format!( + "http://{}/backend-api/wham", + listener + .local_addr() + .expect("listener should have a local addr") + ); + let codex_home = TempDir::new().expect("temp dir should create"); + let (transport_event_tx, mut transport_event_rx) = + mpsc::channel::(CHANNEL_CAPACITY); + let shutdown_token = CancellationToken::new(); + let remote_handle = start_remote_control( + remote_control_url, + codex_home.path().to_path_buf(), + remote_control_auth_manager(), + transport_event_tx, + shutdown_token.clone(), + ) + .await + .expect("remote control should start"); + + 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" + ); + assert_eq!( + enroll_request.headers.get("authorization"), + Some(&"Bearer Access Token".to_string()) + ); + assert_eq!( + enroll_request.headers.get(REMOTE_CONTROL_ACCOUNT_ID_HEADER), + Some(&"account_id".to_string()) + ); + assert_eq!( + serde_json::from_str::(&enroll_request.body) + .expect("enroll body should deserialize"), + json!({ + "name": REMOTE_CONTROL_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_test" })).await; + + let (handshake_request, mut websocket) = + accept_remote_control_backend_connection(&listener).await; + assert_eq!( + handshake_request.path, + "/backend-api/wham/remote/control/server" + ); + assert_eq!( + handshake_request.headers.get("authorization"), + Some(&"Bearer Access Token".to_string()) + ); + assert_eq!( + handshake_request + .headers + .get(REMOTE_CONTROL_ACCOUNT_ID_HEADER), + Some(&"account_id".to_string()) + ); + assert_eq!( + handshake_request.headers.get("x-codex-server-id"), + Some(&"srv_e_test".to_string()) + ); + assert_eq!( + handshake_request.headers.get("x-codex-name"), + Some(&REMOTE_CONTROL_SERVER_NAME.to_string()) + ); + assert_eq!( + handshake_request.headers.get("x-codex-protocol-version"), + Some(&REMOTE_CONTROL_PROTOCOL_VERSION.to_string()) + ); + + let backend_client_id = ClientId("backend-test-client".to_string()); + let writer = { + let initialize_message = + JSONRPCMessage::Request(codex_app_server_protocol::JSONRPCRequest { + id: codex_app_server_protocol::RequestId::Integer(11), + method: "initialize".to_string(), + params: Some(json!({ + "clientInfo": { + "name": "remote-backend-client", + "version": "0.1.0" + } + })), + trace: None, + }); + send_client_event( + &mut websocket, + ClientEvent::ClientMessage { + client_id: backend_client_id.clone(), + message: initialize_message.clone(), + }, + ) + .await; + + let (connection_id, writer) = + match timeout(Duration::from_secs(5), transport_event_rx.recv()) + .await + .expect("connection open should arrive in time") + .expect("connection open should exist") + { + TransportEvent::ConnectionOpened { + connection_id, + writer, + .. + } => (connection_id, writer), + other => panic!("expected connection open event, got {other:?}"), + }; + + match timeout(Duration::from_secs(5), transport_event_rx.recv()) + .await + .expect("initialize message should arrive in time") + .expect("initialize message should exist") + { + TransportEvent::IncomingMessage { + connection_id: incoming_connection_id, + message, + } => { + assert_eq!(incoming_connection_id, connection_id); + assert_eq!(message, initialize_message); + } + other => panic!("expected initialize incoming message, got {other:?}"), + } + writer + }; + + writer + .send(OutgoingMessage::Response( + crate::outgoing_message::OutgoingResponse { + id: codex_app_server_protocol::RequestId::Integer(11), + result: json!({ + "userAgent": "codex-test-agent" + }), + }, + )) + .await + .expect("remote writer should accept initialize response"); + assert_eq!( + read_server_event(&mut websocket).await, + json!({ + "type": "server_message", + "client_id": backend_client_id.0.clone(), + "message": { + "id": 11, + "result": { + "userAgent": "codex-test-agent", + } + } + }) + ); + + writer + .send(OutgoingMessage::Notification( + crate::outgoing_message::OutgoingNotification { + method: "codex/event/test".to_string(), + params: Some(json!({ "backend": true })), + }, + )) + .await + .expect("remote writer should accept outgoing message"); + assert_eq!( + read_server_event(&mut websocket).await, + json!({ + "type": "server_message", + "client_id": backend_client_id.0.clone(), + "message": { + "method": "codex/event/test", + "params": { + "backend": true, + } + } + }) + ); + + shutdown_token.cancel(); + let _ = remote_handle.await; + } + + #[tokio::test] + async fn remote_control_http_mode_reuses_persisted_enrollment_before_reenrolling() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let remote_control_url = format!( + "http://{}/backend-api/wham", + listener + .local_addr() + .expect("listener should have a local addr") + ); + let codex_home = TempDir::new().expect("temp dir should create"); + let remote_control_target = + normalize_remote_control_url(&remote_control_url).expect("target should parse"); + let persisted_enrollment = RemoteControlEnrollment { + server_id: "srv_e_persisted".to_string(), + server_name: REMOTE_CONTROL_SERVER_NAME.to_string(), + }; + update_persisted_remote_control_enrollment( + remote_control_state_path(codex_home.path()).as_path(), + &remote_control_target, + Some("account_id"), + Some(&persisted_enrollment), + ) + .await + .expect("persisted enrollment should save"); + + let (transport_event_tx, _transport_event_rx) = + mpsc::channel::(CHANNEL_CAPACITY); + let shutdown_token = CancellationToken::new(); + let remote_handle = start_remote_control( + remote_control_url, + codex_home.path().to_path_buf(), + remote_control_auth_manager_with_home(&codex_home), + transport_event_tx, + shutdown_token.clone(), + ) + .await + .expect("remote control should start"); + + let (handshake_request, _websocket) = + accept_remote_control_backend_connection(&listener).await; + assert_eq!( + handshake_request.path, + "/backend-api/wham/remote/control/server" + ); + assert_eq!( + handshake_request.headers.get("x-codex-server-id"), + Some(&persisted_enrollment.server_id) + ); + assert_eq!( + load_persisted_remote_control_enrollment( + remote_control_state_path(codex_home.path()).as_path(), + &remote_control_target, + Some("account_id"), + ) + .await, + Some(persisted_enrollment) + ); + + shutdown_token.cancel(); + let _ = remote_handle.await; + } + + #[tokio::test] + async fn remote_control_http_mode_clears_stale_persisted_enrollment_after_404() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let remote_control_url = format!( + "http://{}/backend-api/wham", + listener + .local_addr() + .expect("listener should have a local addr") + ); + let codex_home = TempDir::new().expect("temp dir should create"); + let state_path = remote_control_state_path(codex_home.path()); + let remote_control_target = + normalize_remote_control_url(&remote_control_url).expect("target should parse"); + let stale_enrollment = RemoteControlEnrollment { + server_id: "srv_e_stale".to_string(), + server_name: REMOTE_CONTROL_SERVER_NAME.to_string(), + }; + let refreshed_enrollment = RemoteControlEnrollment { + server_id: "srv_e_refreshed".to_string(), + server_name: REMOTE_CONTROL_SERVER_NAME.to_string(), + }; + update_persisted_remote_control_enrollment( + state_path.as_path(), + &remote_control_target, + Some("account_id"), + Some(&stale_enrollment), + ) + .await + .expect("stale enrollment should save"); + + let (transport_event_tx, _transport_event_rx) = + mpsc::channel::(CHANNEL_CAPACITY); + let shutdown_token = CancellationToken::new(); + let remote_handle = start_remote_control( + remote_control_url, + codex_home.path().to_path_buf(), + remote_control_auth_manager_with_home(&codex_home), + transport_event_tx, + shutdown_token.clone(), + ) + .await + .expect("remote control should start"); + + let websocket_request = accept_http_request(&listener).await; + assert_eq!( + websocket_request.request_line, + "GET /backend-api/wham/remote/control/server HTTP/1.1" + ); + assert_eq!( + websocket_request.headers.get("x-codex-server-id"), + Some(&stale_enrollment.server_id) + ); + respond_with_status(websocket_request.stream, "404 Not Found", "").await; + + let enroll_request = accept_http_request(&listener).await; + assert_eq!( + enroll_request.request_line, + "POST /backend-api/wham/remote/control/server/enroll HTTP/1.1" + ); + respond_with_json( + enroll_request.stream, + json!({ "server_id": refreshed_enrollment.server_id }), + ) + .await; + + let (handshake_request, _websocket) = + accept_remote_control_backend_connection(&listener).await; + assert_eq!( + handshake_request.headers.get("x-codex-server-id"), + Some(&refreshed_enrollment.server_id) + ); + assert_eq!( + load_persisted_remote_control_enrollment( + state_path.as_path(), + &remote_control_target, + Some("account_id"), + ) + .await, + Some(refreshed_enrollment) + ); + + shutdown_token.cancel(); + let _ = remote_handle.await; + } + + #[derive(Debug)] + struct CapturedHttpRequest { + stream: TcpStream, + request_line: String, + headers: BTreeMap, + body: String, + } + + #[derive(Clone, Debug, PartialEq, Eq)] + struct CapturedWebSocketRequest { + path: String, + headers: BTreeMap, + } + async fn accept_remote_control_connection( listener: &TcpListener, ) -> WebSocketStream { @@ -2150,6 +3134,123 @@ mod tests { .expect("websocket handshake should succeed") } + async fn accept_http_request(listener: &TcpListener) -> CapturedHttpRequest { + let (stream, _) = timeout(Duration::from_secs(5), listener.accept()) + .await + .expect("HTTP request should arrive in time") + .expect("listener accept should succeed"); + let mut reader = BufReader::new(stream); + + let mut request_line = String::new(); + reader + .read_line(&mut request_line) + .await + .expect("request line should read"); + let request_line = request_line.trim_end_matches("\r\n").to_string(); + + let mut headers = BTreeMap::new(); + loop { + let mut line = String::new(); + reader + .read_line(&mut line) + .await + .expect("header line should read"); + if line == "\r\n" { + break; + } + let line = line.trim_end_matches("\r\n"); + let (name, value) = line.split_once(':').expect("header should contain colon"); + headers.insert(name.to_ascii_lowercase(), value.trim().to_string()); + } + + let content_length = headers + .get("content-length") + .and_then(|value| value.parse::().ok()) + .unwrap_or(0); + let mut body = vec![0; content_length]; + reader + .read_exact(&mut body) + .await + .expect("request body should read"); + + CapturedHttpRequest { + stream: reader.into_inner(), + request_line, + headers, + body: String::from_utf8(body).expect("body should be utf-8"), + } + } + + async fn respond_with_json(mut stream: TcpStream, body: serde_json::Value) { + let body = body.to_string(); + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ); + stream + .write_all(response.as_bytes()) + .await + .expect("response should write"); + stream.flush().await.expect("response should flush"); + } + + async fn respond_with_status(mut stream: TcpStream, status: &str, body: &str) { + let response = format!( + "HTTP/1.1 {status}\r\ncontent-type: text/plain\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ); + stream + .write_all(response.as_bytes()) + .await + .expect("response should write"); + stream.flush().await.expect("response should flush"); + } + + async fn accept_remote_control_backend_connection( + listener: &TcpListener, + ) -> (CapturedWebSocketRequest, WebSocketStream) { + let (stream, _) = timeout(Duration::from_secs(5), listener.accept()) + .await + .expect("websocket request should arrive in time") + .expect("listener accept should succeed"); + let captured_request = Arc::new(std::sync::Mutex::new(None::)); + let captured_request_for_callback = captured_request.clone(); + let websocket = accept_hdr_async( + stream, + move |request: &tungstenite::handshake::server::Request, + response: tungstenite::handshake::server::Response| { + let headers = request + .headers() + .iter() + .map(|(name, value)| { + ( + name.as_str().to_ascii_lowercase(), + value + .to_str() + .expect("header should be valid utf-8") + .to_string(), + ) + }) + .collect::>(); + *captured_request_for_callback + .lock() + .expect("capture lock should acquire") = Some(CapturedWebSocketRequest { + path: request.uri().path().to_string(), + headers, + }); + Ok(response) + }, + ) + .await + .expect("websocket handshake should succeed"); + let captured_request = captured_request + .lock() + .expect("capture lock should acquire") + .clone() + .expect("websocket request should be captured"); + (captured_request, websocket) + } + async fn send_client_event( websocket: &mut WebSocketStream, client_event: ClientEvent, diff --git a/codex-rs/cloud-requirements/src/lib.rs b/codex-rs/cloud-requirements/src/lib.rs index 8c019f7a04..4c695751bf 100644 --- a/codex-rs/cloud-requirements/src/lib.rs +++ b/codex-rs/cloud-requirements/src/lib.rs @@ -328,6 +328,7 @@ impl CloudRequirementsService { return Ok(None); }; if !auth.is_chatgpt_auth() + || auth.is_external_chatgpt_tokens() || !matches!( auth.account_plan_type(), Some(PlanType::Business | PlanType::Enterprise) @@ -967,6 +968,42 @@ mod tests { auth_manager_with_plan_and_identity(plan_type, Some("user-12345"), Some("account-12345")) } + fn auth_manager_with_external_chatgpt_tokens(plan_type: &str) -> Arc { + let tmp = tempdir().expect("tempdir"); + let header = json!({ "alg": "none", "typ": "JWT" }); + let auth_payload = json!({ + "chatgpt_plan_type": plan_type, + "chatgpt_user_id": "user-12345", + "user_id": "user-12345", + }); + let payload = json!({ + "email": "user@example.com", + "https://api.openai.com/auth": auth_payload, + }); + let header_b64 = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&header).expect("header")); + let payload_b64 = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&payload).expect("payload")); + let signature_b64 = URL_SAFE_NO_PAD.encode(b"sig"); + let fake_jwt = format!("{header_b64}.{payload_b64}.{signature_b64}"); + + let auth_json = json!({ + "auth_mode": "chatgptAuthTokens", + "OPENAI_API_KEY": null, + "tokens": { + "id_token": fake_jwt, + "access_token": "test-access-token", + "refresh_token": "", + "account_id": "account-12345", + }, + "last_refresh": "2025-01-01T00:00:00Z", + }); + write_auth_json(tmp.path(), auth_json).expect("write auth"); + Arc::new(AuthManager::new( + tmp.path().to_path_buf(), + false, + AuthCredentialsStoreMode::File, + )) + } + fn parse_for_fetch(contents: Option<&str>) -> Option { contents.and_then(|contents| parse_cloud_requirements(contents).ok().flatten()) } @@ -1088,6 +1125,21 @@ mod tests { assert_eq!(result, Ok(None)); } + #[tokio::test] + async fn fetch_cloud_requirements_skips_external_chatgpt_tokens_auth() { + let codex_home = tempdir().expect("tempdir"); + let service = CloudRequirementsService::new( + auth_manager_with_external_chatgpt_tokens("enterprise"), + Arc::new(StaticFetcher { + contents: Some("allowed_approval_policies = [\"never\"]".to_string()), + }), + codex_home.path().to_path_buf(), + CLOUD_REQUIREMENTS_TIMEOUT, + ); + let result = service.fetch().await; + assert_eq!(result, Ok(None)); + } + #[tokio::test] async fn fetch_cloud_requirements_skips_non_business_or_enterprise_plan() { let codex_home = tempdir().expect("tempdir");