Skip to main content

codex_state/runtime/
remote_control.rs

1use super::*;
2
3const REMOTE_CONTROL_APP_SERVER_CLIENT_NAME_NONE: &str = "";
4
5/// Persisted remote-control server enrollment, including the lookup key.
6#[derive(Debug, Clone, PartialEq, Eq)]
7pub struct RemoteControlEnrollmentRecord {
8    pub websocket_url: String,
9    pub account_id: String,
10    pub app_server_client_name: Option<String>,
11    pub server_id: String,
12    pub environment_id: String,
13    pub server_name: String,
14    pub remote_control_enabled: Option<bool>,
15}
16
17fn remote_control_app_server_client_name_key(app_server_client_name: Option<&str>) -> &str {
18    app_server_client_name.unwrap_or(REMOTE_CONTROL_APP_SERVER_CLIENT_NAME_NONE)
19}
20
21fn app_server_client_name_from_key(app_server_client_name: String) -> Option<String> {
22    if app_server_client_name.is_empty() {
23        None
24    } else {
25        Some(app_server_client_name)
26    }
27}
28
29impl StateRuntime {
30    pub async fn get_remote_control_enrollment(
31        &self,
32        websocket_url: &str,
33        account_id: &str,
34        app_server_client_name: Option<&str>,
35    ) -> anyhow::Result<Option<RemoteControlEnrollmentRecord>> {
36        let row = sqlx::query(
37            r#"
38SELECT websocket_url, account_id, app_server_client_name, server_id, environment_id, server_name,
39    remote_control_enabled
40FROM remote_control_enrollments
41WHERE websocket_url = ? AND account_id = ? AND app_server_client_name = ?
42            "#,
43        )
44        .bind(websocket_url)
45        .bind(account_id)
46        .bind(remote_control_app_server_client_name_key(
47            app_server_client_name,
48        ))
49        .fetch_optional(self.pool.as_ref())
50        .await?;
51
52        row.map(|row| {
53            let app_server_client_name: String = row.try_get("app_server_client_name")?;
54            Ok(RemoteControlEnrollmentRecord {
55                websocket_url: row.try_get("websocket_url")?,
56                account_id: row.try_get("account_id")?,
57                app_server_client_name: app_server_client_name_from_key(app_server_client_name),
58                server_id: row.try_get("server_id")?,
59                environment_id: row.try_get("environment_id")?,
60                server_name: row.try_get("server_name")?,
61                remote_control_enabled: row.try_get("remote_control_enabled")?,
62            })
63        })
64        .transpose()
65    }
66
67    pub async fn upsert_remote_control_enrollment(
68        &self,
69        enrollment: &RemoteControlEnrollmentRecord,
70    ) -> anyhow::Result<()> {
71        sqlx::query(
72            r#"
73INSERT INTO remote_control_enrollments (
74    websocket_url,
75    account_id,
76    app_server_client_name,
77    server_id,
78    environment_id,
79    server_name,
80    remote_control_enabled,
81    updated_at
82) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
83ON CONFLICT(websocket_url, account_id, app_server_client_name) DO UPDATE SET
84    server_id = excluded.server_id,
85    environment_id = excluded.environment_id,
86    server_name = excluded.server_name,
87    updated_at = excluded.updated_at
88            "#,
89        )
90        .bind(&enrollment.websocket_url)
91        .bind(&enrollment.account_id)
92        .bind(remote_control_app_server_client_name_key(
93            enrollment.app_server_client_name.as_deref(),
94        ))
95        .bind(&enrollment.server_id)
96        .bind(&enrollment.environment_id)
97        .bind(&enrollment.server_name)
98        .bind(enrollment.remote_control_enabled)
99        .bind(Utc::now().timestamp())
100        .execute(self.pool.as_ref())
101        .await?;
102        Ok(())
103    }
104
105    pub async fn set_remote_control_enabled(
106        &self,
107        websocket_url: &str,
108        account_id: &str,
109        app_server_client_name: Option<&str>,
110        remote_control_enabled: bool,
111    ) -> anyhow::Result<u64> {
112        let result = sqlx::query(
113            r#"
114UPDATE remote_control_enrollments
115SET remote_control_enabled = ?, updated_at = ?
116WHERE websocket_url = ? AND account_id = ? AND app_server_client_name = ?
117            "#,
118        )
119        .bind(remote_control_enabled)
120        .bind(Utc::now().timestamp())
121        .bind(websocket_url)
122        .bind(account_id)
123        .bind(remote_control_app_server_client_name_key(
124            app_server_client_name,
125        ))
126        .execute(self.pool.as_ref())
127        .await?;
128        Ok(result.rows_affected())
129    }
130
131    pub async fn delete_remote_control_enrollment(
132        &self,
133        websocket_url: &str,
134        account_id: &str,
135        app_server_client_name: Option<&str>,
136    ) -> anyhow::Result<u64> {
137        let result = sqlx::query(
138            r#"
139DELETE FROM remote_control_enrollments
140WHERE websocket_url = ? AND account_id = ? AND app_server_client_name = ?
141            "#,
142        )
143        .bind(websocket_url)
144        .bind(account_id)
145        .bind(remote_control_app_server_client_name_key(
146            app_server_client_name,
147        ))
148        .execute(self.pool.as_ref())
149        .await?;
150        Ok(result.rows_affected())
151    }
152}
153
154#[cfg(test)]
155mod tests {
156    use super::RemoteControlEnrollmentRecord;
157    use super::StateRuntime;
158    use super::test_support::unique_temp_dir;
159    use crate::migrations::STATE_MIGRATOR;
160    use crate::state_db_path;
161    use codex_utils_absolute_path::test_support::PathExt;
162    use pretty_assertions::assert_eq;
163    use sqlx::migrate::Migrator;
164    use std::borrow::Cow;
165
166    #[tokio::test]
167    async fn remote_control_enrollment_round_trips_by_target_and_account() {
168        let codex_home = unique_temp_dir();
169        let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string())
170            .await
171            .expect("initialize runtime");
172
173        runtime
174            .upsert_remote_control_enrollment(&RemoteControlEnrollmentRecord {
175                websocket_url: "wss://example.com/backend-api/wham/remote/control/server"
176                    .to_string(),
177                account_id: "account-a".to_string(),
178                app_server_client_name: Some("desktop-client".to_string()),
179                server_id: "srv_e_first".to_string(),
180                environment_id: "env_first".to_string(),
181                server_name: "first-server".to_string(),
182                remote_control_enabled: Some(false),
183            })
184            .await
185            .expect("insert first enrollment");
186        runtime
187            .upsert_remote_control_enrollment(&RemoteControlEnrollmentRecord {
188                websocket_url: "wss://example.com/backend-api/wham/remote/control/server"
189                    .to_string(),
190                account_id: "account-b".to_string(),
191                app_server_client_name: Some("desktop-client".to_string()),
192                server_id: "srv_e_second".to_string(),
193                environment_id: "env_second".to_string(),
194                server_name: "second-server".to_string(),
195                remote_control_enabled: Some(true),
196            })
197            .await
198            .expect("insert second enrollment");
199
200        assert_eq!(
201            runtime
202                .get_remote_control_enrollment(
203                    "wss://example.com/backend-api/wham/remote/control/server",
204                    "account-a",
205                    Some("desktop-client"),
206                )
207                .await
208                .expect("load first enrollment"),
209            Some(RemoteControlEnrollmentRecord {
210                websocket_url: "wss://example.com/backend-api/wham/remote/control/server"
211                    .to_string(),
212                account_id: "account-a".to_string(),
213                app_server_client_name: Some("desktop-client".to_string()),
214                server_id: "srv_e_first".to_string(),
215                environment_id: "env_first".to_string(),
216                server_name: "first-server".to_string(),
217                remote_control_enabled: Some(false),
218            })
219        );
220        assert_eq!(
221            runtime
222                .get_remote_control_enrollment(
223                    "wss://example.com/backend-api/wham/remote/control/server",
224                    "account-missing",
225                    Some("desktop-client"),
226                )
227                .await
228                .expect("load missing enrollment"),
229            None
230        );
231        assert_eq!(
232            runtime
233                .get_remote_control_enrollment(
234                    "wss://example.com/backend-api/wham/remote/control/server",
235                    "account-a",
236                    Some("other-client"),
237                )
238                .await
239                .expect("load wrong client enrollment"),
240            None
241        );
242
243        let _ = tokio::fs::remove_dir_all(codex_home).await;
244    }
245
246    #[tokio::test]
247    async fn delete_remote_control_enrollment_removes_only_matching_entry() {
248        let codex_home = unique_temp_dir();
249        let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string())
250            .await
251            .expect("initialize runtime");
252
253        runtime
254            .upsert_remote_control_enrollment(&RemoteControlEnrollmentRecord {
255                websocket_url: "wss://example.com/backend-api/wham/remote/control/server"
256                    .to_string(),
257                account_id: "account-a".to_string(),
258                app_server_client_name: None,
259                server_id: "srv_e_first".to_string(),
260                environment_id: "env_first".to_string(),
261                server_name: "first-server".to_string(),
262                remote_control_enabled: Some(false),
263            })
264            .await
265            .expect("insert first enrollment");
266        runtime
267            .upsert_remote_control_enrollment(&RemoteControlEnrollmentRecord {
268                websocket_url: "wss://example.com/backend-api/wham/remote/control/server"
269                    .to_string(),
270                account_id: "account-b".to_string(),
271                app_server_client_name: None,
272                server_id: "srv_e_second".to_string(),
273                environment_id: "env_second".to_string(),
274                server_name: "second-server".to_string(),
275                remote_control_enabled: Some(true),
276            })
277            .await
278            .expect("insert second enrollment");
279
280        assert_eq!(
281            runtime
282                .delete_remote_control_enrollment(
283                    "wss://example.com/backend-api/wham/remote/control/server",
284                    "account-a",
285                    /*app_server_client_name*/ None,
286                )
287                .await
288                .expect("delete first enrollment"),
289            1
290        );
291        assert_eq!(
292            runtime
293                .get_remote_control_enrollment(
294                    "wss://example.com/backend-api/wham/remote/control/server",
295                    "account-a",
296                    /*app_server_client_name*/ None,
297                )
298                .await
299                .expect("load deleted enrollment"),
300            None
301        );
302        assert_eq!(
303            runtime
304                .get_remote_control_enrollment(
305                    "wss://example.com/backend-api/wham/remote/control/server",
306                    "account-b",
307                    /*app_server_client_name*/ None,
308                )
309                .await
310                .expect("load retained enrollment"),
311            Some(RemoteControlEnrollmentRecord {
312                websocket_url: "wss://example.com/backend-api/wham/remote/control/server"
313                    .to_string(),
314                account_id: "account-b".to_string(),
315                app_server_client_name: None,
316                server_id: "srv_e_second".to_string(),
317                environment_id: "env_second".to_string(),
318                server_name: "second-server".to_string(),
319                remote_control_enabled: Some(true),
320            })
321        );
322
323        let _ = tokio::fs::remove_dir_all(codex_home).await;
324    }
325
326    #[tokio::test]
327    async fn migration_preserves_legacy_remote_control_preference_as_null() {
328        let codex_home = unique_temp_dir();
329        tokio::fs::create_dir_all(&codex_home)
330            .await
331            .expect("create codex home");
332        let old_state_migrator = Migrator {
333            migrations: Cow::Owned(
334                STATE_MIGRATOR
335                    .migrations
336                    .iter()
337                    .filter(|migration| migration.version <= 36)
338                    .cloned()
339                    .collect(),
340            ),
341            ignore_missing: false,
342            locking: true,
343            no_tx: false,
344            table_name: STATE_MIGRATOR.table_name.clone(),
345            create_schemas: STATE_MIGRATOR.create_schemas.clone(),
346        };
347        let pool = crate::SqliteConfig::new_for_testing(codex_home.as_path().abs())
348            .open_read_write_pool(&state_db_path(codex_home.as_path()))
349            .await
350            .expect("open old state db");
351        old_state_migrator
352            .run(&pool)
353            .await
354            .expect("apply old state schema");
355        sqlx::query("INSERT INTO remote_control_enrollments (websocket_url, account_id, app_server_client_name, server_id, environment_id, server_name, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)")
356        .bind("wss://example.com/backend-api/wham/remote/control/server")
357        .bind("account-a")
358        .bind("desktop-client")
359        .bind("srv_e_first")
360        .bind("env_first")
361        .bind("first-server")
362        .bind(1_i64)
363        .execute(&pool)
364        .await
365        .expect("insert legacy enrollment");
366        pool.close().await;
367
368        let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string())
369            .await
370            .expect("initialize runtime");
371        let actual = runtime
372            .get_remote_control_enrollment(
373                "wss://example.com/backend-api/wham/remote/control/server",
374                "account-a",
375                Some("desktop-client"),
376            )
377            .await
378            .expect("load migrated enrollment")
379            .expect("legacy enrollment should remain");
380        assert_eq!(actual.remote_control_enabled, None);
381
382        let _ = tokio::fs::remove_dir_all(codex_home).await;
383    }
384}