Skip to main content

a_agent/context/
stdin.rs

1use std::io::IsTerminal;
2
3/// Whether stdin holds input for this turn.
4///
5/// `a` consumes stdin so `cargo test 2>&1 | a "fix this"` works. Consuming it
6/// whenever it merely is not a terminal is wrong: a supervisor or sandbox hands
7/// its child a socket that may never carry anything, and reading it blocks
8/// forever.
9///
10/// A pipe or a redirected file is read to end of file, however long the producer
11/// takes, because writing the pipe is how the user asked for that. Waiting is
12/// what every filter does; giving up early would silently drop the input of a
13/// producer that is simply slow to start.
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum StdinSource {
16    Terminal,
17    /// Not a pipe or file: an inherited channel that is not this turn's input.
18    Foreign,
19    Stream,
20}
21
22pub fn stdin_source() -> StdinSource {
23    let stdin = std::io::stdin();
24    if stdin.is_terminal() {
25        return StdinSource::Terminal;
26    }
27    #[cfg(unix)]
28    {
29        use std::os::fd::AsRawFd;
30        classify_fd(stdin.as_raw_fd())
31    }
32    #[cfg(not(unix))]
33    StdinSource::Stream
34}
35
36/// A shell pipeline hands over a pipe, and a redirection hands over a file.
37/// Anything else, a socket in particular, came from a supervisor or sandbox that
38/// is still using it for something else.
39#[cfg(unix)]
40pub fn classify_fd(fd: std::os::fd::RawFd) -> StdinSource {
41    let mut stat = unsafe { std::mem::zeroed::<libc::stat>() };
42    if unsafe { libc::fstat(fd, &mut stat) } != 0 {
43        return StdinSource::Foreign;
44    }
45    let kind = stat.st_mode & libc::S_IFMT;
46    if kind == libc::S_IFIFO || kind == libc::S_IFREG {
47        StdinSource::Stream
48    } else {
49        StdinSource::Foreign
50    }
51}
52
53pub fn bound_stdin(input: &[u8], max_bytes: usize) -> String {
54    if input.len() <= max_bytes {
55        return String::from_utf8_lossy(input).into_owned();
56    }
57    let mut start = input.len().saturating_sub(max_bytes);
58    while start < input.len() && (input[start] & 0b1100_0000) == 0b1000_0000 {
59        start += 1;
60    }
61    format!(
62        "[stdin truncated; showing last {max_bytes} bytes]\n{}",
63        String::from_utf8_lossy(&input[start..])
64    )
65}
66
67#[cfg(all(test, unix))]
68mod tests {
69    use super::*;
70    use std::os::fd::AsRawFd;
71
72    #[test]
73    fn only_pipes_and_files_count_as_this_turn_input() {
74        let file = tempfile::NamedTempFile::new().unwrap();
75        assert_eq!(classify_fd(file.as_file().as_raw_fd()), StdinSource::Stream);
76
77        let mut fds = [0_i32; 2];
78        assert_eq!(unsafe { libc::pipe(fds.as_mut_ptr()) }, 0);
79        assert_eq!(classify_fd(fds[0]), StdinSource::Stream);
80
81        // A socket is how a supervisor or sandbox hands over stdio; draining it
82        // would block on a writer that is not talking to us.
83        let mut pair = [0_i32; 2];
84        assert_eq!(
85            unsafe { libc::socketpair(libc::AF_UNIX, libc::SOCK_STREAM, 0, pair.as_mut_ptr()) },
86            0
87        );
88        assert_eq!(classify_fd(pair[0]), StdinSource::Foreign);
89
90        assert_eq!(classify_fd(-1), StdinSource::Foreign);
91        for fd in fds.into_iter().chain(pair) {
92            unsafe { libc::close(fd) };
93        }
94    }
95}