Skip to content

Bound Snapshot/WaitFor tree capture size to prevent hangs and oversized output - #353

Merged
Jeomon merged 2 commits into
CursorTouch:mainfrom
Veloflott:bound-tree-capture-size
Jul 29, 2026
Merged

Bound Snapshot/WaitFor tree capture size to prevent hangs and oversized output#353
Jeomon merged 2 commits into
CursorTouch:mainfrom
Veloflott:bound-tree-capture-size

Conversation

@Veloflott

Copy link
Copy Markdown
Contributor

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:

  • stall UI Automation on the target application for minutes while the traversal fetches 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 finishes.

Observed in practice against a ~9,860-row list view: one Snapshot call blocked the target app for 30+ minutes; another returned a 54k+ character response that the calling client rejected outright.

WaitFor polls the same Tree.get_state capture path (via _iter_nodes/_iter_text_sources), so it's affected identically.

Fix

  • New tree/budget.py: a small, dependency-free TreeElementBudget that tracks how many output-affecting elements a capture has collected, plus resolve_max_tree_elements() reading WINDOWS_MCP_MAX_TREE_ELEMENTS (default 500).
  • Tree.get_state resets a fresh budget per capture.
  • tree_traversal checks 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_nodes skips any remaining windows once the budget is spent.
  • The IA2/Firefox fallback path (which returns its full result in one call, so it can't be bounded incrementally) truncates its already-collected result to the remaining budget.
  • TreeState gains truncated: bool and element_limit: int; semantic_tree_to_string(), interactive_elements_to_string(), and scrollable_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 a display region is requested) passes the new fields through so the flag survives region filtering.

No tool-facing API change — existing Snapshot/WaitFor calls 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 for TreeElementBudget/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 when truncated=True, absent otherwise, omitted when the specific node list is empty).
  • tests/test_tree_service.py: new cases exercising the wiring — tree_traversal stops appending/recursing once the budget is exhausted and leaves it untouched when under budget; get_window_wise_nodes skips remaining windows once exhausted.
  • Ran locally: tests/test_tree_budget.py and tests/test_tree_views.py pass (36 tests) and ruff check is clean on all touched files. I don't have a Windows box handy, so I could not run tests/test_tree_service.py locally (it imports comtypes/pywin32, matching this repo's CI which runs on windows-latest) — those new tests follow the exact MagicMock-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_ELEMENTS env var in CLAUDE.md's environment variable table, consistent with the existing entries there.

…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.
@qodo-code-review

Copy link
Copy Markdown

PR Summary by Qodo

Bound Snapshot/WaitFor tree capture size to prevent hangs and oversized output

🐞 Bug fix ✨ Enhancement 🧪 Tests ⚙️ Configuration changes 🕐 40+ Minutes

Grey Divider

AI Description

• Add TreeElementBudget (tree/budget.py) to cap element count per Snapshot/WaitFor tree capture,
 avoiding minutes-long stalls and oversized responses on huge grids/lists.
• tree_traversal stops descending into children once the budget is exhausted;
 get_window_wise_nodes skips remaining windows; the IA2/Firefox fallback path truncates its
 already-collected result.
• TreeState gains truncated/element_limit fields, and renderers append a clear truncation note
 instead of silently dropping data.
• Limit is configurable via new WINDOWS_MCP_MAX_TREE_ELEMENTS env var (default 500), documented in
 CLAUDE.md.
• Add unit tests for the budget module and truncation behavior in tree service/views.
Diagram

graph TD
    A["Tree.get_state()"] --> B["TreeElementBudget"] --> C["tree_traversal()"]
    C -->|"consume per node"| B
    C -->|"stop descending when exhausted"| D["get_window_wise_nodes()"]
    D -->|"skip remaining windows"| E["IA2/Firefox fallback"]
    E -->|"cap collected result"| F["TreeState truncated/element_limit"]
    F --> G["_filter_tree_state_to_region()"]
    F --> H["semantic/interactive/scrollable renderers"]
    H -->|"append truncation note"| I["MCP tool response"]
Loading
High-Level Assessment

