Skip to main content

adk_ui/tools/
render_page.rs

1use crate::a2ui::{
2    A2uiSchemaVersion, A2uiValidator, column, divider, encode_jsonl, image, row, stable_child_id,
3    stable_id, stable_indexed_id, text,
4};
5use crate::catalog_registry::CatalogRegistry;
6use crate::interop::{AgUiAdapter, McpAppsAdapter, UiProtocol, UiProtocolAdapter, UiSurface};
7use crate::tools::SurfaceProtocolOptions;
8use adk_core::{Result, Tool, ToolContext};
9use async_trait::async_trait;
10use schemars::JsonSchema;
11use serde::{Deserialize, Serialize};
12use serde_json::{Value, json};
13use std::sync::Arc;
14
15fn default_surface_id() -> String {
16    "main".to_string()
17}
18
19fn default_send_data_model() -> bool {
20    true
21}
22
23fn default_validate() -> bool {
24    true
25}
26
27/// Page action button definition.
28#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
29pub struct PageAction {
30    /// Button label
31    pub label: String,
32    /// Action name emitted as A2UI action.event.name
33    pub action: String,
34    /// Button variant: "primary" or "borderless"
35    #[serde(default)]
36    pub variant: Option<String>,
37    /// Optional action context (supports data bindings)
38    #[serde(default)]
39    pub context: Option<Value>,
40}
41
42/// A section in a page.
43#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
44pub struct PageSection {
45    /// Section heading text
46    pub heading: String,
47    /// Optional body text
48    #[serde(default)]
49    pub body: Option<String>,
50    /// Optional bullet list items
51    #[serde(default)]
52    pub bullets: Vec<String>,
53    /// Optional image URL
54    #[serde(default)]
55    pub image_url: Option<String>,
56    /// Optional action buttons
57    #[serde(default)]
58    pub actions: Vec<PageAction>,
59}
60
61/// Parameters for the render_page tool.
62#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
63pub struct RenderPageParams {
64    /// Surface id (default: "main")
65    #[serde(default = "default_surface_id")]
66    pub surface_id: String,
67    /// Catalog id (defaults to the embedded ADK catalog)
68    #[serde(default)]
69    pub catalog_id: Option<String>,
70    /// Page title (rendered as h1)
71    pub title: String,
72    /// Optional description below the title
73    #[serde(default)]
74    pub description: Option<String>,
75    /// Sections to include
76    #[serde(default)]
77    pub sections: Vec<PageSection>,
78    /// Optional initial data model (sent via updateDataModel at path "/")
79    #[serde(default)]
80    pub data_model: Option<Value>,
81    /// Optional theme object for createSurface
82    #[serde(default)]
83    pub theme: Option<Value>,
84    /// If true, the client should include the data model in action metadata (default: true)
85    #[serde(default = "default_send_data_model")]
86    pub send_data_model: bool,
87    /// Validate generated messages against the A2UI v0.9 schema (default: true)
88    #[serde(default = "default_validate")]
89    pub validate: bool,
90    /// Shared protocol output options.
91    #[serde(flatten)]
92    pub protocol_options: SurfaceProtocolOptions,
93}
94
95/// Tool for emitting A2UI JSONL for a multi-section page.
96pub struct RenderPageTool;
97
98impl RenderPageTool {
99    pub fn new() -> Self {
100        Self
101    }
102}
103
104impl Default for RenderPageTool {
105    fn default() -> Self {
106        Self::new()
107    }
108}
109
110#[async_trait]
111impl Tool for RenderPageTool {
112    fn name(&self) -> &str {
113        "render_page"
114    }
115
116    fn description(&self) -> &str {
117        r#"Render a multi-section page as A2UI JSONL. Builds a root column with a title, optional description, and section blocks. Each section can include body text, bullets, images, and action buttons."#
118    }
119
120    fn parameters_schema(&self) -> Option<Value> {
121        Some(super::generate_gemini_schema::<RenderPageParams>())
122    }
123
124    async fn execute(&self, _ctx: Arc<dyn ToolContext>, args: Value) -> Result<Value> {
125        let params: RenderPageParams = serde_json::from_value(args.clone()).map_err(|e| {
126            adk_core::AdkError::Tool(format!("Invalid parameters: {}. Got: {}", e, args))
127        })?;
128
129        let registry = CatalogRegistry::new();
130        let catalog_id =
131            params.catalog_id.unwrap_or_else(|| registry.default_catalog_id().to_string());
132
133        let page_id = stable_id(&format!("page:{}:{}", params.surface_id, params.title));
134        let mut components: Vec<Value> = Vec::new();
135        let mut root_children: Vec<String> = Vec::new();
136
137        let title_id = stable_child_id(&page_id, "title");
138        components.push(text(&title_id, &params.title, Some("h1")));
139        root_children.push(title_id);
140
141        if let Some(description) = params.description {
142            let desc_id = stable_child_id(&page_id, "description");
143            components.push(text(&desc_id, &description, None));
144            root_children.push(desc_id);
145        }
146
147        for (index, section) in params.sections.iter().enumerate() {
148            let section_id = stable_indexed_id(&page_id, "section", index);
149            let mut section_children: Vec<String> = Vec::new();
150
151            let heading_id = stable_child_id(&section_id, "heading");
152            components.push(text(&heading_id, &section.heading, Some("h2")));
153            section_children.push(heading_id);
154
155            if let Some(body) = &section.body {
156                let body_id = stable_child_id(&section_id, "body");
157                components.push(text(&body_id, body, None));
158                section_children.push(body_id);
159            }
160
161            if let Some(image_url) = &section.image_url {
162                let image_id = stable_child_id(&section_id, "image");
163                components.push(image(&image_id, image_url));
164                section_children.push(image_id);
165            }
166
167            if !section.bullets.is_empty() {
168                let list_id = stable_child_id(&section_id, "bullets");
169                let mut bullet_ids = Vec::new();
170                for (idx, bullet) in section.bullets.iter().enumerate() {
171                    let bullet_id = stable_indexed_id(&list_id, "item", idx);
172                    components.push(text(&bullet_id, bullet, None));
173                    bullet_ids.push(bullet_id);
174                }
175                let bullet_ids_str: Vec<&str> = bullet_ids.iter().map(|s| s.as_str()).collect();
176                components.push(column(&list_id, bullet_ids_str));
177                section_children.push(list_id);
178            }
179
180            if !section.actions.is_empty() {
181                let actions_id = stable_child_id(&section_id, "actions");
182                let mut action_ids = Vec::new();
183                for (idx, action) in section.actions.iter().enumerate() {
184                    let button_id = stable_indexed_id(&actions_id, "button", idx);
185                    let label_id = stable_child_id(&button_id, "label");
186                    components.push(text(&label_id, &action.label, None));
187
188                    // Build button with action
189                    let mut button_comp = json!({
190                        "id": button_id,
191                        "component": "Button",
192                        "child": label_id,
193                        "action": {
194                            "event": {
195                                "name": action.action
196                            }
197                        }
198                    });
199
200                    if let Some(variant) = &action.variant {
201                        button_comp["variant"] = json!(variant);
202                    }
203                    if let Some(context) = &action.context {
204                        button_comp["action"]["event"]["context"] = context.clone();
205                    }
206
207                    components.push(button_comp);
208                    action_ids.push(button_id);
209                }
210                let action_ids_str: Vec<&str> = action_ids.iter().map(|s| s.as_str()).collect();
211                components.push(row(&actions_id, action_ids_str));
212                section_children.push(actions_id);
213            }
214
215            let section_children_str: Vec<&str> =
216                section_children.iter().map(|s| s.as_str()).collect();
217            components.push(column(&section_id, section_children_str));
218            root_children.push(section_id);
219
220            if index + 1 < params.sections.len() {
221                let divider_id = stable_indexed_id(&page_id, "divider", index);
222                components.push(divider(&divider_id, "horizontal"));
223                root_children.push(divider_id);
224            }
225        }
226
227        let root_children_str: Vec<&str> = root_children.iter().map(|s| s.as_str()).collect();
228        components.push(column("root", root_children_str));
229
230        let surface = UiSurface::new(params.surface_id.clone(), catalog_id, components)
231            .with_data_model(params.data_model.clone())
232            .with_theme(params.theme.clone())
233            .with_send_data_model(params.send_data_model);
234
235        match params.protocol_options.protocol {
236            UiProtocol::A2ui => {
237                let messages = surface.to_a2ui_messages();
238                if params.validate {
239                    let validator = A2uiValidator::new().map_err(|e| {
240                        adk_core::AdkError::Tool(format!(
241                            "Failed to initialize A2UI validator: {}",
242                            e
243                        ))
244                    })?;
245                    for message in &messages {
246                        if let Err(errors) =
247                            validator.validate_message(message, A2uiSchemaVersion::V0_9)
248                        {
249                            let details = errors
250                                .iter()
251                                .map(|err| format!("{} at {}", err.message, err.instance_path))
252                                .collect::<Vec<_>>()
253                                .join("; ");
254                            return Err(adk_core::AdkError::Tool(format!(
255                                "A2UI validation failed: {}",
256                                details
257                            )));
258                        }
259                    }
260                }
261
262                let jsonl = encode_jsonl(messages).map_err(|e| {
263                    adk_core::AdkError::Tool(format!("Failed to encode A2UI JSONL: {}", e))
264                })?;
265
266                // Keep historical return type for default protocol compatibility.
267                Ok(Value::String(jsonl))
268            }
269            UiProtocol::AgUi => {
270                let thread_id =
271                    params.protocol_options.resolved_ag_ui_thread_id(&params.surface_id);
272                let run_id = params.protocol_options.resolved_ag_ui_run_id(&params.surface_id);
273                let adapter = AgUiAdapter::new(thread_id, run_id);
274                adapter.to_protocol_payload(&surface)
275            }
276            UiProtocol::McpApps => {
277                let options = params.protocol_options.parse_mcp_options()?;
278                let adapter = McpAppsAdapter::new(options);
279                adapter.to_protocol_payload(&surface)
280            }
281        }
282    }
283}
284
285#[cfg(test)]
286mod tests {
287    use super::*;
288    use adk_core::{Content, EventActions, ReadonlyContext};
289    use async_trait::async_trait;
290    use std::sync::{Arc, Mutex};
291
292    struct TestContext {
293        content: Content,
294        actions: Mutex<EventActions>,
295    }
296
297    impl TestContext {
298        fn new() -> Self {
299            Self { content: Content::new("user"), actions: Mutex::new(EventActions::default()) }
300        }
301    }
302
303    #[async_trait]
304    impl ReadonlyContext for TestContext {
305        fn invocation_id(&self) -> &str {
306            "test"
307        }
308        fn agent_name(&self) -> &str {
309            "test"
310        }
311        fn user_id(&self) -> &str {
312            "user"
313        }
314        fn app_name(&self) -> &str {
315            "app"
316        }
317        fn session_id(&self) -> &str {
318            "session"
319        }
320        fn branch(&self) -> &str {
321            ""
322        }
323        fn user_content(&self) -> &Content {
324            &self.content
325        }
326    }
327
328    #[async_trait]
329    impl adk_core::CallbackContext for TestContext {
330        fn artifacts(&self) -> Option<Arc<dyn adk_core::Artifacts>> {
331            None
332        }
333    }
334
335    #[async_trait]
336    impl ToolContext for TestContext {
337        fn function_call_id(&self) -> &str {
338            "call-123"
339        }
340        fn actions(&self) -> EventActions {
341            self.actions.lock().unwrap().clone()
342        }
343        fn set_actions(&self, actions: EventActions) {
344            *self.actions.lock().unwrap() = actions;
345        }
346        async fn search_memory(&self, _query: &str) -> Result<Vec<adk_core::MemoryEntry>> {
347            Ok(vec![])
348        }
349    }
350
351    #[tokio::test]
352    async fn render_page_emits_jsonl() {
353        let tool = RenderPageTool::new();
354        let args = serde_json::json!({
355            "title": "Launch",
356            "sections": [
357                {
358                    "heading": "Features",
359                    "body": "Fast and secure.",
360                    "bullets": ["One", "Two"],
361                    "actions": [
362                        { "label": "Get Started", "action": "start", "variant": "primary" }
363                    ]
364                }
365            ]
366        });
367
368        let ctx: Arc<dyn ToolContext> = Arc::new(TestContext::new());
369        let value = tool.execute(ctx, args).await.unwrap();
370        let jsonl = value.as_str().unwrap();
371        let lines: Vec<Value> =
372            jsonl.trim_end().lines().map(|line| serde_json::from_str(line).unwrap()).collect();
373
374        assert_eq!(lines.len(), 2);
375        assert!(lines[0].get("createSurface").is_some());
376        assert!(lines[1].get("updateComponents").is_some());
377    }
378
379    #[tokio::test]
380    async fn render_page_emits_ag_ui_events() {
381        let tool = RenderPageTool::new();
382        let args = serde_json::json!({
383            "protocol": "ag_ui",
384            "title": "Launch",
385            "sections": [{ "heading": "Features" }]
386        });
387
388        let ctx: Arc<dyn ToolContext> = Arc::new(TestContext::new());
389        let value = tool.execute(ctx, args).await.unwrap();
390        assert_eq!(value["protocol"], "ag_ui");
391        let events = value["events"].as_array().unwrap();
392        assert_eq!(events[1]["type"], "CUSTOM");
393    }
394
395    #[tokio::test]
396    async fn render_page_emits_mcp_apps_payload() {
397        let tool = RenderPageTool::new();
398        let args = serde_json::json!({
399            "protocol": "mcp_apps",
400            "title": "Launch",
401            "sections": [{ "heading": "Features" }],
402            "mcp_apps": {
403                "resource_uri": "ui://tests/page"
404            }
405        });
406
407        let ctx: Arc<dyn ToolContext> = Arc::new(TestContext::new());
408        let value = tool.execute(ctx, args).await.unwrap();
409        assert_eq!(value["protocol"], "mcp_apps");
410        assert_eq!(value["payload"]["toolMeta"]["_meta"]["ui"]["resourceUri"], "ui://tests/page");
411    }
412}