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#[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 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, ¶ms.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(§ion_id, "heading");
153 components.push(text(&heading_id, §ion.heading, Some("h2")));
154 section_children.push(heading_id);
155
156 if let Some(body) = §ion.body {
157 let body_id = stable_child_id(§ion_id, "body");
158 components.push(text(&body_id, body, None));
159 section_children.push(body_id);
160 }
161
162 if let Some(image_url) = §ion.image_url {
163 let image_id = stable_child_id(§ion_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(§ion_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(§ion_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 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(§ion_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(÷r_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 Ok(Value::String(jsonl))
269 }
270 UiProtocol::AgUi => {
271 let thread_id = params
272 .protocol_options
273 .resolved_ag_ui_thread_id(¶ms.surface_id);
274 let run_id = params
275 .protocol_options
276 .resolved_ag_ui_run_id(¶ms.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}