ag_ui_a2ui/toolkit/
parser.rs

1//! Pulling A2UI out of a model's response.
2//!
3//! When A2UI is generated by prompting rather than by structured output, the
4//! model returns prose with one or more A2UI blocks embedded in it, fenced by
5//! [`A2UI_OPEN_TAG`] and [`A2UI_CLOSE_TAG`]. [`parse_response`] splits that into
6//! ordered [`ResponsePart`]s: the conversational text preceding each block, and
7//! the block's parsed JSON.
8//!
9//! # Why this is a scanner and not a string split
10//!
11//! A close tag can appear inside a JSON string literal — a `Text` component
12//! whose content mentions `</a2ui-json>` is perfectly legal — so the scanner
13//! tracks string state and escapes rather than searching for the next
14//! occurrence.
15//!
16//! # Repairing model output
17//!
18//! [`parse_and_fix`] applies the two repairs that are safe because they cannot
19//! change the meaning of valid JSON: normalizing smart quotes, and dropping
20//! trailing commas. It only reaches for them after a straight parse has failed,
21//! and it wraps a lone object in an array, since A2UI payloads are always lists
22//! of messages.
23
24use serde_json::Value;
25
26use crate::constants::{A2UI_CLOSE_TAG, A2UI_OPEN_TAG};
27use crate::error::{Error, Result};
28
29/// One segment of a model response.
30#[derive(Debug, Clone, PartialEq, Default)]
31pub struct ResponsePart {
32    /// Conversational text that preceded this part's A2UI block, trimmed and
33    /// stripped of Markdown fences. Empty when the block came first.
34    pub text: String,
35    /// The raw block content, before parsing.
36    pub raw: Option<String>,
37    /// The parsed A2UI messages, when the block parsed.
38    pub a2ui: Option<Vec<Value>>,
39    /// Whether the block was closed. A truncated stream ends with `false`.
40    pub is_final: bool,
41}
42
43/// Whether the content contains a complete A2UI block.
44///
45/// Both tags must be present; a response that opens a block and stops mid-stream
46/// does not count.
47pub fn has_a2ui_parts(content: &str) -> bool {
48    content.contains(A2UI_OPEN_TAG) && content.contains(A2UI_CLOSE_TAG)
49}
50
51/// Splits a full model response into text and parsed A2UI blocks.
52///
53/// # Errors
54///
55/// Returns [`Error::Parse`] when the response contains no A2UI block, when a
56/// block is empty, when a block is unterminated, or when a block's content is
57/// not JSON even after repair.
58pub fn parse_response(content: &str) -> Result<Vec<ResponsePart>> {
59    let mut parts = unwrap_response(content)?;
60    for part in &mut parts {
61        if let Some(raw) = &part.raw {
62            part.a2ui = Some(parse_and_fix(raw)?);
63        }
64    }
65    Ok(parts)
66}
67
68/// Splits a response into parts without parsing the block contents.
69///
70/// # Errors
71///
72/// See [`parse_response`], minus the JSON failures.
73pub fn unwrap_response(content: &str) -> Result<Vec<ResponsePart>> {
74    let parts = tokenize(content);
75    let mut out = Vec::new();
76    let mut saw_block = false;
77
78    for part in parts {
79        match &part.raw {
80            Some(raw) => {
81                if !part.is_final {
82                    return Err(Error::parse(format!(
83                        "A2UI close tag '{A2UI_CLOSE_TAG}' not found in response."
84                    )));
85                }
86                if raw.is_empty() {
87                    return Err(Error::parse("A2UI JSON part is empty."));
88                }
89                saw_block = true;
90                out.push(part);
91            }
92            None => {
93                if !part.text.is_empty() {
94                    out.push(part);
95                }
96            }
97        }
98    }
99
100    if !saw_block {
101        return Err(Error::parse(format!(
102            "A2UI tags '{A2UI_OPEN_TAG}' and '{A2UI_CLOSE_TAG}' not found in response."
103        )));
104    }
105    Ok(out)
106}
107
108/// Parses a raw A2UI block, repairing the model's usual JSON slips.
109///
110/// Always returns a list: a lone object is wrapped, since an A2UI payload is a
111/// list of messages.
112///
113/// # Errors
114///
115/// Returns [`Error::Parse`] when the content is not JSON even after repair.
116pub fn parse_and_fix(payload: &str) -> Result<Vec<Value>> {
117    let normalized = normalize_smart_quotes(payload);
118    match parse_list(&normalized) {
119        Ok(value) => Ok(value),
120        // Only now try the lossy repair, so valid JSON is never rewritten.
121        Err(first) => parse_list(&remove_trailing_commas(&normalized)).map_err(|_| first),
122    }
123}
124
125fn parse_list(payload: &str) -> Result<Vec<Value>> {
126    let value: Value = serde_json::from_str(payload)
127        .map_err(|e| Error::parse(format!("Failed to parse JSON: {e}")))?;
128    Ok(match value {
129        Value::Array(items) => items,
130        other => vec![other],
131    })
132}
133
134fn normalize_smart_quotes(input: &str) -> String {
135    input
136        .replace(['\u{201C}', '\u{201D}'], "\"")
137        .replace(['\u{2018}', '\u{2019}'], "'")
138}
139
140/// Drops commas that directly precede a closing bracket or brace.
141///
142/// String-aware, so `{"text": "a, b"}` and `["x,"]` are untouched.
143fn remove_trailing_commas(input: &str) -> String {
144    let chars: Vec<char> = input.chars().collect();
145    let mut out = String::with_capacity(input.len());
146    let mut i = 0;
147    let mut in_string = false;
148
149    while i < chars.len() {
150        let ch = chars[i];
151        if in_string {
152            out.push(ch);
153            if ch == '\\' {
154                if let Some(next) = chars.get(i + 1) {
155                    out.push(*next);
156                    i += 2;
157                    continue;
158                }
159            } else if ch == '"' {
160                in_string = false;
161            }
162            i += 1;
163            continue;
164        }
165        match ch {
166            '"' => {
167                in_string = true;
168                out.push(ch);
169            }
170            ',' => {
171                let next = chars[i + 1..].iter().find(|c| !c.is_whitespace()).copied();
172                if !matches!(next, Some(']') | Some('}')) {
173                    out.push(ch);
174                }
175            }
176            _ => out.push(ch),
177        }
178        i += 1;
179    }
180    out
181}
182
183/// Scans a response into alternating text and block parts.
184fn tokenize(content: &str) -> Vec<ResponsePart> {
185    let chars: Vec<char> = content.chars().collect();
186    let mut parts = Vec::new();
187    let mut text = String::new();
188    let mut raw = String::new();
189    let mut i = 0;
190    let mut in_block = false;
191
192    while i < chars.len() {
193        if !in_block {
194            if let Some(end) = match_open_tag(&chars, i) {
195                in_block = true;
196                raw.clear();
197                i = end;
198                continue;
199            }
200            text.push(chars[i]);
201            i += 1;
202            continue;
203        }
204
205        if let Some(end) = match_close_tag(&chars, i) {
206            parts.push(ResponsePart {
207                text: clean_markdown(&text),
208                raw: Some(clean_markdown(&raw)),
209                a2ui: None,
210                is_final: true,
211            });
212            text.clear();
213            raw.clear();
214            in_block = false;
215            i = end;
216            continue;
217        }
218
219        // Inside a JSON string literal a close tag is just text.
220        let ch = chars[i];
221        if ch == '"' || ch == '\'' {
222            raw.push(ch);
223            i += 1;
224            while i < chars.len() {
225                if chars[i] == '\\' {
226                    raw.push(chars[i]);
227                    if let Some(next) = chars.get(i + 1) {
228                        raw.push(*next);
229                    }
230                    i += 2;
231                    continue;
232                }
233                raw.push(chars[i]);
234                let closed = chars[i] == ch;
235                i += 1;
236                if closed {
237                    break;
238                }
239            }
240            continue;
241        }
242        raw.push(ch);
243        i += 1;
244    }
245
246    if in_block {
247        parts.push(ResponsePart {
248            text: clean_markdown(&text),
249            raw: Some(clean_markdown(&raw)),
250            a2ui: None,
251            is_final: false,
252        });
253    } else {
254        let trailing = clean_markdown(&text);
255        if !trailing.is_empty() {
256            parts.push(ResponsePart {
257                text: trailing,
258                raw: None,
259                a2ui: None,
260                is_final: true,
261            });
262        }
263    }
264    parts
265}
266
267/// Matches `<a2ui-json ...>` at `start`, returning the index just past it.
268///
269/// Attributes are tolerated; the tag name must end on a word boundary so
270/// `<a2ui-jsonx>` does not match.
271fn match_open_tag(chars: &[char], start: usize) -> Option<usize> {
272    let name: Vec<char> = A2UI_OPEN_TAG
273        .trim_start_matches('<')
274        .trim_end_matches('>')
275        .chars()
276        .collect();
277    if chars.get(start) != Some(&'<') {
278        return None;
279    }
280    let mut i = start + 1;
281    for expected in &name {
282        let actual = chars.get(i)?;
283        if !actual.eq_ignore_ascii_case(expected) {
284            return None;
285        }
286        i += 1;
287    }
288    if chars
289        .get(i)
290        .is_some_and(|c| c.is_alphanumeric() || *c == '_')
291    {
292        return None;
293    }
294    while let Some(ch) = chars.get(i) {
295        if *ch == '>' {
296            return Some(i + 1);
297        }
298        i += 1;
299    }
300    None
301}
302
303/// Matches `</a2ui-json  >` at `start`, returning the index just past it.
304fn match_close_tag(chars: &[char], start: usize) -> Option<usize> {
305    let name: Vec<char> = A2UI_CLOSE_TAG
306        .trim_start_matches("</")
307        .trim_end_matches('>')
308        .chars()
309        .collect();
310    if chars.get(start) != Some(&'<') || chars.get(start + 1) != Some(&'/') {
311        return None;
312    }
313    let mut i = start + 2;
314    for expected in &name {
315        let actual = chars.get(i)?;
316        if !actual.eq_ignore_ascii_case(expected) {
317            return None;
318        }
319        i += 1;
320    }
321    while chars.get(i).is_some_and(|c| c.is_whitespace()) {
322        i += 1;
323    }
324    (chars.get(i) == Some(&'>')).then_some(i + 1)
325}
326
327/// Trims whitespace and strips a surrounding Markdown code fence.
328fn clean_markdown(text: &str) -> String {
329    let mut text = text.trim();
330    if let Some(rest) = text.strip_prefix("```") {
331        // Drop an optional language tag on the opening fence.
332        let rest = rest.trim_start_matches(|c: char| c.is_alphanumeric() || c == '-');
333        text = rest.trim_start();
334    }
335    if let Some(rest) = text.strip_suffix("```") {
336        text = rest.trim_end();
337    } else if let Some(index) = text.rfind("```") {
338        let tail = &text[index + 3..];
339        if !tail.is_empty() && tail.chars().all(|c| c.is_alphanumeric() || c == '-') {
340            text = text[..index].trim_end();
341        }
342    }
343    text.trim().to_string()
344}
345
346#[cfg(test)]
347mod tests {
348    use super::*;
349    use serde_json::json;
350
351    #[test]
352    fn a_lone_block_yields_one_part_with_empty_text() {
353        let parts = parse_response(r#"<a2ui-json>[{"id": "test"}]</a2ui-json>"#).unwrap();
354        assert_eq!(parts.len(), 1);
355        assert_eq!(parts[0].text, "");
356        assert_eq!(parts[0].a2ui, Some(vec![json!({"id": "test"})]));
357    }
358
359    #[test]
360    fn leading_and_trailing_text_are_separate_parts() {
361        let parts =
362            parse_response("Hello\n<a2ui-json>[{\"id\": \"t\"}]</a2ui-json>\nGoodbye").unwrap();
363        assert_eq!(parts.len(), 2);
364        assert_eq!(parts[0].text, "Hello");
365        assert!(parts[0].a2ui.is_some());
366        assert_eq!(parts[1].text, "Goodbye");
367        assert!(parts[1].a2ui.is_none());
368    }
369
370    #[test]
371    fn markdown_fences_inside_a_block_are_stripped() {
372        let parts =
373            parse_response("Text\n<a2ui-json>\n```json\n[{\"id\": \"t\"}]\n```\n</a2ui-json>")
374                .unwrap();
375        assert_eq!(parts[0].text, "Text");
376        assert_eq!(parts[0].a2ui, Some(vec![json!({"id": "t"})]));
377    }
378
379    #[test]
380    fn a_close_tag_inside_a_string_does_not_end_the_block() {
381        let content = concat!(
382            r#"<a2ui-json>[{"id": "t", "component": "Text", "#,
383            r#""text": "a literal </a2ui-json> inside"}]</a2ui-json>"#
384        );
385        let parts = parse_response(content).unwrap();
386        assert_eq!(parts.len(), 1);
387        let messages = parts[0].a2ui.as_ref().unwrap();
388        assert_eq!(messages[0]["text"], "a literal </a2ui-json> inside");
389    }
390
391    #[test]
392    fn missing_empty_and_unterminated_blocks_all_error() {
393        assert!(parse_response("").is_err());
394        assert!(parse_response("Only text.").is_err());
395        assert!(parse_response("<a2ui-json></a2ui-json>").is_err());
396        assert!(parse_response(r#"<a2ui-json>[{"id": "t"}]"#).is_err());
397        assert!(parse_response("<a2ui-json>\ninvalid_json\n</a2ui-json>").is_err());
398    }
399
400    #[test]
401    fn has_parts_needs_both_tags() {
402        assert!(has_a2ui_parts("Hello <a2ui-json>[]</a2ui-json> World"));
403        assert!(!has_a2ui_parts("Hello World"));
404        assert!(!has_a2ui_parts("Hello <a2ui-json> World"));
405    }
406
407    #[test]
408    fn repairs_are_applied_only_when_needed() {
409        assert_eq!(
410            parse_and_fix(r#"[{"type": "Text", "text": "Hello"},]"#).unwrap(),
411            vec![json!({"type": "Text", "text": "Hello"})]
412        );
413        assert_eq!(
414            parse_and_fix(r#"{"type": "Text"}"#).unwrap(),
415            vec![json!({"type": "Text"})]
416        );
417        assert_eq!(
418            parse_and_fix("{\"type\": \u{201C}Text\u{201D}, \"other\": \"Value\u{2019}s\"}")
419                .unwrap(),
420            vec![json!({"type": "Text", "other": "Value's"})]
421        );
422    }
423
424    #[test]
425    fn commas_inside_strings_survive_the_repair() {
426        assert_eq!(
427            parse_and_fix(r#"{"text": "Hello, world", "array": ["a,b", "c"]}"#).unwrap(),
428            vec![json!({"text": "Hello, world", "array": ["a,b", "c"]})]
429        );
430        // Even when the repair pass has to run because of a real trailing comma.
431        assert_eq!(
432            parse_and_fix(r#"{"text": "a, b", "list": [1, 2,],}"#).unwrap(),
433            vec![json!({"text": "a, b", "list": [1, 2]})]
434        );
435    }
436
437    #[test]
438    fn nested_trailing_commas_are_removed() {
439        assert_eq!(
440            parse_and_fix(r#"{"a": {"b": 1,}}"#).unwrap(),
441            vec![json!({"a": {"b": 1}})]
442        );
443        assert_eq!(
444            parse_and_fix(r#"[{"a": [1, 2, 3,]}]"#).unwrap(),
445            vec![json!({"a": [1, 2, 3]})]
446        );
447    }
448
449    #[test]
450    fn open_tags_with_attributes_match_but_lookalikes_do_not() {
451        let parts =
452            parse_response(r#"<a2ui-json version="v0.9">[{"id":"t"}]</a2ui-json>"#).unwrap();
453        assert_eq!(parts.len(), 1);
454        assert!(parse_response(r#"<a2ui-jsonx>[]</a2ui-jsonx>"#).is_err());
455    }
456
457    /// Two scanners, one grammar.
458    ///
459    /// [`match_open_tag`] matches a tag anchored at a position in a `&[char]`;
460    /// [`crate::toolkit::streaming`]'s searches for one by byte offset and
461    /// additionally has to recognise a tag that a chunk boundary cut in half.
462    /// Neither can be written in terms of the other, so what keeps them honest
463    /// is this: a tolerance added to one and forgotten in the other would mean
464    /// a model's output rendering on one path and vanishing on the other.
465    #[test]
466    fn the_streaming_scanner_opens_a_block_on_exactly_the_tags_this_one_does() {
467        const BODY: &str =
468            r#"[{"version":"v0.9","createSurface":{"surfaceId":"s1","catalogId":"c"}}]"#;
469
470        for tag in [
471            "<a2ui-json>",
472            "<A2UI-JSON>",
473            r#"<a2ui-json version="v0.9">"#,
474            "<a2ui-json  >",
475            "<a2ui-jsonx>",
476            "<a2ui-json_x>",
477            "<a2ui-jso>",
478            "<not-a2ui-json>",
479        ] {
480            let content = format!("hello {tag}{BODY}</a2ui-json>");
481
482            let whole = unwrap_response(&content).is_ok();
483            let streamed =
484                crate::toolkit::streaming::StreamParser::new(crate::catalog::Catalog::basic())
485                    .process_chunk(&content)
486                    .expect("a tag this scanner rejects is conversational text, not an error")
487                    .iter()
488                    .any(|part| part.a2ui.is_some());
489
490            assert_eq!(whole, streamed, "the two scanners disagree about {tag:?}");
491        }
492    }
493
494    #[test]
495    fn multiple_blocks_keep_their_leading_text() {
496        let content = "Part 1\n<a2ui-json>\n[{\"id\": \"1\"}]\n</a2ui-json>\nPart 2\n\
497                       <a2ui-json>\n[{\"id\": \"2\"}]\n</a2ui-json>\nPart 3";
498        let parts = parse_response(content).unwrap();
499        let texts: Vec<&str> = parts.iter().map(|p| p.text.as_str()).collect();
500        assert_eq!(texts, vec!["Part 1", "Part 2", "Part 3"]);
501        assert_eq!(parts[1].a2ui, Some(vec![json!({"id": "2"})]));
502        assert_eq!(parts[2].a2ui, None);
503    }
504}
ag_ui_a2ui/toolkit/parser.rs

ag_ui_a2ui/toolkit/
parser.rs
