Bound Snapshot/WaitFor tree capture size to prevent hangs and oversized output - #353
Conversation
…ed output Tree.tree_traversal walked the full UIA subtree of every visible window with no cap on element count. Against a window that exposes a large flat list or grid (e.g. an unfiltered inventory view with thousands of rows), this could: - stall UI Automation on the target app for minutes while the traversal fetched cached properties for every row, and - produce a serialized response far too large for the calling MCP client to accept, even once the capture itself finished. Add a TreeElementBudget (tree/budget.py) that Tree.get_state resets per capture and threads through tree_traversal: once WINDOWS_MCP_MAX_TREE_ELEMENTS (default 500) output-affecting elements have been captured, the traversal stops descending into further children instead of visiting the rest of the subtree, get_window_wise_nodes skips any remaining windows, and the IA2 (Firefox) fallback path truncates its already-collected result to the remaining budget. TreeState gains `truncated`/`element_limit` fields, and the semantic tree / interactive / scrollable renderers append a note when the capture was cut short, so the truncation is visible in the tool response rather than silent. WaitFor polls the same Tree.get_state path, so it's bounded by the same budget. Since WaitFor/Snapshot are the same capture path, this doesn't require any tool-facing API change — existing calls just get a bounded, clearly-marked partial result instead of a hang or an oversized response.
PR Summary by QodoBound Snapshot/WaitFor tree capture size to prevent hangs and oversized output
AI Description
Diagram
High-Level Assessment
Files changed (8)
|
Code Review by Qodo
Context used✅ Compliance rules (platform):
15 rules 1.
|
|
Can you pls share the name of the apps were the stuck has felt |
Two issues flagged by review on the tree-capture-budget PR: - dom_informative_nodes.append(TextElementNode(...)) in the browser-DOM path never consumed the element budget, unlike the interactive-node branches right above it. Text-heavy pages could accumulate informative nodes without ever tripping truncation. - TreeElementBudget.try_consume(amount) let count overshoot limit when a single batched call exceeded remaining capacity, while still returning True — contradicting its own "hard cap" contract. Latent today (the one batched call site already pre-slices to remaining), but risky as public API. Also folds in the review's smaller nits: missing -> None on TreeElementBudget.__init__, Google-style docstring on resolve_max_tree_elements, and a quoting fix in test_tree_views.py.
|
The stall happened in TriData, a Windows desktop ERP application (Win32 UI Automation, not a browser) — specifically on an unfiltered inventory grid (~9,900 rows). While addressing the automated review comments, I also fixed two related bugs so the budget holds up better against exactly this kind of case:
Also picked up the smaller nits (missing |
|
Thanks |
Problem
Tree.tree_traversal(tree/service.py) walks the full UI Automation subtree of every visible window with no cap on element count. Against a window that exposes a large flat list or grid (e.g. an unfiltered inventory/table view with thousands of rows), this can:Observed in practice against a ~9,860-row list view: one
Snapshotcall blocked the target app for 30+ minutes; another returned a 54k+ character response that the calling client rejected outright.WaitForpolls the sameTree.get_statecapture path (via_iter_nodes/_iter_text_sources), so it's affected identically.Fix
tree/budget.py: a small, dependency-freeTreeElementBudgetthat tracks how many output-affecting elements a capture has collected, plusresolve_max_tree_elements()readingWINDOWS_MCP_MAX_TREE_ELEMENTS(default500).Tree.get_stateresets a fresh budget per capture.tree_traversalchecks the budget before recursing into each child and stops descending once it's exhausted — this is what bounds traversal time, not just the size of the appended lists.get_window_wise_nodesskips any remaining windows once the budget is spent.TreeStategainstruncated: boolandelement_limit: int;semantic_tree_to_string(),interactive_elements_to_string(), andscrollable_elements_to_string()append a clear note when the capture was cut short, so truncation is visible in the tool response instead of silent.desktop/service.py's_filter_tree_state_to_region(used when adisplayregion is requested) passes the new fields through so the flag survives region filtering.No tool-facing API change — existing
Snapshot/WaitForcalls just get a bounded, clearly-marked partial result instead of a hang or an oversized response. Raising the limit (or disabling it in practice by setting a very high value) is available via env var for anyone who wants the old unbounded behavior.Testing
tests/test_tree_budget.py(new): unit tests forTreeElementBudget/resolve_max_tree_elements— pure Python, no UIA dependency.tests/test_tree_views.py: new cases for the truncation note on all three renderers (present whentruncated=True, absent otherwise, omitted when the specific node list is empty).tests/test_tree_service.py: new cases exercising the wiring —tree_traversalstops appending/recursing once the budget is exhausted and leaves it untouched when under budget;get_window_wise_nodesskips remaining windows once exhausted.tests/test_tree_budget.pyandtests/test_tree_views.pypass (36 tests) andruff checkis clean on all touched files. I don't have a Windows box handy, so I could not runtests/test_tree_service.pylocally (it importscomtypes/pywin32, matching this repo's CI which runs onwindows-latest) — those new tests follow the exactMagicMock-based pattern already used in that file and I traced the logic by hand against the traversal code, but please let the CI run confirm.Config
Documented the new
WINDOWS_MCP_MAX_TREE_ELEMENTSenv var inCLAUDE.md's environment variable table, consistent with the existing entries there.