kcode_jsonrpc_stdio/
lib.rs

1use serde_json::{Map, Number, Value, json};
2use std::fmt::{Display, Formatter};
3use std::io;
4use std::process::{ExitStatus, Stdio};
5use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
6use tokio::process::{Child, ChildStdin, ChildStdout, Command};
7use tokio::task::JoinHandle;
8
9pub const MAX_INBOUND_LINE_BYTES: usize = 8 * 1024 * 1024;
10pub const MAX_STDERR_CAPTURE_BYTES: usize = 64 * 1024;
11
12#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
13pub struct RequestId(pub u64);
14
15#[derive(Clone, Debug, Eq, PartialEq)]
16pub enum PeerId {
17    Number(Number),
18    String(String),
19}
20
21impl PeerId {
22    fn into_value(self) -> Value {
23        match self {
24            Self::Number(number) => Value::Number(number),
25            Self::String(string) => Value::String(string),
26        }
27    }
28}
29
30#[derive(Clone, Debug, PartialEq)]
31pub struct RpcError {
32    pub code: i64,
33    pub message: String,
34    pub data: Option<Value>,
35}
36
37#[derive(Clone, Debug, PartialEq)]
38pub enum IncomingMessage {
39    Response {
40        id: PeerId,
41        result: Result<Value, RpcError>,
42    },
43    Request {
44        id: PeerId,
45        method: String,
46        params: Value,
47    },
48    Notification {
49        method: String,
50        params: Value,
51    },
52}
53
54#[derive(Clone, Debug)]
55pub struct ProcessExit {
56    pub status: ExitStatus,
57    pub stderr: String,
58}
59
60#[derive(Debug)]
61pub enum Error {
62    Io(io::Error),
63    Json(serde_json::Error),
64    InvalidMessage(String),
65    InboundLineTooLong,
66    ProcessExited(ProcessExit),
67    StderrTask(String),
68}
69
70impl Display for Error {
71    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
72        match self {
73            Self::Io(error) => write!(formatter, "stdio I/O failed: {error}"),
74            Self::Json(error) => write!(formatter, "JSON failed: {error}"),
75            Self::InvalidMessage(message) => {
76                write!(formatter, "invalid JSON-RPC message: {message}")
77            }
78            Self::InboundLineTooLong => write!(
79                formatter,
80                "JSON-RPC line exceeds {MAX_INBOUND_LINE_BYTES} bytes"
81            ),
82            Self::ProcessExited(exit) => write!(
83                formatter,
84                "child exited with {}: {}",
85                exit.status, exit.stderr
86            ),
87            Self::StderrTask(message) => write!(formatter, "stderr reader failed: {message}"),
88        }
89    }
90}
91
92impl std::error::Error for Error {
93    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
94        match self {
95            Self::Io(error) => Some(error),
96            Self::Json(error) => Some(error),
97            _ => None,
98        }
99    }
100}
101
102impl From<io::Error> for Error {
103    fn from(error: io::Error) -> Self {
104        Self::Io(error)
105    }
106}
107
108impl From<serde_json::Error> for Error {
109    fn from(error: serde_json::Error) -> Self {
110        Self::Json(error)
111    }
112}
113
114pub struct StdioRpc {
115    child: Child,
116    stdin: Option<ChildStdin>,
117    stdout: BufReader<ChildStdout>,
118    stderr_task: Option<JoinHandle<io::Result<Vec<u8>>>>,
119    next_id: u64,
120    exit: Option<ProcessExit>,
121}
122
123impl StdioRpc {
124    pub fn spawn(mut command: Command) -> Result<Self, Error> {
125        command
126            .stdin(Stdio::piped())
127            .stdout(Stdio::piped())
128            .stderr(Stdio::piped())
129            .kill_on_drop(true);
130        let mut child = command.spawn()?;
131        let stdin = child
132            .stdin
133            .take()
134            .ok_or_else(|| Error::InvalidMessage("child stdin was not piped".into()))?;
135        let stdout = child
136            .stdout
137            .take()
138            .ok_or_else(|| Error::InvalidMessage("child stdout was not piped".into()))?;
139        let mut stderr = child
140            .stderr
141            .take()
142            .ok_or_else(|| Error::InvalidMessage("child stderr was not piped".into()))?;
143        let stderr_task = tokio::spawn(async move {
144            let mut captured = Vec::new();
145            let mut buffer = [0_u8; 8192];
146            loop {
147                let read = stderr.read(&mut buffer).await?;
148                if read == 0 {
149                    return Ok(captured);
150                }
151                let remaining = MAX_STDERR_CAPTURE_BYTES.saturating_sub(captured.len());
152                captured.extend_from_slice(&buffer[..read.min(remaining)]);
153            }
154        });
155        Ok(Self {
156            child,
157            stdin: Some(stdin),
158            stdout: BufReader::new(stdout),
159            stderr_task: Some(stderr_task),
160            next_id: 1,
161            exit: None,
162        })
163    }
164
165    pub async fn send_request(
166        &mut self,
167        method: impl Into<String>,
168        params: Value,
169    ) -> Result<RequestId, Error> {
170        let id = RequestId(self.next_id);
171        self.next_id = self
172            .next_id
173            .checked_add(1)
174            .ok_or_else(|| Error::InvalidMessage("local request ID exhausted".into()))?;
175        self.send(json!({"jsonrpc": "2.0", "id": id.0, "method": method.into(), "params": params}))
176            .await?;
177        Ok(id)
178    }
179
180    pub async fn send_notification(
181        &mut self,
182        method: impl Into<String>,
183        params: Value,
184    ) -> Result<(), Error> {
185        self.send(json!({"jsonrpc": "2.0", "method": method.into(), "params": params}))
186            .await
187    }
188
189    pub async fn next(&mut self) -> Result<IncomingMessage, Error> {
190        let Some(line) = self.read_line().await? else {
191            return Err(Error::ProcessExited(self.finish_exit().await?));
192        };
193        parse_message(serde_json::from_slice(&line)?)
194    }
195
196    pub async fn respond(&mut self, id: PeerId, result: Value) -> Result<(), Error> {
197        self.send(json!({"jsonrpc": "2.0", "id": id.into_value(), "result": result}))
198            .await
199    }
200
201    pub async fn respond_error(
202        &mut self,
203        id: PeerId,
204        code: i64,
205        message: impl Into<String>,
206        data: Option<Value>,
207    ) -> Result<(), Error> {
208        let mut error = Map::new();
209        error.insert("code".into(), Value::Number(code.into()));
210        error.insert("message".into(), Value::String(message.into()));
211        if let Some(data) = data {
212            error.insert("data".into(), data);
213        }
214        self.send(json!({"jsonrpc": "2.0", "id": id.into_value(), "error": error}))
215            .await
216    }
217
218    pub async fn shutdown(mut self) -> Result<ProcessExit, Error> {
219        self.stdin.take();
220        if self.child.try_wait()?.is_none()
221            && let Err(error) = self.child.start_kill()
222            && self.child.try_wait()?.is_none()
223        {
224            return Err(Error::Io(error));
225        }
226        self.finish_exit().await
227    }
228
229    async fn send(&mut self, value: Value) -> Result<(), Error> {
230        let mut bytes = serde_json::to_vec(&value)?;
231        bytes.push(b'\n');
232        let stdin = self
233            .stdin
234            .as_mut()
235            .ok_or_else(|| Error::InvalidMessage("child stdin is closed".into()))?;
236        stdin.write_all(&bytes).await?;
237        stdin.flush().await?;
238        Ok(())
239    }
240
241    async fn read_line(&mut self) -> Result<Option<Vec<u8>>, Error> {
242        let mut line = Vec::new();
243        let mut too_long = false;
244        let mut saw_bytes = false;
245        loop {
246            let (consumed, ended, empty) = {
247                let available = self.stdout.fill_buf().await?;
248                if available.is_empty() {
249                    (0, false, true)
250                } else {
251                    saw_bytes = true;
252                    let newline = available.iter().position(|byte| *byte == b'\n');
253                    let content = newline.unwrap_or(available.len());
254                    if !too_long {
255                        if line.len() + content > MAX_INBOUND_LINE_BYTES {
256                            too_long = true;
257                        } else {
258                            line.extend_from_slice(&available[..content]);
259                        }
260                    }
261                    (
262                        content + usize::from(newline.is_some()),
263                        newline.is_some(),
264                        false,
265                    )
266                }
267            };
268            if empty {
269                if too_long {
270                    return Err(Error::InboundLineTooLong);
271                }
272                return if saw_bytes { Ok(Some(line)) } else { Ok(None) };
273            }
274            self.stdout.consume(consumed);
275            if ended {
276                return if too_long {
277                    Err(Error::InboundLineTooLong)
278                } else {
279                    Ok(Some(line))
280                };
281            }
282        }
283    }
284
285    async fn finish_exit(&mut self) -> Result<ProcessExit, Error> {
286        if let Some(exit) = &self.exit {
287            return Ok(exit.clone());
288        }
289        let status = self.child.wait().await?;
290        let task = self
291            .stderr_task
292            .take()
293            .ok_or_else(|| Error::StderrTask("stderr result was already consumed".into()))?;
294        let bytes = task
295            .await
296            .map_err(|error| Error::StderrTask(error.to_string()))??;
297        let exit = ProcessExit {
298            status,
299            stderr: String::from_utf8_lossy(&bytes).into_owned(),
300        };
301        self.exit = Some(exit.clone());
302        Ok(exit)
303    }
304}
305
306fn parse_message(value: Value) -> Result<IncomingMessage, Error> {
307    let Value::Object(mut fields) = value else {
308        return Err(Error::InvalidMessage("top level is not an object".into()));
309    };
310    if fields.remove("jsonrpc") != Some(Value::String("2.0".into())) {
311        return Err(Error::InvalidMessage("jsonrpc is not 2.0".into()));
312    }
313    let has_method = fields.contains_key("method");
314    let has_result = fields.contains_key("result");
315    let has_error = fields.contains_key("error");
316    if has_method {
317        if has_result || has_error {
318            return Err(Error::InvalidMessage(
319                "method message also contains a response".into(),
320            ));
321        }
322        let method = fields
323            .remove("method")
324            .and_then(|value| value.as_str().map(str::to_owned))
325            .ok_or_else(|| Error::InvalidMessage("method is not a string".into()))?;
326        let params = fields.remove("params").unwrap_or(Value::Null);
327        return match fields.remove("id") {
328            Some(id) => Ok(IncomingMessage::Request {
329                id: parse_peer_id(id)?,
330                method,
331                params,
332            }),
333            None => Ok(IncomingMessage::Notification { method, params }),
334        };
335    }
336    if fields.contains_key("params") {
337        return Err(Error::InvalidMessage("response contains params".into()));
338    }
339    if has_result == has_error {
340        return Err(Error::InvalidMessage(
341            "response must contain exactly one of result or error".into(),
342        ));
343    }
344    let id = fields
345        .remove("id")
346        .ok_or_else(|| Error::InvalidMessage("response has no id".into()))?;
347    let result = if has_result {
348        Ok(fields.remove("result").expect("result presence checked"))
349    } else {
350        Err(parse_rpc_error(
351            fields.remove("error").expect("error presence checked"),
352        )?)
353    };
354    Ok(IncomingMessage::Response {
355        id: parse_peer_id(id)?,
356        result,
357    })
358}
359
360fn parse_peer_id(value: Value) -> Result<PeerId, Error> {
361    match value {
362        Value::Number(number) if number.as_i64().is_some() || number.as_u64().is_some() => {
363            Ok(PeerId::Number(number))
364        }
365        Value::String(string) => Ok(PeerId::String(string)),
366        _ => Err(Error::InvalidMessage(
367            "id is not an integer or string".into(),
368        )),
369    }
370}
371
372fn parse_rpc_error(value: Value) -> Result<RpcError, Error> {
373    let Value::Object(mut fields) = value else {
374        return Err(Error::InvalidMessage("error is not an object".into()));
375    };
376    let code = fields
377        .remove("code")
378        .and_then(|value| value.as_i64())
379        .ok_or_else(|| Error::InvalidMessage("error code is not an integer".into()))?;
380    let message = fields
381        .remove("message")
382        .and_then(|value| value.as_str().map(str::to_owned))
383        .ok_or_else(|| Error::InvalidMessage("error message is not a string".into()))?;
384    Ok(RpcError {
385        code,
386        message,
387        data: fields.remove("data"),
388    })
389}
kcode_jsonrpc_stdio/lib.rs

kcode_jsonrpc_stdio/
lib.rs
