codex/codex-rs/app-server-daemon/src/remote_control_client.rs

use std::path::Path;
use std::time::Duration;

use anyhow::Context;
use anyhow::Result;
use anyhow::anyhow;
use codex_app_server_protocol::JSONRPCMessage;
use codex_app_server_protocol::JSONRPCNotification;
use codex_app_server_protocol::JSONRPCRequest;
use codex_app_server_protocol::RemoteControlConnectionStatus;
use codex_app_server_protocol::RemoteControlDisableParams;
use codex_app_server_protocol::RemoteControlDisableResponse;
use codex_app_server_protocol::RemoteControlEnableParams;
use codex_app_server_protocol::RemoteControlEnableResponse;
use codex_app_server_protocol::RemoteControlStatusChangedNotification;
use codex_app_server_protocol::RequestId;
use serde::de::DeserializeOwned;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tokio::time::Instant;
use tokio::time::sleep;
use tokio::time::timeout;
use tokio_tungstenite::WebSocketStream;

use crate::RemoteControlReadyStatus;
use crate::client;

const REMOTE_CONTROL_READY_TIMEOUT: Duration = Duration::from_secs(10);
const REMOTE_CONTROL_REQUEST_ID: RequestId = RequestId::Integer(2);
const INVALID_PARAMS_ERROR_CODE: i64 = -32602;

enum RemoteControlRpcResponse<T> {
    Success(T),
    InvalidParams,
}

pub(crate) async fn enable_remote_control(socket_path: &Path) -> Result<RemoteControlReadyStatus> {
    let mut websocket = client::connect(socket_path).await?;
    enable_remote_control_with_timeout(&mut websocket, REMOTE_CONTROL_READY_TIMEOUT).await
}

pub(crate) async fn disable_remote_control(socket_path: &Path) -> Result<RemoteControlReadyStatus> {
    let mut websocket = client::connect(socket_path).await?;
    initialize_client(&mut websocket).await?;
    let params = serde_json::to_value(RemoteControlDisableParams { ephemeral: true })?;
    let response: RemoteControlDisableResponse = request_remote_control_with_legacy_fallback(
        &mut websocket,
        "remoteControl/disable",
        params,
    )
    .await?;
    websocket.close(None).await.ok();
    Ok(RemoteControlReadyStatus::from(response))
}

pub(crate) async fn enable_remote_control_with_connect_retry(
    socket_path: &Path,
    connect_timeout: Duration,
    connect_retry_delay: Duration,
) -> Result<RemoteControlReadyStatus> {
    let mut websocket =
        connect_with_retry(socket_path, connect_timeout, connect_retry_delay).await?;
    enable_remote_control_with_timeout(&mut websocket, REMOTE_CONTROL_READY_TIMEOUT).await
}

async fn enable_remote_control_with_timeout<S>(
    websocket: &mut WebSocketStream<S>,
    ready_timeout: Duration,
) -> Result<RemoteControlReadyStatus>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    initialize_client(websocket).await?;

    let response: RemoteControlEnableResponse = request_remote_control_with_legacy_fallback(
        websocket,
        "remoteControl/enable",
        serde_json::to_value(RemoteControlEnableParams { ephemeral: true })?,
    )
    .await?;
    let mut latest = RemoteControlReadyStatus::from(response);
    if latest.status == RemoteControlConnectionStatus::Connecting {
        latest = wait_for_remote_control_status(websocket, latest, ready_timeout).await?;
    }
    websocket.close(None).await.ok();
    Ok(latest)
}

async fn initialize_client<S>(websocket: &mut WebSocketStream<S>) -> Result<()>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    client::initialize(websocket, /*experimental_api*/ true).await?;
    let initialized = JSONRPCMessage::Notification(JSONRPCNotification {
        method: "initialized".to_string(),
        params: None,
    });
    client::send_message(websocket, &initialized)
        .await
        .context("failed to send initialized notification")
}

async fn send_remote_control_request<S>(
    websocket: &mut WebSocketStream<S>,
    request_id: RequestId,
    method: &str,
    params: Option<serde_json::Value>,
) -> Result<()>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    let request = JSONRPCMessage::Request(JSONRPCRequest {
        id: request_id,
        method: method.to_string(),
        params,
        trace: None,
    });
    client::send_message(websocket, &request)
        .await
        .with_context(|| format!("failed to send {method} request"))
}

