Skip to main content

lean_ctx/shell/
reentry.rs

1//! Re-entrancy markers that stop nested `lean-ctx` invocations from
2//! double-compressing — decoupled from the user-facing activation flag.
3//!
4//! Two distinct signals were historically conflated in `LEAN_CTX_ACTIVE`
5//! (GH #533), which silently disabled compression for any agent that
6//! *inherited* it:
7//!
8//! - `LEAN_CTX_ACTIVE`: shell-hook re-entry guard. Stops the *shell* hook from
9//!   re-firing inside a command lean-ctx already spawned. It is a plain env var
10//!   an agent's top-level process can legitimately inherit.
11//! - `WRAP_MARKER` (`LEAN_CTX_WRAPPED`): process-level ownership. Set ONLY by
12//!   lean-ctx on the children it spawns, so its presence reliably means "a
13//!   parent lean-ctx already owns compression of this command tree" and a
14//!   nested `lean-ctx -c` must pass through. Because lean-ctx is the only
15//!   writer, an agent cannot trigger it by leaking an env var.
16
17use std::process::Command;
18
19/// Env var lean-ctx stamps on every child it spawns to mark the command tree as
20/// already owned (compression handled by the parent lean-ctx).
21pub(crate) const WRAP_MARKER: &str = "LEAN_CTX_WRAPPED";
22
23/// Env var carrying how many nested `lean-ctx` shell spawns led to this
24/// process. Defense-in-depth against runaway self-recursion: a shell hook
25/// that `exec`s into `lean-ctx -c`, which spawns a shell to run the real
26/// command, is supposed to be a single hop — the shell lean-ctx spawns has its
27/// profile/rc sourcing neutralized so the hook cannot re-fire inside it. But a
28/// third-party shell integration sharing the same rc/profile file (or a
29/// profile-clearing gap specific to a shell/platform) can still make the hook
30/// fire again inside that "clean" shell, and the process tree then recurses
31/// without bound — observed in the wild as thousands of orphaned shell/
32/// lean-ctx processes and multi-GB memory exhaustion within seconds. This
33/// counter survives that failure mode: once nesting exceeds
34/// [`MAX_EXEC_DEPTH`], lean-ctx refuses to spawn another shell at all,
35/// regardless of what is causing the repeated re-entry.
36const EXEC_DEPTH_MARKER: &str = "LEAN_CTX_EXEC_DEPTH";
37
38/// Hard ceiling on nested `lean-ctx -c`/`-t` shell spawns. Legitimate nesting
39/// (an agent wrapper being unwrapped, a script re-invoking itself once) never
40/// comes close to this; it exists purely to bound a runaway loop.
41pub(crate) const MAX_EXEC_DEPTH: u32 = 8;
42
43/// How many nested shell-spawning `exec` calls led to this process (0 if none).
44#[must_use]
45pub(crate) fn exec_depth() -> u32 {
46    std::env::var(EXEC_DEPTH_MARKER)
47        .ok()
48        .and_then(|v| v.parse().ok())
49        .unwrap_or(0)
50}
51
52/// True once nesting is deep enough that lean-ctx must refuse to spawn yet
53/// another shell rather than risk unbounded recursion.
54#[must_use]
55pub(crate) fn exec_depth_exceeded() -> bool {
56    exec_depth() >= MAX_EXEC_DEPTH
57}
58
59/// Stamp the incremented nesting depth on a child about to be spawned. Called
60/// by every function that spawns a shell to run a command, independently of
61/// [`mark_child`]/[`clear_shell_default_markers`] — the depth counter must
62/// keep climbing even on the `exec_shell_default` path, which deliberately
63/// clears the other markers so the delegate shell's own hooks decide fresh.
64pub(crate) fn stamp_exec_depth(cmd: &mut Command) {
65    cmd.env(EXEC_DEPTH_MARKER, (exec_depth() + 1).to_string());
66}
67
68/// Legacy shell-hook re-entry guard. Still stamped on children so the shell
69/// hook (which tests `-z "$LEAN_CTX_ACTIVE"`) does not re-fire inside a wrapped
70/// command, but it no longer gates compression on its own.
71pub(crate) const ACTIVE_MARKER: &str = "LEAN_CTX_ACTIVE";
72
73/// True when `LEAN_CTX_DISABLED` forces `lean-ctx -c` / `-t` to run the
74/// command raw instead of compressing it.
75#[must_use]
76pub(crate) fn is_disabled() -> bool {
77    std::env::var("LEAN_CTX_DISABLED").is_ok()
78}
79
80/// True when a parent lean-ctx already owns this command tree. Nested explicit
81/// `lean-ctx -c` should defer to the ambient shell defaults, not force raw
82/// command bytes, so shell-hook compression can still apply when configured.
83#[must_use]
84pub fn is_wrapped() -> bool {
85    std::env::var(WRAP_MARKER).is_ok()
86}
87
88/// True when `lean-ctx -c` / `-t` must avoid owning compression itself: either
89/// a parent lean-ctx already owns this command tree (`WRAP_MARKER`) or
90/// compression is globally disabled (`LEAN_CTX_DISABLED`).
91///
92/// Deliberately does **not** consult `LEAN_CTX_ACTIVE`: that flag is only a
93/// shell-hook re-entry guard, and an agent's top-level process can inherit it,
94/// which previously suppressed compression for every command it ran (#533).
95#[must_use]
96pub fn should_pass_through() -> bool {
97    is_wrapped() || is_disabled()
98}
99
100/// Stamp a child command so nested lean-ctx invocations pass through and the
101/// shell hook does not re-fire: sets both the ownership marker
102/// (`WRAP_MARKER`) and the legacy hook guard (`LEAN_CTX_ACTIVE`).
103pub fn mark_child(cmd: &mut Command) {
104    cmd.env(ACTIVE_MARKER, "1").env(WRAP_MARKER, "1");
105    stamp_exec_depth(cmd);
106}
107
108/// Clear lean-ctx ownership/force markers before delegating to the user's shell
109/// defaults. This lets a nested explicit wrapper behave like typing the command
110/// in that shell: configured hooks may compress; unconfigured shells run raw.
111pub(crate) fn clear_shell_default_markers(cmd: &mut Command) {
112    cmd.env_remove(ACTIVE_MARKER)
113        .env_remove(WRAP_MARKER)
114        .env_remove("LEAN_CTX_COMPRESS");
115}
116
117#[cfg(test)]
118mod tests {
119    use super::{
120        ACTIVE_MARKER, EXEC_DEPTH_MARKER, MAX_EXEC_DEPTH, WRAP_MARKER, clear_shell_default_markers,
121        exec_depth, exec_depth_exceeded, is_disabled, is_wrapped, mark_child, should_pass_through,
122        stamp_exec_depth,
123    };
124
125    #[test]
126    fn exec_depth_defaults_to_zero() {
127        let _lock = crate::core::data_dir::test_env_lock();
128        crate::test_env::remove_var(EXEC_DEPTH_MARKER);
129        assert_eq!(exec_depth(), 0);
130        assert!(!exec_depth_exceeded());
131    }
132
133    #[test]
134    fn stamp_exec_depth_increments_on_child() {
135        let _lock = crate::core::data_dir::test_env_lock();
136        crate::test_env::set_var(EXEC_DEPTH_MARKER, "3");
137        let mut cmd = std::process::Command::new("true");
138        stamp_exec_depth(&mut cmd);
139        crate::test_env::remove_var(EXEC_DEPTH_MARKER);
140        let envs: std::collections::HashMap<String, String> = cmd
141            .get_envs()
142            .filter_map(|(k, v)| Some((k.to_str()?.to_string(), v?.to_str()?.to_string())))
143            .collect();
144        assert_eq!(envs.get(EXEC_DEPTH_MARKER).map(String::as_str), Some("4"));
145    }
146
147    #[test]
148    fn exec_depth_exceeded_at_ceiling() {
149        let _lock = crate::core::data_dir::test_env_lock();
150        crate::test_env::set_var(EXEC_DEPTH_MARKER, MAX_EXEC_DEPTH.to_string());
151        let exceeded = exec_depth_exceeded();
152        crate::test_env::remove_var(EXEC_DEPTH_MARKER);
153        assert!(
154            exceeded,
155            "depth equal to the ceiling must already be treated as exceeded"
156        );
157    }
158
159    #[test]
160    fn mark_child_also_stamps_exec_depth() {
161        let _lock = crate::core::data_dir::test_env_lock();
162        crate::test_env::remove_var(EXEC_DEPTH_MARKER);
163        let mut cmd = std::process::Command::new("true");
164        mark_child(&mut cmd);
165        let envs: std::collections::HashMap<String, String> = cmd
166            .get_envs()
167            .filter_map(|(k, v)| Some((k.to_str()?.to_string(), v?.to_str()?.to_string())))
168            .collect();
169        assert_eq!(envs.get(EXEC_DEPTH_MARKER).map(String::as_str), Some("1"));
170    }
171
172    /// clear_shell_default_markers intentionally resets the ownership/compress
173    /// markers so a delegate shell's own hooks decide fresh — but the depth
174    /// counter must NOT be one of them, or exactly the runaway loop this guard
175    /// exists for would reset its own trip-wire on every lap.
176    #[test]
177    fn clear_shell_default_markers_does_not_reset_exec_depth() {
178        let mut cmd = std::process::Command::new("true");
179        cmd.env(EXEC_DEPTH_MARKER, "5");
180        clear_shell_default_markers(&mut cmd);
181        let envs: std::collections::HashMap<String, String> = cmd
182            .get_envs()
183            .filter_map(|(k, v)| Some((k.to_str()?.to_string(), v?.to_str()?.to_string())))
184            .collect();
185        assert_eq!(
186            envs.get(EXEC_DEPTH_MARKER).map(String::as_str),
187            Some("5"),
188            "exec depth must survive clear_shell_default_markers"
189        );
190    }
191
192    #[test]
193    fn wrap_marker_triggers_passthrough() {
194        let _lock = crate::core::data_dir::test_env_lock();
195        crate::test_env::remove_var("LEAN_CTX_DISABLED");
196        crate::test_env::remove_var(ACTIVE_MARKER);
197        crate::test_env::set_var(WRAP_MARKER, "1");
198        assert!(is_wrapped());
199        assert!(should_pass_through());
200        crate::test_env::remove_var(WRAP_MARKER);
201    }
202
203    #[test]
204    fn disabled_triggers_passthrough() {
205        let _lock = crate::core::data_dir::test_env_lock();
206        crate::test_env::remove_var(WRAP_MARKER);
207        crate::test_env::remove_var(ACTIVE_MARKER);
208        crate::test_env::set_var("LEAN_CTX_DISABLED", "1");
209        assert!(is_disabled());
210        assert!(should_pass_through());
211        crate::test_env::remove_var("LEAN_CTX_DISABLED");
212    }
213
214    /// #533: an inherited/leaked `LEAN_CTX_ACTIVE` must NOT suppress compression.
215    #[test]
216    fn inherited_active_does_not_trigger_passthrough() {
217        let _lock = crate::core::data_dir::test_env_lock();
218        crate::test_env::remove_var(WRAP_MARKER);
219        crate::test_env::remove_var("LEAN_CTX_DISABLED");
220        crate::test_env::set_var(ACTIVE_MARKER, "1");
221        assert!(
222            !should_pass_through(),
223            "inherited LEAN_CTX_ACTIVE must not disable compression (#533)"
224        );
225        crate::test_env::remove_var(ACTIVE_MARKER);
226    }
227
228    #[test]
229    fn mark_child_sets_both_markers() {
230        let mut cmd = std::process::Command::new("true");
231        mark_child(&mut cmd);
232        let envs: std::collections::HashMap<String, String> = cmd
233            .get_envs()
234            .filter_map(|(k, v)| Some((k.to_str()?.to_string(), v?.to_str()?.to_string())))
235            .collect();
236        assert_eq!(envs.get(WRAP_MARKER).map(String::as_str), Some("1"));
237        assert_eq!(envs.get(ACTIVE_MARKER).map(String::as_str), Some("1"));
238    }
239
240    #[test]
241    fn clear_shell_default_markers_removes_wrapping_state() {
242        let mut cmd = std::process::Command::new("true");
243        mark_child(&mut cmd);
244        cmd.env("LEAN_CTX_COMPRESS", "1");
245
246        clear_shell_default_markers(&mut cmd);
247
248        let envs: std::collections::HashMap<String, Option<String>> = cmd
249            .get_envs()
250            .filter_map(|(k, v)| {
251                Some((
252                    k.to_str()?.to_string(),
253                    v.and_then(|v| v.to_str().map(str::to_string)),
254                ))
255            })
256            .collect();
257
258        assert_eq!(envs.get(WRAP_MARKER), Some(&None));
259        assert_eq!(envs.get(ACTIVE_MARKER), Some(&None));
260        assert_eq!(envs.get("LEAN_CTX_COMPRESS"), Some(&None));
261    }
262}