Codex Handbook
app-server-transport/src/transport/mod.rs 490 lines
pub mod auth;use crate::outgoing_message::ConnectionId;use crate::outgoing_message::OutgoingError;use crate::outgoing_message::OutgoingMessage;use crate::outgoing_message::QueuedOutgoingMessage;use codex_app_server_protocol::JSONRPCErrorError;use codex_app_server_protocol::JSONRPCMessage;use codex_core::config::find_codex_home;use codex_utils_absolute_path::AbsolutePathBuf;use std::net::SocketAddr;use std::path::Path;use std::str::FromStr;use std::sync::atomic::AtomicU64;use std::sync::atomic::Ordering;use tokio::sync::mpsc;use tokio_util::sync::CancellationToken;use tracing::error;use tracing::warn;/// Size of the bounded channels used to communicate between tasks. The value/// is a balance between throughput and memory usage - 128 messages should be/// plenty for an interactive CLI.pub const CHANNEL_CAPACITY: usize = 128;mod remote_control;mod stdio;mod unix_socket;#[cfg(test)]mod unix_socket_tests;mod websocket;pub use remote_control::REMOTE_CONTROL_DISABLED_ENV_VAR;pub use remote_control::RemoteControlDisabledByRequirements;pub use remote_control::RemoteControlEnableError;pub use remote_control::RemoteControlHandle;pub use remote_control::RemoteControlPolicy;pub use remote_control::RemoteControlStartConfig;pub use remote_control::RemoteControlStartupMode;pub use remote_control::RemoteControlUnavailable;pub use remote_control::start_remote_control;pub use remote_control::take_remote_control_disabled_env;pub use stdio::start_stdio_connection;pub use unix_socket::AppServerStartupLock;pub use unix_socket::acquire_app_server_startup_lock;pub use unix_socket::prepare_control_socket_path;pub use unix_socket::start_control_socket_acceptor;pub use websocket::start_websocket_acceptor;const OVERLOADED_ERROR_CODE: i64 = -32001;const APP_SERVER_CONTROL_SOCKET_DIR_NAME: &str = "app-server-control";const APP_SERVER_CONTROL_SOCKET_FILE_NAME: &str = "app-server-control.sock";const APP_SERVER_STARTUP_LOCK_FILE_NAME: &str = "app-server-startup.lock";pub fn app_server_control_socket_path(codex_home: &Path) -> std::io::Result<AbsolutePathBuf> {    AbsolutePathBuf::from_absolute_path(        codex_home            .join(APP_SERVER_CONTROL_SOCKET_DIR_NAME)            .join(APP_SERVER_CONTROL_SOCKET_FILE_NAME),    )}pub fn app_server_startup_lock_path(codex_home: &Path) -> std::io::Result<AbsolutePathBuf> {    AbsolutePathBuf::from_absolute_path(        codex_home            .join(APP_SERVER_CONTROL_SOCKET_DIR_NAME)            .join(APP_SERVER_STARTUP_LOCK_FILE_NAME),    )}#[derive(Clone, Debug, Eq, PartialEq)]pub enum AppServerTransport {    Stdio,    UnixSocket { socket_path: AbsolutePathBuf },    WebSocket { bind_address: SocketAddr },    Off,}#[derive(Debug, Clone, Eq, PartialEq)]pub enum AppServerTransportParseError {    UnsupportedListenUrl(String),    InvalidUnixSocketPath { listen_url: String, message: String },    InvalidWebSocketListenUrl(String),}impl std::fmt::Display for AppServerTransportParseError {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        match self {            AppServerTransportParseError::UnsupportedListenUrl(listen_url) => write!(                f,                "unsupported --listen URL `{listen_url}`; expected `stdio://`, `unix://`, `unix://PATH`, `ws://IP:PORT`, or `off`"            ),            AppServerTransportParseError::InvalidUnixSocketPath {                listen_url,                message,            } => write!(                f,                "invalid unix socket --listen URL `{listen_url}`; failed to resolve socket path: {message}"            ),            AppServerTransportParseError::InvalidWebSocketListenUrl(listen_url) => write!(                f,                "invalid websocket --listen URL `{listen_url}`; expected `ws://IP:PORT`"            ),        }    }}impl std::error::Error for AppServerTransportParseError {}impl AppServerTransport {    pub const DEFAULT_LISTEN_URL: &'static str = "stdio://";    pub fn from_listen_url(listen_url: &str) -> Result<Self, AppServerTransportParseError> {        if listen_url == Self::DEFAULT_LISTEN_URL {            return Ok(Self::Stdio);        }        if let Some(raw_socket_path) = listen_url.strip_prefix("unix://") {            let socket_path = if raw_socket_path.is_empty() {                let codex_home = find_codex_home().map_err(|err| {                    AppServerTransportParseError::InvalidUnixSocketPath {                        listen_url: listen_url.to_string(),                        message: format!("failed to resolve CODEX_HOME: {err}"),                    }                })?;                app_server_control_socket_path(&codex_home).map_err(|err| {                    AppServerTransportParseError::InvalidUnixSocketPath {                        listen_url: listen_url.to_string(),                        message: err.to_string(),                    }                })?            } else {                AbsolutePathBuf::relative_to_current_dir(raw_socket_path).map_err(|err| {                    AppServerTransportParseError::InvalidUnixSocketPath {                        listen_url: listen_url.to_string(),                        message: err.to_string(),                    }                })?            };            return Ok(Self::UnixSocket { socket_path });        }        if listen_url == "off" {            return Ok(Self::Off);        }        if let Some(socket_addr) = listen_url.strip_prefix("ws://") {            let bind_address = socket_addr.parse::<SocketAddr>().map_err(|_| {                AppServerTransportParseError::InvalidWebSocketListenUrl(listen_url.to_string())            })?;            return Ok(Self::WebSocket { bind_address });        }        Err(AppServerTransportParseError::UnsupportedListenUrl(            listen_url.to_string(),        ))    }}impl FromStr for AppServerTransport {    type Err = AppServerTransportParseError;    fn from_str(s: &str) -> Result<Self, Self::Err> {        Self::from_listen_url(s)    }}#[derive(Debug)]pub enum TransportEvent {    ConnectionOpened {        connection_id: ConnectionId,        origin: ConnectionOrigin,        writer: mpsc::Sender<QueuedOutgoingMessage>,        disconnect_sender: Option<CancellationToken>,    },    ConnectionClosed {        connection_id: ConnectionId,    },    IncomingMessage {        connection_id: ConnectionId,        message: JSONRPCMessage,    },}#[derive(Debug, Clone, Copy, PartialEq, Eq)]pub enum ConnectionOrigin {    Stdio,    InProcess,    WebSocket,    RemoteControl,}static CONNECTION_ID_COUNTER: AtomicU64 = AtomicU64::new(0);fn next_connection_id() -> ConnectionId {    ConnectionId(CONNECTION_ID_COUNTER.fetch_add(1, Ordering::Relaxed))}async fn forward_incoming_message(    transport_event_tx: &mpsc::Sender<TransportEvent>,    writer: &mpsc::Sender<QueuedOutgoingMessage>,    connection_id: ConnectionId,    payload: &str,) -> bool {    match serde_json::from_str::<JSONRPCMessage>(payload) {        Ok(message) => {            enqueue_incoming_message(transport_event_tx, writer, connection_id, message).await        }        Err(err) => {            error!("Failed to deserialize JSONRPCMessage: {err}");            true        }    }}async fn enqueue_incoming_message(    transport_event_tx: &mpsc::Sender<TransportEvent>,    writer: &mpsc::Sender<QueuedOutgoingMessage>,    connection_id: ConnectionId,    message: JSONRPCMessage,) -> bool {    let event = TransportEvent::IncomingMessage {        connection_id,        message,    };    match transport_event_tx.try_send(event) {        Ok(()) => true,        Err(mpsc::error::TrySendError::Closed(_)) => false,        Err(mpsc::error::TrySendError::Full(TransportEvent::IncomingMessage {            connection_id,            message: JSONRPCMessage::Request(request),        })) => {            let overload_error = OutgoingMessage::Error(OutgoingError {                id: request.id,                error: JSONRPCErrorError {                    code: OVERLOADED_ERROR_CODE,                    message: "Server overloaded; retry later.".to_string(),                    data: None,                },            });            match writer.try_send(QueuedOutgoingMessage::new(overload_error)) {                Ok(()) => true,                Err(mpsc::error::TrySendError::Closed(_)) => false,                Err(mpsc::error::TrySendError::Full(_overload_error)) => {                    warn!(                        "dropping overload response for connection {:?}: outbound queue is full",                        connection_id                    );                    true                }            }        }        Err(mpsc::error::TrySendError::Full(event)) => transport_event_tx.send(event).await.is_ok(),    }}fn serialize_outgoing_message(outgoing_message: OutgoingMessage) -> Option<String> {    let value = match serde_json::to_value(outgoing_message) {        Ok(value) => value,        Err(err) => {            error!("Failed to convert OutgoingMessage to JSON value: {err}");            return None;        }    };    match serde_json::to_string(&value) {        Ok(json) => Some(json),        Err(err) => {            error!("Failed to serialize JSONRPCMessage: {err}");            None        }    }}#[cfg(test)]mod tests {    use super::*;    use codex_app_server_protocol::ConfigWarningNotification;    use codex_app_server_protocol::JSONRPCNotification;    use codex_app_server_protocol::JSONRPCRequest;    use codex_app_server_protocol::JSONRPCResponse;    use codex_app_server_protocol::RequestId;    use codex_app_server_protocol::ServerNotification;    use pretty_assertions::assert_eq;    use serde_json::json;    use tokio::time::Duration;    use tokio::time::timeout;    #[test]    fn listen_off_parses_as_off_transport() {        assert_eq!(            AppServerTransport::from_listen_url("off"),            Ok(AppServerTransport::Off)        );    }    #[tokio::test]    async fn enqueue_incoming_request_returns_overload_error_when_queue_is_full() {        let connection_id = ConnectionId(42);        let (transport_event_tx, mut transport_event_rx) = mpsc::channel(1);        let (writer_tx, mut writer_rx) = mpsc::channel(1);        let first_message = JSONRPCMessage::Notification(JSONRPCNotification {            method: "initialized".to_string(),            params: None,        });        transport_event_tx            .send(TransportEvent::IncomingMessage {                connection_id,                message: first_message.clone(),            })            .await            .expect("queue should accept first message");        let request = JSONRPCMessage::Request(JSONRPCRequest {            id: RequestId::Integer(7),            method: "config/read".to_string(),            params: Some(json!({ "includeLayers": false })),            trace: None,        });        assert!(            enqueue_incoming_message(&transport_event_tx, &writer_tx, connection_id, request).await        );        let queued_event = transport_event_rx            .recv()            .await            .expect("first event should stay queued");        match queued_event {            TransportEvent::IncomingMessage {                connection_id: queued_connection_id,                message,            } => {                assert_eq!(queued_connection_id, connection_id);                assert_eq!(message, first_message);            }            _ => panic!("expected queued incoming message"),        }        let overload = writer_rx            .recv()            .await            .expect("request should receive overload error");        let overload_json =            serde_json::to_value(overload.message).expect("serialize overload error");        assert_eq!(            overload_json,            json!({                "id": 7,                "error": {                    "code": OVERLOADED_ERROR_CODE,                    "message": "Server overloaded; retry later."                }            })        );    }    #[tokio::test]    async fn enqueue_incoming_response_waits_instead_of_dropping_when_queue_is_full() {        let connection_id = ConnectionId(42);        let (transport_event_tx, mut transport_event_rx) = mpsc::channel(1);        let (writer_tx, _writer_rx) = mpsc::channel(1);        let first_message = JSONRPCMessage::Notification(JSONRPCNotification {            method: "initialized".to_string(),            params: None,        });        transport_event_tx            .send(TransportEvent::IncomingMessage {                connection_id,                message: first_message.clone(),            })            .await            .expect("queue should accept first message");        let response = JSONRPCMessage::Response(JSONRPCResponse {            id: RequestId::Integer(7),            result: json!({"ok": true}),        });        let transport_event_tx_for_enqueue = transport_event_tx.clone();        let writer_tx_for_enqueue = writer_tx.clone();        let enqueue_handle = tokio::spawn(async move {            enqueue_incoming_message(                &transport_event_tx_for_enqueue,                &writer_tx_for_enqueue,                connection_id,                response,            )            .await        });        let queued_event = transport_event_rx            .recv()            .await            .expect("first event should be dequeued");        match queued_event {            TransportEvent::IncomingMessage {                connection_id: queued_connection_id,                message,            } => {                assert_eq!(queued_connection_id, connection_id);                assert_eq!(message, first_message);            }            _ => panic!("expected queued incoming message"),        }        let enqueue_result = enqueue_handle.await.expect("enqueue task should not panic");        assert!(enqueue_result);        let forwarded_event = transport_event_rx            .recv()            .await            .expect("response should be forwarded instead of dropped");        match forwarded_event {            TransportEvent::IncomingMessage {                connection_id: queued_connection_id,                message: JSONRPCMessage::Response(JSONRPCResponse { id, result }),            } => {                assert_eq!(queued_connection_id, connection_id);                assert_eq!(id, RequestId::Integer(7));                assert_eq!(result, json!({"ok": true}));            }            _ => panic!("expected forwarded response message"),        }    }    #[tokio::test]    async fn enqueue_incoming_request_does_not_block_when_writer_queue_is_full() {        let connection_id = ConnectionId(42);        let (transport_event_tx, _transport_event_rx) = mpsc::channel(1);        let (writer_tx, mut writer_rx) = mpsc::channel(1);        transport_event_tx            .send(TransportEvent::IncomingMessage {                connection_id,                message: JSONRPCMessage::Notification(JSONRPCNotification {                    method: "initialized".to_string(),                    params: None,                }),            })            .await            .expect("transport queue should accept first message");        writer_tx            .send(QueuedOutgoingMessage::new(                OutgoingMessage::AppServerNotification(ServerNotification::ConfigWarning(                    ConfigWarningNotification {                        summary: "queued".to_string(),                        details: None,                        path: None,                        range: None,                    },                )),            ))            .await            .expect("writer queue should accept first message");        let request = JSONRPCMessage::Request(JSONRPCRequest {            id: RequestId::Integer(7),            method: "config/read".to_string(),            params: Some(json!({ "includeLayers": false })),            trace: None,        });        let enqueue_result = timeout(            Duration::from_millis(100),            enqueue_incoming_message(&transport_event_tx, &writer_tx, connection_id, request),        )        .await        .expect("enqueue should not block while writer queue is full");        assert!(enqueue_result);        let queued_outgoing = writer_rx            .recv()            .await            .expect("writer queue should still contain original message");        let queued_json =            serde_json::to_value(queued_outgoing.message).expect("serialize queued message");        assert_eq!(            queued_json,            json!({                "method": "configWarning",                "params": {                    "summary": "queued",                    "details": null,                },            })        );    }}