async fn request_remote_control_with_legacy_fallback<S, T>(
    websocket: &mut WebSocketStream<S>,
    method: &str,
    params: serde_json::Value,
) -> Result<T>
where
    S: AsyncRead + AsyncWrite + Unpin,
    T: DeserializeOwned,
{
    send_remote_control_request(
        websocket,
        REMOTE_CONTROL_REQUEST_ID.clone(),
        method,
        Some(params),
    )
    .await?;
    match read_remote_control_response(websocket, &REMOTE_CONTROL_REQUEST_ID, method).await? {
        RemoteControlRpcResponse::Success(response) => Ok(response),
        RemoteControlRpcResponse::InvalidParams => {
            send_remote_control_request(
                websocket,
                REMOTE_CONTROL_REQUEST_ID.clone(),
                method,
                /*params*/ None,
            )
            .await?;
            match read_remote_control_response(websocket, &REMOTE_CONTROL_REQUEST_ID, method)
                .await?
            {
                RemoteControlRpcResponse::Success(response) => Ok(response),
                RemoteControlRpcResponse::InvalidParams => {
                    Err(anyhow!("{method} rejected legacy params"))
                }
            }
        }
    }
}

async fn connect_with_retry(
    socket_path: &Path,
    connect_timeout: Duration,
    connect_retry_delay: Duration,
) -> Result<WebSocketStream<codex_uds::UnixStream>> {
    let deadline = Instant::now() + connect_timeout;
    loop {
        match client::connect(socket_path).await {
            Ok(websocket) => return Ok(websocket),
            Err(_) if Instant::now() < deadline => {
                sleep(connect_retry_delay).await;
            }
            Err(error) => {
                return Err(error).with_context(|| {
                    format!(
                        "app server did not become ready on {}",
                        socket_path.display()
                    )
                });
            }
        }
    }
}

async fn read_remote_control_response<S, T>(
    websocket: &mut WebSocketStream<S>,
    request_id: &RequestId,
    method: &str,
) -> Result<RemoteControlRpcResponse<T>>
where
    S: AsyncRead + AsyncWrite + Unpin,
    T: DeserializeOwned,
{
    loop {
        let message = timeout(
            client::CONTROL_SOCKET_RESPONSE_TIMEOUT,
            client::read_message(websocket),
        )
        .await
        .with_context(|| format!("timed out waiting for {method} response"))??;
        match message {
            JSONRPCMessage::Response(response) if response.id == *request_id => {
                let response = serde_json::from_value::<T>(response.result)
                    .with_context(|| format!("failed to parse {method} response"))?;
                return Ok(RemoteControlRpcResponse::Success(response));
            }
            JSONRPCMessage::Error(err)
                if err.id == *request_id && err.error.code == INVALID_PARAMS_ERROR_CODE =>
            {
                return Ok(RemoteControlRpcResponse::InvalidParams);
            }
            JSONRPCMessage::Error(err) if err.id == *request_id => {
                return Err(anyhow!("{method} failed: {}", err.error.message));
            }
            JSONRPCMessage::Notification(notification)
                if remote_control_status_notification(&notification).is_some() =>
            {
                continue;
            }
            _ => {}
        }
    }
}

async fn wait_for_remote_control_status<S>(
    websocket: &mut WebSocketStream<S>,
    mut latest: RemoteControlReadyStatus,
    ready_timeout: Duration,
) -> Result<RemoteControlReadyStatus>
where
    S: AsyncRead + AsyncWrite + Unpin,
{
    let deadline = tokio::time::Instant::now() + ready_timeout;
    while tokio::time::Instant::now() < deadline {
        let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
        let message = match timeout(remaining, client::read_message(websocket)).await {
            Ok(Ok(message)) => message,
            Ok(Err(err)) => return Err(err),
            Err(_) => {
                latest.timed_out = true;
                return Ok(latest);
            }
        };
        let JSONRPCMessage::Notification(notification) = message else {
            continue;
        };
        let Some(status) = remote_control_status_notification(&notification) else {
            continue;
        };
        latest = RemoteControlReadyStatus::from(status);
        if latest.status != RemoteControlConnectionStatus::Connecting {
            return Ok(latest);
        }
    }
    latest.timed_out = true;
    Ok(latest)
}

fn remote_control_status_notification(
    notification: &JSONRPCNotification,
) -> Option<RemoteControlStatusChangedNotification> {
    if notification.method != "remoteControl/status/changed" {
        return None;
    }
    let params = notification.params.clone()?;
    serde_json::from_value(params).ok()
}

