kcode_jsonrpc_stdio/
lib.rs

1use std::io;
2use std::process::{ExitStatus, Stdio};
3
4use kcode_jsonrpc_wire::{Message, WireError};
5use serde_json::Value;
6use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
7use tokio::process::{Child, ChildStdin, Command};
8use tokio::sync::mpsc;
9use tokio::task::{JoinError, JoinHandle};
10
11pub const MAX_INBOUND_LINE_BYTES: usize = 8 * 1024 * 1024;
12pub const MAX_OUTBOUND_LINE_BYTES: usize = 8 * 1024 * 1024;
13pub const MAX_STDERR_CAPTURE_BYTES: usize = 64 * 1024;
14pub const MAX_INBOUND_CAPACITY: usize = 1024;
15
16#[derive(Clone, Copy, Debug, Eq, PartialEq)]
17pub struct Config {
18    pub inbound_capacity: usize,
19}
20
21#[derive(Clone, Debug)]
22pub struct ProcessExit {
23    pub status: ExitStatus,
24    pub stderr: String,
25}
26
27#[derive(Debug, thiserror::Error)]
28pub enum Error {
29    #[error("stdio I/O failed: {0}")]
30    Io(#[source] io::Error),
31    #[error("child status query failed: {0}")]
32    ProcessQuery(#[source] io::Error),
33    #[error("child kill failed: {0}")]
34    ProcessKill(#[source] io::Error),
35    #[error("child wait failed: {0}")]
36    ProcessWait(#[source] io::Error),
37    #[error("child stderr read failed: {0}")]
38    StderrIo(#[source] io::Error),
39    #[error("JSON failed: {0}")]
40    Json(#[source] serde_json::Error),
41    #[error("inbound line is not UTF-8: {0}")]
42    Utf8(#[source] std::str::Utf8Error),
43    #[error("invalid JSON-RPC message: {0}")]
44    Wire(#[source] WireError),
45    #[error("stdout reader task failed: {0}")]
46    ReaderTask(#[source] JoinError),
47    #[error("stderr task failed: {0}")]
48    StderrTask(#[source] JoinError),
49    #[error("inbound line exceeds {MAX_INBOUND_LINE_BYTES} bytes")]
50    InboundLineTooLong,
51    #[error("stdout ended with an incomplete JSONL frame")]
52    IncompleteFrame,
53    #[error("outbound line exceeds {MAX_OUTBOUND_LINE_BYTES} bytes")]
54    OutboundLineTooLong,
55    #[error("invalid inbound capacity: {0}")]
56    InvalidInboundCapacity(usize),
57    #[error("child stdin is closed")]
58    StdinClosed,
59    #[error("stdout reader is unavailable")]
60    ReaderUnavailable,
61    #[error("child stderr reader is unavailable")]
62    StderrUnavailable,
63    #[error("child stdout closed while the child remains live")]
64    StdoutClosed,
65    #[error("child exited: {0:?}")]
66    ProcessExited(ProcessExit),
67}
68
69pub struct StdioRpc {
70    child: Child,
71    stdin: Option<ChildStdin>,
72    inbound: Option<mpsc::Receiver<Result<Message, Error>>>,
73    reader: Option<JoinHandle<()>>,
74    stderr: Option<JoinHandle<io::Result<String>>>,
75    exit: Option<ProcessExit>,
76}
77
78impl StdioRpc {
79    pub fn spawn(mut command: Command, config: Config) -> Result<Self, Error> {
80        if config.inbound_capacity == 0 || config.inbound_capacity > MAX_INBOUND_CAPACITY {
81            return Err(Error::InvalidInboundCapacity(config.inbound_capacity));
82        }
83        command
84            .stdin(Stdio::piped())
85            .stdout(Stdio::piped())
86            .stderr(Stdio::piped())
87            .kill_on_drop(true);
88        let mut child = command.spawn().map_err(Error::Io)?;
89        let stdin = child.stdin.take().ok_or(Error::StdinClosed)?;
90        let stdout = child.stdout.take().ok_or(Error::StdinClosed)?;
91        let stderr = child.stderr.take().ok_or(Error::StdinClosed)?;
92        let (sender, inbound) = mpsc::channel(config.inbound_capacity);
93        Ok(Self {
94            child,
95            stdin: Some(stdin),
96            inbound: Some(inbound),
97            reader: Some(tokio::spawn(read_stdout(BufReader::new(stdout), sender))),
98            stderr: Some(tokio::spawn(read_stderr(stderr))),
99            exit: None,
100        })
101    }
102
103    pub async fn send(&mut self, value: Value) -> Result<(), Error> {
104        let mut line = serde_json::to_vec(&value).map_err(Error::Json)?;
105        if line.len() > MAX_OUTBOUND_LINE_BYTES {
106            return Err(Error::OutboundLineTooLong);
107        }
108        line.push(b'\n');
109        let result = match self.stdin.as_mut() {
110            Some(stdin) => {
111                async {
112                    stdin.write_all(&line).await?;
113                    stdin.flush().await
114                }
115                .await
116            }
117            None => return Err(Error::StdinClosed),
118        };
119        if result.is_err() {
120            self.stdin.take();
121        }
122        result.map_err(Error::Io)
123    }
124
125    pub async fn next(&mut self) -> Result<Message, Error> {
126        let receiver = self.inbound.as_mut().ok_or(Error::ReaderUnavailable)?;
127        match receiver.recv().await {
128            Some(Ok(message)) => Ok(message),
129            Some(Err(error)) => {
130                self.inbound.take();
131                Err(error)
132            }
133            None => {
134                self.inbound.take();
135                self.join_reader().await?;
136                match self.child.try_wait().map_err(Error::ProcessQuery)? {
137                    Some(_) => Err(Error::ProcessExited(self.reap().await?)),
138                    None => Err(Error::StdoutClosed),
139                }
140            }
141        }
142    }
143
144    pub async fn shutdown(mut self) -> Result<ProcessExit, Error> {
145        self.stdin.take();
146        self.inbound.take();
147        let action = match self.child.try_wait() {
148            Ok(Some(_)) => Ok(()),
149            Ok(None) => self.child.start_kill().map_err(Error::ProcessKill),
150            Err(query) => match self.child.start_kill() {
151                Ok(()) => Err(Error::ProcessQuery(query)),
152                Err(kill) => Err(Error::ProcessKill(kill)),
153            },
154        };
155        let exit = self.reap().await;
156        let reader = self.abort_reader().await;
157        action?;
158        let exit = exit?;
159        reader?;
160        Ok(exit)
161    }
162
163    async fn reap(&mut self) -> Result<ProcessExit, Error> {
164        if let Some(exit) = &self.exit {
165            return Ok(exit.clone());
166        }
167        let status = self.child.wait().await.map_err(Error::ProcessWait)?;
168        self.exit = Some(ProcessExit {
169            status,
170            stderr: String::new(),
171        });
172        let stderr = self.stderr.take().ok_or(Error::StderrUnavailable)?;
173        let stderr = stderr
174            .await
175            .map_err(Error::StderrTask)?
176            .map_err(Error::StderrIo)?;
177        let exit = ProcessExit { status, stderr };
178        self.exit = Some(exit.clone());
179        Ok(exit)
180    }
181
182    async fn join_reader(&mut self) -> Result<(), Error> {
183        if let Some(reader) = self.reader.take() {
184            reader.await.map_err(Error::ReaderTask)?;
185        }
186        Ok(())
187    }
188
189    async fn abort_reader(&mut self) -> Result<(), Error> {
190        if let Some(reader) = self.reader.take() {
191            reader.abort();
192            if let Err(error) = reader.await
193                && !error.is_cancelled()
194            {
195                return Err(Error::ReaderTask(error));
196            }
197        }
198        Ok(())
199    }
200}
201
202async fn read_stdout(
203    mut stdout: BufReader<tokio::process::ChildStdout>,
204    sender: mpsc::Sender<Result<Message, Error>>,
205) {
206    loop {
207        let line = match read_line(&mut stdout, MAX_INBOUND_LINE_BYTES).await {
208            Ok(Some(line)) => line,
209            Ok(None) => return,
210            Err(error) => {
211                let _ = sender.send(Err(error)).await;
212                return;
213            }
214        };
215        let value = match serde_json::from_slice(&line) {
216            Ok(value) => value,
217            Err(error) => {
218                let _ = sender.send(Err(Error::Json(error))).await;
219                return;
220            }
221        };
222        let message = match kcode_jsonrpc_wire::parse(value) {
223            Ok(message) => message,
224            Err(error) => {
225                let _ = sender.send(Err(Error::Wire(error))).await;
226                return;
227            }
228        };
229        if sender.send(Ok(message)).await.is_err() {
230            return;
231        }
232    }
233}
234
235async fn read_line<R: tokio::io::AsyncBufRead + Unpin>(
236    reader: &mut R,
237    limit: usize,
238) -> Result<Option<Vec<u8>>, Error> {
239    let mut line = Vec::new();
240    loop {
241        let available = reader.fill_buf().await.map_err(Error::Io)?;
242        if available.is_empty() {
243            if line.is_empty() {
244                return Ok(None);
245            }
246            std::str::from_utf8(&line).map_err(Error::Utf8)?;
247            return Err(Error::IncompleteFrame);
248        }
249        let end = available.iter().position(|byte| *byte == b'\n');
250        let length = end.unwrap_or(available.len());
251        if line.len() + length > limit {
252            return Err(Error::InboundLineTooLong);
253        }
254        line.extend_from_slice(&available[..length]);
255        reader.consume(length + usize::from(end.is_some()));
256        if end.is_some() {
257            std::str::from_utf8(&line).map_err(Error::Utf8)?;
258            return Ok(Some(line));
259        }
260    }
261}
262
263async fn read_stderr(mut stderr: tokio::process::ChildStderr) -> io::Result<String> {
264    let mut captured = Vec::new();
265    let mut buffer = [0; 8192];
266    loop {
267        let count = stderr.read(&mut buffer).await?;
268        if count == 0 {
269            return Ok(captured.into_iter().map(sanitize).collect());
270        }
271        let keep = (MAX_STDERR_CAPTURE_BYTES - captured.len()).min(count);
272        captured.extend_from_slice(&buffer[..keep]);
273    }
274}
275
276fn sanitize(byte: u8) -> char {
277    if byte.is_ascii_graphic() || matches!(byte, b' ' | b'\n' | b'\r' | b'\t') {
278        byte as char
279    } else {
280        '?'
281    }
282}
283
284#[cfg(test)]
285mod tests {
286    use super::*;
287
288    #[tokio::test]
289    async fn reap_caches_exit_when_stderr_task_fails() {
290        let mut command = Command::new("sh");
291        command.arg("-c").arg("exit 0");
292        let mut rpc = StdioRpc::spawn(
293            command,
294            Config {
295                inbound_capacity: 1,
296            },
297        )
298        .unwrap();
299        rpc.stderr.as_ref().unwrap().abort();
300        assert!(matches!(rpc.reap().await, Err(Error::StderrTask(_))));
301        let exit = rpc.reap().await.unwrap();
302        assert_eq!(exit.status.code(), Some(0));
303        assert!(exit.stderr.is_empty());
304        rpc.join_reader().await.unwrap();
305    }
306}
kcode_jsonrpc_stdio/lib.rs

kcode_jsonrpc_stdio/
lib.rs
