diff --git a/codex-rs/app-server-protocol/src/protocol/common.rs b/codex-rs/app-server-protocol/src/protocol/common.rs index a8bec6b402..1dc6f370ba 100644 --- a/codex-rs/app-server-protocol/src/protocol/common.rs +++ b/codex-rs/app-server-protocol/src/protocol/common.rs @@ -838,6 +838,12 @@ client_request_definitions! { serialization: global_shared_read("remote-control"), response: v2::RemoteControlStatusReadResponse, }, + #[experimental("remoteControl/pairing/start")] + RemoteControlPairingStart => "remoteControl/pairing/start" { + params: v2::RemoteControlPairingStartParams, + serialization: global_shared_read("remote-control"), + response: v2::RemoteControlPairingStartResponse, + }, #[experimental("collaborationMode/list")] /// Lists collaboration mode presets. CollaborationModeList => "collaborationMode/list" { @@ -1961,6 +1967,17 @@ mod tests { }, }; assert_eq!(mcp_resource_read.serialization_scope(), None); + + let remote_control_pairing_start = ClientRequest::RemoteControlPairingStart { + request_id: request_id(), + params: v2::RemoteControlPairingStartParams::default(), + }; + assert_eq!( + remote_control_pairing_start.serialization_scope(), + Some(ClientRequestSerializationScope::GlobalSharedRead( + "remote-control" + )) + ); } #[test] 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 c8ad617e7d..6bd8acf4b6 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 @@ -44,6 +44,25 @@ pub struct RemoteControlStatusReadResponse { pub environment_id: Option, } +#[derive(Serialize, Deserialize, Debug, Clone, Default, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct RemoteControlPairingStartParams { + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub manual_code: bool, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, JsonSchema, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export_to = "v2/")] +pub struct RemoteControlPairingStartResponse { + pub pairing_code: String, + pub manual_pairing_code: Option, + pub environment_id: String, + #[ts(type = "number")] + pub expires_at: i64, +} + #[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq, JsonSchema, TS)] #[serde(rename_all = "camelCase")] #[ts(rename_all = "camelCase", export_to = "v2/")] 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 357a98a30d..8948a69852 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,9 +1,10 @@ mod client_tracker; -mod enroll; mod protocol; mod segment; +mod server; mod websocket; +use crate::transport::remote_control::server::SharedRemoteControlServer; use crate::transport::remote_control::websocket::RemoteControlChannels; use crate::transport::remote_control::websocket::RemoteControlStatusPublisher; use crate::transport::remote_control::websocket::RemoteControlWebsocket; @@ -16,6 +17,8 @@ use super::CHANNEL_CAPACITY; use super::TransportEvent; use super::next_connection_id; use codex_app_server_protocol::RemoteControlConnectionStatus; +use codex_app_server_protocol::RemoteControlPairingStartParams; +use codex_app_server_protocol::RemoteControlPairingStartResponse; use codex_app_server_protocol::RemoteControlStatusChangedNotification; use codex_login::AuthManager; use codex_state::StateRuntime; @@ -52,6 +55,7 @@ pub struct RemoteControlHandle { enabled_tx: Arc>, status_tx: Arc>, state_db_available: bool, + server: SharedRemoteControlServer, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -129,6 +133,16 @@ impl RemoteControlHandle { self.status_tx.subscribe() } + pub async fn start_pairing( + &self, + params: RemoteControlPairingStartParams, + ) -> io::Result { + if !*self.enabled_tx.borrow() { + return Err(remote_control_pairing_disabled_error()); + } + self.server.start_pairing(params).await + } + fn publish_status( &self, connection_status: RemoteControlConnectionStatus, @@ -176,6 +190,20 @@ fn remote_control_status_with_connection_status( } } +pub(super) fn remote_control_pairing_disabled_error() -> io::Error { + io::Error::new( + io::ErrorKind::InvalidInput, + "remote control pairing requires remote control to be enabled", + ) +} + +pub(super) fn remote_control_pairing_unavailable_error() -> io::Error { + io::Error::new( + io::ErrorKind::InvalidInput, + "remote control pairing is unavailable until enrollment completes", + ) +} + pub async fn start_remote_control( config: RemoteControlStartConfig, state_db: Option>, @@ -213,6 +241,13 @@ pub async fn start_remote_control( }; let (status_tx, _status_rx) = watch::channel(initial_status); let status_publisher = RemoteControlStatusPublisher::new(status_tx.clone()); + let server = SharedRemoteControlServer::new( + remote_control_url.clone(), + installation_id.clone(), + server_name.clone(), + state_db, + auth_manager.clone(), + ); info!( remote_control_url = %remote_control_url, installation_id = %installation_id, @@ -225,6 +260,7 @@ pub async fn start_remote_control( let installation_id_for_log = installation_id.clone(); let server_name_for_log = server_name.clone(); let shutdown_token_for_log = shutdown_token.clone(); + let server_for_task = server.clone(); let join_handle = tokio::spawn(async move { info!( remote_control_url = %remote_control_url_for_log, @@ -240,12 +276,12 @@ pub async fn start_remote_control( remote_control_target, server_name, }, - state_db, auth_manager, RemoteControlChannels { transport_event_tx, status_publisher, }, + server_for_task, shutdown_token, enabled_rx, ) @@ -289,6 +325,7 @@ pub async fn start_remote_control( enabled_tx: Arc::new(enabled_tx), status_tx: Arc::new(status_tx), state_db_available, + server, }, )) } diff --git a/codex-rs/app-server-transport/src/transport/remote_control/pairing_integration_tests.rs b/codex-rs/app-server-transport/src/transport/remote_control/pairing_integration_tests.rs new file mode 100644 index 0000000000..7221545633 --- /dev/null +++ b/codex-rs/app-server-transport/src/transport/remote_control/pairing_integration_tests.rs @@ -0,0 +1,386 @@ +use super::super::protocol::RemoteControlTarget; +use super::super::protocol::StartRemoteControlPairingRequest; +use super::super::protocol::normalize_remote_control_url; +use super::super::server::RemoteControlServer; +use super::super::server::SharedRemoteControlServer; +use super::*; +use codex_app_server_protocol::RemoteControlPairingStartResponse; +use pretty_assertions::assert_eq; +use time::OffsetDateTime; + +#[tokio::test] +async fn remote_control_pairing_posts_server_token_and_maps_response() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let remote_control_target = remote_control_target_for_listener(&listener); + assert_eq!( + remote_control_target.pair_url, + format!( + "http://{}/backend-api/wham/remote/control/server/pair", + listener.local_addr().expect("listener should have addr") + ) + ); + let pairing_server = tokio::spawn(async move { + let pairing_request = accept_http_request(&listener).await; + assert_eq!( + pairing_request.request_line, + "POST /backend-api/wham/remote/control/server/pair HTTP/1.1" + ); + assert_eq!( + pairing_request.headers.get("authorization"), + Some(&"Bearer remote-control-token".to_string()) + ); + assert_eq!( + serde_json::from_str::(&pairing_request.body) + .expect("pairing body should deserialize"), + json!({ "manual_code": true }) + ); + respond_with_json( + pairing_request.stream, + json!({ + "pairing_code": "pairing-code", + "manual_pairing_code": "ABCD-EFGH", + "server_id": "server-id", + "environment_id": "environment-id", + "expires_at": "3026-05-22T12:34:56Z", + }), + ) + .await + }); + assert_eq!( + remote_control_server_state( + OffsetDateTime::from_unix_timestamp(33_336_362_096) + .expect("future timestamp should parse"), + ) + .start_pairing( + &remote_control_target, + StartRemoteControlPairingRequest { manual_code: true }, + ) + .await + .expect("pairing should succeed"), + RemoteControlPairingStartResponse { + pairing_code: "pairing-code".to_string(), + manual_pairing_code: Some("ABCD-EFGH".to_string()), + environment_id: "environment-id".to_string(), + expires_at: 33_336_362_096, + } + ); + pairing_server.await.expect("pairing server should join"); +} + +#[tokio::test] +async fn remote_control_pairing_preserves_backend_error_context() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let remote_control_target = remote_control_target_for_listener(&listener); + let expected_pair_url = remote_control_target.pair_url.clone(); + let pairing_server = tokio::spawn(async move { + let pairing_request = accept_http_request(&listener).await; + respond_with_status_and_headers( + pairing_request.stream, + "503 Service Unavailable", + &[("x-request-id", "request-123"), ("cf-ray", "ray-123")], + "pairing unavailable", + ) + .await + }); + + assert_eq!( + remote_control_server_state( + OffsetDateTime::from_unix_timestamp(33_336_362_096) + .expect("future timestamp should parse"), + ) + .start_pairing( + &remote_control_target, + StartRemoteControlPairingRequest { manual_code: false }, + ) + .await + .expect_err("pairing should fail") + .to_string(), + format!( + "remote control pairing failed at `{expected_pair_url}`: HTTP 503 Service Unavailable, request-id: request-123, cf-ray: ray-123, body: pairing unavailable" + ) + ); + pairing_server.await.expect("pairing server should join"); +} + +#[tokio::test] +async fn remote_control_pairing_rejects_expired_server_token() { + let err = remote_control_server_state( + OffsetDateTime::from_unix_timestamp(0).expect("expired timestamp should parse"), + ) + .start_pairing( + &RemoteControlTarget { + websocket_url: "ws://unused".to_string(), + enroll_url: "http://unused".to_string(), + refresh_url: "http://unused".to_string(), + pair_url: "http://unused".to_string(), + }, + StartRemoteControlPairingRequest { manual_code: false }, + ) + .await + .expect_err("expired server token should fail pairing"); + + assert_eq!(err.kind(), std::io::ErrorKind::InvalidInput); + assert_eq!( + err.to_string(), + "remote control pairing is unavailable because the server token expired" + ); +} + +#[tokio::test] +async fn remote_control_pairing_refreshes_server_token_before_posting() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let codex_home = TempDir::new().expect("temp dir should create"); + let server = shared_remote_control_server_for_listener( + &listener, + &codex_home, + remote_control_auth_manager(), + ) + .await; + server + .replace(Some(RemoteControlServer { + account_id: "account_id".to_string(), + environment_id: "env_test".to_string(), + server_id: "srv_e_test".to_string(), + server_name: "server-name".to_string(), + remote_control_token: None, + expires_at: None, + })) + .await; + + let pairing_task = tokio::spawn({ + let server = server.clone(); + async move { + server + .start_pairing(RemoteControlPairingStartParams::default()) + .await + } + }); + 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( + "srv_e_test", + "env_test", + TEST_REFRESHED_REMOTE_CONTROL_SERVER_TOKEN, + ), + ) + .await; + + let pairing_request = accept_http_request(&listener).await; + assert_eq!( + pairing_request.request_line, + "POST /backend-api/wham/remote/control/server/pair HTTP/1.1" + ); + assert_eq!( + pairing_request.headers.get("authorization"), + Some(&format!( + "Bearer {TEST_REFRESHED_REMOTE_CONTROL_SERVER_TOKEN}" + )) + ); + respond_with_json( + pairing_request.stream, + json!({ + "pairing_code": "pairing-code", + "manual_pairing_code": null, + "server_id": "srv_e_test", + "environment_id": "env_test", + "expires_at": "3026-05-22T12:34:56Z", + }), + ) + .await; + pairing_task + .await + .expect("pairing task should join") + .expect("pairing should succeed"); +} + +#[tokio::test] +async fn remote_control_pairing_does_not_require_websocket_connection() { + 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"); + 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(remote_control_state_runtime(&codex_home).await), + remote_control_auth_manager(), + transport_event_tx, + shutdown_token.clone(), + /*app_server_client_name_rx*/ None, + /*initial_enabled*/ true, + ) + .await + .expect("remote control should start"); + + let enroll_request = accept_http_request(&listener).await; + respond_with_json( + enroll_request.stream, + remote_control_server_token_response( + "srv_e_initial", + "env_initial", + TEST_REMOTE_CONTROL_SERVER_TOKEN, + ), + ) + .await; + let stalled_websocket_request = accept_http_request(&listener).await; + assert_eq!( + stalled_websocket_request.request_line, + "GET /backend-api/wham/remote/control/server HTTP/1.1" + ); + let pairing_task = tokio::spawn({ + let remote_handle = remote_handle.clone(); + async move { + remote_handle + .start_pairing(RemoteControlPairingStartParams::default()) + .await + } + }); + let stalled_pairing_request = accept_http_request(&listener).await; + assert_eq!( + stalled_pairing_request.request_line, + "POST /backend-api/wham/remote/control/server/pair HTTP/1.1" + ); + assert_eq!( + remote_handle.status().status, + RemoteControlConnectionStatus::Connecting + ); + respond_with_json( + stalled_pairing_request.stream, + json!({ + "pairing_code": "pairing-code", + "manual_pairing_code": null, + "server_id": "srv_e_initial", + "environment_id": "env_initial", + "expires_at": "3026-05-22T12:34:56Z", + }), + ) + .await; + pairing_task + .await + .expect("pairing task should join") + .expect("pairing should succeed before websocket connects"); + + drop(stalled_websocket_request); + shutdown_token.cancel(); + let _ = remote_task.await; +} + +#[tokio::test] +async fn remote_control_auth_change_cancels_in_flight_pairing() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let codex_home = TempDir::new().expect("temp dir should create"); + save_auth( + codex_home.path(), + &remote_control_auth_dot_json(Some("account_id")), + AuthCredentialsStoreMode::File, + ) + .expect("initial auth should save"); + 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 server = + shared_remote_control_server_for_listener(&listener, &codex_home, auth_manager.clone()) + .await; + let pairing_task = tokio::spawn({ + let server = server.clone(); + async move { + server + .start_pairing(RemoteControlPairingStartParams::default()) + .await + } + }); + let enroll_request = accept_http_request(&listener).await; + respond_with_json( + enroll_request.stream, + remote_control_server_token_response( + "srv_e_initial", + "env_initial", + TEST_REMOTE_CONTROL_SERVER_TOKEN, + ), + ) + .await; + let stalled_pairing_request = accept_http_request(&listener).await; + assert_eq!( + stalled_pairing_request.request_line, + "POST /backend-api/wham/remote/control/server/pair HTTP/1.1" + ); + save_auth( + codex_home.path(), + &remote_control_auth_dot_json(Some("next_account_id")), + AuthCredentialsStoreMode::File, + ) + .expect("next auth should save"); + auth_manager.reload().await; + respond_with_json( + stalled_pairing_request.stream, + json!({ + "pairing_code": "stale-pairing-code", + "manual_pairing_code": null, + "server_id": "srv_e_initial", + "environment_id": "env_initial", + "expires_at": "3026-05-22T12:34:56Z", + }), + ) + .await; + assert_eq!( + pairing_task + .await + .expect("pairing task should join") + .expect_err("stale pairing should be discarded") + .to_string(), + "remote control pairing is unavailable until enrollment completes" + ); +} + +fn remote_control_target_for_listener(listener: &TcpListener) -> RemoteControlTarget { + normalize_remote_control_url(&remote_control_url_for_listener(listener)) + .expect("target should parse") +} + +fn remote_control_server_state(expires_at: OffsetDateTime) -> RemoteControlServer { + RemoteControlServer { + account_id: "account-id".to_string(), + environment_id: "environment-id".to_string(), + server_id: "server-id".to_string(), + server_name: "server-name".to_string(), + remote_control_token: Some("remote-control-token".to_string()), + expires_at: Some(expires_at), + } +} + +async fn shared_remote_control_server_for_listener( + listener: &TcpListener, + codex_home: &TempDir, + auth_manager: Arc, +) -> SharedRemoteControlServer { + SharedRemoteControlServer::new( + remote_control_url_for_listener(listener), + TEST_INSTALLATION_ID.to_string(), + "server-name".to_string(), + Some(remote_control_state_runtime(codex_home).await), + auth_manager, + ) +} diff --git a/codex-rs/app-server-transport/src/transport/remote_control/protocol.rs b/codex-rs/app-server-transport/src/transport/remote_control/protocol.rs index b10af42e86..5856dc42b9 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/protocol.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/protocol.rs @@ -12,6 +12,7 @@ pub(super) struct RemoteControlTarget { pub(super) websocket_url: String, pub(super) enroll_url: String, pub(super) refresh_url: String, + pub(super) pair_url: String, } #[derive(Debug, Serialize)] @@ -37,6 +38,21 @@ pub(super) struct RefreshRemoteServerRequest { pub(super) installation_id: String, } +#[derive(Debug, Serialize)] +pub(super) struct StartRemoteControlPairingRequest { + pub(super) manual_code: bool, +} + +#[derive(Debug, Deserialize)] +pub(super) struct StartRemoteControlPairingResponse { + pub(super) pairing_code: String, + pub(super) manual_pairing_code: Option, + #[serde(rename = "server_id")] + pub(super) _server_id: String, + pub(super) environment_id: String, + pub(super) expires_at: String, +} + #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] #[serde(transparent)] pub struct ClientId(pub String); @@ -189,6 +205,9 @@ pub(super) fn normalize_remote_control_url( let refresh_url = remote_control_url .join("wham/remote/control/server/refresh") .map_err(map_url_parse_error)?; + let pair_url = remote_control_url + .join("wham/remote/control/server/pair") + .map_err(map_url_parse_error)?; let mut websocket_url = remote_control_url .join("wham/remote/control/server") .map_err(map_url_parse_error)?; @@ -207,6 +226,7 @@ pub(super) fn normalize_remote_control_url( websocket_url: websocket_url.to_string(), enroll_url: enroll_url.to_string(), refresh_url: refresh_url.to_string(), + pair_url: pair_url.to_string(), }) } @@ -227,6 +247,8 @@ mod tests { .to_string(), refresh_url: "https://chatgpt.com/backend-api/wham/remote/control/server/refresh" .to_string(), + pair_url: "https://chatgpt.com/backend-api/wham/remote/control/server/pair" + .to_string(), } ); assert_eq!( @@ -242,6 +264,9 @@ mod tests { refresh_url: "https://api.chatgpt-staging.com/backend-api/wham/remote/control/server/refresh" .to_string(), + pair_url: + "https://api.chatgpt-staging.com/backend-api/wham/remote/control/server/pair" + .to_string(), } ); } @@ -258,6 +283,8 @@ mod tests { .to_string(), refresh_url: "http://localhost:8080/backend-api/wham/remote/control/server/refresh" .to_string(), + pair_url: "http://localhost:8080/backend-api/wham/remote/control/server/pair" + .to_string(), } ); assert_eq!( @@ -271,6 +298,8 @@ mod tests { refresh_url: "https://localhost:8443/backend-api/wham/remote/control/server/refresh" .to_string(), + pair_url: "https://localhost:8443/backend-api/wham/remote/control/server/pair" + .to_string(), } ); } diff --git a/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs b/codex-rs/app-server-transport/src/transport/remote_control/server.rs similarity index 52% rename from codex-rs/app-server-transport/src/transport/remote_control/enroll.rs rename to codex-rs/app-server-transport/src/transport/remote_control/server.rs index 6037caa651..100abfa759 100644 --- a/codex-rs/app-server-transport/src/transport/remote_control/enroll.rs +++ b/codex-rs/app-server-transport/src/transport/remote_control/server.rs @@ -2,8 +2,13 @@ use super::protocol::EnrollRemoteServerRequest; use super::protocol::EnrollRemoteServerResponse; use super::protocol::RefreshRemoteServerRequest; use super::protocol::RemoteControlTarget; +use super::protocol::StartRemoteControlPairingRequest; +use super::protocol::StartRemoteControlPairingResponse; +use super::remote_control_pairing_unavailable_error; use axum::http::HeaderMap; use codex_api::SharedAuthProvider; +use codex_app_server_protocol::RemoteControlPairingStartResponse; +use codex_login::AuthManager; use codex_login::default_client::build_reqwest_client; use codex_state::RemoteControlEnrollmentRecord; use codex_state::StateRuntime; @@ -11,12 +16,17 @@ use serde::Serialize; use serde::de::DeserializeOwned; use std::io; use std::io::ErrorKind; +use std::sync::Arc; use time::OffsetDateTime; use time::format_description::well_known::Rfc3339; +use tokio::sync::Mutex; +use tokio::sync::OwnedSemaphorePermit; +use tokio::sync::Semaphore; use tracing::info; use tracing::warn; const REMOTE_CONTROL_ENROLL_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); +const REMOTE_CONTROL_PAIRING_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); const REMOTE_CONTROL_RESPONSE_BODY_MAX_BYTES: usize = 4096; const REMOTE_CONTROL_SERVER_TOKEN_REFRESH_SKEW_SECS: i64 = 30; @@ -27,7 +37,7 @@ pub(super) const REMOTE_CONTROL_ACCOUNT_ID_HEADER: &str = "chatgpt-account-id"; pub(super) const REMOTE_CONTROL_INSTALLATION_ID_HEADER: &str = "x-codex-installation-id"; #[derive(Debug, Clone, PartialEq, Eq)] -pub(super) struct RemoteControlEnrollment { +pub(super) struct RemoteControlServer { pub(super) account_id: String, pub(super) environment_id: String, pub(super) server_id: String, @@ -36,7 +46,7 @@ pub(super) struct RemoteControlEnrollment { pub(super) expires_at: Option, } -impl RemoteControlEnrollment { +impl RemoteControlServer { pub(super) fn should_refresh_server_token(&self) -> bool { self.remote_control_token.is_none() || self.expires_at.is_none_or(|expires_at| { @@ -62,7 +72,7 @@ pub(super) async fn load_persisted_remote_control_enrollment( remote_control_target: &RemoteControlTarget, account_id: &str, app_server_client_name: Option<&str>, -) -> io::Result> { +) -> io::Result> { let Some(state_db) = state_db else { return Err(io::Error::new( ErrorKind::NotFound, @@ -100,7 +110,7 @@ pub(super) async fn load_persisted_remote_control_enrollment( enrollment.server_id, enrollment.environment_id ); - Ok(Some(RemoteControlEnrollment { + Ok(Some(RemoteControlServer { account_id: enrollment.account_id, environment_id: enrollment.environment_id, server_id: enrollment.server_id, @@ -124,7 +134,7 @@ pub(super) async fn update_persisted_remote_control_enrollment( remote_control_target: &RemoteControlTarget, account_id: &str, app_server_client_name: Option<&str>, - enrollment: Option<&RemoteControlEnrollment>, + enrollment: Option<&RemoteControlServer>, ) -> io::Result<()> { let Some(state_db) = state_db else { return Err(io::Error::new( @@ -211,8 +221,14 @@ fn redact_remote_control_response_body(body: &str) -> String { let Some(body_object) = body_json.as_object_mut() else { return body.to_string(); }; - if let Some(remote_control_token) = body_object.get_mut("remote_control_token") { - *remote_control_token = serde_json::Value::String("".to_string()); + for sensitive_field in [ + "remote_control_token", + "pairing_code", + "manual_pairing_code", + ] { + if let Some(value) = body_object.get_mut(sensitive_field) { + *value = serde_json::Value::String("".to_string()); + } } body_json.to_string() } @@ -230,84 +246,526 @@ pub(crate) fn format_headers(headers: &HeaderMap) -> String { format!("request-id: {request_id_str}, cf-ray: {cf_ray_str}") } -pub(super) async fn enroll_remote_control_server( - remote_control_target: &RemoteControlTarget, - auth: &RemoteControlConnectionAuth, - installation_id: &str, - server_name: &str, -) -> io::Result { - let enroll_url = &remote_control_target.enroll_url; - let request = EnrollRemoteServerRequest { - name: server_name.to_string(), - os: std::env::consts::OS, - arch: std::env::consts::ARCH, - app_server_version: env!("CARGO_PKG_VERSION"), - installation_id: installation_id.to_string(), - }; - let enrollment_response = send_remote_control_server_request::<_, EnrollRemoteServerResponse>( - enroll_url, - auth, - installation_id, - &request, - "enroll", - "server enrollment", - ) - .await?; - let mut enrollment = RemoteControlEnrollment { - account_id: auth.account_id.clone(), - environment_id: enrollment_response.environment_id, - server_id: enrollment_response.server_id, - server_name: server_name.to_string(), - remote_control_token: None, - expires_at: None, - }; - update_remote_control_server_token( - &mut enrollment, - enroll_url, - enrollment_response.remote_control_token, - enrollment_response.expires_at, - )?; - Ok(enrollment) -} - -pub(super) async fn refresh_remote_control_server( - remote_control_target: &RemoteControlTarget, - auth: &RemoteControlConnectionAuth, - installation_id: &str, - enrollment: &mut RemoteControlEnrollment, -) -> io::Result<()> { - let refresh_url = &remote_control_target.refresh_url; - let request = RefreshRemoteServerRequest { - server_id: enrollment.server_id.clone(), - installation_id: installation_id.to_string(), - }; - let refreshed = send_remote_control_server_request::<_, EnrollRemoteServerResponse>( - refresh_url, - auth, - installation_id, - &request, - "refresh", - "server refresh", - ) - .await?; - if refreshed.server_id != enrollment.server_id - || refreshed.environment_id != enrollment.environment_id - { - return Err(io::Error::other(format!( - "remote control server refresh returned mismatched enrollment: expected server_id={}, environment_id={}; got server_id={}, environment_id={}", - enrollment.server_id, - enrollment.environment_id, - refreshed.server_id, - refreshed.environment_id - ))); +impl RemoteControlServer { + pub(super) async fn enroll( + remote_control_target: &RemoteControlTarget, + auth: &RemoteControlConnectionAuth, + installation_id: &str, + server_name: &str, + ) -> io::Result { + let enroll_url = &remote_control_target.enroll_url; + let request = EnrollRemoteServerRequest { + name: server_name.to_string(), + os: std::env::consts::OS, + arch: std::env::consts::ARCH, + app_server_version: env!("CARGO_PKG_VERSION"), + installation_id: installation_id.to_string(), + }; + let enrollment_response = + send_remote_control_server_request::<_, EnrollRemoteServerResponse>( + enroll_url, + auth, + installation_id, + &request, + "enroll", + "server enrollment", + ) + .await?; + let mut server = Self { + account_id: auth.account_id.clone(), + environment_id: enrollment_response.environment_id, + server_id: enrollment_response.server_id, + server_name: server_name.to_string(), + remote_control_token: None, + expires_at: None, + }; + update_remote_control_server_token( + &mut server, + enroll_url, + enrollment_response.remote_control_token, + enrollment_response.expires_at, + )?; + Ok(server) } - update_remote_control_server_token( - enrollment, - refresh_url, - refreshed.remote_control_token, - refreshed.expires_at, - ) + pub(super) async fn refresh_server_token( + &mut self, + remote_control_target: &RemoteControlTarget, + auth: &RemoteControlConnectionAuth, + installation_id: &str, + ) -> io::Result<()> { + let refresh_url = &remote_control_target.refresh_url; + let request = RefreshRemoteServerRequest { + server_id: self.server_id.clone(), + installation_id: installation_id.to_string(), + }; + let refreshed = send_remote_control_server_request::<_, EnrollRemoteServerResponse>( + refresh_url, + auth, + installation_id, + &request, + "refresh", + "server refresh", + ) + .await?; + if refreshed.server_id != self.server_id || refreshed.environment_id != self.environment_id + { + return Err(io::Error::other(format!( + "remote control server refresh returned mismatched enrollment: expected server_id={}, environment_id={}; got server_id={}, environment_id={}", + self.server_id, self.environment_id, refreshed.server_id, refreshed.environment_id + ))); + } + + update_remote_control_server_token( + self, + refresh_url, + refreshed.remote_control_token, + refreshed.expires_at, + ) + } + + pub(super) async fn start_pairing( + &self, + remote_control_target: &RemoteControlTarget, + request: StartRemoteControlPairingRequest, + ) -> io::Result { + let remote_control_token = self + .remote_control_token + .as_deref() + .ok_or_else(remote_control_pairing_unavailable_error)?; + let expires_at = self + .expires_at + .ok_or_else(remote_control_pairing_unavailable_error)?; + if expires_at <= OffsetDateTime::now_utc() { + return Err(io::Error::new( + ErrorKind::InvalidInput, + "remote control pairing is unavailable because the server token expired", + )); + } + + let pairing_url = &remote_control_target.pair_url; + let response = build_reqwest_client() + .post(pairing_url) + .timeout(REMOTE_CONTROL_PAIRING_TIMEOUT) + .bearer_auth(remote_control_token) + .json(&request) + .send() + .await + .map_err(|err| { + io::Error::other(format!( + "failed to start remote control pairing at `{pairing_url}`: {err}" + )) + })?; + let headers = response.headers().clone(); + let status = response.status(); + let body = response.bytes().await.map_err(|err| { + io::Error::other(format!( + "failed to read remote control pairing response from `{pairing_url}`: {err}" + )) + })?; + let body_preview = preview_remote_control_response_body(&body); + if !status.is_success() { + let error_kind = match status.as_u16() { + 401 | 403 => ErrorKind::PermissionDenied, + 404 => ErrorKind::NotFound, + _ => ErrorKind::Other, + }; + return Err(io::Error::new( + error_kind, + format!( + "remote control pairing failed at `{pairing_url}`: HTTP {status}, {}, body: {body_preview}", + format_headers(&headers) + ), + )); + } + + let pairing = + serde_json::from_slice::(&body).map_err(|err| { + io::Error::other(format!( + "failed to parse remote control pairing response from `{pairing_url}`: HTTP {status}, {}, body: {body_preview}, decode error: {err}", + format_headers(&headers) + )) + })?; + let StartRemoteControlPairingResponse { + pairing_code, + manual_pairing_code, + _server_id: _, + environment_id, + expires_at, + } = pairing; + let expires_at = OffsetDateTime::parse(&expires_at, &Rfc3339) + .map_err(|err| { + io::Error::new( + ErrorKind::InvalidData, + format!("invalid remote control pairing expires_at: {err}"), + ) + })? + .unix_timestamp(); + + Ok(RemoteControlPairingStartResponse { + pairing_code, + manual_pairing_code, + environment_id, + expires_at, + }) + } +} + +#[derive(Clone)] +pub(super) struct SharedRemoteControlServer { + remote_control_url: String, + installation_id: String, + server_name: String, + state_db: Option>, + auth_manager: Arc, + app_server_client_name: Arc>>, + server: Arc>>, + server_mutations: Arc, +} + +impl SharedRemoteControlServer { + pub(super) fn new( + remote_control_url: String, + installation_id: String, + server_name: String, + state_db: Option>, + auth_manager: Arc, + ) -> Self { + Self { + remote_control_url, + installation_id, + server_name, + state_db, + auth_manager, + app_server_client_name: Arc::new(std::sync::Mutex::new(None)), + server: Arc::new(Mutex::new(None)), + server_mutations: Arc::new(Semaphore::new(1)), + } + } + + pub(super) fn set_app_server_client_name(&self, app_server_client_name: Option) { + *self + .app_server_client_name + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = app_server_client_name; + } + + pub(super) async fn snapshot(&self) -> Option { + self.server.lock().await.clone() + } + + pub(super) async fn clear(&self) { + let Ok(_server_mutation) = self.server_mutation().await else { + return; + }; + *self.server.lock().await = None; + } + + #[cfg(test)] + pub(super) async fn replace(&self, server: Option) { + let Ok(_server_mutation) = self.server_mutation().await else { + return; + }; + *self.server.lock().await = server; + } + + pub(super) async fn prepare( + &self, + auth: &RemoteControlConnectionAuth, + ) -> io::Result<(RemoteControlTarget, RemoteControlServer)> { + self.prepare_with_environment_updates(auth, |_| {}).await + } + + pub(super) async fn prepare_with_environment_updates( + &self, + auth: &RemoteControlConnectionAuth, + mut publish_environment_id: impl FnMut(Option), + ) -> io::Result<(RemoteControlTarget, RemoteControlServer)> { + let remote_control_target = + super::protocol::normalize_remote_control_url(&self.remote_control_url)?; + let app_server_client_name = self + .app_server_client_name + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(); + let _server_mutation = self.server_mutation().await?; + let Some(state_db) = self.state_db.as_deref() else { + *self.server.lock().await = None; + return Err(io::Error::new( + ErrorKind::NotFound, + "remote control requires sqlite state db", + )); + }; + let server = self.snapshot().await; + let server = prepare_remote_control_server( + &RemoteControlServerPrepareContext { + state_db, + remote_control_target: &remote_control_target, + auth, + installation_id: &self.installation_id, + server_name: &self.server_name, + app_server_client_name: app_server_client_name.as_deref(), + }, + server, + &mut publish_environment_id, + ) + .await?; + *self.server.lock().await = server.clone(); + Ok(( + remote_control_target, + server + .ok_or_else(|| io::Error::other("missing remote control server after prepare"))?, + )) + } + + pub(super) async fn start_pairing( + &self, + params: codex_app_server_protocol::RemoteControlPairingStartParams, + ) -> io::Result { + let auth_change_rx = self.auth_manager.auth_change_receiver(); + let auth_change_revision = *auth_change_rx.borrow(); + let auth = load_remote_control_auth(&self.auth_manager) + .await + .map_err(|err| { + if matches!( + err.kind(), + ErrorKind::PermissionDenied | ErrorKind::WouldBlock + ) { + super::remote_control_pairing_unavailable_error() + } else { + err + } + })?; + let (remote_control_target, server) = self.prepare(&auth).await?; + let pairing_response = server + .start_pairing( + &remote_control_target, + StartRemoteControlPairingRequest { + manual_code: params.manual_code, + }, + ) + .await; + if *auth_change_rx.borrow() != auth_change_revision { + return Err(super::remote_control_pairing_unavailable_error()); + } + if pairing_response + .as_ref() + .is_err_and(|err| err.kind() == ErrorKind::PermissionDenied) + { + self.clear_server_token_if_current(&server).await; + } + if pairing_response + .as_ref() + .is_err_and(|err| err.kind() == ErrorKind::NotFound) + { + self.clear_server_if_current(&remote_control_target, &server) + .await; + } + pairing_response + } + + pub(super) async fn clear_server_token_if_current( + &self, + expected_server: &RemoteControlServer, + ) { + let Ok(_server_mutation) = self.server_mutation().await else { + return; + }; + let mut server = self.server.lock().await; + if server.as_ref() != Some(expected_server) { + return; + } + if let Some(server) = server.as_mut() { + server.clear_server_token(); + } + } + + pub(super) async fn clear_server_if_current( + &self, + remote_control_target: &RemoteControlTarget, + expected_server: &RemoteControlServer, + ) { + let app_server_client_name = self + .app_server_client_name + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(); + let Ok(_server_mutation) = self.server_mutation().await else { + return; + }; + if self.snapshot().await.as_ref() != Some(expected_server) { + return; + } + if update_persisted_remote_control_enrollment( + self.state_db.as_deref(), + remote_control_target, + &expected_server.account_id, + app_server_client_name.as_deref(), + None, + ) + .await + .is_ok() + { + *self.server.lock().await = None; + } + } + + async fn server_mutation(&self) -> io::Result { + self.server_mutations + .clone() + .acquire_owned() + .await + .map_err(|_| io::Error::other("remote control server mutation semaphore closed")) + } +} + +pub(super) async fn load_remote_control_auth( + auth_manager: &AuthManager, +) -> io::Result { + let mut reloaded = false; + let auth = loop { + let Some(auth) = auth_manager.auth().await else { + if reloaded { + return Err(io::Error::new( + ErrorKind::PermissionDenied, + "remote control requires ChatGPT authentication", + )); + } + auth_manager.reload().await; + reloaded = true; + continue; + }; + if !auth.uses_codex_backend() { + break auth; + } + if auth.get_account_id().is_none() && !reloaded { + auth_manager.reload().await; + reloaded = true; + continue; + } + break auth; + }; + + if !auth.uses_codex_backend() { + return Err(io::Error::new( + ErrorKind::PermissionDenied, + "remote control requires ChatGPT authentication; API key auth is not supported", + )); + } + + Ok(RemoteControlConnectionAuth { + auth_provider: codex_model_provider::auth_provider_from_auth(&auth), + account_id: auth.get_account_id().ok_or_else(|| { + io::Error::new( + ErrorKind::WouldBlock, + "remote control enrollment is waiting for a ChatGPT account id", + ) + })?, + }) +} + +struct RemoteControlServerPrepareContext<'a> { + state_db: &'a StateRuntime, + remote_control_target: &'a RemoteControlTarget, + auth: &'a RemoteControlConnectionAuth, + installation_id: &'a str, + server_name: &'a str, + app_server_client_name: Option<&'a str>, +} + +async fn prepare_remote_control_server( + context: &RemoteControlServerPrepareContext<'_>, + mut server: Option, + publish_environment_id: &mut impl FnMut(Option), +) -> io::Result> { + if server + .as_ref() + .is_some_and(|server| server.account_id != context.auth.account_id) + { + server = None; + publish_environment_id(None); + } + if server.is_none() { + server = load_persisted_remote_control_enrollment( + Some(context.state_db), + context.remote_control_target, + &context.auth.account_id, + context.app_server_client_name, + ) + .await? + .map(|mut server| { + server.server_name = context.server_name.to_string(); + server + }); + } + if let Some(server) = server.as_ref() { + publish_environment_id(Some(server.environment_id.clone())); + } + if server.is_none() { + let new_server = RemoteControlServer::enroll( + context.remote_control_target, + context.auth, + context.installation_id, + context.server_name, + ) + .await?; + update_persisted_remote_control_enrollment( + Some(context.state_db), + context.remote_control_target, + &context.auth.account_id, + context.app_server_client_name, + Some(&new_server), + ) + .await?; + publish_environment_id(Some(new_server.environment_id.clone())); + server = Some(new_server); + } + if server + .as_ref() + .is_some_and(RemoteControlServer::should_refresh_server_token) + { + let refresh_result = server + .as_mut() + .ok_or_else(|| io::Error::other("missing remote control server before refresh"))? + .refresh_server_token( + context.remote_control_target, + context.auth, + context.installation_id, + ) + .await; + if let Err(err) = refresh_result { + if err.kind() != ErrorKind::NotFound { + return Err(err); + } + update_persisted_remote_control_enrollment( + Some(context.state_db), + context.remote_control_target, + &context.auth.account_id, + context.app_server_client_name, + None, + ) + .await?; + publish_environment_id(None); + let new_server = RemoteControlServer::enroll( + context.remote_control_target, + context.auth, + context.installation_id, + context.server_name, + ) + .await?; + update_persisted_remote_control_enrollment( + Some(context.state_db), + context.remote_control_target, + &context.auth.account_id, + context.app_server_client_name, + Some(&new_server), + ) + .await?; + publish_environment_id(Some(new_server.environment_id.clone())); + server = Some(new_server); + } + } + Ok(server) } async fn send_remote_control_server_request( @@ -371,7 +829,7 @@ where } fn update_remote_control_server_token( - enrollment: &mut RemoteControlEnrollment, + server: &mut RemoteControlServer, url: &str, token: String, expires_at: String, @@ -381,8 +839,8 @@ fn update_remote_control_server_token( "failed to parse remote control server token expiry from `{url}`: {err}" )) })?; - enrollment.remote_control_token = Some(token); - enrollment.expires_at = Some(expires_at); + server.remote_control_token = Some(token); + server.expires_at = Some(expires_at); Ok(()) } @@ -411,7 +869,7 @@ mod tests { #[test] fn remote_control_enrollment_refreshes_server_token_before_expiry() { - let expires_soon = RemoteControlEnrollment { + let expires_soon = RemoteControlServer { account_id: "account-a".to_string(), environment_id: "env_first".to_string(), server_id: "srv_e_first".to_string(), @@ -419,7 +877,7 @@ mod tests { remote_control_token: Some("expires-soon".to_string()), expires_at: Some(OffsetDateTime::now_utc() + time::Duration::seconds(29)), }; - let expires_later = RemoteControlEnrollment { + let expires_later = RemoteControlServer { expires_at: Some(OffsetDateTime::now_utc() + time::Duration::seconds(31)), remote_control_token: Some("expires-later".to_string()), ..expires_soon.clone() @@ -443,6 +901,20 @@ mod tests { ); } + #[test] + fn preview_remote_control_response_body_redacts_pairing_codes() { + assert_eq!( + serde_json::from_str::(&preview_remote_control_response_body( + br#"{"pairing_code":"pairing-code","manual_pairing_code":"ABCD-EFGH"}"# + )) + .expect("redacted response preview should stay valid json"), + json!({ + "pairing_code": "", + "manual_pairing_code": "", + }) + ); + } + #[tokio::test] async fn persisted_remote_control_enrollment_round_trips_by_target_and_account() { let codex_home = TempDir::new().expect("temp dir should create"); @@ -452,7 +924,7 @@ mod tests { let second_target = normalize_remote_control_url("https://api.chatgpt-staging.com/other/control") .expect("second target should parse"); - let first_enrollment = RemoteControlEnrollment { + let first_enrollment = RemoteControlServer { account_id: "account-a".to_string(), environment_id: "env_first".to_string(), server_id: "srv_e_first".to_string(), @@ -460,7 +932,7 @@ mod tests { remote_control_token: None, expires_at: None, }; - let second_enrollment = RemoteControlEnrollment { + let second_enrollment = RemoteControlServer { account_id: "account-a".to_string(), environment_id: "env_second".to_string(), server_id: "srv_e_second".to_string(), @@ -532,7 +1004,7 @@ mod tests { let second_target = normalize_remote_control_url("https://api.chatgpt-staging.com/other/control") .expect("second target should parse"); - let first_enrollment = RemoteControlEnrollment { + let first_enrollment = RemoteControlServer { account_id: "account-a".to_string(), environment_id: "env_first".to_string(), server_id: "srv_e_first".to_string(), @@ -540,7 +1012,7 @@ mod tests { remote_control_token: None, expires_at: None, }; - let second_enrollment = RemoteControlEnrollment { + let second_enrollment = RemoteControlServer { account_id: "account-a".to_string(), environment_id: "env_second".to_string(), server_id: "srv_e_second".to_string(), @@ -627,7 +1099,7 @@ mod tests { respond_with_json(stream, response_body).await; }); - let err = enroll_remote_control_server( + let err = RemoteControlServer::enroll( &remote_control_target, &RemoteControlConnectionAuth { auth_provider: codex_model_provider::unauthenticated_auth_provider(), 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 85bcf09222..179ec5daab 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 @@ -1,13 +1,13 @@ -use super::enroll::REMOTE_CONTROL_ACCOUNT_ID_HEADER; -use super::enroll::REMOTE_CONTROL_INSTALLATION_ID_HEADER; -use super::enroll::RemoteControlEnrollment; -use super::enroll::load_persisted_remote_control_enrollment; -use super::enroll::update_persisted_remote_control_enrollment; use super::protocol::ClientEnvelope; use super::protocol::ClientEvent; use super::protocol::ClientId; use super::protocol::StreamId; use super::protocol::normalize_remote_control_url; +use super::server::REMOTE_CONTROL_ACCOUNT_ID_HEADER; +use super::server::REMOTE_CONTROL_INSTALLATION_ID_HEADER; +use super::server::RemoteControlServer; +use super::server::load_persisted_remote_control_enrollment; +use super::server::update_persisted_remote_control_enrollment; use super::websocket::REMOTE_CONTROL_PROTOCOL_VERSION; use super::*; use crate::outgoing_message::OutgoingMessage; @@ -20,6 +20,7 @@ use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::ConfigWarningNotification; use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::RemoteControlConnectionStatus; +use codex_app_server_protocol::RemoteControlPairingStartParams; use codex_app_server_protocol::RemoteControlStatusChangedNotification; use codex_app_server_protocol::ServerNotification; use codex_config::types::AuthCredentialsStoreMode; @@ -895,6 +896,9 @@ async fn remote_control_handle_enable_disable_stops_and_restarts_connections() { let _ = remote_task.await; } +#[path = "pairing_integration_tests.rs"] +mod pairing_integration_tests; + #[tokio::test] async fn remote_control_transport_clears_outgoing_buffer_when_backend_acks() { let listener = TcpListener::bind("127.0.0.1:0") @@ -1312,7 +1316,7 @@ async fn remote_control_http_mode_refreshes_persisted_enrollment_before_connecti let state_db = remote_control_state_runtime(&codex_home).await; let remote_control_target = normalize_remote_control_url(&remote_control_url).expect("target should parse"); - let persisted_enrollment = RemoteControlEnrollment { + let persisted_enrollment = RemoteControlServer { account_id: "account_id".to_string(), environment_id: "env_persisted".to_string(), server_id: "srv_e_persisted".to_string(), @@ -1417,7 +1421,7 @@ async fn remote_control_stdio_mode_waits_for_client_name_before_connecting() { let remote_control_target = normalize_remote_control_url(&remote_control_url).expect("target should parse"); let app_server_client_name = "stdio-client"; - let persisted_enrollment = RemoteControlEnrollment { + let persisted_enrollment = RemoteControlServer { account_id: "account_id".to_string(), environment_id: "env_persisted".to_string(), server_id: "srv_e_persisted".to_string(), @@ -1505,7 +1509,7 @@ async fn remote_control_waits_for_account_id_before_enrolling() { ) .await; let expected_server_name = gethostname().to_string_lossy().trim().to_string(); - let expected_enrollment = RemoteControlEnrollment { + let expected_enrollment = RemoteControlServer { account_id: "account_id".to_string(), environment_id: "env_ready".to_string(), server_id: "srv_e_ready".to_string(), @@ -1582,7 +1586,7 @@ async fn remote_control_http_mode_reenrolls_when_refresh_reports_stale_enrollmen let remote_control_target = normalize_remote_control_url(&remote_control_url).expect("target should parse"); let expected_server_name = gethostname().to_string_lossy().trim().to_string(); - let stale_enrollment = RemoteControlEnrollment { + let stale_enrollment = RemoteControlServer { account_id: "account_id".to_string(), environment_id: "env_stale".to_string(), server_id: "srv_e_stale".to_string(), @@ -1590,7 +1594,7 @@ async fn remote_control_http_mode_reenrolls_when_refresh_reports_stale_enrollmen remote_control_token: None, expires_at: None, }; - let refreshed_enrollment = RemoteControlEnrollment { + let refreshed_enrollment = RemoteControlServer { account_id: "account_id".to_string(), environment_id: "env_refreshed".to_string(), server_id: "srv_e_refreshed".to_string(), @@ -1699,7 +1703,7 @@ async fn remote_control_http_mode_clears_stale_persisted_enrollment_after_404() let remote_control_target = normalize_remote_control_url(&remote_control_url).expect("target should parse"); let expected_server_name = gethostname().to_string_lossy().trim().to_string(); - let stale_enrollment = RemoteControlEnrollment { + let stale_enrollment = RemoteControlServer { account_id: "account_id".to_string(), environment_id: "env_stale".to_string(), server_id: "srv_e_stale".to_string(), @@ -1707,7 +1711,7 @@ async fn remote_control_http_mode_clears_stale_persisted_enrollment_after_404() remote_control_token: None, expires_at: None, }; - let refreshed_enrollment = RemoteControlEnrollment { + let refreshed_enrollment = RemoteControlServer { account_id: "account_id".to_string(), environment_id: "env_refreshed".to_string(), server_id: "srv_e_refreshed".to_string(), 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 21710b5e26..73a024a94b 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,14 +1,11 @@ use crate::transport::TransportEvent; use crate::transport::remote_control::client_tracker::ClientTracker; use crate::transport::remote_control::client_tracker::REMOTE_CONTROL_IDLE_SWEEP_INTERVAL; -use crate::transport::remote_control::enroll::RemoteControlConnectionAuth; -use crate::transport::remote_control::enroll::RemoteControlEnrollment; -use crate::transport::remote_control::enroll::enroll_remote_control_server; -use crate::transport::remote_control::enroll::format_headers; -use crate::transport::remote_control::enroll::load_persisted_remote_control_enrollment; -use crate::transport::remote_control::enroll::preview_remote_control_response_body; -use crate::transport::remote_control::enroll::refresh_remote_control_server; -use crate::transport::remote_control::enroll::update_persisted_remote_control_enrollment; +use crate::transport::remote_control::server::RemoteControlServer; +use crate::transport::remote_control::server::SharedRemoteControlServer; +use crate::transport::remote_control::server::format_headers; +use crate::transport::remote_control::server::load_remote_control_auth; +use crate::transport::remote_control::server::preview_remote_control_response_body; use super::protocol::ClientEnvelope; use super::protocol::ClientEvent; @@ -28,7 +25,6 @@ use codex_app_server_protocol::RemoteControlStatusChangedNotification; use codex_core::util::backoff; use codex_login::AuthManager; use codex_login::UnauthorizedRecovery; -use codex_state::StateRuntime; use codex_utils_rustls_provider::ensure_rustls_crypto_provider; use futures::SinkExt; use futures::StreamExt; @@ -243,12 +239,11 @@ pub(crate) struct RemoteControlWebsocket { installation_id: String, server_name: String, remote_control_target: Option, - state_db: Option>, auth_manager: Arc, status_publisher: RemoteControlStatusPublisher, shutdown_token: CancellationToken, reconnect_attempt: u64, - enrollment: Option, + server: SharedRemoteControlServer, auth_recovery: UnauthorizedRecovery, auth_change_rx: watch::Receiver, client_tracker: Arc>, @@ -282,6 +277,8 @@ enum ConnectionEndReason { Shutdown, Disabled, EnabledWatchClosed, + AuthChanged, + AuthWatchClosed, ConnectionWorkerStopped, } @@ -366,17 +363,15 @@ impl RemoteControlStatusPublisher { #[derive(Clone, Copy)] pub(super) struct RemoteControlConnectOptions<'a> { installation_id: &'a str, - server_name: &'a str, subscribe_cursor: Option<&'a str>, - app_server_client_name: Option<&'a str>, } impl RemoteControlWebsocket { pub(crate) fn new( config: RemoteControlWebsocketConfig, - state_db: Option>, auth_manager: Arc, channels: RemoteControlChannels, + server: SharedRemoteControlServer, shutdown_token: CancellationToken, enabled_rx: watch::Receiver, ) -> Self { @@ -396,12 +391,11 @@ impl RemoteControlWebsocket { installation_id: config.installation_id, server_name: config.server_name, remote_control_target: config.remote_control_target, - state_db, auth_manager, status_publisher: channels.status_publisher, shutdown_token, reconnect_attempt: 0, - enrollment: None, + server, auth_recovery, auth_change_rx, client_tracker: Arc::new(Mutex::new(client_tracker)), @@ -450,6 +444,8 @@ impl RemoteControlWebsocket { return; } }; + self.server + .set_app_server_client_name(app_server_client_name.clone()); loop { if !self.wait_until_enabled().await { @@ -572,24 +568,22 @@ impl RemoteControlWebsocket { loop { let subscribe_cursor = self.state.lock().await.subscribe_cursor.clone(); - let enrollment = self.enrollment.as_ref(); + let server = self.server.snapshot().await; info!( websocket_url = %remote_control_target.websocket_url, installation_id = %self.installation_id, server_name = %self.server_name, reconnect_attempt = self.reconnect_attempt.saturating_add(1), - has_enrollment = enrollment.is_some(), - server_id = ?enrollment.map(|enrollment| enrollment.server_id.as_str()), - environment_id = ?enrollment.map(|enrollment| enrollment.environment_id.as_str()), + has_enrollment = server.is_some(), + server_id = ?server.as_ref().map(|server| server.server_id.as_str()), + environment_id = ?server.as_ref().map(|server| server.environment_id.as_str()), subscribe_cursor_present = subscribe_cursor.is_some(), app_server_client_name = ?app_server_client_name, "connecting to app-server remote control websocket" ); let connect_options = RemoteControlConnectOptions { installation_id: &self.installation_id, - server_name: &self.server_name, subscribe_cursor: subscribe_cursor.as_deref(), - app_server_client_name, }; let auth_context = RemoteControlAuthContext { auth_manager: &self.auth_manager, @@ -606,9 +600,8 @@ impl RemoteControlWebsocket { } connect_result = connect_remote_control_websocket( &remote_control_target, - self.state_db.as_deref(), auth_context, - &mut self.enrollment, + &self.server, connect_options, &self.status_publisher, ) => connect_result, @@ -623,13 +616,13 @@ impl RemoteControlWebsocket { self.auth_recovery = self.auth_manager.unauthorized_recovery(); self.status_publisher .publish_status(RemoteControlConnectionStatus::Connected); - let enrollment = self.enrollment.as_ref(); + let server = self.server.snapshot().await; info!( websocket_url = %remote_control_target.websocket_url, installation_id = %self.installation_id, server_name = %self.server_name, - server_id = ?enrollment.map(|enrollment| enrollment.server_id.as_str()), - environment_id = ?enrollment.map(|enrollment| enrollment.environment_id.as_str()), + server_id = ?server.as_ref().map(|server| server.server_id.as_str()), + environment_id = ?server.as_ref().map(|server| server.environment_id.as_str()), subscribe_cursor_present = subscribe_cursor.is_some(), response_headers = %format_headers(response.headers()), "connected to app-server remote control websocket" @@ -648,7 +641,7 @@ impl RemoteControlWebsocket { let reconnect_attempt = self.reconnect_attempt.saturating_add(1); let (reconnect_delay, reconnect_backoff_reset) = next_reconnect_delay(&mut self.reconnect_attempt); - let enrollment = self.enrollment.as_ref(); + let server = self.server.snapshot().await; warn!( websocket_url = %remote_control_target.websocket_url, installation_id = %self.installation_id, @@ -658,9 +651,9 @@ impl RemoteControlWebsocket { reconnect_attempt, reconnect_delay = ?reconnect_delay, reconnect_backoff_reset, - has_enrollment = enrollment.is_some(), - server_id = ?enrollment.map(|enrollment| enrollment.server_id.as_str()), - environment_id = ?enrollment.map(|enrollment| enrollment.environment_id.as_str()), + has_enrollment = server.is_some(), + server_id = ?server.as_ref().map(|server| server.server_id.as_str()), + environment_id = ?server.as_ref().map(|server| server.environment_id.as_str()), subscribe_cursor_present = subscribe_cursor.is_some(), "failed to connect to app-server remote control websocket" ); @@ -696,7 +689,7 @@ impl RemoteControlWebsocket { } async fn run_connection( - &self, + &mut self, websocket_connection: WebSocketStream>, shutdown_token: CancellationToken, ) -> ConnectionEndReason { @@ -731,6 +724,14 @@ impl RemoteControlWebsocket { ConnectionEndReason::EnabledWatchClosed } } + changed = self.auth_change_rx.changed() => { + if changed.is_ok() { + self.auth_recovery = self.auth_manager.unauthorized_recovery(); + ConnectionEndReason::AuthChanged + } else { + ConnectionEndReason::AuthWatchClosed + } + } _ = join_set.join_next() => ConnectionEndReason::ConnectionWorkerStopped, }; shutdown_token.cancel(); @@ -1118,7 +1119,7 @@ fn set_remote_control_header( fn build_remote_control_websocket_request( websocket_url: &str, - enrollment: &RemoteControlEnrollment, + server: &RemoteControlServer, installation_id: &str, subscribe_cursor: Option<&str>, ) -> io::Result> { @@ -1129,11 +1130,11 @@ fn build_remote_control_websocket_request( ) })?; let headers = request.headers_mut(); - set_remote_control_header(headers, "x-codex-server-id", &enrollment.server_id)?; + set_remote_control_header(headers, "x-codex-server-id", &server.server_id)?; set_remote_control_header( headers, "x-codex-name", - &base64::engine::general_purpose::STANDARD.encode(&enrollment.server_name), + &base64::engine::general_purpose::STANDARD.encode(&server.server_name), )?; set_remote_control_header( headers, @@ -1145,7 +1146,7 @@ fn build_remote_control_websocket_request( "authorization", &format!( "Bearer {}", - enrollment + server .remote_control_token .as_deref() .ok_or_else(|| io::Error::other("missing remote control server token"))? @@ -1166,51 +1167,6 @@ fn build_remote_control_websocket_request( Ok(request) } -pub(crate) async fn load_remote_control_auth( - auth_manager: &Arc, -) -> io::Result { - let mut reloaded = false; - let auth = loop { - let Some(auth) = auth_manager.auth().await else { - if reloaded { - return Err(io::Error::new( - ErrorKind::PermissionDenied, - "remote control requires ChatGPT authentication", - )); - } - auth_manager.reload().await; - reloaded = true; - continue; - }; - if !auth.uses_codex_backend() { - break auth; - } - if auth.get_account_id().is_none() && !reloaded { - auth_manager.reload().await; - reloaded = true; - continue; - } - break auth; - }; - - if !auth.uses_codex_backend() { - return Err(io::Error::new( - ErrorKind::PermissionDenied, - "remote control requires ChatGPT authentication; API key auth is not supported", - )); - } - - Ok(RemoteControlConnectionAuth { - auth_provider: codex_model_provider::auth_provider_from_auth(&auth), - account_id: auth.get_account_id().ok_or_else(|| { - io::Error::new( - ErrorKind::WouldBlock, - "remote control enrollment is waiting for a ChatGPT account id", - ) - })?, - }) -} - fn next_reconnect_delay(reconnect_attempt: &mut u64) -> (std::time::Duration, bool) { let reconnect_delay = backoff(*reconnect_attempt).min(REMOTE_CONTROL_RECONNECT_BACKOFF_CAP); let reconnect_backoff_reset = reconnect_delay == REMOTE_CONTROL_RECONNECT_BACKOFF_CAP; @@ -1224,9 +1180,8 @@ fn next_reconnect_delay(reconnect_attempt: &mut u64) -> (std::time::Duration, bo pub(super) async fn connect_remote_control_websocket( remote_control_target: &RemoteControlTarget, - state_db: Option<&StateRuntime>, - mut auth_context: RemoteControlAuthContext<'_>, - enrollment: &mut Option, + auth_context: RemoteControlAuthContext<'_>, + server: &SharedRemoteControlServer, connect_options: RemoteControlConnectOptions<'_>, status_publisher: &RemoteControlStatusPublisher, ) -> io::Result<( @@ -1235,148 +1190,42 @@ pub(super) async fn connect_remote_control_websocket( )> { ensure_rustls_crypto_provider(); - let Some(state_db) = state_db else { - *enrollment = None; - return Err(io::Error::new( - ErrorKind::NotFound, - "remote control requires sqlite state db", - )); - }; - let auth = match load_remote_control_auth(auth_context.auth_manager).await { Ok(auth) => auth, Err(err) => { if err.kind() == ErrorKind::PermissionDenied { - *enrollment = None; + server.clear().await; status_publisher.publish_environment_id(/*environment_id*/ None); } return Err(err); } }; - 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) { - info!( - "clearing in-memory remote control enrollment because account id changed: websocket_url={}, previous_account_id={:?}, current_account_id={:?}", - remote_control_target.websocket_url, - enrollment - .as_ref() - .map(|enrollment| enrollment.account_id.as_str()), - auth.account_id - ); - *enrollment = None; - status_publisher.publish_environment_id(/*environment_id*/ None); - } - - if let Some(enrollment) = enrollment.as_ref() { - status_publisher.publish_environment_id(Some(enrollment.environment_id.clone())); - } - - if enrollment.is_none() { - let loaded_enrollment = load_persisted_remote_control_enrollment( - Some(state_db), - remote_control_target, - &auth.account_id, - connect_options.app_server_client_name, - ) - .await?; - if let Some(loaded_enrollment) = loaded_enrollment.as_ref() { - status_publisher.publish_environment_id(Some(loaded_enrollment.environment_id.clone())); - } - *enrollment = loaded_enrollment.map(|mut enrollment| { - enrollment.server_name = connect_options.server_name.to_string(); - enrollment - }); - } - - enroll_remote_control_server_if_missing( - remote_control_target, - state_db, - &auth, - &mut auth_context, - enrollment, - connect_options, - status_publisher, - ) - .await?; - - if enrollment - .as_ref() - .ok_or_else(|| io::Error::other("missing remote control enrollment after enrollment step"))? - .should_refresh_server_token() - { - let enrollment_ref = enrollment.as_ref().ok_or_else(|| { - io::Error::other("missing remote control enrollment after enrollment step") - })?; - let server_id = enrollment_ref.server_id.clone(); - let environment_id = enrollment_ref.environment_id.clone(); - - info!( - "refreshing remote control server token: websocket_url={}, refresh_url={}, account_id={}, server_id={}, environment_id={}", - remote_control_target.websocket_url, - remote_control_target.refresh_url, - auth.account_id, - server_id, - environment_id - ); - let enrollment_ref = enrollment.as_mut().ok_or_else(|| { - io::Error::other("missing remote control enrollment before server refresh") - })?; - match refresh_remote_control_server( - remote_control_target, - &auth, - connect_options.installation_id, - enrollment_ref, - ) + auth_context.auth_change_rx.borrow_and_update(); + let (_, server_snapshot) = match server + .prepare_with_environment_updates(&auth, |environment_id| { + status_publisher.publish_environment_id(environment_id); + }) .await + { + Ok(server) => server, + Err(err) + if err.kind() == ErrorKind::PermissionDenied + && recover_remote_control_auth( + auth_context.auth_recovery, + auth_context.auth_change_rx, + ) + .await => { - 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_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( - remote_control_target, - state_db, - &auth, - &mut auth_context, - enrollment, - connect_options, - status_publisher, - ) - .await?; - } - Err(err) - if err.kind() == ErrorKind::PermissionDenied - && recover_remote_control_auth( - auth_context.auth_recovery, - auth_context.auth_change_rx, - ) - .await => - { - return Err(io::Error::other(format!( - "{err}; retrying after auth recovery" - ))); - } - Err(err) => return Err(err), + return Err(io::Error::other(format!( + "{err}; retrying after auth recovery" + ))); } - } - - let enrollment_ref = enrollment.as_ref().ok_or_else(|| { - io::Error::other("missing remote control enrollment after enrollment step") - })?; + Err(err) => return Err(err), + }; + status_publisher.publish_environment_id(Some(server_snapshot.environment_id.clone())); let request = build_remote_control_websocket_request( &remote_control_target.websocket_url, - enrollment_ref, + &server_snapshot, connect_options.installation_id, connect_options.subscribe_cursor, )?; @@ -1405,30 +1254,18 @@ pub(super) async fn connect_remote_control_websocket( "remote control websocket returned HTTP 404; clearing stale enrollment before re-enrolling: websocket_url={}, account_id={}, server_id={}, environment_id={}", remote_control_target.websocket_url, auth.account_id, - enrollment_ref.server_id, - enrollment_ref.environment_id + server_snapshot.server_id, + server_snapshot.environment_id ); - clear_remote_control_enrollment( - state_db, - remote_control_target, - &auth.account_id, - connect_options.app_server_client_name, - enrollment, - status_publisher, - ) - .await; + server + .clear_server_if_current(remote_control_target, &server_snapshot) + .await; + status_publisher.publish_environment_id(/*environment_id*/ None); } tungstenite::Error::Http(response) if matches!(response.status().as_u16(), 401 | 403) => { - enrollment - .as_mut() - .ok_or_else(|| { - io::Error::other( - "missing remote control enrollment after websocket auth failure", - ) - })? - .clear_server_token(); + server.clear_server_token_if_current(&server_snapshot).await; return Err(io::Error::other(format!( "remote control websocket auth failed with HTTP {}; refreshing server token before reconnect", response.status() @@ -1446,94 +1283,6 @@ pub(super) async fn connect_remote_control_websocket( } } -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 enroll_remote_control_server_if_missing( - remote_control_target: &RemoteControlTarget, - state_db: &StateRuntime, - auth: &RemoteControlConnectionAuth, - auth_context: &mut RemoteControlAuthContext<'_>, - enrollment: &mut Option, - connect_options: RemoteControlConnectOptions<'_>, - status_publisher: &RemoteControlStatusPublisher, -) -> io::Result<()> { - if enrollment.is_some() { - return Ok(()); - } - - 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 - ); - let new_enrollment = match enroll_remote_control_server( - remote_control_target, - auth, - connect_options.installation_id, - connect_options.server_name, - ) - .await - { - Ok(new_enrollment) => new_enrollment, - Err(err) - if err.kind() == ErrorKind::PermissionDenied - && recover_remote_control_auth( - auth_context.auth_recovery, - auth_context.auth_change_rx, - ) - .await => - { - return Err(io::Error::other(format!( - "{err}; retrying after auth recovery" - ))); - } - Err(err) => return Err(err), - }; - if let Err(err) = update_persisted_remote_control_enrollment( - Some(state_db), - remote_control_target, - &auth.account_id, - connect_options.app_server_client_name, - Some(&new_enrollment), - ) - .await - { - return Err(io::Error::other(format!( - "failed to persist remote control enrollment in sqlite state db: {err}" - ))); - } - info!( - "created new remote control enrollment: websocket_url={}, account_id={}, server_id={}, environment_id={}", - remote_control_target.websocket_url, - new_enrollment.account_id, - new_enrollment.server_id, - new_enrollment.environment_id - ); - status_publisher.publish_environment_id(Some(new_enrollment.environment_id.clone())); - *enrollment = Some(new_enrollment); - Ok(()) -} - async fn recover_remote_control_auth( auth_recovery: &mut UnauthorizedRecovery, auth_change_rx: &mut watch::Receiver, @@ -1645,8 +1394,8 @@ mod tests { const TEST_INSTALLATION_ID: &str = "11111111-1111-4111-8111-111111111111"; const TEST_REMOTE_CONTROL_SERVER_TOKEN: &str = "Remote Control Token"; - fn remote_control_enrollment(remote_control_token: Option<&str>) -> RemoteControlEnrollment { - RemoteControlEnrollment { + fn remote_control_server_state(remote_control_token: Option<&str>) -> RemoteControlServer { + RemoteControlServer { account_id: "account_id".to_string(), environment_id: "env_test".to_string(), server_id: "srv_e_test".to_string(), @@ -1731,6 +1480,23 @@ mod tests { auth_manager_from_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) } + async fn shared_remote_control_server( + remote_control_url: String, + state_db: Option>, + auth_manager: Arc, + server: Option, + ) -> SharedRemoteControlServer { + let shared_server = SharedRemoteControlServer::new( + remote_control_url, + TEST_INSTALLATION_ID.to_string(), + "test-server".to_string(), + state_db, + auth_manager, + ); + shared_server.replace(server).await; + shared_server + } + fn remote_control_url_for_listener(listener: &TcpListener) -> String { let addr = listener .local_addr() @@ -1807,25 +1573,28 @@ mod tests { let auth_manager = remote_control_auth_manager(); let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut auth_change_rx = auth_manager.auth_change_receiver(); - let mut enrollment = Some(remote_control_enrollment(Some( - TEST_REMOTE_CONTROL_SERVER_TOKEN, - ))); + let server = shared_remote_control_server( + remote_control_url, + Some(state_db.clone()), + auth_manager.clone(), + Some(remote_control_server_state(Some( + TEST_REMOTE_CONTROL_SERVER_TOKEN, + ))), + ) + .await; let (status_publisher, status_rx) = remote_control_status_channel(); let err = match connect_remote_control_websocket( &remote_control_target, - Some(state_db.as_ref()), RemoteControlAuthContext { auth_manager: &auth_manager, auth_recovery: &mut auth_recovery, auth_change_rx: &mut auth_change_rx, }, - &mut enrollment, + &server, RemoteControlConnectOptions { installation_id: TEST_INSTALLATION_ID, - server_name: "test-server", subscribe_cursor: None, - app_server_client_name: None, }, &status_publisher, ) @@ -1861,9 +1630,15 @@ mod tests { let auth_manager = remote_control_auth_manager(); let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut auth_change_rx = auth_manager.auth_change_receiver(); - let mut enrollment = Some(remote_control_enrollment(Some( - TEST_REMOTE_CONTROL_SERVER_TOKEN, - ))); + let server = shared_remote_control_server( + remote_control_url, + Some(state_db.clone()), + auth_manager.clone(), + Some(remote_control_server_state(Some( + TEST_REMOTE_CONTROL_SERVER_TOKEN, + ))), + ) + .await; let (status_publisher, status_rx) = remote_control_status_channel(); let server_task = tokio::spawn(async move { @@ -1877,18 +1652,15 @@ mod tests { let err = connect_remote_control_websocket( &remote_control_target, - Some(state_db.as_ref()), RemoteControlAuthContext { auth_manager: &auth_manager, auth_recovery: &mut auth_recovery, auth_change_rx: &mut auth_change_rx, }, - &mut enrollment, + &server, RemoteControlConnectOptions { installation_id: TEST_INSTALLATION_ID, - server_name: "test-server", subscribe_cursor: None, - app_server_client_name: None, }, &status_publisher, ) @@ -1910,8 +1682,8 @@ mod tests { "remote control websocket auth failed with HTTP 401 Unauthorized; refreshing server token before reconnect" ); assert_eq!( - enrollment, - Some(remote_control_enrollment( + server.snapshot().await, + Some(remote_control_server_state( /*remote_control_token*/ None )) ); @@ -1951,7 +1723,13 @@ mod tests { .await; let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut auth_change_rx = auth_manager.auth_change_receiver(); - let mut enrollment = None; + let server = shared_remote_control_server( + remote_control_url, + Some(state_db.clone()), + auth_manager.clone(), + None, + ) + .await; let (status_publisher, status_rx) = remote_control_status_channel(); save_auth( codex_home.path(), @@ -1962,18 +1740,15 @@ mod tests { let err = connect_remote_control_websocket( &remote_control_target, - Some(state_db.as_ref()), RemoteControlAuthContext { auth_manager: &auth_manager, auth_recovery: &mut auth_recovery, auth_change_rx: &mut auth_change_rx, }, - &mut enrollment, + &server, RemoteControlConnectOptions { installation_id: TEST_INSTALLATION_ID, - server_name: "test-server", subscribe_cursor: None, - app_server_client_name: None, }, &status_publisher, ) @@ -2043,9 +1818,15 @@ mod tests { .await; let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut auth_change_rx = auth_manager.auth_change_receiver(); - let mut enrollment = Some(remote_control_enrollment( - /*remote_control_token*/ None, - )); + let server = shared_remote_control_server( + remote_control_url, + Some(state_db.clone()), + auth_manager.clone(), + Some(remote_control_server_state( + /*remote_control_token*/ None, + )), + ) + .await; let (status_publisher, status_rx) = remote_control_status_channel(); save_auth( codex_home.path(), @@ -2056,18 +1837,15 @@ mod tests { let err = connect_remote_control_websocket( &remote_control_target, - Some(state_db.as_ref()), RemoteControlAuthContext { auth_manager: &auth_manager, auth_recovery: &mut auth_recovery, auth_change_rx: &mut auth_change_rx, }, - &mut enrollment, + &server, RemoteControlConnectOptions { installation_id: TEST_INSTALLATION_ID, - server_name: "test-server", subscribe_cursor: None, - app_server_client_name: None, }, &status_publisher, ) @@ -2114,25 +1892,28 @@ mod tests { let auth_manager = remote_control_auth_manager(); let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut auth_change_rx = auth_manager.auth_change_receiver(); - let mut enrollment = Some(remote_control_enrollment(Some( - TEST_REMOTE_CONTROL_SERVER_TOKEN, - ))); + let server = shared_remote_control_server( + "http://127.0.0.1:9/backend-api/".to_string(), + /*state_db*/ None, + auth_manager.clone(), + Some(remote_control_server_state(Some( + TEST_REMOTE_CONTROL_SERVER_TOKEN, + ))), + ) + .await; let (status_publisher, _status_rx) = remote_control_status_channel(); let err = connect_remote_control_websocket( &remote_control_target, - /*state_db*/ None, RemoteControlAuthContext { auth_manager: &auth_manager, auth_recovery: &mut auth_recovery, auth_change_rx: &mut auth_change_rx, }, - &mut enrollment, + &server, RemoteControlConnectOptions { installation_id: TEST_INSTALLATION_ID, - server_name: "test-server", subscribe_cursor: None, - app_server_client_name: None, }, &status_publisher, ) @@ -2141,7 +1922,7 @@ mod tests { assert_eq!(err.kind(), ErrorKind::NotFound); assert_eq!(err.to_string(), "remote control requires sqlite state db"); - assert_eq!(enrollment, None); + assert_eq!(server.snapshot().await, None); } #[tokio::test] @@ -2159,9 +1940,15 @@ mod tests { .await; let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut auth_change_rx = auth_manager.auth_change_receiver(); - let mut enrollment = Some(remote_control_enrollment(Some( - TEST_REMOTE_CONTROL_SERVER_TOKEN, - ))); + let server = shared_remote_control_server( + "http://127.0.0.1:9/backend-api/".to_string(), + Some(state_db.clone()), + auth_manager.clone(), + Some(remote_control_server_state(Some( + TEST_REMOTE_CONTROL_SERVER_TOKEN, + ))), + ) + .await; let (status_publisher, mut status_rx) = remote_control_status_channel(); status_publisher.publish_environment_id(Some("env_test".to_string())); status_rx @@ -2171,18 +1958,15 @@ mod tests { let err = connect_remote_control_websocket( &remote_control_target, - Some(state_db.as_ref()), RemoteControlAuthContext { auth_manager: &auth_manager, auth_recovery: &mut auth_recovery, auth_change_rx: &mut auth_change_rx, }, - &mut enrollment, + &server, RemoteControlConnectOptions { installation_id: TEST_INSTALLATION_ID, - server_name: "test-server", subscribe_cursor: None, - app_server_client_name: None, }, &status_publisher, ) @@ -2194,7 +1978,7 @@ mod tests { err.to_string(), "remote control requires ChatGPT authentication" ); - assert_eq!(enrollment, None); + assert_eq!(server.snapshot().await, None); assert_eq!( status_rx.borrow().clone(), RemoteControlStatusChangedNotification { @@ -2221,6 +2005,14 @@ mod tests { let (status_publisher, _status_rx) = remote_control_status_channel(); let shutdown_token = CancellationToken::new(); let (_enabled_tx, enabled_rx) = watch::channel(true); + let auth_manager = remote_control_auth_manager(); + let server = SharedRemoteControlServer::new( + remote_control_url.clone(), + TEST_INSTALLATION_ID.to_string(), + "test-server".to_string(), + /*state_db*/ None, + auth_manager.clone(), + ); let websocket_task = tokio::spawn({ let shutdown_token = shutdown_token.clone(); async move { @@ -2231,12 +2023,12 @@ mod tests { remote_control_target: Some(remote_control_target), server_name: "test-server".to_string(), }, - /*state_db*/ None, - remote_control_auth_manager(), + auth_manager, RemoteControlChannels { transport_event_tx, status_publisher, }, + server, shutdown_token, enabled_rx, ) diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index ad3767bb77..2ee01f2035 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -209,6 +209,7 @@ Example with notification opt-out: - `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/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/status/changed` — notification emitted when the remote-control status or client-visible environment id changes. `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. Newly initialized app-server clients always receive the current status snapshot. - `skills/config/write` — write user-level skill config by name or absolute path. - `plugin/install` — install a plugin from a discovered marketplace entry, rejecting marketplace entries marked unavailable for install, install MCPs if any, and return the effective plugin auth policy plus any apps that still need auth (**under development; do not call from production clients yet**). diff --git a/codex-rs/app-server/src/message_processor.rs b/codex-rs/app-server/src/message_processor.rs index 33b9c56d4a..57a3e1d5f5 100644 --- a/codex-rs/app-server/src/message_processor.rs +++ b/codex-rs/app-server/src/message_processor.rs @@ -813,7 +813,6 @@ impl MessageProcessor { connection_request_id.request_id.clone(), &codex_request, ); - let serialization_scope = codex_request.serialization_scope(); let app_server_client_name = session.app_server_client_name().map(str::to_string); let client_version = session.client_version().map(str::to_string); @@ -915,6 +914,11 @@ impl MessageProcessor { .remote_control_processor .status_read() .map(|response| Some(response.into())), + ClientRequest::RemoteControlPairingStart { params, .. } => self + .remote_control_processor + .pairing_start(params) + .await + .map(|response| Some(response.into())), ClientRequest::ConfigRequirementsRead { params: _, .. } => self .config_processor .config_requirements_read() diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 8408a3f67f..e0e1189456 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -470,6 +470,8 @@ mod mcp_processor; mod plugins; mod process_exec_processor; mod remote_control_processor; +#[cfg(test)] +mod remote_control_processor_tests; mod search; mod thread_processor; mod token_usage_replay; 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 2fceeb5030..18fb9d2511 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 @@ -5,7 +5,10 @@ use crate::transport::RemoteControlUnavailable; use codex_app_server_protocol::JSONRPCErrorError; use codex_app_server_protocol::RemoteControlDisableResponse; use codex_app_server_protocol::RemoteControlEnableResponse; +use codex_app_server_protocol::RemoteControlPairingStartParams; +use codex_app_server_protocol::RemoteControlPairingStartResponse; use codex_app_server_protocol::RemoteControlStatusReadResponse; +use std::io; #[derive(Clone)] pub(crate) struct RemoteControlRequestProcessor { @@ -42,6 +45,16 @@ impl RemoteControlRequestProcessor { }) } + pub(crate) async fn pairing_start( + &self, + params: RemoteControlPairingStartParams, + ) -> Result { + self.handle()? + .start_pairing(params) + .await + .map_err(map_pairing_start_error) + } + fn handle(&self) -> Result<&RemoteControlHandle, JSONRPCErrorError> { self.remote_control_handle .as_ref() @@ -52,3 +65,11 @@ impl RemoteControlRequestProcessor { fn map_unavailable(err: RemoteControlUnavailable) -> JSONRPCErrorError { invalid_request(err.to_string()) } + +pub(super) fn map_pairing_start_error(err: io::Error) -> JSONRPCErrorError { + if err.kind() == io::ErrorKind::InvalidInput { + invalid_request(err.to_string()) + } else { + internal_error(err.to_string()) + } +} diff --git a/codex-rs/app-server/src/request_processors/remote_control_processor_tests.rs b/codex-rs/app-server/src/request_processors/remote_control_processor_tests.rs new file mode 100644 index 0000000000..48ddb15923 --- /dev/null +++ b/codex-rs/app-server/src/request_processors/remote_control_processor_tests.rs @@ -0,0 +1,52 @@ +use super::remote_control_processor::RemoteControlRequestProcessor; +use super::remote_control_processor::map_pairing_start_error; +use crate::error_code::INTERNAL_ERROR_CODE; +use crate::error_code::INVALID_REQUEST_ERROR_CODE; +use codex_app_server_protocol::JSONRPCErrorError; +use codex_app_server_protocol::RemoteControlPairingStartParams; +use pretty_assertions::assert_eq; +use std::io; + +#[tokio::test] +async fn pairing_start_returns_internal_error_when_remote_control_is_unavailable() { + let err = RemoteControlRequestProcessor::new(/*remote_control_handle*/ None) + .pairing_start(RemoteControlPairingStartParams::default()) + .await + .expect_err("missing remote control should fail pairing"); + + assert_eq!( + err, + JSONRPCErrorError { + code: INTERNAL_ERROR_CODE, + data: None, + message: "remote control is unavailable for this app-server".to_string(), + } + ); +} + +#[test] +fn pairing_start_maps_invalid_input_to_invalid_request() { + assert_eq!( + map_pairing_start_error(io::Error::new( + io::ErrorKind::InvalidInput, + "remote control pairing is unavailable", + )), + JSONRPCErrorError { + code: INVALID_REQUEST_ERROR_CODE, + data: None, + message: "remote control pairing is unavailable".to_string(), + } + ); +} + +#[test] +fn pairing_start_maps_backend_failures_to_internal_error() { + assert_eq!( + map_pairing_start_error(io::Error::other("remote control pairing failed")), + JSONRPCErrorError { + code: INTERNAL_ERROR_CODE, + data: None, + message: "remote control pairing failed".to_string(), + } + ); +} diff --git a/codex-rs/app-server/tests/common/mcp_process.rs b/codex-rs/app-server/tests/common/mcp_process.rs index 2d2c26b5a6..e862c1db4d 100644 --- a/codex-rs/app-server/tests/common/mcp_process.rs +++ b/codex-rs/app-server/tests/common/mcp_process.rs @@ -636,6 +636,16 @@ impl McpProcess { .await } + /// Send a `remoteControl/pairing/start` JSON-RPC request. + pub async fn send_remote_control_pairing_start_request( + &mut self, + params: codex_app_server_protocol::RemoteControlPairingStartParams, + ) -> anyhow::Result { + let params = Some(serde_json::to_value(params)?); + self.send_request("remoteControl/pairing/start", params) + .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 d71f05bf58..8c00db695c 100644 --- a/codex-rs/app-server/tests/suite/v2/remote_control.rs +++ b/codex-rs/app-server/tests/suite/v2/remote_control.rs @@ -1,3 +1,4 @@ +use std::path::Path; use std::time::Duration; use anyhow::Context; @@ -11,18 +12,22 @@ use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RemoteControlConnectionStatus; use codex_app_server_protocol::RemoteControlDisableResponse; use codex_app_server_protocol::RemoteControlEnableResponse; +use codex_app_server_protocol::RemoteControlPairingStartParams; +use codex_app_server_protocol::RemoteControlPairingStartResponse; use codex_app_server_protocol::RemoteControlStatusReadResponse; use codex_app_server_protocol::RequestId; use codex_config::types::AuthCredentialsStoreMode; use pretty_assertions::assert_eq; use tempfile::TempDir; use tokio::io::AsyncBufReadExt; +use tokio::io::AsyncWriteExt; use tokio::io::BufReader; use tokio::net::TcpListener; use tokio::net::TcpStream; use tokio::sync::oneshot; use tokio::task::JoinHandle; use tokio::time::timeout; +use tokio_tungstenite::accept_hdr_async; const DEFAULT_TIMEOUT: Duration = Duration::from_secs(10); @@ -125,6 +130,192 @@ async fn remote_control_status_read_returns_connecting_status_after_enable() -> Ok(()) } +#[tokio::test] +async fn remote_control_pairing_start_returns_pairing_artifacts() -> Result<()> { + let codex_home = TempDir::new()?; + let mut backend = PairingRemoteControlBackend::start(codex_home.path()).await?; + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_TIMEOUT, mcp.initialize()).await??; + enable_connected_remote_control(&mut mcp, &mut backend).await?; + + let request_id = mcp + .send_remote_control_pairing_start_request(RemoteControlPairingStartParams { + manual_code: true, + }) + .await?; + let response: JSONRPCResponse = timeout( + DEFAULT_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(request_id)), + ) + .await??; + let received: RemoteControlPairingStartResponse = to_response(response)?; + + assert_eq!( + received, + RemoteControlPairingStartResponse { + pairing_code: "pairing-code".to_string(), + manual_pairing_code: Some("ABCD-EFGH".to_string()), + environment_id: "environment-id".to_string(), + expires_at: 33_336_362_096, + } + ); + Ok(()) +} + +struct PairingRemoteControlBackend { + enroll_request_rx: Option>>, + websocket_request_rx: Option>>, + server_task: JoinHandle<()>, +} + +async fn enable_connected_remote_control( + mcp: &mut McpProcess, + backend: &mut PairingRemoteControlBackend, +) -> Result<()> { + 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??; + assert_eq!( + timeout(DEFAULT_TIMEOUT, backend.wait_for_enroll_request()).await??, + "POST /backend-api/wham/remote/control/server/enroll HTTP/1.1" + ); + assert_eq!( + timeout(DEFAULT_TIMEOUT, backend.wait_for_websocket_request()).await??, + "GET /backend-api/wham/remote/control/server HTTP/1.1" + ); + timeout(DEFAULT_TIMEOUT, async { + loop { + let request_id = mcp.send_remote_control_status_read_request().await?; + let response: JSONRPCResponse = mcp + .read_stream_until_response_message(RequestId::Integer(request_id)) + .await?; + let received: RemoteControlStatusReadResponse = to_response(response)?; + if received.status == RemoteControlConnectionStatus::Connected { + return Ok::<(), anyhow::Error>(()); + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await? +} + +impl PairingRemoteControlBackend { + async fn start(codex_home: &Path) -> Result { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let remote_control_url = format!("http://{}/backend-api/", listener.local_addr()?); + write_mock_responses_config_toml_with_chatgpt_base_url( + codex_home, + &remote_control_url, + &remote_control_url, + )?; + write_chatgpt_auth( + codex_home, + ChatGptAuthFixture::new("chatgpt-token") + .account_id("account_id") + .chatgpt_account_id("account_id"), + AuthCredentialsStoreMode::File, + )?; + + let (enroll_request_tx, enroll_request_rx) = oneshot::channel(); + let (websocket_request_tx, websocket_request_rx) = oneshot::channel(); + let server_task = tokio::spawn(async move { + let mut enroll_request_tx = Some(enroll_request_tx); + let mut websocket_request_tx = Some(websocket_request_tx); + let result = async { + let (enroll_request_line, enroll_reader) = read_enroll_request(&listener).await?; + if let Some(enroll_request_tx) = enroll_request_tx.take() { + let _ = enroll_request_tx.send(Ok(enroll_request_line)); + } + respond_with_json( + enroll_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", + "scopes": ["remote_control_server_websocket"], + }), + ) + .await?; + + let (websocket_stream, _) = listener.accept().await?; + let websocket_request_tx = websocket_request_tx.take(); + let mut websocket = accept_hdr_async( + websocket_stream, + move |request: &tokio_tungstenite::tungstenite::handshake::server::Request, + response: tokio_tungstenite::tungstenite::handshake::server::Response| { + let method = request.method(); + let uri = request.uri(); + let request_line = format!("{method} {uri} HTTP/1.1"); + if let Some(websocket_request_tx) = websocket_request_tx { + let _ = websocket_request_tx.send(Ok(request_line)); + } + Ok(response) + }, + ) + .await?; + + let (request_line, request_reader) = read_enroll_request(&listener).await?; + if request_line != "POST /backend-api/wham/remote/control/server/pair HTTP/1.1" { + anyhow::bail!("unexpected remote control request: {request_line}"); + } + respond_with_json( + request_reader.into_inner(), + serde_json::json!({ + "pairing_code": "pairing-code", + "manual_pairing_code": "ABCD-EFGH", + "server_id": "server-id", + "environment_id": "environment-id", + "expires_at": "3026-05-22T12:34:56Z", + }), + ) + .await?; + websocket.close(None).await?; + Ok::<(), anyhow::Error>(()) + } + .await; + + if let Err(err) = result { + if let Some(enroll_request_tx) = enroll_request_tx { + let _ = enroll_request_tx.send(Err(anyhow::anyhow!(err.to_string()))); + } + if let Some(websocket_request_tx) = websocket_request_tx { + let _ = websocket_request_tx.send(Err(anyhow::anyhow!(err.to_string()))); + } + } + }); + + Ok(Self { + enroll_request_rx: Some(enroll_request_rx), + websocket_request_rx: Some(websocket_request_rx), + server_task, + }) + } + + async fn wait_for_enroll_request(&mut self) -> Result { + self.enroll_request_rx + .take() + .context("enroll request should only be awaited once")? + .await? + } + + async fn wait_for_websocket_request(&mut self) -> Result { + self.websocket_request_rx + .take() + .context("websocket request should only be awaited once")? + .await? + } +} + +impl Drop for PairingRemoteControlBackend { + fn drop(&mut self) { + self.server_task.abort(); + } +} + struct BlockingRemoteControlBackend { enroll_request_rx: Option>>, server_task: JoinHandle<()>, @@ -149,7 +340,7 @@ impl BlockingRemoteControlBackend { let (enroll_request_tx, enroll_request_rx) = oneshot::channel(); let server_task = tokio::spawn(async move { - match read_enroll_request(listener).await { + match read_enroll_request(&listener).await { Ok((request_line, _reader)) => { let _ = enroll_request_tx.send(Ok(request_line)); std::future::pending::<()>().await; @@ -181,7 +372,7 @@ impl Drop for BlockingRemoteControlBackend { } } -async fn read_enroll_request(listener: TcpListener) -> Result<(String, BufReader)> { +async fn read_enroll_request(listener: &TcpListener) -> Result<(String, BufReader)> { let (stream, _) = listener.accept().await?; let mut reader = BufReader::new(stream); @@ -198,3 +389,25 @@ async fn read_enroll_request(listener: TcpListener) -> Result<(String, BufReader Ok((request_line.trim_end().to_string(), reader)) } + +async fn respond_with_json(stream: TcpStream, body: serde_json::Value) -> Result<()> { + let body = body.to_string(); + respond_with_status_and_body(stream, "200 OK", &body).await +} + +async fn respond_with_status_and_body( + mut stream: TcpStream, + status: &str, + body: &str, +) -> Result<()> { + stream + .write_all( + format!( + "HTTP/1.1 {status}\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ) + .as_bytes(), + ) + .await?; + Ok(()) +}