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
112/// Normalize command separators for the detected Windows shell.
113///
114/// - cmd.exe (`/C`): `;` is invalid → convert to `&&`
115/// - PowerShell 5.x (`-Command`): `&&` is invalid → convert to `;`
116/// - POSIX / PS7+ via Git Bash (`-c`): no changes needed
117pub fn normalize_command_for_shell(command: &str) -> String {
118    if !cfg!(windows) {
119        return command.to_string();
120    }
121    let (_, flag) = crate::shell::shell_and_flag();
122    match flag.as_str() {
123        "/C" => replace_unquoted(command, b";", b" && "),
124        "-Command" => replace_unquoted(command, b"&&", b"; "),
125        _ => command.to_string(),
126    }
127}
128
129fn replace_unquoted(command: &str, needle: &[u8], replacement: &[u8]) -> String {
130    let bytes = command.as_bytes();
131    let mut result = Vec::with_capacity(bytes.len() + 16);
132    let mut in_single = false;
133    let mut in_double = false;
134    let mut i = 0;
135    while i < bytes.len() {
136        if bytes[i] == b'\'' && !in_double {
137            in_single = !in_single;
138        } else if bytes[i] == b'"' && !in_single {
139            in_double = !in_double;
140        } else if !in_single && !in_double && bytes[i..].starts_with(needle) {
141            result.extend_from_slice(replacement);
142            i += needle.len();
143            continue;
144        }
145        result.push(bytes[i]);
146        i += 1;
147    }
148    String::from_utf8(result).unwrap_or_else(|_| command.to_string())
149}
150
151pub fn handle(command: &str, output: &str, crp_mode: CrpMode) -> String {
152    let original_tokens = count_tokens(output);
153
154    if contains_auth_flow(output) {
155        let savings = protocol::format_savings(original_tokens, original_tokens);
156        return format!(
157            "{output}\n[lean-ctx: auth/device-code flow detected — output preserved uncompressed]\n{savings}"
158        );
159    }
160
161    let compressed = match patterns::compress_output(command, output) {
162        Some(c) => c,
163        None => generic_compress(output),
164    };
165
166    if crp_mode.is_tdd() && looks_like_code(&compressed) {
167        let ext = detect_ext_from_command(command);
168        let mut sym = SymbolMap::new();
169        let idents = symbol_map::extract_identifiers(&compressed, ext);
170        for ident in &idents {
171            sym.register(ident);
172        }
173        if !sym.is_empty() {
174            let mapped = sym.apply(&compressed);
175            let sym_table = sym.format_table();
176            let result = format!("{mapped}{sym_table}");
177            let sent = count_tokens(&result);
178            let savings = protocol::format_savings(original_tokens, sent);
179            return format!("{result}\n{savings}");
180        }
181    }
182
183    let sent = count_tokens(&compressed);
184    let savings = protocol::format_savings(original_tokens, sent);
185
186    format!("{compressed}\n{savings}")
187}
188
189fn generic_compress(output: &str) -> String {
190    let output = crate::core::compressor::strip_ansi(output);
191    let lines: Vec<&str> = output
192        .lines()
193        .filter(|l| {
194            let t = l.trim();
195            !t.is_empty()
196        })
197        .collect();
198
199    if lines.len() <= 10 {
200        return lines.join("\n");
201    }
202
203    let first_3 = &lines[..3];
204    let last_3 = &lines[lines.len() - 3..];
205    let omitted = lines.len() - 6;
206    format!(
207        "{}\n[truncated: showing 6/{} lines, {} omitted. Use raw=true for full output.]\n{}",
208        first_3.join("\n"),
209        lines.len(),
210        omitted,
211        last_3.join("\n")
212    )
213}
214
215fn looks_like_code(text: &str) -> bool {
216    let indicators = [
217        "fn ",
218        "pub ",
219        "let ",
220        "const ",
221        "impl ",
222        "struct ",
223        "enum ",
224        "function ",
225        "class ",
226        "import ",
227        "export ",
228        "def ",
229        "async ",
230        "=>",
231        "->",
232        "::",
233        "self.",
234        "this.",
235    ];
236    let total_lines = text.lines().count();
237    if total_lines < 3 {
238        return false;
239    }
240    let code_lines = text
241        .lines()
242        .filter(|l| indicators.iter().any(|i| l.contains(i)))
243        .count();
244    code_lines as f64 / total_lines as f64 > 0.15
245}
246
247fn detect_ext_from_command(command: &str) -> &str {
248    let cmd = command.to_lowercase();
249    if cmd.contains("cargo") || cmd.contains(".rs") {
250        "rs"
251    } else if cmd.contains("npm")
252        || cmd.contains("node")
253        || cmd.contains(".ts")
254        || cmd.contains(".js")
255    {
256        "ts"
257    } else if cmd.contains("python") || cmd.contains("pip") || cmd.contains(".py") {
258        "py"
259    } else if cmd.contains("go ") || cmd.contains(".go") {
260        "go"
261    } else {
262        "rs"
263    }
264}
265
266/// Detects OAuth device code flow output that must not be compressed.
267/// Uses a two-tier approach: strong signals match alone (very specific to
268/// device code flows), weak signals require a URL/domain in the same output.
269pub fn contains_auth_flow(output: &str) -> bool {
270    let lower = output.to_lowercase();
271
272    const STRONG_SIGNALS: &[&str] = &[
273        "devicelogin",
274        "deviceauth",
275        "device_code",
276        "device code",
277        "device-code",
278        "verification_uri",
279        "user_code",
280        "one-time code",
281    ];
282
283    if STRONG_SIGNALS.iter().any(|s| lower.contains(s)) {
284        return true;
285    }
286
287    const WEAK_SIGNALS: &[&str] = &[
288        "enter the code",
289        "enter this code",
290        "enter code:",
291        "use the code",
292        "use a web browser to open",
293        "open the page",
294        "authenticate by visiting",
295        "sign in with the code",
296        "sign in using a code",
297        "verification code",
298        "authorize this device",
299        "waiting for authentication",
300        "waiting for login",
301        "waiting for you to authenticate",
302        "open your browser",
303        "open in your browser",
304    ];
305
306    let has_weak_signal = WEAK_SIGNALS.iter().any(|s| lower.contains(s));
307    if !has_weak_signal {
308        return false;
309    }
310
311    lower.contains("http://") || lower.contains("https://")
312}
313
314#[cfg(test)]
315mod tests {
316    use super::*;
317
318    #[test]
319    fn normalize_cmd_no_change_on_unix() {
320        if cfg!(windows) {
321            return;
322        }
323        assert_eq!(
324            normalize_command_for_shell("cd /tmp; ls -la"),
325            "cd /tmp; ls -la"
326        );
327    }
328
329    #[test]
330    fn replace_unquoted_semicolons_to_ampersand() {
331        assert_eq!(
332            replace_unquoted("cd /tmp; ls -la", b";", b" && "),
333            "cd /tmp &&  ls -la"
334        );
335    }
336
337    #[test]
338    fn replace_unquoted_ampersand_to_semicolons() {
339        assert_eq!(
340            replace_unquoted("cd backend && git status", b"&&", b"; "),
341            "cd backend ;  git status"
342        );
343    }
344
345    #[test]
346    fn replace_unquoted_preserves_quoted_strings() {
347        assert_eq!(
348            replace_unquoted(r#"echo "a && b" && ls"#, b"&&", b"; "),
349            r#"echo "a && b" ;  ls"#
350        );
351        assert_eq!(
352            replace_unquoted("echo 'a; b'; ls", b";", b" && "),
353            "echo 'a; b' &&  ls"
354        );
355    }
356
357    #[test]
358    fn validate_allows_safe_commands() {
359        assert!(validate_command("git status").is_none());
360        assert!(validate_command("cargo test").is_none());
361        assert!(validate_command("npm run build").is_none());
362        assert!(validate_command("ls -la").is_none());
363    }
364
365    #[test]
366    fn validate_blocks_file_writes() {
367        assert!(validate_command("cat > file.py << 'EOF'\nprint('hi')\nEOF").is_some());
368        assert!(validate_command("echo 'data' > output.txt").is_some());
369        assert!(validate_command("tee /tmp/file.txt").is_some());
370        assert!(validate_command("printf 'hello' > test.txt").is_some());
371        assert!(validate_command("cat << EOF\ncontent\nEOF").is_some());
372    }
373
374    #[test]
375    fn validate_blocks_oversized_commands() {
376        let huge = "x".repeat(MAX_COMMAND_BYTES + 1);
377        let result = validate_command(&huge);
378        assert!(result.is_some());
379        assert!(result.unwrap().contains("too large"));
380    }
381
382    #[test]
383    fn validate_allows_cat_without_redirect() {
384        assert!(validate_command("cat file.txt").is_none());
385    }
386
387    // --- Auth flow detection: strong signals (no URL needed) ---
388
389    #[test]
390    fn auth_flow_detects_azure_device_code() {
391        let output = "To sign in, use a web browser to open the page https://microsoft.com/devicelogin and enter the code ABCD1234 to authenticate.";
392        assert!(contains_auth_flow(output));
393    }
394
395    #[test]
396    fn auth_flow_detects_gh_auth_one_time_code() {
397        let output =
398            "! First copy your one-time code: ABCD-1234\n- Press Enter to open github.com in your browser...";
399        assert!(contains_auth_flow(output));
400    }
401
402    #[test]
403    fn auth_flow_detects_device_code_json() {
404        let output = r#"{"device_code":"abc123","user_code":"ABCD-1234","verification_uri":"https://example.com/activate"}"#;
405        assert!(contains_auth_flow(output));
406    }
407
408    #[test]
409    fn auth_flow_detects_verification_uri_field() {
410        let output =
411            r#"{"verification_uri": "https://login.microsoftonline.com/common/oauth2/deviceauth"}"#;
412        assert!(contains_auth_flow(output));
413    }
414
415    #[test]
416    fn auth_flow_detects_user_code_field() {
417        let output = r#"{"user_code": "FGHJK-LMNOP", "expires_in": 900}"#;
418        assert!(contains_auth_flow(output));
419    }
420
421    // --- Auth flow detection: weak signals (require URL) ---
422
423    #[test]
424    fn auth_flow_detects_gcloud_with_url() {
425        let output = "Go to the following link in your browser:\n\n    https://accounts.google.com/o/oauth2/auth?response_type=code\n\nEnter verification code: ";
426        assert!(contains_auth_flow(output));
427    }
428
429    #[test]
430    fn auth_flow_detects_aws_sso_with_url() {
431        let output = "If the browser does not open, open the following URL:\nhttps://device.sso.us-east-1.amazonaws.com/\n\nThen enter the code:\nABCD-EFGH";
432        assert!(contains_auth_flow(output));
433    }
434
435    #[test]
436    fn auth_flow_detects_firebase_with_url() {
437        let output = "Visit this URL on this device to log in:\nhttps://accounts.google.com/o/oauth2/auth?...\n\nWaiting for authentication...";
438        assert!(contains_auth_flow(output));
439    }
440
441    #[test]
442    fn auth_flow_detects_generic_browser_open_with_url() {
443        let output =
444            "Open your browser to https://login.example.com/device and enter the code XYZW-1234";
445        assert!(contains_auth_flow(output));
446    }
447
448    // --- False positive protection ---
449
450    #[test]
451    fn auth_flow_ignores_normal_build_output() {
452        let output = "Compiling lean-ctx v2.21.9\nFinished release profile\n";
453        assert!(!contains_auth_flow(output));
454    }
455
456    #[test]
457    fn auth_flow_ignores_git_output() {
458        let output = "On branch main\nYour branch is up to date with 'origin/main'.\nnothing to commit, working tree clean";
459        assert!(!contains_auth_flow(output));
460    }
461
462    #[test]
463    fn auth_flow_ignores_npm_install_output() {
464        let output = "added 150 packages in 3s\n\n24 packages are looking for funding\n  run `npm fund` for details\nhttps://npmjs.com/package/lean-ctx";
465        assert!(!contains_auth_flow(output));
466    }
467
468    #[test]
469    fn auth_flow_ignores_docs_mentioning_auth() {
470        let output = "The authorization code grant type is the most common OAuth flow.\nSee https://oauth.net/2/grant-types/ for details.";
471        assert!(!contains_auth_flow(output));
472    }
473
474    #[test]
475    fn auth_flow_weak_signal_requires_url() {
476        let output = "Please enter the code ABC123 in the terminal";
477        assert!(!contains_auth_flow(output));
478    }
479
480    #[test]
481    fn auth_flow_weak_signal_without_url_is_ignored() {
482        let output = "Waiting for authentication to complete... done!";
483        assert!(!contains_auth_flow(output));
484    }
485
486    #[test]
487    fn auth_flow_ignores_virtualenv_activate() {
488        let output = "Created virtualenv at .venv\nRun: source .venv/bin/activate";
489        assert!(!contains_auth_flow(output));
490    }
491
492    #[test]
493    fn auth_flow_ignores_api_response_with_code_field() {
494        let output = r#"{"status": "ok", "code": 200, "message": "success"}"#;
495        assert!(!contains_auth_flow(output));
496    }
497
498    // --- Integration: handle() preserves auth flow ---
499
500    #[test]
501    fn handle_preserves_auth_flow_output_fully() {
502        let output = "To sign in, use a web browser to open the page https://microsoft.com/devicelogin and enter the code ABCD1234 to authenticate.\nWaiting for you...\nLine 3\nLine 4\nLine 5\nLine 6\nLine 7\nLine 8\nLine 9\nLine 10\nLine 11\nLine 12\nLine 13";
503        let result = handle("az login --use-device-code", output, CrpMode::Off);
504        assert!(result.contains("ABCD1234"), "auth code must be preserved");
505        assert!(result.contains("devicelogin"), "URL must be preserved");
506        assert!(
507            result.contains("auth/device-code flow detected"),
508            "detection note must be present"
509        );
510        assert!(
511            result.contains("Line 13"),
512            "all lines must be preserved (no truncation)"
513        );
514    }
515
516    #[test]
517    fn handle_compresses_normal_output_not_auth() {
518        let lines: Vec<String> = (1..=20).map(|i| format!("Line {i} of output")).collect();
519        let output = lines.join("\n");
520        let result = handle("some-tool check", &output, CrpMode::Off);
521        assert!(
522            !result.contains("auth/device-code flow detected"),
523            "normal output must not trigger auth detection"
524        );
525        assert!(
526            result.len() < output.len() + 100,
527            "normal output should be compressed, not inflated"
528        );
529    }
530}
lean_ctx/tools/ctx_shell.rs

lean_ctx/tools/
ctx_shell.rs
