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