Codex Handbook
exec-server/src/client_transport.rs 239 lines
use std::process::Stdio;use tokio::io::AsyncBufReadExt;use tokio::io::BufReader;use tokio::process::Command;use tokio::time::timeout;use tokio_tungstenite::connect_async;use tokio_tungstenite::connect_async_with_config;use tracing::debug;use tracing::warn;use codex_utils_rustls_provider::ensure_rustls_crypto_provider;use crate::ExecServerClient;use crate::ExecServerError;use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT;use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT;use crate::client_api::NoiseRendezvousConnectArgs;use crate::client_api::NoiseRendezvousConnectBundle;use crate::client_api::RemoteExecServerConnectArgs;use crate::client_api::StdioExecServerCommand;use crate::client_api::StdioExecServerConnectArgs;use crate::connection::JsonRpcConnection;use crate::noise_relay::NoiseHarnessConnectionArgs;use crate::noise_relay::noise_harness_connection_from_websocket;use crate::noise_relay::noise_relay_websocket_config;use crate::relay::harness_connection_from_websocket;const ENVIRONMENT_CLIENT_NAME: &str = "codex-environment";impl ExecServerClient {    /// Open the selected transport and run the common JSON-RPC initialization.    /// Noise connection details are fetched here so reconnects get a fresh URL    /// and authorization without replacing the harness identity.    pub(crate) async fn connect_for_transport(        transport_params: crate::client_api::ExecServerTransportParams,    ) -> Result<Self, ExecServerError> {        match transport_params {            crate::client_api::ExecServerTransportParams::WebSocketUrl {                websocket_url,                connect_timeout,                initialize_timeout,            } => {                Self::connect_websocket(RemoteExecServerConnectArgs {                    websocket_url,                    client_name: ENVIRONMENT_CLIENT_NAME.to_string(),                    connect_timeout,                    initialize_timeout,                    resume_session_id: None,                })                .await            }            crate::client_api::ExecServerTransportParams::NoiseRendezvous {                provider,                identity,            } => {                let bundle = provider.connect_bundle(identity.public_key()).await?;                Self::connect_noise_rendezvous(NoiseRendezvousConnectArgs {                    bundle,                    harness_identity: identity,                    client_name: ENVIRONMENT_CLIENT_NAME.to_string(),                    connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT,                    initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT,                    resume_session_id: None,                })                .await            }            crate::client_api::ExecServerTransportParams::StdioCommand {                command,                initialize_timeout,            } => {                Self::connect_stdio_command(StdioExecServerConnectArgs {                    command,                    client_name: ENVIRONMENT_CLIENT_NAME.to_string(),                    initialize_timeout,                    resume_session_id: None,                })                .await            }        }    }    pub async fn connect_websocket(        args: RemoteExecServerConnectArgs,    ) -> Result<Self, ExecServerError> {        ensure_rustls_crypto_provider();        let websocket_url = args.websocket_url.clone();        let connect_timeout = args.connect_timeout;        let (stream, _) = timeout(connect_timeout, connect_async(websocket_url.as_str()))            .await            .map_err(|_| ExecServerError::WebSocketConnectTimeout {                url: websocket_url.clone(),                timeout: connect_timeout,            })?            .map_err(|source| ExecServerError::WebSocketConnect {                url: websocket_url.clone(),                source,            })?;        let connection_label = format!("exec-server websocket {websocket_url}");        let connection = if is_rendezvous_harness_url(&websocket_url) {            harness_connection_from_websocket(stream, connection_label)        } else {            JsonRpcConnection::from_websocket(stream, connection_label)        };        Self::connect(connection, args.into()).await    }    /// Connect to one exec-server through an authenticated rendezvous stream.    /// The executor key is pinned before JSON-RPC starts; the websocket carries    /// only ciphertext after that.    pub async fn connect_noise_rendezvous(        args: NoiseRendezvousConnectArgs,    ) -> Result<Self, ExecServerError> {        ensure_rustls_crypto_provider();        // Keep the registry-issued URL, key, and authorization together for this        // connection attempt.        let NoiseRendezvousConnectArgs {            bundle,            harness_identity,            client_name,            connect_timeout,            initialize_timeout,            resume_session_id,        } = args;        let NoiseRendezvousConnectBundle {            websocket_url,            environment_id,            executor_registration_id,            executor_public_key,            harness_key_authorization,        } = bundle;        let diagnostic_url = websocket_url            .split(['?', '#'])            .next()            .unwrap_or(websocket_url.as_str())            .to_string();        let (stream, _) = timeout(            connect_timeout,            connect_async_with_config(                websocket_url.as_str(),                Some(noise_relay_websocket_config()),                /*disable_nagle*/ false,            ),        )        .await        .map_err(|_| ExecServerError::WebSocketConnectTimeout {            url: diagnostic_url.clone(),            timeout: connect_timeout,        })?        .map_err(|source| ExecServerError::WebSocketConnect {            url: diagnostic_url.clone(),            source,        })?;        let connection_label = format!("Noise exec-server rendezvous websocket {diagnostic_url}");        let connection = noise_harness_connection_from_websocket(            stream,            NoiseHarnessConnectionArgs {                connection_label,                environment_id,                executor_registration_id,                identity: harness_identity,                responder_public_key: executor_public_key,                harness_key_authorization,            },        );        Self::connect(            connection,            crate::client_api::ExecServerClientConnectOptions {                client_name,                initialize_timeout,                resume_session_id,            },        )        .await    }    pub(crate) async fn connect_stdio_command(        args: StdioExecServerConnectArgs,    ) -> Result<Self, ExecServerError> {        let mut child = stdio_command_process(&args.command)            .stdin(Stdio::piped())            .stdout(Stdio::piped())            .stderr(Stdio::piped())            .spawn()            .map_err(ExecServerError::Spawn)?;        let stdin = child.stdin.take().ok_or_else(|| {            ExecServerError::Protocol("spawned exec-server command has no stdin".to_string())        })?;        let stdout = child.stdout.take().ok_or_else(|| {            ExecServerError::Protocol("spawned exec-server command has no stdout".to_string())        })?;        if let Some(stderr) = child.stderr.take() {            tokio::spawn(async move {                let mut lines = BufReader::new(stderr).lines();                loop {                    match lines.next_line().await {                        Ok(Some(line)) => debug!("exec-server stdio stderr: {line}"),                        Ok(None) => break,                        Err(err) => {                            warn!("failed to read exec-server stdio stderr: {err}");                            break;                        }                    }                }            });        }        Self::connect(            JsonRpcConnection::from_stdio(stdout, stdin, "exec-server stdio command".to_string())                .with_child_process(child),            args.into(),        )        .await    }}fn is_rendezvous_harness_url(websocket_url: &str) -> bool {    let Some((_path, query)) = websocket_url.split_once('?') else {        return false;    };    query        .split('&')        .filter_map(|pair| pair.split_once('='))        .any(|(key, value)| key == "role" && value == "harness")}fn stdio_command_process(stdio_command: &StdioExecServerCommand) -> Command {    let mut command = Command::new(&stdio_command.program);    command.args(&stdio_command.args);    command.envs(&stdio_command.env);    if let Some(cwd) = &stdio_command.cwd {        command.current_dir(cwd);    }    #[cfg(unix)]    command.process_group(0);    command}