systemprompt_models/ai/
template_validation.rs

1//! Static validation of cross-tool template references in an execution plan.
2//!
3//! [`TemplateValidator`] walks each [`PlannedToolCall`]'s arguments for
4//! `$N.output.field` templates and checks them against the referenced tools'
5//! output schemas, surfacing each problem as a [`PlanValidationError`] tagged
6//! with a [`ValidationErrorKind`] (bad syntax, self/forward reference,
7//! out-of-bounds index, or a field absent from the target schema).
8//!
9//! Copyright (c) systemprompt.io — Business Source License 1.1.
10//! See <https://systemprompt.io> for licensing details.
11
12use serde::{Deserialize, Serialize};
13use serde_json::Value;
14
15use super::execution_plan::{PlannedToolCall, TemplateRef};
16use super::tools::McpTool;
17
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct PlanValidationError {
20    pub tool_index: usize,
21    pub argument: String,
22    pub template: String,
23    pub error: ValidationErrorKind,
24}
25
26#[derive(Debug, Clone, Serialize, Deserialize)]
27#[serde(tag = "type", rename_all = "snake_case")]
28pub enum ValidationErrorKind {
29    InvalidTemplateSyntax,
30    IndexOutOfBounds {
31        referenced_index: usize,
32        max_valid_index: usize,
33    },
34    SelfReference,
35    ForwardReference {
36        referenced_index: usize,
37    },
38    FieldNotFound {
39        tool_name: String,
40        field: String,
41        available_fields: Vec<String>,
42    },
43    NoOutputSchema {
44        tool_name: String,
45    },
46}
47
48impl std::fmt::Display for PlanValidationError {
49    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
50        match &self.error {
51            ValidationErrorKind::InvalidTemplateSyntax => {
52                write!(
53                    f,
54                    "Tool {}: Invalid template syntax '{}' for argument '{}'",
55                    self.tool_index, self.template, self.argument
56                )
57            },
58            ValidationErrorKind::IndexOutOfBounds {
59                referenced_index,
60                max_valid_index,
61            } => {
62                write!(
63                    f,
64                    "Tool {}: Template '{}' references tool {} but only tools 0-{} are available",
65                    self.tool_index, self.template, referenced_index, max_valid_index
66                )
67            },
68            ValidationErrorKind::SelfReference => {
69                write!(
70                    f,
71                    "Tool {}: Template '{}' cannot reference itself",
72                    self.tool_index, self.template
73                )
74            },
75            ValidationErrorKind::ForwardReference { referenced_index } => {
76                write!(
77                    f,
78                    "Tool {}: Template '{}' references tool {} which hasn't executed yet",
79                    self.tool_index, self.template, referenced_index
80                )
81            },
82            ValidationErrorKind::FieldNotFound {
83                tool_name,
84                field,
85                available_fields,
86            } => {
87                write!(
88                    f,
89                    "Tool {}: Template '{}' references field '{}' but tool '{}' outputs: [{}]",
90                    self.tool_index,
91                    self.template,
92                    field,
93                    tool_name,
94                    available_fields.join(", ")
95                )
96            },
97            ValidationErrorKind::NoOutputSchema { tool_name } => {
98                write!(
99                    f,
100                    "Tool {}: Template '{}' references '{}' which has no output schema",
101                    self.tool_index, self.template, tool_name
102                )
103            },
104        }
105    }
106}
107
108impl std::error::Error for PlanValidationError {}
109
110#[derive(Debug, Clone, Copy)]
111pub struct TemplateValidator;
112
113impl TemplateValidator {
114    pub fn get_tool_output_schemas(
115        calls: &[PlannedToolCall],
116        tools: &[McpTool],
117    ) -> Vec<(String, Option<Value>)> {
118        calls
119            .iter()
120            .map(|call| {
121                let output_schema = tools
122                    .iter()
123                    .find(|t| t.name == call.tool_name)
124                    .and_then(|t| t.output_schema.clone());
125                (call.tool_name.clone(), output_schema)
126            })
127            .collect()
128    }
129
130    pub fn find_templates_in_value(value: &Value) -> Vec<String> {
131        let mut templates = Vec::new();
132        Self::collect_templates(value, &mut templates);
133        templates
134    }
135
136    fn collect_templates(value: &Value, templates: &mut Vec<String>) {
137        match value {
138            Value::String(s) if s.starts_with('$') && s.contains(".output.") => {
139                templates.push(s.clone());
140            },
141            Value::Array(arr) => {
142                for v in arr {
143                    Self::collect_templates(v, templates);
144                }
145            },
146            Value::Object(obj) => {
147                for v in obj.values() {
148                    Self::collect_templates(v, templates);
149                }
150            },
151            Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => {},
152        }
153    }
154
155    pub fn validate_plan(
156        calls: &[PlannedToolCall],
157        tool_output_schemas: &[(String, Option<Value>)],
158    ) -> Result<(), Vec<PlanValidationError>> {
159        let mut errors = Vec::new();
160
161        for (tool_index, call) in calls.iter().enumerate() {
162            for template in Self::find_templates_in_value(&call.arguments) {
163                if let Some(err) =
164                    Self::validate_template(tool_index, call, &template, tool_output_schemas)
165                {
166                    errors.push(err);
167                }
168            }
169        }
170
171        if errors.is_empty() {
172            Ok(())
173        } else {
174            Err(errors)
175        }
176    }
177
178    fn validate_template(
179        tool_index: usize,
180        call: &PlannedToolCall,
181        template: &str,
182        tool_output_schemas: &[(String, Option<Value>)],
183    ) -> Option<PlanValidationError> {
184        let make_error = |error: ValidationErrorKind| PlanValidationError {
185            tool_index,
186            argument: Self::find_argument_for_template(&call.arguments, template),
187            template: template.to_owned(),
188            error,
189        };
190
191        let Some(template_ref) = TemplateRef::parse(template) else {
192            return Some(make_error(ValidationErrorKind::InvalidTemplateSyntax));
193        };
194
195        if template_ref.tool_index == tool_index {
196            return Some(make_error(ValidationErrorKind::SelfReference));
197        }
198        if template_ref.tool_index > tool_index {
199            return Some(make_error(ValidationErrorKind::ForwardReference {
200                referenced_index: template_ref.tool_index,
201            }));
202        }
203        if template_ref.tool_index >= tool_output_schemas.len() {
204            return Some(make_error(ValidationErrorKind::IndexOutOfBounds {
205                referenced_index: template_ref.tool_index,
206                max_valid_index: tool_output_schemas.len().saturating_sub(1),
207            }));
208        }
209
210        let (ref_tool_name, ref_output_schema) = &tool_output_schemas[template_ref.tool_index];
211
212        ref_output_schema.as_ref().map_or_else(
213            || {
214                Some(make_error(ValidationErrorKind::NoOutputSchema {
215                    tool_name: ref_tool_name.clone(),
216                }))
217            },
218            |schema| {
219                Self::validate_field_access(&template_ref, schema, ref_tool_name).map(make_error)
220            },
221        )
222    }
223
224    fn validate_field_access(
225        template_ref: &TemplateRef,
226        schema: &Value,
227        tool_name: &str,
228    ) -> Option<ValidationErrorKind> {
229        let first_field = template_ref.field_path.first()?;
230        let available_fields = Self::get_schema_fields(schema);
231
232        if available_fields.contains(first_field) {
233            None
234        } else {
235            Some(ValidationErrorKind::FieldNotFound {
236                tool_name: tool_name.to_owned(),
237                field: first_field.clone(),
238                available_fields,
239            })
240        }
241    }
242
243    fn find_argument_for_template(value: &Value, template: &str) -> String {
244        if let Value::Object(obj) = value {
245            for (key, val) in obj {
246                if let Value::String(s) = val
247                    && s == template
248                {
249                    return key.clone();
250                }
251                let nested = Self::find_argument_for_template(val, template);
252                if !nested.is_empty() {
253                    return format!("{key}.{nested}");
254                }
255            }
256        }
257        String::new()
258    }
259
260    fn get_schema_fields(schema: &Value) -> Vec<String> {
261        schema
262            .get("properties")
263            .and_then(|p| p.as_object())
264            .map_or_else(Vec::new, |obj| obj.keys().cloned().collect())
265    }
266}
systemprompt_models/ai/template_validation.rs

systemprompt_models/ai/
template_validation.rs