A soft element-count budget checked during descent is the right lightweight fix: it directly targets the reported failure mode (unbounded traversal time and response size) without requiring UIA-side filtering APIs or a rewrite of the traversal algorithm. Alternatives like a wall-clock timeout would bound time but not response size, and pre-filtering by control type would risk silently dropping relevant elements without a visible truncation signal. The chosen approach is minimal, testable in pure Python, configurable via env var, and preserves backward compatibility by defaulting to a generous limit.

Files changed (8) +401 / -6

Enhancement (2) +100 / -3
budget.pyNew TreeElementBudget and resolve_max_tree_elements helpers +77/-0

New TreeElementBudget and resolve_max_tree_elements helpers

• New dependency-free module providing 'TreeElementBudget' (tracks consumed/remaining/exhausted state) and 'resolve_max_tree_elements()' which reads 'WINDOWS_MCP_MAX_TREE_ELEMENTS' (default 500) with validation and fallback on invalid/blank/non-positive values.

src/windows_mcp/tree/budget.py

views.pyAdd truncation note and TreeState fields +23/-3

Add truncation note and TreeState fields

• Adds 'truncated: bool' and 'element_limit: int' fields to 'TreeState', plus a '_truncation_note()' helper appended by 'semantic_tree_to_string()', 'interactive_elements_to_string()', and 'scrollable_elements_to_string()' when the capture was cut short.

src/windows_mcp/tree/views.py

Bug fix (2) +48 / -3
service.pyPropagate truncated/element_limit through region filtering +2/-0

Propagate truncated/element_limit through region filtering

• '_filter_tree_state_to_region' now passes 'truncated' and 'element_limit' through when constructing the filtered 'TreeState', so the truncation flag survives display-region filtering.

src/windows_mcp/desktop/service.py

service.pyWire element budget into tree traversal, window loop, and IA2 fallback +46/-3

Wire element budget into tree traversal, window loop, and IA2 fallback

• 'Tree' now resets a fresh 'TreeElementBudget' per 'get_state' call; 'tree_traversal' consumes budget for each appended node and stops descending into children once exhausted; 'get_window_wise_nodes' breaks out of the window loop once the budget is spent; the IA2/Firefox fallback path caps its already-collected interactive/informative nodes to the remaining budget; 'TreeState' construction and profiling logs now include truncation info.

src/windows_mcp/tree/service.py

Tests (3) +252 / -0
test_tree_budget.pyNew unit tests for TreeElementBudget and resolve_max_tree_elements +77/-0

New unit tests for TreeElementBudget and resolve_max_tree_elements

• Pure-Python tests covering budget consumption, exhaustion/truncation flags, limit clamping, and env-var parsing (missing, blank, valid, invalid, zero, negative).

tests/test_tree_budget.py

test_tree_service.pyAdd tests exercising budget wiring in tree_traversal/get_window_wise_nodes +132/-0

Add tests exercising budget wiring in tree_traversal/get_window_wise_nodes

• New test class verifies that traversal stops appending/recursing once the budget is exhausted, that all elements are captured when under budget, and that 'get_window_wise_nodes' skips remaining windows once the budget is spent.

tests/test_tree_service.py

test_tree_views.pyAdd tests for truncation note rendering +43/-0

Add tests for truncation note rendering

• New tests verify the truncation note is appended to interactive/scrollable/semantic tree string renderers when 'truncated=True', omitted when 'False', and still shown even when the specific node list is empty.

tests/test_tree_views.py

Documentation (1) +1 / -0
CLAUDE.mdDocument WINDOWS_MCP_MAX_TREE_ELEMENTS env var +1/-0

Document WINDOWS_MCP_MAX_TREE_ELEMENTS env var

• Adds a row to the environment variable reference table describing the new tree capture element limit, its default, and where it's resolved.

CLAUDE.md

@qodo-code-review

qodo-code-review Bot commented Jul 27, 2026

Copy link
Copy Markdown

Code Review by Qodo

🐞 Bugs (0) 📘 Rule violations (0) 📜 Skill insights (0)

