Skip to main content

agpu/ontology/
registry.rs

1//! Ontology registry, UI tree, and node types.
2//!
3//! The [`OntologyRegistry`] catalogs widget schemas, while [`UiTree`] and
4//! [`UiNode`] represent the live widget tree that agents can inspect.
5
6use super::{AgentCapability, Discoverable, SemanticRole, WidgetSchema};
7use crate::core::Rect;
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10
11/// Global registry of widget types and the live UI tree.
12///
13/// 1. **Type catalog**: Agents can list all available widget types, search by
14///    name/role/tag, and read schemas before interacting.
15/// 2. **UI tree**: The current widget hierarchy exposed as a navigable tree.
16#[derive(Debug, Default)]
17pub struct OntologyRegistry {
18    schemas: HashMap<String, WidgetSchema>,
19    tree: Option<UiTree>,
20}
21
22impl OntologyRegistry {
23    pub fn new() -> Self {
24        Self::default()
25    }
26
27    // ── Type Catalog ─────────────────────────────────────────────────
28
29    /// Register a widget type's schema.
30    pub fn register_schema(&mut self, schema: WidgetSchema) {
31        self.schemas.insert(schema.name.clone(), schema);
32    }
33
34    /// Register a discoverable widget type (convenience).
35    pub fn register<W: Discoverable>(&mut self, instance: &W) {
36        self.register_schema(instance.schema());
37    }
38
39    /// List all registered widget type names.
40    pub fn list_types(&self) -> Vec<&str> {
41        self.schemas.keys().map(|s| s.as_str()).collect()
42    }
43
44    /// Get the schema for a widget type by name.
45    pub fn get_schema(&self, name: &str) -> Option<&WidgetSchema> {
46        self.schemas.get(name)
47    }
48
49    /// Find widget types matching a semantic role.
50    pub fn find_by_role(&self, role: SemanticRole) -> Vec<&WidgetSchema> {
51        self.schemas
52            .values()
53            .filter(|s| s.default_role == role)
54            .collect()
55    }
56
57    /// Search widget types by tag (case-insensitive substring match).
58    pub fn search(&self, query: &str) -> Vec<&WidgetSchema> {
59        let query_lower = query.to_lowercase();
60        self.schemas
61            .values()
62            .filter(|s| {
63                s.name.to_lowercase().contains(&query_lower)
64                    || s.description.to_lowercase().contains(&query_lower)
65                    || s.tags
66                        .iter()
67                        .any(|t| t.to_lowercase().contains(&query_lower))
68            })
69            .collect()
70    }
71
72    /// Export the full type catalog as JSON.
73    pub fn export_catalog(&self) -> serde_json::Value {
74        serde_json::to_value(&self.schemas).unwrap_or_default()
75    }
76
77    /// Validate params against a declared action schema.
78    pub fn validate_action_params(
79        &self,
80        widget_type: &str,
81        action: &str,
82        params: &serde_json::Value,
83    ) -> Result<(), String> {
84        let Some(schema) = self.schemas.get(widget_type) else {
85            return Ok(());
86        };
87        let Some(declared) = schema.actions.iter().find(|a| a.name == action) else {
88            return Ok(());
89        };
90        declared.validate_params(params)
91    }
92
93    // ── Live UI Tree ─────────────────────────────────────────────────
94
95    /// Set the current UI tree snapshot.
96    pub fn set_tree(&mut self, tree: UiTree) {
97        self.tree = Some(tree);
98    }
99
100    /// Get the current UI tree.
101    pub fn tree(&self) -> Option<&UiTree> {
102        self.tree.as_ref()
103    }
104
105    /// Find a node in the UI tree by its agent ID.
106    pub fn find_node(&self, agent_id: &str) -> Option<&UiNode> {
107        self.tree.as_ref().and_then(|t| t.find(agent_id))
108    }
109
110    /// Export the UI tree as JSON for agent consumption.
111    pub fn export_tree(&self) -> serde_json::Value {
112        match &self.tree {
113            Some(tree) => serde_json::to_value(tree).unwrap_or_default(),
114            None => serde_json::Value::Null,
115        }
116    }
117}
118
119/// A snapshot of the live UI widget tree.
120#[derive(Debug, Clone, Serialize, Deserialize)]
121pub struct UiTree {
122    pub root: UiNode,
123}
124
125impl UiTree {
126    pub fn new(root: UiNode) -> Self {
127        Self { root }
128    }
129
130    /// Depth-first search for a node by agent_id.
131    pub fn find(&self, agent_id: &str) -> Option<&UiNode> {
132        self.root.find(agent_id)
133    }
134
135    /// Collect all nodes matching a role.
136    pub fn find_by_role(&self, role: SemanticRole) -> Vec<&UiNode> {
137        let mut results = Vec::new();
138        self.root.collect_by_role(role, &mut results);
139        results
140    }
141
142    /// Collect all focusable nodes.
143    pub fn focusable_nodes(&self) -> Vec<&UiNode> {
144        let mut results = Vec::new();
145        self.root.collect_by_capability("focusable", &mut results);
146        results
147    }
148}
149
150/// A node in the UI tree representing a single widget instance.
151#[derive(Debug, Clone, Serialize, Deserialize)]
152pub struct UiNode {
153    /// Optional unique agent-addressable ID.
154    pub agent_id: Option<String>,
155    /// The widget type name (matches a registered schema).
156    pub widget_type: String,
157    /// Semantic role of this instance.
158    pub role: SemanticRole,
159    /// Capabilities of this instance.
160    pub capabilities: Vec<AgentCapability>,
161    /// Current state snapshot as JSON.
162    pub state: serde_json::Value,
163    /// Accessibility label.
164    pub label: Option<String>,
165    /// Bounding rectangle in logical pixel coordinates.
166    pub bounds: Option<NodeBounds>,
167    /// Accessibility attributes for screen readers and assistive agents.
168    #[serde(default, skip_serializing_if = "Accessibility::is_empty")]
169    pub accessibility: Accessibility,
170    /// Child nodes.
171    pub children: Vec<UiNode>,
172}
173
174/// Bounding rectangle of a UI node in logical pixel coordinates.
175#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
176pub struct NodeBounds {
177    pub x: f32,
178    pub y: f32,
179    pub width: f32,
180    pub height: f32,
181}
182
183impl From<Rect> for NodeBounds {
184    fn from(r: Rect) -> Self {
185        Self {
186            x: r.x,
187            y: r.y,
188            width: r.width,
189            height: r.height,
190        }
191    }
192}
193
194/// Accessibility attributes following ARIA conventions.
195///
196/// Agents and screen readers use these to understand widget semantics.
197#[derive(Debug, Clone, Default, Serialize, Deserialize)]
198pub struct Accessibility {
199    #[serde(default, skip_serializing_if = "Option::is_none")]
200    pub role: Option<String>,
201    #[serde(default, skip_serializing_if = "Option::is_none")]
202    pub description: Option<String>,
203    #[serde(default, skip_serializing_if = "Option::is_none")]
204    pub disabled: Option<bool>,
205    #[serde(default, skip_serializing_if = "Option::is_none")]
206    pub value_text: Option<String>,
207    #[serde(default, skip_serializing_if = "Option::is_none")]
208    pub expanded: Option<bool>,
209    #[serde(default, skip_serializing_if = "Option::is_none")]
210    pub selected: Option<bool>,
211    #[serde(default, skip_serializing_if = "Option::is_none")]
212    pub required: Option<bool>,
213    #[serde(default, skip_serializing_if = "Option::is_none")]
214    pub shortcut: Option<String>,
215    #[serde(default, skip_serializing_if = "Option::is_none")]
216    pub tab_index: Option<i32>,
217    #[serde(default, skip_serializing_if = "Option::is_none")]
218    pub live: Option<String>,
219}
220
221impl Accessibility {
222    pub fn is_empty(&self) -> bool {
223        self.role.is_none()
224            && self.description.is_none()
225            && self.disabled.is_none()
226            && self.value_text.is_none()
227            && self.expanded.is_none()
228            && self.selected.is_none()
229            && self.required.is_none()
230            && self.shortcut.is_none()
231            && self.tab_index.is_none()
232            && self.live.is_none()
233    }
234}
235
236impl UiNode {
237    #[must_use]
238    pub fn new(widget_type: impl Into<String>, role: SemanticRole) -> Self {
239        Self {
240            agent_id: None,
241            widget_type: widget_type.into(),
242            role,
243            capabilities: Vec::new(),
244            state: serde_json::Value::Null,
245            label: None,
246            bounds: None,
247            accessibility: Accessibility::default(),
248            children: Vec::new(),
249        }
250    }
251
252    #[must_use]
253    pub fn with_id(mut self, id: impl Into<String>) -> Self {
254        self.agent_id = Some(id.into());
255        self
256    }
257
258    #[must_use]
259    pub fn with_label(mut self, label: impl Into<String>) -> Self {
260        self.label = Some(label.into());
261        self
262    }
263
264    #[must_use]
265    pub fn with_bounds(mut self, bounds: NodeBounds) -> Self {
266        self.bounds = Some(bounds);
267        self
268    }
269
270    #[must_use]
271    pub fn with_state(mut self, state: serde_json::Value) -> Self {
272        self.state = state;
273        self
274    }
275
276    #[must_use]
277    pub fn with_capability(mut self, cap: AgentCapability) -> Self {
278        self.capabilities.push(cap);
279        self
280    }
281
282    #[must_use]
283    pub fn with_child(mut self, child: UiNode) -> Self {
284        self.children.push(child);
285        self
286    }
287
288    /// Set accessibility attributes.
289    #[must_use]
290    pub fn with_accessibility(mut self, acc: Accessibility) -> Self {
291        self.accessibility = acc;
292        self
293    }
294
295    /// Convenience: set a named property in the state JSON object.
296    #[must_use]
297    pub fn with_property(mut self, key: &str, value: serde_json::Value) -> Self {
298        if self.state.is_null() {
299            self.state = serde_json::json!({});
300        }
301        if let Some(obj) = self.state.as_object_mut() {
302            obj.insert(key.to_string(), value);
303        }
304        self
305    }
306
307    /// Depth-first search by agent_id.
308    pub fn find(&self, agent_id: &str) -> Option<&UiNode> {
309        if self.agent_id.as_deref() == Some(agent_id) {
310            return Some(self);
311        }
312        for child in &self.children {
313            if let Some(node) = child.find(agent_id) {
314                return Some(node);
315            }
316        }
317        None
318    }
319
320    /// Collect nodes matching a semantic role.
321    pub fn collect_by_role<'a>(&'a self, role: SemanticRole, results: &mut Vec<&'a UiNode>) {
322        if self.role == role {
323            results.push(self);
324        }
325        for child in &self.children {
326            child.collect_by_role(role, results);
327        }
328    }
329
330    /// Collect nodes with a specific capability.
331    pub fn collect_by_capability<'a>(&'a self, cap_name: &str, results: &mut Vec<&'a UiNode>) {
332        if self.capabilities.iter().any(|c| c.name() == cap_name) {
333            results.push(self);
334        }
335        for child in &self.children {
336            child.collect_by_capability(cap_name, results);
337        }
338    }
339}
340
341#[cfg(test)]
342mod tests {
343    use super::*;
344
345    #[test]
346    fn registry_search() {
347        let mut reg = OntologyRegistry::new();
348        reg.register_schema(WidgetSchema {
349            name: "Button".into(),
350            description: "A clickable button".into(),
351            default_role: SemanticRole::Action,
352            properties: vec![],
353            actions: vec![],
354            usage_hint: None,
355            tags: vec!["button".into(), "action".into()],
356        });
357        assert_eq!(reg.search("button").len(), 1);
358        assert_eq!(reg.search("nonexistent").len(), 0);
359    }
360
361    #[test]
362    fn ui_tree_find() {
363        let tree = UiTree::new(
364            UiNode::new("Panel", SemanticRole::Container)
365                .with_id("root")
366                .with_child(
367                    UiNode::new("Button", SemanticRole::Action)
368                        .with_id("btn-1")
369                        .with_capability(AgentCapability::Focusable),
370                ),
371        );
372        assert!(tree.find("btn-1").is_some());
373        assert!(tree.find("missing").is_none());
374        assert_eq!(tree.focusable_nodes().len(), 1);
375    }
376
377    #[test]
378    fn registry_list_types() {
379        let mut reg = OntologyRegistry::new();
380        reg.register_schema(WidgetSchema::new("Button", "btn", SemanticRole::Action));
381        reg.register_schema(WidgetSchema::new("Label", "lbl", SemanticRole::Display));
382        let types = reg.list_types();
383        assert_eq!(types.len(), 2);
384        assert!(types.contains(&"Button"));
385        assert!(types.contains(&"Label"));
386    }
387
388    #[test]
389    fn registry_get_schema() {
390        let mut reg = OntologyRegistry::new();
391        reg.register_schema(WidgetSchema::new("Button", "btn", SemanticRole::Action));
392        assert!(reg.get_schema("Button").is_some());
393        assert!(reg.get_schema("Missing").is_none());
394    }
395
396    #[test]
397    fn registry_find_by_role() {
398        let mut reg = OntologyRegistry::new();
399        reg.register_schema(WidgetSchema::new("Button", "btn", SemanticRole::Action));
400        reg.register_schema(WidgetSchema::new("Label", "lbl", SemanticRole::Display));
401        reg.register_schema(WidgetSchema::new("Link", "link", SemanticRole::Action));
402        assert_eq!(reg.find_by_role(SemanticRole::Action).len(), 2);
403        assert_eq!(reg.find_by_role(SemanticRole::Display).len(), 1);
404        assert_eq!(reg.find_by_role(SemanticRole::Container).len(), 0);
405    }
406
407    #[test]
408    fn registry_export_catalog() {
409        let mut reg = OntologyRegistry::new();
410        reg.register_schema(WidgetSchema::new("Button", "btn", SemanticRole::Action));
411        let catalog = reg.export_catalog();
412        assert!(catalog.is_object());
413        assert!(catalog.get("Button").is_some());
414    }
415
416    #[test]
417    fn registry_validate_action_params_unknown_type() {
418        let reg = OntologyRegistry::new();
419        // Unknown widget type should pass validation (permissive)
420        assert!(
421            reg.validate_action_params("Unknown", "click", &serde_json::json!({}))
422                .is_ok()
423        );
424    }
425
426    #[test]
427    fn ui_tree_find_by_role() {
428        let tree = UiTree::new(
429            UiNode::new("Panel", SemanticRole::Container)
430                .with_id("root")
431                .with_child(UiNode::new("Button", SemanticRole::Action).with_id("b1"))
432                .with_child(UiNode::new("Button", SemanticRole::Action).with_id("b2"))
433                .with_child(UiNode::new("Label", SemanticRole::Display).with_id("l1")),
434        );
435        assert_eq!(tree.find_by_role(SemanticRole::Action).len(), 2);
436        assert_eq!(tree.find_by_role(SemanticRole::Display).len(), 1);
437        assert_eq!(tree.find_by_role(SemanticRole::Container).len(), 1);
438    }
439
440    #[test]
441    fn ui_node_builder_chain() {
442        let node = UiNode::new("TextInput", SemanticRole::Input)
443            .with_id("input-1")
444            .with_label("Username")
445            .with_bounds(NodeBounds {
446                x: 10.0,
447                y: 20.0,
448                width: 200.0,
449                height: 30.0,
450            })
451            .with_state(serde_json::json!({"value": ""}))
452            .with_capability(AgentCapability::Focusable)
453            .with_capability(AgentCapability::TextInput {
454                multiline: false,
455                max_length: Some(100),
456            });
457
458        assert_eq!(node.agent_id.as_deref(), Some("input-1"));
459        assert_eq!(node.label.as_deref(), Some("Username"));
460        assert!(node.bounds.is_some());
461        assert_eq!(node.capabilities.len(), 2);
462    }
463
464    #[test]
465    fn ui_node_with_property() {
466        let node = UiNode::new("Slider", SemanticRole::Input)
467            .with_property("value", serde_json::json!(50))
468            .with_property("min", serde_json::json!(0));
469        assert_eq!(node.state["value"], 50);
470        assert_eq!(node.state["min"], 0);
471    }
472
473    #[test]
474    fn node_bounds_from_rect() {
475        let rect = Rect::new(1.0, 2.0, 3.0, 4.0);
476        let nb: NodeBounds = rect.into();
477        assert_eq!(nb.x, 1.0);
478        assert_eq!(nb.y, 2.0);
479        assert_eq!(nb.width, 3.0);
480        assert_eq!(nb.height, 4.0);
481    }
482
483    #[test]
484    fn accessibility_is_empty() {
485        let a = Accessibility::default();
486        assert!(a.is_empty());
487
488        let a2 = Accessibility {
489            role: Some("button".into()),
490            ..Default::default()
491        };
492        assert!(!a2.is_empty());
493    }
494
495    #[test]
496    fn registry_set_and_get_tree() {
497        let mut reg = OntologyRegistry::new();
498        assert!(reg.tree().is_none());
499
500        let tree = UiTree::new(UiNode::new("Root", SemanticRole::Container));
501        reg.set_tree(tree);
502        assert!(reg.tree().is_some());
503    }
504
505    #[test]
506    fn registry_find_node() {
507        let mut reg = OntologyRegistry::new();
508        let tree = UiTree::new(
509            UiNode::new("Root", SemanticRole::Container)
510                .with_id("root")
511                .with_child(UiNode::new("Button", SemanticRole::Action).with_id("btn")),
512        );
513        reg.set_tree(tree);
514        assert!(reg.find_node("btn").is_some());
515        assert!(reg.find_node("missing").is_none());
516    }
517
518    #[test]
519    fn registry_export_tree() {
520        let mut reg = OntologyRegistry::new();
521        assert!(reg.export_tree().is_null());
522
523        let tree = UiTree::new(UiNode::new("Root", SemanticRole::Container).with_id("root"));
524        reg.set_tree(tree);
525        let exported = reg.export_tree();
526        assert!(exported.is_object());
527    }
528}