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#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
29pub struct PageAction {
30 pub label: String,
32 pub action: String,
34 #[serde(default)]
36 pub variant: Option<String>,
37 #[serde(default)]
39 pub context: Option<Value>,
40}
41
42#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
44pub struct PageSection {
45 pub heading: String,
47 #[serde(default)]
49 pub body: Option<String>,
50 #[serde(default)]
52 pub bullets: Vec<String>,
53 #[serde(default)]
55 pub image_url: Option<String>,
56 #[serde(default)]
58 pub actions: Vec<PageAction>,
59}
60
61#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
63pub struct RenderPageParams {
64 #[serde(default = "default_surface_id")]
66 pub surface_id: String,
67 #[serde(default)]
69 pub catalog_id: Option<String>,
70 pub title: String,
72 #[serde(default)]
74 pub description: Option<String>,
75 #[serde(default)]
77 pub sections: Vec<PageSection>,
78 #[serde(default)]
80 pub data_model: Option<Value>,
81 #[serde(default)]
83 pub theme: Option<Value>,
84 #[serde(default = "default_send_data_model")]
86 pub send_data_model: bool,
87 #[serde(default = "default_validate")]
89 pub validate: bool,
90 #[serde(flatten)]
92 pub protocol_options: SurfaceProtocolOptions,
93}
94
95pub 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, ¶ms.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(§ion_id, "heading");
152 components.push(text(&heading_id, §ion.heading, Some("h2")));
153 section_children.push(heading_id);
154
155 if let Some(body) = §ion.body {
156 let body_id = stable_child_id(§ion_id, "body");
157 components.push(text(&body_id, body, None));
158 section_children.push(body_id);
159 }
160
161 if let Some(image_url) = §ion.image_url {
162 let image_id = stable_child_id(§ion_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(§ion_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(§ion_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 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(§ion_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(÷r_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 Ok(Value::String(jsonl))
268 }
269 UiProtocol::AgUi => {
270 let thread_id =
271 params.protocol_options.resolved_ag_ui_thread_id(¶ms.surface_id);
272 let run_id = params.protocol_options.resolved_ag_ui_run_id(¶ms.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}