diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index 00c2dc2c89..6466749cdb 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -531,13 +531,13 @@ pub async fn run_main_with_transport( let feedback_layer = feedback.logger_layer(); let feedback_metadata_layer = feedback.metadata_layer(); - let log_db = codex_state::StateRuntime::init( + let state_db = codex_state::StateRuntime::init( config.sqlite_home.clone(), config.model_provider_id.clone(), ) .await - .ok() - .map(log_db::start); + .ok(); + let log_db = state_db.clone().map(log_db::start); let log_db_layer = log_db .clone() .map(|layer| layer.with_filter(Targets::new().with_default(Level::TRACE))); @@ -595,7 +595,7 @@ pub async fn run_main_with_transport( validate_remote_control_auth(auth_manager.as_ref()).await?; let accept_handle = start_remote_control( remote_control_config.base_url, - config.codex_home.clone(), + state_db.clone(), auth_manager.clone(), transport_event_tx.clone(), transport_shutdown_token.clone(), diff --git a/codex-rs/app-server/src/transport/remote_control/enroll.rs b/codex-rs/app-server/src/transport/remote_control/enroll.rs index b32a94fe71..56f309a7d0 100644 --- a/codex-rs/app-server/src/transport/remote_control/enroll.rs +++ b/codex-rs/app-server/src/transport/remote_control/enroll.rs @@ -1,18 +1,14 @@ use super::protocol::EnrollRemoteServerRequest; use super::protocol::EnrollRemoteServerResponse; -use super::protocol::PersistedRemoteControlEnrollment; -use super::protocol::RemoteControlStateToml; use super::protocol::RemoteControlTarget; use base64::Engine; use codex_core::AuthManager; use codex_core::default_client::build_reqwest_client; -use codex_core::path_utils::write_atomically; +use codex_state::StateRuntime; use codex_utils_rustls_provider::ensure_rustls_crypto_provider; use gethostname::gethostname; use io::ErrorKind; use std::io; -use std::path::Path; -use std::path::PathBuf; use tokio::net::TcpStream; use tokio_tungstenite::MaybeTlsStream; use tokio_tungstenite::WebSocketStream; @@ -27,7 +23,6 @@ const REMOTE_CONTROL_RESPONSE_BODY_MAX_BYTES: usize = 4096; pub(super) const REMOTE_CONTROL_PROTOCOL_VERSION: &str = "2"; pub(super) const REMOTE_CONTROL_ACCOUNT_ID_HEADER: &str = "chatgpt-account-id"; const REMOTE_CONTROL_SUBSCRIBE_CURSOR_HEADER: &str = "x-codex-subscribe-cursor"; -const REMOTE_CONTROL_STATE_FILE: &str = "remote_control.toml"; #[derive(Debug, Clone, PartialEq, Eq)] pub(super) struct RemoteControlEnrollment { @@ -45,112 +40,55 @@ pub(super) struct RemoteControlWebsocketConnection { pub(super) websocket_stream: WebSocketStream>, } -pub(super) 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) -> io::Result { - 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| { - 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, -) -> io::Result<()> { - if let Some(parent) = state_path.parent() { - tokio::fs::create_dir_all(parent).await?; - } - let state_path = state_path.to_owned(); - let contents: String = toml::to_string(state).map_err(io::Error::other)?; - tokio::task::spawn_blocking(move || write_atomically(&state_path, &contents)).await? -} - pub(super) async fn load_persisted_remote_control_enrollment( - state_path: &Path, + state_db: Option<&StateRuntime>, remote_control_target: &RemoteControlTarget, account_id: Option<&str>, ) -> Option { - let state = match load_remote_control_state(state_path).await { - Ok(state) => state, + let state_db = state_db?; + let enrollment = match state_db + .get_remote_control_enrollment(&remote_control_target.websocket_url, account_id) + .await + { + Ok(enrollment) => enrollment, 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, - }) + enrollment.map(|(server_id, server_name)| RemoteControlEnrollment { + server_id, + server_name, + }) } pub(super) async fn update_persisted_remote_control_enrollment( - state_path: &Path, + state_db: Option<&StateRuntime>, remote_control_target: &RemoteControlTarget, account_id: Option<&str>, enrollment: Option<&RemoteControlEnrollment>, ) -> io::Result<()> { - 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), + let Some(state_db) = state_db else { + return Ok(()); }; - 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), - } + state_db + .upsert_remote_control_enrollment( + &remote_control_target.websocket_url, + account_id, + &enrollment.server_id, + &enrollment.server_name, + ) + .await + .map_err(io::Error::other) } else { - write_remote_control_state(state_path, &state).await + state_db + .delete_remote_control_enrollment(&remote_control_target.websocket_url, account_id) + .await + .map(|_| ()) + .map_err(io::Error::other) } } @@ -346,7 +284,7 @@ fn build_remote_control_websocket_request( pub(super) async fn connect_remote_control_websocket( remote_control_target: &RemoteControlTarget, - remote_control_state_path: &Path, + state_db: Option<&StateRuntime>, auth_manager: &AuthManager, enrollment: &mut Option, subscribe_cursor: Option<&str>, @@ -356,7 +294,7 @@ pub(super) async fn connect_remote_control_websocket( let auth = load_remote_control_auth(auth_manager).await?; if enrollment.is_none() { *enrollment = load_persisted_remote_control_enrollment( - remote_control_state_path, + state_db, remote_control_target, auth.account_id.as_deref(), ) @@ -366,17 +304,14 @@ pub(super) async fn connect_remote_control_websocket( 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, + state_db, 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() - ); + warn!("failed to persist remote control enrollment in sqlite state db: {err}"); } *enrollment = Some(new_enrollment); } @@ -399,7 +334,7 @@ pub(super) async fn connect_remote_control_websocket( tungstenite::Error::Http(response) if response.status().as_u16() == 404 ) { if let Err(clear_err) = update_persisted_remote_control_enrollment( - remote_control_state_path, + state_db, remote_control_target, auth.account_id.as_deref(), /*enrollment*/ None, @@ -407,8 +342,7 @@ pub(super) async fn connect_remote_control_websocket( .await { warn!( - "failed to clear stale remote control enrollment in `{}`: {clear_err}", - remote_control_state_path.display() + "failed to clear stale remote control enrollment in sqlite state db: {clear_err}" ); } *enrollment = None; diff --git a/codex-rs/app-server/src/transport/remote_control/mod.rs b/codex-rs/app-server/src/transport/remote_control/mod.rs index a056f64043..a85406deb4 100644 --- a/codex-rs/app-server/src/transport/remote_control/mod.rs +++ b/codex-rs/app-server/src/transport/remote_control/mod.rs @@ -3,7 +3,6 @@ mod protocol; mod websocket; use self::enroll::load_remote_control_auth; -use self::enroll::remote_control_state_path; use self::protocol::ClientEnvelope; pub use self::protocol::ClientEvent; pub use self::protocol::ClientId; @@ -19,9 +18,9 @@ use crate::outgoing_message::ConnectionId; use crate::outgoing_message::QueuedOutgoingMessage; use codex_app_server_protocol::JSONRPCMessage; use codex_core::AuthManager; +use codex_state::StateRuntime; use std::collections::HashMap; use std::io; -use std::path::PathBuf; use std::sync::Arc; use tokio::sync::mpsc; use tokio::sync::oneshot; @@ -48,13 +47,12 @@ pub(super) struct RemoteControlQueuedServerEnvelope { pub(crate) async fn start_remote_control( remote_control_url: String, - codex_home: PathBuf, + state_db: Option>, auth_manager: Arc, transport_event_tx: mpsc::Sender, shutdown_token: CancellationToken, ) -> io::Result> { 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); @@ -63,7 +61,7 @@ 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, + state_db, auth_manager, client_event_tx, server_event_rx, diff --git a/codex-rs/app-server/src/transport/remote_control/protocol.rs b/codex-rs/app-server/src/transport/remote_control/protocol.rs index 1a9191369e..67a394f174 100644 --- a/codex-rs/app-server/src/transport/remote_control/protocol.rs +++ b/codex-rs/app-server/src/transport/remote_control/protocol.rs @@ -12,20 +12,6 @@ pub(super) struct RemoteControlTarget { pub(super) enroll_url: String, } -#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] -pub(super) struct RemoteControlStateToml { - #[serde(default)] - pub(super) enrollments: Vec, -} - -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -pub(super) struct PersistedRemoteControlEnrollment { - pub(super) websocket_url: String, - pub(super) account_id: Option, - pub(super) server_id: String, - pub(super) server_name: String, -} - #[derive(Debug, Serialize)] pub(super) struct EnrollRemoteServerRequest { pub(super) name: String, diff --git a/codex-rs/app-server/src/transport/remote_control/tests.rs b/codex-rs/app-server/src/transport/remote_control/tests.rs index 1c0a86f3f0..5c75a2a609 100644 --- a/codex-rs/app-server/src/transport/remote_control/tests.rs +++ b/codex-rs/app-server/src/transport/remote_control/tests.rs @@ -28,6 +28,7 @@ use codex_core::AuthManager; 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_state::StateRuntime; use codex_utils_absolute_path::AbsolutePathBuf; use futures::SinkExt; use futures::StreamExt; @@ -73,6 +74,12 @@ fn remote_control_auth_manager_with_home(codex_home: &TempDir) -> Arc Arc { + StateRuntime::init(codex_home.path().to_path_buf(), "test-provider".to_string()) + .await + .expect("state runtime should initialize") +} + #[test] fn app_server_transport_parses_stdio_listen_url() { let transport = AppServerTransport::from_listen_url(AppServerTransport::DEFAULT_LISTEN_URL) @@ -795,6 +802,7 @@ async fn connect_remote_control_websocket_includes_http_error_details() { .await; }); let codex_home = TempDir::new().expect("temp dir should create"); + let state_db = remote_control_state_runtime(&codex_home).await; let auth_manager = remote_control_auth_manager(); let mut enrollment = Some(RemoteControlEnrollment { server_id: "srv_e_test".to_string(), @@ -803,7 +811,7 @@ async fn connect_remote_control_websocket_includes_http_error_details() { let err = match connect_remote_control_websocket( &remote_control_target, - remote_control_state_path(codex_home.path()).as_path(), + Some(state_db.as_ref()), auth_manager.as_ref(), &mut enrollment, None, @@ -821,7 +829,7 @@ async fn connect_remote_control_websocket_includes_http_error_details() { #[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 state_db = remote_control_state_runtime(&codex_home).await; 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") @@ -836,7 +844,7 @@ async fn persisted_remote_control_enrollment_round_trips_by_target_and_account() }; update_persisted_remote_control_enrollment( - state_path.as_path(), + Some(state_db.as_ref()), &first_target, Some("account-a"), Some(&first_enrollment), @@ -844,7 +852,7 @@ async fn persisted_remote_control_enrollment_round_trips_by_target_and_account() .await .expect("first enrollment should persist"); update_persisted_remote_control_enrollment( - state_path.as_path(), + Some(state_db.as_ref()), &second_target, Some("account-a"), Some(&second_enrollment), @@ -854,7 +862,7 @@ async fn persisted_remote_control_enrollment_round_trips_by_target_and_account() assert_eq!( load_persisted_remote_control_enrollment( - state_path.as_path(), + Some(state_db.as_ref()), &first_target, Some("account-a"), ) @@ -863,7 +871,7 @@ async fn persisted_remote_control_enrollment_round_trips_by_target_and_account() ); assert_eq!( load_persisted_remote_control_enrollment( - state_path.as_path(), + Some(state_db.as_ref()), &first_target, Some("account-b"), ) @@ -872,7 +880,7 @@ async fn persisted_remote_control_enrollment_round_trips_by_target_and_account() ); assert_eq!( load_persisted_remote_control_enrollment( - state_path.as_path(), + Some(state_db.as_ref()), &second_target, Some("account-a"), ) @@ -884,7 +892,7 @@ async fn persisted_remote_control_enrollment_round_trips_by_target_and_account() #[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 state_db = remote_control_state_runtime(&codex_home).await; 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") @@ -899,7 +907,7 @@ async fn clearing_persisted_remote_control_enrollment_removes_only_matching_entr }; update_persisted_remote_control_enrollment( - state_path.as_path(), + Some(state_db.as_ref()), &first_target, Some("account-a"), Some(&first_enrollment), @@ -907,7 +915,7 @@ async fn clearing_persisted_remote_control_enrollment_removes_only_matching_entr .await .expect("first enrollment should persist"); update_persisted_remote_control_enrollment( - state_path.as_path(), + Some(state_db.as_ref()), &second_target, Some("account-a"), Some(&second_enrollment), @@ -916,7 +924,7 @@ async fn clearing_persisted_remote_control_enrollment_removes_only_matching_entr .expect("second enrollment should persist"); update_persisted_remote_control_enrollment( - state_path.as_path(), + Some(state_db.as_ref()), &first_target, Some("account-a"), None, @@ -926,7 +934,7 @@ async fn clearing_persisted_remote_control_enrollment_removes_only_matching_entr assert_eq!( load_persisted_remote_control_enrollment( - state_path.as_path(), + Some(state_db.as_ref()), &first_target, Some("account-a"), ) @@ -935,7 +943,7 @@ async fn clearing_persisted_remote_control_enrollment_removes_only_matching_entr ); assert_eq!( load_persisted_remote_control_enrollment( - state_path.as_path(), + Some(state_db.as_ref()), &second_target, Some("account-a"), ) @@ -981,7 +989,7 @@ async fn remote_control_transport_manages_virtual_clients_and_routes_messages() let shutdown_token = CancellationToken::new(); let remote_handle = start_remote_control( remote_control_url, - codex_home.path().to_path_buf(), + Some(remote_control_state_runtime(&codex_home).await), remote_control_auth_manager(), transport_event_tx, shutdown_token.clone(), @@ -1235,7 +1243,7 @@ async fn remote_control_transport_reconnects_after_disconnect() { let shutdown_token = CancellationToken::new(); let remote_handle = start_remote_control( remote_control_url, - codex_home.path().to_path_buf(), + Some(remote_control_state_runtime(&codex_home).await), remote_control_auth_manager(), transport_event_tx, shutdown_token.clone(), @@ -1311,7 +1319,7 @@ async fn remote_control_http_mode_enrolls_before_connecting() { let shutdown_token = CancellationToken::new(); let remote_handle = start_remote_control( remote_control_url, - codex_home.path().to_path_buf(), + Some(remote_control_state_runtime(&codex_home).await), remote_control_auth_manager(), transport_event_tx, shutdown_token.clone(), @@ -1545,6 +1553,7 @@ async fn remote_control_http_mode_reuses_persisted_enrollment_before_reenrolling .expect("listener should have a local addr") ); 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 persisted_enrollment = RemoteControlEnrollment { @@ -1552,7 +1561,7 @@ async fn remote_control_http_mode_reuses_persisted_enrollment_before_reenrolling server_name: "persisted-server".to_string(), }; update_persisted_remote_control_enrollment( - remote_control_state_path(codex_home.path()).as_path(), + Some(state_db.as_ref()), &remote_control_target, Some("account_id"), Some(&persisted_enrollment), @@ -1565,7 +1574,7 @@ async fn remote_control_http_mode_reuses_persisted_enrollment_before_reenrolling let shutdown_token = CancellationToken::new(); let remote_handle = start_remote_control( remote_control_url, - codex_home.path().to_path_buf(), + Some(state_db.clone()), remote_control_auth_manager_with_home(&codex_home), transport_event_tx, shutdown_token.clone(), @@ -1584,7 +1593,7 @@ async fn remote_control_http_mode_reuses_persisted_enrollment_before_reenrolling ); assert_eq!( load_persisted_remote_control_enrollment( - remote_control_state_path(codex_home.path()).as_path(), + Some(state_db.as_ref()), &remote_control_target, Some("account_id"), ) @@ -1608,7 +1617,7 @@ async fn remote_control_http_mode_clears_stale_persisted_enrollment_after_404() .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 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 expected_server_name = gethostname().to_string_lossy().trim().to_string(); @@ -1621,7 +1630,7 @@ async fn remote_control_http_mode_clears_stale_persisted_enrollment_after_404() server_name: expected_server_name, }; update_persisted_remote_control_enrollment( - state_path.as_path(), + Some(state_db.as_ref()), &remote_control_target, Some("account_id"), Some(&stale_enrollment), @@ -1634,7 +1643,7 @@ async fn remote_control_http_mode_clears_stale_persisted_enrollment_after_404() let shutdown_token = CancellationToken::new(); let remote_handle = start_remote_control( remote_control_url, - codex_home.path().to_path_buf(), + Some(state_db.clone()), remote_control_auth_manager_with_home(&codex_home), transport_event_tx, shutdown_token.clone(), @@ -1671,7 +1680,7 @@ async fn remote_control_http_mode_clears_stale_persisted_enrollment_after_404() ); assert_eq!( load_persisted_remote_control_enrollment( - state_path.as_path(), + Some(state_db.as_ref()), &remote_control_target, Some("account_id"), ) diff --git a/codex-rs/app-server/src/transport/remote_control/websocket.rs b/codex-rs/app-server/src/transport/remote_control/websocket.rs index 8872593ddf..256c229083 100644 --- a/codex-rs/app-server/src/transport/remote_control/websocket.rs +++ b/codex-rs/app-server/src/transport/remote_control/websocket.rs @@ -8,11 +8,11 @@ use super::protocol::RemoteControlTarget; use super::protocol::ServerEnvelope; use super::protocol::ServerEvent; use codex_core::AuthManager; +use codex_state::StateRuntime; use futures::SinkExt; use futures::StreamExt; use std::collections::BTreeMap; use std::collections::HashMap; -use std::path::PathBuf; use std::sync::Arc; use tokio::sync::mpsc; use tokio::sync::oneshot; @@ -39,7 +39,7 @@ struct BufferedServerEvent { #[allow(clippy::print_stderr)] pub(super) async fn run_remote_control_websocket_loop( remote_control_target: RemoteControlTarget, - remote_control_state_path: PathBuf, + state_db: Option>, auth_manager: Arc, client_event_tx: mpsc::Sender, mut server_event_rx: mpsc::Receiver, @@ -73,7 +73,7 @@ pub(super) async fn run_remote_control_websocket_loop( _ = shutdown_token.cancelled() => break, connect_result = connect_remote_control_websocket( &remote_control_target, - remote_control_state_path.as_path(), + state_db.as_deref(), auth_manager.as_ref(), &mut enrollment, subscribe_cursor.as_deref(), diff --git a/codex-rs/app-server/src/transport/websocket.rs b/codex-rs/app-server/src/transport/websocket.rs index 48d945d458..20d8ce10db 100644 --- a/codex-rs/app-server/src/transport/websocket.rs +++ b/codex-rs/app-server/src/transport/websocket.rs @@ -1,8 +1,8 @@ +use super::CHANNEL_CAPACITY; +use super::TransportEvent; use super::auth::WebsocketAuthPolicy; use super::auth::authorize_upgrade; use super::auth::should_warn_about_unauthenticated_non_loopback_listener; -use super::CHANNEL_CAPACITY; -use super::TransportEvent; use super::forward_incoming_message; use super::serialize_outgoing_message; use crate::outgoing_message::ConnectionId; diff --git a/codex-rs/state/migrations/0023_remote_control_enrollments.sql b/codex-rs/state/migrations/0023_remote_control_enrollments.sql new file mode 100644 index 0000000000..9a2081dd8f --- /dev/null +++ b/codex-rs/state/migrations/0023_remote_control_enrollments.sql @@ -0,0 +1,8 @@ +CREATE TABLE remote_control_enrollments ( + websocket_url TEXT NOT NULL, + account_id TEXT NOT NULL, + server_id TEXT NOT NULL, + server_name TEXT NOT NULL, + updated_at INTEGER NOT NULL, + PRIMARY KEY (websocket_url, account_id) +); diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index 645aa42695..04eed8e0e2 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -53,6 +53,7 @@ mod agent_jobs; mod backfill; mod logs; mod memories; +mod remote_control; #[cfg(test)] mod test_support; mod threads; diff --git a/codex-rs/state/src/runtime/remote_control.rs b/codex-rs/state/src/runtime/remote_control.rs new file mode 100644 index 0000000000..307dac8184 --- /dev/null +++ b/codex-rs/state/src/runtime/remote_control.rs @@ -0,0 +1,197 @@ +use super::*; + +const REMOTE_CONTROL_ACCOUNT_ID_NONE: &str = ""; + +fn remote_control_account_id_key(account_id: Option<&str>) -> &str { + account_id.unwrap_or(REMOTE_CONTROL_ACCOUNT_ID_NONE) +} + +impl StateRuntime { + pub async fn get_remote_control_enrollment( + &self, + websocket_url: &str, + account_id: Option<&str>, + ) -> anyhow::Result> { + let row = sqlx::query( + r#" +SELECT server_id, server_name +FROM remote_control_enrollments +WHERE websocket_url = ? AND account_id = ? + "#, + ) + .bind(websocket_url) + .bind(remote_control_account_id_key(account_id)) + .fetch_optional(self.pool.as_ref()) + .await?; + + row.map(|row| Ok((row.try_get("server_id")?, row.try_get("server_name")?))) + .transpose() + } + + pub async fn upsert_remote_control_enrollment( + &self, + websocket_url: &str, + account_id: Option<&str>, + server_id: &str, + server_name: &str, + ) -> anyhow::Result<()> { + sqlx::query( + r#" +INSERT INTO remote_control_enrollments ( + websocket_url, + account_id, + server_id, + server_name, + updated_at +) VALUES (?, ?, ?, ?, ?) +ON CONFLICT(websocket_url, account_id) DO UPDATE SET + server_id = excluded.server_id, + server_name = excluded.server_name, + updated_at = excluded.updated_at + "#, + ) + .bind(websocket_url) + .bind(remote_control_account_id_key(account_id)) + .bind(server_id) + .bind(server_name) + .bind(Utc::now().timestamp()) + .execute(self.pool.as_ref()) + .await?; + Ok(()) + } + + pub async fn delete_remote_control_enrollment( + &self, + websocket_url: &str, + account_id: Option<&str>, + ) -> anyhow::Result { + let result = sqlx::query( + r#" +DELETE FROM remote_control_enrollments +WHERE websocket_url = ? AND account_id = ? + "#, + ) + .bind(websocket_url) + .bind(remote_control_account_id_key(account_id)) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected()) + } +} + +#[cfg(test)] +mod tests { + use super::StateRuntime; + use super::test_support::unique_temp_dir; + use pretty_assertions::assert_eq; + + #[tokio::test] + async fn remote_control_enrollment_round_trips_by_target_and_account() { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) + .await + .expect("initialize runtime"); + + runtime + .upsert_remote_control_enrollment( + "wss://example.com/backend-api/wham/remote/control/server", + Some("account-a"), + "srv_e_first", + "first-server", + ) + .await + .expect("insert first enrollment"); + runtime + .upsert_remote_control_enrollment( + "wss://example.com/backend-api/wham/remote/control/server", + Some("account-b"), + "srv_e_second", + "second-server", + ) + .await + .expect("insert second enrollment"); + + assert_eq!( + runtime + .get_remote_control_enrollment( + "wss://example.com/backend-api/wham/remote/control/server", + Some("account-a"), + ) + .await + .expect("load first enrollment"), + Some(("srv_e_first".to_string(), "first-server".to_string())) + ); + assert_eq!( + runtime + .get_remote_control_enrollment( + "wss://example.com/backend-api/wham/remote/control/server", + None, + ) + .await + .expect("load missing enrollment"), + None + ); + + let _ = tokio::fs::remove_dir_all(codex_home).await; + } + + #[tokio::test] + async fn delete_remote_control_enrollment_removes_only_matching_entry() { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) + .await + .expect("initialize runtime"); + + runtime + .upsert_remote_control_enrollment( + "wss://example.com/backend-api/wham/remote/control/server", + None, + "srv_e_first", + "first-server", + ) + .await + .expect("insert first enrollment"); + runtime + .upsert_remote_control_enrollment( + "wss://example.com/backend-api/wham/remote/control/server", + Some("account-a"), + "srv_e_second", + "second-server", + ) + .await + .expect("insert second enrollment"); + + assert_eq!( + runtime + .delete_remote_control_enrollment( + "wss://example.com/backend-api/wham/remote/control/server", + None, + ) + .await + .expect("delete first enrollment"), + 1 + ); + assert_eq!( + runtime + .get_remote_control_enrollment( + "wss://example.com/backend-api/wham/remote/control/server", + None, + ) + .await + .expect("load deleted enrollment"), + None + ); + assert_eq!( + runtime + .get_remote_control_enrollment( + "wss://example.com/backend-api/wham/remote/control/server", + Some("account-a"), + ) + .await + .expect("load retained enrollment"), + Some(("srv_e_second".to_string(), "second-server".to_string())) + ); + + let _ = tokio::fs::remove_dir_all(codex_home).await; + } +}