impl From<RemoteControlEnableResponse> for RemoteControlReadyStatus {
    fn from(response: RemoteControlEnableResponse) -> Self {
        let RemoteControlEnableResponse {
            status,
            server_name,
            installation_id: _,
            environment_id,
        } = response;
        Self {
            status,
            server_name,
            environment_id,
            timed_out: false,
        }
    }
}

impl From<RemoteControlDisableResponse> for RemoteControlReadyStatus {
    fn from(response: RemoteControlDisableResponse) -> Self {
        let RemoteControlDisableResponse {
            status,
            server_name,
            installation_id: _,
            environment_id,
        } = response;
        Self {
            status,
            server_name,
            environment_id,
            timed_out: false,
        }
    }
}

impl From<RemoteControlStatusChangedNotification> for RemoteControlReadyStatus {
    fn from(notification: RemoteControlStatusChangedNotification) -> Self {
        let RemoteControlStatusChangedNotification {
            status,
            server_name,
            installation_id: _,
            environment_id,
        } = notification;
        Self {
            status,
            server_name,
            environment_id,
            timed_out: false,
        }
    }
}

#[cfg(all(test, unix))]
mod tests {
    use anyhow::Result;
    use codex_app_server_protocol::JSONRPCError;
    use codex_app_server_protocol::JSONRPCErrorError;
    use codex_app_server_protocol::JSONRPCResponse;
    use codex_uds::UnixListener;
    use pretty_assertions::assert_eq;
    use tempfile::TempDir;
    use tokio_tungstenite::accept_async;

    use super::*;

    const INITIALIZE_REQUEST_ID: RequestId = RequestId::Integer(1);
    const TEST_INSTALLATION_ID: &str = "11111111-1111-4111-8111-111111111111";
    const TEST_SERVER_NAME: &str = "owen-mbp";
    const TEST_CODEX_HOME: &str = "/tmp/codex-home";

    #[tokio::test]
    async fn enable_remote_control_uses_connected_enable_response_without_later_notification()
    -> Result<()> {
        let status = run_enable_remote_control_scenario(EnableScenario {
            initial_notification: Some(remote_control_status(
                RemoteControlConnectionStatus::Connected,
                Some("env_test"),
            )),
            enable_response: remote_control_status(
                RemoteControlConnectionStatus::Connected,
                Some("env_test"),
            ),
            after_enable_notification: None,
            ready_timeout: Duration::from_millis(20),
            reject_ephemeral_params: false,
        })
        .await?;

        assert_eq!(
            status,
            RemoteControlReadyStatus {
                status: RemoteControlConnectionStatus::Connected,
                server_name: TEST_SERVER_NAME.to_string(),
                environment_id: Some("env_test".to_string()),
                timed_out: false,
            }
        );
        Ok(())
    }

    #[tokio::test]
    async fn enable_remote_control_waits_for_connected_notification() -> Result<()> {
        let status = run_enable_remote_control_scenario(EnableScenario {
            initial_notification: None,
            enable_response: remote_control_status(
                RemoteControlConnectionStatus::Connecting,
                /*environment_id*/ None,
            ),
            after_enable_notification: Some(remote_control_status(
                RemoteControlConnectionStatus::Connected,
                Some("env_test"),
            )),
            ready_timeout: Duration::from_secs(1),
            reject_ephemeral_params: false,
        })
        .await?;

        assert_eq!(
            status,
            RemoteControlReadyStatus {
                status: RemoteControlConnectionStatus::Connected,
                server_name: TEST_SERVER_NAME.to_string(),
                environment_id: Some("env_test".to_string()),
                timed_out: false,
            }
        );
        Ok(())
    }

    #[tokio::test]
    async fn enable_remote_control_reports_connecting_after_timeout() -> Result<()> {
        let status = run_enable_remote_control_scenario(EnableScenario {
            initial_notification: None,
            enable_response: remote_control_status(
                RemoteControlConnectionStatus::Connecting,
                /*environment_id*/ None,
            ),
            after_enable_notification: None,
            ready_timeout: Duration::from_millis(20),
            reject_ephemeral_params: false,
        })
        .await?;

        assert_eq!(
            status,
            RemoteControlReadyStatus {
                status: RemoteControlConnectionStatus::Connecting,
                server_name: TEST_SERVER_NAME.to_string(),
                environment_id: None,
                timed_out: true,
            }
        );
        Ok(())
    }

