1use 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#[derive(Debug, Clone, Copy, PartialEq, Default)]
15pub enum OutputMode {
16 #[default]
18 Stream,
19
20 Capture,
23
24 Structured,
27}
28
29#[derive(Debug, Clone, Serialize, Deserialize)]
31pub struct CommandOutput {
32 pub command: String,
34
35 pub stdout: String,
37
38 pub stderr: String,
40
41 pub exit_code: Option<i32>,
43
44 pub duration_ms: u128,
46
47 pub started_at: u128,
49}
50
51#[derive(Debug, Clone, Serialize, Deserialize)]
53pub struct ExecutionContext {
54 pub function_name: String,
56
57 pub remote_host: Option<String>,
59
60 pub remote_user: Option<String>,
62
63 pub interpreter: String,
65
66 pub working_directory: Option<String>,
68}
69
70#[derive(Debug, Clone, Serialize, Deserialize)]
72pub struct StructuredResult {
73 pub context: ExecutionContext,
75
76 pub outputs: Vec<CommandOutput>,
78
79 pub success: bool,
81
82 pub total_duration_ms: u128,
84
85 pub summary: String,
87}
88
89impl StructuredResult {
90 #[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 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 #[must_use]
135 pub fn to_json(&self) -> String {
136 serde_json::to_string_pretty(self).unwrap_or_default()
137 }
138
139 #[must_use]
141 pub fn to_markdown(&self) -> String {
142 let mut md = String::new();
143
144 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 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 #[must_use]
198 pub fn to_mcp_format(&self) -> String {
199 let mut md = String::new();
200
201 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 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 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
273static SSH_REGEX: OnceLock<Regex> = OnceLock::new();
275
276impl ExecutionContext {
277 pub fn extract_ssh_context(command: &str) -> Option<(String, String)> {
279 let regex = SSH_REGEX.get_or_init(|| {
285 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(), caps.get(2)?.as_str().to_string(), ))
297 }
298}
299
300#[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, }
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 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 assert!(mcp.contains("a\n"));
650 assert!(mcp.contains("b\n"));
651 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}