Context used
✅ Compliance rules (platform): 15 rules

Grey Divider


Action required

1. DOM text nodes unbudgeted ✓ Resolved 🐞 Bug ☼ Reliability
Description
The new traversal stop condition depends on element_budget.exhausted, but dom_informative_nodes
(TextElementNode) are appended without consuming budget. Text-heavy DOM captures can therefore still
accumulate and be serialized/iterated unboundedly even when the element budget is low.
Code

src/windows_mcp/tree/service.py[R680-684]

+                if self.element_budget.exhausted:
+                    # Stop descending once the element budget is spent — this is what
+                    # bounds traversal time on huge flat lists/grids (thousands of rows),
+                    # not just the size of the appended node lists.
+                    break
Relevance

⭐⭐ Medium

Potential reliability issue but no prior budget/text-node precedent found; may be considered
out-of-scope for this PR.

ⓘ Recommendations generated based on similar findings in past PRs

Evidence
The traversal cut-off only checks element_budget.exhausted, but the informative-node append path
never calls try_consume(), so budget exhaustion may never occur on text-heavy DOMs. Those
unbounded informative nodes are explicitly consumed by WaitFor and Scrape, increasing runtime and
response size.

src/windows_mcp/tree/service.py[678-685]
src/windows_mcp/tree/service.py[629-656]
src/windows_mcp/tools/input.py[92-121]
src/windows_mcp/tools/scrape.py[35-46]

Agent prompt
The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

## Issue description
Traversal now stops descending when `element_budget.exhausted`, but the budget is only consumed for interactive/scrollable nodes. The DOM informative text path (`dom_informative_nodes.append(TextElementNode(...))`) does not consume from the budget, so on text-heavy pages (many `TextControl` nodes), the traversal may keep walking and the output can still balloon.

## Issue Context
Downstream tools iterate and/or join `dom_informative_nodes`:
- `WaitFor` includes them in `_iter_text_sources()`
- `Scrape(use_dom=True)` concatenates all `node.text`
So an unbounded `dom_informative_nodes` defeats the PR’s goal of bounding traversal cost and response size.

## Fix Focus Areas
- src/windows_mcp/tree/service.py[678-685]
- src/windows_mcp/tree/service.py[629-656]
- src/windows_mcp/tools/input.py[92-121]
- src/windows_mcp/tools/scrape.py[35-46]

## Suggested fix
- Treat informative text nodes as budget-consuming output elements:
 - Before appending a `TextElementNode`, call `if not self.element_budget.try_consume(): return/skip` (or conditionally stop descending).
- Ensure the traversal exits promptly when the budget becomes exhausted in the informative-text-heavy case.
- Add/adjust a unit test to construct a DOM-like tree with many informative nodes and few/no interactive nodes, and assert traversal stops and output is truncated.

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools



Remediation recommended

2. Single-quoted string in test ✓ Resolved 📘 Rule violation ✧ Quality
Description
A newly added string literal uses single quotes instead of double quotes. This violates the
project’s enforced string-literal quoting convention.
Code

tests/test_tree_views.py[101]

+        assert 'window "Notepad"' in result
Relevance

⭐⭐⭐ High

Close precedent: reviewers accepted switching single-quoted assertions to double quotes to match
quoting rule.

PR-#307

ⓘ Recommendations generated based on similar findings in past PRs

Evidence
The checklist requires double quotes for string literals. The assertion at
tests/test_tree_views.py:101 uses a single-quoted string literal.

Rule 222799: Enforce double quotes for all string literals
tests/test_tree_views.py[101-101]

Agent prompt
The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

## Issue description
A newly added string literal uses single quotes.

## Issue Context
Compliance requires double quotes for all string literals where possible.

## Fix Focus Areas
- tests/test_tree_views.py[101-101]

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools


3. TreeElementBudget.__init__ missing return hint ✓ Resolved 📘 Rule violation ✧ Quality
Description
TreeElementBudget.__init__ is missing an explicit return type annotation (-> None). This
violates the requirement that all function signatures (including constructors) must have return type
hints.
Code

