ag_ui_a2ui/toolkit/
recovery.rs

1//! The validate-and-retry loop around a generating model.
2//!
3//! Prompted generation is not schema-constrained, so a model will sometimes
4//! return a surface that does not hold together: a child id it never defined, a
5//! missing root, a loop. That is recoverable — the validator says exactly what
6//! is wrong in words the model can act on, so the fix is to hand the errors back
7//! and ask again.
8//!
9//! [`generate_with_recovery`] runs that loop up to
10//! [`MAX_A2UI_ATTEMPTS`] times, appending
11//! the formatted errors to the prompt between attempts and reporting each step
12//! under the [`A2UI_RECOVERY_ACTIVITY_TYPE`] activity type so a caller can show
13//! progress rather than a stall.
14//!
15//! The loop is synchronous and takes the model as a closure, so it imposes no
16//! async runtime: wrap a blocking call directly, or drive an async client with
17//! whatever executor the host already uses.
18//!
19//! ```
20//! use ag_ui_a2ui::catalog::Catalog;
21//! use ag_ui_a2ui::toolkit::recovery::{generate_with_recovery, RecoveryOptions};
22//!
23//! fn response(components: &str) -> String {
24//!     format!(
25//!         r#"<a2ui-json>[
26//!              {{"version":"v0.9","createSurface":{{"surfaceId":"s","catalogId":"c"}}}},
27//!              {{"version":"v0.9","updateComponents":{{"surfaceId":"s","components":{components}}}}}
28//!            ]</a2ui-json>"#
29//!     )
30//! }
31//!
32//! let catalog = Catalog::basic();
33//! let mut attempt = 0;
34//! let outcome = generate_with_recovery(
35//!     "build a greeting card",
36//!     &catalog,
37//!     &RecoveryOptions::default(),
38//!     |prompt, _n| {
39//!         attempt += 1;
40//!         Ok(if attempt == 1 {
41//!             // First try references a component that was never defined.
42//!             assert!(!prompt.contains("Correction required"));
43//!             response(r#"[{"id":"root","component":"Card","child":"missing"}]"#)
44//!         } else {
45//!             // The retry prompt now carries the validator's complaint.
46//!             assert!(prompt.contains("unresolved_child"));
47//!             response(r#"[{"id":"root","component":"Text","text":"hi"}]"#)
48//!         })
49//!     },
50//!     |_activity| {},
51//! )
52//! .unwrap();
53//!
54//! assert_eq!(outcome.attempts, 2);
55//! assert_eq!(outcome.components.len(), 1);
56//! ```
57
58use serde_json::Value;
59
60use crate::catalog::Catalog;
61use crate::constants::{A2UI_RECOVERY_ACTIVITY_TYPE, MAX_A2UI_ATTEMPTS};
62use crate::error::{Error, Result, ValidationErrors};
63use crate::message::{AgentMessage, AgentPayload, Component};
64use crate::toolkit::parser::parse_response;
65use crate::toolkit::prompt::augment_prompt_with_errors;
66use crate::validate::{ValidateOptions, ValidationError, Validator};
67
68/// How the recovery loop should behave.
69#[derive(Debug, Clone, PartialEq, Eq)]
70pub struct RecoveryOptions {
71    /// Total attempts before giving up. Defaults to [`MAX_A2UI_ATTEMPTS`].
72    pub max_attempts: u32,
73    /// The contract each attempt is held to.
74    pub validate: ValidateOptions,
75}
76
77impl Default for RecoveryOptions {
78    fn default() -> Self {
79        Self {
80            max_attempts: MAX_A2UI_ATTEMPTS,
81            validate: ValidateOptions::full_surface(),
82        }
83    }
84}
85
86impl RecoveryOptions {
87    /// Options for editing a surface that already exists.
88    pub fn for_update() -> Self {
89        Self {
90            validate: ValidateOptions::incremental_update(),
91            ..Self::default()
92        }
93    }
94}
95
96/// What happened on one attempt.
97#[derive(Debug, Clone, PartialEq, Eq)]
98pub struct RecoveryActivity {
99    /// Always [`A2UI_RECOVERY_ACTIVITY_TYPE`], so callers can route on it.
100    pub activity_type: &'static str,
101    /// 1-based attempt number.
102    pub attempt: u32,
103    /// Total attempts allowed.
104    pub max_attempts: u32,
105    /// How the attempt ended.
106    pub status: RecoveryStatus,
107    /// A sentence describing this step, suitable for showing to a user.
108    pub message: String,
109    /// Validation failures from this attempt, empty on success.
110    pub errors: Vec<ValidationError>,
111}
112
113/// How one attempt ended.
114#[derive(Debug, Clone, Copy, PartialEq, Eq)]
115pub enum RecoveryStatus {
116    /// The model is being asked to generate.
117    Started,
118    /// The output validated; the loop is done.
119    Succeeded,
120    /// The output failed validation and another attempt follows.
121    Retrying,
122    /// The attempts are used up.
123    Failed,
124}
125
126impl RecoveryStatus {
127    /// The wire string for this status.
128    pub fn as_str(self) -> &'static str {
129        match self {
130            RecoveryStatus::Started => "started",
131            RecoveryStatus::Succeeded => "succeeded",
132            RecoveryStatus::Retrying => "retrying",
133            RecoveryStatus::Failed => "failed",
134        }
135    }
136}
137
138/// A surface that survived validation.
139#[derive(Debug, Clone, PartialEq)]
140pub struct RecoveredSurface {
141    /// The operations the model produced, in order.
142    pub operations: Vec<AgentMessage>,
143    /// The components those operations define, folded together.
144    pub components: Vec<Component>,
145    /// The data model those operations build.
146    pub data_model: Value,
147    /// Conversational text the model wrote around the A2UI blocks.
148    pub text: String,
149    /// How many attempts it took, 1 when the first try was clean.
150    pub attempts: u32,
151}
152
153/// Generates a surface, retrying with the validator's complaints on failure.
154///
155/// `generate` receives the prompt for this attempt and the 1-based attempt
156/// number, and returns the model's raw response. `on_activity` is called for
157/// every step, including the initial start.
158///
159/// # Errors
160///
161/// Returns [`Error::RecoveryExhausted`] when every attempt fails validation, or
162/// whatever error `generate` itself returns. A response that cannot be parsed at
163/// all is treated as a failed attempt and retried, since that is exactly the
164/// case retrying tends to fix.
165pub fn generate_with_recovery(
166    prompt: &str,
167    catalog: &Catalog,
168    options: &RecoveryOptions,
169    mut generate: impl FnMut(&str, u32) -> Result<String>,
170    mut on_activity: impl FnMut(&RecoveryActivity),
171) -> Result<RecoveredSurface> {
172    let max_attempts = options.max_attempts.max(1);
173    let validator = Validator::with_options(catalog, options.validate.clone());
174    let mut errors: Vec<ValidationError> = Vec::new();
175
176    for attempt in 1..=max_attempts {
177        let attempt_prompt = augment_prompt_with_errors(prompt, &errors);
178        on_activity(&RecoveryActivity {
179            activity_type: A2UI_RECOVERY_ACTIVITY_TYPE,
180            attempt,
181            max_attempts,
182            status: RecoveryStatus::Started,
183            message: if attempt == 1 {
184                "Generating the A2UI surface.".to_string()
185            } else {
186                format!(
187                    "Retrying the A2UI surface after {} validation error(s).",
188                    errors.len()
189                )
190            },
191            errors: Vec::new(),
192        });
193
194        let response = generate(&attempt_prompt, attempt)?;
195
196        match interpret(&response, &validator) {
197            Ok(mut surface) => {
198                surface.attempts = attempt;
199                on_activity(&RecoveryActivity {
200                    activity_type: A2UI_RECOVERY_ACTIVITY_TYPE,
201                    attempt,
202                    max_attempts,
203                    status: RecoveryStatus::Succeeded,
204                    message: format!(
205                        "A2UI surface validated on attempt {attempt} with {} component(s).",
206                        surface.components.len()
207                    ),
208                    errors: Vec::new(),
209                });
210                return Ok(surface);
211            }
212            Err(attempt_errors) => {
213                errors = attempt_errors;
214                let last = attempt == max_attempts;
215                on_activity(&RecoveryActivity {
216                    activity_type: A2UI_RECOVERY_ACTIVITY_TYPE,
217                    attempt,
218                    max_attempts,
219                    status: if last {
220                        RecoveryStatus::Failed
221                    } else {
222                        RecoveryStatus::Retrying
223                    },
224                    message: if last {
225                        format!(
226                            "Gave up after {max_attempts} attempt(s); {} error(s) remain.",
227                            errors.len()
228                        )
229                    } else {
230                        format!(
231                            "Attempt {attempt} produced {} validation error(s); retrying.",
232                            errors.len()
233                        )
234                    },
235                    errors: errors.clone(),
236                });
237            }
238        }
239    }
240
241    Err(Error::RecoveryExhausted {
242        attempts: max_attempts,
243        last: ValidationErrors(errors),
244    })
245}
246
247/// Parses and validates one model response.
248///
249/// A parse failure is reported as a validation error rather than a hard error,
250/// so the loop treats "the model wrote something unparseable" the same way it
251/// treats "the model wrote something inconsistent": tell it, and ask again.
252///
253/// Response-level failures — unparseable output, or a message that matches no
254/// envelope — are reported as [`ErrorCode::EmptyComponents`] at path `response`,
255/// since the outcome is the same in each case: no components could be extracted.
256/// The error-code set is a fixed contract, so no new code is invented for them.
257///
258/// [`ErrorCode::EmptyComponents`]: crate::validate::ErrorCode::EmptyComponents
259fn interpret(
260    response: &str,
261    validator: &Validator<'_>,
262) -> std::result::Result<RecoveredSurface, Vec<ValidationError>> {
263    let parts = match parse_response(response) {
264        Ok(parts) => parts,
265        Err(error) => {
266            return Err(vec![ValidationError::new(
267                crate::validate::ErrorCode::EmptyComponents,
268                "response",
269                format!(
270                    "{error} Return the A2UI messages as a JSON array wrapped in the required \
271                     tags."
272                ),
273            )]);
274        }
275    };
276
277    let mut operations: Vec<AgentMessage> = Vec::new();
278    let mut text_parts: Vec<String> = Vec::new();
279    for part in &parts {
280        if !part.text.is_empty() {
281            text_parts.push(part.text.clone());
282        }
283        let Some(messages) = &part.a2ui else { continue };
284        for message in messages {
285            match serde_json::from_value::<AgentMessage>(message.clone()) {
286                Ok(operation) => operations.push(operation),
287                Err(error) => {
288                    return Err(vec![ValidationError::new(
289                        crate::validate::ErrorCode::EmptyComponents,
290                        "response",
291                        format!(
292                            "A message did not match any A2UI envelope ({error}). Each message \
293                             must be an object with a 'version' and exactly one of createSurface, \
294                             updateComponents, updateDataModel or deleteSurface."
295                        ),
296                    )]);
297                }
298            }
299        }
300    }
301
302    let report = validator.validate_messages(&operations);
303    if !report.is_valid() {
304        return Err(report.errors);
305    }
306
307    let mut components: Vec<Component> = Vec::new();
308    let mut data_model = Value::Null;
309    for operation in &operations {
310        match &operation.payload {
311            AgentPayload::UpdateComponents(update) => {
312                components.extend(update.components.iter().cloned());
313            }
314            AgentPayload::UpdateDataModel(update) => {
315                let _ = update.apply(&mut data_model);
316            }
317            _ => {}
318        }
319    }
320
321    Ok(RecoveredSurface {
322        operations,
323        components,
324        data_model,
325        text: text_parts.join("\n"),
326        attempts: 1,
327    })
328}
329
330#[cfg(test)]
331mod tests {
332    use super::*;
333    use serde_json::json;
334
335    fn wrap(components: Value) -> String {
336        format!(
337            "<a2ui-json>[{{\"version\":\"v0.9\",\"createSurface\":{{\"surfaceId\":\"s\",\
338             \"catalogId\":\"c\"}}}},{{\"version\":\"v0.9\",\"updateComponents\":\
339             {{\"surfaceId\":\"s\",\"components\":{components}}}}}]</a2ui-json>"
340        )
341    }
342
343    fn good() -> String {
344        wrap(json!([{"id": "root", "component": "Text", "text": "hi"}]))
345    }
346
347    fn bad() -> String {
348        wrap(json!([{"id": "root", "component": "Card", "child": "nope"}]))
349    }
350
351    #[test]
352    fn a_clean_first_attempt_does_not_retry() {
353        let catalog = Catalog::basic();
354        let mut calls = 0;
355        let mut activities = Vec::new();
356        let surface = generate_with_recovery(
357            "prompt",
358            &catalog,
359            &RecoveryOptions::default(),
360            |_, _| {
361                calls += 1;
362                Ok(good())
363            },
364            |activity| activities.push(activity.clone()),
365        )
366        .unwrap();
367
368        assert_eq!(calls, 1);
369        assert_eq!(surface.attempts, 1);
370        assert_eq!(surface.components.len(), 1);
371        assert_eq!(surface.operations.len(), 2);
372        assert_eq!(
373            activities.iter().map(|a| a.status).collect::<Vec<_>>(),
374            vec![RecoveryStatus::Started, RecoveryStatus::Succeeded]
375        );
376        assert!(
377            activities
378                .iter()
379                .all(|a| a.activity_type == A2UI_RECOVERY_ACTIVITY_TYPE)
380        );
381    }
382
383    #[test]
384    fn a_failed_attempt_puts_the_errors_in_the_next_prompt() {
385        let catalog = Catalog::basic();
386        let mut prompts = Vec::new();
387        let mut calls = 0;
388        let surface = generate_with_recovery(
389            "base prompt",
390            &catalog,
391            &RecoveryOptions::default(),
392            |prompt, attempt| {
393                prompts.push(prompt.to_string());
394                calls += 1;
395                assert_eq!(attempt, calls);
396                Ok(if calls == 1 { bad() } else { good() })
397            },
398            |_| {},
399        )
400        .unwrap();
401
402        assert_eq!(surface.attempts, 2);
403        assert_eq!(prompts.len(), 2);
404        assert!(!prompts[0].contains("Correction required"));
405        assert!(prompts[1].contains("Correction required"));
406        assert!(prompts[1].contains("unresolved_child"));
407        assert!(prompts[1].contains("components[0].child"));
408    }
409
410    #[test]
411    fn three_failures_exhaust_the_loop_and_report_the_last_errors() {
412        let catalog = Catalog::basic();
413        let mut calls = 0;
414        let mut activities = Vec::new();
415        let error = generate_with_recovery(
416            "prompt",
417            &catalog,
418            &RecoveryOptions::default(),
419            |_, _| {
420                calls += 1;
421                Ok(bad())
422            },
423            |activity| activities.push(activity.clone()),
424        )
425        .unwrap_err();
426
427        assert_eq!(calls, MAX_A2UI_ATTEMPTS);
428        let Error::RecoveryExhausted { attempts, last } = error else {
429            panic!("expected RecoveryExhausted");
430        };
431        assert_eq!(attempts, MAX_A2UI_ATTEMPTS);
432        assert!(last.to_string().contains("unresolved_child"));
433
434        let statuses: Vec<_> = activities.iter().map(|a| a.status).collect();
435        assert_eq!(
436            statuses,
437            vec![
438                RecoveryStatus::Started,
439                RecoveryStatus::Retrying,
440                RecoveryStatus::Started,
441                RecoveryStatus::Retrying,
442                RecoveryStatus::Started,
443                RecoveryStatus::Failed,
444            ]
445        );
446    }
447
448    #[test]
449    fn unparseable_output_is_retried_rather_than_raised() {
450        let catalog = Catalog::basic();
451        let mut calls = 0;
452        let surface = generate_with_recovery(
453            "prompt",
454            &catalog,
455            &RecoveryOptions::default(),
456            |_, _| {
457                calls += 1;
458                Ok(if calls == 1 {
459                    "I'm afraid I can't do that.".to_string()
460                } else {
461                    good()
462                })
463            },
464            |_| {},
465        )
466        .unwrap();
467        assert_eq!(surface.attempts, 2);
468    }
469
470    #[test]
471    fn a_generator_error_stops_the_loop_immediately() {
472        let catalog = Catalog::basic();
473        let mut calls = 0;
474        let error = generate_with_recovery(
475            "prompt",
476            &catalog,
477            &RecoveryOptions::default(),
478            |_, _| {
479                calls += 1;
480                Err(Error::parse("model is offline"))
481            },
482            |_| {},
483        )
484        .unwrap_err();
485        assert_eq!(calls, 1);
486        assert!(matches!(error, Error::Parse(_)));
487    }
488
489    #[test]
490    fn conversational_text_and_data_survive_the_loop() {
491        let catalog = Catalog::basic();
492        let response = "Here is your card.\n<a2ui-json>[{\"version\":\"v0.9\",\"createSurface\":\
493             {\"surfaceId\":\"s\",\"catalogId\":\"c\"}},{\"version\":\"v0.9\",\
494             \"updateComponents\":{\"surfaceId\":\"s\",\"components\":[{\"id\":\"root\",\
495             \"component\":\"Text\",\"text\":{\"path\":\"/name\"}}]}},{\"version\":\
496             \"v0.9\",\"updateDataModel\":{\"surfaceId\":\"s\",\"path\":\"/name\",\
497             \"value\":\"Ada\"}}]</a2ui-json>\nAnything else?";
498        let surface = generate_with_recovery(
499            "prompt",
500            &catalog,
501            &RecoveryOptions::default(),
502            |_, _| Ok(response.to_string()),
503            |_| {},
504        )
505        .unwrap();
506        assert_eq!(surface.text, "Here is your card.\nAnything else?");
507        assert_eq!(surface.data_model, json!({"name": "Ada"}));
508    }
509
510    #[test]
511    fn update_options_accept_an_incremental_payload() {
512        let catalog = Catalog::basic();
513        let response = "<a2ui-json>[{\"version\":\"v0.9\",\"updateComponents\":\
514                        {\"surfaceId\":\"s\",\"components\":[{\"id\":\"label\",\
515                        \"component\":\"Text\",\"text\":\"updated\"}]}}]</a2ui-json>";
516        let surface = generate_with_recovery(
517            "prompt",
518            &catalog,
519            &RecoveryOptions::for_update(),
520            |_, _| Ok(response.to_string()),
521            |_| {},
522        )
523        .unwrap();
524        assert_eq!(surface.components.len(), 1);
525        assert!(
526            !surface
527                .operations
528                .iter()
529                .any(|op| matches!(op.payload, AgentPayload::CreateSurface(_)))
530        );
531    }
532
533    #[test]
534    fn max_attempts_is_honoured_and_never_zero() {
535        let catalog = Catalog::basic();
536        let mut calls = 0;
537        let options = RecoveryOptions {
538            max_attempts: 0,
539            ..RecoveryOptions::default()
540        };
541        let _ = generate_with_recovery(
542            "prompt",
543            &catalog,
544            &options,
545            |_, _| {
546                calls += 1;
547                Ok(bad())
548            },
549            |_| {},
550        );
551        assert_eq!(calls, 1, "a zero budget must still make one attempt");
552    }
553
554    #[test]
555    fn statuses_have_stable_wire_strings() {
556        assert_eq!(RecoveryStatus::Started.as_str(), "started");
557        assert_eq!(RecoveryStatus::Succeeded.as_str(), "succeeded");
558        assert_eq!(RecoveryStatus::Retrying.as_str(), "retrying");
559        assert_eq!(RecoveryStatus::Failed.as_str(), "failed");
560    }
561}
ag_ui_a2ui/toolkit/recovery.rs

ag_ui_a2ui/toolkit/
recovery.rs
