adk_ui/tools/
render_screen.rs

1use crate::a2ui::{A2uiSchemaVersion, A2uiValidator};
2use crate::catalog_registry::CatalogRegistry;
3use crate::compat::{Result, Tool, ToolContext};
4use crate::interop::{
5    A2uiAdapter, AgUiAdapter, McpAppsAdapter, UiProtocol, UiProtocolAdapter, UiSurface,
6};
7use crate::tools::SurfaceProtocolOptions;
8use async_trait::async_trait;
9use schemars::JsonSchema;
10use serde::{Deserialize, Serialize};
11use serde_json::Value;
12use std::sync::Arc;
13
14fn default_surface_id() -> String {
15    "main".to_string()
16}
17
18fn default_send_data_model() -> bool {
19    true
20}
21
22fn default_validate() -> bool {
23    true
24}
25
26/// Parameters for the render_screen tool.
27#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
28pub struct RenderScreenParams {
29    /// Surface id (default: "main")
30    #[serde(default = "default_surface_id")]
31    pub surface_id: String,
32    /// Catalog id (defaults to the embedded ADK catalog)
33    #[serde(default)]
34    pub catalog_id: Option<String>,
35    /// A2UI component definitions (must include a component with id "root")
36    pub components: Vec<Value>,
37    /// Optional initial data model (sent via updateDataModel at path "/")
38    #[serde(default)]
39    pub data_model: Option<Value>,
40    /// Optional theme object for createSurface
41    #[serde(default)]
42    pub theme: Option<Value>,
43    /// If true, the client should include the data model in action metadata (default: true)
44    #[serde(default = "default_send_data_model")]
45    pub send_data_model: bool,
46    /// Validate generated messages against the A2UI v0.9 schema (default: true)
47    #[serde(default = "default_validate")]
48    pub validate: bool,
49    /// Shared protocol output options.
50    #[serde(flatten)]
51    pub protocol_options: SurfaceProtocolOptions,
52}
53
54/// Tool for emitting A2UI JSONL for a single screen (surface).
55///
56/// This tool wraps a list of A2UI components with the standard envelope messages:
57/// - createSurface
58/// - updateDataModel (optional)
59/// - updateComponents
60pub struct RenderScreenTool;
61
62impl RenderScreenTool {
63    pub fn new() -> Self {
64        Self
65    }
66}
67
68impl Default for RenderScreenTool {
69    fn default() -> Self {
70        Self::new()
71    }
72}
73
74#[async_trait]
75impl Tool for RenderScreenTool {
76    fn name(&self) -> &str {
77        "render_screen"
78    }
79
80    fn description(&self) -> &str {
81        r#"Emit A2UI JSONL for a single screen (surface). Input must include A2UI component objects with ids, including a root component with id "root".
82Returns a JSONL string with createSurface/updateDataModel/updateComponents messages."#
83    }
84
85    fn parameters_schema(&self) -> Option<Value> {
86        Some(super::generate_gemini_schema::<RenderScreenParams>())
87    }
88
89    async fn execute(&self, ctx: Arc<dyn ToolContext>, args: Value) -> Result<Value> {
90        let params: RenderScreenParams = serde_json::from_value(args.clone()).map_err(|e| {
91            crate::compat::AdkError::tool(format!("Invalid parameters: {}. Got: {}", e, args))
92        })?;
93
94        if params.components.is_empty() {
95            return Err(crate::compat::AdkError::tool(
96                "Invalid parameters: components must not be empty.".to_string(),
97            ));
98        }
99
100        let has_root = params.components.iter().any(|component| {
101            component
102                .get("id")
103                .and_then(Value::as_str)
104                .map(|id| id == "root")
105                .unwrap_or(false)
106        });
107
108        if !has_root {
109            return Err(crate::compat::AdkError::tool(
110                "Invalid parameters: components must include a root component with id \"root\"."
111                    .to_string(),
112            ));
113        }
114
115        let registry = CatalogRegistry::new();
116        let catalog_id = params
117            .catalog_id
118            .unwrap_or_else(|| registry.default_catalog_id().to_string());
119
120        let surface = UiSurface::new(
121            params.surface_id.clone(),
122            catalog_id,
123            params.components.clone(),
124        )
125        .with_data_model(params.data_model.clone())
126        .with_theme(params.theme.clone())
127        .with_send_data_model(params.send_data_model);
128
129        let output = match params.protocol_options.protocol {
130            UiProtocol::A2ui => {
131                let messages = surface.to_a2ui_messages();
132                if params.validate {
133                    let validator = A2uiValidator::new().map_err(|e| {
134                        crate::compat::AdkError::tool(format!(
135                            "Failed to initialize A2UI validator: {}",
136                            e
137                        ))
138                    })?;
139                    for message in &messages {
140                        if let Err(errors) =
141                            validator.validate_message(message, A2uiSchemaVersion::V0_9)
142                        {
143                            let details = errors
144                                .iter()
145                                .map(|err| format!("{} at {}", err.message, err.instance_path))
146                                .collect::<Vec<_>>()
147                                .join("; ");
148                            return Err(crate::compat::AdkError::tool(format!(
149                                "A2UI validation failed: {}",
150                                details
151                            )));
152                        }
153                    }
154                }
155
156                let adapter = A2uiAdapter;
157                let payload = adapter.to_protocol_payload(&surface)?;
158                adapter.validate(&payload)?;
159                Ok(payload)
160            }
161            UiProtocol::AgUi => {
162                let thread_id = params
163                    .protocol_options
164                    .resolved_ag_ui_thread_id(&params.surface_id);
165                let run_id = params
166                    .protocol_options
167                    .resolved_ag_ui_run_id(&params.surface_id);
168                let adapter = AgUiAdapter::new(thread_id, run_id);
169                adapter.to_protocol_payload(&surface)
170            }
171            UiProtocol::McpApps => {
172                let options = params.protocol_options.parse_mcp_options()?;
173                let adapter = McpAppsAdapter::new(options);
174                adapter.to_protocol_payload(&surface)
175            }
176            #[cfg(feature = "awp")]
177            UiProtocol::Awp => {
178                let adapter = crate::interop::AwpAdapter::new();
179                adapter.to_protocol_payload(&surface)
180            }
181        }?;
182        let surface_ref = crate::surface_runtime::next_surface_ref(&ctx, params.surface_id);
183        crate::surface_runtime::record_surface_ref(&ctx, &surface_ref);
184        Ok(output)
185    }
186}
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191    use crate::compat::{Content, EventActions, ReadonlyContext};
192    use async_trait::async_trait;
193    use std::sync::{Arc, Mutex};
194
195    struct TestContext {
196        content: Content,
197        actions: Mutex<EventActions>,
198    }
199
200    impl TestContext {
201        fn new() -> Self {
202            Self {
203                content: Content::new("user"),
204                actions: Mutex::new(EventActions::default()),
205            }
206        }
207    }
208
209    #[async_trait]
210    impl ReadonlyContext for TestContext {
211        fn invocation_id(&self) -> &str {
212            "test"
213        }
214        fn agent_name(&self) -> &str {
215            "test"
216        }
217        fn user_id(&self) -> &str {
218            "user"
219        }
220        fn app_name(&self) -> &str {
221            "app"
222        }
223        fn session_id(&self) -> &str {
224            "session"
225        }
226        fn branch(&self) -> &str {
227            ""
228        }
229        fn user_content(&self) -> &Content {
230            &self.content
231        }
232    }
233
234    #[async_trait]
235    impl crate::compat::CallbackContext for TestContext {
236        fn artifacts(&self) -> Option<Arc<dyn crate::compat::Artifacts>> {
237            None
238        }
239    }
240
241    #[async_trait]
242    impl ToolContext for TestContext {
243        fn function_call_id(&self) -> &str {
244            "call-123"
245        }
246        fn actions(&self) -> EventActions {
247            self.actions.lock().unwrap().clone()
248        }
249        fn set_actions(&self, actions: EventActions) {
250            *self.actions.lock().unwrap() = actions;
251        }
252        async fn search_memory(&self, _query: &str) -> Result<Vec<crate::compat::MemoryEntry>> {
253            Ok(vec![])
254        }
255    }
256
257    #[tokio::test]
258    async fn render_screen_emits_jsonl() {
259        use crate::a2ui::{column, text};
260
261        let tool = RenderScreenTool::new();
262        let args = serde_json::json!({
263            "components": [
264                text("title", "Hello World", Some("h1")),
265                text("desc", "Welcome", None),
266                column("root", vec!["title", "desc"])
267            ],
268            "data_model": { "title": "Hello" }
269        });
270
271        let ctx: Arc<dyn ToolContext> = Arc::new(TestContext::new());
272        let value = tool.execute(ctx, args).await.unwrap();
273
274        // The tool now returns a JSON object with components, data_model, and jsonl
275        assert!(value.is_object());
276        assert!(value.get("surface_id").is_some());
277        assert!(value.get("components").is_some());
278        assert!(value.get("jsonl").is_some());
279
280        // Verify JSONL is still generated
281        let jsonl = value["jsonl"].as_str().unwrap();
282        let lines: Vec<Value> = jsonl
283            .trim_end()
284            .lines()
285            .map(|line| serde_json::from_str(line).unwrap())
286            .collect();
287
288        assert_eq!(lines.len(), 3);
289        assert!(lines[0].get("createSurface").is_some());
290        assert!(lines[1].get("updateDataModel").is_some());
291        assert!(lines[2].get("updateComponents").is_some());
292
293        // Verify component structure in the returned JSON
294        let components = value["components"].as_array().unwrap();
295        assert_eq!(components.len(), 3);
296        let root = &components[2];
297        assert_eq!(root["id"], "root");
298        assert_eq!(root["component"], "Column");
299    }
300
301    #[tokio::test]
302    async fn render_screen_emits_ag_ui_events() {
303        use crate::a2ui::{column, text};
304
305        let tool = RenderScreenTool::new();
306        let args = serde_json::json!({
307            "protocol": "ag_ui",
308            "components": [
309                text("title", "Hello World", Some("h1")),
310                column("root", vec!["title"])
311            ]
312        });
313
314        let ctx: Arc<dyn ToolContext> = Arc::new(TestContext::new());
315        let value = tool.execute(ctx, args).await.unwrap();
316
317        assert_eq!(value["protocol"], "ag_ui");
318        let events = value["events"].as_array().unwrap();
319        assert_eq!(events.len(), 4);
320        assert_eq!(events[0]["type"], "RUN_STARTED");
321        assert_eq!(events[1]["type"], "ACTIVITY_SNAPSHOT");
322        assert_eq!(events[2]["type"], "CUSTOM");
323        assert_eq!(events[3]["type"], "RUN_FINISHED");
324    }
325
326    #[tokio::test]
327    async fn render_screen_emits_mcp_apps_payload() {
328        use crate::a2ui::{column, text};
329
330        let tool = RenderScreenTool::new();
331        let args = serde_json::json!({
332            "protocol": "mcp_apps",
333            "components": [
334                text("title", "Hello World", Some("h1")),
335                column("root", vec!["title"])
336            ],
337            "mcp_apps": {
338                "resource_uri": "ui://tests/screen"
339            }
340        });
341
342        let ctx: Arc<dyn ToolContext> = Arc::new(TestContext::new());
343        let value = tool.execute(ctx, args).await.unwrap();
344
345        assert_eq!(value["protocol"], "mcp_apps");
346        assert_eq!(value["payload"]["resource"]["uri"], "ui://tests/screen");
347        assert_eq!(
348            value["payload"]["toolMeta"]["_meta"]["ui"]["resourceUri"],
349            "ui://tests/screen"
350        );
351    }
352}
adk_ui/tools/render_screen.rs

adk_ui/tools/
render_screen.rs
