Skip to main content

lean_ctx/tools/
ctx_execute.rs

1use crate::core::sandbox::{self, SandboxResult};
2use crate::core::tokens::count_tokens;
3use crate::server::tool_trait::ShellOutcome;
4
5/// Executes a code snippet in a sandboxed environment.
6/// Returns the formatted output plus the structured outcome so the MCP layer
7/// can set `isError`/`structuredContent` on failures (GitHub #389).
8pub fn handle(
9    language: &str,
10    code: &str,
11    intent: Option<&str>,
12    timeout: Option<u64>,
13) -> (String, ShellOutcome) {
14    let result = sandbox::execute(language, code, timeout);
15    (
16        format_result(&result, intent),
17        ShellOutcome::Exit(result.exit_code),
18    )
19}
20
21/// Reads a file from disk, detects its language, and executes a processing script.
22///
23/// `project_root` is used for pathjail validation. If `None`, the current
24/// directory is used as the jail root. Precondition failures (path rejected,
25/// unreadable, too large) never execute anything and report `Blocked`.
26pub fn handle_file(
27    path: &str,
28    intent: Option<&str>,
29    project_root: Option<&str>,
30) -> (String, ShellOutcome) {
31    let jail_root = match project_root {
32        Some(r) => std::path::PathBuf::from(r),
33        None => std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from(".")),
34    };
35    let candidate = std::path::Path::new(path);
36    let jailed = match crate::core::pathjail::jail_path(candidate, &jail_root) {
37        Ok(p) => p,
38        Err(e) => return (format!("Path rejected: {e}"), ShellOutcome::Blocked),
39    };
40    let path_str = jailed.to_string_lossy();
41
42    let cap = crate::core::limits::max_read_bytes();
43    let meta = match std::fs::metadata(&*jailed) {
44        Ok(m) => m,
45        Err(e) => {
46            return (
47                format!("Error reading {path_str}: {e}"),
48                ShellOutcome::Blocked,
49            );
50        }
51    };
52    if meta.len() > cap as u64 {
53        return (
54            format!(
55                "File too large ({} bytes, limit {cap} bytes). Use a line-range read instead.",
56                meta.len()
57            ),
58            ShellOutcome::Blocked,
59        );
60    }
61    let content = match std::fs::read_to_string(&*jailed) {
62        Ok(c) => c,
63        Err(e) => {
64            return (
65                format!("Error reading {path_str}: {e}"),
66                ShellOutcome::Blocked,
67            );
68        }
69    };
70
71    let language = detect_language_from_extension(path);
72    let code = build_file_processing_script(&language, &content, intent);
73    let result = sandbox::execute(&language, &code, None);
74    (
75        format_result(&result, intent),
76        ShellOutcome::Exit(result.exit_code),
77    )
78}
79
80/// Executes multiple (language, code) pairs in parallel and returns aggregated
81/// results. The outcome carries the first non-zero exit code (0 when all
82/// tasks succeeded), so one failing task marks the whole batch as failed.
83pub fn handle_batch(items: &[(String, String)]) -> (String, ShellOutcome) {
84    let results = sandbox::batch_execute(items);
85    let mut output = Vec::new();
86
87    for (i, result) in results.iter().enumerate() {
88        let label = format!("[{}/{}] {}", i + 1, results.len(), result.language);
89        if result.exit_code == 0 {
90            let stdout = result.stdout.trim();
91            if stdout.is_empty() {
92                output.push(format!("{label}: (no output) [{} ms]", result.duration_ms));
93            } else {
94                output.push(format!("{label}: {stdout} [{} ms]", result.duration_ms));
95            }
96        } else {
97            let stderr = result.stderr.trim();
98            output.push(format!(
99                "{label}: EXIT {} — {stderr} [{} ms]",
100                result.exit_code, result.duration_ms
101            ));
102        }
103    }
104
105    let total_ms: u64 = results.iter().map(|r| r.duration_ms).sum();
106    output.push(format!("\n{} tasks, {} ms total", results.len(), total_ms));
107    let first_failure = results
108        .iter()
109        .map(|r| r.exit_code)
110        .find(|c| *c != 0)
111        .unwrap_or(0);
112    (output.join("\n"), ShellOutcome::Exit(first_failure))
113}
114
115fn format_result(result: &SandboxResult, intent: Option<&str>) -> String {
116    let mut parts = Vec::new();
117
118    if result.exit_code == 0 {
119        let stdout = result.stdout.trim();
120        if stdout.is_empty() {
121            parts.push("(no output)".to_string());
122        } else {
123            let raw_tokens = count_tokens(stdout);
124            parts.push(stdout.to_string());
125
126            if let Some(intent_desc) = intent
127                && raw_tokens > 50
128            {
129                parts.push(format!("[intent: {intent_desc}]"));
130            }
131        }
132    } else {
133        if !result.stdout.is_empty() {
134            parts.push(result.stdout.trim().to_string());
135        }
136        parts.push(format!(
137            "EXIT {} — {}",
138            result.exit_code,
139            result.stderr.trim()
140        ));
141    }
142
143    parts.push(format!("[{} | {} ms]", result.language, result.duration_ms));
144    parts.join("\n")
145}
146
147fn detect_language_from_extension(path: &str) -> String {
148    let ext = path.rsplit('.').next().unwrap_or("");
149    match ext {
150        "js" | "mjs" | "cjs" => "javascript",
151        "ts" | "mts" | "cts" => "typescript",
152        "py" | "json" | "csv" | "log" | "txt" | "xml" | "yaml" | "yml" | "md" | "html" => "python",
153        "rb" => "ruby",
154        "go" => "go",
155        "rs" => "rust",
156        "php" => "php",
157        "pl" => "perl",
158        "r" | "R" => "r",
159        "ex" | "exs" => "elixir",
160        _ => "shell",
161    }
162    .to_string()
163}
164
165fn sanitize_intent(raw: &str) -> String {
166    raw.chars()
167        .filter(|c| c.is_alphanumeric() || *c == ' ' || *c == '-' || *c == '_' || *c == '.')
168        .take(200)
169        .collect()
170}
171
172/// Escape a path for safe embedding inside a Python raw double-quoted string.
173/// Handles embedded double-quotes which would break `r"..."`.
174fn escape_for_python_raw(path: &str) -> String {
175    path.replace('"', r#"\" + '"' + r""#)
176}
177
178/// Escape a path for safe embedding inside a shell double-quoted string.
179/// Handles `$`, backtick, `\`, and `"` which are special inside double quotes.
180fn escape_for_shell_dq(path: &str) -> String {
181    let mut out = String::with_capacity(path.len());
182    for ch in path.chars() {
183        match ch {
184            '$' | '`' | '"' | '\\' => {
185                out.push('\\');
186                out.push(ch);
187            }
188            _ => out.push(ch),
189        }
190    }
191    out
192}
193
194fn build_file_processing_script(language: &str, content: &str, intent: Option<&str>) -> String {
195    let Ok(tmp) = tempfile::Builder::new()
196        .prefix("lean-ctx-exec-")
197        .suffix(".dat")
198        .tempfile()
199    else {
200        return "echo 'lean-ctx: failed to create temp file'".to_string();
201    };
202    let _ = std::fs::write(tmp.path(), content);
203    let tmp_path = tmp.path().to_string_lossy().to_string();
204    let _keep = tmp.into_temp_path();
205    let intent_str = sanitize_intent(intent.unwrap_or("summarize the content"));
206
207    if language == "python" {
208        let py_path = escape_for_python_raw(&tmp_path);
209        format!(
210            r#"
211    import os
212
213    with open(r"{py_path}", "r", encoding="utf-8") as f:
214        data = f.read()
215    os.remove(r"{py_path}")
216
217    lines = data.strip().split('\n')
218    total_lines = len(lines)
219    total_bytes = len(data.encode('utf-8'))
220
221    word_count = sum(len(line.split()) for line in lines)
222
223    print(f"{{total_lines}} lines, {{total_bytes}} bytes, {{word_count}} words")
224    print("Intent: {intent_str}")
225
226    if total_lines > 10:
227        print(f"First 3: {{lines[:3]}}")
228        print(f"Last 3: {{lines[-3:]}}")
229    "#
230        )
231    } else {
232        let sh_path = escape_for_shell_dq(&tmp_path);
233        format!(
234            r#"
235    data=$(cat "{sh_path}")
236    rm -f "{sh_path}"
237    lines=$(echo "$data" | wc -l | tr -d ' ')
238    bytes=$(echo "$data" | wc -c | tr -d ' ')
239    echo "$lines lines, $bytes bytes"
240    echo 'Intent: {intent_str}'
241    echo "$data" | head -3
242    echo "..."
243    echo "$data" | tail -3
244    "#
245        )
246    }
247}
248
249#[cfg(test)]
250mod tests {
251    use super::*;
252
253    #[test]
254    fn handle_simple_python() {
255        let (result, outcome) = handle("python", "print(2 + 2)", None, None);
256        assert!(result.contains('4'));
257        assert!(result.contains("python"));
258        assert_eq!(outcome, ShellOutcome::Exit(0), "success must report exit 0");
259    }
260
261    #[test]
262    fn handle_with_intent() {
263        let (result, _) = handle(
264            "python",
265            "print('found 5 errors')",
266            Some("count errors"),
267            None,
268        );
269        assert!(result.contains("found 5 errors"));
270    }
271
272    #[test]
273    fn handle_error_shows_stderr() {
274        let (result, outcome) = handle("python", "raise Exception('boom')", None, None);
275        assert!(result.contains("EXIT"));
276        assert!(result.contains("boom"));
277        assert!(
278            outcome.is_error(),
279            "non-zero sandbox exit must surface as a tool error (#389)"
280        );
281    }
282
283    #[test]
284    fn detect_language_from_path() {
285        assert_eq!(detect_language_from_extension("test.py"), "python");
286        assert_eq!(detect_language_from_extension("test.js"), "javascript");
287        assert_eq!(detect_language_from_extension("test.rs"), "rust");
288        assert_eq!(detect_language_from_extension("test.csv"), "python");
289        assert_eq!(detect_language_from_extension("test.log"), "python");
290    }
291
292    #[test]
293    fn escape_shell_dq_handles_special_chars() {
294        assert_eq!(escape_for_shell_dq(r"C:\tmp\file"), r"C:\\tmp\\file");
295        assert_eq!(escape_for_shell_dq("/tmp/normal"), "/tmp/normal");
296        assert_eq!(escape_for_shell_dq("path with $VAR"), r"path with \$VAR");
297        assert_eq!(escape_for_shell_dq(r#"path"quote"#), r#"path\"quote"#);
298        assert_eq!(escape_for_shell_dq("has `backtick`"), r"has \`backtick\`");
299    }
300
301    #[test]
302    fn escape_python_raw_handles_quotes() {
303        assert_eq!(escape_for_python_raw("/tmp/normal"), "/tmp/normal");
304        assert_eq!(escape_for_python_raw(r"C:\Users\test"), r"C:\Users\test");
305    }
306
307    #[test]
308    fn script_with_spaces_in_path() {
309        let script = build_file_processing_script("shell", "test data", None);
310        let lines: Vec<&str> = script.lines().collect();
311        for line in &lines {
312            if line.contains("cat ") || line.contains("rm -f") {
313                assert!(
314                    line.contains('"'),
315                    "path must be double-quoted in shell script: {line}"
316                );
317            }
318        }
319    }
320
321    #[test]
322    #[cfg(not(target_os = "windows"))]
323    fn batch_multiple_tasks() {
324        let items = vec![
325            ("python".to_string(), "print('task1')".to_string()),
326            ("shell".to_string(), "echo task2".to_string()),
327        ];
328        let (result, outcome) = handle_batch(&items);
329        assert!(result.contains("task1"));
330        assert!(result.contains("task2"));
331        assert!(result.contains("2 tasks"));
332        assert_eq!(outcome, ShellOutcome::Exit(0), "all tasks succeeded");
333    }
334
335    #[test]
336    #[cfg(not(target_os = "windows"))]
337    fn batch_with_failing_task_reports_failure() {
338        let items = vec![
339            ("shell".to_string(), "echo ok".to_string()),
340            ("shell".to_string(), "exit 3".to_string()),
341        ];
342        let (_, outcome) = handle_batch(&items);
343        assert_eq!(
344            outcome,
345            ShellOutcome::Exit(3),
346            "one failing task marks the whole batch as failed (#389)"
347        );
348    }
349
350    #[test]
351    fn handle_file_precondition_failure_is_blocked() {
352        let (result, outcome) = handle_file("/nonexistent/definitely-missing.py", None, None);
353        assert!(outcome.is_error(), "precondition failures are tool errors");
354        assert_eq!(outcome, ShellOutcome::Blocked, "nothing was executed");
355        assert!(!result.is_empty());
356    }
357}