molo_agent/message_channel/
cli.rs

1//! A command-line message channel: prints messages to the terminal and reads replies from stdin.
2
3use super::{ChannelError, MessageChannel};
4use std::fmt;
5use tokio::io::{
6    AsyncBufRead, AsyncBufReadExt, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader,
7};
8use tokio::sync::Mutex;
9
10/// Upper bound for a single input line (1 MiB): unbounded appends would exhaust memory when piped
11/// input has no newline for a long time.
12const MAX_LINE: usize = 1 << 20;
13
14/// A command-line message channel: prints messages to the terminal and reads one line from stdin as the reply.
15///
16/// Binds to standard input / standard output by default; alternatively, inject custom read/write sources
17/// with [`CliMessageChannel::with_io`] (files, network, in-memory buffers — common in tests).
18///
19/// Messages travel line by line: each message gets its own line on output (a newline is appended
20/// automatically); `ask` reads one line from the input, trims leading and trailing whitespace,
21/// and returns it as the reply.
22///
23/// All operations are serialized through an internal mutex: `ask` presents only one request at a time,
24/// and other `ask` / `notify` calls queue while it waits for the reply; input is read asynchronously,
25/// yielding while waiting, so other tasks are not blocked on a single-threaded runtime.
26///
27/// # Example
28///
29/// ```rust
30/// # extern crate molo_agent as molo;
31/// use molo::{CliMessageChannel, MessageChannel};
32/// use tokio::io::{AsyncWriteExt, BufReader};
33///
34/// # #[tokio::main]
35/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
36/// // Inject in-memory read/write sources instead of touching a real terminal.
37/// let (mut seed, rx) = tokio::io::duplex(64);
38/// seed.write_all("yes\n".as_bytes()).await?;
39/// let channel = CliMessageChannel::with_io(
40///     Box::new(BufReader::new(rx)),
41///     Box::new(tokio::io::sink()),
42/// );
43///
44/// let answer = channel.ask("continue? [y/n]").await?;
45/// assert_eq!(answer, "yes");
46/// # Ok(())
47/// # }
48/// ```
49///
50/// # Notes
51///
52/// - when input reaches end-of-file (Ctrl-D / EOF), `ask` returns
53///   [`ChannelError::Closed`](crate::ChannelError::Closed);
54/// - replies are read line by line and trimmed: an empty line still counts as a valid reply (an empty string).
55pub struct CliMessageChannel {
56    io: Mutex<CliIo>,
57}
58
59impl fmt::Debug for CliMessageChannel {
60    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
61        // The read/write sources are trait objects and can't implement Debug; print the bound shapes
62        // so logs can tell how the channel was constructed.
63        f.debug_struct("CliMessageChannel")
64            .field("io", &"Box<dyn AsyncBufRead> + Box<dyn AsyncWrite>")
65            .finish()
66    }
67}
68
69struct CliIo {
70    reader: Box<dyn AsyncBufRead + Unpin + Send>,
71    writer: Box<dyn AsyncWrite + Unpin + Send>,
72}
73
74impl CliMessageChannel {
75    /// Binds to standard input / standard output.
76    pub fn new() -> Self {
77        Self::with_io(
78            Box::new(BufReader::new(tokio::io::stdin())),
79            Box::new(tokio::io::stdout()),
80        )
81    }
82
83    /// Binds custom read/write sources.
84    ///
85    /// `reader` supplies the reply input (read line by line), and `writer` receives the output messages
86    /// (written line by line and flushed). Tests can inject an in-memory buffer, or the channel can be
87    /// attached to a file or the network.
88    pub fn with_io(
89        reader: Box<dyn AsyncBufRead + Unpin + Send>,
90        writer: Box<dyn AsyncWrite + Unpin + Send>,
91    ) -> Self {
92        Self {
93            io: Mutex::new(CliIo { reader, writer }),
94        }
95    }
96}
97
98impl Default for CliMessageChannel {
99    fn default() -> Self {
100        Self::new()
101    }
102}
103
104/// Writes one line of message to the writer and flushes it.
105async fn write_line(
106    writer: &mut (dyn AsyncWrite + Unpin + Send),
107    message: &str,
108) -> Result<(), ChannelError> {
109    writer
110        .write_all(message.as_bytes())
111        .await
112        .map_err(ChannelError::from)?;
113    writer.write_all(b"\n").await.map_err(ChannelError::from)?;
114    writer.flush().await.map_err(ChannelError::from)
115}
116
117#[async_trait::async_trait]
118impl MessageChannel for CliMessageChannel {
119    async fn ask(&self, message: &str) -> Result<String, ChannelError> {
120        let mut io = self.io.lock().await;
121        write_line(io.writer.as_mut(), message).await?;
122        // Bounded read: lines longer than MAX_LINE bytes are truncated and rejected (bounded memory when
123        // piped input has no newline; the error includes the length for diagnosis).
124        let mut line = String::new();
125        let read = io
126            .reader
127            .as_mut()
128            .take((MAX_LINE + 1) as u64)
129            .read_line(&mut line)
130            .await
131            .map_err(ChannelError::from)?;
132        if line.len() > MAX_LINE {
133            return Err(ChannelError::Io(format!(
134                "input line exceeds {MAX_LINE} bytes"
135            )));
136        }
137        if read == 0 {
138            return Err(ChannelError::Closed);
139        }
140        Ok(line.trim().to_string())
141    }
142
143    async fn notify(&self, message: &str) -> Result<(), ChannelError> {
144        let mut io = self.io.lock().await;
145        write_line(io.writer.as_mut(), message).await
146    }
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152    use std::pin::Pin;
153    use std::sync::{Arc, Mutex as StdMutex};
154    use std::task::{Context, Poll};
155
156    /// Shared output buffer: collects written content across tasks in concurrency tests.
157    #[derive(Clone, Default)]
158    struct SharedBuf(Arc<StdMutex<Vec<u8>>>);
159
160    impl AsyncWrite for SharedBuf {
161        fn poll_write(
162            self: Pin<&mut Self>,
163            _cx: &mut Context<'_>,
164            buf: &[u8],
165        ) -> Poll<std::io::Result<usize>> {
166            self.0
167                .lock()
168                .expect("internal lock poisoned")
169                .extend_from_slice(buf);
170            Poll::Ready(Ok(buf.len()))
171        }
172
173        fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
174            Poll::Ready(Ok(()))
175        }
176
177        fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
178            Poll::Ready(Ok(()))
179        }
180    }
181
182    /// Builds a channel from the injected reader / writer; seed is the pre-written input content.
183    async fn channel(seed: &str) -> (CliMessageChannel, SharedBuf) {
184        let buf = SharedBuf(Arc::new(StdMutex::new(Vec::new())));
185        let (mut seed_tx, rx) = tokio::io::duplex(64);
186        seed_tx.write_all(seed.as_bytes()).await.unwrap();
187        let channel =
188            CliMessageChannel::with_io(Box::new(BufReader::new(rx)), Box::new(buf.clone()));
189        (channel, buf)
190    }
191
192    #[tokio::test]
193    async fn ask_reads_line_and_trims() {
194        let (channel, buf) = channel("  reply text  \n").await;
195        assert_eq!(channel.ask("question").await.unwrap(), "reply text");
196        let out = String::from_utf8(buf.0.lock().unwrap().clone()).unwrap();
197        assert_eq!(out, "question\n");
198    }
199
200    #[tokio::test]
201    async fn ask_closed_input_returns_closed() {
202        let channel = CliMessageChannel::with_io(
203            Box::new(BufReader::new(tokio::io::empty())),
204            Box::new(tokio::io::sink()),
205        );
206        assert!(matches!(
207            channel.ask("question").await,
208            Err(ChannelError::Closed)
209        ));
210    }
211
212    #[tokio::test]
213    async fn notify_writes_message() {
214        let (channel, buf) = channel("").await;
215        channel.notify("notice").await.unwrap();
216        let out = String::from_utf8(buf.0.lock().unwrap().clone()).unwrap();
217        assert_eq!(out, "notice\n");
218    }
219
220    #[tokio::test]
221    async fn concurrent_asks_serialized() {
222        let (channel, buf) = channel("r1\nr2\n").await;
223        // Two asks run concurrently; the channel presents them serially, so replies line up one-to-one
224        // without interleaving.
225        let (a, b) = tokio::join!(channel.ask("m1"), channel.ask("m2"));
226        let mut replies = vec![a.unwrap(), b.unwrap()];
227        replies.sort();
228        assert_eq!(replies, vec!["r1", "r2"]);
229        // The output is also non-interleaved: two lines, each a complete message.
230        let out = String::from_utf8(buf.0.lock().unwrap().clone()).unwrap();
231        let lines: Vec<&str> = out.lines().collect();
232        assert_eq!(lines.len(), 2);
233        assert!(lines.contains(&"m1") && lines.contains(&"m2"));
234    }
235}
molo_agent/message_channel/cli.rs

molo_agent/message_channel/
cli.rs
