Skip to main content

lean_ctx/shell/
exec.rs

1use std::io::{self, IsTerminal, Write};
2use std::process::{Command, Stdio};
3
4use crate::core::config;
5use crate::core::slow_log;
6use crate::core::stats;
7use crate::core::tokens::count_tokens;
8
9/// Execute a command from pre-split argv without going through `sh -c`.
10/// Used by `-t` mode when the shell hook passes `"$@"` — arguments are
11/// already correctly split by the user's shell, so re-serializing them
12/// into a string and re-parsing via `sh -c` would risk mangling complex
13/// quoted arguments (em-dashes, `#`, nested quotes, etc.).
14pub fn exec_argv(args: &[String]) -> i32 {
15    if args.is_empty() {
16        return 127;
17    }
18
19    if std::env::var("LEAN_CTX_DISABLED").is_ok() || std::env::var("LEAN_CTX_ACTIVE").is_ok() {
20        return exec_direct(args);
21    }
22
23    let joined = super::platform::join_command(args);
24    let cfg = config::Config::load();
25
26    if super::compress::is_excluded_command(&joined, &cfg.excluded_commands) {
27        return exec_direct(args);
28    }
29
30    let code = exec_direct(args);
31    stats::record(&joined, 0, 0);
32    code
33}
34
35fn exec_direct(args: &[String]) -> i32 {
36    let status = Command::new(&args[0])
37        .args(&args[1..])
38        .env("LEAN_CTX_ACTIVE", "1")
39        .stdin(Stdio::inherit())
40        .stdout(Stdio::inherit())
41        .stderr(Stdio::inherit())
42        .status();
43
44    match status {
45        Ok(s) => s.code().unwrap_or(1),
46        Err(e) => {
47            tracing::error!("lean-ctx: failed to execute: {e}");
48            127
49        }
50    }
51}
52
53pub fn exec(command: &str) -> i32 {
54    let (shell, shell_flag) = super::platform::shell_and_flag();
55    let command = crate::tools::ctx_shell::normalize_command_for_shell(command);
56    let command = command.as_str();
57
58    if std::env::var("LEAN_CTX_DISABLED").is_ok() || std::env::var("LEAN_CTX_ACTIVE").is_ok() {
59        return exec_inherit(command, &shell, &shell_flag);
60    }
61
62    let cfg = config::Config::load();
63    let force_compress = std::env::var("LEAN_CTX_COMPRESS").is_ok();
64    let raw_mode = std::env::var("LEAN_CTX_RAW").is_ok();
65
66    if raw_mode
67        || (!force_compress
68            && super::compress::is_excluded_command(command, &cfg.excluded_commands))
69    {
70        return exec_inherit(command, &shell, &shell_flag);
71    }
72
73    if !force_compress {
74        if io::stdout().is_terminal() {
75            return exec_inherit_tracked(command, &shell, &shell_flag);
76        }
77        return exec_inherit(command, &shell, &shell_flag);
78    }
79
80    exec_buffered(command, &shell, &shell_flag, &cfg)
81}
82
83fn exec_inherit(command: &str, shell: &str, shell_flag: &str) -> i32 {
84    let status = Command::new(shell)
85        .arg(shell_flag)
86        .arg(command)
87        .env("LEAN_CTX_ACTIVE", "1")
88        .stdin(Stdio::inherit())
89        .stdout(Stdio::inherit())
90        .stderr(Stdio::inherit())
91        .status();
92
93    match status {
94        Ok(s) => s.code().unwrap_or(1),
95        Err(e) => {
96            tracing::error!("lean-ctx: failed to execute: {e}");
97            127
98        }
99    }
100}
101
102fn exec_inherit_tracked(command: &str, shell: &str, shell_flag: &str) -> i32 {
103    let code = exec_inherit(command, shell, shell_flag);
104    stats::record(command, 0, 0);
105    code
106}
107
108fn combine_output(stdout: &str, stderr: &str) -> String {
109    if stderr.is_empty() {
110        stdout.to_string()
111    } else if stdout.is_empty() {
112        stderr.to_string()
113    } else {
114        format!("{stdout}\n{stderr}")
115    }
116}
117
118fn exec_buffered(command: &str, shell: &str, shell_flag: &str, cfg: &config::Config) -> i32 {
119    #[cfg(windows)]
120    super::platform::set_console_utf8();
121
122    let start = std::time::Instant::now();
123
124    let mut cmd = Command::new(shell);
125    cmd.arg(shell_flag);
126
127    #[cfg(windows)]
128    {
129        let is_powershell =
130            shell.to_lowercase().contains("powershell") || shell.to_lowercase().contains("pwsh");
131        if is_powershell {
132            cmd.arg(format!(
133                "[Console]::OutputEncoding = [System.Text.Encoding]::UTF8; {command}"
134            ));
135        } else {
136            cmd.arg(command);
137        }
138    }
139    #[cfg(not(windows))]
140    cmd.arg(command);
141
142    let child = cmd
143        .env("LEAN_CTX_ACTIVE", "1")
144        .env_remove("DISPLAY")
145        .env_remove("XAUTHORITY")
146        .env_remove("WAYLAND_DISPLAY")
147        .stdout(Stdio::piped())
148        .stderr(Stdio::piped())
149        .spawn();
150
151    let child = match child {
152        Ok(c) => c,
153        Err(e) => {
154            tracing::error!("lean-ctx: failed to execute: {e}");
155            return 127;
156        }
157    };
158
159    let output = match child.wait_with_output() {
160        Ok(o) => o,
161        Err(e) => {
162            tracing::error!("lean-ctx: failed to wait: {e}");
163            return 127;
164        }
165    };
166
167    let duration_ms = start.elapsed().as_millis();
168    let exit_code = output.status.code().unwrap_or(1);
169    let stdout = super::platform::decode_output(&output.stdout);
170    let stderr = super::platform::decode_output(&output.stderr);
171
172    let full_output = combine_output(&stdout, &stderr);
173    let input_tokens = count_tokens(&full_output);
174
175    let (compressed, output_tokens) =
176        super::compress::compress_and_measure(command, &stdout, &stderr);
177
178    stats::record(command, input_tokens, output_tokens);
179
180    if !compressed.is_empty() {
181        let _ = io::stdout().write_all(compressed.as_bytes());
182        if !compressed.ends_with('\n') {
183            let _ = io::stdout().write_all(b"\n");
184        }
185    }
186    let should_tee = match cfg.tee_mode {
187        config::TeeMode::Always => !full_output.trim().is_empty(),
188        config::TeeMode::Failures => exit_code != 0 && !full_output.trim().is_empty(),
189        config::TeeMode::Never => false,
190    };
191    if should_tee {
192        if let Some(path) = super::redact::save_tee(command, &full_output) {
193            eprintln!("[lean-ctx: full output -> {path} (redacted, 24h TTL)]");
194        }
195    }
196
197    let threshold = cfg.slow_command_threshold_ms;
198    if threshold > 0 && duration_ms >= threshold as u128 {
199        slow_log::record(command, duration_ms, exit_code);
200    }
201
202    exit_code
203}
204
205#[cfg(test)]
206mod exec_tests {
207    #[test]
208    fn exec_direct_runs_true() {
209        let code = super::exec_direct(&["true".to_string()]);
210        assert_eq!(code, 0);
211    }
212
213    #[test]
214    fn exec_direct_runs_false() {
215        let code = super::exec_direct(&["false".to_string()]);
216        assert_ne!(code, 0);
217    }
218
219    #[test]
220    fn exec_direct_preserves_args_with_special_chars() {
221        let code = super::exec_direct(&[
222            "echo".to_string(),
223            "hello world".to_string(),
224            "it's here".to_string(),
225            "a \"quoted\" thing".to_string(),
226        ]);
227        assert_eq!(code, 0);
228    }
229
230    #[test]
231    fn exec_direct_nonexistent_returns_127() {
232        let code = super::exec_direct(&["__nonexistent_binary_12345__".to_string()]);
233        assert_eq!(code, 127);
234    }
235
236    #[test]
237    fn exec_argv_empty_returns_127() {
238        let code = super::exec_argv(&[]);
239        assert_eq!(code, 127);
240    }
241
242    #[test]
243    fn exec_argv_runs_simple_command() {
244        let code = super::exec_argv(&["true".to_string()]);
245        assert_eq!(code, 0);
246    }
247
248    #[test]
249    fn exec_argv_passes_through_when_disabled() {
250        std::env::set_var("LEAN_CTX_DISABLED", "1");
251        let code = super::exec_argv(&["true".to_string()]);
252        std::env::remove_var("LEAN_CTX_DISABLED");
253        assert_eq!(code, 0);
254    }
255}