llm_inspect/
output.rs

1//! Output validation: JSON Schema, allow-lists and HTML checks for model responses.
2//!
3//! The typical use is a model that is asked to return a JSON plan. Before the
4//! plan is executed, the application checks that it is well-formed JSON, that
5//! it matches the schema the model was given, that every `action` field is one
6//! of the actions the application actually permits, and that no HTML is hiding
7//! in string values that will later be rendered.
8
9use std::sync::OnceLock;
10
11use regex::Regex;
12use serde::{Deserialize, Serialize};
13use serde_json::Value;
14
15use crate::Error;
16
17/// An allow-list applied to the values selected by a JSON pointer.
18///
19/// The pointer uses `/` separated segments like RFC 6901, with the addition
20/// that `*` matches every element of an array (or every member of an object).
21/// Example: `/steps/*/action`.
22#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
23pub struct AllowList {
24    /// Pointer to the values to check, e.g. `/steps/*/action`.
25    pub pointer: String,
26    /// Permitted values. Non-string values are compared by their JSON text.
27    pub values: Vec<String>,
28}
29
30/// What an output must satisfy.
31#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
32#[serde(default)]
33pub struct OutputPolicy {
34    /// JSON Schema the output must match. When set, the output must be JSON.
35    #[serde(skip_serializing_if = "Option::is_none")]
36    pub schema: Option<Value>,
37    /// Allow-lists for selected fields. Requires the output to be JSON.
38    pub allow: Vec<AllowList>,
39    /// Reject HTML tags in the output (in string values if JSON, else in the raw text).
40    pub forbid_html: bool,
41    /// Maximum size of the output in bytes.
42    #[serde(skip_serializing_if = "Option::is_none")]
43    pub max_bytes: Option<usize>,
44}
45
46/// Kind of an output violation.
47#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
48#[serde(rename_all = "snake_case")]
49#[non_exhaustive]
50pub enum ViolationKind {
51    /// The output exceeds `max_bytes`.
52    TooLong,
53    /// The policy needs JSON but the output is not valid JSON.
54    NotJson,
55    /// The output does not match the JSON Schema.
56    Schema,
57    /// A value is not in its allow-list.
58    NotAllowed,
59    /// HTML was found.
60    Html,
61}
62
63/// One problem found in an output.
64#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
65pub struct Violation {
66    /// Kind of violation.
67    pub kind: ViolationKind,
68    /// JSON pointer to the offending value (empty for the whole document).
69    pub path: String,
70    /// Human-readable explanation. Never contains the offending value itself.
71    pub message: String,
72}
73
74/// Result of [`OutputValidator::validate`].
75#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
76pub struct OutputReport {
77    /// Violations found; empty means the output is acceptable.
78    pub violations: Vec<Violation>,
79    /// The parsed JSON, when the output parsed as JSON.
80    #[serde(skip_serializing_if = "Option::is_none")]
81    pub json: Option<Value>,
82}
83
84impl OutputReport {
85    /// `true` when there are no violations.
86    pub fn is_valid(&self) -> bool {
87        self.violations.is_empty()
88    }
89}
90
91/// Validates model outputs against an [`OutputPolicy`].
92pub struct OutputValidator {
93    policy: OutputPolicy,
94    schema: Option<jsonschema::Validator>,
95}
96
97impl std::fmt::Debug for OutputValidator {
98    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
99        f.debug_struct("OutputValidator")
100            .field("policy", &self.policy)
101            .finish_non_exhaustive()
102    }
103}
104
105impl OutputValidator {
106    /// Compiles the schema (if any) and checks the allow-list pointers.
107    pub fn new(policy: OutputPolicy) -> Result<Self, Error> {
108        for a in &policy.allow {
109            if !a.pointer.starts_with('/') {
110                return Err(Error::InvalidPointer(a.pointer.clone()));
111            }
112        }
113        let schema = match &policy.schema {
114            Some(s) => Some(
115                jsonschema::validator_for(s).map_err(|e| Error::InvalidSchema(e.to_string()))?,
116            ),
117            None => None,
118        };
119        Ok(Self { policy, schema })
120    }
121
122    /// The policy in effect.
123    pub fn policy(&self) -> &OutputPolicy {
124        &self.policy
125    }
126
127    /// Validates `text`.
128    pub fn validate(&self, text: &str) -> OutputReport {
129        let mut violations = Vec::new();
130
131        if let Some(max) = self.policy.max_bytes.filter(|&max| text.len() > max) {
132            violations.push(Violation {
133                kind: ViolationKind::TooLong,
134                path: String::new(),
135                message: format!("output is {} bytes, limit is {max}", text.len()),
136            });
137        }
138
139        let needs_json = self.schema.is_some() || !self.policy.allow.is_empty();
140        let json: Option<Value> = serde_json::from_str(text).ok();
141
142        match &json {
143            None if needs_json => violations.push(Violation {
144                kind: ViolationKind::NotJson,
145                path: String::new(),
146                message: "output is not valid JSON".to_owned(),
147            }),
148            None => {
149                if self.policy.forbid_html && contains_html(text) {
150                    violations.push(Violation {
151                        kind: ViolationKind::Html,
152                        path: String::new(),
153                        message: "output contains HTML".to_owned(),
154                    });
155                }
156            }
157            Some(v) => {
158                if let Some(schema) = &self.schema {
159                    for err in schema.iter_errors(v) {
160                        violations.push(Violation {
161                            kind: ViolationKind::Schema,
162                            path: err.instance_path().to_string(),
163                            message: schema_message(&err),
164                        });
165                    }
166                }
167                for a in &self.policy.allow {
168                    let mut hits = Vec::new();
169                    select(v, &a.pointer[1..], String::new(), &mut hits);
170                    for (path, val) in hits {
171                        let s = match val {
172                            Value::String(s) => s.clone(),
173                            other => other.to_string(),
174                        };
175                        if !a.values.contains(&s) {
176                            violations.push(Violation {
177                                kind: ViolationKind::NotAllowed,
178                                path,
179                                message: format!(
180                                    "value is not in the allow-list for {}",
181                                    a.pointer
182                                ),
183                            });
184                        }
185                    }
186                }
187                if self.policy.forbid_html {
188                    let mut hits = Vec::new();
189                    strings(v, String::new(), &mut hits);
190                    for (path, s) in hits {
191                        if contains_html(s) {
192                            violations.push(Violation {
193                                kind: ViolationKind::Html,
194                                path,
195                                message: "string value contains HTML".to_owned(),
196                            });
197                        }
198                    }
199                }
200            }
201        }
202
203        OutputReport { violations, json }
204    }
205}
206
207/// A schema error message without the offending instance value.
208fn schema_message(err: &jsonschema::ValidationError<'_>) -> String {
209    let keyword = err.kind();
210    format!("schema violation at `{}`: {keyword:?}", err.schema_path())
211}
212
213/// Walks `pointer` (without its leading `/`) through `v`, expanding `*`.
214fn select<'a>(v: &'a Value, pointer: &str, path: String, out: &mut Vec<(String, &'a Value)>) {
215    if pointer.is_empty() {
216        out.push((path, v));
217        return;
218    }
219    let (head, rest) = pointer.split_once('/').unwrap_or((pointer, ""));
220    let head = head.replace("~1", "/").replace("~0", "~");
221    match v {
222        Value::Array(items) if head == "*" => {
223            for (i, item) in items.iter().enumerate() {
224                select(item, rest, format!("{path}/{i}"), out);
225            }
226        }
227        Value::Array(items) => {
228            if let Some((i, item)) = head
229                .parse::<usize>()
230                .ok()
231                .and_then(|i| Some((i, items.get(i)?)))
232            {
233                select(item, rest, format!("{path}/{i}"), out);
234            }
235        }
236        Value::Object(map) if head == "*" => {
237            for (k, item) in map {
238                select(item, rest, format!("{path}/{}", escape(k)), out);
239            }
240        }
241        Value::Object(map) => {
242            if let Some(item) = map.get(&head) {
243                select(item, rest, format!("{path}/{}", escape(&head)), out);
244            }
245        }
246        _ => {}
247    }
248}
249
250fn escape(segment: &str) -> String {
251    segment.replace('~', "~0").replace('/', "~1")
252}
253
254/// Collects every string value in `v` with its JSON pointer.
255fn strings<'a>(v: &'a Value, path: String, out: &mut Vec<(String, &'a str)>) {
256    match v {
257        Value::String(s) => out.push((path, s)),
258        Value::Array(items) => {
259            for (i, item) in items.iter().enumerate() {
260                strings(item, format!("{path}/{i}"), out);
261            }
262        }
263        Value::Object(map) => {
264            for (k, item) in map {
265                strings(item, format!("{path}/{}", escape(k)), out);
266            }
267        }
268        _ => {}
269    }
270}
271
272fn html_regex() -> &'static Regex {
273    static RE: OnceLock<Regex> = OnceLock::new();
274    RE.get_or_init(|| {
275        Regex::new(r"(?is)</?[a-z][a-z0-9-]*(?:\s[^<>]*)?/?>|<!--|<!\[cdata\[")
276            .expect("valid regex")
277    })
278}
279
280/// Returns `true` if `text` contains something that looks like an HTML tag,
281/// comment or CDATA section.
282///
283/// This is a detector, not a parser: it errs on the side of reporting.
284pub fn contains_html(text: &str) -> bool {
285    html_regex().is_match(text)
286}
287
288/// Removes HTML tags, comments and CDATA markers from `text`.
289///
290/// **This is not an HTML sanitizer.** It does not decode entities, handle
291/// malformed markup the way browsers do, or protect against every XSS vector.
292/// Use it to clean up plain-text display, and a real sanitizer for anything
293/// rendered as HTML.
294pub fn strip_html(text: &str) -> String {
295    html_regex().replace_all(text, "").into_owned()
296}
297
298#[cfg(test)]
299mod tests {
300    use super::*;
301    use serde_json::json;
302
303    #[test]
304    fn detects_tags_comments_and_cdata() {
305        assert!(contains_html("hi <b>there</b>"));
306        assert!(contains_html("<SCRIPT src=x>"));
307        assert!(contains_html("<!-- hidden -->"));
308        assert!(contains_html("<![CDATA[x]]>"));
309        assert!(!contains_html("a < b and c > d"));
310        assert!(!contains_html("1 <2> 3")); // tag names start with a letter
311    }
312
313    #[test]
314    fn strip_removes_tags_only() {
315        assert_eq!(strip_html("<p>Hello <b>world</b></p>"), "Hello world");
316    }
317
318    #[test]
319    fn pointer_wildcards_select_array_elements() {
320        let v = json!({"steps": [{"action": "read"}, {"action": "delete"}]});
321        let mut out = Vec::new();
322        select(&v, "steps/*/action", String::new(), &mut out);
323        let paths: Vec<_> = out.iter().map(|(p, _)| p.as_str()).collect();
324        assert_eq!(paths, ["/steps/0/action", "/steps/1/action"]);
325    }
326
327    #[test]
328    fn pointer_escapes_are_decoded() {
329        let v = json!({"a/b": {"~": 1}});
330        let mut out = Vec::new();
331        select(&v, "a~1b/~0", String::new(), &mut out);
332        assert_eq!(out.len(), 1);
333        assert_eq!(out[0].0, "/a~1b/~0");
334    }
335
336    #[test]
337    fn rejects_bad_pointer() {
338        let err = OutputValidator::new(OutputPolicy {
339            allow: vec![AllowList {
340                pointer: "steps".into(),
341                values: vec![],
342            }],
343            ..Default::default()
344        })
345        .unwrap_err();
346        assert!(matches!(err, Error::InvalidPointer(_)));
347    }
348}
llm_inspect/output.rs

llm_inspect/
output.rs
