1use crate::compiler::CompiledConfig;
14use serde::{Deserialize, Serialize};
15use std::collections::BTreeMap;
16
17pub const A2UI_SCHEMA_VERSION: &str = "cellstate.a2ui.schema.v1";
19
20#[derive(Debug, Clone, Serialize, Deserialize)]
24#[serde(rename_all = "camelCase")]
25pub struct A2uiResponse {
26 pub schema_version: String,
28 pub version: String,
30 pub components: Vec<A2uiComponent>,
32 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
34 pub data_model: BTreeMap<String, A2uiDataField>,
35}
36
37#[derive(Debug, Clone, Serialize, Deserialize)]
39#[serde(rename_all = "camelCase")]
40pub struct A2uiComponent {
41 pub id: String,
43 #[serde(rename = "type")]
45 pub component_type: String,
46 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
48 pub properties: BTreeMap<String, serde_json::Value>,
49 #[serde(default, skip_serializing_if = "Vec::is_empty")]
51 pub children: Vec<String>,
52}
53
54#[derive(Debug, Clone, Serialize, Deserialize)]
56#[serde(rename_all = "camelCase")]
57pub struct A2uiDataField {
58 #[serde(rename = "type")]
60 pub field_type: String,
61 #[serde(skip_serializing_if = "Option::is_none")]
63 pub label: Option<String>,
64 #[serde(default)]
66 pub required: bool,
67}
68
69pub fn cellstate_to_a2ui_type(cellstate_type: &str) -> &str {
71 match cellstate_type {
72 "cellstate.memory-stream" => "data-stream",
73 "cellstate.pack-editor" => "code-editor",
74 "cellstate.agent-chat" => "chat",
75 "cellstate.browser-preview" => "web-view",
76 "cellstate.event-forensics" => "data-table",
77 "cellstate.policy-graph" => "graph",
78 "cellstate.terminal" => "terminal",
79 "cellstate.deploy-bar" => "action-bar",
80 _ => "container",
81 }
82}
83
84pub fn generate_a2ui_schema(config: &CompiledConfig) -> A2uiResponse {
88 let mut components = Vec::new();
89 let mut counter = 0u32;
90
91 components.push(A2uiComponent {
93 id: format!("comp_{}", counter),
94 component_type: "data-stream".to_string(),
95 properties: BTreeMap::from([
96 (
97 "title".to_string(),
98 serde_json::Value::String("Memory Stream".to_string()),
99 ),
100 (
101 "source".to_string(),
102 serde_json::Value::String("cellstate.memory-stream".to_string()),
103 ),
104 ]),
105 children: vec![],
106 });
107 counter += 1;
108
109 components.push(A2uiComponent {
111 id: format!("comp_{}", counter),
112 component_type: "data-table".to_string(),
113 properties: BTreeMap::from([(
114 "title".to_string(),
115 serde_json::Value::String("Event Forensics".to_string()),
116 )]),
117 children: vec![],
118 });
119 counter += 1;
120
121 if !config.tools.is_empty() {
125 components.push(A2uiComponent {
126 id: format!("comp_{}", counter),
127 component_type: "chat".to_string(),
128 properties: BTreeMap::from([(
129 "title".to_string(),
130 serde_json::Value::String("Agent Chat".to_string()),
131 )]),
132 children: vec![],
133 });
134 counter += 1;
135 }
136
137 let has_bash = config
139 .tools
140 .iter()
141 .any(|t| t.kind == crate::compiler::CompiledToolKind::Bash);
142 if has_bash {
143 components.push(A2uiComponent {
144 id: format!("comp_{}", counter),
145 component_type: "terminal".to_string(),
146 properties: BTreeMap::from([(
147 "title".to_string(),
148 serde_json::Value::String("Terminal".to_string()),
149 )]),
150 children: vec![],
151 });
152 counter += 1;
153 }
154
155 let has_browser = config
157 .tools
158 .iter()
159 .any(|t| t.kind == crate::compiler::CompiledToolKind::Browser);
160 if has_browser {
161 components.push(A2uiComponent {
162 id: format!("comp_{}", counter),
163 component_type: "web-view".to_string(),
164 properties: BTreeMap::from([(
165 "title".to_string(),
166 serde_json::Value::String("Browser Preview".to_string()),
167 )]),
168 children: vec![],
169 });
170 counter += 1;
171 }
172
173 if !config.pack_injections.is_empty() {
175 components.push(A2uiComponent {
176 id: format!("comp_{}", counter),
177 component_type: "graph".to_string(),
178 properties: BTreeMap::from([(
179 "title".to_string(),
180 serde_json::Value::String("Policy Graph".to_string()),
181 )]),
182 children: vec![],
183 });
184 let _ = counter; }
186
187 let mut data_model = BTreeMap::new();
189 data_model.insert(
190 "trajectory_id".to_string(),
191 A2uiDataField {
192 field_type: "string".to_string(),
193 label: Some("Trajectory ID".to_string()),
194 required: false,
195 },
196 );
197 data_model.insert(
198 "scope_id".to_string(),
199 A2uiDataField {
200 field_type: "string".to_string(),
201 label: Some("Scope ID".to_string()),
202 required: false,
203 },
204 );
205 data_model.insert(
206 "agent_id".to_string(),
207 A2uiDataField {
208 field_type: "string".to_string(),
209 label: Some("Agent ID".to_string()),
210 required: false,
211 },
212 );
213
214 A2uiResponse {
215 schema_version: A2UI_SCHEMA_VERSION.to_string(),
216 version: "0.8".to_string(),
217 components,
218 data_model,
219 }
220}
221
222pub fn to_a2ui_json(response: &A2uiResponse) -> Result<String, serde_json::Error> {
224 serde_json::to_string_pretty(response)
225}
226
227#[cfg(test)]
228mod tests {
229 use super::*;
230
231 #[test]
232 fn test_generate_a2ui_schema_defaults() {
233 let config = CompiledConfig::default();
234 let schema = generate_a2ui_schema(&config);
235
236 assert_eq!(schema.schema_version, A2UI_SCHEMA_VERSION);
237 assert_eq!(schema.version, "0.8");
238 assert!(schema.components.len() >= 2);
240 assert_eq!(schema.components[0].component_type, "data-stream");
241 assert_eq!(schema.components[1].component_type, "data-table");
242 }
243
244 #[test]
245 fn test_cellstate_to_a2ui_type_mapping() {
246 assert_eq!(
247 cellstate_to_a2ui_type("cellstate.memory-stream"),
248 "data-stream"
249 );
250 assert_eq!(cellstate_to_a2ui_type("cellstate.agent-chat"), "chat");
251 assert_eq!(cellstate_to_a2ui_type("cellstate.terminal"), "terminal");
252 assert_eq!(
253 cellstate_to_a2ui_type("cellstate.browser-preview"),
254 "web-view"
255 );
256 assert_eq!(cellstate_to_a2ui_type("unknown"), "container");
257 }
258
259 #[test]
260 fn test_a2ui_json_serialization() {
261 let config = CompiledConfig::default();
262 let schema = generate_a2ui_schema(&config);
263 let json = to_a2ui_json(&schema).expect("should serialize");
264 assert!(json.contains("\"schemaVersion\""));
265 assert!(json.contains("\"version\""));
266 assert!(json.contains("\"components\""));
267 assert!(json.contains("\"dataModel\""));
268 }
269
270 #[test]
271 fn test_a2ui_data_model_includes_core_fields() {
272 let config = CompiledConfig::default();
273 let schema = generate_a2ui_schema(&config);
274 assert!(schema.data_model.contains_key("trajectory_id"));
275 assert!(schema.data_model.contains_key("scope_id"));
276 assert!(schema.data_model.contains_key("agent_id"));
277 }
278}