Skip to main content

lean_ctx/
shell_hook.rs

1use crate::marked_block;
2
3const MARKER_START: &str = "# >>> lean-ctx shell hook >>>";
4const MARKER_END: &str = "# <<< lean-ctx shell hook <<<";
5const ALIAS_START: &str = "# >>> lean-ctx agent aliases >>>";
6const ALIAS_END: &str = "# <<< lean-ctx agent aliases <<<";
7
8const KNOWN_AGENT_ENV_VARS: &[&str] = &[
9    "LEAN_CTX_AGENT",
10    "CLAUDECODE",
11    "CODEX_CLI_SESSION",
12    "GEMINI_SESSION",
13];
14
15const AGENT_ALIASES: &[(&str, &str)] = &[
16    ("claude", "claude"),
17    ("codex", "codex"),
18    ("gemini", "gemini"),
19];
20
21pub fn install_all(quiet: bool) {
22    let home = match dirs::home_dir() {
23        Some(h) => h,
24        None => {
25            eprintln!("Cannot resolve home directory");
26            return;
27        }
28    };
29
30    install_zshenv(&home, quiet);
31    install_bashenv(&home, quiet);
32    install_aliases(&home, quiet);
33}
34
35pub fn uninstall_all(quiet: bool) {
36    let home = match dirs::home_dir() {
37        Some(h) => h,
38        None => return,
39    };
40
41    marked_block::remove_from_file(
42        &home.join(".zshenv"),
43        MARKER_START,
44        MARKER_END,
45        quiet,
46        "shell hook from ~/.zshenv",
47    );
48    marked_block::remove_from_file(
49        &home.join(".bashenv"),
50        MARKER_START,
51        MARKER_END,
52        quiet,
53        "shell hook from ~/.bashenv",
54    );
55    marked_block::remove_from_file(
56        &home.join(".zshrc"),
57        ALIAS_START,
58        ALIAS_END,
59        quiet,
60        "agent aliases from ~/.zshrc",
61    );
62    marked_block::remove_from_file(
63        &home.join(".bashrc"),
64        ALIAS_START,
65        ALIAS_END,
66        quiet,
67        "agent aliases from ~/.bashrc",
68    );
69}
70
71fn install_zshenv(home: &std::path::Path, quiet: bool) {
72    let path = home.join(".zshenv");
73    let env_check = build_env_check();
74
75    let hook = format!(
76        r#"{MARKER_START}
77if [[ -z "$LEAN_CTX_ACTIVE" && -n "$ZSH_EXECUTION_STRING" ]] && command -v lean-ctx &>/dev/null; then
78  if {env_check}; then
79    export LEAN_CTX_ACTIVE=1
80    exec lean-ctx -c "$ZSH_EXECUTION_STRING"
81  fi
82fi
83{MARKER_END}"#
84    );
85
86    marked_block::upsert(
87        &path,
88        MARKER_START,
89        MARKER_END,
90        &hook,
91        quiet,
92        "shell hook in ~/.zshenv",
93    );
94}
95
96fn install_bashenv(home: &std::path::Path, quiet: bool) {
97    let path = home.join(".bashenv");
98    let env_check = build_env_check();
99
100    let hook = format!(
101        r#"{MARKER_START}
102if [[ -z "$LEAN_CTX_ACTIVE" && -n "$BASH_EXECUTION_STRING" ]] && command -v lean-ctx &>/dev/null; then
103  if {env_check}; then
104    export LEAN_CTX_ACTIVE=1
105    exec lean-ctx -c "$BASH_EXECUTION_STRING"
106  fi
107fi
108{MARKER_END}"#
109    );
110
111    marked_block::upsert(
112        &path,
113        MARKER_START,
114        MARKER_END,
115        &hook,
116        quiet,
117        "shell hook in ~/.bashenv",
118    );
119}
120
121fn install_aliases(home: &std::path::Path, quiet: bool) {
122    let mut lines = Vec::new();
123    lines.push(ALIAS_START.to_string());
124    for (alias_name, bin_name) in AGENT_ALIASES {
125        lines.push(format!(
126            "alias {alias_name}='LEAN_CTX_AGENT=1 BASH_ENV=\"$HOME/.bashenv\" {bin_name}'"
127        ));
128    }
129    lines.push(ALIAS_END.to_string());
130    let block = lines.join("\n");
131
132    for rc in &[home.join(".zshrc"), home.join(".bashrc")] {
133        if rc.exists() {
134            let label = format!(
135                "agent aliases in ~/{}",
136                rc.file_name().unwrap_or_default().to_string_lossy()
137            );
138            marked_block::upsert(rc, ALIAS_START, ALIAS_END, &block, quiet, &label);
139        }
140    }
141}
142
143fn build_env_check() -> String {
144    let checks: Vec<String> = KNOWN_AGENT_ENV_VARS
145        .iter()
146        .map(|v| format!("-n \"${v}\""))
147        .collect();
148    format!("[[ {} ]]", checks.join(" || "))
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154
155    #[test]
156    fn env_check_format() {
157        let check = build_env_check();
158        assert!(check.contains("LEAN_CTX_AGENT"));
159        assert!(check.contains("CLAUDECODE"));
160        assert!(check.contains("||"));
161    }
162}