Skip to main content

statespace_tool_runtime/
validation.rs

1//! Command validation and placeholder expansion.
2
3use crate::error::Error;
4use crate::frontmatter::Frontmatter;
5use crate::spec::{ToolPart, ToolSpec, find_matching_spec, is_valid_tool_call};
6use std::collections::HashMap;
7
8/// # Errors
9///
10/// Returns an error when the command is empty or not present in frontmatter.
11pub fn validate_command(frontmatter: &Frontmatter, command: &[String]) -> Result<(), Error> {
12    if command.is_empty() {
13        return Err(Error::InvalidCommand("command cannot be empty".to_string()));
14    }
15
16    if !frontmatter.has_tool(command) {
17        return Err(Error::CommandNotFound {
18            command: command.join(" "),
19        });
20    }
21
22    Ok(())
23}
24
25/// # Errors
26///
27/// Returns an error when the command is empty or does not match any spec.
28pub fn validate_command_with_specs(specs: &[ToolSpec], command: &[String]) -> Result<(), Error> {
29    if command.is_empty() {
30        return Err(Error::InvalidCommand("command cannot be empty".to_string()));
31    }
32
33    if !is_valid_tool_call(command, specs) {
34        return Err(Error::CommandNotFound {
35            command: command.join(" "),
36        });
37    }
38
39    Ok(())
40}
41
42#[must_use]
43pub fn expand_placeholders<S: std::hash::BuildHasher>(
44    command: &[String],
45    args: &HashMap<String, String, S>,
46) -> Vec<String> {
47    command
48        .iter()
49        .map(|part| {
50            let mut result = part.clone();
51
52            for (key, value) in args {
53                let placeholder = format!("{{{key}}}");
54                result = result.replace(&placeholder, value);
55            }
56
57            result
58        })
59        .collect()
60}
61
62#[must_use]
63pub fn expand_env_vars<S: std::hash::BuildHasher>(
64    command: &[String],
65    env: &HashMap<String, String, S>,
66) -> Vec<String> {
67    command
68        .iter()
69        .map(|part| {
70            let mut result = part.clone();
71
72            for (key, value) in env {
73                let var = format!("${key}");
74                result = result.replace(&var, value);
75            }
76
77            result
78        })
79        .collect()
80}
81
82fn expand_literal_segment<S: std::hash::BuildHasher>(
83    segment: &str,
84    env: &HashMap<String, String, S>,
85) -> String {
86    let mut expanded = segment.to_string();
87    for (key, value) in env {
88        let variable = format!("${key}");
89        expanded = expanded.replace(&variable, value);
90    }
91    expanded
92}
93
94/// Expand trusted env only in spec-declared literal `$VAR` segments.
95///
96/// Placeholder-derived arguments remain opaque to prevent secret expansion in
97/// caller-controlled values.
98#[must_use]
99pub fn expand_command_for_execution<S: std::hash::BuildHasher>(
100    command: &[String],
101    specs: &[ToolSpec],
102    env: &HashMap<String, String, S>,
103) -> Vec<String> {
104    if let Some(spec) = find_matching_spec(command, specs) {
105        return command
106            .iter()
107            .enumerate()
108            .map(|(index, part)| match spec.parts.get(index) {
109                Some(ToolPart::Literal(literal)) if literal == part && literal.contains('$') => {
110                    expand_literal_segment(part, env)
111                }
112                _ => part.clone(),
113            })
114            .collect();
115    }
116
117    command.to_vec()
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123    use crate::spec::ToolPart;
124
125    fn legacy_frontmatter(tools: Vec<Vec<String>>) -> Frontmatter {
126        Frontmatter {
127            specs: vec![],
128            tools,
129        }
130    }
131
132    #[test]
133    fn test_validate_command_empty() {
134        let fm = legacy_frontmatter(vec![]);
135        let result = validate_command(&fm, &[]);
136        assert!(matches!(result, Err(Error::InvalidCommand(_))));
137    }
138
139    #[test]
140    fn test_validate_command_not_found() {
141        let fm = legacy_frontmatter(vec![vec!["ls".to_string()]]);
142
143        let result = validate_command(&fm, &["cat".to_string(), "file.md".to_string()]);
144        assert!(matches!(result, Err(Error::CommandNotFound { .. })));
145    }
146
147    #[test]
148    fn test_validate_command_success() {
149        let fm = legacy_frontmatter(vec![
150            vec!["ls".to_string(), "{path}".to_string()],
151            vec!["cat".to_string(), "{path}".to_string()],
152        ]);
153
154        let result = validate_command(&fm, &["ls".to_string(), "docs/".to_string()]);
155        assert!(result.is_ok());
156
157        let result = validate_command(&fm, &["cat".to_string(), "index.md".to_string()]);
158        assert!(result.is_ok());
159    }
160
161    #[test]
162    fn test_expand_placeholders() {
163        let command = vec![
164            "curl".to_string(),
165            "-X".to_string(),
166            "GET".to_string(),
167            "https://api.com/{endpoint}".to_string(),
168        ];
169
170        let mut args = HashMap::new();
171        args.insert("endpoint".to_string(), "orders".to_string());
172
173        let expanded = expand_placeholders(&command, &args);
174        assert_eq!(
175            expanded,
176            vec!["curl", "-X", "GET", "https://api.com/orders"]
177        );
178    }
179
180    #[test]
181    fn test_expand_env_vars() {
182        let command = vec![
183            "curl".to_string(),
184            "-H".to_string(),
185            "Authorization: Bearer $API_KEY".to_string(),
186        ];
187
188        let mut env = HashMap::new();
189        env.insert("API_KEY".to_string(), "secret123".to_string());
190
191        let expanded = expand_env_vars(&command, &env);
192        assert_eq!(
193            expanded,
194            vec!["curl", "-H", "Authorization: Bearer secret123"]
195        );
196    }
197
198    #[test]
199    fn test_expand_command_for_execution_expands_matching_literal_segments() {
200        let specs = vec![ToolSpec {
201            parts: vec![
202                ToolPart::Literal("echo".to_string()),
203                ToolPart::Literal("$SECRET".to_string()),
204            ],
205            options_disabled: false,
206        }];
207        let command = vec!["echo".to_string(), "$SECRET".to_string()];
208        let env = HashMap::from([("SECRET".to_string(), "trusted".to_string())]);
209
210        let expanded = expand_command_for_execution(&command, &specs, &env);
211
212        assert_eq!(expanded, vec!["echo", "trusted"]);
213    }
214
215    #[test]
216    fn test_expand_command_for_execution_does_not_expand_placeholder_segments() {
217        let specs = vec![ToolSpec {
218            parts: vec![
219                ToolPart::Literal("echo".to_string()),
220                ToolPart::Placeholder { regex: None },
221            ],
222            options_disabled: false,
223        }];
224        let command = vec!["echo".to_string(), "$SECRET".to_string()];
225        let env = HashMap::from([("SECRET".to_string(), "trusted".to_string())]);
226
227        let expanded = expand_command_for_execution(&command, &specs, &env);
228
229        assert_eq!(expanded, command);
230    }
231
232    #[test]
233    fn test_expand_command_for_execution_leaves_missing_literal_var_opaque() {
234        let specs = vec![ToolSpec {
235            parts: vec![
236                ToolPart::Literal("echo".to_string()),
237                ToolPart::Literal("$MISSING".to_string()),
238            ],
239            options_disabled: false,
240        }];
241        let command = vec!["echo".to_string(), "$MISSING".to_string()];
242        let env = HashMap::from([("OTHER".to_string(), "value".to_string())]);
243
244        let expanded = expand_command_for_execution(&command, &specs, &env);
245
246        assert_eq!(expanded, command);
247    }
248}