Skip to main content

lean_ctx/tools/
ctx_shell.rs

1use crate::core::patterns;
2use crate::core::protocol;
3use crate::core::symbol_map::{self, SymbolMap};
4use crate::core::tokens::count_tokens;
5use crate::tools::CrpMode;
6
7const MAX_COMMAND_BYTES: usize = 8192;
8
9const HEREDOC_PATTERNS: &[&str] = &[
10    "<< 'EOF'",
11    "<<'EOF'",
12    "<< 'ENDOFFILE'",
13    "<<'ENDOFFILE'",
14    "<< 'END'",
15    "<<'END'",
16    "<< EOF",
17    "<<EOF",
18    "cat <<",
19];
20
21/// Validates a shell command before execution. Returns Some(error_message) if
22/// the command should be rejected, None if it's safe to run.
23pub fn validate_command(command: &str) -> Option<String> {
24    if command.len() > MAX_COMMAND_BYTES {
25        return Some(format!(
26            "ERROR: Command too large ({} bytes, limit {}). \
27             If you're writing file content, use the native Write/Edit tool instead. \
28             ctx_shell is for reading command output only (git, cargo, npm, etc.).",
29            command.len(),
30            MAX_COMMAND_BYTES
31        ));
32    }
33
34    if has_file_write_redirect(command) {
35        return Some(
36            "ERROR: ctx_shell detected a file-write command (shell redirect > or >>). \
37             Use the native Write tool to create/modify files. \
38             ctx_shell is ONLY for reading command output (git status, cargo test, npm run, etc.). \
39             File writes via shell cause MCP protocol corruption on large payloads."
40                .to_string(),
41        );
42    }
43
44    let cmd_lower = command.to_lowercase();
45
46    if cmd_lower.starts_with("tee ") || cmd_lower.contains("| tee ") {
47        return Some(
48            "ERROR: ctx_shell detected a file-write command (tee). \
49             Use the native Write tool to create/modify files. \
50             ctx_shell is ONLY for reading command output."
51                .to_string(),
52        );
53    }
54
55    for pattern in HEREDOC_PATTERNS {
56        if cmd_lower.contains(&pattern.to_lowercase()) {
57            return Some(
58                "ERROR: ctx_shell detected a heredoc file-write command. \
59                 Use the native Write tool to create/modify files. \
60                 ctx_shell is ONLY for reading command output."
61                    .to_string(),
62            );
63        }
64    }
65
66    None
67}
68
69/// Detects shell redirect operators (`>` or `>>`) that write to files.
70/// Ignores `>` inside quotes, `2>` (stderr), `/dev/null`, and comparison operators.
71fn has_file_write_redirect(command: &str) -> bool {
72    let bytes = command.as_bytes();
73    let len = bytes.len();
74    let mut i = 0;
75    let mut in_single_quote = false;
76    let mut in_double_quote = false;
77
78    while i < len {
79        let c = bytes[i];
80        if c == b'\'' && !in_double_quote {
81            in_single_quote = !in_single_quote;
82        } else if c == b'"' && !in_single_quote {
83            in_double_quote = !in_double_quote;
84        } else if c == b'>' && !in_single_quote && !in_double_quote {
85            if i > 0 && bytes[i - 1] == b'2' {
86                i += 1;
87                continue;
88            }
89            let target_start = if i + 1 < len && bytes[i + 1] == b'>' {
90                i + 2
91            } else {
92                i + 1
93            };
94            let target: String = command[target_start..]
95                .trim_start()
96                .chars()
97                .take_while(|c| !c.is_whitespace())
98                .collect();
99            if target == "/dev/null" {
100                i += 1;
101                continue;
102            }
103            if !target.is_empty() {
104                return true;
105            }
106        }
107        i += 1;
108    }
109    false
110}
111
112pub fn handle(command: &str, output: &str, crp_mode: CrpMode) -> String {
113    let original_tokens = count_tokens(output);
114
115    let compressed = match patterns::compress_output(command, output) {
116        Some(c) => c,
117        None => generic_compress(output),
118    };
119
120    if crp_mode.is_tdd() && looks_like_code(&compressed) {
121        let ext = detect_ext_from_command(command);
122        let mut sym = SymbolMap::new();
123        let idents = symbol_map::extract_identifiers(&compressed, ext);
124        for ident in &idents {
125            sym.register(ident);
126        }
127        if !sym.is_empty() {
128            let mapped = sym.apply(&compressed);
129            let sym_table = sym.format_table();
130            let result = format!("{mapped}{sym_table}");
131            let sent = count_tokens(&result);
132            let savings = protocol::format_savings(original_tokens, sent);
133            return format!("{result}\n{savings}");
134        }
135    }
136
137    let sent = count_tokens(&compressed);
138    let savings = protocol::format_savings(original_tokens, sent);
139
140    format!("{compressed}\n{savings}")
141}
142
143fn generic_compress(output: &str) -> String {
144    let output = crate::core::compressor::strip_ansi(output);
145    let lines: Vec<&str> = output
146        .lines()
147        .filter(|l| {
148            let t = l.trim();
149            !t.is_empty()
150        })
151        .collect();
152
153    if lines.len() <= 10 {
154        return lines.join("\n");
155    }
156
157    let first_3 = &lines[..3];
158    let last_3 = &lines[lines.len() - 3..];
159    let omitted = lines.len() - 6;
160    format!(
161        "{}\n[truncated: showing 6/{} lines, {} omitted. Use raw=true for full output.]\n{}",
162        first_3.join("\n"),
163        lines.len(),
164        omitted,
165        last_3.join("\n")
166    )
167}
168
169fn looks_like_code(text: &str) -> bool {
170    let indicators = [
171        "fn ",
172        "pub ",
173        "let ",
174        "const ",
175        "impl ",
176        "struct ",
177        "enum ",
178        "function ",
179        "class ",
180        "import ",
181        "export ",
182        "def ",
183        "async ",
184        "=>",
185        "->",
186        "::",
187        "self.",
188        "this.",
189    ];
190    let total_lines = text.lines().count();
191    if total_lines < 3 {
192        return false;
193    }
194    let code_lines = text
195        .lines()
196        .filter(|l| indicators.iter().any(|i| l.contains(i)))
197        .count();
198    code_lines as f64 / total_lines as f64 > 0.15
199}
200
201fn detect_ext_from_command(command: &str) -> &str {
202    let cmd = command.to_lowercase();
203    if cmd.contains("cargo") || cmd.contains(".rs") {
204        "rs"
205    } else if cmd.contains("npm")
206        || cmd.contains("node")
207        || cmd.contains(".ts")
208        || cmd.contains(".js")
209    {
210        "ts"
211    } else if cmd.contains("python") || cmd.contains("pip") || cmd.contains(".py") {
212        "py"
213    } else if cmd.contains("go ") || cmd.contains(".go") {
214        "go"
215    } else {
216        "rs"
217    }
218}
219
220#[cfg(test)]
221mod tests {
222    use super::*;
223
224    #[test]
225    fn validate_allows_safe_commands() {
226        assert!(validate_command("git status").is_none());
227        assert!(validate_command("cargo test").is_none());
228        assert!(validate_command("npm run build").is_none());
229        assert!(validate_command("ls -la").is_none());
230    }
231
232    #[test]
233    fn validate_blocks_file_writes() {
234        assert!(validate_command("cat > file.py << 'EOF'\nprint('hi')\nEOF").is_some());
235        assert!(validate_command("echo 'data' > output.txt").is_some());
236        assert!(validate_command("tee /tmp/file.txt").is_some());
237        assert!(validate_command("printf 'hello' > test.txt").is_some());
238        assert!(validate_command("cat << EOF\ncontent\nEOF").is_some());
239    }
240
241    #[test]
242    fn validate_blocks_oversized_commands() {
243        let huge = "x".repeat(MAX_COMMAND_BYTES + 1);
244        let result = validate_command(&huge);
245        assert!(result.is_some());
246        assert!(result.unwrap().contains("too large"));
247    }
248
249    #[test]
250    fn validate_allows_cat_without_redirect() {
251        assert!(validate_command("cat file.txt").is_none());
252    }
253}