src/windows_mcp/tree/budget.py[R50-53]

+    def __init__(self, limit: int = DEFAULT_MAX_TREE_ELEMENTS):
+        self.limit = max(1, limit)
+        self.count = 0
+        self.truncated = False
Relevance

⭐⭐⭐ High

Team previously accepted adding explicit return/type annotations (including -> None) to satisfy
type-hint rules.

PR-#307

ⓘ Recommendations generated based on similar findings in past PRs

Evidence
The type-hints rule requires explicit return annotations for all functions, including __init__.
The new __init__ definition omits -> None.

Rule 222805: Require type hints for all function signatures
src/windows_mcp/tree/budget.py[50-53]

Agent prompt
The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

## Issue description
`TreeElementBudget.__init__` lacks an explicit return type annotation.

## Issue Context
Compliance requires type hints for all function signatures; constructors should explicitly declare `-> None`.

## Fix Focus Areas
- src/windows_mcp/tree/budget.py[50-53]

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools


4. Budget can exceed limit ✓ Resolved 🐞 Bug ≡ Correctness
Description
TreeElementBudget.try_consume() can increase count beyond limit and still returns True, violating
the intended “max elements collected” invariant and making batch consumers unsafe. This is codified
by a new unit test, so future batch use can silently overshoot the configured cap and inflate
logs/output accounting.
Code

src/windows_mcp/tree/budget.py[R63-77]

+    def try_consume(self, amount: int = 1) -> bool:
+        """Register `amount` newly captured elements.
+
+        Returns False (and marks the budget truncated) once the limit is
+        reached; the caller should then stop appending/recursing further.
+        """
+        if amount <= 0:
+            return not self.exhausted
+        if self.exhausted:
+            self.truncated = True
+            return False
+        self.count += amount
+        if self.exhausted:
+            self.truncated = True
+        return True
Relevance

⭐⭐ Medium

Semantic behavior change (hard-cap vs overshoot) with no historical precedent; PR’s tests suggest
intent is unclear.

ⓘ Recommendations generated based on similar findings in past PRs

Evidence
The implementation allows overshooting the limit in a single call (no clamp) and still returns
success; the added test asserts this overshoot behavior, confirming it’s intentional in the PR but
contradicts the stated purpose of a hard cap.

src/windows_mcp/tree/budget.py[63-77]
tests/test_tree_budget.py[35-40]

Agent prompt
The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

## Issue description
`TreeElementBudget.try_consume()` currently increments `count` by the full requested `amount` even if that exceeds the remaining budget, and it returns `True` in that case. This violates the class’s own contract (“Returns False once the limit is reached”) and undermines the core invariant that the budget represents a hard cap.

## Issue Context
Today most production call sites consume one-at-a-time, and the IA2 batch path pre-slices to `remaining`, so the bug is latent. But the abstraction is now part of the public internal API for capture bounding, and tests explicitly lock in the overshoot behavior.

## Fix Focus Areas
- src/windows_mcp/tree/budget.py[63-77]
- tests/test_tree_budget.py[35-40]

## Suggested fix
- Make `try_consume(amount)` strict:
 - If `amount <= 0`: return `not exhausted` (optionally set `truncated=True` if exhausted).
 - If `amount > remaining`: set `count = limit`, set `truncated = True`, and return `False`.
 - Else: increment `count` by `amount`, set `truncated` if `count == limit`, and return `True`.
- Update the unit test to reflect that `try_consume(10)` when only 2 remain should return `False` and cap `count` at `limit` (or whatever strict semantics you choose).

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools


View more (2)
5. Over-100 char lines added ✓ Resolved 📘 Rule violation ✧ Quality
Description
Newly added/modified lines exceed 100 characters, which violates the repository line-length
requirement and reduces readability/maintainability. This appears in both documentation and tests,
indicating the style rule is not consistently applied in this PR.
Code

tests/test_tree_budget.py[62]

+        assert resolve_max_tree_elements({"WINDOWS_MCP_MAX_TREE_ELEMENTS": "  "}) == DEFAULT_MAX_TREE_ELEMENTS
Relevance

