awaken_ext_generative_ui/a2ui/
plugin.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use schemars::JsonSchema;
5use serde::{Deserialize, Serialize};
6
7use awaken_contract::PluginConfigKey;
8use awaken_contract::StateError;
9use awaken_contract::contract::context_message::ContextMessage;
10use awaken_contract::model::Phase;
11use awaken_contract::registry_spec::AgentSpec;
12
13use awaken_runtime::agent::state::AddContextMessage;
14use awaken_runtime::plugins::{ConfigSchema, Plugin, PluginDescriptor, PluginRegistrar};
15use awaken_runtime::state::MutationBatch;
16use awaken_runtime::{PhaseContext, PhaseHook, StateCommand};
17
18use super::tool::A2uiRenderTool;
19use super::{A2UI_PLUGIN_ID, A2UI_TOOL_ID};
20
21pub const DEFAULT_A2UI_CATALOG_ID: &str =
22    "https://a2ui.org/specification/v0_8/standard_catalog_definition.json";
23
24const A2UI_INSTRUCTION_CONTEXT_KEY: &str = "generative_ui.instructions";
25
26/// Prompt customization for the A2UI plugin.
27///
28/// Stored in `AgentSpec.sections["generative-ui"]` and resolved on each
29/// inference step, so a caller can override the catalog hint, append
30/// examples, or replace the injected instructions entirely.
31#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
32#[serde(default)]
33pub struct A2uiPromptConfig {
34    /// Override the catalog identifier mentioned in the default instructions.
35    #[serde(default, skip_serializing_if = "Option::is_none")]
36    pub catalog_id: Option<String>,
37    /// Extra examples appended after the built-in instructions.
38    #[serde(default, skip_serializing_if = "Option::is_none")]
39    pub examples: Option<String>,
40    /// Full instruction override. When set, `catalog_id` and `examples`
41    /// are ignored for prompt construction.
42    #[serde(default, skip_serializing_if = "Option::is_none")]
43    pub instructions: Option<String>,
44}
45
46impl A2uiPromptConfig {
47    #[must_use]
48    pub fn with_catalog_id(mut self, catalog_id: impl Into<String>) -> Self {
49        self.catalog_id = Some(catalog_id.into());
50        self
51    }
52
53    #[must_use]
54    pub fn with_examples(mut self, examples: impl Into<String>) -> Self {
55        self.examples = Some(examples.into());
56        self
57    }
58
59    #[must_use]
60    pub fn with_instructions(mut self, instructions: impl Into<String>) -> Self {
61        self.instructions = Some(instructions.into());
62        self
63    }
64
65    #[must_use]
66    pub fn is_empty(&self) -> bool {
67        self.catalog_id.is_none() && self.examples.is_none() && self.instructions.is_none()
68    }
69
70    fn overlay(&self, overrides: Self) -> Self {
71        Self {
72            catalog_id: overrides.catalog_id.or_else(|| self.catalog_id.clone()),
73            examples: overrides.examples.or_else(|| self.examples.clone()),
74            instructions: overrides.instructions.or_else(|| self.instructions.clone()),
75        }
76    }
77
78    fn instructions_text(&self) -> String {
79        if let Some(instructions) = &self.instructions {
80            return instructions.clone();
81        }
82
83        build_instructions(
84            self.catalog_id
85                .as_deref()
86                .unwrap_or(DEFAULT_A2UI_CATALOG_ID),
87            self.examples.as_deref(),
88        )
89    }
90}
91
92/// [`PluginConfigKey`] binding for A2UI prompt overrides in agent specs.
93pub struct A2uiPromptConfigKey;
94
95impl PluginConfigKey for A2uiPromptConfigKey {
96    const KEY: &'static str = A2UI_PLUGIN_ID;
97    type Config = A2uiPromptConfig;
98}
99
100#[derive(Clone)]
101pub(crate) struct A2uiInstructionHook {
102    defaults: A2uiPromptConfig,
103}
104
105impl A2uiInstructionHook {
106    pub(crate) fn new(defaults: A2uiPromptConfig) -> Self {
107        Self { defaults }
108    }
109
110    pub(crate) fn instructions_for(&self, agent_spec: &AgentSpec) -> Result<String, StateError> {
111        let overrides = agent_spec.config::<A2uiPromptConfigKey>()?;
112        Ok(self.defaults.overlay(overrides).instructions_text())
113    }
114}
115
116#[async_trait]
117impl PhaseHook for A2uiInstructionHook {
118    async fn run(&self, ctx: &PhaseContext) -> Result<StateCommand, StateError> {
119        let instructions = self.instructions_for(ctx.agent_spec.as_ref())?;
120        if instructions.trim().is_empty() {
121            return Ok(StateCommand::new());
122        }
123
124        let mut cmd = StateCommand::new();
125        cmd.schedule_action::<AddContextMessage>(ContextMessage::system(
126            A2UI_INSTRUCTION_CONTEXT_KEY,
127            instructions,
128        ))?;
129        Ok(cmd)
130    }
131}
132
133/// A2UI plugin that provides the render tool and prompt instructions.
134pub struct A2uiPlugin {
135    default_prompt_config: A2uiPromptConfig,
136    instructions: String,
137}
138
139impl A2uiPlugin {
140    /// Create with the default standard v0.8 catalog guidance.
141    pub fn new() -> Self {
142        Self::with_prompt_config(A2uiPromptConfig::default())
143    }
144
145    /// Create with a default prompt configuration that can still be overridden
146    /// by the active agent's `A2uiPromptConfigKey` section.
147    pub fn with_prompt_config(prompt_config: A2uiPromptConfig) -> Self {
148        let instructions = prompt_config.instructions_text();
149        Self {
150            default_prompt_config: prompt_config,
151            instructions,
152        }
153    }
154
155    /// Create with a specific catalog URI description for prompt guidance.
156    pub fn with_catalog_id(catalog_id: &str) -> Self {
157        Self::with_prompt_config(A2uiPromptConfig::default().with_catalog_id(catalog_id))
158    }
159
160    /// Create with a catalog ID and custom examples.
161    pub fn with_catalog_and_examples(catalog_id: &str, examples: &str) -> Self {
162        Self::with_prompt_config(
163            A2uiPromptConfig::default()
164                .with_catalog_id(catalog_id)
165                .with_examples(examples),
166        )
167    }
168
169    /// Create with fully custom instructions.
170    pub fn with_custom_instructions(instructions: String) -> Self {
171        Self::with_prompt_config(A2uiPromptConfig::default().with_instructions(instructions))
172    }
173
174    /// Returns the default prompt config resolved by the plugin when the
175    /// active agent does not provide overrides.
176    pub fn prompt_config(&self) -> &A2uiPromptConfig {
177        &self.default_prompt_config
178    }
179
180    /// Returns the default instructions that will be injected when there is no
181    /// per-agent override.
182    pub fn instructions(&self) -> &str {
183        &self.instructions
184    }
185}
186
187impl Default for A2uiPlugin {
188    fn default() -> Self {
189        Self::new()
190    }
191}
192
193impl Plugin for A2uiPlugin {
194    fn descriptor(&self) -> PluginDescriptor {
195        PluginDescriptor {
196            name: A2UI_PLUGIN_ID,
197        }
198    }
199
200    fn register(&self, registrar: &mut PluginRegistrar) -> Result<(), StateError> {
201        registrar.register_phase_hook(
202            A2UI_PLUGIN_ID,
203            Phase::BeforeInference,
204            A2uiInstructionHook::new(self.default_prompt_config.clone()),
205        )?;
206        registrar.register_tool(A2UI_TOOL_ID, Arc::new(A2uiRenderTool::new()))?;
207        Ok(())
208    }
209
210    fn config_schemas(&self) -> Vec<ConfigSchema> {
211        vec![ConfigSchema::for_key::<A2uiPromptConfigKey>()]
212    }
213
214    fn on_activate(
215        &self,
216        _agent_spec: &AgentSpec,
217        _patch: &mut MutationBatch,
218    ) -> Result<(), StateError> {
219        Ok(())
220    }
221}
222
223fn build_instructions(catalog_id: &str, examples: Option<&str>) -> String {
224    let mut s = String::from(A2UI_SCHEMA_INSTRUCTIONS);
225    s = s.replace("{{CATALOG_ID}}", catalog_id);
226    if let Some(examples) = examples {
227        s.push_str("\n\n---BEGIN A2UI EXAMPLES---\n");
228        s.push_str(examples);
229        s.push_str("\n---END A2UI EXAMPLES---");
230    }
231    s
232}
233
234const A2UI_SCHEMA_INSTRUCTIONS: &str = r#"
235## A2UI Declarative UI
236
237You have access to the `render_a2ui` tool. Pass ONE A2UI message key directly
238as the tool argument. Use the official A2UI v0.8 message format and do NOT add
239wrapper objects beyond the protocol itself.
240
241Call the tool once per message in this order: surfaceUpdate → dataModelUpdate → beginRendering.
242
243### 1. surfaceUpdate
244```json
245{"surfaceUpdate": {"surfaceId": "ops_request", "components": [
246  {"id": "root", "component": {"Card": {"child": "content"}}},
247  {"id": "content", "component": {"Column": {"children": {"explicitList": ["title", "requester", "notes", "priority_label", "priority", "submit_label", "submit_button"]}}}},
248  {"id": "title", "component": {"Text": {"usageHint": "h2", "text": {"literalString": "Operations Request"}}}},
249  {"id": "requester", "component": {"TextField": {"label": {"literalString": "Requester"}, "text": {"path": "/request/requester"}, "textFieldType": "shortText"}}},
250  {"id": "notes", "component": {"TextField": {"label": {"literalString": "Request details"}, "text": {"path": "/request/notes"}, "textFieldType": "longText"}}},
251  {"id": "priority_label", "component": {"Text": {"usageHint": "caption", "text": {"literalString": "Priority"}}}},
252  {"id": "priority", "component": {"MultipleChoice": {"selections": {"path": "/request/priority"}, "options": [{"label": {"literalString": "Standard"}, "value": "standard"}, {"label": {"literalString": "Urgent"}, "value": "urgent"}], "maxAllowedSelections": 1}}},
253  {"id": "submit_label", "component": {"Text": {"text": {"literalString": "Submit request"}}}},
254  {"id": "submit_button", "component": {"Button": {"child": "submit_label", "primary": true, "action": {"name": "ops_request.submit", "context": [{"key": "requester", "value": {"path": "/request/requester"}}, {"key": "notes", "value": {"path": "/request/notes"}}, {"key": "priority", "value": {"path": "/request/priority"}}]}}}}
255]}}
256```
257
258### 2. dataModelUpdate
259```json
260{"dataModelUpdate": {"surfaceId": "ops_request", "path": "/request", "contents": [
261  {"key": "requester", "valueString": ""},
262  {"key": "notes", "valueString": ""},
263  {"key": "priority", "valueMap": [{"key": "0", "valueString": "standard"}]}
264]}}
265```
266
267### 3. beginRendering
268```json
269{"beginRendering": {"surfaceId": "ops_request", "root": "root"}}
270```
271
272### Rules
273- The client defaults to the standard v0.8 catalog: `{{CATALOG_ID}}`.
274- Components are a flat list. Relationships are expressed by component IDs inside nested component props.
275- Each component must have `"id"` and `"component"`, where `"component"` is an object like `{"Text": {...}}`.
276- `Text` components must always nest copy under the `text` field, for example `{"Text":{"text":{"literalString":"Submit"}}}`.
277- For container children, use `{"children": {"explicitList": ["id1", "id2"]}}`.
278- For text input binding, v0.8 `TextField` uses the `text` property, not `value`.
279- For selection inputs, v0.8 `MultipleChoice` uses `options`, `selections`, and `maxAllowedSelections`.
280- Each `MultipleChoice.options` item must look like `{"label":{"literalString":"High"},"value":"high"}`.
281- `MultipleChoice` does not have a `label` field. If you need a visible label, render a nearby `Text` component.
282- `MultipleChoice.selections` binds to an array path. Initialize defaults with `valueMap`, for example `{"key":"priority","valueMap":[{"key":"0","valueString":"standard"}]}`.
283- For `DateTimeInput` defaults, use browser-compatible local strings like `2026-04-10T09:00`. Do not include a trailing `Z` or timezone offset.
284- For buttons, use `"action": {"name": "...", "context": [{"key":"field","value":{"path":"/request/field"}}]}`; do not use the v0.9 `event` wrapper.
285- `dataModelUpdate.contents` must be an array of `{key, valueString|valueNumber|valueBoolean|valueMap}` entries.
286- Call `beginRendering` only after the referenced root component already exists in `surfaceUpdate`.
287- Use `deleteSurface` when the workflow is complete or should be dismissed.
288- When updating an existing surface after a user interaction, resend the current values for any bound fields that remain visible on the surface. Do not send only the newly changed status fields if the next UI still displays older inputs.
289- If the conversation includes an `A2UI action:` payload, treat its `context` object as the authoritative source for the user's current bound field values unless the user explicitly changed them in the new turn.
290
291### Update reminder
292If a follow-up state still shows the original form fields, the next `dataModelUpdate` should carry both the retained field values and the new status fields, for example:
293```json
294{"dataModelUpdate":{"surfaceId":"surface","path":"/request","contents":[
295  {"key":"requester","valueString":"Jordan Patel"},
296  {"key":"department","valueString":"Operations"},
297  {"key":"status","valueString":"submitted"}
298]}}
299```
300"#;
awaken_ext_generative_ui/a2ui/plugin.rs

awaken_ext_generative_ui/a2ui/
plugin.rs
