Codex Handbook
app-server-daemon/src/client.rs 181 lines
use std::path::Path;use std::time::Duration;use anyhow::Context;use anyhow::Result;use anyhow::anyhow;use codex_app_server_protocol::ClientInfo;use codex_app_server_protocol::InitializeCapabilities;use codex_app_server_protocol::InitializeParams;use codex_app_server_protocol::InitializeResponse;use codex_app_server_protocol::JSONRPCMessage;use codex_app_server_protocol::JSONRPCNotification;use codex_app_server_protocol::JSONRPCRequest;use codex_app_server_protocol::RequestId;use codex_uds::UnixStream;use futures::SinkExt;use futures::StreamExt;use tokio::io::AsyncRead;use tokio::io::AsyncWrite;use tokio::time::timeout;use tokio_tungstenite::WebSocketStream;use tokio_tungstenite::client_async;use tokio_tungstenite::tungstenite::Message;pub(crate) const CONTROL_SOCKET_RESPONSE_TIMEOUT: Duration = Duration::from_secs(2);const CLIENT_NAME: &str = "codex_app_server_daemon";const INITIALIZE_REQUEST_ID: RequestId = RequestId::Integer(1);#[derive(Debug, Clone, PartialEq, Eq)]pub(crate) struct ProbeInfo {    pub(crate) app_server_version: String,}pub(crate) async fn probe(socket_path: &Path) -> Result<ProbeInfo> {    timeout(CONTROL_SOCKET_RESPONSE_TIMEOUT, probe_inner(socket_path))        .await        .with_context(|| {            format!(                "timed out probing app-server control socket {}",                socket_path.display()            )        })?}async fn probe_inner(socket_path: &Path) -> Result<ProbeInfo> {    let mut websocket = connect(socket_path).await?;    let initialize_response = initialize(&mut websocket, /*experimental_api*/ false).await?;    let initialized = JSONRPCMessage::Notification(JSONRPCNotification {        method: "initialized".to_string(),        params: None,    });    send_message(&mut websocket, &initialized)        .await        .context("failed to send initialized notification")?;    websocket.close(None).await.ok();    Ok(ProbeInfo {        app_server_version: parse_version_from_user_agent(&initialize_response.user_agent)?,    })}pub(crate) async fn connect(socket_path: &Path) -> Result<WebSocketStream<UnixStream>> {    let stream = UnixStream::connect(socket_path)        .await        .with_context(|| format!("failed to connect to {}", socket_path.display()))?;    let (websocket, _response) = client_async("ws://localhost/", stream)        .await        .with_context(|| format!("failed to upgrade {}", socket_path.display()))?;    Ok(websocket)}pub(crate) async fn initialize<S>(    websocket: &mut WebSocketStream<S>,    experimental_api: bool,) -> Result<InitializeResponse>where    S: AsyncRead + AsyncWrite + Unpin,{    let initialize = JSONRPCMessage::Request(JSONRPCRequest {        id: INITIALIZE_REQUEST_ID,        method: "initialize".to_string(),        params: Some(serde_json::to_value(InitializeParams {            client_info: ClientInfo {                name: CLIENT_NAME.to_string(),                title: Some("Codex App Server Daemon".to_string()),                version: env!("CARGO_PKG_VERSION").to_string(),            },            capabilities: if experimental_api {                Some(InitializeCapabilities {                    experimental_api: true,                    ..Default::default()                })            } else {                None            },        })?),        trace: None,    });    send_message(websocket, &initialize)        .await        .context("failed to send initialize request")?;    let response = loop {        let message = timeout(CONTROL_SOCKET_RESPONSE_TIMEOUT, read_message(websocket))            .await            .context("timed out waiting for initialize response")??;        if let JSONRPCMessage::Response(response) = message            && response.id == INITIALIZE_REQUEST_ID        {            break response;        }    };    serde_json::from_value::<InitializeResponse>(response.result)        .context("failed to parse initialize response")}pub(crate) async fn send_message<S>(    websocket: &mut WebSocketStream<S>,    message: &JSONRPCMessage,) -> Result<()>where    S: AsyncRead + AsyncWrite + Unpin,{    websocket        .send(Message::Text(serde_json::to_string(message)?.into()))        .await?;    Ok(())}pub(crate) async fn read_message<S>(websocket: &mut WebSocketStream<S>) -> Result<JSONRPCMessage>where    S: AsyncRead + AsyncWrite + Unpin,{    loop {        let frame = websocket            .next()            .await            .ok_or_else(|| anyhow!("app-server closed the control socket"))??;        let Message::Text(payload) = frame else {            continue;        };        return serde_json::from_str::<JSONRPCMessage>(&payload)            .context("failed to parse app-server JSON-RPC message");    }}fn parse_version_from_user_agent(user_agent: &str) -> Result<String> {    let (_originator, rest) = user_agent        .split_once('/')        .ok_or_else(|| anyhow!("app-server user-agent omitted version separator"))?;    let version = rest        .split_whitespace()        .next()        .filter(|version| !version.is_empty())        .ok_or_else(|| anyhow!("app-server user-agent omitted version"))?;    Ok(version.to_string())}#[cfg(all(test, unix))]mod tests {    use pretty_assertions::assert_eq;    use super::parse_version_from_user_agent;    #[test]    fn parses_version_from_codex_user_agent() {        assert_eq!(            parse_version_from_user_agent(                "codex_app_server_daemon/1.2.3 (Linux 6.8.0; x86_64) codex_cli_rs/1.2.3",            )            .expect("version"),            "1.2.3"        );    }    #[test]    fn rejects_user_agent_without_version() {        assert!(parse_version_from_user_agent("codex_app_server_daemon").is_err());    }}