codex_codes/
client_raw_async.rs1use 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 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 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}