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::compat::{Result, Tool, ToolContext};
7use crate::interop::{AgUiAdapter, McpAppsAdapter, UiProtocol, UiProtocolAdapter, UiSurface};
8use crate::tools::SurfaceProtocolOptions;
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            crate::compat::AdkError::tool(format!("Invalid parameters: {}. Got: {}", e, args))
127        })?;
128
129        let registry = CatalogRegistry::new();
130        let catalog_id = params
131            .catalog_id
132            .unwrap_or_else(|| registry.default_catalog_id().to_string());
133
134        let page_id = stable_id(&format!("page:{}:{}", params.surface_id, params.title));
135        let mut components: Vec<Value> = Vec::new();
136        let mut root_children: Vec<String> = Vec::new();
137
138        let title_id = stable_child_id(&page_id, "title");
139        components.push(text(&title_id, &params.title, Some("h1")));
140        root_children.push(title_id);
141
142        if let Some(description) = params.description {
143            let desc_id = stable_child_id(&page_id, "description");
144            components.push(text(&desc_id, &description, None));
145            root_children.push(desc_id);
146        }
147
148        for (index, section) in params.sections.iter().enumerate() {
149            let section_id = stable_indexed_id(&page_id, "section", index);
150            let mut section_children: Vec<String> = Vec::new();
151
152            let heading_id = stable_child_id(&section_id, "heading");
153            components.push(text(&heading_id, &section.heading, Some("h2")));
154            section_children.push(heading_id);
155
156            if let Some(body) = &section.body {
157                let body_id = stable_child_id(&section_id, "body");
158                components.push(text(&body_id, body, None));
159                section_children.push(body_id);
160            }
161
162            if let Some(image_url) = &section.image_url {
163                let image_id = stable_child_id(&section_id, "image");
164                components.push(image(&image_id, image_url));
165                section_children.push(image_id);
166            }
167
168            if !section.bullets.is_empty() {
169                let list_id = stable_child_id(&section_id, "bullets");
170                let mut bullet_ids = Vec::new();
171                for (idx, bullet) in section.bullets.iter().enumerate() {
172                    let bullet_id = stable_indexed_id(&list_id, "item", idx);
173                    components.push(text(&bullet_id, bullet, None));
174                    bullet_ids.push(bullet_id);
175                }
176                let bullet_ids_str: Vec<&str> = bullet_ids.iter().map(|s| s.as_str()).collect();
177                components.push(column(&list_id, bullet_ids_str));
178                section_children.push(list_id);
179            }
180
181            if !section.actions.is_empty() {
182                let actions_id = stable_child_id(&section_id, "actions");
183                let mut action_ids = Vec::new();
184                for (idx, action) in section.actions.iter().enumerate() {
185                    let button_id = stable_indexed_id(&actions_id, "button", idx);
186                    let label_id = stable_child_id(&button_id, "label");
187                    components.push(text(&label_id, &action.label, None));
188
189                    // Build button with action
190                    let mut button_comp = json!({
191                        "id": button_id,
192                        "component": "Button",
193                        "child": label_id,
194                        "action": {
195                            "event": {
196                                "name": action.action
197                            }
198                        }
199                    });
200
201                    if let Some(variant) = &action.variant {
202                        button_comp["variant"] = json!(variant);
203                    }
204                    if let Some(context) = &action.context {
205                        button_comp["action"]["event"]["context"] = context.clone();
206                    }
207
208                    components.push(button_comp);
209                    action_ids.push(button_id);
210                }
211                let action_ids_str: Vec<&str> = action_ids.iter().map(|s| s.as_str()).collect();
212                components.push(row(&actions_id, action_ids_str));
213                section_children.push(actions_id);
214            }
215
216            let section_children_str: Vec<&str> =
217                section_children.iter().map(|s| s.as_str()).collect();
218            components.push(column(&section_id, section_children_str));
219            root_children.push(section_id);
220
221            if index + 1 < params.sections.len() {
222                let divider_id = stable_indexed_id(&page_id, "divider", index);
223                components.push(divider(&divider_id, "horizontal"));
224                root_children.push(divider_id);
225            }
226        }
227
228        let root_children_str: Vec<&str> = root_children.iter().map(|s| s.as_str()).collect();
229        components.push(column("root", root_children_str));
230
231        let surface = UiSurface::new(params.surface_id.clone(), catalog_id, components)
232            .with_data_model(params.data_model.clone())
233            .with_theme(params.theme.clone())
234            .with_send_data_model(params.send_data_model);
235
236        let output = match params.protocol_options.protocol {
237            UiProtocol::A2ui => {
238                let messages = surface.to_a2ui_messages();
239                if params.validate {
240                    let validator = A2uiValidator::new().map_err(|e| {
241                        crate::compat::AdkError::tool(format!(
242                            "Failed to initialize A2UI validator: {}",
243                            e
244                        ))
245                    })?;
246                    for message in &messages {
247                        if let Err(errors) =
248                            validator.validate_message(message, A2uiSchemaVersion::V0_9)
249                        {
250                            let details = errors
251                                .iter()
252                                .map(|err| format!("{} at {}", err.message, err.instance_path))
253                                .collect::<Vec<_>>()
254                                .join("; ");
255                            return Err(crate::compat::AdkError::tool(format!(
256                                "A2UI validation failed: {}",
257                                details
258                            )));
259                        }
260                    }
261                }
262
263                let jsonl = encode_jsonl(messages).map_err(|e| {
264                    crate::compat::AdkError::tool(format!("Failed to encode A2UI JSONL: {}", e))
265                })?;
266
267                // Keep historical return type for default protocol compatibility.
268                Ok(Value::String(jsonl))
269            }
270            UiProtocol::AgUi => {
271                let thread_id = params
272                    .protocol_options
273                    .resolved_ag_ui_thread_id(&params.surface_id);
274                let run_id = params
275                    .protocol_options
276                    .resolved_ag_ui_run_id(&params.surface_id);
277                let adapter = AgUiAdapter::new(thread_id, run_id);
278                adapter.to_protocol_payload(&surface)
279            }
280            UiProtocol::McpApps => {
281                let options = params.protocol_options.parse_mcp_options()?;
282                let adapter = McpAppsAdapter::new(options);
283                adapter.to_protocol_payload(&surface)
284            }
285            #[cfg(feature = "awp")]
286            UiProtocol::Awp => {
287                let adapter = crate::interop::AwpAdapter::new();
288                adapter.to_protocol_payload(&surface)
289            }
290        }?;
291        let surface_ref = crate::surface_runtime::next_surface_ref(&ctx, params.surface_id);
292        crate::surface_runtime::record_surface_ref(&ctx, &surface_ref);
293        Ok(output)
294    }
295}
296
297#[cfg(test)]
298mod tests {
299    use super::*;
300    use crate::compat::{Content, EventActions, ReadonlyContext};
301    use async_trait::async_trait;
302    use std::sync::{Arc, Mutex};
303
304    struct TestContext {
305        content: Content,
306        actions: Mutex<EventActions>,
307    }
308
309    impl TestContext {
310        fn new() -> Self {
311            Self {
312                content: Content::new("user"),
313                actions: Mutex::new(EventActions::default()),
314            }
315        }
316    }
317
318    #[async_trait]
319    impl ReadonlyContext for TestContext {
320        fn invocation_id(&self) -> &str {
321            "test"
322        }
323        fn agent_name(&self) -> &str {
324            "test"
325        }
326        fn user_id(&self) -> &str {
327            "user"
328        }
329        fn app_name(&self) -> &str {
330            "app"
331        }
332        fn session_id(&self) -> &str {
333            "session"
334        }
335        fn branch(&self) -> &str {
336            ""
337        }
338        fn user_content(&self) -> &Content {
339            &self.content
340        }
341    }
342
343    #[async_trait]
344    impl crate::compat::CallbackContext for TestContext {
345        fn artifacts(&self) -> Option<Arc<dyn crate::compat::Artifacts>> {
346            None
347        }
348    }
349
350    #[async_trait]
351    impl ToolContext for TestContext {
352        fn function_call_id(&self) -> &str {
353            "call-123"
354        }
355        fn actions(&self) -> EventActions {
356            self.actions.lock().unwrap().clone()
357        }
358        fn set_actions(&self, actions: EventActions) {
359            *self.actions.lock().unwrap() = actions;
360        }
361        async fn search_memory(&self, _query: &str) -> Result<Vec<crate::compat::MemoryEntry>> {
362            Ok(vec![])
363        }
364    }
365
366    #[tokio::test]
367    async fn render_page_emits_jsonl() {
368        let tool = RenderPageTool::new();
369        let args = serde_json::json!({
370            "title": "Launch",
371            "sections": [
372                {
373                    "heading": "Features",
374                    "body": "Fast and secure.",
375                    "bullets": ["One", "Two"],
376                    "actions": [
377                        { "label": "Get Started", "action": "start", "variant": "primary" }
378                    ]
379                }
380            ]
381        });
382
383        let ctx: Arc<dyn ToolContext> = Arc::new(TestContext::new());
384        let value = tool.execute(ctx, args).await.unwrap();
385        let jsonl = value.as_str().unwrap();
386        let lines: Vec<Value> = jsonl
387            .trim_end()
388            .lines()
389            .map(|line| serde_json::from_str(line).unwrap())
390            .collect();
391
392        assert_eq!(lines.len(), 2);
393        assert!(lines[0].get("createSurface").is_some());
394        assert!(lines[1].get("updateComponents").is_some());
395    }
396
397    #[tokio::test]
398    async fn render_page_emits_ag_ui_events() {
399        let tool = RenderPageTool::new();
400        let args = serde_json::json!({
401            "protocol": "ag_ui",
402            "title": "Launch",
403            "sections": [{ "heading": "Features" }]
404        });
405
406        let ctx: Arc<dyn ToolContext> = Arc::new(TestContext::new());
407        let value = tool.execute(ctx, args).await.unwrap();
408        assert_eq!(value["protocol"], "ag_ui");
409        let events = value["events"].as_array().unwrap();
410        assert_eq!(events[0]["type"], "RUN_STARTED");
411        assert_eq!(events[1]["type"], "ACTIVITY_SNAPSHOT");
412        assert_eq!(events[2]["type"], "CUSTOM");
413        assert_eq!(events[3]["type"], "RUN_FINISHED");
414    }
415
416    #[tokio::test]
417    async fn render_page_emits_mcp_apps_payload() {
418        let tool = RenderPageTool::new();
419        let args = serde_json::json!({
420            "protocol": "mcp_apps",
421            "title": "Launch",
422            "sections": [{ "heading": "Features" }],
423            "mcp_apps": {
424                "resource_uri": "ui://tests/page"
425            }
426        });
427
428        let ctx: Arc<dyn ToolContext> = Arc::new(TestContext::new());
429        let value = tool.execute(ctx, args).await.unwrap();
430        assert_eq!(value["protocol"], "mcp_apps");
431        assert_eq!(
432            value["payload"]["toolMeta"]["_meta"]["ui"]["resourceUri"],
433            "ui://tests/page"
434        );
435    }
436}