⭐⭐ Medium

Mixed precedent: line-length wraps accepted in some tests but explicitly rejected in another PR.

PR-#307
PR-#329

ⓘ Recommendations generated based on similar findings in past PRs

Evidence
The checklist mandates a 100-character max line length. The assertion line in
tests/test_tree_budget.py:62 and the new environment-variable table row in CLAUDE.md:64 are both
single lines well over 100 characters.

Rule 222796: Enforce maximum line length of 100 characters
tests/test_tree_budget.py[62-62]
CLAUDE.md[64-64]

Agent prompt
The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

## Issue description
Several newly added/modified lines exceed the 100-character maximum.

## Issue Context
The PR compliance checklist requires all non-comment, non-whitespace lines to be <= 100 characters. Long lines should be wrapped at reasonable boundaries (commas/operators) rather than kept as single long statements.

## Fix Focus Areas
- tests/test_tree_budget.py[62-62]
- CLAUDE.md[64-64]

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools


6. resolve_max_tree_elements docstring not Google ✓ Resolved 📘 Rule violation ✧ Quality
Description
resolve_max_tree_elements is a public function but its docstring is not Google-style (missing
Args:/Returns: sections). This reduces consistency and makes the API contract less clear for
readers and tooling.
Code

src/windows_mcp/tree/budget.py[R22-24]

+def resolve_max_tree_elements(env: "os._Environ[str] | dict[str, str] | None" = None) -> int:
+    """Read WINDOWS_MCP_MAX_TREE_ELEMENTS, falling back to the default on any invalid value."""
+    source = os.environ if env is None else env
Relevance

⭐⭐ Medium

No close repo precedent found for enforcing Google-style docstrings on public functions.

ⓘ Recommendations generated based on similar findings in past PRs

Evidence
The rule requires Google-style docstrings for public functions. resolve_max_tree_elements is
public and has only a one-line docstring without the required structured sections.

Rule 222802: Require Google-style docstrings on public functions and classes
src/windows_mcp/tree/budget.py[22-24]

Agent prompt
The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

## Issue description
The public function `resolve_max_tree_elements` has a docstring that does not follow Google-style formatting (no `Args:` / `Returns:`).

## Issue Context
Compliance requires Google-style docstrings on public functions/classes in changed files.

## Fix Focus Areas
- src/windows_mcp/tree/budget.py[22-44]

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools


Grey Divider

To customize comments, go to the Qodo configuration screen, or learn more in the docs.

Qodo Logo

Comment thread tests/test_tree_budget.py Outdated
Comment thread src/windows_mcp/tree/budget.py
Comment thread src/windows_mcp/tree/budget.py Outdated
Comment thread tests/test_tree_views.py Outdated
Comment thread src/windows_mcp/tree/budget.py
Comment thread src/windows_mcp/tree/service.py
@Jeomon

Jeomon commented Jul 28, 2026

Copy link
Copy Markdown
Member

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.
@Veloflott

Copy link
Copy Markdown
Contributor Author

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). Snapshot on that window took 607s and the target app itself became unresponsive ("Not Responding") until we restarted it manually.

While addressing the automated review comments, I also fixed two related bugs so the budget holds up better against exactly this kind of case:

  • dom_informative_nodes.append(TextElementNode(...)) in the browser-DOM path wasn't consuming the element budget (only the interactive-node branches were) — text-heavy pages could bypass truncation entirely. Fixed by consuming budget there too. This path is browser-DOM-only though, so it isn't what caused the TriData stall above — just a real gap the review caught.
  • TreeElementBudget.try_consume(amount) could push count past limit on a single batched call while still returning True, contradicting its own hard-cap contract. Made it strict (clamps to limit, returns False, marks truncated).

Also picked up the smaller nits (missing -> None, docstring style, quoting). Pushed as a new commit on this branch.

@Jeomon
Jeomon merged commit 4be75f3 into CursorTouch:main Jul 29, 2026
1 check passed
@Jeomon

Jeomon commented Jul 29, 2026

Copy link
Copy Markdown
Member

Thanks

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants