Skip to main content

statespace_tool_runtime/
validation.rs

1//! Command validation and env 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_env_vars<S: std::hash::BuildHasher>(
44    command: &[String],
45    env: &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 env {
53                let var = format!("${key}");
54                result = result.replace(&var, value);
55            }
56
57            result
58        })
59        .collect()
60}
61
62fn expand_literal_segment<S: std::hash::BuildHasher>(
63    segment: &str,
64    env: &HashMap<String, String, S>,
65) -> String {
66    let mut expanded = segment.to_string();
67    for (key, value) in env {
68        let variable = format!("${key}");
69        expanded = expanded.replace(&variable, value);
70    }
71    expanded
72}
73
74/// Expand trusted env only in spec-declared literal `$VAR` segments.
75///
76/// Placeholder-derived arguments remain opaque to prevent secret expansion in
77/// caller-controlled values.
78#[must_use]
79pub fn expand_command_for_execution<S: std::hash::BuildHasher>(
80    command: &[String],
81    specs: &[ToolSpec],
82    env: &HashMap<String, String, S>,
83) -> Vec<String> {
84    if let Some(spec) = find_matching_spec(command, specs) {
85        return command
86            .iter()
87            .enumerate()
88            .map(|(index, part)| match spec.parts.get(index) {
89                Some(ToolPart::Literal(literal)) if literal == part && literal.contains('$') => {
90                    expand_literal_segment(part, env)
91                }
92                _ => part.clone(),
93            })
94            .collect();
95    }
96
97    command.to_vec()
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103    use crate::spec::ToolPart;
104
105    fn legacy_frontmatter(tools: Vec<Vec<String>>) -> Frontmatter {
106        Frontmatter {
107            specs: vec![],
108            tools,
109        }
110    }
111
112    #[test]
113    fn test_validate_command_empty() {
114        let fm = legacy_frontmatter(vec![]);
115        let result = validate_command(&fm, &[]);
116        assert!(matches!(result, Err(Error::InvalidCommand(_))));
117    }
118
119    #[test]
120    fn test_validate_command_not_found() {
121        let fm = legacy_frontmatter(vec![vec!["ls".to_string()]]);
122
123        let result = validate_command(&fm, &["cat".to_string(), "file.md".to_string()]);
124        assert!(matches!(result, Err(Error::CommandNotFound { .. })));
125    }
126
127    #[test]
128    fn test_validate_command_success() {
129        let fm = legacy_frontmatter(vec![
130            vec!["ls".to_string(), "{path}".to_string()],
131            vec!["cat".to_string(), "{path}".to_string()],
132        ]);
133
134        let result = validate_command(&fm, &["ls".to_string(), "docs/".to_string()]);
135        assert!(result.is_ok());
136
137        let result = validate_command(&fm, &["cat".to_string(), "index.md".to_string()]);
138        assert!(result.is_ok());
139    }
140
141    #[test]
142    fn test_expand_env_vars() {
143        let command = vec![
144            "curl".to_string(),
145            "-H".to_string(),
146            "Authorization: Bearer $API_KEY".to_string(),
147        ];
148
149        let mut env = HashMap::new();
150        env.insert("API_KEY".to_string(), "secret123".to_string());
151
152        let expanded = expand_env_vars(&command, &env);
153        assert_eq!(
154            expanded,
155            vec!["curl", "-H", "Authorization: Bearer secret123"]
156        );
157    }
158
159    #[test]
160    fn test_expand_command_for_execution_expands_matching_literal_segments() {
161        let specs = vec![ToolSpec {
162            parts: vec![
163                ToolPart::Literal("echo".to_string()),
164                ToolPart::Literal("$SECRET".to_string()),
165            ],
166            options_disabled: false,
167        }];
168        let command = vec!["echo".to_string(), "$SECRET".to_string()];
169        let env = HashMap::from([("SECRET".to_string(), "trusted".to_string())]);
170
171        let expanded = expand_command_for_execution(&command, &specs, &env);
172
173        assert_eq!(expanded, vec!["echo", "trusted"]);
174    }
175
176    #[test]
177    fn test_expand_command_for_execution_does_not_expand_placeholder_segments() {
178        let specs = vec![ToolSpec {
179            parts: vec![
180                ToolPart::Literal("echo".to_string()),
181                ToolPart::Placeholder { regex: None },
182            ],
183            options_disabled: false,
184        }];
185        let command = vec!["echo".to_string(), "$SECRET".to_string()];
186        let env = HashMap::from([("SECRET".to_string(), "trusted".to_string())]);
187
188        let expanded = expand_command_for_execution(&command, &specs, &env);
189
190        assert_eq!(expanded, command);
191    }
192
193    #[test]
194    fn test_expand_command_for_execution_leaves_missing_literal_var_opaque() {
195        let specs = vec![ToolSpec {
196            parts: vec![
197                ToolPart::Literal("echo".to_string()),
198                ToolPart::Literal("$MISSING".to_string()),
199            ],
200            options_disabled: false,
201        }];
202        let command = vec!["echo".to_string(), "$MISSING".to_string()];
203        let env = HashMap::from([("OTHER".to_string(), "value".to_string())]);
204
205        let expanded = expand_command_for_execution(&command, &specs, &env);
206
207        assert_eq!(expanded, command);
208    }
209}