kroki_rs/diagrams/providers/
plugin.rs

1use crate::diagrams::{DiagramError, DiagramProvider, DiagramResult};
2use async_trait::async_trait;
3use std::process::Stdio;
4use tokio::process::Command;
5
6/// A diagram provider that executes an external plugin via a subprocess protocol.
7///
8/// It communicates with the plugin using stdin (optional) and captures stdout
9/// for the rendered diagram. It supports argument templating for formats.
10pub struct PluginProvider {
11    /// Human-readable name of the plugin (e.g. "mytool").
12    pub name: String,
13    /// Path to the executable binary.
14    pub command: String,
15    /// Arguments passed to the binary. Supports `{format}` substitution.
16    pub args: Vec<String>,
17    /// Whether to pipe the diagram source to the binary's stdin.
18    pub stdin: bool,
19    /// Maximum time allowed for the plugin to run.
20    pub timeout_ms: u64,
21}
22
23impl PluginProvider {
24    /// Creates a new plugin provider from the given configuration.
25    pub fn new(config: &crate::config::PluginConfig) -> Self {
26        Self {
27            name: config.name.clone(),
28            command: config.command.clone(),
29            args: config.args.clone(),
30            stdin: config.stdin,
31            timeout_ms: config.timeout_ms.unwrap_or(5000),
32        }
33    }
34}
35
36#[async_trait]
37impl DiagramProvider for PluginProvider {
38    fn validate(&self, source: &str) -> DiagramResult<()> {
39        if source.trim().is_empty() {
40            return Err(DiagramError::ValidationFailed(
41                "Diagram source is empty".into(),
42            ));
43        }
44        Ok(())
45    }
46
47    async fn generate(&self, source: &str, format: &str) -> DiagramResult<Vec<u8>> {
48        let mut cmd = Command::new(&self.command);
49
50        // Template substitution for {format}
51        for arg in &self.args {
52            cmd.arg(arg.replace("{format}", format));
53        }
54
55        if self.stdin {
56            cmd.stdin(Stdio::piped());
57        } else {
58            cmd.stdin(Stdio::null());
59        }
60
61        cmd.stdout(Stdio::piped());
62        cmd.stderr(Stdio::piped());
63
64        let input_bytes = if self.stdin {
65            Some(source.as_bytes())
66        } else {
67            None
68        };
69
70        let output = crate::diagrams::run_process_with_timeout(
71            &self.name,
72            cmd,
73            input_bytes,
74            Some(self.timeout_ms),
75            source.len(),
76        )
77        .await?;
78
79        if output.status.success() {
80            if output.stdout.is_empty() {
81                return Err(DiagramError::ProcessFailed(format!(
82                    "Plugin '{}' succeeded but returned empty output",
83                    self.name
84                )));
85            }
86            Ok(output.stdout)
87        } else {
88            let stderr = String::from_utf8_lossy(&output.stderr);
89            Err(DiagramError::ProcessFailed(format!(
90                "Plugin '{}' failed: {}",
91                self.name, stderr
92            )))
93        }
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[tokio::test]
102    async fn test_plugin_provider_stdout() {
103        let config = crate::config::PluginConfig {
104            name: "test-echo".to_string(),
105            command: "echo".to_string(),
106            args: vec!["hello".to_string()],
107            stdin: false,
108            formats: vec!["txt".to_string()],
109            timeout_ms: Some(1000),
110        };
111        let provider = PluginProvider::new(&config);
112        let result = provider.generate("", "txt").await.unwrap();
113        assert_eq!(String::from_utf8_lossy(&result).trim(), "hello");
114    }
115
116    #[tokio::test]
117    async fn test_plugin_provider_stdin() {
118        // Use 'cat' as a simple plugin that echoes stdin to stdout
119        let config = crate::config::PluginConfig {
120            name: "test-cat".to_string(),
121            command: "cat".to_string(),
122            args: vec![],
123            stdin: true,
124            formats: vec!["txt".to_string()],
125            timeout_ms: Some(1000),
126        };
127        let provider = PluginProvider::new(&config);
128        let result = provider.generate("input-text", "txt").await.unwrap();
129        assert_eq!(String::from_utf8_lossy(&result), "input-text");
130    }
131}
kroki_rs/diagrams/providers/plugin.rs

kroki_rs/diagrams/providers/
plugin.rs