    #[tokio::test]
    async fn enable_remote_control_returns_errored_enable_response() -> Result<()> {
        let status = run_enable_remote_control_scenario(EnableScenario {
            initial_notification: None,
            enable_response: remote_control_status(
                RemoteControlConnectionStatus::Errored,
                /*environment_id*/ None,
            ),
            after_enable_notification: None,
            ready_timeout: Duration::from_millis(20),
            reject_ephemeral_params: false,
        })
        .await?;

        assert_eq!(
            status,
            RemoteControlReadyStatus {
                status: RemoteControlConnectionStatus::Errored,
                server_name: TEST_SERVER_NAME.to_string(),
                environment_id: None,
                timed_out: false,
            }
        );
        Ok(())
    }

    #[tokio::test]
    async fn enable_remote_control_retries_without_params_for_older_servers() -> Result<()> {
        let status = run_enable_remote_control_scenario(EnableScenario {
            initial_notification: None,
            enable_response: remote_control_status(
                RemoteControlConnectionStatus::Connected,
                Some("env_test"),
            ),
            after_enable_notification: None,
            ready_timeout: Duration::from_millis(20),
            reject_ephemeral_params: true,
        })
        .await?;

        assert_eq!(
            status,
            RemoteControlReadyStatus {
                status: RemoteControlConnectionStatus::Connected,
                server_name: TEST_SERVER_NAME.to_string(),
                environment_id: Some("env_test".to_string()),
                timed_out: false,
            }
        );
        Ok(())
    }

    #[tokio::test]
    async fn disable_remote_control_retries_without_params_for_older_servers() -> Result<()> {
        let dir = TempDir::new()?;
        let socket_path = dir.path().join("app-server.sock");
        let listener = UnixListener::bind(&socket_path).await?;
        let server_task = tokio::spawn(async move {
            let mut websocket = accept_initialized_client(listener).await?;
            let disable = client::read_message(&mut websocket).await?;
            let JSONRPCMessage::Request(disable) = disable else {
                panic!("expected remoteControl/disable request");
            };
            assert_eq!(disable.id, REMOTE_CONTROL_REQUEST_ID);
            assert_eq!(disable.method, "remoteControl/disable");
            assert_eq!(
                disable.params,
                Some(serde_json::json!({ "ephemeral": true }))
            );
            client::send_message(
                &mut websocket,
                &JSONRPCMessage::Error(JSONRPCError {
                    id: REMOTE_CONTROL_REQUEST_ID,
                    error: JSONRPCErrorError {
                        code: INVALID_PARAMS_ERROR_CODE,
                        message: "Invalid params".to_string(),
                        data: None,
                    },
                }),
            )
            .await?;
            let fallback = client::read_message(&mut websocket).await?;
            let JSONRPCMessage::Request(fallback) = fallback else {
                panic!("expected fallback remoteControl/disable request");
            };
            assert_eq!(fallback.id, REMOTE_CONTROL_REQUEST_ID);
            assert_eq!(fallback.method, "remoteControl/disable");
            assert_eq!(fallback.params, None);
            client::send_message(
                &mut websocket,
                &JSONRPCMessage::Response(JSONRPCResponse {
                    id: REMOTE_CONTROL_REQUEST_ID,
                    result: serde_json::to_value(RemoteControlDisableResponse::from(
                        remote_control_status(
                            RemoteControlConnectionStatus::Disabled,
                            /*environment_id*/ None,
                        ),
                    ))?,
                }),
            )
            .await?;
            Ok::<_, anyhow::Error>(())
        });

        let status = disable_remote_control(&socket_path).await?;
        server_task.await??;
        assert_eq!(
            status,
            RemoteControlReadyStatus {
                status: RemoteControlConnectionStatus::Disabled,
                server_name: TEST_SERVER_NAME.to_string(),
                environment_id: None,
                timed_out: false,
            }
        );
        Ok(())
    }

    struct EnableScenario {
        initial_notification: Option<RemoteControlStatusChangedNotification>,
        enable_response: RemoteControlStatusChangedNotification,
        after_enable_notification: Option<RemoteControlStatusChangedNotification>,
        ready_timeout: Duration,
        reject_ephemeral_params: bool,
    }

    async fn run_enable_remote_control_scenario(
        scenario: EnableScenario,
    ) -> Result<RemoteControlReadyStatus> {
        let dir = TempDir::new()?;
        let socket_path = dir.path().join("app-server.sock");
        let listener = UnixListener::bind(&socket_path).await?;
        let ready_timeout = scenario.ready_timeout;
        let server_task = tokio::spawn(serve_enable_remote_control_scenario(listener, scenario));

        let mut websocket = client::connect(&socket_path).await?;
        let status = enable_remote_control_with_timeout(&mut websocket, ready_timeout).await?;
        server_task.await??;
        Ok(status)
    }

