From 6f31fb04dd5c96011540ed8f2a1b49cda99d7dc3 Mon Sep 17 00:00:00 2001
From: Tryanks <tryanks@outlook.com>
Date: Fri, 2 Oct 2026 02:50:56 +0800
Subject: [PATCH] fix(computer-use): tell visited AX elements apart by
 identity, not address

The walk deduplicated elements by the address of the AXUIElement it was
handed. Each child array is released once its subtree is walked, and the
next array the system hands out reuses freed addresses, so a node whose
reference landed on a released sibling's address was skipped: the first
child of a group, typically. Visited elements are now retained for the
walk and compared with CFEqual under CFHash.
---
 .../computer-use-mcp/src/backend/macos/ax.rs  | 66 +++++++++++++++++--
 1 file changed, 62 insertions(+), 4 deletions(-)

diff --git a/crates/computer-use-mcp/src/backend/macos/ax.rs b/crates/computer-use-mcp/src/backend/macos/ax.rs
index 5fe5adeac..e851affc7 100644
--- a/crates/computer-use-mcp/src/backend/macos/ax.rs
+++ b/crates/computer-use-mcp/src/backend/macos/ax.rs
@@ -7,7 +7,9 @@ use std::sync::{Mutex, OnceLock};
 use std::time::Duration;
 
 use core_foundation::array::{CFArrayGetCount, CFArrayGetValueAtIndex, CFArrayRef};
-use core_foundation::base::{CFGetTypeID, CFRelease, CFRetain, CFTypeID, CFTypeRef, TCFType};
+use core_foundation::base::{
+    CFEqual, CFGetTypeID, CFHash, CFRelease, CFRetain, CFTypeID, CFTypeRef, TCFType,
+};
 use core_foundation::boolean::CFBoolean;
 use core_foundation::number::CFNumber;
 use core_foundation::runloop::{
@@ -610,7 +612,7 @@ pub(super) fn observe_tree(root: &RootInfo) -> Result<UiNode, BackendError> {
     let window = locate_window_in_application(application.as_ax(), root)?;
     let mut context = WalkContext {
         count: 0,
-        visited: HashSet::new(),
+        visited: VisitedElements::default(),
         root_frame: root.frame,
     };
     let mut tree = walk_element(window.as_ax(), 0, &mut context).ok_or_else(|| {
@@ -732,16 +734,47 @@ fn window_match_score(window: AXUIElementRef, root: &RootInfo) -> f64 {
 
 struct WalkContext {
     count: usize,
-    visited: HashSet<usize>,
+    visited: VisitedElements,
     root_frame: Frame,
 }
 
+/// The elements a walk has entered, compared as AX identities.
+///
+/// Each element is retained for the walk: a child array is released once
+/// its subtree is walked, and the next array the system hands out can reuse
+/// a freed element's address, so an address alone cannot stand for a node.
+#[derive(Default)]
+struct VisitedElements {
+    buckets: HashMap<usize, Vec<OwnedCf>>,
+}
+
+impl VisitedElements {
+    /// Records `element`; `false` when an equal element was already recorded.
+    fn insert(&mut self, element: AXUIElementRef) -> bool {
+        let value = element as CFTypeRef;
+        // SAFETY: the caller holds `element` live for the call.
+        let hash = unsafe { CFHash(value) };
+        let bucket = self.buckets.entry(hash).or_default();
+        if bucket
+            .iter()
+            .any(|seen| unsafe { CFEqual(seen.0, value) } != 0)
+        {
+            return false;
+        }
+        // SAFETY: `element` is live; retaining keeps it so for the walk.
+        if let Some(owned) = unsafe { OwnedCf::from_borrowed(value) } {
+            bucket.push(owned);
+        }
+        true
+    }
+}
+
 fn walk_element(
     element: AXUIElementRef,
     depth: usize,
     context: &mut WalkContext,
 ) -> Option<UiNode> {
-    if context.count >= MAX_NODES || !context.visited.insert(element as usize) {
+    if context.count >= MAX_NODES || !context.visited.insert(element) {
         return None;
     }
     context.count += 1;
@@ -1031,6 +1064,31 @@ fn ax_error_name(code: AXError) -> &'static str {
 mod tests {
     use super::*;
 
+    /// The system hands out a fresh object for every element reference, and
+    /// reuses a released one's address: a walk that told nodes apart by
+    /// address dropped the first child of a group whose address a freed
+    /// sibling's child had used.
+    #[test]
+    fn visited_elements_are_told_apart_by_identity_not_address() {
+        let element = |pid: i32| {
+            // SAFETY: AXUIElementCreateApplication follows the create rule and
+            // needs no permission to make a reference.
+            unsafe { OwnedCf::from_create(AXUIElementCreateApplication(pid) as CFTypeRef) }
+                .expect("an application element")
+        };
+        let mut visited = VisitedElements::default();
+        let own = std::process::id() as i32;
+        assert!(visited.insert(element(own).as_ax()));
+        assert!(
+            !visited.insert(element(own).as_ax()),
+            "a second reference to the same element is the same node"
+        );
+        for pid in 1..=256 {
+            let other = element(pid);
+            assert!(visited.insert(other.as_ax()), "pid {pid} is a new element");
+        }
+    }
+
     #[test]
     fn chromium_bundle_detection_is_case_insensitive_and_narrow() {
         for bundle_id in [