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