Skip to main content

sqlite_graphrag/
stdin_helper.rs

1//! Stdin reader with timeout to prevent indefinite blocking when the
2//! upstream pipe is held open without sending data.
3//!
4//! Used by `remember --body-stdin` and `edit` body input to enforce a
5//! deadline (default 60s). When the timeout fires, the spawned reader
6//! thread is leaked because `std::io::stdin()` cannot be cancelled
7//! from outside; this is acceptable in error scenarios because the
8//! process is about to exit anyway.
9//!
10//! When stdin is attached to a terminal (interactive TTY), the function
11//! returns an `AppError::Internal` immediately with an actionable message
12//! instead of blocking for up to `secs` seconds waiting for EOF.
13
14use crate::errors::AppError;
15use std::io::{IsTerminal, Read};
16use std::sync::mpsc;
17use std::thread;
18use std::time::Duration;
19
20/// Reads stdin to a `String` with a hard deadline.
21///
22/// Returns `AppError::Internal` immediately when stdin is attached to a
23/// terminal (TTY) — the caller must redirect data via a pipe or file.
24///
25/// # Errors
26/// Returns `AppError::Internal` when stdin is a TTY, when the read does
27/// not finish within `secs` seconds, or `AppError::Io` when the
28/// underlying read fails.
29pub fn read_stdin_with_timeout(secs: u64) -> Result<String, AppError> {
30    if std::io::stdin().is_terminal() {
31        return Err(AppError::Internal(anyhow::anyhow!(
32            "stdin is attached to a terminal; pipe data via stdin \
33             (e.g. `echo ... | sqlite-graphrag ...` or `... < file`) \
34             or use --body instead of --body-stdin"
35        )));
36    }
37    let (tx, rx) = mpsc::channel::<std::io::Result<String>>();
38    thread::spawn(move || {
39        let mut buf = String::new();
40        let result = std::io::stdin().read_to_string(&mut buf).map(|_| buf);
41        let _ = tx.send(result);
42    });
43    match rx.recv_timeout(Duration::from_secs(secs)) {
44        Ok(Ok(buf)) => Ok(buf),
45        Ok(Err(e)) => Err(AppError::Io(e)),
46        Err(mpsc::RecvTimeoutError::Timeout) => Err(AppError::Internal(anyhow::anyhow!(
47            "stdin read timed out after {secs}s; pipe must close within timeout window"
48        ))),
49        Err(mpsc::RecvTimeoutError::Disconnected) => Err(AppError::Internal(anyhow::anyhow!(
50            "stdin reader thread disconnected unexpectedly"
51        ))),
52    }
53}
54
55#[cfg(test)]
56mod tests {
57    use super::*;
58    use std::time::Instant;
59
60    // Note: we cannot easily test the success path because tests inherit stdin
61    // from the test runner. We only assert the timeout path here.
62    #[test]
63    fn read_stdin_with_timeout_returns_internal_error_on_timeout() {
64        // 1s is enough — stdin in test runner is typically a tty or pipe with no input.
65        let start = Instant::now();
66        let result = read_stdin_with_timeout(1);
67        let elapsed = start.elapsed();
68        // We expect either a timeout (most cases), an immediate TTY error, or a
69        // successful EOF read (rare in CI environments).
70        match result {
71            Err(AppError::Internal(e)) => {
72                let msg = e.to_string();
73                // Accept both the TTY-detected error and the timeout error.
74                assert!(
75                    msg.contains("timed out") || msg.contains("terminal"),
76                    "unexpected internal error: {msg}"
77                );
78                // TTY path exits immediately; timeout path takes ~1s.
79                assert!(elapsed.as_secs_f64() < 2.5);
80            }
81            Ok(_) | Err(AppError::Io(_)) => {
82                // EOF reached before timeout — also acceptable in CI environments.
83            }
84            Err(other) => panic!("unexpected error variant: {other:?}"),
85        }
86    }
87
88    // TTY detection cannot be simulated in unit tests because the test runner
89    // always provides a non-TTY stdin (pipe). Empirical validation:
90    //   cargo run --release -- remember --body-stdin --name h1-test
91    // Expected: exits in <2s with "stdin is attached to a terminal" message.
92}