Skip to main content

eoka_server/
snapshot.rs

1use std::collections::HashMap;
2use std::fmt::Write;
3
4use crate::eoka::Page;
5use anyhow::Result;
6use serde::{Deserialize, Serialize};
7
8#[derive(Serialize)]
9#[serde(rename_all = "camelCase")]
10struct GetFullAXTreeParams {}
11
12#[derive(Deserialize)]
13#[serde(rename_all = "camelCase")]
14struct AXTreeResult {
15    nodes: Vec<AXNode>,
16}
17
18#[derive(Deserialize)]
19#[serde(rename_all = "camelCase")]
20struct AXNode {
21    node_id: String,
22    #[serde(default)]
23    ignored: bool,
24    role: Option<AXValue>,
25    name: Option<AXValue>,
26    value: Option<AXValue>,
27    #[serde(default)]
28    properties: Vec<AXProperty>,
29    parent_id: Option<String>,
30    #[serde(rename = "backendDOMNodeId")]
31    backend_dom_node_id: Option<i64>,
32}
33
34#[derive(Deserialize)]
35#[serde(rename_all = "camelCase")]
36struct AXValue {
37    value: Option<serde_json::Value>,
38}
39
40impl AXValue {
41    fn as_str(&self) -> &str {
42        self.value.as_ref().and_then(|v| v.as_str()).unwrap_or("")
43    }
44
45    fn is_truthy(&self) -> bool {
46        self.value
47            .as_ref()
48            .map(|v| v.as_bool().unwrap_or(false) || v.as_str() == Some("true"))
49            .unwrap_or(false)
50    }
51
52    fn as_bool(&self) -> Option<bool> {
53        match self.value.as_ref()? {
54            serde_json::Value::Bool(value) => Some(*value),
55            serde_json::Value::String(value) if value == "true" => Some(true),
56            serde_json::Value::String(value) if value == "false" => Some(false),
57            _ => None,
58        }
59    }
60}
61
62#[derive(Deserialize)]
63#[serde(rename_all = "camelCase")]
64struct AXProperty {
65    name: String,
66    value: AXValue,
67}
68
69#[derive(Debug, Clone)]
70pub struct SnapshotRef {
71    pub ref_label: String,
72    pub backend_node_id: i64,
73    pub role: String,
74    pub name: String,
75    pub disabled: bool,
76    pub checked: Option<bool>,
77    pub expanded: Option<bool>,
78}
79
80#[derive(Debug)]
81pub struct SnapshotResult {
82    pub tree_text: String,
83    pub refs: Vec<SnapshotRef>,
84}
85
86const NOISE_ROLES: &[&str] = &[
87    "generic",
88    "none",
89    "presentation",
90    "InlineTextBox",
91    "LineBreak",
92];
93
94fn should_skip(role: &str, name: &str, ignored: bool, include_all: bool) -> bool {
95    if ignored {
96        return true;
97    }
98    if include_all {
99        return false;
100    }
101    if NOISE_ROLES.contains(&role) {
102        return true;
103    }
104    role == "StaticText" && name.is_empty()
105}
106
107fn opt_str(v: &Option<AXValue>) -> &str {
108    v.as_ref().map(|av| av.as_str()).unwrap_or("")
109}
110
111fn format_props(node: &AXNode, out: &mut String) {
112    let val = opt_str(&node.value);
113    if !val.is_empty() {
114        let _ = write!(out, " value=\"{}\"", val);
115    }
116
117    for prop in &node.properties {
118        let token: Option<&str> = match prop.name.as_str() {
119            "disabled" if prop.value.is_truthy() => Some("disabled"),
120            "required" if prop.value.is_truthy() => Some("required"),
121            "selected" if prop.value.is_truthy() => Some("selected"),
122            "focused" if prop.value.is_truthy() => Some("focused"),
123            "checked" => match prop.value.as_str() {
124                "true" => Some("checked"),
125                "mixed" => Some("checked=mixed"),
126                _ => None,
127            },
128            "expanded" => Some(if prop.value.is_truthy() {
129                "expanded"
130            } else {
131                "collapsed"
132            }),
133            "invalid" => match prop.value.as_str() {
134                "true" | "spelling" | "grammar" => Some("invalid"),
135                _ => None,
136            },
137            _ => None,
138        };
139        if let Some(t) = token {
140            let _ = write!(out, " {}", t);
141        }
142    }
143}
144
145struct Walker<'a> {
146    nodes_by_id: HashMap<&'a str, &'a AXNode>,
147    children_of: HashMap<&'a str, Vec<&'a str>>,
148    include_all: bool,
149    ref_counter: usize,
150    refs: Vec<SnapshotRef>,
151    output: String,
152}
153
154impl<'a> Walker<'a> {
155    fn walk(&mut self, node_id: &str, depth: usize) {
156        let node = match self.nodes_by_id.get(node_id) {
157            Some(n) => *n,
158            None => return,
159        };
160
161        let role = opt_str(&node.role);
162        let name = opt_str(&node.name);
163        let skip = should_skip(role, name, node.ignored, self.include_all);
164
165        if !skip {
166            self.emit_node(node, role, name, depth);
167        }
168
169        let child_depth = if skip { depth } else { depth + 1 };
170        if let Some(children) = self.children_of.get(node_id) {
171            for child_id in children.clone() {
172                self.walk(child_id, child_depth);
173            }
174        }
175    }
176
177    fn emit_node(&mut self, node: &AXNode, role: &str, name: &str, depth: usize) {
178        let ref_label = match node.backend_dom_node_id {
179            Some(id) if id > 0 => {
180                self.ref_counter += 1;
181                let label = format!("@e{}", self.ref_counter);
182                self.refs.push(SnapshotRef {
183                    ref_label: label.clone(),
184                    backend_node_id: id,
185                    role: role.to_string(),
186                    name: name.to_string(),
187                    disabled: node
188                        .properties
189                        .iter()
190                        .any(|p| p.name == "disabled" && p.value.is_truthy()),
191                    checked: node
192                        .properties
193                        .iter()
194                        .find(|p| p.name == "checked")
195                        .and_then(|p| p.value.as_bool()),
196                    expanded: node
197                        .properties
198                        .iter()
199                        .find(|p| p.name == "expanded")
200                        .and_then(|p| p.value.as_bool()),
201                });
202                Some(label)
203            }
204            _ => None,
205        };
206
207        for _ in 0..depth {
208            self.output.push_str("  ");
209        }
210        if let Some(ref label) = ref_label {
211            let _ = write!(self.output, "{} ", label);
212        }
213        self.output.push_str(role);
214        if !name.is_empty() {
215            let _ = write!(self.output, " \"{}\"", name);
216        }
217        format_props(node, &mut self.output);
218        self.output.push('\n');
219    }
220}
221
222pub async fn snapshot(page: &Page, include_all: bool) -> Result<SnapshotResult> {
223    let _ = page
224        .session()
225        .send::<_, serde_json::Value>("DOM.enable", &serde_json::json!({}))
226        .await;
227    let _ = page
228        .session()
229        .send::<_, serde_json::Value>(
230            "DOM.getDocument",
231            &serde_json::json!({"depth": -1, "pierce": true}),
232        )
233        .await;
234
235    let result: AXTreeResult = page
236        .session()
237        .send("Accessibility.getFullAXTree", &GetFullAXTreeParams {})
238        .await
239        .map_err(|e| anyhow::anyhow!("CDP Accessibility.getFullAXTree failed: {}", e))?;
240
241    if result.nodes.is_empty() {
242        return Ok(SnapshotResult {
243            tree_text: "No accessibility tree available.".into(),
244            refs: Vec::new(),
245        });
246    }
247
248    let (nodes_by_id, children_of, roots) = build_tree_index(&result.nodes);
249
250    let mut walker = Walker {
251        nodes_by_id,
252        children_of,
253        include_all,
254        ref_counter: 0,
255        refs: Vec::new(),
256        output: String::with_capacity(4096),
257    };
258
259    for root in roots {
260        walker.walk(root, 0);
261    }
262
263    Ok(SnapshotResult {
264        tree_text: walker.output,
265        refs: walker.refs,
266    })
267}
268
269type TreeIndex<'a> = (
270    HashMap<&'a str, &'a AXNode>,
271    HashMap<&'a str, Vec<&'a str>>,
272    Vec<&'a str>,
273);
274
275fn build_tree_index(nodes: &[AXNode]) -> TreeIndex<'_> {
276    let mut by_id: HashMap<&str, &AXNode> = HashMap::with_capacity(nodes.len());
277    let mut children: HashMap<&str, Vec<&str>> = HashMap::new();
278    let mut roots: Vec<&str> = Vec::new();
279
280    for node in nodes {
281        by_id.insert(&node.node_id, node);
282    }
283
284    for node in nodes {
285        match node.parent_id.as_deref() {
286            Some(pid) if by_id.contains_key(pid) => {
287                children.entry(pid).or_default().push(&node.node_id);
288            }
289            _ => roots.push(&node.node_id),
290        }
291    }
292
293    (by_id, children, roots)
294}
295
296#[cfg(test)]
297mod tests {
298    use super::*;
299
300    #[test]
301    fn skip_ignored_nodes() {
302        assert!(should_skip("button", "Click", true, false));
303        assert!(should_skip("button", "Click", true, true));
304    }
305
306    #[test]
307    fn skip_noise_roles() {
308        for role in NOISE_ROLES {
309            assert!(should_skip(role, "", false, false), "should skip {}", role);
310        }
311    }
312
313    #[test]
314    fn skip_unnamed_static_text() {
315        assert!(should_skip("StaticText", "", false, false));
316        assert!(!should_skip("StaticText", "Hello", false, false));
317    }
318
319    #[test]
320    fn keep_normal_roles() {
321        for (role, name) in [
322            ("button", "Submit"),
323            ("link", "Home"),
324            ("textbox", "Email"),
325            ("heading", "Title"),
326        ] {
327            assert!(!should_skip(role, name, false, false));
328        }
329    }
330
331    #[test]
332    fn include_all_keeps_noise() {
333        assert!(!should_skip("generic", "", false, true));
334        assert!(!should_skip("none", "", false, true));
335    }
336
337    #[test]
338    fn ax_value_helpers() {
339        let truthy = AXValue {
340            value: Some(serde_json::json!(true)),
341        };
342        assert!(truthy.is_truthy());
343        assert_eq!(truthy.as_str(), "");
344
345        let str_true = AXValue {
346            value: Some(serde_json::json!("true")),
347        };
348        assert!(str_true.is_truthy());
349        assert_eq!(str_true.as_str(), "true");
350
351        let falsy = AXValue {
352            value: Some(serde_json::json!(false)),
353        };
354        assert!(!falsy.is_truthy());
355
356        let none = AXValue { value: None };
357        assert!(!none.is_truthy());
358        assert_eq!(none.as_str(), "");
359    }
360
361    #[test]
362    fn format_props_value_and_disabled() {
363        let node = AXNode {
364            node_id: "1".into(),
365            ignored: false,
366            role: None,
367            name: None,
368            value: Some(AXValue {
369                value: Some(serde_json::json!("hello")),
370            }),
371            properties: vec![AXProperty {
372                name: "disabled".into(),
373                value: AXValue {
374                    value: Some(serde_json::json!(true)),
375                },
376            }],
377            parent_id: None,
378            backend_dom_node_id: None,
379        };
380        let mut out = String::new();
381        format_props(&node, &mut out);
382        assert!(out.contains("value=\"hello\""));
383        assert!(out.contains("disabled"));
384    }
385}