Skip to main content

statespace_tool_runtime/
frontmatter.rs

1//! Frontmatter parsing for YAML (`---`) and TOML (`+++`) formats.
2
3use crate::error::Error;
4use crate::spec::ToolSpec;
5use serde::Deserialize;
6
7#[derive(Debug, Clone, Deserialize)]
8struct RawFrontmatter {
9    #[serde(default)]
10    tools: Vec<Vec<serde_json::Value>>,
11}
12
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub struct Frontmatter {
15    pub specs: Vec<ToolSpec>,
16    pub tools: Vec<Vec<String>>,
17}
18
19impl Frontmatter {
20    #[must_use]
21    pub fn has_tool(&self, command: &[String]) -> bool {
22        if command.is_empty() {
23            return false;
24        }
25
26        self.tools.iter().any(|tool| {
27            if tool.is_empty() {
28                return false;
29            }
30
31            if command.len() != tool.len() {
32                return false;
33            }
34
35            if tool[0] != command[0] {
36                return false;
37            }
38
39            true
40        })
41    }
42
43    #[must_use]
44    pub fn tool_names(&self) -> Vec<&str> {
45        self.tools
46            .iter()
47            .filter_map(|tool| tool.first().map(String::as_str))
48            .collect()
49    }
50}
51
52/// # Errors
53///
54/// Returns errors when frontmatter is missing or malformed.
55pub fn parse_frontmatter(content: &str) -> Result<Frontmatter, Error> {
56    if let Some(yaml_content) = extract_yaml_frontmatter(content) {
57        return parse_yaml(&yaml_content);
58    }
59
60    if let Some(toml_content) = extract_toml_frontmatter(content) {
61        return parse_toml(&toml_content);
62    }
63
64    Err(Error::NoFrontmatter)
65}
66
67fn convert_raw(raw: &RawFrontmatter) -> Result<Frontmatter, Error> {
68    let mut specs = Vec::new();
69    let mut tools = Vec::new();
70
71    for tool_parts in &raw.tools {
72        match ToolSpec::parse(tool_parts) {
73            Ok(spec) => specs.push(spec),
74            Err(e) => {
75                return Err(Error::FrontmatterParse(format!("Invalid tool spec: {e}")));
76            }
77        }
78
79        let legacy: Vec<String> = tool_parts
80            .iter()
81            .filter_map(|v| match v {
82                serde_json::Value::String(s) if s != ";" => Some(s.clone()),
83                _ => None,
84            })
85            .collect();
86        if !legacy.is_empty() {
87            tools.push(legacy);
88        }
89    }
90
91    Ok(Frontmatter { specs, tools })
92}
93
94fn extract_yaml_frontmatter(content: &str) -> Option<String> {
95    let trimmed = content.trim_start();
96
97    if !trimmed.starts_with("---") {
98        return None;
99    }
100
101    let after_open = &trimmed[3..];
102    let close_pos = after_open.find("\n---")?;
103
104    Some(after_open[..close_pos].trim().to_string())
105}
106
107fn extract_toml_frontmatter(content: &str) -> Option<String> {
108    let trimmed = content.trim_start();
109
110    if !trimmed.starts_with("+++") {
111        return None;
112    }
113
114    let after_open = &trimmed[3..];
115    let close_pos = after_open.find("\n+++")?;
116
117    Some(after_open[..close_pos].trim().to_string())
118}
119
120fn parse_yaml(content: &str) -> Result<Frontmatter, Error> {
121    let raw: RawFrontmatter = serde_yaml::from_str(content)
122        .map_err(|e| Error::FrontmatterParse(format!("YAML parse error: {e}")))?;
123    convert_raw(&raw)
124}
125
126fn parse_toml(content: &str) -> Result<Frontmatter, Error> {
127    let raw: RawFrontmatter = toml::from_str(content)
128        .map_err(|e| Error::FrontmatterParse(format!("TOML parse error: {e}")))?;
129    convert_raw(&raw)
130}
131
132#[cfg(test)]
133#[allow(clippy::unwrap_used)]
134mod tests {
135    use super::*;
136    use crate::spec::is_valid_tool_call;
137
138    fn legacy_frontmatter(tools: Vec<Vec<String>>) -> Frontmatter {
139        Frontmatter {
140            specs: vec![],
141            tools,
142        }
143    }
144
145    #[test]
146    fn test_parse_yaml_frontmatter() {
147        let markdown = r#"---
148tools:
149  - ["ls", "{path}"]
150  - ["cat", "{path}"]
151---
152
153# Documentation
154"#;
155
156        let fm = parse_frontmatter(markdown).unwrap();
157        assert_eq!(fm.tools.len(), 2);
158        assert_eq!(fm.tools[0], vec!["ls", "{path}"]);
159        assert_eq!(fm.tools[1], vec!["cat", "{path}"]);
160        assert_eq!(fm.specs.len(), 2);
161    }
162
163    #[test]
164    fn test_parse_toml_frontmatter() {
165        let markdown = r#"+++
166tools = [
167  ["ls", "{path}"],
168  ["cat", "{path}"],
169]
170+++
171
172# Documentation
173"#;
174
175        let fm = parse_frontmatter(markdown).unwrap();
176        assert_eq!(fm.tools.len(), 2);
177        assert_eq!(fm.tools[0], vec!["ls", "{path}"]);
178    }
179
180    #[test]
181    fn test_no_frontmatter() {
182        let markdown = "# Just a regular markdown file";
183        let result = parse_frontmatter(markdown);
184        assert!(matches!(result, Err(Error::NoFrontmatter)));
185    }
186
187    #[test]
188    fn test_has_tool() {
189        let fm = legacy_frontmatter(vec![
190            vec!["ls".to_string(), "{path}".to_string()],
191            vec!["cat".to_string(), "{path}".to_string()],
192            vec!["search".to_string()],
193        ]);
194
195        assert!(fm.has_tool(&["search".to_string()]));
196        assert!(fm.has_tool(&["ls".to_string(), "docs/".to_string()]));
197        assert!(fm.has_tool(&["cat".to_string(), "index.md".to_string()]));
198        assert!(!fm.has_tool(&["grep".to_string(), "pattern".to_string()]));
199        assert!(!fm.has_tool(&[]));
200    }
201
202    #[test]
203    fn test_tool_names() {
204        let fm = legacy_frontmatter(vec![
205            vec!["ls".to_string()],
206            vec!["cat".to_string()],
207            vec!["search".to_string()],
208        ]);
209
210        let names = fm.tool_names();
211        assert_eq!(names, vec!["ls", "cat", "search"]);
212    }
213
214    #[test]
215    fn test_e2e_regex_constraint() {
216        let markdown = r#"---
217tools:
218  - [psql, -c, { regex: "^SELECT" }, ";"]
219---
220"#;
221
222        let fm = parse_frontmatter(markdown).unwrap();
223
224        assert!(is_valid_tool_call(
225            &[
226                "psql".to_string(),
227                "-c".to_string(),
228                "SELECT * FROM users".to_string()
229            ],
230            &fm.specs
231        ));
232
233        assert!(!is_valid_tool_call(
234            &[
235                "psql".to_string(),
236                "-c".to_string(),
237                "INSERT INTO users VALUES (1)".to_string()
238            ],
239            &fm.specs
240        ));
241
242        assert!(!is_valid_tool_call(
243            &[
244                "psql".to_string(),
245                "-c".to_string(),
246                "SELECT 1".to_string(),
247                "--extra".to_string()
248            ],
249            &fm.specs
250        ));
251    }
252
253    #[test]
254    fn test_e2e_options_control() {
255        let markdown = r#"---
256tools:
257  - [ls]
258  - [cat, { }, ";"]
259---
260"#;
261
262        let fm = parse_frontmatter(markdown).unwrap();
263
264        assert!(is_valid_tool_call(&["ls".to_string()], &fm.specs));
265        assert!(is_valid_tool_call(
266            &["ls".to_string(), "-la".to_string()],
267            &fm.specs
268        ));
269        assert!(is_valid_tool_call(
270            &["ls".to_string(), "-la".to_string(), "docs/".to_string()],
271            &fm.specs
272        ));
273
274        assert!(is_valid_tool_call(
275            &["cat".to_string(), "file.txt".to_string()],
276            &fm.specs
277        ));
278        assert!(!is_valid_tool_call(
279            &["cat".to_string(), "file.txt".to_string(), "-n".to_string()],
280            &fm.specs
281        ));
282    }
283}