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-from-stdin and `edit` body input to enforce a
5//! deadline ([`crate::constants::DEFAULT_STDIN_READ_TIMEOUT_SECS`], override
6//! via XDG `cli.stdin_timeout_secs`). When the timeout fires, the spawned
7//! reader thread is leaked because `std::io::stdin()` cannot be cancelled
8//! from outside; this is acceptable in error scenarios because the
9//! process is about to exit anyway.
10//!
11//! Two refusals happen before any read is attempted:
12//!
13//! * `--no-input` (or XDG `cli.no_input`) makes the refusal DECLARATIVE — it
14//!   holds even when a pipe is attached and would have supplied data;
15//! * an interactive TTY makes it EMERGENT — there is no producer, so waiting
16//!   for EOF would just burn the whole deadline.
17
18use crate::errors::AppError;
19use std::io::{IsTerminal, Read};
20use std::sync::atomic::{AtomicBool, Ordering};
21use std::sync::mpsc;
22use std::thread;
23use std::time::Duration;
24
25/// Whether this process refuses to read stdin (`--no-input` / XDG `cli.no_input`).
26///
27/// An atomic rather than a `OnceLock` so tests can exercise both branches in
28/// the same binary; production installs it exactly once from CLI bootstrap.
29static NO_INPUT: AtomicBool = AtomicBool::new(false);
30
31/// Installs the resolved `--no-input` decision for the whole process.
32///
33/// Called from [`crate::cli::Cli::validate_flags`], the one bootstrap hook that
34/// runs after XDG initialisation and before any subcommand dispatch.
35pub fn install_no_input(enabled: bool) {
36    NO_INPUT.store(enabled, Ordering::Release);
37}
38
39/// Whether stdin reads are refused for this invocation.
40pub fn no_input() -> bool {
41    NO_INPUT.load(Ordering::Acquire)
42}
43
44/// Reads stdin to a `String` with the configured deadline.
45///
46/// Resolves the timeout from XDG `cli.stdin_timeout_secs`, falling back to
47/// [`crate::constants::DEFAULT_STDIN_READ_TIMEOUT_SECS`]. Prefer this over
48/// [`read_stdin_with_timeout`] at call sites that have no reason to pick a
49/// different budget.
50///
51/// # Errors
52/// Same as [`read_stdin_with_timeout`].
53pub fn read_stdin() -> Result<String, AppError> {
54    read_stdin_with_timeout(crate::runtime_config::stdin_timeout_secs())
55}
56
57/// Reads stdin to a `String` with a hard deadline.
58///
59/// Returns `AppError::Validation` immediately when `--no-input` is in force,
60/// and `AppError::Internal` immediately when stdin is attached to a terminal
61/// (TTY) — the caller must redirect data via a pipe or file.
62///
63/// # Errors
64/// Returns `AppError::Validation` when `--no-input` is in force,
65/// `AppError::Internal` when stdin is a TTY, when the read does
66/// not finish within `secs` seconds, or `AppError::Io` when the
67/// underlying read fails.
68pub fn read_stdin_with_timeout(secs: u64) -> Result<String, AppError> {
69    if no_input() {
70        return Err(AppError::Validation(
71            crate::i18n::validation::no_input_blocks_stdin(),
72        ));
73    }
74    if std::io::stdin().is_terminal() {
75        return Err(AppError::Internal(anyhow::anyhow!(
76            "stdin is attached to a terminal; pipe data via stdin \
77             (e.g. `echo ... | sqlite-graphrag ...` or `... < file`) \
78             or use --body instead of the stdin body flag"
79        )));
80    }
81    let (tx, rx) = mpsc::channel::<std::io::Result<String>>();
82    thread::spawn(move || {
83        let mut buf = String::new();
84        let result = std::io::stdin().read_to_string(&mut buf).map(|_| buf);
85        let _ = tx.send(result);
86    });
87    match rx.recv_timeout(Duration::from_secs(secs)) {
88        Ok(Ok(buf)) => Ok(buf),
89        Ok(Err(e)) => Err(AppError::Io(e)),
90        Err(mpsc::RecvTimeoutError::Timeout) => Err(AppError::Internal(anyhow::anyhow!(
91            "stdin read timed out after {secs}s; pipe must close within timeout window"
92        ))),
93        Err(mpsc::RecvTimeoutError::Disconnected) => Err(AppError::Internal(anyhow::anyhow!(
94            "stdin reader thread disconnected unexpectedly"
95        ))),
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102    use std::time::Instant;
103
104    /// Serialises the tests that flip the process-wide `NO_INPUT` flag.
105    static NO_INPUT_GUARD: std::sync::Mutex<()> = std::sync::Mutex::new(());
106
107    // Note: we cannot easily test the success path because tests inherit stdin
108    // from the test runner. We only assert the timeout path here.
109    #[test]
110    fn read_stdin_with_timeout_returns_internal_error_on_timeout() {
111        let _guard = NO_INPUT_GUARD.lock().unwrap_or_else(|e| e.into_inner());
112        install_no_input(false);
113        // 1s is enough — stdin in test runner is typically a tty or pipe with no input.
114        let start = Instant::now();
115        let result = read_stdin_with_timeout(1);
116        let elapsed = start.elapsed();
117        // We expect either a timeout (most cases), an immediate TTY error, or a
118        // successful EOF read (rare in CI environments).
119        match result {
120            Err(AppError::Internal(e)) => {
121                let msg = e.to_string();
122                // Accept both the TTY-detected error and the timeout error.
123                assert!(
124                    msg.contains("timed out") || msg.contains("terminal"),
125                    "unexpected internal error: {msg}"
126                );
127                // TTY path exits immediately; timeout path takes ~1s.
128                assert!(elapsed.as_secs_f64() < 2.5);
129            }
130            Ok(_) | Err(AppError::Io(_)) => {
131                // EOF reached before timeout — also acceptable in CI environments.
132            }
133            Err(other) => unreachable!("stdin test: expected Internal/Io, got {other:?}"),
134        }
135    }
136
137    #[test]
138    fn no_input_refuses_before_the_read_is_attempted() {
139        let _guard = NO_INPUT_GUARD.lock().unwrap_or_else(|e| e.into_inner());
140        install_no_input(true);
141        // A 600-second budget would dominate the elapsed time if the refusal
142        // happened after the read rather than before it.
143        let start = Instant::now();
144        let result = read_stdin_with_timeout(600);
145        let elapsed = start.elapsed();
146        install_no_input(false);
147        match result {
148            Err(AppError::Validation(msg)) => {
149                assert!(msg.contains("--no-input"), "unexpected message: {msg}");
150            }
151            other => unreachable!("expected Validation under --no-input, got {other:?}"),
152        }
153        assert!(
154            elapsed.as_secs_f64() < 1.0,
155            "refusal must precede the read, took {elapsed:?}"
156        );
157    }
158
159    #[test]
160    fn install_no_input_round_trips() {
161        let _guard = NO_INPUT_GUARD.lock().unwrap_or_else(|e| e.into_inner());
162        install_no_input(true);
163        assert!(no_input());
164        install_no_input(false);
165        assert!(!no_input());
166    }
167
168    // TTY detection cannot be simulated in unit tests because the test runner
169    // always provides a non-TTY stdin (pipe). Empirical validation:
170    //   cargo run --release -- remember --name h1-test  (with the stdin body flag)
171    // Expected: exits in <2s with "stdin is attached to a terminal" message.
172}