    async fn serve_enable_remote_control_scenario(
        listener: UnixListener,
        scenario: EnableScenario,
    ) -> Result<()> {
        let mut websocket = accept_initialized_client(listener).await?;
        if let Some(status) = scenario.initial_notification {
            send_remote_control_status(&mut websocket, status).await?;
        }

        let enable = client::read_message(&mut websocket).await?;
        let JSONRPCMessage::Request(enable) = enable else {
            panic!("expected remoteControl/enable request");
        };
        assert_eq!(enable.id, REMOTE_CONTROL_REQUEST_ID);
        assert_eq!(enable.method, "remoteControl/enable");
        assert_eq!(
            enable.params,
            Some(serde_json::json!({ "ephemeral": true }))
        );
        if scenario.reject_ephemeral_params {
            client::send_message(
                &mut websocket,
                &JSONRPCMessage::Error(JSONRPCError {
                    id: REMOTE_CONTROL_REQUEST_ID,
                    error: JSONRPCErrorError {
                        code: INVALID_PARAMS_ERROR_CODE,
                        message: "Invalid params".to_string(),
                        data: None,
                    },
                }),
            )
            .await?;
            let fallback = client::read_message(&mut websocket).await?;
            let JSONRPCMessage::Request(fallback) = fallback else {
                panic!("expected fallback remoteControl/enable request");
            };
            assert_eq!(fallback.id, REMOTE_CONTROL_REQUEST_ID);
            assert_eq!(fallback.method, "remoteControl/enable");
            assert_eq!(fallback.params, None);
        }
        client::send_message(
            &mut websocket,
            &JSONRPCMessage::Response(JSONRPCResponse {
                id: REMOTE_CONTROL_REQUEST_ID,
                result: serde_json::to_value(RemoteControlEnableResponse::from(
                    scenario.enable_response,
                ))?,
            }),
        )
        .await?;

        if let Some(status) = scenario.after_enable_notification {
            send_remote_control_status(&mut websocket, status).await?;
        } else {
            tokio::time::sleep(Duration::from_millis(50)).await;
        }

        Ok(())
    }

    async fn accept_initialized_client(
        mut listener: UnixListener,
    ) -> Result<WebSocketStream<codex_uds::UnixStream>> {
        let stream = listener.accept().await?;
        let mut websocket = accept_async(stream).await?;
        let initialize = client::read_message(&mut websocket).await?;
        let JSONRPCMessage::Request(initialize) = initialize else {
            panic!("expected initialize request");
        };
        assert_eq!(initialize.id, INITIALIZE_REQUEST_ID);
        assert_eq!(initialize.method, "initialize");
        let Some(initialize_params) = initialize.params else {
            panic!("expected initialize params");
        };
        assert_eq!(
            initialize_params["capabilities"]["experimentalApi"],
            serde_json::Value::Bool(true)
        );
        client::send_message(
            &mut websocket,
            &JSONRPCMessage::Response(JSONRPCResponse {
                id: INITIALIZE_REQUEST_ID,
                result: serde_json::json!({
                    "userAgent": "codex_app_server/1.2.3",
                    "codexHome": TEST_CODEX_HOME,
                    "platformFamily": "unix",
                    "platformOs": "macos",
                }),
            }),
        )
        .await?;

        let initialized = client::read_message(&mut websocket).await?;
        let JSONRPCMessage::Notification(initialized) = initialized else {
            panic!("expected initialized notification");
        };
        assert_eq!(initialized.method, "initialized");
        Ok(websocket)
    }

    async fn send_remote_control_status<S>(
        websocket: &mut WebSocketStream<S>,
        status: RemoteControlStatusChangedNotification,
    ) -> Result<()>
    where
        S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
    {
        client::send_message(
            websocket,
            &JSONRPCMessage::Notification(JSONRPCNotification {
                method: "remoteControl/status/changed".to_string(),
                params: Some(serde_json::to_value(status)?),
            }),
        )
        .await
    }

    fn remote_control_status(
        status: RemoteControlConnectionStatus,
        environment_id: Option<&str>,
    ) -> RemoteControlStatusChangedNotification {
        RemoteControlStatusChangedNotification {
            status,
            server_name: TEST_SERVER_NAME.to_string(),
            installation_id: TEST_INSTALLATION_ID.to_string(),
            environment_id: environment_id.map(str::to_string),
        }
    }
}