Skip to main content

run/
ast.rs

1// Abstract Syntax Tree definitions
2
3use regex::Regex;
4use serde::{Deserialize, Serialize};
5use std::fmt::Write as _;
6use std::sync::OnceLock;
7
8#[derive(Debug, Clone, PartialEq)]
9pub struct Program {
10    pub statements: Vec<Statement>,
11}
12
13/// Output capture mode for command execution
14#[derive(Debug, Clone, Copy, PartialEq, Default)]
15pub enum OutputMode {
16    /// Stream directly to terminal (current behavior, default for CLI)
17    #[default]
18    Stream,
19
20    /// Capture output and also print to terminal (for programmatic access with live output)
21    /// Not exposed via CLI, but available for library use
22    Capture,
23
24    /// Capture output silently and format as structured result (for MCP/JSON/Markdown)
25    /// Output is suppressed during execution and only the formatted result is printed at the end
26    Structured,
27}
28
29/// Result of a single command execution
30#[derive(Debug, Clone, Serialize, Deserialize)]
31pub struct CommandOutput {
32    /// The command that was executed
33    pub command: String,
34
35    /// Captured standard output
36    pub stdout: String,
37
38    /// Captured standard error
39    pub stderr: String,
40
41    /// Process exit code (None if killed by signal)
42    pub exit_code: Option<i32>,
43
44    /// Execution duration in milliseconds
45    pub duration_ms: u128,
46
47    /// Timestamp when execution started (Unix epoch ms)
48    pub started_at: u128,
49}
50
51/// Context information about command execution
52#[derive(Debug, Clone, Serialize, Deserialize)]
53pub struct ExecutionContext {
54    /// The Runfile function name that was invoked
55    pub function_name: String,
56
57    /// Remote host if SSH detected (extracted from command)
58    pub remote_host: Option<String>,
59
60    /// Remote user if SSH detected
61    pub remote_user: Option<String>,
62
63    /// Shell/interpreter used (@shell attribute value)
64    pub interpreter: String,
65
66    /// Working directory
67    pub working_directory: Option<String>,
68}
69
70/// Complete structured result for MCP responses
71#[derive(Debug, Clone, Serialize, Deserialize)]
72pub struct StructuredResult {
73    /// Execution context
74    pub context: ExecutionContext,
75
76    /// Individual command outputs (in execution order)
77    pub outputs: Vec<CommandOutput>,
78
79    /// Overall success (all commands exited 0)
80    pub success: bool,
81
82    /// Total execution time
83    pub total_duration_ms: u128,
84
85    /// Human-readable summary
86    pub summary: String,
87}
88
89impl StructuredResult {
90    /// Create from a collection of command outputs
91    #[must_use]
92    pub fn from_outputs(
93        function_name: &str,
94        outputs: Vec<CommandOutput>,
95        interpreter: &str,
96    ) -> Self {
97        let success = outputs.iter().all(|o| o.exit_code == Some(0));
98        let total_duration_ms = outputs.iter().map(|o| o.duration_ms).sum();
99
100        let summary = if success {
101            format!(
102                "Successfully executed {} with {} command(s)",
103                function_name,
104                outputs.len()
105            )
106        } else {
107            format!("Execution of {function_name} failed")
108        };
109
110        // Try to extract SSH context from any of the commands
111        let (remote_user, remote_host) = outputs
112            .iter()
113            .find_map(|o| ExecutionContext::extract_ssh_context(&o.command))
114            .map_or((None, None), |(user, host)| (Some(user), Some(host)));
115
116        Self {
117            context: ExecutionContext {
118                function_name: function_name.to_string(),
119                remote_host,
120                remote_user,
121                interpreter: interpreter.to_string(),
122                working_directory: std::env::current_dir()
123                    .ok()
124                    .and_then(|p| p.to_str().map(String::from)),
125            },
126            outputs,
127            success,
128            total_duration_ms,
129            summary,
130        }
131    }
132
133    /// Format as JSON for programmatic consumption
134    #[must_use]
135    pub fn to_json(&self) -> String {
136        serde_json::to_string_pretty(self).unwrap_or_default()
137    }
138
139    /// Format as Markdown for LLM readability
140    #[must_use]
141    pub fn to_markdown(&self) -> String {
142        let mut md = String::new();
143
144        // Header with context
145        let _ = write!(md, "## Execution: `{}`\n\n", self.context.function_name);
146
147        if let Some(host) = &self.context.remote_host {
148            let _ = writeln!(
149                md,
150                "**Host:** {}@{}",
151                self.context.remote_user.as_deref().unwrap_or("?"),
152                host
153            );
154        }
155
156        let _ = writeln!(
157            md,
158            "**Status:** {}",
159            if self.success {
160                "✓ Success"
161            } else {
162                "✗ Failed"
163            }
164        );
165        let _ = write!(md, "**Duration:** {}ms\n\n", self.total_duration_ms);
166
167        // Individual command outputs
168        for (i, output) in self.outputs.iter().enumerate() {
169            let _ = writeln!(md, "### Step {} ({}ms)", i + 1, output.duration_ms);
170            let _ = write!(md, "`{}`\n\n", output.command);
171
172            if !output.stdout.is_empty() {
173                md.push_str("**Output:**\n```\n");
174                md.push_str(&output.stdout);
175                md.push_str("```\n\n");
176            }
177
178            if !output.stderr.is_empty() {
179                md.push_str("**Errors:**\n```\n");
180                md.push_str(&output.stderr);
181                md.push_str("```\n\n");
182            }
183
184            if let Some(code) = output.exit_code
185                && code != 0
186            {
187                let _ = writeln!(md, "**Exit Code:** {code}");
188            }
189        }
190
191        md
192    }
193
194    /// Format optimized for MCP tool response (clean markdown, no implementation details)
195    /// This intentionally hides the command source code to protect sensitive information
196    /// like database connection strings, API keys, etc.
197    #[must_use]
198    pub fn to_mcp_format(&self) -> String {
199        let mut md = String::new();
200
201        // Header with context
202        let _ = write!(md, "## Execution: `{}`\n\n", self.context.function_name);
203
204        if let Some(host) = &self.context.remote_host {
205            let _ = writeln!(
206                md,
207                "**Host:** {}@{}",
208                self.context.remote_user.as_deref().unwrap_or("?"),
209                host
210            );
211        }
212
213        let _ = writeln!(
214            md,
215            "**Status:** {}",
216            if self.success {
217                "✓ Success"
218            } else {
219                "✗ Failed"
220            }
221        );
222        let _ = write!(md, "**Duration:** {}ms\n\n", self.total_duration_ms);
223
224        // For MCP, we only show output, not implementation
225        // Combine all outputs into a single section
226        let all_stdout: String = self
227            .outputs
228            .iter()
229            .filter(|o| !o.stdout.is_empty())
230            .map(|o| o.stdout.as_str())
231            .collect::<Vec<_>>()
232            .join("");
233
234        let all_stderr: String = self
235            .outputs
236            .iter()
237            .filter(|o| !o.stderr.is_empty())
238            .map(|o| o.stderr.as_str())
239            .collect::<Vec<_>>()
240            .join("");
241
242        if !all_stdout.is_empty() {
243            md.push_str("**Output:**\n```\n");
244            md.push_str(&all_stdout);
245            if !all_stdout.ends_with('\n') {
246                md.push('\n');
247            }
248            md.push_str("```\n\n");
249        }
250
251        if !all_stderr.is_empty() {
252            md.push_str("**Errors:**\n```\n");
253            md.push_str(&all_stderr);
254            if !all_stderr.ends_with('\n') {
255                md.push('\n');
256            }
257            md.push_str("```\n\n");
258        }
259
260        // Show exit code if failed
261        if !self.success
262            && let Some(output) = self.outputs.last()
263            && let Some(code) = output.exit_code
264            && code != 0
265        {
266            let _ = writeln!(md, "**Exit Code:** {code}");
267        }
268
269        md
270    }
271}
272
273/// Static regex for SSH context extraction (compiled once)
274static SSH_REGEX: OnceLock<Regex> = OnceLock::new();
275
276impl ExecutionContext {
277    /// Parse SSH commands to extract remote execution context
278    pub fn extract_ssh_context(command: &str) -> Option<(String, String)> {
279        // Match patterns like:
280        //   ssh user@host
281        //   ssh -i key.pem user@host
282        //   ssh -T -o LogLevel=QUIET user@host
283        // The regex looks for "ssh" followed by optional flags, then user@host
284        let regex = SSH_REGEX.get_or_init(|| {
285            // This regex pattern is hardcoded and known to be valid
286            match Regex::new(r"ssh\s+(?:-\S+\s+(?:\S+\s+)?)*(\w+)@([\w.-]+)") {
287                Ok(r) => r,
288                Err(_) => unreachable!("SSH regex pattern is hardcoded and valid"),
289            }
290        });
291        let caps = regex.captures(command)?;
292
293        Some((
294            caps.get(1)?.as_str().to_string(), // user
295            caps.get(2)?.as_str().to_string(), // host
296        ))
297    }
298}
299
300/// Function parameter definition
301#[derive(Debug, Clone, PartialEq)]
302pub struct Parameter {
303    pub name: String,
304    pub param_type: ArgType,
305    pub default_value: Option<String>,
306    pub is_rest: bool,
307}
308
309#[derive(Debug, Clone, PartialEq)]
310pub enum Statement {
311    Assignment {
312        name: String,
313        value: Expression,
314    },
315    SimpleFunctionDef {
316        name: String,
317        params: Vec<Parameter>,
318        command_template: String,
319        attributes: Vec<Attribute>,
320    },
321    BlockFunctionDef {
322        name: String,
323        params: Vec<Parameter>,
324        commands: Vec<String>,
325        attributes: Vec<Attribute>,
326        shebang: Option<String>,
327    },
328    FunctionCall {
329        name: String,
330        args: Vec<String>,
331    },
332    Command {
333        command: String,
334    },
335}
336
337#[derive(Debug, Clone, PartialEq)]
338pub enum Expression {
339    String(String),
340}
341
342#[derive(Debug, Clone, PartialEq)]
343pub enum Attribute {
344    Os(OsPlatform),
345    Shell(ShellType),
346    Desc(String),
347    Arg(ArgMetadata),
348}
349
350#[derive(Debug, Clone, PartialEq)]
351pub struct ArgMetadata {
352    pub position: usize,
353    pub name: String,
354    pub arg_type: ArgType,
355    pub description: String,
356}
357
358#[derive(Debug, Clone, PartialEq)]
359pub enum ArgType {
360    String,
361    Integer,
362    Float,
363    Boolean,
364    Object,
365}
366
367#[derive(Debug, Clone, PartialEq)]
368pub enum OsPlatform {
369    Windows,
370    Linux,
371    MacOS,
372    Unix, // Matches both Linux and MacOS
373}
374
375#[derive(Debug, Clone, PartialEq)]
376pub enum ShellType {
377    Python,
378    Python3,
379    Node,
380    Ruby,
381    Pwsh,
382    Bash,
383    Sh,
384}
385
386#[cfg(test)]
387#[allow(clippy::expect_used, clippy::unwrap_used)]
388mod tests {
389    use super::*;
390
391    #[test]
392    fn test_extract_ssh_context_basic() {
393        let result = ExecutionContext::extract_ssh_context("ssh admin@webserver.example.com");
394        assert!(result.is_some(), "Failed to match basic SSH command");
395        let (user, host) = result.expect("Expected SSH context to be extracted");
396        assert_eq!(user, "admin");
397        assert_eq!(host, "webserver.example.com");
398    }
399
400    #[test]
401    fn test_extract_ssh_context_with_key() {
402        let result =
403            ExecutionContext::extract_ssh_context("ssh -i ~/.ssh/key.pem ubuntu@192.168.1.1");
404        assert!(result.is_some(), "Failed to match SSH with -i flag");
405        let (user, host) = result.expect("Expected SSH context to be extracted");
406        assert_eq!(user, "ubuntu");
407        assert_eq!(host, "192.168.1.1");
408    }
409
410    #[test]
411    fn test_extract_ssh_context_multiple_options() {
412        let result =
413            ExecutionContext::extract_ssh_context("ssh -T -o LogLevel=QUIET root@server.local");
414        assert!(
415            result.is_some(),
416            "Failed to match SSH with multiple options"
417        );
418        let (user, host) = result.expect("Expected SSH context to be extracted");
419        assert_eq!(user, "root");
420        assert_eq!(host, "server.local");
421    }
422
423    #[test]
424    fn test_extract_ssh_context_no_match() {
425        let result = ExecutionContext::extract_ssh_context("echo hello");
426        assert!(result.is_none());
427    }
428
429    #[test]
430    fn test_output_mode_default() {
431        assert_eq!(OutputMode::default(), OutputMode::Stream);
432    }
433
434    #[test]
435    fn test_structured_result_from_outputs_success() {
436        let outputs = vec![CommandOutput {
437            command: "echo hello".to_string(),
438            stdout: "hello\n".to_string(),
439            stderr: String::new(),
440            exit_code: Some(0),
441            duration_ms: 10,
442            started_at: 1000,
443        }];
444
445        let result = StructuredResult::from_outputs("test_fn", outputs, "sh");
446        assert!(result.success);
447        assert_eq!(result.total_duration_ms, 10);
448        assert_eq!(result.context.function_name, "test_fn");
449        assert_eq!(result.context.interpreter, "sh");
450        assert!(result.context.remote_host.is_none());
451        assert!(result.context.remote_user.is_none());
452        assert!(result.summary.contains("Successfully executed"));
453        assert_eq!(result.outputs.len(), 1);
454    }
455
456    #[test]
457    fn test_structured_result_from_outputs_failure() {
458        let outputs = vec![CommandOutput {
459            command: "false".to_string(),
460            stdout: String::new(),
461            stderr: "error\n".to_string(),
462            exit_code: Some(1),
463            duration_ms: 5,
464            started_at: 1000,
465        }];
466
467        let result = StructuredResult::from_outputs("failing_fn", outputs, "bash");
468        assert!(!result.success);
469        assert!(result.summary.contains("failed"));
470    }
471
472    #[test]
473    fn test_structured_result_from_outputs_with_ssh() {
474        let outputs = vec![CommandOutput {
475            command: "ssh deploy@prod.server.com 'uptime'".to_string(),
476            stdout: "up 10 days\n".to_string(),
477            stderr: String::new(),
478            exit_code: Some(0),
479            duration_ms: 100,
480            started_at: 1000,
481        }];
482
483        let result = StructuredResult::from_outputs("check_uptime", outputs, "sh");
484        assert_eq!(result.context.remote_user.as_deref(), Some("deploy"));
485        assert_eq!(
486            result.context.remote_host.as_deref(),
487            Some("prod.server.com")
488        );
489    }
490
491    #[test]
492    fn test_structured_result_from_outputs_multiple() {
493        let outputs = vec![
494            CommandOutput {
495                command: "echo step1".to_string(),
496                stdout: "step1\n".to_string(),
497                stderr: String::new(),
498                exit_code: Some(0),
499                duration_ms: 5,
500                started_at: 1000,
501            },
502            CommandOutput {
503                command: "echo step2".to_string(),
504                stdout: "step2\n".to_string(),
505                stderr: String::new(),
506                exit_code: Some(0),
507                duration_ms: 10,
508                started_at: 1005,
509            },
510        ];
511
512        let result = StructuredResult::from_outputs("multi", outputs, "sh");
513        assert!(result.success);
514        assert_eq!(result.total_duration_ms, 15);
515        assert_eq!(result.outputs.len(), 2);
516        assert!(result.summary.contains("2 command(s)"));
517    }
518
519    #[test]
520    fn test_structured_result_to_json() {
521        let result = StructuredResult {
522            context: ExecutionContext {
523                function_name: "test".to_string(),
524                remote_host: None,
525                remote_user: None,
526                interpreter: "sh".to_string(),
527                working_directory: None,
528            },
529            outputs: vec![CommandOutput {
530                command: "echo hi".to_string(),
531                stdout: "hi\n".to_string(),
532                stderr: String::new(),
533                exit_code: Some(0),
534                duration_ms: 5,
535                started_at: 1000,
536            }],
537            success: true,
538            total_duration_ms: 5,
539            summary: "ok".to_string(),
540        };
541
542        let json = result.to_json();
543        assert!(json.contains("\"function_name\": \"test\""));
544        assert!(json.contains("\"success\": true"));
545        assert!(json.contains("\"stdout\": \"hi\\n\""));
546        // Verify it's valid JSON
547        let parsed: serde_json::Value = serde_json::from_str(&json).expect("Valid JSON");
548        assert_eq!(parsed["success"], true);
549    }
550
551    #[test]
552    fn test_structured_result_to_markdown() {
553        let result = StructuredResult {
554            context: ExecutionContext {
555                function_name: "deploy".to_string(),
556                remote_host: Some("server.com".to_string()),
557                remote_user: Some("admin".to_string()),
558                interpreter: "bash".to_string(),
559                working_directory: None,
560            },
561            outputs: vec![CommandOutput {
562                command: "deploy.sh".to_string(),
563                stdout: "deployed\n".to_string(),
564                stderr: "warning: slow\n".to_string(),
565                exit_code: Some(0),
566                duration_ms: 100,
567                started_at: 1000,
568            }],
569            success: true,
570            total_duration_ms: 100,
571            summary: "ok".to_string(),
572        };
573
574        let md = result.to_markdown();
575        assert!(md.contains("## Execution: `deploy`"));
576        assert!(md.contains("**Host:** admin@server.com"));
577        assert!(md.contains("✓ Success"));
578        assert!(md.contains("**Duration:** 100ms"));
579        assert!(md.contains("### Step 1"));
580        assert!(md.contains("deployed"));
581        assert!(md.contains("warning: slow"));
582    }
583
584    #[test]
585    fn test_structured_result_to_markdown_failed_with_exit_code() {
586        let result = StructuredResult {
587            context: ExecutionContext {
588                function_name: "fail".to_string(),
589                remote_host: None,
590                remote_user: None,
591                interpreter: "sh".to_string(),
592                working_directory: None,
593            },
594            outputs: vec![CommandOutput {
595                command: "exit 42".to_string(),
596                stdout: String::new(),
597                stderr: "error\n".to_string(),
598                exit_code: Some(42),
599                duration_ms: 1,
600                started_at: 1000,
601            }],
602            success: false,
603            total_duration_ms: 1,
604            summary: "failed".to_string(),
605        };
606
607        let md = result.to_markdown();
608        assert!(md.contains("✗ Failed"));
609        assert!(md.contains("**Exit Code:** 42"));
610    }
611
612    #[test]
613    fn test_structured_result_to_mcp_format() {
614        let result = StructuredResult {
615            context: ExecutionContext {
616                function_name: "test".to_string(),
617                remote_host: None,
618                remote_user: None,
619                interpreter: "sh".to_string(),
620                working_directory: None,
621            },
622            outputs: vec![
623                CommandOutput {
624                    command: "echo a".to_string(),
625                    stdout: "a\n".to_string(),
626                    stderr: String::new(),
627                    exit_code: Some(0),
628                    duration_ms: 5,
629                    started_at: 1000,
630                },
631                CommandOutput {
632                    command: "echo b".to_string(),
633                    stdout: "b\n".to_string(),
634                    stderr: String::new(),
635                    exit_code: Some(0),
636                    duration_ms: 5,
637                    started_at: 1005,
638                },
639            ],
640            success: true,
641            total_duration_ms: 10,
642            summary: "ok".to_string(),
643        };
644
645        let mcp = result.to_mcp_format();
646        assert!(mcp.contains("## Execution: `test`"));
647        assert!(mcp.contains("✓ Success"));
648        // MCP format combines all stdout
649        assert!(mcp.contains("a\n"));
650        assert!(mcp.contains("b\n"));
651        // MCP format should NOT show individual steps
652        assert!(!mcp.contains("### Step"));
653    }
654
655    #[test]
656    fn test_structured_result_to_mcp_format_failed() {
657        let result = StructuredResult {
658            context: ExecutionContext {
659                function_name: "fail".to_string(),
660                remote_host: None,
661                remote_user: None,
662                interpreter: "sh".to_string(),
663                working_directory: None,
664            },
665            outputs: vec![CommandOutput {
666                command: "false".to_string(),
667                stdout: String::new(),
668                stderr: "oh no\n".to_string(),
669                exit_code: Some(1),
670                duration_ms: 1,
671                started_at: 1000,
672            }],
673            success: false,
674            total_duration_ms: 1,
675            summary: "failed".to_string(),
676        };
677
678        let mcp = result.to_mcp_format();
679        assert!(mcp.contains("✗ Failed"));
680        assert!(mcp.contains("oh no"));
681        assert!(mcp.contains("**Exit Code:** 1"));
682    }
683
684    #[test]
685    fn test_structured_result_to_markdown_no_host() {
686        let result = StructuredResult {
687            context: ExecutionContext {
688                function_name: "local".to_string(),
689                remote_host: None,
690                remote_user: None,
691                interpreter: "sh".to_string(),
692                working_directory: None,
693            },
694            outputs: vec![],
695            success: true,
696            total_duration_ms: 0,
697            summary: "ok".to_string(),
698        };
699
700        let md = result.to_markdown();
701        assert!(!md.contains("**Host:**"));
702    }
703}