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