Skip to main content

computer_use/
state.rs

1use std::fmt::{self, Write as _};
2use std::time::SystemTime;
3
4use crate::app::AppInfo;
5use crate::element::{Element, ElementStates};
6use crate::geometry::Rect;
7use crate::screenshot::Screenshot;
8
9/// The window an [`AppState`] was captured from.
10#[derive(Debug, Clone, PartialEq, Default)]
11#[non_exhaustive]
12pub struct WindowInfo {
13    /// Platform window ID (`CGWindowID`, `HWND`, X11 window), when known.
14    pub id: Option<u64>,
15    /// Window title.
16    pub title: Option<String>,
17    /// Frame in screen points.
18    pub frame: Rect,
19}
20
21impl WindowInfo {
22    /// Creates window info.
23    #[must_use]
24    pub const fn new(id: Option<u64>, title: Option<String>, frame: Rect) -> Self {
25        Self { id, title, frame }
26    }
27}
28
29/// Limits and toggles for [`Computer::get_app_state`](crate::Computer::get_app_state).
30#[derive(Debug, Clone, PartialEq, Eq)]
31#[non_exhaustive]
32pub struct CaptureOptions {
33    /// Walk the accessibility tree. Disabling this invalidates element indexes.
34    pub include_tree: bool,
35    /// Capture a screenshot. Disabling this prevents coordinate-based actions.
36    pub include_screenshot: bool,
37    /// Maximum characters kept per text attribute; `None` keeps everything.
38    pub text_limit: Option<usize>,
39    /// Maximum number of raw accessibility nodes visited.
40    pub max_nodes: usize,
41    /// Maximum tree depth visited.
42    pub max_depth: usize,
43    /// Screenshots larger than this on either side are downscaled; `None` keeps full resolution.
44    pub max_screenshot_dimension: Option<u32>,
45    /// Launch the app (without activating it) if it is installed but not running.
46    pub launch_if_needed: bool,
47}
48
49impl Default for CaptureOptions {
50    fn default() -> Self {
51        Self {
52            include_tree: true,
53            include_screenshot: true,
54            text_limit: Some(500),
55            max_nodes: 1200,
56            max_depth: 64,
57            max_screenshot_dimension: Some(1280),
58            launch_if_needed: true,
59        }
60    }
61}
62
63impl CaptureOptions {
64    /// Sets [`include_tree`](Self::include_tree).
65    #[must_use]
66    pub const fn include_tree(mut self, yes: bool) -> Self {
67        self.include_tree = yes;
68        self
69    }
70
71    /// Sets [`include_screenshot`](Self::include_screenshot).
72    #[must_use]
73    pub const fn include_screenshot(mut self, yes: bool) -> Self {
74        self.include_screenshot = yes;
75        self
76    }
77
78    /// Sets [`text_limit`](Self::text_limit).
79    #[must_use]
80    pub const fn text_limit(mut self, limit: Option<usize>) -> Self {
81        self.text_limit = limit;
82        self
83    }
84
85    /// Sets [`max_nodes`](Self::max_nodes).
86    #[must_use]
87    pub const fn max_nodes(mut self, max: usize) -> Self {
88        self.max_nodes = max;
89        self
90    }
91
92    /// Sets [`max_depth`](Self::max_depth).
93    #[must_use]
94    pub const fn max_depth(mut self, max: usize) -> Self {
95        self.max_depth = max;
96        self
97    }
98
99    /// Sets [`max_screenshot_dimension`](Self::max_screenshot_dimension).
100    #[must_use]
101    pub const fn max_screenshot_dimension(mut self, max: Option<u32>) -> Self {
102        self.max_screenshot_dimension = max;
103        self
104    }
105
106    /// Sets [`launch_if_needed`](Self::launch_if_needed).
107    #[must_use]
108    pub const fn launch_if_needed(mut self, yes: bool) -> Self {
109        self.launch_if_needed = yes;
110        self
111    }
112}
113
114/// A snapshot of one app window.
115#[derive(Debug, Clone)]
116#[non_exhaustive]
117pub struct AppState {
118    /// The app.
119    pub app: AppInfo,
120    /// The captured window.
121    pub window: WindowInfo,
122    /// Accessibility elements in depth-first order; `elements[i].index == i`.
123    pub elements: Vec<Element>,
124    /// Whether node or depth limits cut the tree short.
125    pub truncated: bool,
126    /// Whether the tree was requested at all.
127    pub tree_included: bool,
128    /// Index of the focused element, if it is part of the tree.
129    pub focused: Option<usize>,
130    /// Currently selected text in the focused element.
131    pub selected_text: Option<String>,
132    /// Screenshot of the window.
133    pub screenshot: Option<Screenshot>,
134    /// When the state was captured.
135    pub captured_at: SystemTime,
136    /// Non-fatal problems encountered while capturing, e.g. why the screenshot is missing.
137    pub notes: Vec<String>,
138}
139
140impl AppState {
141    /// Looks up an element by index.
142    #[must_use]
143    pub fn element(&self, index: usize) -> Option<&Element> {
144        self.elements.get(index)
145    }
146
147    /// The focused element.
148    #[must_use]
149    pub fn focused_element(&self) -> Option<&Element> {
150        self.focused.and_then(|i| self.element(i))
151    }
152
153    /// Elements matching `query`, in tree order.
154    pub fn find(&self, query: &ElementQuery) -> impl Iterator<Item = &Element> + use<'_> {
155        let query = query.clone();
156        let limit = query.limit;
157        self.elements.iter().filter(move |e| query.matches(e)).take(limit)
158    }
159
160    /// The deepest element whose frame contains the screen point, preferring actionable ones.
161    #[must_use]
162    pub fn element_at(&self, point: crate::geometry::Point) -> Option<&Element> {
163        self.elements.iter().filter(|e| e.frame.is_some_and(|f| f.contains(point))).min_by(|a, b| {
164            let rank = |e: &Element| (e.primary_action().is_none(), e.frame.map_or(f64::MAX, |f| f.area()));
165            let (ra, rb) = (rank(a), rank(b));
166            ra.0.cmp(&rb.0).then(ra.1.total_cmp(&rb.1))
167        })
168    }
169
170    /// Renders the state as compact, human-readable text (also handy as language-model input).
171    ///
172    /// Each element is one tab-indented line beginning with its index:
173    ///
174    /// ```text
175    /// App=com.apple.TextEdit (pid 812)
176    /// Window: "Untitled", App: TextEdit.
177    /// 0 standard window Untitled
178    ///     1 text area (focused, settable, editable) Value: Hello
179    ///     2 button Close Secondary Actions: Show Menu
180    /// ```
181    #[must_use]
182    pub fn render(&self) -> String {
183        let mut out = String::new();
184        let pid = self.app.pid.map_or_else(String::new, |p| format!(" (pid {p})"));
185        let _ = writeln!(out, "App={}{pid}", self.app.reference());
186        let title = self.window.title.as_deref().filter(|t| !t.is_empty()).unwrap_or(&self.app.name);
187        let _ = write!(out, "Window: {title:?}, App: {}.", self.app.name);
188        for element in &self.elements {
189            out.push('\n');
190            for _ in 0..element.depth {
191                out.push('\t');
192            }
193            render_element(&mut out, element);
194        }
195        if !self.tree_included {
196            out.push_str("\nAccessibility tree omitted (include_tree = false); element indexes are unavailable.");
197        } else if self.truncated {
198            out.push_str("\n… tree truncated by node or depth limits.");
199        }
200        for note in &self.notes {
201            let _ = write!(out, "\n{note}");
202        }
203        if let Some(selected) = self.selected_text.as_deref().filter(|s| !s.is_empty()) {
204            let _ = write!(out, "\n\nSelected text: [{selected}]");
205        } else if let Some(focused) = self.focused_element() {
206            out.push_str("\n\nThe focused UI element is ");
207            render_element(&mut out, focused);
208            out.push('.');
209        }
210        out
211    }
212}
213
214impl fmt::Display for AppState {
215    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
216        f.write_str(&self.render())
217    }
218}
219
220fn render_element(out: &mut String, e: &Element) {
221    let _ = write!(out, "{} {}", e.index, e.role);
222    let states = (e.states - ElementStates::OFFSCREEN).labels();
223    if !states.is_empty() {
224        let _ = write!(out, " ({})", states.join(", "));
225    }
226    if let Some(title) = &e.title {
227        let _ = write!(out, " {title}");
228    }
229    let mut field = |name: &str, value: Option<&String>| {
230        if let Some(v) = value.filter(|v| Some(*v) != e.title.as_ref()) {
231            let _ = write!(out, " {name}: {v}");
232        }
233    };
234    field("Description", e.description.as_ref());
235    field("Help", e.help.as_ref().filter(|h| Some(*h) != e.description.as_ref()));
236    field("URL", e.url.as_ref());
237    field("ID", e.identifier.as_ref());
238    field("Value", e.value.as_ref());
239    if e.value.is_none() {
240        field("Placeholder", e.placeholder.as_ref());
241    }
242    let secondary: Vec<&str> = e.actions.iter().filter(|a| !a.primary).map(|a| a.label.as_str()).collect();
243    if !secondary.is_empty() {
244        let _ = write!(out, " Secondary Actions: {}", secondary.join(", "));
245    }
246}
247
248/// Criteria for [`Computer::query`](crate::Computer::query) and [`AppState::find`].
249#[derive(Debug, Clone, PartialEq, Eq)]
250pub struct ElementQuery {
251    /// Text matched against title, description, value, placeholder, identifier, and help.
252    pub text: Option<String>,
253    /// Role matched against the display role or platform role (an `AX` prefix is optional).
254    pub role: Option<String>,
255    /// Require an exact (case-insensitive) text match instead of a substring.
256    pub exact: bool,
257    /// Maximum number of matches.
258    pub limit: usize,
259}
260
261impl Default for ElementQuery {
262    fn default() -> Self {
263        Self { text: None, role: None, exact: false, limit: 20 }
264    }
265}
266
267impl ElementQuery {
268    /// Matches elements containing `text`.
269    #[must_use]
270    pub fn text(text: impl Into<String>) -> Self {
271        Self { text: Some(text.into()), ..Self::default() }
272    }
273
274    /// Matches elements with `role`.
275    #[must_use]
276    pub fn role(role: impl Into<String>) -> Self {
277        Self { role: Some(role.into()), ..Self::default() }
278    }
279
280    /// Adds a role constraint.
281    #[must_use]
282    pub fn with_role(mut self, role: impl Into<String>) -> Self {
283        self.role = Some(role.into());
284        self
285    }
286
287    /// Requires exact text matches.
288    #[must_use]
289    pub const fn exact(mut self, exact: bool) -> Self {
290        self.exact = exact;
291        self
292    }
293
294    /// Caps the number of results.
295    #[must_use]
296    pub const fn limit(mut self, limit: usize) -> Self {
297        self.limit = limit;
298        self
299    }
300
301    /// Whether `element` satisfies the query.
302    #[must_use]
303    pub fn matches(&self, element: &Element) -> bool {
304        if let Some(role) = &self.role {
305            let want = normalize_role(role);
306            if normalize_role(&element.role) != want && normalize_role(&element.platform_role) != want {
307                return false;
308            }
309        }
310        if let Some(text) = &self.text {
311            let needle = text.to_lowercase();
312            let hit = element.text_fields().any(|field| {
313                let field = field.to_lowercase();
314                if self.exact { field == needle } else { field.contains(&needle) }
315            });
316            if !hit {
317                return false;
318            }
319        }
320        true
321    }
322}
323
324fn normalize_role(role: &str) -> String {
325    let lowered: String = role.chars().filter(|c| c.is_alphanumeric()).flat_map(char::to_lowercase).collect();
326    lowered.strip_prefix("ax").map(str::to_owned).unwrap_or(lowered)
327}
328
329#[cfg(test)]
330mod tests {
331    use super::*;
332    use crate::element::{Action, ElementData, RawNode, build_tree};
333
334    pub(crate) fn sample_state() -> AppState {
335        let mut window = ElementData::new("standard window", "AXWindow");
336        window.title = Some("Doc".into());
337        let mut text = ElementData::new("text area", "AXTextArea");
338        text.value = Some("Hello world".into());
339        text.states = ElementStates::FOCUSED | ElementStates::SETTABLE;
340        text.frame = Some(Rect::new(0.0, 20.0, 200.0, 100.0));
341        let mut button = ElementData::new("button", "AXButton");
342        button.title = Some("Save".into());
343        button.actions = vec![Action::new("AXPress", true), Action::new("AXShowMenu", false)];
344        button.frame = Some(Rect::new(10.0, 130.0, 50.0, 20.0));
345        let nodes = vec![
346            RawNode { handle: (), parent: None, data: window },
347            RawNode { handle: (), parent: Some(0), data: text },
348            RawNode { handle: (), parent: Some(0), data: button },
349        ];
350        let (elements, _, _) = build_tree(nodes, Some(500));
351        AppState {
352            app: AppInfo::new("TextEdit").with_bundle_id("com.apple.TextEdit").with_pid(9),
353            window: WindowInfo::new(Some(1), Some("Doc".into()), Rect::new(0.0, 0.0, 300.0, 200.0)),
354            elements,
355            truncated: false,
356            tree_included: true,
357            focused: Some(1),
358            selected_text: None,
359            screenshot: None,
360            captured_at: SystemTime::UNIX_EPOCH,
361            notes: Vec::new(),
362        }
363    }
364
365    #[test]
366    fn renders_tree() {
367        let text = sample_state().render();
368        let expected = "App=com.apple.TextEdit (pid 9)\n\
369            Window: \"Doc\", App: TextEdit.\n\
370            0 standard window Doc\n\
371            \t1 text area (focused, settable) Value: Hello world\n\
372            \t2 button Save Secondary Actions: Show Menu\n\
373            \n\
374            The focused UI element is 1 text area (focused, settable) Value: Hello world.";
375        assert_eq!(text, expected);
376    }
377
378    #[test]
379    fn queries_by_text_and_role() {
380        let state = sample_state();
381        let q = ElementQuery::text("save");
382        assert_eq!(state.find(&q).map(|e| e.index).collect::<Vec<_>>(), vec![2]);
383        let q = ElementQuery::role("AXTextArea");
384        assert_eq!(state.find(&q).map(|e| e.index).collect::<Vec<_>>(), vec![1]);
385        let q = ElementQuery::role("button").exact(true);
386        assert_eq!(state.find(&q).count(), 1);
387        let q = ElementQuery::text("hello").exact(true);
388        assert_eq!(state.find(&q).count(), 0);
389    }
390
391    #[test]
392    fn hit_tests_prefer_actionable_elements() {
393        let state = sample_state();
394        let hit = state.element_at(crate::geometry::Point::new(20.0, 140.0));
395        assert_eq!(hit.map(|e| e.index), Some(2));
396    }
397}