openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-13 11:47:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

sqlite

Browse Source
This commit is contained in:
Ruslan Nigmatullin
2026-03-25 15:02:30 -07:00
parent 8c644a154b
commit 8d62dd3257
10 changed files with 285 additions and 152 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
codex-rs/app-server/src/lib.rs
View File

@@ -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(),

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

@@ -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<MaybeTlsStream<TcpStream>>,
}
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<RemoteControlStateToml> {
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<RemoteControlEnrollment> {
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<RemoteControlEnrollment>,
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;

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

@@ -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<Arc<StateRuntime>>,
auth_manager: Arc<AuthManager>,
transport_event_tx: mpsc::Sender<TransportEvent>,
shutdown_token: CancellationToken,
) -> io::Result<JoinHandle<()>> {
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,

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

@@ -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<PersistedRemoteControlEnrollment>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub(super) struct PersistedRemoteControlEnrollment {
pub(super) websocket_url: String,
pub(super) account_id: Option<String>,
pub(super) server_id: String,
pub(super) server_name: String,
}
#[derive(Debug, Serialize)]
pub(super) struct EnrollRemoteServerRequest {
pub(super) name: String,

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

@@ -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<AuthManage
)
}
async fn remote_control_state_runtime(codex_home: &TempDir) -> Arc<StateRuntime> {
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"),
)

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

@@ -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<Arc<StateRuntime>>,
auth_manager: Arc<AuthManager>,
client_event_tx: mpsc::Sender<ClientEnvelope>,
mut server_event_rx: mpsc::Receiver<RemoteControlQueuedServerEnvelope>,
@@ -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(),

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

@@ -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;

8
codex-rs/state/migrations/0023_remote_control_enrollments.sql Normal file
View File

@@ -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)
);

1
codex-rs/state/src/runtime.rs
View File

@@ -53,6 +53,7 @@ mod agent_jobs;
mod backfill;
mod logs;
mod memories;
mod remote_control;
#[cfg(test)]
mod test_support;
mod threads;

197
codex-rs/state/src/runtime/remote_control.rs Normal file
View File

@@ -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<Option<(String, String)>> {
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<u64> {
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;
}
}