intelli_shell/utils/
process.rs

1use std::{
2    cmp::Ordering,
3    collections::BTreeMap,
4    env,
5    ffi::OsStr,
6    io::{self, Read, Write},
7    ops::Deref,
8    path::{Path, PathBuf},
9    process::{self, ExitStatus, Stdio},
10    sync::LazyLock,
11    time::Duration,
12};
13
14use color_eyre::eyre::Context;
15use ignore::WalkBuilder;
16use os_info::Info;
17use sysinfo::{Pid, System};
18use tokio::{
19    io::{AsyncBufReadExt, BufReader},
20    signal,
21};
22use wait_timeout::ChildExt;
23
24#[derive(Debug)]
25pub struct ShellInfo {
26    pub kind: ShellType,
27    pub version: Option<String>,
28}
29
30#[derive(Clone, Debug, PartialEq, Eq, strum::Display, strum::EnumString)]
31pub enum ShellType {
32    #[strum(serialize = "cmd", serialize = "cmd.exe")]
33    Cmd,
34    #[strum(serialize = "powershell", serialize = "powershell.exe")]
35    WindowsPowerShell,
36    #[strum(serialize = "pwsh", serialize = "pwsh.exe")]
37    PowerShellCore,
38    #[strum(to_string = "bash", serialize = "bash.exe")]
39    Bash,
40    #[strum(serialize = "sh")]
41    Sh,
42    #[strum(serialize = "fish")]
43    Fish,
44    #[strum(serialize = "zsh")]
45    Zsh,
46    #[strum(default, to_string = "{0}")]
47    Other(String),
48}
49
50static PARENT_SHELL_INFO: LazyLock<ShellInfo> = LazyLock::new(|| {
51    let pid = Pid::from_u32(process::id());
52
53    tracing::debug!("Retrieving info for pid {pid}");
54    let sys = System::new_all();
55
56    let parent_process = sys
57        .process(Pid::from_u32(process::id()))
58        .expect("Couldn't retrieve current process from pid")
59        .parent()
60        .and_then(|parent_pid| sys.process(parent_pid));
61
62    let Some(parent) = parent_process else {
63        let default = if cfg!(target_os = "windows") {
64            ShellType::WindowsPowerShell
65        } else {
66            ShellType::Sh
67        };
68        tracing::warn!("Couldn't detect shell, assuming {default}");
69        return ShellInfo {
70            kind: default,
71            version: None,
72        };
73    };
74
75    let parent_name = parent
76        .name()
77        .to_str()
78        .expect("Invalid parent shell name")
79        .trim()
80        .to_lowercase();
81
82    let kind = ShellType::try_from(parent_name.as_str()).expect("infallible");
83    tracing::info!("Detected shell: {kind}");
84
85    let exe_path = parent
86        .exe()
87        .map(|p| p.as_os_str())
88        .filter(|p| !p.is_empty())
89        .unwrap_or_else(|| parent_name.as_ref());
90    let version = get_shell_version(&kind, exe_path).inspect(|v| tracing::info!("Detected shell version: {v}"));
91
92    ShellInfo { kind, version }
93});
94
95/// A helper function to get the version from a shell's executable path
96fn get_shell_version(shell_kind: &ShellType, shell_path: impl AsRef<OsStr>) -> Option<String> {
97    // `cmd.exe` version is tied to the OS version, so we don't query it
98    if *shell_kind == ShellType::Cmd {
99        return None;
100    }
101
102    // Most shells respond to `--version`, except PowerShell
103    let mut command = std::process::Command::new(shell_path);
104    if matches!(shell_kind, ShellType::PowerShellCore | ShellType::WindowsPowerShell) {
105        command.args([
106            "-Command",
107            "'PowerShell {0} ({1} Edition)' -f $PSVersionTable.PSVersion, $PSVersionTable.PSEdition",
108        ]);
109    } else {
110        command.arg("--version");
111    }
112
113    // Configure pipes for stdout and stderr to capture the output manually
114    let mut child = match command.stdout(Stdio::piped()).stderr(Stdio::piped()).spawn() {
115        Ok(child) => child,
116        Err(err) => {
117            tracing::warn!("Failed to spawn shell process: {err}");
118            return None;
119        }
120    };
121
122    // Wait for the process to exit, with a timeout
123    match child.wait_timeout(Duration::from_millis(250)) {
124        // The command finished within the timeout period
125        Ok(Some(status)) => {
126            if status.success() {
127                let mut output = String::new();
128                // Read the output from the stdout pipe
129                if let Some(mut stdout) = child.stdout {
130                    stdout.read_to_string(&mut output).unwrap_or_default();
131                }
132                // Return just the first line of the output
133                Some(output.lines().next().unwrap_or("").trim().to_string()).filter(|v| !v.is_empty())
134            } else {
135                tracing::warn!("Shell version command failed with status: {}", status);
136                None
137            }
138        }
139        // The command timed out
140        Ok(None) => {
141            // Kill the child process to prevent it from running forever
142            if let Err(err) = child.kill() {
143                tracing::warn!("Failed to kill timed-out process: {err}");
144            }
145            tracing::warn!("Shell version command timed out");
146            None
147        }
148        // An error occurred while waiting
149        Err(err) => {
150            tracing::warn!("Error waiting for shell version command: {err}");
151            None
152        }
153    }
154}
155
156/// Retrieves information about the current shell, including its type and version
157pub fn get_shell_info() -> &'static ShellInfo {
158    PARENT_SHELL_INFO.deref()
159}
160
161/// Retrieves the current shell type
162pub fn get_shell_type() -> &'static ShellType {
163    &get_shell_info().kind
164}
165
166/// A helper function to get the version from an executable (e.g. git)
167pub fn get_executable_version(root_cmd: impl AsRef<OsStr>) -> Option<String> {
168    if root_cmd.as_ref().is_empty() {
169        return None;
170    }
171
172    // Most shells commands respond to `--version`
173    let mut child = std::process::Command::new(root_cmd)
174        .arg("--version")
175        .stdout(Stdio::piped())
176        .stderr(Stdio::piped())
177        .spawn()
178        .ok()?;
179
180    // Wait for the process to exit, with a timeout
181    match child.wait_timeout(Duration::from_millis(250)) {
182        Ok(Some(status)) if status.success() => {
183            let mut output = String::new();
184            if let Some(mut stdout) = child.stdout {
185                stdout.read_to_string(&mut output).unwrap_or_default();
186            }
187            Some(output.lines().next().unwrap_or("").trim().to_string()).filter(|v| !v.is_empty())
188        }
189        Ok(None) => {
190            if let Err(err) = child.kill() {
191                tracing::warn!("Failed to kill timed-out process: {err}");
192            }
193            None
194        }
195        _ => None,
196    }
197}
198
199static OS_INFO: LazyLock<Info> = LazyLock::new(|| {
200    let info = os_info::get();
201    tracing::info!("Detected OS: {info}");
202    info
203});
204
205/// Retrieves the operating system information
206pub fn get_os_info() -> &'static Info {
207    &OS_INFO
208}
209
210static WORING_DIR: LazyLock<String> = LazyLock::new(|| {
211    std::env::current_dir()
212        .inspect_err(|err| tracing::warn!("Couldn't retrieve current dir: {err}"))
213        .ok()
214        .and_then(|p| p.to_str().map(|s| s.to_owned()))
215        .unwrap_or_default()
216});
217
218/// Retrieves the working directory
219pub fn get_working_dir() -> &'static str {
220    WORING_DIR.deref()
221}
222
223/// Formats an env var name into its shell representation, based on the current shell
224pub fn format_env_var(var: impl AsRef<str>) -> String {
225    let var = var.as_ref();
226    match get_shell_type() {
227        ShellType::Cmd => format!("%{var}%"),
228        ShellType::WindowsPowerShell | ShellType::PowerShellCore => format!("$env:{var}"),
229        _ => format!("${var}"),
230    }
231}
232
233/// Generates a string representation of the current working directory tree, respecting .gitignore files
234pub fn generate_working_dir_tree(max_depth: usize, entry_limit: usize) -> Option<String> {
235    let root = PathBuf::from(get_working_dir());
236    if !root.is_dir() {
237        return None;
238    }
239
240    let root_canon = root.canonicalize().ok()?;
241
242    // Phase 1: Collect all entries by depth and also get total child counts for every directory
243    let mut entries_by_depth: BTreeMap<usize, Vec<ignore::DirEntry>> = BTreeMap::new();
244    let mut total_child_counts: BTreeMap<PathBuf, usize> = BTreeMap::new();
245    let walker = WalkBuilder::new(&root_canon).max_depth(Some(max_depth + 1)).build();
246
247    for entry in walker.flatten() {
248        if entry.depth() == 0 {
249            continue;
250        }
251        if let Some(parent_path) = entry.path().parent() {
252            *total_child_counts.entry(parent_path.to_path_buf()).or_default() += 1;
253        }
254        entries_by_depth.entry(entry.depth()).or_default().push(entry);
255    }
256
257    // Phase 2: Create a limited list of entries using the breadth-first approach
258    let mut limited_entries: Vec<ignore::DirEntry> = Vec::with_capacity(entry_limit);
259    'outer: for (_depth, entries) in entries_by_depth {
260        for entry in entries {
261            if limited_entries.len() >= entry_limit {
262                break 'outer;
263            }
264            limited_entries.push(entry);
265        }
266    }
267
268    // Phase 3: Populate the display tree and add "..." where contents are truncated
269    let mut dir_children: BTreeMap<PathBuf, Vec<(String, bool)>> = BTreeMap::new();
270    for entry in limited_entries {
271        let is_dir = entry.path().is_dir();
272        if let Some(parent_path) = entry.path().parent() {
273            let file_name = entry.file_name().to_string_lossy().to_string();
274            dir_children
275                .entry(parent_path.to_path_buf())
276                .or_default()
277                .push((file_name, is_dir));
278        }
279    }
280    for (path, total_count) in total_child_counts {
281        let displayed_count = dir_children.get(&path).map_or(0, |v| v.len());
282        if displayed_count < total_count {
283            dir_children.entry(path).or_default().push(("...".to_string(), false));
284        }
285    }
286
287    // Sort the children in each directory alphabetically for consistent output
288    for children in dir_children.values_mut() {
289        children.sort_by(|a, b| {
290            // "..." is always last
291            if a.0 == "..." {
292                Ordering::Greater
293            } else if b.0 == "..." {
294                Ordering::Less
295            } else {
296                // Otherwise, sort alphabetically
297                a.0.cmp(&b.0)
298            }
299        });
300    }
301
302    // Phase 4: Build the final string
303    let mut tree_string = format!("{} (current working dir)\n", root_canon.display());
304    build_tree_from_map(&root_canon, "", &mut tree_string, &dir_children);
305    Some(tree_string)
306}
307
308/// Recursively builds the tree string from the pre-compiled map of directory children
309fn build_tree_from_map(
310    dir_path: &Path,
311    prefix: &str,
312    output: &mut String,
313    dir_children: &BTreeMap<PathBuf, Vec<(String, bool)>>,
314) {
315    let Some(entries) = dir_children.get(dir_path) else {
316        return;
317    };
318
319    let mut iter = entries.iter().peekable();
320    while let Some((name, is_dir)) = iter.next() {
321        let is_last = iter.peek().is_none();
322        let connector = if is_last { "└── " } else { "├── " };
323        let new_prefix = format!("{prefix}{}", if is_last { "    " } else { "│   " });
324
325        if *is_dir {
326            // This is a directory; let's see if we can collapse it
327            let mut path_components = vec![name.clone()];
328            let mut current_path = dir_path.join(name);
329
330            // Keep collapsing as long as the current directory has only one child, which is also a directory
331            while let Some(children) = dir_children.get(&current_path) {
332                if children.len() == 1 {
333                    let (child_name, child_is_dir) = &children[0];
334                    if *child_is_dir {
335                        path_components.push(child_name.clone());
336                        current_path.push(child_name);
337                        // Continue to the next level of nesting
338                        continue;
339                    }
340                }
341                // Stop collapsing
342                break;
343            }
344
345            // Print the combined, collapsed path.
346            let collapsed_name = path_components.join("/");
347            output.push_str(&format!("{prefix}{connector}{collapsed_name}/\n"));
348
349            // Recurse using the final path in the chain
350            build_tree_from_map(&current_path, &new_prefix, output, dir_children);
351        } else {
352            // This is a file or "...", print it normally.
353            output.push_str(&format!("{prefix}{connector}{name}\n"));
354        }
355    }
356}
357
358/// Executes a shell command, inheriting the parent's `stdout` and `stderr`
359pub async fn execute_shell_command_inherit(command: &str, include_prompt: bool) -> color_eyre::Result<ExitStatus> {
360    let mut cmd = prepare_command_execution(command, include_prompt)?;
361
362    // Spawn the child process to get a handle to it
363    let mut child = cmd
364        .spawn()
365        .with_context(|| format!("Failed to spawn command: `{command}`"))?;
366
367    // Race the child process against a Ctrl+C signal
368    let status = tokio::select! {
369        // Prioritize Ctrl+C handler
370        biased;
371        // User presses Ctrl+C
372        _ = signal::ctrl_c() => {
373            tracing::info!("Received Ctrl+C, terminating child process...");
374            // Send a kill signal to the child process
375            child.kill().await.with_context(|| format!("Failed to kill child process for command: `{command}`"))?;
376            // Wait for the process to exit and get its status
377            child.wait().await.with_context(|| "Failed to await child process after kill")?
378        }
379        // The child process completes on its own
380        status = child.wait() => {
381            status.with_context(|| format!("Child process for command `{command}` failed"))?
382        }
383    };
384
385    Ok(status)
386}
387
388/// Executes a shell command, capturing `stdout` and `stderr`.
389///
390/// While capturing, it simultaneously prints both streams to the parent's `stderr` in real-time.
391pub async fn execute_shell_command_capture(
392    command: &str,
393    include_prompt: bool,
394) -> color_eyre::Result<(ExitStatus, String, bool)> {
395    let mut cmd = prepare_command_execution(command, include_prompt)?;
396
397    // Configure the command to capture output streams by creating pipes
398    cmd.stdout(Stdio::piped());
399    cmd.stderr(Stdio::piped());
400
401    let mut child = cmd
402        .spawn()
403        .with_context(|| format!("Failed to spawn command: `{command}`"))?;
404
405    // Create buffered readers for the child's output streams
406    let mut stdout_reader = BufReader::new(child.stdout.take().unwrap()).lines();
407    let mut stderr_reader = BufReader::new(child.stderr.take().unwrap()).lines();
408
409    let mut output_capture = String::new();
410
411    // Flag to track if the process was terminated by our signal handler
412    let mut terminated_by_signal = false;
413
414    // Use boolean flags to track when each stream is finished
415    let mut stdout_done = false;
416    let mut stderr_done = false;
417
418    // Loop until both stdout and stderr streams have been completely read
419    while !stdout_done || !stderr_done {
420        tokio::select! {
421            // Prioritize Ctrl+C handler
422            biased;
423            // User presses Ctrl+C
424            _ = signal::ctrl_c() => {
425                tracing::info!("Received Ctrl+C, terminating child process...");
426                // Kill the child process, this will also cause the stdout/stderr streams to close
427                child.kill().await.with_context(|| format!("Failed to kill child process for command: `{command}`"))?;
428                // Set the flag to true since we handled the signal
429                terminated_by_signal = true;
430                // Break the loop to proceed to the final `child.wait()`
431                break;
432            },
433            // Read from stdout if it's not done yet
434            res = stdout_reader.next_line(), if !stdout_done => {
435                match res {
436                    Ok(Some(line)) => {
437                        writeln!(io::stderr(), "{line}")?;
438                        output_capture.push_str(&line);
439                        output_capture.push('\n');
440                    },
441                    _ => stdout_done = true,
442                }
443            },
444            // Read from stderr if it's not done yet
445            res = stderr_reader.next_line(), if !stderr_done => {
446                match res {
447                    Ok(Some(line)) => {
448                        writeln!(io::stderr(), "{line}")?;
449                        output_capture.push_str(&line);
450                        output_capture.push('\n');
451                    },
452                    _ => stderr_done = true,
453                }
454            },
455            // This branch is taken once both output streams are done
456            else => break,
457        }
458    }
459
460    // Wait for the process to fully exit to get its final status
461    let status = child.wait().await.wrap_err("Failed to wait for command")?;
462
463    Ok((status, output_capture, terminated_by_signal))
464}
465
466/// Builds a base `Command` object for executing a command string via the OS shell
467fn prepare_command_execution(command: &str, include_prompt: bool) -> color_eyre::Result<tokio::process::Command> {
468    // Let the OS shell parse the command, supporting complex commands, arguments, and pipelines
469    let shell = get_shell_type();
470    let shell_arg = match shell {
471        ShellType::Cmd => "/c",
472        ShellType::WindowsPowerShell => "-Command",
473        _ => "-c",
474    };
475
476    tracing::info!("Executing command: {shell} {shell_arg} -- {command}");
477
478    // Print the command on stderr
479    let write_result = if include_prompt {
480        writeln!(
481            io::stderr(),
482            "{}{command}",
483            env::var("INTELLI_EXEC_PROMPT").as_deref().unwrap_or("> "),
484        )
485    } else {
486        writeln!(io::stderr(), "{command}")
487    };
488    // Handle broken pipe
489    if let Err(err) = write_result {
490        if err.kind() != io::ErrorKind::BrokenPipe {
491            return Err(err).wrap_err("Failed writing to stderr");
492        }
493        tracing::error!("Failed writing to stderr: Broken pipe");
494    };
495
496    // Build the base command object
497    let mut cmd = tokio::process::Command::new(shell.to_string());
498    cmd.arg(shell_arg).arg(command).kill_on_drop(true);
499    Ok(cmd)
500}
intelli_shell/utils/process.rs

intelli_shell/utils/
process.rs
