From d61dfeb23a13be748c2d58409ca5e16d281267b4 Mon Sep 17 00:00:00 2001 From: Anton Panasenko Date: Thu, 11 Jun 2026 21:28:52 -0700 Subject: [PATCH] feat(app-server): persist remote-control desired state (#27445) ## Why Remote-control runtime enablement and persisted enrollment preference were represented by separate flags. That made startup rehydration, RPC persistence, and new-enrollment seeding race with one another, and it did not cleanly distinguish runtime-only CLI or daemon starts from durable app-server RPC changes. ## What Changed - Replace the parallel enablement, seed, and rehydration flags with one transport-owned `RemoteControlDesiredState`. - Add nullable enrollment-scoped persistence and preserve existing preferences during enrollment upserts. - Rehydrate plain startup only after auth and client scope resolve, without overwriting a concurrent RPC transition. - Make ordinary `remoteControl/enable` and `remoteControl/disable` durable while retaining `ephemeral: true` for runtime-only callers. - Have the daemon explicitly request ephemeral enablement and regenerate the app-server schemas. ## Verification - Covered migration and `NULL`/`0`/`1` persistence round trips. - Covered plain-start rehydration and runtime-only versus durable enrollment seeding. - Covered durable enable, durable disable, and ephemeral enable through app-server RPC. - Covered the daemon's exact `{ "ephemeral": true }` request payload. Related issue: N/A (internal remote-control persistence architecture change). --- codex-rs/app-server-daemon/src/backend/pid.rs | 18 + .../src/backend/pid_tests.rs | 20 + codex-rs/app-server-daemon/src/lib.rs | 10 + .../src/remote_control_client.rs | 351 ++++++++++-- .../schema/json/ClientRequest.json | 16 + .../codex_app_server_protocol.schemas.json | 16 + .../codex_app_server_protocol.v2.schemas.json | 16 + .../v2/RemoteControlDisableParams.ts | 5 + .../v2/RemoteControlEnableParams.ts | 5 + .../schema/typescript/v2/index.ts | 2 + .../src/protocol/common.rs | 4 +- .../src/protocol/v2/remote_control.rs | 20 + codex-rs/app-server-transport/src/lib.rs | 3 + .../app-server-transport/src/transport/mod.rs | 3 + .../transport/remote_control/desired_state.rs | 158 ++++++ .../src/transport/remote_control/enroll.rs | 7 + .../src/transport/remote_control/mod.rs | 435 ++++++++------ .../src/transport/remote_control/tests.rs | 531 ++++++++++++++++-- .../remote_control/tests/clients_tests.rs | 7 +- .../remote_control/tests/pairing_tests.rs | 31 +- .../src/transport/remote_control/websocket.rs | 369 +++++++++--- codex-rs/app-server/README.md | 4 +- codex-rs/app-server/src/lib.rs | 46 +- codex-rs/app-server/src/main.rs | 12 +- codex-rs/app-server/src/message_processor.rs | 16 +- .../remote_control_processor.rs | 43 +- codex-rs/app-server/src/transport.rs | 2 + .../tests/common/test_app_server.rs | 23 + .../tests/suite/v2/remote_control.rs | 259 ++++++++- codex-rs/cli/src/main.rs | 31 +- codex-rs/cli/src/remote_control_cmd.rs | 2 +- ...037_remote_control_enrollments_enabled.sql | 2 + codex-rs/state/src/runtime/remote_control.rs | 102 +++- 33 files changed, 2157 insertions(+), 412 deletions(-) create mode 100644 codex-rs/app-server-protocol/schema/typescript/v2/RemoteControlDisableParams.ts create mode 100644 codex-rs/app-server-protocol/schema/typescript/v2/RemoteControlEnableParams.ts create mode 100644 codex-rs/app-server-transport/src/transport/remote_control/desired_state.rs create mode 100644 codex-rs/state/migrations/0037_remote_control_enrollments_enabled.sql diff --git a/codex-rs/app-server-daemon/src/backend/pid.rs b/codex-rs/app-server-daemon/src/backend/pid.rs index f5f4fc5b6..b525e6125 100644 --- a/codex-rs/app-server-daemon/src/backend/pid.rs +++ b/codex-rs/app-server-daemon/src/backend/pid.rs @@ -8,6 +8,8 @@ use std::time::Duration; use anyhow::Context; use anyhow::Result; use anyhow::bail; +#[cfg(unix)] +use codex_app_server_transport::REMOTE_CONTROL_DISABLED_ENV_VAR; use serde::Deserialize; use serde::Serialize; use tokio::fs; @@ -164,6 +166,9 @@ impl PidBackend { .stdin(Stdio::null()) .stdout(Stdio::null()) .stderr(Stdio::from(stderr_log.into_std().await)); + if let Some((key, value)) = self.command_env() { + command.env(key, value); + } #[cfg(unix)] { @@ -407,6 +412,19 @@ impl PidBackend { } } + #[cfg(unix)] + fn command_env(&self) -> Option<(&'static str, &'static str)> { + match self.command_kind { + PidCommandKind::AppServer { + remote_control_enabled: false, + } => Some((REMOTE_CONTROL_DISABLED_ENV_VAR, "1")), + PidCommandKind::AppServer { + remote_control_enabled: true, + } + | PidCommandKind::UpdateLoop => None, + } + } + fn terminate_process(&self, pid: u32) -> Result<()> { match self.command_kind { PidCommandKind::AppServer { .. } => terminate_process(pid), diff --git a/codex-rs/app-server-daemon/src/backend/pid_tests.rs b/codex-rs/app-server-daemon/src/backend/pid_tests.rs index 4c3a0e441..57b3d227e 100644 --- a/codex-rs/app-server-daemon/src/backend/pid_tests.rs +++ b/codex-rs/app-server-daemon/src/backend/pid_tests.rs @@ -3,6 +3,8 @@ use std::time::Duration; use pretty_assertions::assert_eq; use tempfile::TempDir; +use codex_app_server_transport::REMOTE_CONTROL_DISABLED_ENV_VAR; + use super::PidBackend; use super::PidCommandKind; use super::PidFileState; @@ -174,6 +176,24 @@ fn app_server_remote_control_uses_runtime_flag() { ); } +#[test] +fn app_server_disabled_remote_control_uses_compatible_args_and_runtime_env() { + let backend = PidBackend::new( + "codex".into(), + "app-server.pid".into(), + /*remote_control_enabled*/ false, + ); + + assert_eq!( + backend.command_args(), + vec!["app-server", "--listen", "unix://"] + ); + assert_eq!( + backend.command_env(), + Some((REMOTE_CONTROL_DISABLED_ENV_VAR, "1")) + ); +} + #[tokio::test] async fn read_stderr_log_tail_returns_recent_complete_lines() { let temp_dir = TempDir::new().expect("temp dir"); diff --git a/codex-rs/app-server-daemon/src/lib.rs b/codex-rs/app-server-daemon/src/lib.rs index 354c08ec1..21fecef1c 100644 --- a/codex-rs/app-server-daemon/src/lib.rs +++ b/codex-rs/app-server-daemon/src/lib.rs @@ -543,6 +543,16 @@ impl Daemon { } else { None }; + if info.is_some() { + match mode { + RemoteControlMode::Enabled => { + remote_control_client::enable_remote_control(&self.socket_path).await?; + } + RemoteControlMode::Disabled => { + remote_control_client::disable_remote_control(&self.socket_path).await?; + } + } + } return Ok(self.remote_control_output( already_remote_control_status(mode), backend.map(|_| BackendKind::Pid), diff --git a/codex-rs/app-server-daemon/src/remote_control_client.rs b/codex-rs/app-server-daemon/src/remote_control_client.rs index cc32e50dd..447e22ac6 100644 --- a/codex-rs/app-server-daemon/src/remote_control_client.rs +++ b/codex-rs/app-server-daemon/src/remote_control_client.rs @@ -8,9 +8,13 @@ use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCNotification; use codex_app_server_protocol::JSONRPCRequest; use codex_app_server_protocol::RemoteControlConnectionStatus; +use codex_app_server_protocol::RemoteControlDisableParams; +use codex_app_server_protocol::RemoteControlDisableResponse; +use codex_app_server_protocol::RemoteControlEnableParams; use codex_app_server_protocol::RemoteControlEnableResponse; use codex_app_server_protocol::RemoteControlStatusChangedNotification; use codex_app_server_protocol::RequestId; +use serde::de::DeserializeOwned; use tokio::io::AsyncRead; use tokio::io::AsyncWrite; use tokio::time::Instant; @@ -22,13 +26,33 @@ use crate::RemoteControlReadyStatus; use crate::client; const REMOTE_CONTROL_READY_TIMEOUT: Duration = Duration::from_secs(10); -const REMOTE_CONTROL_ENABLE_REQUEST_ID: RequestId = RequestId::Integer(2); +const REMOTE_CONTROL_REQUEST_ID: RequestId = RequestId::Integer(2); +const INVALID_PARAMS_ERROR_CODE: i64 = -32602; + +enum RemoteControlRpcResponse { + Success(T), + InvalidParams, +} pub(crate) async fn enable_remote_control(socket_path: &Path) -> Result { let mut websocket = client::connect(socket_path).await?; enable_remote_control_with_timeout(&mut websocket, REMOTE_CONTROL_READY_TIMEOUT).await } +pub(crate) async fn disable_remote_control(socket_path: &Path) -> Result { + let mut websocket = client::connect(socket_path).await?; + initialize_client(&mut websocket).await?; + let params = serde_json::to_value(RemoteControlDisableParams { ephemeral: true })?; + let response: RemoteControlDisableResponse = request_remote_control_with_legacy_fallback( + &mut websocket, + "remoteControl/disable", + params, + ) + .await?; + websocket.close(None).await.ok(); + Ok(RemoteControlReadyStatus::from(response)) +} + pub(crate) async fn enable_remote_control_with_connect_retry( socket_path: &Path, connect_timeout: Duration, @@ -43,6 +67,26 @@ async fn enable_remote_control_with_timeout( websocket: &mut WebSocketStream, ready_timeout: Duration, ) -> Result +where + S: AsyncRead + AsyncWrite + Unpin, +{ + initialize_client(websocket).await?; + + let response: RemoteControlEnableResponse = request_remote_control_with_legacy_fallback( + websocket, + "remoteControl/enable", + serde_json::to_value(RemoteControlEnableParams { ephemeral: true })?, + ) + .await?; + let mut latest = RemoteControlReadyStatus::from(response); + if latest.status == RemoteControlConnectionStatus::Connecting { + latest = wait_for_remote_control_status(websocket, latest, ready_timeout).await?; + } + websocket.close(None).await.ok(); + Ok(latest) +} + +async fn initialize_client(websocket: &mut WebSocketStream) -> Result<()> where S: AsyncRead + AsyncWrite + Unpin, { @@ -53,24 +97,65 @@ where }); client::send_message(websocket, &initialized) .await - .context("failed to send initialized notification")?; + .context("failed to send initialized notification") +} - let enable = JSONRPCMessage::Request(JSONRPCRequest { - id: REMOTE_CONTROL_ENABLE_REQUEST_ID, - method: "remoteControl/enable".to_string(), - params: None, +async fn send_remote_control_request( + websocket: &mut WebSocketStream, + request_id: RequestId, + method: &str, + params: Option, +) -> Result<()> +where + S: AsyncRead + AsyncWrite + Unpin, +{ + let request = JSONRPCMessage::Request(JSONRPCRequest { + id: request_id, + method: method.to_string(), + params, trace: None, }); - client::send_message(websocket, &enable) + client::send_message(websocket, &request) .await - .context("failed to send remoteControl/enable request")?; + .with_context(|| format!("failed to send {method} request")) +} - let mut latest = read_enable_response(websocket).await?; - if latest.status == RemoteControlConnectionStatus::Connecting { - latest = wait_for_remote_control_status(websocket, latest, ready_timeout).await?; +async fn request_remote_control_with_legacy_fallback( + websocket: &mut WebSocketStream, + method: &str, + params: serde_json::Value, +) -> Result +where + S: AsyncRead + AsyncWrite + Unpin, + T: DeserializeOwned, +{ + send_remote_control_request( + websocket, + REMOTE_CONTROL_REQUEST_ID.clone(), + method, + Some(params), + ) + .await?; + match read_remote_control_response(websocket, &REMOTE_CONTROL_REQUEST_ID, method).await? { + RemoteControlRpcResponse::Success(response) => Ok(response), + RemoteControlRpcResponse::InvalidParams => { + send_remote_control_request( + websocket, + REMOTE_CONTROL_REQUEST_ID.clone(), + method, + /*params*/ None, + ) + .await?; + match read_remote_control_response(websocket, &REMOTE_CONTROL_REQUEST_ID, method) + .await? + { + RemoteControlRpcResponse::Success(response) => Ok(response), + RemoteControlRpcResponse::InvalidParams => { + Err(anyhow!("{method} rejected legacy params")) + } + } + } } - websocket.close(None).await.ok(); - Ok(latest) } async fn connect_with_retry( @@ -97,11 +182,14 @@ async fn connect_with_retry( } } -async fn read_enable_response( +async fn read_remote_control_response( websocket: &mut WebSocketStream, -) -> Result + request_id: &RequestId, + method: &str, +) -> Result> where S: AsyncRead + AsyncWrite + Unpin, + T: DeserializeOwned, { loop { let message = timeout( @@ -109,21 +197,20 @@ where client::read_message(websocket), ) .await - .context("timed out waiting for remoteControl/enable response")??; + .with_context(|| format!("timed out waiting for {method} response"))??; match message { - JSONRPCMessage::Response(response) - if response.id == REMOTE_CONTROL_ENABLE_REQUEST_ID => - { - let response = - serde_json::from_value::(response.result) - .context("failed to parse remoteControl/enable response")?; - return Ok(RemoteControlReadyStatus::from(response)); + JSONRPCMessage::Response(response) if response.id == *request_id => { + let response = serde_json::from_value::(response.result) + .with_context(|| format!("failed to parse {method} response"))?; + return Ok(RemoteControlRpcResponse::Success(response)); } - JSONRPCMessage::Error(err) if err.id == REMOTE_CONTROL_ENABLE_REQUEST_ID => { - return Err(anyhow!( - "remoteControl/enable failed: {}", - err.error.message - )); + JSONRPCMessage::Error(err) + if err.id == *request_id && err.error.code == INVALID_PARAMS_ERROR_CODE => + { + return Ok(RemoteControlRpcResponse::InvalidParams); + } + JSONRPCMessage::Error(err) if err.id == *request_id => { + return Err(anyhow!("{method} failed: {}", err.error.message)); } JSONRPCMessage::Notification(notification) if remote_control_status_notification(¬ification).is_some() => @@ -196,6 +283,23 @@ impl From for RemoteControlReadyStatus { } } +impl From for RemoteControlReadyStatus { + fn from(response: RemoteControlDisableResponse) -> Self { + let RemoteControlDisableResponse { + status, + server_name, + installation_id: _, + environment_id, + } = response; + Self { + status, + server_name, + environment_id, + timed_out: false, + } + } +} + impl From for RemoteControlReadyStatus { fn from(notification: RemoteControlStatusChangedNotification) -> Self { let RemoteControlStatusChangedNotification { @@ -216,6 +320,8 @@ impl From for RemoteControlReadyStatus { #[cfg(all(test, unix))] mod tests { use anyhow::Result; + use codex_app_server_protocol::JSONRPCError; + use codex_app_server_protocol::JSONRPCErrorError; use codex_app_server_protocol::JSONRPCResponse; use codex_uds::UnixListener; use pretty_assertions::assert_eq; @@ -243,6 +349,7 @@ mod tests { ), after_enable_notification: None, ready_timeout: Duration::from_millis(20), + reject_ephemeral_params: false, }) .await?; @@ -271,6 +378,7 @@ mod tests { Some("env_test"), )), ready_timeout: Duration::from_secs(1), + reject_ephemeral_params: false, }) .await?; @@ -296,6 +404,7 @@ mod tests { ), after_enable_notification: None, ready_timeout: Duration::from_millis(20), + reject_ephemeral_params: false, }) .await?; @@ -321,6 +430,7 @@ mod tests { ), after_enable_notification: None, ready_timeout: Duration::from_millis(20), + reject_ephemeral_params: false, }) .await?; @@ -336,11 +446,104 @@ mod tests { Ok(()) } + #[tokio::test] + async fn enable_remote_control_retries_without_params_for_older_servers() -> Result<()> { + let status = run_enable_remote_control_scenario(EnableScenario { + initial_notification: None, + enable_response: remote_control_status( + RemoteControlConnectionStatus::Connected, + Some("env_test"), + ), + after_enable_notification: None, + ready_timeout: Duration::from_millis(20), + reject_ephemeral_params: true, + }) + .await?; + + assert_eq!( + status, + RemoteControlReadyStatus { + status: RemoteControlConnectionStatus::Connected, + server_name: TEST_SERVER_NAME.to_string(), + environment_id: Some("env_test".to_string()), + timed_out: false, + } + ); + Ok(()) + } + + #[tokio::test] + async fn disable_remote_control_retries_without_params_for_older_servers() -> Result<()> { + let dir = TempDir::new()?; + let socket_path = dir.path().join("app-server.sock"); + let listener = UnixListener::bind(&socket_path).await?; + let server_task = tokio::spawn(async move { + let mut websocket = accept_initialized_client(listener).await?; + let disable = client::read_message(&mut websocket).await?; + let JSONRPCMessage::Request(disable) = disable else { + panic!("expected remoteControl/disable request"); + }; + assert_eq!(disable.id, REMOTE_CONTROL_REQUEST_ID); + assert_eq!(disable.method, "remoteControl/disable"); + assert_eq!( + disable.params, + Some(serde_json::json!({ "ephemeral": true })) + ); + client::send_message( + &mut websocket, + &JSONRPCMessage::Error(JSONRPCError { + id: REMOTE_CONTROL_REQUEST_ID, + error: JSONRPCErrorError { + code: INVALID_PARAMS_ERROR_CODE, + message: "Invalid params".to_string(), + data: None, + }, + }), + ) + .await?; + let fallback = client::read_message(&mut websocket).await?; + let JSONRPCMessage::Request(fallback) = fallback else { + panic!("expected fallback remoteControl/disable request"); + }; + assert_eq!(fallback.id, REMOTE_CONTROL_REQUEST_ID); + assert_eq!(fallback.method, "remoteControl/disable"); + assert_eq!(fallback.params, None); + client::send_message( + &mut websocket, + &JSONRPCMessage::Response(JSONRPCResponse { + id: REMOTE_CONTROL_REQUEST_ID, + result: serde_json::to_value(RemoteControlDisableResponse::from( + remote_control_status( + RemoteControlConnectionStatus::Disabled, + /*environment_id*/ None, + ), + ))?, + }), + ) + .await?; + Ok::<_, anyhow::Error>(()) + }); + + let status = disable_remote_control(&socket_path).await?; + server_task.await??; + assert_eq!( + status, + RemoteControlReadyStatus { + status: RemoteControlConnectionStatus::Disabled, + server_name: TEST_SERVER_NAME.to_string(), + environment_id: None, + timed_out: false, + } + ); + Ok(()) + } + struct EnableScenario { initial_notification: Option, enable_response: RemoteControlStatusChangedNotification, after_enable_notification: Option, ready_timeout: Duration, + reject_ephemeral_params: bool, } async fn run_enable_remote_control_scenario( @@ -359,12 +562,70 @@ mod tests { } async fn serve_enable_remote_control_scenario( - mut listener: UnixListener, + listener: UnixListener, scenario: EnableScenario, ) -> Result<()> { + let mut websocket = accept_initialized_client(listener).await?; + if let Some(status) = scenario.initial_notification { + send_remote_control_status(&mut websocket, status).await?; + } + + let enable = client::read_message(&mut websocket).await?; + let JSONRPCMessage::Request(enable) = enable else { + panic!("expected remoteControl/enable request"); + }; + assert_eq!(enable.id, REMOTE_CONTROL_REQUEST_ID); + assert_eq!(enable.method, "remoteControl/enable"); + assert_eq!( + enable.params, + Some(serde_json::json!({ "ephemeral": true })) + ); + if scenario.reject_ephemeral_params { + client::send_message( + &mut websocket, + &JSONRPCMessage::Error(JSONRPCError { + id: REMOTE_CONTROL_REQUEST_ID, + error: JSONRPCErrorError { + code: INVALID_PARAMS_ERROR_CODE, + message: "Invalid params".to_string(), + data: None, + }, + }), + ) + .await?; + let fallback = client::read_message(&mut websocket).await?; + let JSONRPCMessage::Request(fallback) = fallback else { + panic!("expected fallback remoteControl/enable request"); + }; + assert_eq!(fallback.id, REMOTE_CONTROL_REQUEST_ID); + assert_eq!(fallback.method, "remoteControl/enable"); + assert_eq!(fallback.params, None); + } + client::send_message( + &mut websocket, + &JSONRPCMessage::Response(JSONRPCResponse { + id: REMOTE_CONTROL_REQUEST_ID, + result: serde_json::to_value(RemoteControlEnableResponse::from( + scenario.enable_response, + ))?, + }), + ) + .await?; + + if let Some(status) = scenario.after_enable_notification { + send_remote_control_status(&mut websocket, status).await?; + } else { + tokio::time::sleep(Duration::from_millis(50)).await; + } + + Ok(()) + } + + async fn accept_initialized_client( + mut listener: UnixListener, + ) -> Result> { let stream = listener.accept().await?; let mut websocket = accept_async(stream).await?; - let initialize = client::read_message(&mut websocket).await?; let JSONRPCMessage::Request(initialize) = initialize else { panic!("expected initialize request"); @@ -397,35 +658,7 @@ mod tests { panic!("expected initialized notification"); }; assert_eq!(initialized.method, "initialized"); - - if let Some(status) = scenario.initial_notification { - send_remote_control_status(&mut websocket, status).await?; - } - - let enable = client::read_message(&mut websocket).await?; - let JSONRPCMessage::Request(enable) = enable else { - panic!("expected remoteControl/enable request"); - }; - assert_eq!(enable.id, REMOTE_CONTROL_ENABLE_REQUEST_ID); - assert_eq!(enable.method, "remoteControl/enable"); - client::send_message( - &mut websocket, - &JSONRPCMessage::Response(JSONRPCResponse { - id: REMOTE_CONTROL_ENABLE_REQUEST_ID, - result: serde_json::to_value(RemoteControlEnableResponse::from( - scenario.enable_response, - ))?, - }), - ) - .await?; - - if let Some(status) = scenario.after_enable_notification { - send_remote_control_status(&mut websocket, status).await?; - } else { - tokio::time::sleep(Duration::from_millis(50)).await; - } - - Ok(()) + Ok(websocket) } async fn send_remote_control_status( diff --git a/codex-rs/app-server-protocol/schema/json/ClientRequest.json b/codex-rs/app-server-protocol/schema/json/ClientRequest.json index 6bdf19536..7247794ea 100644 --- a/codex-rs/app-server-protocol/schema/json/ClientRequest.json +++ b/codex-rs/app-server-protocol/schema/json/ClientRequest.json @@ -2124,6 +2124,22 @@ } ] }, + "RemoteControlDisableParams": { + "properties": { + "ephemeral": { + "type": "boolean" + } + }, + "type": "object" + }, + "RemoteControlEnableParams": { + "properties": { + "ephemeral": { + "type": "boolean" + } + }, + "type": "object" + }, "RequestId": { "anyOf": [ { diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json index 62b071769..fadec6acb 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.schemas.json @@ -14050,6 +14050,22 @@ ], "type": "string" }, + "RemoteControlDisableParams": { + "properties": { + "ephemeral": { + "type": "boolean" + } + }, + "type": "object" + }, + "RemoteControlEnableParams": { + "properties": { + "ephemeral": { + "type": "boolean" + } + }, + "type": "object" + }, "RemoteControlStatusChangedNotification": { "$schema": "http://json-schema.org/draft-07/schema#", "description": "Current remote-control connection status and remote identity exposed to clients.", diff --git a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json index d9f0b9179..ebe1efe26 100644 --- a/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json +++ b/codex-rs/app-server-protocol/schema/json/codex_app_server_protocol.v2.schemas.json @@ -10532,6 +10532,22 @@ ], "type": "string" }, + "RemoteControlDisableParams": { + "properties": { + "ephemeral": { + "type": "boolean" + } + }, + "type": "object" + }, + "RemoteControlEnableParams": { + "properties": { + "ephemeral": { + "type": "boolean" + } + }, + "type": "object" + }, "RemoteControlStatusChangedNotification": { "$schema": "http://json-schema.org/draft-07/schema#", "description": "Current remote-control connection status and remote identity exposed to clients.", diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/RemoteControlDisableParams.ts b/codex-rs/app-server-protocol/schema/typescript/v2/RemoteControlDisableParams.ts new file mode 100644 index 000000000..30a59d357 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/RemoteControlDisableParams.ts @@ -0,0 +1,5 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. + +export type RemoteControlDisableParams = { ephemeral?: boolean, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/RemoteControlEnableParams.ts b/codex-rs/app-server-protocol/schema/typescript/v2/RemoteControlEnableParams.ts new file mode 100644 index 000000000..3848982d6 --- /dev/null +++ b/codex-rs/app-server-protocol/schema/typescript/v2/RemoteControlEnableParams.ts @@ -0,0 +1,5 @@ +// GENERATED CODE! DO NOT MODIFY BY HAND! + +// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually. + +export type RemoteControlEnableParams = { ephemeral?: boolean, }; diff --git a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts index b4cd88aad..c46e95f95 100644 --- a/codex-rs/app-server-protocol/schema/typescript/v2/index.ts +++ b/codex-rs/app-server-protocol/schema/typescript/v2/index.ts @@ -325,6 +325,8 @@ export type { ReasoningSummaryPartAddedNotification } from "./ReasoningSummaryPa export type { ReasoningSummaryTextDeltaNotification } from "./ReasoningSummaryTextDeltaNotification"; export type { ReasoningTextDeltaNotification } from "./ReasoningTextDeltaNotification"; export type { RemoteControlConnectionStatus } from "./RemoteControlConnectionStatus"; +export type { RemoteControlDisableParams } from "./RemoteControlDisableParams"; +export type { RemoteControlEnableParams } from "./RemoteControlEnableParams"; export type { RemoteControlStatusChangedNotification } from "./RemoteControlStatusChangedNotification"; export type { RequestPermissionProfile } from "./RequestPermissionProfile"; export type { ResidencyRequirement } from "./ResidencyRequirement"; diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index 9311fe905..0c2c88c74 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -864,13 +864,13 @@ client_request_definitions! { }, #[experimental("remoteControl/enable")] RemoteControlEnable => "remoteControl/enable" { - params: #[ts(type = "undefined")] #[serde(skip_serializing_if = "Option::is_none")] Option<()>, + params: #[serde(skip_serializing_if = "Option::is_none")] v2::NullableRemoteControlEnableParams, serialization: global("remote-control"), response: v2::RemoteControlEnableResponse, }, #[experimental("remoteControl/disable")] RemoteControlDisable => "remoteControl/disable" { - params: #[ts(type = "undefined")] #[serde(skip_serializing_if = "Option::is_none")] Option<()>, + params: #[serde(skip_serializing_if = "Option::is_none")] v2::NullableRemoteControlDisableParams, serialization: global("remote-control"), response: v2::RemoteControlDisableResponse, }, diff --git a/codex-rs/app-server-protocol/src/protocol/v2/remote_control.rs b/codex-rs/app-server-protocol/src/protocol/v2/remote_control.rs index 0e1809805..f69bbdc00 100644 --- a/codex-rs/app-server-protocol/src/protocol/v2/remote_control.rs +++ b/codex-rs/app-server-protocol/src/protocol/v2/remote_control.rs @@ -3,6 +3,26 @@ use serde::Deserialize; use serde::Serialize; use ts_rs::TS; +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct RemoteControlEnableParams { + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub ephemeral: bool, +} + +pub type NullableRemoteControlEnableParams = Option; + +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct RemoteControlDisableParams { + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub ephemeral: bool, +} + +pub type NullableRemoteControlDisableParams = Option; + /// Current remote-control connection status and remote identity exposed to clients. #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)] #[serde(rename_all = "camelCase")] diff --git a/codex-rs/app-server-transport/src/lib.rs b/codex-rs/app-server-transport/src/lib.rs index 8608102f3..d0c4d8faa 100644 --- a/codex-rs/app-server-transport/src/lib.rs +++ b/codex-rs/app-server-transport/src/lib.rs @@ -11,8 +11,10 @@ pub use transport::AppServerTransport; pub use transport::AppServerTransportParseError; pub use transport::CHANNEL_CAPACITY; pub use transport::ConnectionOrigin; +pub use transport::REMOTE_CONTROL_DISABLED_ENV_VAR; pub use transport::RemoteControlHandle; pub use transport::RemoteControlStartConfig; +pub use transport::RemoteControlStartupMode; pub use transport::RemoteControlUnavailable; pub use transport::TransportEvent; pub use transport::acquire_app_server_startup_lock; @@ -24,3 +26,4 @@ pub use transport::start_control_socket_acceptor; pub use transport::start_remote_control; pub use transport::start_stdio_connection; pub use transport::start_websocket_acceptor; +pub use transport::take_remote_control_disabled_env; diff --git a/codex-rs/app-server-transport/src/transport/mod.rs b/codex-rs/app-server-transport/src/transport/mod.rs index 7452f0693..b8f867a90 100644 --- a/codex-rs/app-server-transport/src/transport/mod.rs +++ b/codex-rs/app-server-transport/src/transport/mod.rs @@ -30,10 +30,13 @@ mod unix_socket; mod unix_socket_tests; mod websocket; +pub use remote_control::REMOTE_CONTROL_DISABLED_ENV_VAR; pub use remote_control::RemoteControlHandle; pub use remote_control::RemoteControlStartConfig; +pub use remote_control::RemoteControlStartupMode; pub use remote_control::RemoteControlUnavailable; pub use remote_control::start_remote_control; +pub use remote_control::take_remote_control_disabled_env; pub use stdio::start_stdio_connection; pub use unix_socket::AppServerStartupLock; pub use unix_socket::acquire_app_server_startup_lock; diff --git a/codex-rs/app-server-transport/src/transport/remote_control/desired_state.rs b/codex-rs/app-server-transport/src/transport/remote_control/desired_state.rs new file mode 100644 index 000000000..5d4896247 --- /dev/null +++ b/codex-rs/app-server-transport/src/transport/remote_control/desired_state.rs @@ -0,0 +1,158 @@ +use super::RemoteControlHandle; +use super::RemoteControlUnavailable; +use super::enroll::update_persisted_remote_control_enrollment; +use super::protocol::normalize_remote_control_url; +use super::publish_current_enrollment; +use super::websocket::RemoteControlStatusPublisher; +use codex_app_server_protocol::RemoteControlStatusChangedNotification; +use codex_state::RemoteControlEnrollmentRecord; +use std::io; +use tokio::sync::Semaphore; +use tokio::sync::SemaphorePermit; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum RemoteControlDesiredState { + // `Unknown` exists only on plain startup before auth and enrollment scope resolve. Persisted + // `1` is `Enabled { persistence_preference: Some(true) }`; `0`, `NULL`, or no row are + // `Disabled`. Runtime-only enable is `Enabled { persistence_preference: None }`, so new rows + // keep `NULL`; durable RPC enable uses `Some(true)`, so new rows get `1`. Durable disable writes + // `0` before entering `Disabled`; runtime-only disable does not write. `Disabled` carries no + // preference because disabled sessions do not create enrollments. + Unknown, + Disabled, + Enabled { + persistence_preference: Option, + }, +} +impl RemoteControlDesiredState { + pub(super) fn is_enabled(self) -> bool { + matches!(self, Self::Enabled { .. }) + } +} + +pub(super) async fn acquire_persistence_lock(lock: &Semaphore) -> SemaphorePermit<'_> { + lock.acquire().await.unwrap_or_else(|_| unreachable!()) +} + +pub(super) fn desired_state_from_persisted_enrollment( + enrollment: Option, +) -> RemoteControlDesiredState { + if enrollment.and_then(|enrollment| enrollment.remote_control_enabled) == Some(true) { + RemoteControlDesiredState::Enabled { + persistence_preference: Some(true), + } + } else { + RemoteControlDesiredState::Disabled + } +} + +impl RemoteControlHandle { + pub async fn resolve_persisted_preference( + &self, + app_server_client_name: Option<&str>, + ) -> io::Result { + let _transition = self + .desired_state_rpc_lock + .acquire() + .await + .unwrap_or_else(|_| unreachable!()); + if !matches!( + *self.desired_state_tx.borrow(), + RemoteControlDesiredState::Unknown + ) { + return Ok(self.desired_state_tx.borrow().is_enabled()); + } + + let state_db = self + .state_db + .as_deref() + .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, RemoteControlUnavailable))?; + let auth = super::auth::load_remote_control_auth(&self.auth_manager).await?; + let remote_control_target = normalize_remote_control_url(&self.remote_control_url)?; + let app_server_client_name = self.pairing_persistence_key(app_server_client_name)?; + let enrollment = state_db + .get_remote_control_enrollment( + &remote_control_target.websocket_url, + &auth.account_id, + app_server_client_name.as_deref(), + ) + .await + .map_err(io::Error::other)?; + let desired_state = desired_state_from_persisted_enrollment(enrollment); + self.desired_state_tx.send_if_modified(|state| { + if !matches!(*state, RemoteControlDesiredState::Unknown) { + return false; + } + *state = desired_state; + true + }); + Ok(self.desired_state_tx.borrow().is_enabled()) + } + + pub async fn enable( + &self, + app_server_client_name: Option<&str>, + ) -> io::Result { + let _transition = self + .desired_state_rpc_lock + .acquire() + .await + .unwrap_or_else(|_| unreachable!()); + let state_db = self + .state_db + .as_deref() + .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, RemoteControlUnavailable))?; + let mut auth = super::auth::load_remote_control_auth(&self.auth_manager).await?; + let remote_control_target = normalize_remote_control_url(&self.remote_control_url)?; + let app_server_client_name = self.pairing_persistence_key(app_server_client_name)?; + let app_server_client_name = app_server_client_name.as_deref(); + let status = self.status(); + let mut current_enrollment = self.current_enrollment.lock().await; + let (enrollment, _) = self + .load_or_enroll_server( + ¤t_enrollment, + &mut auth, + &status.installation_id, + &status.server_name, + app_server_client_name, + super::RemoteControlEnrollmentSelection::ReuseOrCreate, + ) + .await?; + + let current_auth = super::auth::load_remote_control_auth(&self.auth_manager).await?; + if current_auth.account_id != auth.account_id { + return Err(io::Error::new( + io::ErrorKind::Interrupted, + "remote control account changed during enrollment", + )); + } + + let _persistence = acquire_persistence_lock(&self.desired_state_persistence_lock).await; + let updated = state_db + .set_remote_control_enabled( + &remote_control_target.websocket_url, + &auth.account_id, + app_server_client_name, + /*remote_control_enabled*/ true, + ) + .await + .map_err(io::Error::other)?; + if updated == 0 { + update_persisted_remote_control_enrollment( + Some(state_db), + &remote_control_target, + &auth.account_id, + app_server_client_name, + Some(&enrollment), + Some(true), + ) + .await?; + } + publish_current_enrollment(&mut current_enrollment, &enrollment); + self.enable_with_preference(Some(true)) + .map_err(|err| io::Error::new(io::ErrorKind::NotFound, err))?; + RemoteControlStatusPublisher::new(self.status_tx.as_ref().clone()) + .publish_environment_id(Some(enrollment.environment_id)); + Ok(self.status()) + } +} diff --git a/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs b/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs index b4506a3d4..3d6c74cce 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs @@ -286,6 +286,7 @@ pub(super) async fn update_persisted_remote_control_enrollment( account_id: &str, app_server_client_name: Option<&str>, enrollment: Option<&RemoteControlEnrollment>, + remote_control_enabled: Option, ) -> io::Result<()> { let Some(state_db) = state_db else { return Err(io::Error::new( @@ -316,6 +317,7 @@ pub(super) async fn update_persisted_remote_control_enrollment( server_id: enrollment.server_id.clone(), environment_id: enrollment.environment_id.clone(), server_name: enrollment.server_name.clone(), + remote_control_enabled, }) .await .map_err(io::Error::other)?; @@ -648,6 +650,7 @@ mod tests { "account-a", Some("desktop-client"), Some(&first_enrollment), + /*remote_control_enabled*/ None, ) .await .expect("first enrollment should persist"); @@ -657,6 +660,7 @@ mod tests { "account-a", Some("desktop-client"), Some(&second_enrollment), + /*remote_control_enabled*/ None, ) .await .expect("second enrollment should persist"); @@ -730,6 +734,7 @@ mod tests { "account-a", /*app_server_client_name*/ None, Some(&first_enrollment), + /*remote_control_enabled*/ None, ) .await .expect("first enrollment should persist"); @@ -739,6 +744,7 @@ mod tests { "account-a", /*app_server_client_name*/ None, Some(&second_enrollment), + /*remote_control_enabled*/ None, ) .await .expect("second enrollment should persist"); @@ -749,6 +755,7 @@ mod tests { "account-a", /*app_server_client_name*/ None, /*enrollment*/ None, + /*remote_control_enabled*/ None, ) .await .expect("matching enrollment should clear"); diff --git a/codex-rs/app-server-transport/src/transport/remote_control/mod.rs b/codex-rs/app-server-transport/src/transport/remote_control/mod.rs index 6e8f31d45..d7d147c0d 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/mod.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/mod.rs @@ -1,6 +1,7 @@ mod auth; mod client_tracker; mod clients; +mod desired_state; mod enroll; mod protocol; mod segment; @@ -8,6 +9,8 @@ mod websocket; use self::auth::load_remote_control_auth; use self::auth::recover_remote_control_auth; +use self::desired_state::RemoteControlDesiredState; +use self::desired_state::acquire_persistence_lock; use self::enroll::RemoteControlEnrollment; use self::enroll::enroll_remote_control_server; use self::enroll::load_persisted_remote_control_enrollment; @@ -63,6 +66,26 @@ pub struct RemoteControlStartConfig { pub installation_id: String, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RemoteControlStartupMode { + ResolvePersisted, + DisabledEphemeral, + EnabledEphemeral, +} + +/// Internal marker used by the daemon to disable remote control without requiring a new CLI flag. +pub const REMOTE_CONTROL_DISABLED_ENV_VAR: &str = + "CODEX_INTERNAL_APP_SERVER_REMOTE_CONTROL_DISABLED"; + +/// Reads and removes the daemon's internal disabled-start marker before worker threads start. +pub fn take_remote_control_disabled_env() -> bool { + let disabled = + std::env::var_os(REMOTE_CONTROL_DISABLED_ENV_VAR).is_some_and(|value| value == "1"); + // SAFETY: app-server calls this synchronously at process startup, before spawning threads. + unsafe { std::env::remove_var(REMOTE_CONTROL_DISABLED_ENV_VAR) }; + disabled +} + pub(super) struct QueuedServerEnvelope { pub(super) event: ServerEvent, pub(super) client_id: ClientId, @@ -72,9 +95,10 @@ pub(super) struct QueuedServerEnvelope { #[derive(Clone)] pub struct RemoteControlHandle { - enabled_tx: Arc>, + desired_state_tx: Arc>, + desired_state_rpc_lock: Arc, + desired_state_persistence_lock: Arc, status_tx: Arc>, - state_db_available: bool, state_db: Option>, remote_control_url: String, current_enrollment: CurrentRemoteControlEnrollment, @@ -83,11 +107,17 @@ pub struct RemoteControlHandle { auth_manager: Arc, } -// Pairing and websocket connect share one selected server so they cannot enroll or clear +// Pairing and websocket connect share one selected server so they cannot enroll or replace // different persisted rows while either path is awaiting backend I/O. type CurrentRemoteControlEnrollment = Arc; type RemoteControlPairingPersistenceKey = watch::Sender>; +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum RemoteControlEnrollmentSelection { + ReuseOrCreate, + ReplaceExisting, +} + struct RemoteControlEnrollmentState { enrollment: StdMutex>, lock: Semaphore, @@ -166,23 +196,45 @@ impl fmt::Display for RemoteControlUnavailable { impl Error for RemoteControlUnavailable {} impl RemoteControlHandle { - pub fn enable( + pub fn enable_ephemeral( &self, ) -> Result { - if !self.state_db_available { + self.enable_with_preference(/*persistence_preference*/ None) + } + + fn enable_with_preference( + &self, + persistence_preference: Option, + ) -> Result { + if self.state_db.is_none() { warn!("remote control cannot be enabled because sqlite state db is unavailable"); return Err(RemoteControlUnavailable); } - let enabled_changed = self.enabled_tx.send_if_modified(|state| { - let changed = !*state; - *state = true; + let mut effective_persistence_preference = persistence_preference; + let desired_state_changed = self.desired_state_tx.send_if_modified(|state| { + if effective_persistence_preference.is_none() + && matches!( + *state, + RemoteControlDesiredState::Enabled { + persistence_preference: Some(true) + } + ) + { + effective_persistence_preference = Some(true); + } + let next_state = RemoteControlDesiredState::Enabled { + persistence_preference: effective_persistence_preference, + }; + let changed = *state != next_state; + *state = next_state; changed }); let status = self.status(); info!( - enabled_changed, + desired_state_changed, + ?effective_persistence_preference, current_status = ?status.status, environment_id = ?status.environment_id, installation_id = %status.installation_id, @@ -199,15 +251,43 @@ impl RemoteControlHandle { Ok(self.publish_status(RemoteControlConnectionStatus::Connecting)) } - pub fn disable(&self) -> RemoteControlStatusChangedNotification { - let enabled_changed = self.enabled_tx.send_if_modified(|state| { - let changed = *state; - *state = false; + pub async fn disable( + &self, + app_server_client_name: Option<&str>, + ) -> io::Result { + let _transition = self + .desired_state_rpc_lock + .acquire() + .await + .unwrap_or_else(|_| unreachable!()); + let _persistence = acquire_persistence_lock(&self.desired_state_persistence_lock).await; + self.persist_preference( + app_server_client_name, + /*remote_control_enabled*/ false, + ) + .await?; + Ok(self.transition_disabled()) + } + + pub async fn disable_ephemeral(&self) -> RemoteControlStatusChangedNotification { + let _transition = self + .desired_state_rpc_lock + .acquire() + .await + .unwrap_or_else(|_| unreachable!()); + let _persistence = acquire_persistence_lock(&self.desired_state_persistence_lock).await; + self.transition_disabled() + } + + fn transition_disabled(&self) -> RemoteControlStatusChangedNotification { + let desired_state_changed = self.desired_state_tx.send_if_modified(|state| { + let changed = *state != RemoteControlDesiredState::Disabled; + *state = RemoteControlDesiredState::Disabled; changed }); let status = self.status(); info!( - enabled_changed, + desired_state_changed, current_status = ?status.status, environment_id = ?status.environment_id, installation_id = %status.installation_id, @@ -217,6 +297,30 @@ impl RemoteControlHandle { self.publish_status(RemoteControlConnectionStatus::Disabled) } + async fn persist_preference( + &self, + app_server_client_name: Option<&str>, + remote_control_enabled: bool, + ) -> io::Result<()> { + let state_db = self + .state_db + .as_deref() + .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, RemoteControlUnavailable))?; + let auth = load_remote_control_auth(&self.auth_manager).await?; + let remote_control_target = normalize_remote_control_url(&self.remote_control_url)?; + let app_server_client_name = self.pairing_persistence_key(app_server_client_name)?; + state_db + .set_remote_control_enabled( + &remote_control_target.websocket_url, + &auth.account_id, + app_server_client_name.as_deref(), + remote_control_enabled, + ) + .await + .map_err(io::Error::other)?; + Ok(()) + } + pub fn status(&self) -> RemoteControlStatusChangedNotification { self.status_tx.borrow().clone() } @@ -230,6 +334,9 @@ impl RemoteControlHandle { params: RemoteControlPairingStartParams, app_server_client_name: Option<&str>, ) -> io::Result { + if !self.desired_state_tx.borrow().is_enabled() { + return Err(Self::pairing_disabled_error()); + } let mut auth = load_remote_control_auth(&self.auth_manager) .await .map_err(|_| pairing_unavailable_error())?; @@ -245,19 +352,35 @@ impl RemoteControlHandle { &installation_id, &status.server_name, app_server_client_name, + RemoteControlEnrollmentSelection::ReuseOrCreate, ) .await?; if enrollment.should_refresh_server_token() { - refresh_pairing_enrollment( + let refresh_result = refresh_pairing_enrollment( &mut current_enrollment, - self.state_db.as_deref(), - app_server_client_name, &self.auth_manager, &mut auth, &installation_id, &mut enrollment, ) - .await?; + .await; + if refresh_result + .as_ref() + .is_err_and(|err| err.kind() == io::ErrorKind::NotFound) + { + enrollment = self + .load_or_enroll_pairing_server( + &mut current_enrollment, + &mut auth, + &installation_id, + &status.server_name, + app_server_client_name, + RemoteControlEnrollmentSelection::ReplaceExisting, + ) + .await?; + } else { + refresh_result?; + } } let pairing_request = || protocol::StartRemoteControlPairingRequest { manual_code: params.manual_code, @@ -267,8 +390,6 @@ impl RemoteControlHandle { clear_pairing_server_token(&mut current_enrollment, &mut enrollment)?; refresh_pairing_enrollment( &mut current_enrollment, - self.state_db.as_deref(), - app_server_client_name, &self.auth_manager, &mut auth, &installation_id, @@ -278,13 +399,6 @@ impl RemoteControlHandle { enrollment.start_pairing(pairing_request()).await } Err(err) if err.kind() == io::ErrorKind::NotFound => { - clear_pairing_enrollment( - &mut current_enrollment, - self.state_db.as_deref(), - app_server_client_name, - &enrollment, - ) - .await; enrollment = self .load_or_enroll_pairing_server( &mut current_enrollment, @@ -292,6 +406,7 @@ impl RemoteControlHandle { &installation_id, &status.server_name, app_server_client_name, + RemoteControlEnrollmentSelection::ReplaceExisting, ) .await?; enrollment.start_pairing(pairing_request()).await @@ -301,13 +416,15 @@ impl RemoteControlHandle { if let Err(err) = &pairing_response { match err.kind() { io::ErrorKind::NotFound => { - clear_pairing_enrollment( + self.load_or_enroll_pairing_server( &mut current_enrollment, - self.state_db.as_deref(), + &mut auth, + &installation_id, + &status.server_name, app_server_client_name, - &enrollment, + RemoteControlEnrollmentSelection::ReplaceExisting, ) - .await; + .await?; return Err(pairing_unavailable_error()); } io::ErrorKind::PermissionDenied => { @@ -323,6 +440,9 @@ impl RemoteControlHandle { if current_auth.account_id != auth.account_id { return Err(pairing_unavailable_error()); } + if !self.desired_state_tx.borrow().is_enabled() { + return Err(Self::pairing_disabled_error()); + } pairing_response } @@ -333,31 +453,86 @@ impl RemoteControlHandle { installation_id: &str, server_name: &str, app_server_client_name: Option<&str>, + selection: RemoteControlEnrollmentSelection, ) -> io::Result { - if let Some(enrollment) = current_enrollment - .as_ref() - .filter(|enrollment| enrollment.account_id == auth.account_id) - .cloned() - { + let (enrollment, created) = self + .load_or_enroll_server( + current_enrollment, + auth, + installation_id, + server_name, + app_server_client_name, + selection, + ) + .await?; + if !created { + publish_current_enrollment(current_enrollment, &enrollment); return Ok(enrollment); } - let remote_control_target = normalize_remote_control_url(&self.remote_control_url)?; let state_db = self .state_db .as_deref() .ok_or_else(pairing_unavailable_error)?; - if let Some(mut enrollment) = load_persisted_remote_control_enrollment( + let _persistence = acquire_persistence_lock(&self.desired_state_persistence_lock).await; + let persistence_preference = match *self.desired_state_tx.borrow() { + RemoteControlDesiredState::Enabled { + persistence_preference, + } => persistence_preference, + RemoteControlDesiredState::Unknown | RemoteControlDesiredState::Disabled => { + return Err(Self::pairing_disabled_error()); + } + }; + update_persisted_remote_control_enrollment( Some(state_db), - &remote_control_target, + &enrollment.remote_control_target, &auth.account_id, app_server_client_name, + Some(&enrollment), + persistence_preference, ) - .await? - { - enrollment.server_name = server_name.to_string(); - publish_current_enrollment(current_enrollment, &enrollment); - return Ok(enrollment); + .await?; + publish_current_enrollment(current_enrollment, &enrollment); + Ok(enrollment) + } + + async fn load_or_enroll_server( + &self, + current_enrollment: &Option, + auth: &mut auth::RemoteControlConnectionAuth, + installation_id: &str, + server_name: &str, + app_server_client_name: Option<&str>, + selection: RemoteControlEnrollmentSelection, + ) -> io::Result<(RemoteControlEnrollment, bool)> { + let remote_control_target = normalize_remote_control_url(&self.remote_control_url)?; + match selection { + RemoteControlEnrollmentSelection::ReuseOrCreate => { + if let Some(enrollment) = current_enrollment + .as_ref() + .filter(|enrollment| enrollment.account_id == auth.account_id) + .cloned() + { + return Ok((enrollment, false)); + } + + let state_db = self + .state_db + .as_deref() + .ok_or_else(pairing_unavailable_error)?; + if let Some(mut enrollment) = load_persisted_remote_control_enrollment( + Some(state_db), + &remote_control_target, + &auth.account_id, + app_server_client_name, + ) + .await? + { + enrollment.server_name = server_name.to_string(); + return Ok((enrollment, false)); + } + } + RemoteControlEnrollmentSelection::ReplaceExisting => {} } let enrollment = enroll_pairing_server( @@ -368,16 +543,7 @@ impl RemoteControlHandle { server_name, ) .await?; - update_persisted_remote_control_enrollment( - Some(state_db), - &remote_control_target, - &auth.account_id, - app_server_client_name, - Some(&enrollment), - ) - .await?; - publish_current_enrollment(current_enrollment, &enrollment); - Ok(enrollment) + Ok((enrollment, true)) } fn pairing_persistence_key( @@ -398,7 +564,7 @@ impl RemoteControlHandle { &self, params: RemoteControlPairingStatusParams, ) -> io::Result { - if !*self.enabled_tx.borrow() { + if !self.desired_state_tx.borrow().is_enabled() { return Err(Self::pairing_disabled_error()); } let mut auth = load_remote_control_auth(&self.auth_manager) @@ -412,18 +578,34 @@ impl RemoteControlHandle { .filter(|enrollment| enrollment.account_id == auth.account_id) .cloned() .ok_or_else(pairing_unavailable_error)?; - let installation_id = self.status().installation_id; + let status = self.status(); + let installation_id = status.installation_id; + let server_name = status.server_name; if enrollment.should_refresh_server_token() { - refresh_pairing_enrollment( + let refresh_result = refresh_pairing_enrollment( &mut current_enrollment, - self.state_db.as_deref(), - app_server_client_name, &self.auth_manager, &mut auth, &installation_id, &mut enrollment, ) - .await?; + .await; + if refresh_result + .as_ref() + .is_err_and(|err| err.kind() == io::ErrorKind::NotFound) + { + self.load_or_enroll_pairing_server( + &mut current_enrollment, + &mut auth, + &installation_id, + &server_name, + app_server_client_name, + RemoteControlEnrollmentSelection::ReplaceExisting, + ) + .await?; + return Err(pairing_unavailable_error()); + } + refresh_result?; } let status_code = remote_control_pairing_status_code(¶ms)?; let pairing_status_request = @@ -434,8 +616,6 @@ impl RemoteControlHandle { clear_pairing_server_token(&mut current_enrollment, &mut enrollment)?; refresh_pairing_enrollment( &mut current_enrollment, - self.state_db.as_deref(), - app_server_client_name, &self.auth_manager, &mut auth, &installation_id, @@ -449,13 +629,15 @@ impl RemoteControlHandle { if let Err(err) = &pairing_status_response { match err.kind() { io::ErrorKind::NotFound => { - clear_pairing_enrollment( + self.load_or_enroll_pairing_server( &mut current_enrollment, - self.state_db.as_deref(), + &mut auth, + &installation_id, + &server_name, app_server_client_name, - &enrollment, + RemoteControlEnrollmentSelection::ReplaceExisting, ) - .await; + .await?; return Err(pairing_unavailable_error()); } io::ErrorKind::PermissionDenied => { @@ -465,7 +647,7 @@ impl RemoteControlHandle { _ => {} } } - if !*self.enabled_tx.borrow() { + if !self.desired_state_tx.borrow().is_enabled() { return Err(Self::pairing_disabled_error()); } let current_auth = load_remote_control_auth(&self.auth_manager) @@ -580,8 +762,6 @@ fn remote_control_pairing_status_code( async fn refresh_pairing_enrollment( current_enrollment: &mut Option, - state_db: Option<&StateRuntime>, - app_server_client_name: Option<&str>, auth_manager: &Arc, auth: &mut auth::RemoteControlConnectionAuth, installation_id: &str, @@ -589,14 +769,7 @@ async fn refresh_pairing_enrollment( ) -> io::Result<()> { if let Err(err) = refresh_remote_control_server(auth, installation_id, enrollment).await { if err.kind() != io::ErrorKind::PermissionDenied { - return handle_pairing_refresh_error( - current_enrollment, - state_db, - app_server_client_name, - enrollment, - err, - ) - .await; + return Err(err); } let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut auth_change_rx = auth_manager.auth_change_receiver(); @@ -609,16 +782,7 @@ async fn refresh_pairing_enrollment( if auth.account_id != enrollment.account_id { return Err(pairing_unavailable_error()); } - if let Err(err) = refresh_remote_control_server(auth, installation_id, enrollment).await { - return handle_pairing_refresh_error( - current_enrollment, - state_db, - app_server_client_name, - enrollment, - err, - ) - .await; - } + refresh_remote_control_server(auth, installation_id, enrollment).await? } if replace_current_enrollment(current_enrollment, enrollment) { Ok(()) @@ -627,52 +791,6 @@ async fn refresh_pairing_enrollment( } } -async fn handle_pairing_refresh_error( - current_enrollment: &mut Option, - state_db: Option<&StateRuntime>, - app_server_client_name: Option<&str>, - enrollment: &RemoteControlEnrollment, - err: io::Error, -) -> io::Result<()> { - if err.kind() == io::ErrorKind::NotFound { - clear_pairing_enrollment( - current_enrollment, - state_db, - app_server_client_name, - enrollment, - ) - .await; - Err(pairing_unavailable_error()) - } else { - Err(err) - } -} - -async fn clear_pairing_enrollment( - current_enrollment: &mut Option, - state_db: Option<&StateRuntime>, - app_server_client_name: Option<&str>, - enrollment: &RemoteControlEnrollment, -) { - if !clear_current_enrollment_if_matches(current_enrollment, enrollment) { - return; - } - let Some(state_db) = state_db else { - return; - }; - if let Err(err) = update_persisted_remote_control_enrollment( - Some(state_db), - &enrollment.remote_control_target, - &enrollment.account_id, - app_server_client_name, - /*enrollment*/ None, - ) - .await - { - warn!("failed to clear stale pairing enrollment: {err}"); - } -} - fn clear_pairing_server_token( current_enrollment: &mut Option, enrollment: &mut RemoteControlEnrollment, @@ -729,21 +847,6 @@ fn replace_current_enrollment( true } -fn clear_current_enrollment_if_matches( - current_enrollment: &mut Option, - enrollment: &RemoteControlEnrollment, -) -> bool { - if current_enrollment - .as_ref() - .is_some_and(|current| same_remote_control_enrollment(current, enrollment)) - { - *current_enrollment = None; - true - } else { - false - } -} - fn same_remote_control_enrollment( left: &RemoteControlEnrollment, right: &RemoteControlEnrollment, @@ -762,11 +865,22 @@ pub async fn start_remote_control( transport_event_tx: mpsc::Sender, shutdown_token: CancellationToken, app_server_client_name_rx: Option>, - initial_enabled: bool, + startup_mode: RemoteControlStartupMode, ) -> io::Result<(JoinHandle<()>, RemoteControlHandle)> { let state_db_available = state_db.is_some(); - let requested_initial_enabled = initial_enabled; - let initial_enabled = initial_enabled && state_db_available; + let requested_initial_enabled = startup_mode == RemoteControlStartupMode::EnabledEphemeral; + let desired_state = if !state_db_available { + RemoteControlDesiredState::Disabled + } else { + match startup_mode { + RemoteControlStartupMode::ResolvePersisted => RemoteControlDesiredState::Unknown, + RemoteControlStartupMode::DisabledEphemeral => RemoteControlDesiredState::Disabled, + RemoteControlStartupMode::EnabledEphemeral => RemoteControlDesiredState::Enabled { + persistence_preference: None, + }, + } + }; + let initial_enabled = desired_state.is_enabled(); if requested_initial_enabled && !state_db_available { warn!("remote control disabled because sqlite state db is unavailable"); } @@ -776,7 +890,12 @@ pub async fn start_remote_control( None }; - let (enabled_tx, enabled_rx) = watch::channel(initial_enabled); + let (desired_state_tx, _desired_state_rx) = watch::channel(desired_state); + let desired_state_tx = Arc::new(desired_state_tx); + let desired_state_rpc_lock = Arc::new(Semaphore::new(1)); + let desired_state_persistence_lock = Arc::new(Semaphore::new(1)); + let websocket_desired_state_tx = desired_state_tx.clone(); + let websocket_desired_state_persistence_lock = desired_state_persistence_lock.clone(); let current_enrollment = Arc::new(RemoteControlEnrollmentState::new(/*enrollment*/ None)); let websocket_current_enrollment = current_enrollment.clone(); let pairing_persistence_key_required = app_server_client_name_rx.is_some(); @@ -804,7 +923,7 @@ pub async fn start_remote_control( installation_id = %installation_id, server_name = %server_name, state_db_available, - initial_enabled, + ?desired_state, "starting app-server remote control websocket task" ); let remote_control_url_for_log = remote_control_url.clone(); @@ -817,7 +936,7 @@ pub async fn start_remote_control( remote_control_url = %remote_control_url_for_log, installation_id = %installation_id_for_log, server_name = %server_name_for_log, - initial_enabled, + ?desired_state, "app-server remote control websocket task started" ); let websocket_task = RemoteControlWebsocket::new( @@ -834,9 +953,10 @@ pub async fn start_remote_control( status_publisher, current_enrollment: websocket_current_enrollment, pairing_persistence_key: websocket_pairing_persistence_key, + desired_state_persistence_lock: websocket_desired_state_persistence_lock, }, shutdown_token, - enabled_rx, + websocket_desired_state_tx, ) .run(app_server_client_name_rx); match AssertUnwindSafe(websocket_task).catch_unwind().await { @@ -875,9 +995,10 @@ pub async fn start_remote_control( Ok(( join_handle, RemoteControlHandle { - enabled_tx: Arc::new(enabled_tx), + desired_state_tx, + desired_state_rpc_lock, + desired_state_persistence_lock, status_tx: Arc::new(status_tx), - state_db_available, state_db: handle_state_db, remote_control_url: handle_remote_control_url, current_enrollment, diff --git a/codex-rs/app-server-transport/src/transport/remote_control/tests.rs b/codex-rs/app-server-transport/src/transport/remote_control/tests.rs index 1d4640485..ecc9acafa 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/tests.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/tests.rs @@ -9,6 +9,8 @@ use super::protocol::ClientId; use super::protocol::StreamId; use super::protocol::normalize_remote_control_url; use super::websocket::REMOTE_CONTROL_PROTOCOL_VERSION; +use super::websocket::RemoteControlWebsocket; +use super::websocket::RemoteControlWebsocketConfig; use super::*; use crate::outgoing_message::OutgoingMessage; use crate::outgoing_message::QueuedOutgoingMessage; @@ -33,6 +35,7 @@ use codex_login::CodexAuth; use codex_login::save_auth; use codex_login::token_data::TokenData; use codex_login::token_data::parse_chatgpt_jwt_claims; +use codex_state::RemoteControlEnrollmentRecord; use codex_state::StateRuntime; use futures::SinkExt; use futures::StreamExt; @@ -126,6 +129,137 @@ async fn remote_control_state_runtime(codex_home: &TempDir) -> Arc .expect("state runtime should initialize") } +#[tokio::test] +async fn plain_start_resolves_persisted_remote_control_preference() { + let cases = [ + ("enabled", Some(Some(true))), + ("disabled", Some(Some(false))), + ("unset", Some(None)), + ("missing", None), + ]; + 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(TEST_REMOTE_CONTROL_URL) + .expect("remote control target should normalize"); + for (name, stored_preference) in cases { + let Some(remote_control_enabled) = stored_preference else { + continue; + }; + state_db + .upsert_remote_control_enrollment(&RemoteControlEnrollmentRecord { + websocket_url: remote_control_target.websocket_url.clone(), + account_id: "account_id".to_string(), + app_server_client_name: Some(name.to_string()), + server_id: format!("server-{name}"), + environment_id: format!("environment-{name}"), + server_name: format!("server-name-{name}"), + remote_control_enabled, + }) + .await + .expect("enrollment should persist"); + } + let (transport_event_tx, _transport_event_rx) = mpsc::channel(CHANNEL_CAPACITY); + let (status_tx, _status_rx) = watch::channel(RemoteControlStatusChangedNotification { + status: RemoteControlConnectionStatus::Disabled, + server_name: test_server_name(), + installation_id: TEST_INSTALLATION_ID.to_string(), + environment_id: None, + }); + let (desired_state_tx, _desired_state_rx) = watch::channel(RemoteControlDesiredState::Unknown); + let desired_state_tx = Arc::new(desired_state_tx); + let mut websocket = RemoteControlWebsocket::new( + RemoteControlWebsocketConfig { + remote_control_url: TEST_REMOTE_CONTROL_URL.to_string(), + installation_id: TEST_INSTALLATION_ID.to_string(), + remote_control_target: None, + server_name: test_server_name(), + }, + Some(state_db), + remote_control_auth_manager(), + RemoteControlChannels { + transport_event_tx, + status_publisher: RemoteControlStatusPublisher::new(status_tx), + current_enrollment: Arc::new(RemoteControlEnrollmentState::new( + /*enrollment*/ None, + )), + pairing_persistence_key: watch::channel(None).0, + desired_state_persistence_lock: Arc::new(Semaphore::new(1)), + }, + CancellationToken::new(), + desired_state_tx.clone(), + ); + + for (name, stored_preference) in cases { + desired_state_tx.send_replace(RemoteControlDesiredState::Unknown); + assert!(websocket.resolve_unknown_desired_state(Some(name)).await); + let expected = if stored_preference == Some(Some(true)) { + RemoteControlDesiredState::Enabled { + persistence_preference: Some(true), + } + } else { + RemoteControlDesiredState::Disabled + }; + assert_eq!(*desired_state_tx.borrow(), expected, "case {name}"); + } +} + +#[tokio::test] +async fn explicit_disabled_start_ignores_persisted_enable() { + 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(TEST_REMOTE_CONTROL_URL) + .expect("remote control target should normalize"); + let enrollment = RemoteControlEnrollmentRecord { + websocket_url: remote_control_target.websocket_url, + account_id: "account_id".to_string(), + app_server_client_name: None, + server_id: "server-id".to_string(), + environment_id: "environment-id".to_string(), + server_name: "server-name".to_string(), + remote_control_enabled: Some(true), + }; + state_db + .upsert_remote_control_enrollment(&enrollment) + .await + .expect("enrollment should persist"); + let (transport_event_tx, _transport_event_rx) = mpsc::channel(CHANNEL_CAPACITY); + let shutdown_token = CancellationToken::new(); + + let (remote_task, remote_handle) = start_remote_control( + RemoteControlStartConfig { + remote_control_url: TEST_REMOTE_CONTROL_URL.to_string(), + installation_id: TEST_INSTALLATION_ID.to_string(), + }, + Some(state_db.clone()), + remote_control_auth_manager(), + transport_event_tx, + shutdown_token.clone(), + /*app_server_client_name_rx*/ None, + RemoteControlStartupMode::DisabledEphemeral, + ) + .await + .expect("remote control should start disabled"); + + assert_eq!( + *remote_handle.desired_state_tx.borrow(), + RemoteControlDesiredState::Disabled + ); + assert_eq!( + state_db + .get_remote_control_enrollment( + &enrollment.websocket_url, + &enrollment.account_id, + /*app_server_client_name*/ None, + ) + .await + .expect("enrollment should load"), + Some(enrollment) + ); + + shutdown_token.cancel(); + remote_task.await.expect("remote control task should join"); +} + fn remote_control_url_for_listener(listener: &TcpListener) -> String { let addr = listener .local_addr() @@ -141,7 +275,10 @@ fn remote_control_handle_with_current_enrollment( remote_control_url: &str, auth_manager: Arc, ) -> RemoteControlHandle { - let (enabled_tx, _enabled_rx) = watch::channel(/*init*/ true); + let (desired_state_tx, _desired_state_rx) = + watch::channel(RemoteControlDesiredState::Enabled { + persistence_preference: None, + }); let (status_tx, _status_rx) = watch::channel(RemoteControlStatusChangedNotification { status: RemoteControlConnectionStatus::Connecting, server_name: test_server_name(), @@ -165,9 +302,10 @@ fn remote_control_handle_with_current_enrollment( }, ))); RemoteControlHandle { - enabled_tx: Arc::new(enabled_tx), + desired_state_tx: Arc::new(desired_state_tx), + desired_state_rpc_lock: Arc::new(Semaphore::new(1)), + desired_state_persistence_lock: Arc::new(Semaphore::new(1)), status_tx: Arc::new(status_tx), - state_db_available: true, state_db: None, remote_control_url: remote_control_url.to_string(), current_enrollment, @@ -177,6 +315,44 @@ fn remote_control_handle_with_current_enrollment( } } +#[tokio::test] +async fn ephemeral_enable_preserves_durable_preference() { + let codex_home = TempDir::new().expect("temp dir should create"); + let mut remote_handle = remote_control_handle_with_current_enrollment( + TEST_REMOTE_CONTROL_URL, + remote_control_auth_manager(), + ); + remote_handle.state_db = Some(remote_control_state_runtime(&codex_home).await); + remote_handle + .desired_state_tx + .send_replace(RemoteControlDesiredState::Enabled { + persistence_preference: Some(true), + }); + + remote_handle + .enable_ephemeral() + .expect("ephemeral enable should succeed"); + assert_eq!( + *remote_handle.desired_state_tx.borrow(), + RemoteControlDesiredState::Enabled { + persistence_preference: Some(true), + } + ); + + remote_handle + .desired_state_tx + .send_replace(RemoteControlDesiredState::Disabled); + remote_handle + .enable_ephemeral() + .expect("ephemeral enable should succeed"); + assert_eq!( + *remote_handle.desired_state_tx.borrow(), + RemoteControlDesiredState::Enabled { + persistence_preference: None, + } + ); +} + fn remote_control_server_token_response( server_id: &str, environment_id: &str, @@ -242,7 +418,10 @@ async fn remote_control_transport_manages_virtual_clients_and_routes_messages() .await .expect("listener should bind"); let remote_control_url = remote_control_url_for_listener(&listener); + let remote_control_target = normalize_remote_control_url(&remote_control_url) + .expect("remote control target should normalize"); let codex_home = TempDir::new().expect("temp dir should create"); + let state_db = remote_control_state_runtime(&codex_home).await; let (transport_event_tx, mut transport_event_rx) = mpsc::channel::(CHANNEL_CAPACITY); let shutdown_token = CancellationToken::new(); @@ -251,12 +430,12 @@ async fn remote_control_transport_manages_virtual_clients_and_routes_messages() remote_control_url, installation_id: TEST_INSTALLATION_ID.to_string(), }, - Some(remote_control_state_runtime(&codex_home).await), + Some(state_db.clone()), remote_control_auth_manager(), transport_event_tx, shutdown_token.clone(), /*app_server_client_name_rx*/ None, - /*initial_enabled*/ true, + RemoteControlStartupMode::EnabledEphemeral, ) .await .expect("remote control should start"); @@ -276,6 +455,16 @@ async fn remote_control_transport_manages_virtual_clients_and_routes_messages() ) .await; let mut websocket = accept_remote_control_connection(&listener).await; + let enrollment = state_db + .get_remote_control_enrollment( + &remote_control_target.websocket_url, + "account_id", + /*app_server_client_name*/ None, + ) + .await + .expect("new enrollment should load") + .expect("new enrollment should exist"); + assert_eq!(enrollment.remote_control_enabled, None); expect_remote_control_status( &mut status_rx, /*expected_status*/ None, @@ -539,7 +728,7 @@ async fn remote_control_transport_reconnects_after_disconnect() { transport_event_tx, shutdown_token.clone(), /*app_server_client_name_rx*/ None, - /*initial_enabled*/ true, + RemoteControlStartupMode::EnabledEphemeral, ) .await .expect("remote control should start"); @@ -640,7 +829,7 @@ async fn remote_control_transport_refreshes_server_token_after_websocket_unautho transport_event_tx, shutdown_token.clone(), /*app_server_client_name_rx*/ None, - /*initial_enabled*/ true, + RemoteControlStartupMode::EnabledEphemeral, ) .await .expect("remote control should start"); @@ -720,7 +909,7 @@ async fn remote_control_start_allows_remote_control_invalid_url_when_disabled() transport_event_tx, shutdown_token.clone(), /*app_server_client_name_rx*/ None, - /*initial_enabled*/ false, + RemoteControlStartupMode::ResolvePersisted, ) .await .expect("disabled remote control should not validate the URL at startup"); @@ -759,7 +948,7 @@ async fn remote_control_start_allows_missing_auth_when_enabled() { transport_event_tx, shutdown_token.clone(), /*app_server_client_name_rx*/ None, - /*initial_enabled*/ true, + RemoteControlStartupMode::EnabledEphemeral, ) .await .expect("remote control should start before ChatGPT auth is available"); @@ -794,7 +983,7 @@ async fn remote_control_start_reports_missing_state_db_as_disabled_when_enabled( transport_event_tx, shutdown_token.clone(), /*app_server_client_name_rx*/ None, - /*initial_enabled*/ true, + RemoteControlStartupMode::EnabledEphemeral, ) .await .expect("remote control should start disabled without sqlite state db"); @@ -814,7 +1003,9 @@ async fn remote_control_start_reports_missing_state_db_as_disabled_when_enabled( .expect_err("remote control should not connect without sqlite state db"); assert_eq!( - remote_handle.enable().expect_err("enable should fail"), + remote_handle + .enable_ephemeral() + .expect_err("enable should fail"), super::RemoteControlUnavailable ); timeout(Duration::from_millis(100), listener.accept()) @@ -851,7 +1042,7 @@ async fn remote_control_handle_enable_disable_stops_and_restarts_connections() { transport_event_tx, shutdown_token.clone(), /*app_server_client_name_rx*/ None, - /*initial_enabled*/ true, + RemoteControlStartupMode::EnabledEphemeral, ) .await .expect("remote control should start"); @@ -884,7 +1075,10 @@ async fn remote_control_handle_enable_disable_stops_and_restarts_connections() { .await; assert_eq!( - remote_handle.disable(), + remote_handle + .disable(Some("rpc-client")) + .await + .expect("disable should succeed"), RemoteControlStatusChangedNotification { status: RemoteControlConnectionStatus::Disabled, server_name: test_server_name(), @@ -910,12 +1104,15 @@ async fn remote_control_handle_enable_disable_stops_and_restarts_connections() { .expect_err("disabled remote control should not reconnect"); assert_eq!( - remote_handle.enable().expect("enable should succeed"), + remote_handle + .enable(Some("rpc-client")) + .await + .expect("enable should succeed"), RemoteControlStatusChangedNotification { status: RemoteControlConnectionStatus::Connecting, server_name: test_server_name(), installation_id: TEST_INSTALLATION_ID.to_string(), - environment_id: None, + environment_id: Some("env_test".to_string()), } ); expect_remote_control_status_snapshot( @@ -924,7 +1121,7 @@ async fn remote_control_handle_enable_disable_stops_and_restarts_connections() { status: RemoteControlConnectionStatus::Connecting, server_name: test_server_name(), installation_id: TEST_INSTALLATION_ID.to_string(), - environment_id: None, + environment_id: Some("env_test".to_string()), }, ) .await; @@ -964,7 +1161,7 @@ async fn remote_control_transport_clears_outgoing_buffer_when_backend_acks() { transport_event_tx, shutdown_token.clone(), /*app_server_client_name_rx*/ None, - /*initial_enabled*/ true, + RemoteControlStartupMode::EnabledEphemeral, ) .await .expect("remote control should start"); @@ -1146,7 +1343,7 @@ async fn remote_control_http_mode_enrolls_before_connecting() { transport_event_tx, shutdown_token.clone(), /*app_server_client_name_rx*/ None, - /*initial_enabled*/ true, + RemoteControlStartupMode::EnabledEphemeral, ) .await .expect("remote control should start"); @@ -1376,6 +1573,7 @@ async fn remote_control_http_mode_refreshes_persisted_enrollment_before_connecti "account_id", /*app_server_client_name*/ None, Some(&persisted_enrollment), + /*remote_control_enabled*/ None, ) .await .expect("persisted enrollment should save"); @@ -1393,7 +1591,7 @@ async fn remote_control_http_mode_refreshes_persisted_enrollment_before_connecti transport_event_tx, shutdown_token.clone(), /*app_server_client_name_rx*/ None, - /*initial_enabled*/ true, + RemoteControlStartupMode::EnabledEphemeral, ) .await .expect("remote control should start"); @@ -1482,6 +1680,7 @@ async fn remote_control_stdio_mode_waits_for_client_name_before_connecting() { "account_id", Some(app_server_client_name), Some(&persisted_enrollment), + /*remote_control_enabled*/ None, ) .await .expect("persisted enrollment should save"); @@ -1500,7 +1699,7 @@ async fn remote_control_stdio_mode_waits_for_client_name_before_connecting() { transport_event_tx, shutdown_token.clone(), Some(app_server_client_name_rx), - /*initial_enabled*/ true, + RemoteControlStartupMode::EnabledEphemeral, ) .await .expect("remote control should start"); @@ -1581,7 +1780,7 @@ async fn remote_control_waits_for_account_id_before_enrolling() { transport_event_tx, shutdown_token.clone(), /*app_server_client_name_rx*/ None, - /*initial_enabled*/ true, + RemoteControlStartupMode::EnabledEphemeral, ) .await .expect("remote control should start before account id is available"); @@ -1625,6 +1824,123 @@ async fn remote_control_waits_for_account_id_before_enrolling() { let _ = remote_task.await; } +#[tokio::test] +async fn persisted_enable_does_not_follow_auth_to_an_account_without_a_preference() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let remote_control_url = remote_control_url_for_listener(&listener); + let codex_home = TempDir::new().expect("temp dir should create"); + save_auth( + codex_home.path(), + &remote_control_auth_dot_json(Some("account_a")), + AuthCredentialsStoreMode::File, + ) + .expect("account A auth should save"); + let state_db = remote_control_state_runtime(&codex_home).await; + let auth_manager = AuthManager::shared( + codex_home.path().to_path_buf(), + /*enable_codex_api_key_env*/ false, + AuthCredentialsStoreMode::File, + /*chatgpt_base_url*/ None, + ) + .await; + let remote_control_target = + normalize_remote_control_url(&remote_control_url).expect("target should parse"); + let enrollment = RemoteControlEnrollment { + remote_control_target: remote_control_target.clone(), + account_id: "account_a".to_string(), + environment_id: "env_a".to_string(), + server_id: "srv_e_a".to_string(), + server_name: "server-a".to_string(), + remote_control_token: None, + expires_at: None, + }; + update_persisted_remote_control_enrollment( + Some(state_db.as_ref()), + &remote_control_target, + "account_a", + /*app_server_client_name*/ None, + Some(&enrollment), + /*remote_control_enabled*/ Some(true), + ) + .await + .expect("account A enrollment should save"); + + let (transport_event_tx, _transport_event_rx) = + mpsc::channel::(CHANNEL_CAPACITY); + let shutdown_token = CancellationToken::new(); + let (remote_task, remote_handle) = start_remote_control( + RemoteControlStartConfig { + remote_control_url, + installation_id: TEST_INSTALLATION_ID.to_string(), + }, + Some(state_db.clone()), + auth_manager.clone(), + transport_event_tx, + shutdown_token.clone(), + /*app_server_client_name_rx*/ None, + RemoteControlStartupMode::ResolvePersisted, + ) + .await + .expect("remote control should start"); + + let refresh_request = accept_http_request(&listener).await; + assert_eq!( + refresh_request.request_line, + "POST /backend-api/wham/remote/control/server/refresh HTTP/1.1" + ); + respond_with_json( + refresh_request.stream, + remote_control_server_token_response( + &enrollment.server_id, + &enrollment.environment_id, + TEST_REFRESHED_REMOTE_CONTROL_SERVER_TOKEN, + ), + ) + .await; + let (_handshake_request, mut websocket) = + accept_remote_control_backend_connection(&listener).await; + + save_auth( + codex_home.path(), + &remote_control_auth_dot_json(Some("account_b")), + AuthCredentialsStoreMode::File, + ) + .expect("account B auth should save"); + auth_manager.reload().await; + websocket + .close(None) + .await + .expect("backend websocket should close"); + + let mut desired_state_rx = remote_handle.desired_state_tx.subscribe(); + timeout( + Duration::from_secs(1), + desired_state_rx.wait_for(|state| *state == RemoteControlDesiredState::Disabled), + ) + .await + .expect("account B missing preference should disable remote control") + .expect("desired state channel should stay open"); + timeout(Duration::from_millis(100), listener.accept()) + .await + .expect_err("disabled account B should not enroll"); + assert_eq!( + state_db + .get_remote_control_enrollment( + &remote_control_target.websocket_url, + "account_b", + /*app_server_client_name*/ None, + ) + .await + .expect("account B enrollment should load"), + None + ); + + shutdown_token.cancel(); + let _ = remote_task.await; +} + #[tokio::test] async fn remote_control_http_mode_reenrolls_when_refresh_reports_stale_enrollment() { let listener = TcpListener::bind("127.0.0.1:0") @@ -1660,6 +1976,7 @@ async fn remote_control_http_mode_reenrolls_when_refresh_reports_stale_enrollmen "account_id", /*app_server_client_name*/ None, Some(&stale_enrollment), + /*remote_control_enabled*/ Some(true), ) .await .expect("stale enrollment should save"); @@ -1677,7 +1994,7 @@ async fn remote_control_http_mode_reenrolls_when_refresh_reports_stale_enrollmen transport_event_tx, shutdown_token.clone(), /*app_server_client_name_rx*/ None, - /*initial_enabled*/ true, + RemoteControlStartupMode::ResolvePersisted, ) .await .expect("remote control should start"); @@ -1695,12 +2012,6 @@ async fn remote_control_http_mode_reenrolls_when_refresh_reports_stale_enrollmen ) .await; respond_with_status(refresh_request.stream, "404 Not Found", "").await; - expect_remote_control_status( - &mut status_rx, - /*expected_status*/ None, - /*expected_environment_id*/ None, - ) - .await; let enroll_request = accept_http_request(&listener).await; assert_eq!( @@ -1729,15 +2040,23 @@ async fn remote_control_http_mode_reenrolls_when_refresh_reports_stale_enrollmen Some(&refreshed_enrollment.server_id) ); assert_eq!( - load_persisted_remote_control_enrollment( - Some(state_db.as_ref()), - &remote_control_target, - "account_id", - /*app_server_client_name*/ None, - ) - .await - .expect("refreshed enrollment should load"), - Some(refreshed_enrollment) + state_db + .get_remote_control_enrollment( + &remote_control_target.websocket_url, + "account_id", + /*app_server_client_name*/ None, + ) + .await + .expect("refreshed enrollment should load"), + Some(RemoteControlEnrollmentRecord { + websocket_url: remote_control_target.websocket_url.clone(), + account_id: "account_id".to_string(), + app_server_client_name: None, + server_id: refreshed_enrollment.server_id.clone(), + environment_id: refreshed_enrollment.environment_id.clone(), + server_name: refreshed_enrollment.server_name.clone(), + remote_control_enabled: Some(true), + }) ); shutdown_token.cancel(); @@ -1779,6 +2098,7 @@ async fn remote_control_http_mode_reenrolls_after_explicit_missing_server_404() "account_id", /*app_server_client_name*/ None, Some(&stale_enrollment), + /*remote_control_enabled*/ Some(true), ) .await .expect("stale enrollment should save"); @@ -1796,7 +2116,7 @@ async fn remote_control_http_mode_reenrolls_after_explicit_missing_server_404() transport_event_tx, shutdown_token.clone(), /*app_server_client_name_rx*/ None, - /*initial_enabled*/ true, + RemoteControlStartupMode::ResolvePersisted, ) .await .expect("remote control should start"); @@ -1838,12 +2158,6 @@ async fn remote_control_http_mode_reenrolls_after_explicit_missing_server_404() &json!({"detail": "Remote app server not found"}).to_string(), ) .await; - expect_remote_control_status( - &mut status_rx, - /*expected_status*/ None, - /*expected_environment_id*/ None, - ) - .await; let enroll_request = accept_http_request(&listener).await; assert_eq!( @@ -1872,15 +2186,125 @@ async fn remote_control_http_mode_reenrolls_after_explicit_missing_server_404() Some(&refreshed_enrollment.server_id) ); assert_eq!( - load_persisted_remote_control_enrollment( - Some(state_db.as_ref()), - &remote_control_target, - "account_id", - /*app_server_client_name*/ None, - ) + state_db + .get_remote_control_enrollment( + &remote_control_target.websocket_url, + "account_id", + /*app_server_client_name*/ None, + ) + .await + .expect("refreshed enrollment should load"), + Some(RemoteControlEnrollmentRecord { + websocket_url: remote_control_target.websocket_url.clone(), + account_id: "account_id".to_string(), + app_server_client_name: None, + server_id: refreshed_enrollment.server_id.clone(), + environment_id: refreshed_enrollment.environment_id.clone(), + server_name: refreshed_enrollment.server_name.clone(), + remote_control_enabled: Some(true), + }) + ); + + shutdown_token.cancel(); + let _ = remote_task.await; +} + +#[tokio::test] +async fn remote_control_http_mode_preserves_stale_enrollment_when_reenrollment_fails() { + let listener = TcpListener::bind("127.0.0.1:0") .await - .expect("refreshed enrollment should load"), - Some(refreshed_enrollment) + .expect("listener should bind"); + let remote_control_url = remote_control_url_for_listener(&listener); + let codex_home = TempDir::new().expect("temp dir should create"); + let state_db = remote_control_state_runtime(&codex_home).await; + let remote_control_target = + normalize_remote_control_url(&remote_control_url).expect("target should parse"); + let stale_enrollment = RemoteControlEnrollment { + remote_control_target: remote_control_target.clone(), + account_id: "account_id".to_string(), + environment_id: "env_stale".to_string(), + server_id: "srv_e_stale".to_string(), + server_name: test_server_name(), + remote_control_token: None, + expires_at: None, + }; + update_persisted_remote_control_enrollment( + Some(state_db.as_ref()), + &remote_control_target, + "account_id", + /*app_server_client_name*/ None, + Some(&stale_enrollment), + /*remote_control_enabled*/ Some(true), + ) + .await + .expect("stale enrollment should save"); + + let (transport_event_tx, _transport_event_rx) = + mpsc::channel::(CHANNEL_CAPACITY); + let shutdown_token = CancellationToken::new(); + let (remote_task, remote_handle) = start_remote_control( + RemoteControlStartConfig { + remote_control_url, + installation_id: TEST_INSTALLATION_ID.to_string(), + }, + Some(state_db.clone()), + remote_control_auth_manager_with_home(&codex_home), + transport_event_tx, + shutdown_token.clone(), + /*app_server_client_name_rx*/ None, + RemoteControlStartupMode::ResolvePersisted, + ) + .await + .expect("remote control should start"); + + let refresh_request = accept_http_request(&listener).await; + assert_eq!( + refresh_request.request_line, + "POST /backend-api/wham/remote/control/server/refresh HTTP/1.1" + ); + respond_with_status(refresh_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_status(enroll_request.stream, "500 Internal Server Error", "failed").await; + + let retry_refresh_request = accept_http_request(&listener).await; + assert_eq!( + retry_refresh_request.request_line, + "POST /backend-api/wham/remote/control/server/refresh HTTP/1.1" + ); + respond_with_status( + retry_refresh_request.stream, + "500 Internal Server Error", + "failed", + ) + .await; + + assert_eq!( + *remote_handle.current_enrollment.lock().await, + Some(stale_enrollment.clone()) + ); + assert_eq!( + state_db + .get_remote_control_enrollment( + &remote_control_target.websocket_url, + "account_id", + /*app_server_client_name*/ None, + ) + .await + .expect("stale enrollment should load"), + Some(RemoteControlEnrollmentRecord { + websocket_url: remote_control_target.websocket_url, + account_id: "account_id".to_string(), + app_server_client_name: None, + server_id: stale_enrollment.server_id, + environment_id: stale_enrollment.environment_id, + server_name: stale_enrollment.server_name, + remote_control_enabled: Some(true), + }) ); shutdown_token.cancel(); @@ -1912,6 +2336,7 @@ async fn remote_control_http_mode_preserves_enrollment_after_generic_websocket_4 "account_id", /*app_server_client_name*/ None, Some(&stale_enrollment), + /*remote_control_enabled*/ None, ) .await .expect("stale enrollment should save"); @@ -1929,7 +2354,7 @@ async fn remote_control_http_mode_preserves_enrollment_after_generic_websocket_4 transport_event_tx, shutdown_token.clone(), /*app_server_client_name_rx*/ None, - /*initial_enabled*/ true, + RemoteControlStartupMode::EnabledEphemeral, ) .await .expect("remote control should start"); diff --git a/codex-rs/app-server-transport/src/transport/remote_control/tests/clients_tests.rs b/codex-rs/app-server-transport/src/transport/remote_control/tests/clients_tests.rs index 060063909..fe4408a84 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/tests/clients_tests.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/tests/clients_tests.rs @@ -13,7 +13,7 @@ fn client_management_handle( remote_control_url: String, auth_manager: Arc, ) -> RemoteControlHandle { - let (enabled_tx, _enabled_rx) = watch::channel(/*init*/ false); + let desired_state_tx = watch::channel(RemoteControlDesiredState::Disabled).0; let (status_tx, _status_rx) = watch::channel(RemoteControlStatusChangedNotification { status: RemoteControlConnectionStatus::Disabled, server_name: test_server_name(), @@ -21,9 +21,10 @@ fn client_management_handle( environment_id: None, }); RemoteControlHandle { - enabled_tx: Arc::new(enabled_tx), + desired_state_tx: Arc::new(desired_state_tx), + desired_state_rpc_lock: Arc::new(Semaphore::new(1)), + desired_state_persistence_lock: Arc::new(Semaphore::new(1)), status_tx: Arc::new(status_tx), - state_db_available: false, state_db: None, remote_control_url, current_enrollment: Arc::new(RemoteControlEnrollmentState::new(/*enrollment*/ None)), diff --git a/codex-rs/app-server-transport/src/transport/remote_control/tests/pairing_tests.rs b/codex-rs/app-server-transport/src/transport/remote_control/tests/pairing_tests.rs index fcdd5e30f..4dbb4c6d4 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/tests/pairing_tests.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/tests/pairing_tests.rs @@ -645,7 +645,9 @@ async fn remote_control_handle_disable_keeps_current_enrollment() { remote_control_auth_manager(), ); - remote_handle.disable(); + remote_handle + .desired_state_tx + .send_replace(RemoteControlDesiredState::Disabled); assert!( remote_handle.current_enrollment.lock().await.is_some(), "disabled remote control should keep the selected pairing server" @@ -665,7 +667,6 @@ async fn remote_control_handle_reenrolls_after_stale_pairing_enrollment() { remote_control_auth_manager_with_home(&codex_home), ); remote_handle.state_db = Some(state_db.clone()); - remote_handle.disable(); let stale_enrollment = remote_handle .current_enrollment .lock() @@ -688,9 +689,15 @@ async fn remote_control_handle_reenrolls_after_stale_pairing_enrollment() { "account_id", /*app_server_client_name*/ None, Some(&stale_enrollment), + /*remote_control_enabled*/ Some(true), ) .await .expect("stale enrollment should save"); + remote_handle + .desired_state_tx + .send_replace(RemoteControlDesiredState::Enabled { + persistence_preference: Some(true), + }); let server_refreshed_enrollment = refreshed_enrollment.clone(); let server_task = tokio::spawn(async move { let stale_pairing_request = accept_http_request(&listener).await; @@ -761,15 +768,17 @@ async fn remote_control_handle_reenrolls_after_stale_pairing_enrollment() { } ); assert_eq!( - load_persisted_remote_control_enrollment( - Some(state_db.as_ref()), - &remote_control_target, - "account_id", - /*app_server_client_name*/ None, - ) - .await - .expect("refreshed enrollment should load"), - Some(refreshed_enrollment) + state_db + .get_remote_control_enrollment( + &remote_control_target.websocket_url, + "account_id", + /*app_server_client_name*/ None, + ) + .await + .expect("refreshed enrollment should load") + .expect("refreshed enrollment should exist") + .remote_control_enabled, + Some(true) ); } diff --git a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs index ef74527e0..6d58b2862 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/websocket.rs @@ -1,5 +1,9 @@ use super::CurrentRemoteControlEnrollment; +use super::RemoteControlEnrollmentSelection; use super::RemoteControlPairingPersistenceKey; +use super::desired_state::RemoteControlDesiredState; +use super::desired_state::acquire_persistence_lock; +use super::desired_state::desired_state_from_persisted_enrollment; use super::protocol::ClientEnvelope; use super::protocol::ClientEvent; use super::protocol::ClientId; @@ -45,6 +49,7 @@ use std::io::ErrorKind; use std::sync::Arc; use tokio::net::TcpStream; use tokio::sync::Mutex; +use tokio::sync::Semaphore; use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio::sync::watch; @@ -264,7 +269,9 @@ pub(crate) struct RemoteControlWebsocket { state: Arc>, server_event_rx: Arc>>, used_rx: watch::Receiver, - enabled_rx: watch::Receiver, + desired_state_tx: Arc>, + desired_state_rx: watch::Receiver, + desired_state_persistence_lock: Arc, } pub(crate) struct RemoteControlWebsocketConfig { @@ -280,6 +287,11 @@ pub(super) struct RemoteControlAuthContext<'a> { auth_change_rx: &'a mut watch::Receiver, } +struct RemoteControlEnrollmentAuthContext<'a, 'b> { + auth: &'a RemoteControlConnectionAuth, + recovery: &'a mut RemoteControlAuthContext<'b>, +} + enum ConnectOutcome { Connected(Box>>), Disabled, @@ -299,6 +311,7 @@ pub(super) struct RemoteControlChannels { pub(super) status_publisher: RemoteControlStatusPublisher, pub(super) current_enrollment: CurrentRemoteControlEnrollment, pub(super) pairing_persistence_key: RemoteControlPairingPersistenceKey, + pub(super) desired_state_persistence_lock: Arc, } #[derive(Clone)] @@ -341,7 +354,7 @@ impl RemoteControlStatusPublisher { } } - fn publish_environment_id(&self, environment_id: Option) { + pub(super) fn publish_environment_id(&self, environment_id: Option) { let mut status_change = None; self.tx.send_if_modified(|status| { if status.status == RemoteControlConnectionStatus::Disabled { @@ -380,6 +393,8 @@ pub(super) struct RemoteControlConnectOptions<'a> { server_name: &'a str, subscribe_cursor: Option<&'a str>, app_server_client_name: Option<&'a str>, + desired_state_tx: &'a watch::Sender, + desired_state_persistence_lock: &'a Semaphore, } impl RemoteControlWebsocket { @@ -389,7 +404,7 @@ impl RemoteControlWebsocket { auth_manager: Arc, channels: RemoteControlChannels, shutdown_token: CancellationToken, - enabled_rx: watch::Receiver, + desired_state_tx: Arc>, ) -> Self { let shutdown_token = shutdown_token.child_token(); let (server_event_tx, server_event_rx) = mpsc::channel(super::CHANNEL_CAPACITY); @@ -402,6 +417,7 @@ impl RemoteControlWebsocket { let auth_recovery = auth_manager.unauthorized_recovery(); let auth_change_rx = auth_manager.auth_change_receiver(); + let desired_state_rx = desired_state_tx.subscribe(); Self { remote_control_url: config.remote_control_url, installation_id: config.installation_id, @@ -426,7 +442,9 @@ impl RemoteControlWebsocket { })), server_event_rx: Arc::new(Mutex::new(server_event_rx)), used_rx, - enabled_rx, + desired_state_tx, + desired_state_rx, + desired_state_persistence_lock: channels.desired_state_persistence_lock, } } @@ -442,7 +460,7 @@ impl RemoteControlWebsocket { remote_control_url = %self.remote_control_url, installation_id = %self.installation_id, server_name = %self.server_name, - initial_enabled = *self.enabled_rx.borrow(), + initial_desired_state = ?*self.desired_state_rx.borrow(), "app-server remote control websocket loop started" ); let app_server_client_name = match self @@ -464,6 +482,16 @@ impl RemoteControlWebsocket { }; self.pairing_persistence_key .send_replace(app_server_client_name.clone()); + if matches!( + *self.desired_state_rx.borrow(), + RemoteControlDesiredState::Unknown + ) && !self + .resolve_unknown_desired_state(app_server_client_name.as_deref()) + .await + { + self.client_tracker.lock().await.shutdown().await; + return; + } loop { if !self.wait_until_enabled().await { @@ -513,7 +541,7 @@ impl RemoteControlWebsocket { connection_end_reason = ?connection_end_reason, current_status = ?status.status, environment_id = ?status.environment_id, - enabled = *self.enabled_rx.borrow(), + desired_state = ?*self.desired_state_rx.borrow(), "app-server remote control websocket connection cycle ended" ); } @@ -546,10 +574,96 @@ impl RemoteControlWebsocket { } } + pub(super) async fn resolve_unknown_desired_state( + &mut self, + app_server_client_name: Option<&str>, + ) -> bool { + let remote_control_target = match super::protocol::normalize_remote_control_url( + &self.remote_control_url, + ) { + Ok(remote_control_target) => remote_control_target, + Err(err) => { + warn!( + "remote control preference cannot be resolved because the URL is invalid: {err}" + ); + self.transition_unknown_to(RemoteControlDesiredState::Disabled); + return true; + } + }; + self.remote_control_target = Some(remote_control_target.clone()); + let Some(state_db) = self.state_db.clone() else { + self.transition_unknown_to(RemoteControlDesiredState::Disabled); + return true; + }; + + loop { + if !matches!( + *self.desired_state_rx.borrow(), + RemoteControlDesiredState::Unknown + ) { + return true; + } + let auth = match load_remote_control_auth(&self.auth_manager).await { + Ok(auth) => auth, + Err(err) => { + info!( + error = %err, + "waiting to resolve remote control preference until authentication is available" + ); + if !self.wait_for_preference_resolution_retry().await { + return false; + } + continue; + } + }; + let enrollment = match state_db + .get_remote_control_enrollment( + &remote_control_target.websocket_url, + &auth.account_id, + app_server_client_name, + ) + .await + { + Ok(enrollment) => enrollment, + Err(err) => { + warn!( + error = %err, + "failed to resolve persisted remote control preference; retrying" + ); + if !self.wait_for_preference_resolution_retry().await { + return false; + } + continue; + } + }; + let desired_state = desired_state_from_persisted_enrollment(enrollment); + self.transition_unknown_to(desired_state); + return true; + } + } + + fn transition_unknown_to(&self, desired_state: RemoteControlDesiredState) { + self.desired_state_tx.send_if_modified(|state| { + if !matches!(*state, RemoteControlDesiredState::Unknown) { + return false; + } + *state = desired_state; + true + }); + } + + async fn wait_for_preference_resolution_retry(&mut self) -> bool { + tokio::select! { + _ = self.shutdown_token.cancelled() => false, + changed = self.desired_state_rx.changed() => changed.is_ok(), + _ = tokio::time::sleep(REMOTE_CONTROL_ACCOUNT_ID_RETRY_INTERVAL) => true, + } + } + async fn wait_until_enabled(&mut self) -> bool { tokio::select! { _ = self.shutdown_token.cancelled() => false, - enabled = self.enabled_rx.wait_for(|enabled| *enabled) => enabled.is_ok(), + desired_state = self.desired_state_rx.wait_for(|state| state.is_enabled()) => desired_state.is_ok(), } } @@ -573,7 +687,7 @@ impl RemoteControlWebsocket { warn!("remote control is enabled but the URL is invalid: {err}"); tokio::select! { _ = shutdown_token.cancelled() => return ConnectOutcome::Shutdown, - changed = self.enabled_rx.wait_for(|enabled| !*enabled) => { + changed = self.desired_state_rx.wait_for(|state| !state.is_enabled()) => { if changed.is_err() { return ConnectOutcome::Shutdown; } @@ -604,15 +718,18 @@ impl RemoteControlWebsocket { server_name: &self.server_name, subscribe_cursor: subscribe_cursor.as_deref(), app_server_client_name, + desired_state_tx: &self.desired_state_tx, + desired_state_persistence_lock: &self.desired_state_persistence_lock, }; let auth_context = RemoteControlAuthContext { auth_manager: &self.auth_manager, auth_recovery: &mut self.auth_recovery, auth_change_rx: &mut self.auth_change_rx, }; + let mut disabled_rx = self.desired_state_rx.clone(); let connect_result = tokio::select! { _ = shutdown_token.cancelled() => return ConnectOutcome::Shutdown, - changed = self.enabled_rx.wait_for(|enabled| !*enabled) => { + changed = disabled_rx.wait_for(|state| !state.is_enabled()) => { if changed.is_err() { return ConnectOutcome::Shutdown; } @@ -633,7 +750,7 @@ impl RemoteControlWebsocket { match connect_result { Ok((websocket_connection, response)) => { - if !*self.enabled_rx.borrow() { + if !self.desired_state_rx.borrow().is_enabled() { return ConnectOutcome::Disabled; } self.reconnect_attempt = 0; @@ -654,7 +771,7 @@ impl RemoteControlWebsocket { return ConnectOutcome::Connected(Box::new(websocket_connection)); } Err(err) => { - if !*self.enabled_rx.borrow() { + if !self.desired_state_rx.borrow().is_enabled() { return ConnectOutcome::Disabled; } let reconnect_delay = if err.kind() == ErrorKind::WouldBlock { @@ -691,7 +808,7 @@ impl RemoteControlWebsocket { }; tokio::select! { _ = shutdown_token.cancelled() => return ConnectOutcome::Shutdown, - changed = self.enabled_rx.wait_for(|enabled| !*enabled) => { + changed = self.desired_state_rx.wait_for(|state| !state.is_enabled()) => { if changed.is_err() { return ConnectOutcome::Shutdown; } @@ -736,10 +853,10 @@ impl RemoteControlWebsocket { shutdown_token.clone(), )); - let mut enabled_rx = self.enabled_rx.clone(); + let mut desired_state_rx = self.desired_state_rx.clone(); let connection_end_reason = tokio::select! { _ = shutdown_token.cancelled() => ConnectionEndReason::Shutdown, - changed = enabled_rx.wait_for(|enabled| !*enabled) => { + changed = desired_state_rx.wait_for(|state| !state.is_enabled()) => { if changed.is_ok() { self.status_publisher .publish_status(RemoteControlConnectionStatus::Disabled); @@ -1253,22 +1370,25 @@ pub(super) async fn connect_remote_control_websocket( if websocket_response_reports_missing_remote_app_server(response) => { info!( - "remote control websocket returned HTTP 404; clearing stale enrollment before re-enrolling: websocket_url={}, account_id={}, server_id={}, environment_id={}", + "remote control websocket returned HTTP 404; replacing stale enrollment: websocket_url={}, account_id={}, server_id={}, environment_id={}", remote_control_target.websocket_url, auth.account_id, enrollment.server_id, enrollment.environment_id ); - clear_remote_control_enrollment_if_matches( + replace_remote_control_enrollment_if_matches( state_db, remote_control_target, - &auth.account_id, - connect_options.app_server_client_name, + RemoteControlEnrollmentAuthContext { + auth: &auth, + recovery: &mut auth_context, + }, current_enrollment, &enrollment, + connect_options, status_publisher, ) - .await; + .await?; } tungstenite::Error::Http(response) if response.status().as_u16() == 404 => { let response_body = response @@ -1337,6 +1457,14 @@ async fn prepare_remote_control_enrollment( }; let enrollment_account_id = enrollment.as_ref().map(|enrollment| &enrollment.account_id); if enrollment_account_id.is_some_and(|account_id| account_id != &auth.account_id) { + resolve_desired_state_after_account_change( + state_db, + remote_control_target, + auth_context.auth_manager, + &auth.account_id, + connect_options, + ) + .await?; info!( "clearing in-memory remote control enrollment because account id changed: websocket_url={}, previous_account_id={:?}, current_account_id={:?}", remote_control_target.websocket_url, @@ -1347,6 +1475,12 @@ async fn prepare_remote_control_enrollment( ); *enrollment = None; status_publisher.publish_environment_id(/*environment_id*/ None); + if !connect_options.desired_state_tx.borrow().is_enabled() { + return Err(io::Error::new( + ErrorKind::Interrupted, + "remote control disabled after account changed", + )); + } } if let Some(enrollment) = enrollment.as_mut() { enrollment.remote_control_target = remote_control_target.clone(); @@ -1373,14 +1507,17 @@ async fn prepare_remote_control_enrollment( }); } - enroll_remote_control_server_if_missing( + enroll_and_persist_remote_control_server( remote_control_target, state_db, - &auth, - auth_context, + RemoteControlEnrollmentAuthContext { + auth: &auth, + recovery: auth_context, + }, enrollment, connect_options, status_publisher, + RemoteControlEnrollmentSelection::ReuseOrCreate, ) .await?; @@ -1412,26 +1549,20 @@ async fn prepare_remote_control_enrollment( Ok(()) => {} Err(err) if err.kind() == ErrorKind::NotFound => { info!( - "remote control server refresh returned HTTP 404; clearing stale enrollment before re-enrolling: websocket_url={}, account_id={}, server_id={}, environment_id={}", + "remote control server refresh returned HTTP 404; replacing stale enrollment: websocket_url={}, account_id={}, server_id={}, environment_id={}", remote_control_target.websocket_url, auth.account_id, server_id, environment_id ); - clear_remote_control_enrollment( - state_db, - remote_control_target, - &auth.account_id, - connect_options.app_server_client_name, - enrollment, - status_publisher, - ) - .await; - enroll_remote_control_server_if_missing( + enroll_and_persist_remote_control_server( remote_control_target, state_db, - &auth, - auth_context, + RemoteControlEnrollmentAuthContext { + auth: &auth, + recovery: auth_context, + }, enrollment, connect_options, status_publisher, + RemoteControlEnrollmentSelection::ReplaceExisting, ) .await?; } @@ -1454,6 +1585,51 @@ async fn prepare_remote_control_enrollment( Ok(auth) } +async fn resolve_desired_state_after_account_change( + state_db: &StateRuntime, + remote_control_target: &RemoteControlTarget, + auth_manager: &Arc, + account_id: &str, + connect_options: RemoteControlConnectOptions<'_>, +) -> io::Result<()> { + let durable_enabled = RemoteControlDesiredState::Enabled { + persistence_preference: Some(true), + }; + if *connect_options.desired_state_tx.borrow() != durable_enabled { + return Ok(()); + } + + let _persistence = + acquire_persistence_lock(connect_options.desired_state_persistence_lock).await; + if *connect_options.desired_state_tx.borrow() != durable_enabled { + return Ok(()); + } + let enrollment = state_db + .get_remote_control_enrollment( + &remote_control_target.websocket_url, + account_id, + connect_options.app_server_client_name, + ) + .await + .map_err(io::Error::other)?; + let current_auth = load_remote_control_auth(auth_manager).await?; + if current_auth.account_id != account_id { + return Err(io::Error::new( + ErrorKind::WouldBlock, + "remote control account changed while resolving persisted preference", + )); + } + let resolved_state = desired_state_from_persisted_enrollment(enrollment); + connect_options.desired_state_tx.send_if_modified(|state| { + if *state != durable_enabled || *state == resolved_state { + return false; + } + *state = resolved_state; + true + }); + Ok(()) +} + fn websocket_response_reports_missing_remote_app_server( response: &tungstenite::http::Response>>, ) -> bool { @@ -1466,57 +1642,38 @@ fn websocket_response_reports_missing_remote_app_server( }) } -async fn clear_remote_control_enrollment( - state_db: &StateRuntime, - remote_control_target: &RemoteControlTarget, - account_id: &str, - app_server_client_name: Option<&str>, - enrollment: &mut Option, - status_publisher: &RemoteControlStatusPublisher, -) { - if let Err(clear_err) = update_persisted_remote_control_enrollment( - Some(state_db), - remote_control_target, - account_id, - app_server_client_name, - /*enrollment*/ None, - ) - .await - { - warn!("failed to clear stale remote control enrollment in sqlite state db: {clear_err}"); - } - *enrollment = None; - status_publisher.publish_environment_id(/*environment_id*/ None); -} - -async fn clear_remote_control_enrollment_if_matches( +async fn replace_remote_control_enrollment_if_matches( state_db: Option<&StateRuntime>, remote_control_target: &RemoteControlTarget, - account_id: &str, - app_server_client_name: Option<&str>, + auth_context: RemoteControlEnrollmentAuthContext<'_, '_>, current_enrollment: &CurrentRemoteControlEnrollment, enrollment: &RemoteControlEnrollment, + connect_options: RemoteControlConnectOptions<'_>, status_publisher: &RemoteControlStatusPublisher, -) { +) -> io::Result<()> { let Some(state_db) = state_db else { - return; + return Err(io::Error::new( + ErrorKind::NotFound, + "remote control requires sqlite state db", + )); }; let mut current_enrollment = current_enrollment.lock().await; if !current_enrollment .as_ref() .is_some_and(|current| same_remote_control_enrollment(current, enrollment)) { - return; + return Ok(()); } - clear_remote_control_enrollment( - state_db, + enroll_and_persist_remote_control_server( remote_control_target, - account_id, - app_server_client_name, + state_db, + auth_context, &mut current_enrollment, + connect_options, status_publisher, + RemoteControlEnrollmentSelection::ReplaceExisting, ) - .await; + .await } async fn clear_remote_control_server_token_if_matches( @@ -1534,26 +1691,39 @@ async fn clear_remote_control_server_token_if_matches( Ok(()) } -async fn enroll_remote_control_server_if_missing( +async fn enroll_and_persist_remote_control_server( remote_control_target: &RemoteControlTarget, state_db: &StateRuntime, - auth: &RemoteControlConnectionAuth, - auth_context: &mut RemoteControlAuthContext<'_>, + auth_context: RemoteControlEnrollmentAuthContext<'_, '_>, enrollment: &mut Option, connect_options: RemoteControlConnectOptions<'_>, status_publisher: &RemoteControlStatusPublisher, + selection: RemoteControlEnrollmentSelection, ) -> io::Result<()> { - if enrollment.is_some() { - return Ok(()); + match selection { + RemoteControlEnrollmentSelection::ReuseOrCreate => { + if enrollment.is_some() { + return Ok(()); + } + } + RemoteControlEnrollmentSelection::ReplaceExisting => {} + } + if !connect_options.desired_state_tx.borrow().is_enabled() { + return Err(io::Error::new( + ErrorKind::Interrupted, + "remote control disabled before enrollment", + )); } info!( "creating new remote control enrollment: websocket_url={}, enroll_url={}, account_id={}", - remote_control_target.websocket_url, remote_control_target.enroll_url, auth.account_id + remote_control_target.websocket_url, + remote_control_target.enroll_url, + auth_context.auth.account_id ); let new_enrollment = match enroll_remote_control_server( remote_control_target, - auth, + auth_context.auth, connect_options.installation_id, connect_options.server_name, ) @@ -1563,8 +1733,8 @@ async fn enroll_remote_control_server_if_missing( Err(err) if err.kind() == ErrorKind::PermissionDenied && recover_remote_control_auth( - auth_context.auth_recovery, - auth_context.auth_change_rx, + auth_context.recovery.auth_recovery, + auth_context.recovery.auth_change_rx, ) .await => { @@ -1574,12 +1744,26 @@ async fn enroll_remote_control_server_if_missing( } Err(err) => return Err(err), }; + let _persistence = + acquire_persistence_lock(connect_options.desired_state_persistence_lock).await; + let persistence_preference = match *connect_options.desired_state_tx.borrow() { + RemoteControlDesiredState::Enabled { + persistence_preference, + } => persistence_preference, + RemoteControlDesiredState::Unknown | RemoteControlDesiredState::Disabled => { + return Err(io::Error::new( + ErrorKind::Interrupted, + "remote control disabled during enrollment", + )); + } + }; if let Err(err) = update_persisted_remote_control_enrollment( Some(state_db), remote_control_target, - &auth.account_id, + &auth_context.auth.account_id, connect_options.app_server_client_name, Some(&new_enrollment), + persistence_preference, ) .await { @@ -1758,6 +1942,13 @@ mod tests { (RemoteControlStatusPublisher::new(status_tx), status_rx) } + fn enabled_desired_state_sender() -> watch::Sender { + watch::channel(RemoteControlDesiredState::Enabled { + persistence_preference: None, + }) + .0 + } + #[test] fn mark_recovery_auth_change_seen_marks_only_recovery_revision_seen() { let (auth_change_tx, mut auth_change_rx) = watch::channel(0u64); @@ -1896,6 +2087,8 @@ mod tests { server_name: "test-server", subscribe_cursor: None, app_server_client_name: None, + desired_state_tx: &enabled_desired_state_sender(), + desired_state_persistence_lock: &Semaphore::new(1), }, &status_publisher, ) @@ -1960,6 +2153,8 @@ mod tests { server_name: "test-server", subscribe_cursor: None, app_server_client_name: None, + desired_state_tx: &enabled_desired_state_sender(), + desired_state_persistence_lock: &Semaphore::new(1), }, &status_publisher, ) @@ -2042,6 +2237,8 @@ mod tests { server_name: "test-server", subscribe_cursor: None, app_server_client_name: None, + desired_state_tx: &enabled_desired_state_sender(), + desired_state_persistence_lock: &Semaphore::new(1), }, &status_publisher, ) @@ -2136,6 +2333,8 @@ mod tests { server_name: "test-server", subscribe_cursor: None, app_server_client_name: None, + desired_state_tx: &enabled_desired_state_sender(), + desired_state_persistence_lock: &Semaphore::new(1), }, &status_publisher, ) @@ -2201,6 +2400,8 @@ mod tests { server_name: "test-server", subscribe_cursor: None, app_server_client_name: None, + desired_state_tx: &enabled_desired_state_sender(), + desired_state_persistence_lock: &Semaphore::new(1), }, &status_publisher, ) @@ -2251,6 +2452,8 @@ mod tests { server_name: "test-server", subscribe_cursor: None, app_server_client_name: None, + desired_state_tx: &enabled_desired_state_sender(), + desired_state_persistence_lock: &Semaphore::new(1), }, &status_publisher, ) @@ -2288,7 +2491,10 @@ mod tests { drop(transport_event_rx); let (status_publisher, _status_rx) = remote_control_status_channel(); let shutdown_token = CancellationToken::new(); - let (_enabled_tx, enabled_rx) = watch::channel(true); + let (desired_state_tx, _desired_state_rx) = + watch::channel(RemoteControlDesiredState::Enabled { + persistence_preference: None, + }); let websocket_task = tokio::spawn({ let shutdown_token = shutdown_token.clone(); async move { @@ -2306,9 +2512,10 @@ mod tests { status_publisher, current_enrollment: test_current_enrollment(/*enrollment*/ None), pairing_persistence_key: watch::channel(None).0, + desired_state_persistence_lock: Arc::new(Semaphore::new(1)), }, shutdown_token, - enabled_rx, + Arc::new(desired_state_tx), ) .run(/*app_server_client_name_rx*/ None) .await diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 42ac510a8..cc36bb451 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -210,8 +210,8 @@ Example with notification opt-out: - `plugin/skill/read` — read remote plugin skill markdown on demand by `remoteMarketplaceName`, `remotePluginId`, and `skillName`. This lets clients preview uninstalled remote plugin skills without downloading the plugin bundle. - `skills/changed` — notification emitted when watched local skill files change. - `app/list` — list available apps. -- `remoteControl/enable` — experimental; enable remote control for the current app-server process and return the current remote-control status snapshot. The caller is responsible for persisting the desired setting outside app-server. -- `remoteControl/disable` — experimental; disable remote control for the current app-server process and return the current remote-control status snapshot. This does not revoke already enrolled controller devices. +- `remoteControl/enable` — experimental; enable remote control for the current app-server process and return the current remote-control status snapshot. By default, any missing enrollment is completed before the response and the preference is persisted for the current app-server client scope. Pass `ephemeral: true` to enable remote control only for the current process without changing the persisted preference. +- `remoteControl/disable` — experimental; disable remote control for the current app-server process and return the current remote-control status snapshot. By default, the disabled preference is persisted for the current app-server client scope. Pass `ephemeral: true` to disable only for the current process without changing the persisted preference. This does not revoke already enrolled controller devices. - `remoteControl/status/read` — experimental; read the current remote-control status snapshot. `status` is one of `disabled`, `connecting`, `connected`, or `errored`; `serverName` is the local machine name used by this app-server process; `environmentId` is a string when the app-server has a current enrollment and `null` when that enrollment is cleared, invalidated, or remote control is disabled. - `remoteControl/pairing/start` — experimental; start a short-lived remote-control pairing artifact for the current app-server process. Pass `manualCode: true` to also request a manual pairing code. Returns `pairingCode`, `manualPairingCode`, `environmentId`, and Unix-seconds `expiresAt`; app-server intentionally does not expose the backend `serverId`. - `remoteControl/pairing/status` — experimental; poll whether a remote-control `pairingCode` or `manualPairingCode` has been claimed. Pass exactly one of the two fields. Returns `claimed`. diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index e2d02cde4..b5ca8938c 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -110,10 +110,12 @@ mod transport; pub use crate::error_code::INPUT_TOO_LARGE_ERROR_CODE; pub use crate::error_code::INVALID_PARAMS_ERROR_CODE; pub use crate::transport::AppServerTransport; +pub use crate::transport::RemoteControlStartupMode; pub use crate::transport::app_server_control_socket_path; pub use crate::transport::auth::AppServerWebsocketAuthArgs; pub use crate::transport::auth::AppServerWebsocketAuthSettings; pub use crate::transport::auth::WebsocketAuthCliMode; +pub use crate::transport::take_remote_control_disabled_env; const LOG_FORMAT_ENV_VAR: &str = "LOG_FORMAT"; const OTEL_SERVICE_NAME: &str = "codex-app-server"; @@ -408,7 +410,7 @@ pub enum PluginStartupTasks { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct AppServerRuntimeOptions { pub plugin_startup_tasks: PluginStartupTasks, - pub remote_control_enabled: bool, + pub remote_control_startup_mode: RemoteControlStartupMode, pub install_shutdown_signal_handler: bool, } @@ -416,7 +418,7 @@ impl Default for AppServerRuntimeOptions { fn default() -> Self { Self { plugin_startup_tasks: PluginStartupTasks::Start, - remote_control_enabled: false, + remote_control_startup_mode: RemoteControlStartupMode::ResolvePersisted, install_shutdown_signal_handler: true, } } @@ -717,15 +719,21 @@ pub async fn run_main_with_transport_options( let auth_manager = AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false).await; - let remote_control_requested = runtime_options.remote_control_enabled; - let remote_control_enabled = remote_control_requested && state_db.is_some(); - if remote_control_requested && state_db.is_none() { + let remote_control_startup_mode = runtime_options.remote_control_startup_mode; + let remote_control_explicitly_requested = + remote_control_startup_mode == RemoteControlStartupMode::EnabledEphemeral; + let remote_control_enabled = remote_control_explicitly_requested && state_db.is_some(); + if remote_control_explicitly_requested && state_db.is_none() { error!("remote control disabled because sqlite state db is unavailable"); } - if transport_accept_handles.is_empty() && !remote_control_enabled { + let no_local_transport = transport_accept_handles.is_empty(); + if no_local_transport + && remote_control_startup_mode != RemoteControlStartupMode::ResolvePersisted + && !remote_control_enabled + { return Err(std::io::Error::new( ErrorKind::InvalidInput, - if remote_control_requested && state_db.is_none() { + if remote_control_explicitly_requested && state_db.is_none() { "no transport configured; remote control disabled because sqlite state db is unavailable" } else { "no transport configured; use --listen or enable remote control" @@ -743,9 +751,31 @@ pub async fn run_main_with_transport_options( transport_event_tx.clone(), transport_shutdown_token.clone(), app_server_client_name_rx, - remote_control_enabled, + remote_control_startup_mode, ) .await?; + if no_local_transport + && remote_control_startup_mode == RemoteControlStartupMode::ResolvePersisted + { + let persisted_enabled = match remote_control_handle + .resolve_persisted_preference(/*app_server_client_name*/ None) + .await + { + Ok(persisted_enabled) => persisted_enabled, + Err(err) => { + warn!("failed to resolve persisted remote control preference: {err}"); + false + } + }; + if !persisted_enabled { + transport_shutdown_token.cancel(); + let _ = remote_control_accept_handle.await; + return Err(std::io::Error::new( + ErrorKind::InvalidInput, + "no transport configured; use --listen or enable remote control", + )); + } + } transport_accept_handles.push(remote_control_accept_handle); let outbound_handle = tokio::spawn(async move { diff --git a/codex-rs/app-server/src/main.rs b/codex-rs/app-server/src/main.rs index 2412f8c3c..7599e498d 100644 --- a/codex-rs/app-server/src/main.rs +++ b/codex-rs/app-server/src/main.rs @@ -53,13 +53,14 @@ struct AppServerArgs { #[arg(long = "disable-plugin-startup-tasks-for-tests", hide = true)] disable_plugin_startup_tasks_for_tests: bool, - /// Enable remote control for this app-server process. + /// Enable remote control for this app-server process without changing persistence. #[arg(long = "remote-control", hide = true)] remote_control: bool, } fn main() -> anyhow::Result<()> { - arg0_dispatch_or_else(|arg0_paths: Arg0DispatchPaths| async move { + let remote_control_disabled = codex_app_server::take_remote_control_disabled_env(); + arg0_dispatch_or_else(move |arg0_paths: Arg0DispatchPaths| async move { let AppServerArgs { config_overrides, listen, @@ -84,7 +85,12 @@ fn main() -> anyhow::Result<()> { if disable_plugin_startup_tasks_for_tests { runtime_options.plugin_startup_tasks = PluginStartupTasks::Skip; } - runtime_options.remote_control_enabled = remote_control; + runtime_options.remote_control_startup_mode = + match (remote_control, remote_control_disabled) { + (true, _) => codex_app_server::RemoteControlStartupMode::EnabledEphemeral, + (false, true) => codex_app_server::RemoteControlStartupMode::DisabledEphemeral, + (false, false) => codex_app_server::RemoteControlStartupMode::ResolvePersisted, + }; run_main_with_transport_options( arg0_paths, diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index d55b266dc..6747830b8 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -951,13 +951,21 @@ impl MessageProcessor { .experimental_feature_enablement_set(request_id.clone(), params) .await } - ClientRequest::RemoteControlEnable { .. } => self + ClientRequest::RemoteControlEnable { params, .. } => self .remote_control_processor - .enable() + .enable( + params.is_some_and(|params| params.ephemeral), + app_server_client_name.as_deref(), + ) + .await .map(|response| Some(response.into())), - ClientRequest::RemoteControlDisable { .. } => self + ClientRequest::RemoteControlDisable { params, .. } => self .remote_control_processor - .disable() + .disable( + params.is_some_and(|params| params.ephemeral), + app_server_client_name.as_deref(), + ) + .await .map(|response| Some(response.into())), ClientRequest::RemoteControlStatusRead { .. } => self .remote_control_processor diff --git a/codex-rs/app-server/src/request_processors/remote_control_processor.rs b/codex-rs/app-server/src/request_processors/remote_control_processor.rs index 31b5fd6b7..fb94bb6b2 100644 --- a/codex-rs/app-server/src/request_processors/remote_control_processor.rs +++ b/codex-rs/app-server/src/request_processors/remote_control_processor.rs @@ -28,17 +28,38 @@ impl RemoteControlRequestProcessor { } } - pub(crate) fn enable(&self) -> Result { + pub(crate) async fn enable( + &self, + ephemeral: bool, + app_server_client_name: Option<&str>, + ) -> Result { let handle = self.handle()?; - handle - .enable() - .map(RemoteControlEnableResponse::from) - .map_err(map_unavailable) + let status = if ephemeral { + handle.enable_ephemeral().map_err(map_unavailable)? + } else { + handle + .enable(app_server_client_name) + .await + .map_err(map_update_error)? + }; + Ok(RemoteControlEnableResponse::from(status)) } - pub(crate) fn disable(&self) -> Result { + pub(crate) async fn disable( + &self, + ephemeral: bool, + app_server_client_name: Option<&str>, + ) -> Result { let handle = self.handle()?; - Ok(RemoteControlDisableResponse::from(handle.disable())) + let status = if ephemeral { + handle.disable_ephemeral().await + } else { + handle + .disable(app_server_client_name) + .await + .map_err(map_update_error)? + }; + Ok(RemoteControlDisableResponse::from(status)) } pub(crate) fn status_read(&self) -> Result { @@ -104,6 +125,14 @@ fn map_unavailable(err: RemoteControlUnavailable) -> JSONRPCErrorError { invalid_request(err.to_string()) } +fn map_update_error(err: io::Error) -> JSONRPCErrorError { + if err.kind() == io::ErrorKind::NotFound { + invalid_request(err.to_string()) + } else { + internal_error(err.to_string()) + } +} + fn map_pairing_start_error(err: io::Error) -> JSONRPCErrorError { if err.kind() == io::ErrorKind::InvalidInput { invalid_request(err.to_string()) diff --git a/codex-rs/app-server/src/transport.rs b/codex-rs/app-server/src/transport.rs index 2be5b5541..94c1ec05f 100644 --- a/codex-rs/app-server/src/transport.rs +++ b/codex-rs/app-server/src/transport.rs @@ -20,6 +20,7 @@ pub(crate) use codex_app_server_transport::OutgoingMessage; pub(crate) use codex_app_server_transport::QueuedOutgoingMessage; pub(crate) use codex_app_server_transport::RemoteControlHandle; pub(crate) use codex_app_server_transport::RemoteControlStartConfig; +pub use codex_app_server_transport::RemoteControlStartupMode; pub(crate) use codex_app_server_transport::RemoteControlUnavailable; pub(crate) use codex_app_server_transport::TransportEvent; pub(crate) use codex_app_server_transport::acquire_app_server_startup_lock; @@ -31,6 +32,7 @@ pub(crate) use codex_app_server_transport::start_control_socket_acceptor; pub(crate) use codex_app_server_transport::start_remote_control; pub(crate) use codex_app_server_transport::start_stdio_connection; pub(crate) use codex_app_server_transport::start_websocket_acceptor; +pub use codex_app_server_transport::take_remote_control_disabled_env; pub(crate) struct ConnectionState { pub(crate) outbound_initialized: Arc, diff --git a/codex-rs/app-server/tests/common/test_app_server.rs b/codex-rs/app-server/tests/common/test_app_server.rs index 91c42e2fe..b2ecb6473 100644 --- a/codex-rs/app-server/tests/common/test_app_server.rs +++ b/codex-rs/app-server/tests/common/test_app_server.rs @@ -1,5 +1,6 @@ use std::collections::VecDeque; use std::path::Path; +use std::process::ExitStatus; use std::process::Stdio; use std::sync::atomic::AtomicI64; use std::sync::atomic::Ordering; @@ -128,6 +129,10 @@ pub const DISABLE_PLUGIN_STARTUP_TASKS_ARG: &str = "--disable-plugin-startup-tas const DISABLE_MANAGED_CONFIG_ENV_VAR: &str = "CODEX_APP_SERVER_DISABLE_MANAGED_CONFIG"; impl TestAppServer { + pub async fn wait_for_exit(&mut self) -> std::io::Result { + self.process.wait().await + } + pub async fn new(codex_home: &Path) -> anyhow::Result { Self::new_with_env_and_args(codex_home, &[], &[DISABLE_PLUGIN_STARTUP_TASKS_ARG]).await } @@ -645,12 +650,30 @@ impl TestAppServer { .await } + /// Send a runtime-only `remoteControl/enable` JSON-RPC request. + pub async fn send_remote_control_ephemeral_enable_request(&mut self) -> anyhow::Result { + self.send_request( + "remoteControl/enable", + Some(serde_json::json!({ "ephemeral": true })), + ) + .await + } + /// Send a `remoteControl/disable` JSON-RPC request. pub async fn send_remote_control_disable_request(&mut self) -> anyhow::Result { self.send_request("remoteControl/disable", /*params*/ None) .await } + /// Send a runtime-only `remoteControl/disable` JSON-RPC request. + pub async fn send_remote_control_ephemeral_disable_request(&mut self) -> anyhow::Result { + self.send_request( + "remoteControl/disable", + Some(serde_json::json!({ "ephemeral": true })), + ) + .await + } + /// Send a `remoteControl/status/read` JSON-RPC request. pub async fn send_remote_control_status_read_request(&mut self) -> anyhow::Result { self.send_request("remoteControl/status/read", /*params*/ None) diff --git a/codex-rs/app-server/tests/suite/v2/remote_control.rs b/codex-rs/app-server/tests/suite/v2/remote_control.rs index 0229b0e5a..920e9bfa6 100644 --- a/codex-rs/app-server/tests/suite/v2/remote_control.rs +++ b/codex-rs/app-server/tests/suite/v2/remote_control.rs @@ -3,6 +3,7 @@ use std::time::Duration; use anyhow::Context; use anyhow::Result; use app_test_support::ChatGptAuthFixture; +use app_test_support::DEFAULT_CLIENT_NAME; use app_test_support::TestAppServer; use app_test_support::to_response; use app_test_support::write_chatgpt_auth; @@ -24,6 +25,8 @@ use codex_app_server_protocol::RemoteControlPairingStatusResponse; use codex_app_server_protocol::RemoteControlStatusReadResponse; use codex_app_server_protocol::RequestId; use codex_config::types::AuthCredentialsStoreMode; +use codex_state::RemoteControlEnrollmentRecord; +use codex_state::StateRuntime; use pretty_assertions::assert_eq; use tempfile::TempDir; use tokio::io::AsyncBufReadExt; @@ -37,10 +40,92 @@ use tokio::task::JoinHandle; use tokio::time::timeout; const DEFAULT_TIMEOUT: Duration = Duration::from_secs(10); +const STARTUP_TIMEOUT: Duration = Duration::from_secs(30); + +async fn remote_control_preference( + state_db: &StateRuntime, + websocket_url: &str, +) -> Result> { + Ok(state_db + .get_remote_control_enrollment(websocket_url, "account_id", Some(DEFAULT_CLIENT_NAME)) + .await? + .context("enrollment should exist")? + .remote_control_enabled) +} + +async fn wait_for_response(mcp: &mut TestAppServer, request_id: i64) -> Result { + timeout( + DEFAULT_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await? +} + +#[tokio::test] +async fn listen_off_honors_persisted_remote_control_enable() -> Result<()> { + let codex_home = TempDir::new()?; + let listener = configured_remote_control_listener(codex_home.path()).await?; + let websocket_url = format!( + "ws://{}/backend-api/wham/remote/control/server", + listener.local_addr()? + ); + let state_db = + StateRuntime::init(codex_home.path().to_path_buf(), "test-provider".to_string()).await?; + state_db + .upsert_remote_control_enrollment(&RemoteControlEnrollmentRecord { + websocket_url, + account_id: "account_id".to_string(), + app_server_client_name: None, + server_id: "server-id".to_string(), + environment_id: "environment-id".to_string(), + server_name: "server-name".to_string(), + remote_control_enabled: Some(true), + }) + .await?; + + let _app_server = TestAppServer::new_with_args(codex_home.path(), &["--listen", "off"]).await?; + timeout(STARTUP_TIMEOUT, listener.accept()).await??; + Ok(()) +} + +#[tokio::test] +async fn listen_off_exits_without_persisted_remote_control_enable() -> Result<()> { + for persisted_preference in [None, Some(false)] { + let codex_home = TempDir::new()?; + let listener = configured_remote_control_listener(codex_home.path()).await?; + if let Some(remote_control_enabled) = persisted_preference { + let websocket_url = format!( + "ws://{}/backend-api/wham/remote/control/server", + listener.local_addr()? + ); + let state_db = + StateRuntime::init(codex_home.path().to_path_buf(), "test-provider".to_string()) + .await?; + state_db + .upsert_remote_control_enrollment(&RemoteControlEnrollmentRecord { + websocket_url, + account_id: "account_id".to_string(), + app_server_client_name: None, + server_id: "server-id".to_string(), + environment_id: "environment-id".to_string(), + server_name: "server-name".to_string(), + remote_control_enabled: Some(remote_control_enabled), + }) + .await?; + } + + let mut app_server = + TestAppServer::new_with_args(codex_home.path(), &["--listen", "off"]).await?; + let status = timeout(STARTUP_TIMEOUT, app_server.wait_for_exit()).await??; + assert!(!status.success()); + } + Ok(()) +} #[tokio::test] async fn remote_control_disable_returns_disabled_status() -> Result<()> { let codex_home = TempDir::new()?; + let _listener = configured_remote_control_listener(codex_home.path()).await?; let mut mcp = TestAppServer::new(codex_home.path()).await?; timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; @@ -83,11 +168,22 @@ async fn remote_control_status_read_returns_disabled_status() -> Result<()> { #[tokio::test] async fn remote_control_enable_returns_connecting_status() -> Result<()> { let codex_home = TempDir::new()?; - let _backend = BlockingRemoteControlBackend::start(codex_home.path()).await?; + let mut backend = BlockingRemoteControlBackend::start(codex_home.path()).await?; let mut mcp = TestAppServer::new(codex_home.path()).await?; timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_remote_control_enable_request().await?; + assert_eq!( + timeout(DEFAULT_TIMEOUT, backend.wait_for_enroll_request()).await??, + "POST /backend-api/wham/remote/control/server/enroll HTTP/1.1" + ); + timeout( + Duration::from_millis(100), + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await + .expect_err("enable response should wait for enrollment"); + backend.complete_enrollment()?; let response: JSONRPCResponse = timeout( DEFAULT_TIMEOUT, mcp.read_stream_until_response_message(RequestId::Integer(request_id)), @@ -97,11 +193,103 @@ async fn remote_control_enable_returns_connecting_status() -> Result<()> { assert_eq!(received.status, RemoteControlConnectionStatus::Connecting); assert!(!received.server_name.is_empty()); - assert_eq!(received.environment_id, None); + assert_eq!(received.environment_id.as_deref(), Some("environment-id")); assert!(!received.installation_id.is_empty()); Ok(()) } +#[tokio::test] +async fn disable_waits_for_in_flight_durable_enable() -> Result<()> { + let codex_home = TempDir::new()?; + let mut backend = BlockingRemoteControlBackend::start(codex_home.path()).await?; + let websocket_url = backend.websocket_url().to_string(); + let state_db = + StateRuntime::init(codex_home.path().to_path_buf(), "test-provider".to_string()).await?; + let mut mcp = TestAppServer::new(codex_home.path()).await?; + timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; + + mcp.send_remote_control_enable_request().await?; + timeout(DEFAULT_TIMEOUT, backend.wait_for_enroll_request()).await??; + let disable_request_id = mcp.send_remote_control_disable_request().await?; + timeout( + Duration::from_millis(100), + mcp.read_stream_until_response_message(RequestId::Integer(disable_request_id)), + ) + .await + .expect_err("disable response should wait for the in-flight enable"); + + backend.complete_enrollment()?; + let response = wait_for_response(&mut mcp, disable_request_id).await?; + let received: RemoteControlDisableResponse = to_response(response)?; + assert_eq!(received.status, RemoteControlConnectionStatus::Disabled); + assert_eq!( + remote_control_preference(&state_db, &websocket_url).await?, + Some(false) + ); + Ok(()) +} + +#[tokio::test] +async fn rpc_updates_durable_preference_but_ephemeral_does_not() -> Result<()> { + let codex_home = TempDir::new()?; + let mut backend = BlockingRemoteControlBackend::start(codex_home.path()).await?; + let websocket_url = backend.websocket_url().to_string(); + let state_db = + StateRuntime::init(codex_home.path().to_path_buf(), "test-provider".to_string()).await?; + + let mut mcp = TestAppServer::new(codex_home.path()).await?; + timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; + + let request_id = mcp.send_remote_control_enable_request().await?; + assert_eq!( + timeout(DEFAULT_TIMEOUT, backend.wait_for_enroll_request()).await??, + "POST /backend-api/wham/remote/control/server/enroll HTTP/1.1" + ); + backend.complete_enrollment()?; + wait_for_response(&mut mcp, request_id).await?; + assert_eq!( + remote_control_preference(&state_db, &websocket_url).await?, + Some(true) + ); + + let request_id = mcp.send_remote_control_ephemeral_disable_request().await?; + wait_for_response(&mut mcp, request_id).await?; + assert_eq!( + remote_control_preference(&state_db, &websocket_url).await?, + Some(true) + ); + + let request_id = mcp.send_remote_control_disable_request().await?; + wait_for_response(&mut mcp, request_id).await?; + assert_eq!( + remote_control_preference(&state_db, &websocket_url).await?, + Some(false) + ); + + let request_id = mcp.send_remote_control_enable_request().await?; + wait_for_response(&mut mcp, request_id).await?; + assert_eq!( + remote_control_preference(&state_db, &websocket_url).await?, + Some(true) + ); + + let request_id = mcp.send_remote_control_disable_request().await?; + wait_for_response(&mut mcp, request_id).await?; + assert_eq!( + remote_control_preference(&state_db, &websocket_url).await?, + Some(false) + ); + + let request_id = mcp.send_remote_control_ephemeral_enable_request().await?; + wait_for_response(&mut mcp, request_id).await?; + assert_eq!( + remote_control_preference(&state_db, &websocket_url).await?, + Some(false) + ); + + Ok(()) +} + #[tokio::test] async fn remote_control_status_read_returns_connecting_status_after_enable() -> Result<()> { let codex_home = TempDir::new()?; @@ -110,17 +298,17 @@ async fn remote_control_status_read_returns_connecting_status_after_enable() -> timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; let request_id = mcp.send_remote_control_enable_request().await?; - let _: JSONRPCResponse = timeout( - DEFAULT_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(request_id)), - ) - .await??; - let enroll_request = timeout(DEFAULT_TIMEOUT, backend.wait_for_enroll_request()).await??; assert_eq!( enroll_request, "POST /backend-api/wham/remote/control/server/enroll HTTP/1.1" ); + backend.complete_enrollment()?; + let _: JSONRPCResponse = timeout( + DEFAULT_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; let request_id = mcp.send_remote_control_status_read_request().await?; let response: JSONRPCResponse = timeout( @@ -132,7 +320,7 @@ async fn remote_control_status_read_returns_connecting_status_after_enable() -> assert_eq!(received.status, RemoteControlConnectionStatus::Connecting); assert!(!received.server_name.is_empty()); - assert_eq!(received.environment_id, None); + assert_eq!(received.environment_id.as_deref(), Some("environment-id")); assert!(!received.installation_id.is_empty()); Ok(()) } @@ -235,11 +423,13 @@ async fn remote_control_pairing_start_returns_pairing_artifacts() -> Result<()> } #[tokio::test] -async fn remote_control_pairing_start_returns_pairing_artifacts_while_disabled() -> Result<()> { +async fn pairing_start_works_after_ephemeral_enable() -> Result<()> { let codex_home = TempDir::new()?; let mut backend = PairingRemoteControlBackend::start(codex_home.path()).await?; let mut mcp = TestAppServer::new(codex_home.path()).await?; timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; + let request_id = mcp.send_remote_control_ephemeral_enable_request().await?; + wait_for_response(&mut mcp, request_id).await?; let request_id = mcp .send_remote_control_pairing_start_request(RemoteControlPairingStartParams { @@ -333,6 +523,8 @@ async fn remote_control_client_management_works_while_disabled() -> Result<()> { struct BlockingRemoteControlBackend { enroll_request_rx: Option>>, + enroll_response_tx: Option>, + websocket_url: String, server_task: JoinHandle<()>, } @@ -397,12 +589,37 @@ impl ClientManagementRemoteControlBackend { impl BlockingRemoteControlBackend { async fn start(codex_home: &std::path::Path) -> Result { let listener = configured_remote_control_listener(codex_home).await?; + let websocket_url = format!( + "ws://{}/backend-api/wham/remote/control/server", + listener.local_addr()? + ); let (enroll_request_tx, enroll_request_rx) = oneshot::channel(); + let (enroll_response_tx, enroll_response_rx) = oneshot::channel(); let server_task = tokio::spawn(async move { - match read_enroll_request(listener).await { - Ok((request_line, _reader)) => { + match read_enroll_request(&listener).await { + Ok((request_line, reader)) => { let _ = enroll_request_tx.send(Ok(request_line)); + if enroll_response_rx.await.is_err() { + return; + } + if respond_with_json( + reader.into_inner(), + serde_json::json!({ + "server_id": "server-id", + "environment_id": "environment-id", + "remote_control_token": "remote-control-token", + "expires_at": "3026-05-22T12:34:56Z", + }), + ) + .await + .is_err() + { + return; + } + let Ok(_websocket) = listener.accept().await else { + return; + }; std::future::pending::<()>().await; } Err(err) => { @@ -413,6 +630,8 @@ impl BlockingRemoteControlBackend { Ok(Self { enroll_request_rx: Some(enroll_request_rx), + enroll_response_tx: Some(enroll_response_tx), + websocket_url, server_task, }) } @@ -424,6 +643,18 @@ impl BlockingRemoteControlBackend { .context("enroll request should only be awaited once")?; rx.await? } + + fn complete_enrollment(&mut self) -> Result<()> { + self.enroll_response_tx + .take() + .context("enrollment should only complete once")? + .send(()) + .map_err(|()| anyhow::anyhow!("enrollment response receiver dropped")) + } + + fn websocket_url(&self) -> &str { + &self.websocket_url + } } struct PairingRemoteControlBackend { @@ -558,8 +789,8 @@ async fn configured_remote_control_listener(codex_home: &std::path::Path) -> Res Ok(listener) } -async fn read_enroll_request(listener: TcpListener) -> Result<(String, BufReader)> { - let request = read_http_request(&listener).await?; +async fn read_enroll_request(listener: &TcpListener) -> Result<(String, BufReader)> { + let request = read_http_request(listener).await?; Ok((request.request_line, request.reader)) } diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 8b2c02ad1..1d06ee787 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -533,7 +533,7 @@ struct AppServerCommand { #[arg(long = "stdio", conflicts_with = "listen")] stdio: bool, - /// Enable remote control for this app-server process. + /// Enable remote control for this app-server process without changing persistence. #[arg(long = "remote-control", hide = true)] remote_control: bool, @@ -953,13 +953,17 @@ fn stage_str(stage: Stage) -> &'static str { } fn main() -> anyhow::Result<()> { - arg0_dispatch_or_else(|arg0_paths: Arg0DispatchPaths| async move { - cli_main(arg0_paths).await?; + let remote_control_disabled = codex_app_server::take_remote_control_disabled_env(); + arg0_dispatch_or_else(move |arg0_paths: Arg0DispatchPaths| async move { + cli_main(arg0_paths, remote_control_disabled).await?; Ok(()) }) } -async fn cli_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { +async fn cli_main( + arg0_paths: Arg0DispatchPaths, + remote_control_disabled: bool, +) -> anyhow::Result<()> { let MultitoolCli { config_overrides: mut root_config_overrides, feature_toggles, @@ -1118,7 +1122,18 @@ async fn cli_main(arg0_paths: Arg0DispatchPaths) -> anyhow::Result<()> { }; let auth = auth.try_into_settings()?; let runtime_options = codex_app_server::AppServerRuntimeOptions { - remote_control_enabled: remote_control, + remote_control_startup_mode: match (remote_control, remote_control_disabled) + { + (true, _) => { + codex_app_server::RemoteControlStartupMode::EnabledEphemeral + } + (false, true) => { + codex_app_server::RemoteControlStartupMode::DisabledEphemeral + } + (false, false) => { + codex_app_server::RemoteControlStartupMode::ResolvePersisted + } + }, ..Default::default() }; codex_app_server::run_main_with_transport_options( @@ -3411,6 +3426,12 @@ mod tests { ); } + #[test] + fn app_server_remote_control_startup_flag_enables_remote_control() { + let enabled = app_server_from_args(["codex", "app-server", "--remote-control"].as_ref()); + assert!(enabled.remote_control); + } + #[test] fn app_server_analytics_default_enabled_with_flag() { let app_server = diff --git a/codex-rs/cli/src/remote_control_cmd.rs b/codex-rs/cli/src/remote_control_cmd.rs index 1ff8bd759..85434dbb6 100644 --- a/codex-rs/cli/src/remote_control_cmd.rs +++ b/codex-rs/cli/src/remote_control_cmd.rs @@ -115,7 +115,7 @@ async fn run_foreground_remote_control( socket_path: socket_path.clone(), }; let runtime_options = AppServerRuntimeOptions { - remote_control_enabled: true, + remote_control_startup_mode: codex_app_server::RemoteControlStartupMode::EnabledEphemeral, install_shutdown_signal_handler: false, ..Default::default() }; diff --git a/codex-rs/state/migrations/0037_remote_control_enrollments_enabled.sql b/codex-rs/state/migrations/0037_remote_control_enrollments_enabled.sql new file mode 100644 index 000000000..abb422384 --- /dev/null +++ b/codex-rs/state/migrations/0037_remote_control_enrollments_enabled.sql @@ -0,0 +1,2 @@ +ALTER TABLE remote_control_enrollments +ADD COLUMN remote_control_enabled INTEGER; diff --git a/codex-rs/state/src/runtime/remote_control.rs b/codex-rs/state/src/runtime/remote_control.rs index fa0b1823f..3482dbd68 100644 --- a/codex-rs/state/src/runtime/remote_control.rs +++ b/codex-rs/state/src/runtime/remote_control.rs @@ -11,6 +11,7 @@ pub struct RemoteControlEnrollmentRecord { pub server_id: String, pub environment_id: String, pub server_name: String, + pub remote_control_enabled: Option, } fn remote_control_app_server_client_name_key(app_server_client_name: Option<&str>) -> &str { @@ -34,7 +35,8 @@ impl StateRuntime { ) -> anyhow::Result> { let row = sqlx::query( r#" -SELECT websocket_url, account_id, app_server_client_name, server_id, environment_id, server_name +SELECT websocket_url, account_id, app_server_client_name, server_id, environment_id, server_name, + remote_control_enabled FROM remote_control_enrollments WHERE websocket_url = ? AND account_id = ? AND app_server_client_name = ? "#, @@ -56,6 +58,7 @@ WHERE websocket_url = ? AND account_id = ? AND app_server_client_name = ? server_id: row.try_get("server_id")?, environment_id: row.try_get("environment_id")?, server_name: row.try_get("server_name")?, + remote_control_enabled: row.try_get("remote_control_enabled")?, }) }) .transpose() @@ -74,8 +77,9 @@ INSERT INTO remote_control_enrollments ( server_id, environment_id, server_name, + remote_control_enabled, updated_at -) VALUES (?, ?, ?, ?, ?, ?, ?) +) VALUES (?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(websocket_url, account_id, app_server_client_name) DO UPDATE SET server_id = excluded.server_id, environment_id = excluded.environment_id, @@ -91,12 +95,39 @@ ON CONFLICT(websocket_url, account_id, app_server_client_name) DO UPDATE SET .bind(&enrollment.server_id) .bind(&enrollment.environment_id) .bind(&enrollment.server_name) + .bind(enrollment.remote_control_enabled) .bind(Utc::now().timestamp()) .execute(self.pool.as_ref()) .await?; Ok(()) } + pub async fn set_remote_control_enabled( + &self, + websocket_url: &str, + account_id: &str, + app_server_client_name: Option<&str>, + remote_control_enabled: bool, + ) -> anyhow::Result { + let result = sqlx::query( + r#" +UPDATE remote_control_enrollments +SET remote_control_enabled = ?, updated_at = ? +WHERE websocket_url = ? AND account_id = ? AND app_server_client_name = ? + "#, + ) + .bind(remote_control_enabled) + .bind(Utc::now().timestamp()) + .bind(websocket_url) + .bind(account_id) + .bind(remote_control_app_server_client_name_key( + app_server_client_name, + )) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected()) + } + pub async fn delete_remote_control_enrollment( &self, websocket_url: &str, @@ -125,7 +156,13 @@ mod tests { use super::RemoteControlEnrollmentRecord; use super::StateRuntime; use super::test_support::unique_temp_dir; + use crate::migrations::STATE_MIGRATOR; + use crate::state_db_path; use pretty_assertions::assert_eq; + use sqlx::SqlitePool; + use sqlx::migrate::Migrator; + use sqlx::sqlite::SqliteConnectOptions; + use std::borrow::Cow; #[tokio::test] async fn remote_control_enrollment_round_trips_by_target_and_account() { @@ -143,6 +180,7 @@ mod tests { server_id: "srv_e_first".to_string(), environment_id: "env_first".to_string(), server_name: "first-server".to_string(), + remote_control_enabled: Some(false), }) .await .expect("insert first enrollment"); @@ -155,6 +193,7 @@ mod tests { server_id: "srv_e_second".to_string(), environment_id: "env_second".to_string(), server_name: "second-server".to_string(), + remote_control_enabled: Some(true), }) .await .expect("insert second enrollment"); @@ -176,6 +215,7 @@ mod tests { server_id: "srv_e_first".to_string(), environment_id: "env_first".to_string(), server_name: "first-server".to_string(), + remote_control_enabled: Some(false), }) ); assert_eq!( @@ -220,6 +260,7 @@ mod tests { server_id: "srv_e_first".to_string(), environment_id: "env_first".to_string(), server_name: "first-server".to_string(), + remote_control_enabled: Some(false), }) .await .expect("insert first enrollment"); @@ -232,6 +273,7 @@ mod tests { server_id: "srv_e_second".to_string(), environment_id: "env_second".to_string(), server_name: "second-server".to_string(), + remote_control_enabled: Some(true), }) .await .expect("insert second enrollment"); @@ -275,9 +317,65 @@ mod tests { server_id: "srv_e_second".to_string(), environment_id: "env_second".to_string(), server_name: "second-server".to_string(), + remote_control_enabled: Some(true), }) ); let _ = tokio::fs::remove_dir_all(codex_home).await; } + + #[tokio::test] + async fn migration_preserves_legacy_remote_control_preference_as_null() { + let codex_home = unique_temp_dir(); + tokio::fs::create_dir_all(&codex_home) + .await + .expect("create codex home"); + let old_state_migrator = Migrator { + migrations: Cow::Owned(STATE_MIGRATOR.migrations[..36].to_vec()), + ignore_missing: false, + locking: true, + no_tx: false, + table_name: STATE_MIGRATOR.table_name.clone(), + create_schemas: STATE_MIGRATOR.create_schemas.clone(), + }; + let pool = SqlitePool::connect_with( + SqliteConnectOptions::new() + .filename(state_db_path(codex_home.as_path())) + .create_if_missing(true), + ) + .await + .expect("open old state db"); + old_state_migrator + .run(&pool) + .await + .expect("apply old state schema"); + sqlx::query("INSERT INTO remote_control_enrollments (websocket_url, account_id, app_server_client_name, server_id, environment_id, server_name, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?)") + .bind("wss://example.com/backend-api/wham/remote/control/server") + .bind("account-a") + .bind("desktop-client") + .bind("srv_e_first") + .bind("env_first") + .bind("first-server") + .bind(1_i64) + .execute(&pool) + .await + .expect("insert legacy enrollment"); + pool.close().await; + + let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string()) + .await + .expect("initialize runtime"); + let actual = runtime + .get_remote_control_enrollment( + "wss://example.com/backend-api/wham/remote/control/server", + "account-a", + Some("desktop-client"), + ) + .await + .expect("load migrated enrollment") + .expect("legacy enrollment should remain"); + assert_eq!(actual.remote_control_enabled, None); + + let _ = tokio::fs::remove_dir_all(codex_home).await; + } }