Skip to main content

codex_codes/
client_raw_async.rs

1//! Raw asynchronous transport for the Codex app-server.
2//!
3//! This client only frames newline-delimited messages. It does not decode JSON
4//! or correlate JSON-RPC requests and responses.
5
6use crate::cli::AppServerBuilder;
7use crate::error::{Error, Result};
8use log::{debug, error};
9use serde::Serialize;
10use std::process::ExitStatus;
11use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, BufWriter};
12use tokio::process::Child;
13
14const STDOUT_BUFFER_SIZE: usize = 10 * 1024 * 1024;
15
16pub struct RawAsyncClient {
17    child: Child,
18    writer: BufWriter<tokio::process::ChildStdin>,
19    reader: BufReader<tokio::process::ChildStdout>,
20    _stderr_drain: tokio::task::JoinHandle<()>,
21}
22
23impl RawAsyncClient {
24    pub async fn start() -> Result<Self> {
25        Self::start_with(AppServerBuilder::new()).await
26    }
27
28    pub async fn start_with(builder: AppServerBuilder) -> Result<Self> {
29        crate::version::check_codex_version_async().await?;
30        Self::new(builder.spawn().await?)
31    }
32
33    pub fn new(mut child: Child) -> Result<Self> {
34        let stdin = child
35            .stdin
36            .take()
37            .ok_or_else(|| Error::Protocol("Failed to get stdin".to_string()))?;
38        let stdout = child
39            .stdout
40            .take()
41            .ok_or_else(|| Error::Protocol("Failed to get stdout".to_string()))?;
42        let stderr = child
43            .stderr
44            .take()
45            .ok_or_else(|| Error::Protocol("Failed to get stderr".to_string()))?;
46
47        Ok(Self {
48            child,
49            writer: BufWriter::new(stdin),
50            reader: BufReader::with_capacity(STDOUT_BUFFER_SIZE, stdout),
51            _stderr_drain: crate::stderr_drain::spawn_async(stderr),
52        })
53    }
54
55    /// Serialize and write one typed SDK message.
56    pub async fn send<T: Serialize>(&mut self, message: &T) -> Result<()> {
57        let line = serde_json::to_string(message).map_err(Error::Json)?;
58        self.write_line(&line).await
59    }
60
61    async fn write_line(&mut self, line: &str) -> Result<()> {
62        let line = single_line(line)?;
63        debug!("[RAW CLIENT] Sending {} bytes", line.len());
64        self.writer
65            .write_all(line.as_bytes())
66            .await
67            .map_err(Error::Io)?;
68        self.writer.write_all(b"\n").await.map_err(Error::Io)?;
69        self.writer.flush().await.map_err(Error::Io)
70    }
71
72    /// Read one complete JSONL frame without inspecting its contents.
73    pub async fn next_line(&mut self) -> Result<Option<String>> {
74        let mut line = String::new();
75        loop {
76            line.clear();
77            if self.reader.read_line(&mut line).await.map_err(Error::Io)? == 0 {
78                return Ok(None);
79            }
80            remove_line_ending(&mut line);
81            if line.trim().is_empty() {
82                continue;
83            }
84            debug!("[RAW CLIENT] Received {} bytes", line.len());
85            return Ok(Some(line));
86        }
87    }
88
89    pub fn pid(&self) -> Option<u32> {
90        self.child.id()
91    }
92
93    pub fn is_alive(&mut self) -> bool {
94        self.child.try_wait().ok().flatten().is_none()
95    }
96
97    pub async fn wait_for_exit(&mut self) -> Result<ExitStatus> {
98        self.child.wait().await.map_err(Error::Io)
99    }
100
101    pub async fn shutdown(mut self) -> Result<()> {
102        self.child.kill().await.map_err(Error::Io)
103    }
104}
105
106impl Drop for RawAsyncClient {
107    fn drop(&mut self) {
108        if self.is_alive() {
109            if let Err(error) = self.child.start_kill() {
110                error!("Failed to kill raw app-server process on drop: {error}");
111            }
112        }
113    }
114}
115
116fn single_line(line: &str) -> Result<&str> {
117    let line = line.trim_end_matches(['\r', '\n']);
118    if line.contains(['\r', '\n']) {
119        return Err(Error::Protocol(
120            "raw app-server frame contains an embedded line break".to_string(),
121        ));
122    }
123    Ok(line)
124}
125
126fn remove_line_ending(line: &mut String) {
127    if line.ends_with('\n') {
128        line.pop();
129        if line.ends_with('\r') {
130            line.pop();
131        }
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138
139    #[test]
140    fn accepts_one_frame_with_a_line_ending() {
141        assert_eq!(single_line("{\"id\":1}\r\n").unwrap(), "{\"id\":1}");
142    }
143
144    #[test]
145    fn rejects_multiple_frames() {
146        assert!(single_line("{\"id\":1}\n{\"id\":2}").is_err());
147    }
148}