1use std::io;
17
18use agent_client_protocol::{Lines, schema::v1::Error};
19use futures::{AsyncBufReadExt, AsyncWriteExt, Sink, Stream, StreamExt};
20use serde_json::Value;
21use thiserror::Error as ThisError;
22
23use crate::server::{ServeConfig, serve};
24
25#[derive(Debug, ThisError)]
27pub enum StdioError {
28 #[error("the first line on stdin was not a JSON-RPC message")]
34 NotAClient,
35
36 #[error(transparent)]
37 Protocol(#[from] Error),
38
39 #[error("failed to read stdin: {0}")]
40 Stdin(#[from] io::Error),
41}
42
43pub async fn serve_stdio(config: ServeConfig) -> Result<(), StdioError> {
53 let stdin = blocking::Unblock::new(io::stdin());
54 let mut lines = Box::pin(futures::io::BufReader::new(stdin).lines());
55
56 let (said, opening) = opening_lines(&mut lines).await?;
57 if opening == Opening::NotAClient {
58 return Err(StdioError::NotAClient);
59 }
60
61 let incoming = futures::stream::iter(said.into_iter().map(Ok)).chain(lines);
62 serve(config, Lines::new(stdout_lines(), incoming)).await?;
63
64 Ok(())
65}
66
67#[derive(Debug, Clone, Copy, PartialEq, Eq)]
69enum Opening {
70 Client,
73 NotAClient,
75 Silence,
79}
80
81async fn opening_lines<S>(lines: &mut S) -> io::Result<(Vec<String>, Opening)>
87where
88 S: Stream<Item = io::Result<String>> + Unpin,
89{
90 let mut said = Vec::new();
91
92 while let Some(line) = lines.next().await {
93 let line = line?;
94 let opening = (!line.trim().is_empty()).then(|| opening_for(&line));
98 said.push(line);
99
100 if let Some(opening) = opening {
101 return Ok((said, opening));
102 }
103 }
104
105 Ok((said, Opening::Silence))
106}
107
108fn opening_for(line: &str) -> Opening {
116 match serde_json::from_str::<Value>(line) {
117 Ok(Value::Object(_) | Value::Array(_)) => Opening::Client,
118 _ => Opening::NotAClient,
119 }
120}
121
122fn stdout_lines() -> impl Sink<String, Error = io::Error> + Send + 'static {
127 futures::sink::unfold(
128 blocking::Unblock::new(io::stdout()),
129 async move |mut out: blocking::Unblock<io::Stdout>, line: String| {
130 let mut bytes = line.into_bytes();
131 bytes.push(b'\n');
132 out.write_all(&bytes).await?;
133 out.flush().await?;
134 Ok::<_, io::Error>(out)
135 },
136 )
137}
138
139#[cfg(test)]
140mod tests {
141 use super::*;
142
143 async fn opening_of(lines: &[&str]) -> (Vec<String>, Opening) {
144 let mut stream = futures::stream::iter(
145 lines
146 .iter()
147 .map(|line| Ok((*line).to_string()))
148 .collect::<Vec<_>>(),
149 );
150
151 opening_lines(&mut stream)
152 .await
153 .expect("an in-memory stream does not fail")
154 }
155
156 #[test]
157 fn a_json_rpc_message_is_an_object_or_a_batch() {
158 assert_eq!(
159 opening_for(r#"{"jsonrpc":"2.0","id":1,"method":"initialize"}"#),
160 Opening::Client
161 );
162 assert_eq!(
163 opening_for(r#"[{"jsonrpc":"2.0","id":1}]"#),
164 Opening::Client
165 );
166 }
167
168 #[test]
169 fn a_client_is_not_refused_over_a_field_lan_does_not_own() {
170 assert_eq!(opening_for(r#"{"method":"initialize"}"#), Opening::Client);
175 assert_eq!(opening_for("{}"), Opening::Client);
176 }
177
178 #[test]
179 fn prose_is_not_a_client() {
180 for line in [
181 "fix the failing test",
182 "run the tests and summarize",
183 "why is CI red?",
184 ] {
185 assert_eq!(opening_for(line), Opening::NotAClient, "{line}");
186 }
187 }
188
189 #[tokio::test]
190 async fn the_first_line_is_handed_to_the_server_unread() {
191 let message = r#"{"jsonrpc":"2.0","id":1,"method":"initialize"}"#;
194
195 let (said, opening) = opening_of(&[message, r#"{"jsonrpc":"2.0","id":2}"#]).await;
196
197 assert_eq!(opening, Opening::Client);
198 assert_eq!(said, vec![message.to_string()]);
199 }
200
201 #[tokio::test]
202 async fn blank_lines_before_a_message_are_kept_and_skipped_over() {
203 let message = r#"{"jsonrpc":"2.0","id":1}"#;
204
205 let (said, opening) = opening_of(&["", " ", message]).await;
206
207 assert_eq!(opening, Opening::Client);
208 assert_eq!(
209 said,
210 vec!["".to_string(), " ".to_string(), message.to_string()],
211 "the layer that answered blank lines before must still see them"
212 );
213 }
214
215 #[tokio::test]
216 async fn a_peer_that_says_nothing_is_not_accused_of_anything() {
217 let (said, opening) = opening_of(&[]).await;
218
219 assert_eq!(opening, Opening::Silence);
220 assert!(said.is_empty());
221 }
222
223 #[tokio::test]
224 async fn prose_on_the_first_line_stops_before_the_server_starts() {
225 let (said, opening) = opening_of(&["fix the failing test", "and push"]).await;
226
227 assert_eq!(opening, Opening::NotAClient);
228 assert_eq!(said, vec!["fix the failing test".to_string()]);
229 }
230}