Skip to content

Harden UIA control collection #335

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

Open
rodion-m wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Harden UIA control collection #335

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
7 changes: 6 additions & 1 deletion documents/docs/mcp/servers/ui_collector.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -188,6 +188,8 @@ Scans the currently selected window and retrieves information about all interact
| Parameter | Type | Required | Default | Description |
|-----------|------|----------|---------|-------------|
| `field_list` | `List[str]` | Yes | - | List of field names to retrieve for each control |
| `depth` | `int` | No | `0` | Descendant depth to search. Use `1` to inspect direct children only and avoid expensive full-subtree scans. |
| `best_effort` | `bool` | No | `True` | Return partial or empty results instead of failing when a UIA element cannot be read. |

#### Supported Fields

Expand All @@ -212,7 +214,8 @@ result = await computer.run_actions([
tool_key="data_collection::get_app_window_controls_info",
tool_name="get_app_window_controls_info",
parameters={
"field_list": ["label", "control_text", "control_type"]
"field_list": ["label", "control_text", "control_type"],
"depth": 1
}
)
])
Expand Down Expand Up @@ -247,6 +250,8 @@ Similar to `get_app_window_controls_info`, but returns structured `TargetInfo` o
| Parameter | Type | Required | Default | Description |
|-----------|------|----------|---------|-------------|
| `field_list` | `List[str]` | Yes | - | List of field names to retrieve |
| `depth` | `int` | No | `0` | Descendant depth to search. Use `1` to inspect direct children only and avoid expensive full-subtree scans. |
| `best_effort` | `bool` | No | `True` | Return partial or empty results instead of failing when a UIA element cannot be read. |

#### Returns

Expand Down
51 changes: 36 additions & 15 deletions ufo/automator/ui_control/inspector.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,13 +4,16 @@
from __future__ import annotations

import functools
import logging
import platform
import time
from abc import ABC, abstractmethod
from typing import Callable, Dict, List, Optional, cast, TYPE_CHECKING, Any

import psutil

logger = logging.getLogger(__name__)

# Conditional imports for Windows-specific packages
if TYPE_CHECKING or platform.system() == "Windows":
import comtypes.gen.UIAutomationClient as UIAutomationClient_dll
Expand Down Expand Up @@ -248,24 +251,42 @@ def find_control_elements_in_descendants(

cache_request = UIABackendStrategy._get_cache_request()

com_elem_array = window_elem_com_ref.FindAllBuildCache(
scope=iuia_dll.TreeScope_Descendants,
condition=condition,
cacheRequest=cache_request,
tree_scope = (
iuia_dll.TreeScope_Children if depth == 1 else iuia_dll.TreeScope_Descendants
)

elem_info_list = [
(
elem,
elem.CachedControlType,
elem.CachedName,
elem.CachedBoundingRectangle,
)
for elem in (
com_elem_array.GetElement(n)
for n in range(min(com_elem_array.Length, 500))
try:
com_elem_array = window_elem_com_ref.FindAllBuildCache(
scope=tree_scope,
condition=condition,
cacheRequest=cache_request,
)
]
except Exception as exc:
logger.warning("UIA FindAllBuildCache failed: %s", exc)
return []

elem_info_list = []
try:
elem_count = min(com_elem_array.Length, 500)
except Exception as exc:
logger.warning("UIA element array length lookup failed: %s", exc)
return []

for n in range(elem_count):
try:
elem = com_elem_array.GetElement(n)
elem_info_list.append(
(
elem,
elem.CachedControlType,
elem.CachedName,
elem.CachedBoundingRectangle,
)
)
except Exception as exc:
logger.debug(
"Skipping UIA element %s after property lookup failed: %s", n, exc
)

control_elements: List[UIAWrapper] = []

Expand Down
58 changes: 43 additions & 15 deletions ufo/client/mcp/local_servers/ui_mcp_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -744,46 +744,74 @@ def get_app_window_info(field_list: List[str]) -> Dict[str, Any]:
return window_info

@data_mcp.tool()
def get_app_window_controls_info(field_list: List[str]) -> List:
def get_app_window_controls_info(
field_list: List[str], depth: int = 0, best_effort: bool = True
) -> List:
"""
Get information about controls in the currently selected application window.
:param field_list: List of fields to retrieve from the control info.
:param depth: Descendant depth to search. Use 1 to avoid expensive full-subtree scans.
:param best_effort: If true, return partial/empty results instead of failing on UIA errors.
:return: Dictionary containing the requested control information.
"""
if not ui_state.selected_app_window:
raise ToolError("No window is selected， please select a window first.")

controls_list = ui_state.control_inspector.find_control_elements_in_descendants(
ui_state.selected_app_window,
control_type_list=configs.get("CONTROL_LIST", []),
class_name_list=configs.get("CONTROL_LIST", []),
)
try:
controls_list = ui_state.control_inspector.find_control_elements_in_descendants(
ui_state.selected_app_window,
control_type_list=configs.get("CONTROL_LIST", []),
class_name_list=configs.get("CONTROL_LIST", []),
depth=depth,
)
except Exception as exc:
if not best_effort:
raise
logger.warning("get_app_window_controls_info failed: %s", exc)
controls_list = []

control_dict = {str(i + 1): control for i, control in enumerate(controls_list)}

result = ui_state.control_inspector.get_control_info_list_of_dict(
control_dict, field_list=field_list
)
try:
result = ui_state.control_inspector.get_control_info_list_of_dict(
control_dict, field_list=field_list
)
except Exception as exc:
if not best_effort:
raise
logger.warning("get_app_window_controls_info serialization failed: %s", exc)
result = []

ui_state.control_dict = control_dict

return result

@data_mcp.tool()
def get_app_window_controls_target_info(field_list: List[str]) -> List:
def get_app_window_controls_target_info(
field_list: List[str], depth: int = 0, best_effort: bool = True
) -> List:
"""
Get information about controls in the currently selected application window.
:param field_list: List of fields to retrieve from the control info.
:param depth: Descendant depth to search. Use 1 to avoid expensive full-subtree scans.
:param best_effort: If true, return partial/empty results instead of failing on UIA errors.
:return: Dictionary containing the requested control information.
"""
if not ui_state.selected_app_window:
raise ToolError("No window is selected， please select a window first.")

controls_list = ui_state.control_inspector.find_control_elements_in_descendants(
ui_state.selected_app_window,
control_type_list=configs.get("CONTROL_LIST", []),
class_name_list=configs.get("CONTROL_LIST", []),
)
try:
controls_list = ui_state.control_inspector.find_control_elements_in_descendants(
ui_state.selected_app_window,
control_type_list=configs.get("CONTROL_LIST", []),
class_name_list=configs.get("CONTROL_LIST", []),
depth=depth,
)
except Exception as exc:
if not best_effort:
raise
logger.warning("get_app_window_controls_target_info failed: %s", exc)
controls_list = []

control_dict = {str(i + 1): control for i, control in enumerate(controls_list)}

Expand Down