Skip to content

fix(get_window_state): 30s timeout + 2000-node cap for heavy webview apps #1609

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
hippoley wants to merge 2 commits into from
Closed

fix(get_window_state): 30s timeout + 2000-node cap for heavy webview apps #1609

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8d3d3e4
fix(get_window_state): 30s timeout + 2000-node cap for heavy webview …
May 20, 2026
3f06aea
fix(get_window_state): explicit truncation flag + remove misleading q…
May 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 40 additions & 7 deletions libs/cua-driver-rs/crates/platform-macos/src/ax/tree.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,13 @@ use core_foundation::base::{CFRelease, CFRetain, CFTypeRef};
/// pathological trees (mirrors Swift reference implementation).
const MAX_DEPTH: usize = 25;

/// Maximum total nodes visited during a single AX walk. Chromium-family apps
/// (Arc, VS Code, Chrome) can expose thousands of nodes; capping at 2 000
/// keeps the walk bounded while still covering realistic app chrome.
/// When the cap is hit the walk stops early and the partial tree is returned
/// with a warning line appended (mirrors Swift reference implementation).
const MAX_ELEMENTS: usize = 2_000;

/// A single node in the AX tree.
#[derive(Debug, Clone)]
pub struct AXNode {
Expand All @@ -42,6 +49,8 @@ pub struct AXNode {
pub struct TreeWalkResult {
pub tree_markdown: String,
pub nodes: Vec<AXNode>,
/// True when the walk was cut short by the MAX_ELEMENTS cap.
pub truncated: bool,
}

/// Walk the AX tree of `pid`, optionally filtered to a specific window.
Expand All @@ -62,11 +71,16 @@ pub fn walk_tree(pid: i32, window_id: Option<u32>, query: Option<&str>) -> TreeW
let mut nodes: Vec<AXNode> = Vec::new();
let mut lines: Vec<(usize, String)> = Vec::new(); // (depth, line)
let mut index_counter = 0usize;
// Shared visited-node counter passed into walk_element to enforce MAX_ELEMENTS.
let mut visited_count = 0usize;
// Set to true only when walk_element actually stops early due to the cap —
// avoids a false-positive when the tree naturally ends on exactly MAX_ELEMENTS.
let mut truncated = false;

unsafe {
let app_elem = AXUIElementCreateApplication(pid);
if app_elem.is_null() {
return TreeWalkResult { tree_markdown: String::new(), nodes };
return TreeWalkResult { tree_markdown: String::new(), nodes, truncated: false };
}

// Union AXChildren + AXWindows — the only way to see background windows.
Expand Down Expand Up @@ -102,7 +116,7 @@ pub fn walk_tree(pid: i32, window_id: Option<u32>, query: Option<&str>) -> TreeW

// Walk each top-level child at depth 0.
for child in walk_these {
walk_element(child, 0, &mut nodes, &mut lines, &mut index_counter);
walk_element(child, 0, &mut nodes, &mut lines, &mut index_counter, &mut visited_count, &mut truncated);
}

// Release all top-level elements (copy_children / copy_ax_windows both retain).
Expand All @@ -113,14 +127,24 @@ pub fn walk_tree(pid: i32, window_id: Option<u32>, query: Option<&str>) -> TreeW
CFRelease(app_elem as CFTypeRef);
}

let truncated_flag = truncated;
let raw_markdown = render_lines(&lines);
let tree_markdown = if let Some(q) = query {
let mut tree_markdown = if let Some(q) = query {
filter_tree(&raw_markdown, q)
} else {
raw_markdown
};

TreeWalkResult { tree_markdown, nodes }
if truncated_flag {
tree_markdown.push_str(&format!(
"\n⚠️ AX tree truncated at {MAX_ELEMENTS} nodes \
(app has a very large accessibility tree — Arc, Electron, or similar). \
Element indices above are still valid. Use pixel clicks for elements \
not visible in this partial tree."
));
}

TreeWalkResult { tree_markdown, nodes, truncated: truncated_flag }
}

unsafe fn walk_element(
Expand All @@ -129,8 +153,17 @@ unsafe fn walk_element(
nodes: &mut Vec<AXNode>,
lines: &mut Vec<(usize, String)>,
counter: &mut usize,
visited_count: &mut usize,
truncated: &mut bool,
) {
if depth > MAX_DEPTH { return; }
// Enforce total-node cap — mirrors Swift's maxElements guard.
// Set the truncated flag only when we actually stop early.
if *visited_count >= MAX_ELEMENTS {
*truncated = true;
return;
}
*visited_count += 1;

let role = copy_string_attr(element, "AXRole")
.unwrap_or_else(|| "AXUnknown".into());
Expand All @@ -140,7 +173,7 @@ unsafe fn walk_element(
// Still recurse — children may be interesting.
let children = copy_children(element);
for child in children {
walk_element(child, depth, nodes, lines, counter);
walk_element(child, depth, nodes, lines, counter, visited_count, truncated);
CFRelease(child as CFTypeRef);
}
return;
Expand Down Expand Up @@ -173,7 +206,7 @@ unsafe fn walk_element(
if !is_actionable && !has_content && role != "AXWindow" && role != "AXSheet" {
let children = copy_children(element);
for child in children {
walk_element(child, depth + 1, nodes, lines, counter);
walk_element(child, depth + 1, nodes, lines, counter, visited_count, truncated);
CFRelease(child as CFTypeRef);
}
return;
Expand Down Expand Up @@ -217,7 +250,7 @@ unsafe fn walk_element(

let children = copy_children(element);
for child in children {
walk_element(child, depth + 1, nodes, lines, counter);
walk_element(child, depth + 1, nodes, lines, counter, visited_count, truncated);
CFRelease(child as CFTypeRef);
}
}
Expand Down
23 changes: 18 additions & 5 deletions libs/cua-driver-rs/crates/platform-macos/src/tools/get_window_state.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,12 +83,25 @@ impl Tool for GetWindowStateTool {
let capture_mode = if capture_mode == "tree" { "ax" } else { capture_mode };
let tree_result = if capture_mode != "vision" {
let q = query.clone();
let result = tokio::task::spawn_blocking(move || {
// Wrap the blocking AX walk in a 30-second timeout. Heavy webview apps
// (Arc, Safari with many tabs, Electron) can block
// AXUIElementCopyAttributeValue indefinitely via XPC — without a
// deadline the MCP server hangs forever (issue #1537).
let walk_future = tokio::task::spawn_blocking(move || {
crate::ax::tree::walk_tree(pid, Some(window_id), q.as_deref())
}).await;
match result {
Ok(r) => Some(r),
Err(e) => return ToolResult::error(format!("AX tree walk failed: {e}")),
});
match tokio::time::timeout(std::time::Duration::from_secs(30), walk_future).await {
Ok(Ok(r)) => Some(r),
Ok(Err(e)) => return ToolResult::error(format!("AX tree walk failed: {e}")),
Err(_elapsed) => {
return ToolResult::error(format!(
"AX tree walk for pid={pid} timed out after 30 s. \
The app (likely Arc, Electron, or Safari with many tabs) has a \
pathologically large accessibility tree. \
Workarounds: switch to capture_mode=vision for pixel-click \
workflows, or use capture_mode=ax with a depth-limited scan."
));
}
}
Comment on lines +90 to 105

@coderabbitai coderabbitai Bot May 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | 🏗️ Heavy lift

🧩 Analysis chain

🌐 Web query:

In Tokio, if tokio::time::timeoutexpires while awaiting aspawn_blocking JoinHandle, does the underlying blocking task stop, or does it continue running after the handle is dropped?

💡 Result:

If tokio::time::timeout expires while you’re awaiting a JoinHandle returned by tokio::task::spawn_blocking, the underlying blocking closure does not get stopped; it continues running in the background after the JoinHandle is dropped. Tokio’s docs state that spawn_blocking closures “cannot be aborted because they are not async” and that even calling abort() on a spawn_blocking task “will not have any effect” (once it has started). It also explains that when a Tokio runtime is shut down it “will attempt to abort all tasks including spawn_blocking tasks,” but blocking tasks “cannot be aborted once they start running,” and shutdown may still wait for them to finish. [1][2] Separately, Tokio’s JoinHandle behavior is that dropping the handle “detaches the associated task,” and “the task continues running in the background and its return value is lost.” [3][4] So, in the specific scenario you described: - timeout(...).await timing out typically means the await is dropped (i.e., the JoinHandle variable goes out of scope / is dropped). - Dropping the JoinHandle detaches the task, so the blocking closure keeps running. - Additionally, even if you were to try to abort the handle, spawn_blocking tasks can’t be aborted once they’ve started. [1][3]

Citations:

Timeout does not stop the blocking task—the AX walk continues running in the background.

When tokio::time::timeout expires on the spawn_blocking JoinHandle, the underlying blocking closure is not aborted; it continues consuming a blocking thread in the pool. Repeated timeouts can accumulate idle threads and starve the blocking pool for future work.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@libs/cua-driver-rs/crates/platform-macos/src/tools/get_window_state.rs`
around lines 90 - 106, The timeout currently only drops the JoinHandle but does
not stop the blocking AX walk (spawn_blocking + crate::ax::tree::walk_tree), so
repeated timeouts leak blocking threads; modify the approach so the blocking
work can be cooperatively cancelled: change walk_tree to accept a cancellation
token/flag (e.g., tokio_util::sync::CancellationToken or an Arc<AtomicBool>) and
have walk_tree periodically check it and return early, then when
tokio::time::timeout fires set/trigger that token before returning the timeout
ToolResult; update the call site (the walk_future spawn_blocking closure and its
caller) to create and pass the token, signal it on Err(_elapsed), and
await/handle the task completion or a bounded grace period to ensure the
blocking thread finishes cleanly.

} else {
None
Expand Down
103 changes: 80 additions & 23 deletions libs/cua-driver/Sources/CuaDriverCore/AppState/AppState.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,7 @@ public enum AppStateError: Error, CustomStringConvertible, Sendable {
case appNotFound(Int32)
case noCachedState(pid: Int32, windowId: UInt32)
case invalidElementIndex(pid: Int32, windowId: UInt32, index: Int)
case axWalkTimedOut(pid: Int32)

public var description: String {
switch self {
Expand All @@ -167,6 +168,13 @@ public enum AppStateError: Error, CustomStringConvertible, Sendable {
+ "out of range in the cached snapshot. Re-run "
+ "`get_window_state({pid: \(pid), window_id: \(windowId)})` and use "
+ "an index from the fresh tree."
case .axWalkTimedOut(let pid):
return "AX tree walk for pid \(pid) timed out after 30 s. "
+ "The app (likely Arc, Electron, or Safari with many tabs) has a "
+ "pathologically large accessibility tree. "
+ "Workarounds: switch to `capture_mode: vision` for pixel-click "
+ "workflows, or use `capture_mode: ax` with a `query` filter to "
+ "limit the walk scope."
}
}
}
Expand All @@ -193,6 +201,13 @@ public actor AppStateEngine {
/// 25 covers realistic app chrome without exploding on pathological trees.
public static let maxDepth = 25

/// Maximum total elements visited during a single AX walk. Chromium-family
/// apps (Arc, VS Code, Chrome) can expose thousands of nodes; capping at
/// 2 000 keeps the walk bounded while still covering realistic app chrome.
/// When the cap is hit the walk stops and the partial tree is returned with
/// a warning appended to the markdown.
public static let maxElements = 2_000

/// Shared assertion that tracks which pids accept AXManualAccessibility /
/// AXEnhancedUserInterface. Extracted into its own actor so a call-site-
/// level ``FocusGuard`` can consult and update the same caches.
Expand Down Expand Up @@ -252,34 +267,65 @@ public actor AppStateEngine {
throw AppStateError.appNotFound(pid)
}

let root = AXUIElementCreateApplication(pid)

// Cue Chromium/Electron apps to turn on their web accessibility tree.
// Non-Chromium apps ignore these attribute writes — safe no-op.
try await activateAccessibilityIfNeeded(pid: pid, root: root)
// Must be called before the task group so it runs on the actor.
try await activateAccessibilityIfNeeded(pid: pid, root: AXUIElementCreateApplication(pid))

nextTurnId += 1
let turnId = nextTurnId

var elements: [Int: AXUIElement] = [:]
var nextIndex = 0
var markdown = ""
// Wrap the AX walk in a 30-second timeout. Heavy webview apps (Arc,
// Safari with many tabs, Electron) can block AXUIElementCopyAttributeValue
// indefinitely via XPC. The walk runs on a detached task so the timeout
// can cancel it; the main snapshot task throws axWalkTimedOut if the
// deadline is exceeded.
let (snapshotMarkdown, snapshotElements, snapshotDidTruncate) = try await withThrowingTaskGroup(
of: (String, [Int: AXUIElement], Bool).self
) { group in
group.addTask {
let root = AXUIElementCreateApplication(pid)
var els: [Int: AXUIElement] = [:]
var idx = 0
var md = ""
var visited = 0
var didTruncate = false
let layer0WindowIds: Set<CGWindowID> = Set(
WindowEnumerator.allWindows()
.filter { $0.layer == 0 }
.map { CGWindowID($0.id) }
)
self.renderTree(
root,
depth: 0,
targetWindowId: windowId,
layer0WindowIds: layer0WindowIds,
elements: &els,
nextIndex: &idx,
output: &md,
visitedCount: &visited,
didTruncate: &didTruncate
)
return (md, els, didTruncate)
}
group.addTask {
try await Task.sleep(nanoseconds: 30_000_000_000)
throw AppStateError.axWalkTimedOut(pid: pid)
}
let result = try await group.next()!
group.cancelAll()
return result
}

let layer0WindowIds: Set<CGWindowID> = Set(
WindowEnumerator.allWindows()
.filter { $0.layer == 0 }
.map { CGWindowID($0.id) }
)
var markdown = snapshotMarkdown
let elements = snapshotElements

renderTree(
root,
depth: 0,
targetWindowId: windowId,
layer0WindowIds: layer0WindowIds,
elements: &elements,
nextIndex: &nextIndex,
output: &markdown
)
if snapshotDidTruncate {
markdown += "\n⚠️ AX tree truncated at \(AppStateEngine.maxElements) nodes"
+ " (app has a very large accessibility tree — Arc, Electron, or similar)."
+ " Element indices above are still valid. Use pixel clicks for elements"
+ " not visible in this partial tree."
}

sessions[SessionKey(pid: pid, windowId: windowId)] =
SessionState(turnId: turnId, elements: elements)
Expand Down Expand Up @@ -513,16 +559,25 @@ public actor AppStateEngine {

// MARK: - Walking

private func renderTree(
private nonisolated func renderTree(
_ element: AXUIElement,
depth: Int,
targetWindowId: UInt32?,
layer0WindowIds: Set<CGWindowID> = [],
elements: inout [Int: AXUIElement],
nextIndex: inout Int,
output: inout String
output: inout String,
visitedCount: inout Int,
didTruncate: inout Bool
) {
guard depth <= AppStateEngine.maxDepth else { return }
guard visitedCount < AppStateEngine.maxElements else {
// Set the flag only when we actually stop early due to the cap —
// avoids a false-positive when the tree naturally ends on exactly maxElements.
didTruncate = true
return
}
visitedCount += 1

let role = attributeString(element, "AXRole") ?? "?"
let title = attributeString(element, "AXTitle")
Expand Down Expand Up @@ -616,7 +671,9 @@ public actor AppStateEngine {
layer0WindowIds: layer0WindowIds,
elements: &elements,
nextIndex: &nextIndex,
output: &output
output: &output,
visitedCount: &visitedCount,
didTruncate: &didTruncate
)
}
}
Expand Down