feat: indexed UI tree with SoM annotations + improved element resolution

- Flat indexed UI element list with numeric IDs and center coordinates
- Screenshot annotations: numbered colored boxes (Set of Marks)
- click_element resolves by index from cached element list
- Cursor crosshair + focused element highlight in screenshots

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

easybot_windows_agent/main.py

@@ -36,6 +36,40 @@ RECONNECT_DELAY_BASE = 5
 RECONNECT_DELAY_MAX = 120
 TOOL_TIMEOUT = 60  # max seconds for a single tool call
 
+# ── UI Element Cache (Set of Marks) ─────────────────────────────────────────
+
+_last_ui_elements: list[dict] = []
+
+# Control types considered interactive for SoM indexing
+_INTERACTIVE_TYPES = {
+    "ButtonControl", "EditControl", "CheckBoxControl", "RadioButtonControl",
+    "ListItemControl", "MenuItemControl", "ComboBoxControl", "TextControl",
+    "HyperlinkControl", "TabItemControl", "TreeItemControl", "DataItemControl",
+    "SliderControl", "MenuBarItemControl", "SpinnerControl", "SplitButtonControl",
+    "ToggleButtonControl",
+}
+
+# Color palette for SoM annotations (by control type category)
+_SOM_COLORS: dict[str, str] = {
+    "ButtonControl": "#FF6B6B",
+    "EditControl": "#4ECDC4",
+    "CheckBoxControl": "#45B7D1",
+    "RadioButtonControl": "#45B7D1",
+    "ListItemControl": "#96CEB4",
+    "MenuItemControl": "#FFEAA7",
+    "ComboBoxControl": "#DDA0DD",
+    "TextControl": "#98D8C8",
+    "HyperlinkControl": "#87CEEB",
+    "TabItemControl": "#F7DC6F",
+    "TreeItemControl": "#82E0AA",
+    "DataItemControl": "#AED6F1",
+    "SliderControl": "#F1948A",
+    "MenuBarItemControl": "#FDEBD0",
+    "SpinnerControl": "#D2B4DE",
+    "SplitButtonControl": "#F5B041",
+    "ToggleButtonControl": "#FF6B6B",
+}
+
 
 # ── Tool Handlers ────────────────────────────────────────────────────────────
 
@@ -74,11 +108,12 @@ async def tool_run_command(args: dict[str, Any]) -> dict[str, Any]:
 
 
 async def tool_get_ui_tree(args: dict[str, Any]) -> dict[str, Any]:
-    """Get the accessibility tree of a window or the desktop."""
+    """Get indexed interactive UI elements with coordinates (flat list for SoM)."""
+    global _last_ui_elements
     import uiautomation as auto
 
     window_title = args.get("window")
-    max_depth = args.get("depth", 3)
+    max_depth = args.get("depth", 8)
 
     try:
         if window_title:
@@ -93,46 +128,50 @@ async def tool_get_ui_tree(args: dict[str, Any]) -> dict[str, Any]:
             if not root:
                 root = auto.GetRootControl()
 
-        def walk(ctrl: Any, depth: int, max_d: int) -> list[dict]:
+        elements: list[dict] = []
-            if depth > max_d:
+
-                return []
+        def collect(ctrl: Any, depth: int) -> None:
-            items = []
+            if depth > max_depth:
+                return
             try:
-                name = ctrl.Name or ""
                 ctrl_type = ctrl.ControlTypeName or ""
+                if ctrl_type in _INTERACTIVE_TYPES:
+                    name = ctrl.Name or ""
                     auto_id = ctrl.AutomationId or ""
                     rect = ctrl.BoundingRectangle
-                value = ""
+                    if rect and rect.width() > 0 and rect.height() > 0:
-                try:
+                        cx = rect.left + rect.width() // 2
-                    value = ctrl.GetValuePattern().Value if ctrl.GetValuePattern() else ""
+                        cy = rect.top + rect.height() // 2
-                except Exception:
+                        elem: dict[str, Any] = {
-                    pass
+                            "id": len(elements),
-
-                item: dict[str, Any] = {
                             "type": ctrl_type,
                             "name": name,
+                            "coords": [cx, cy],
+                            "rect": {"x": rect.left, "y": rect.top, "w": rect.width(), "h": rect.height()},
                         }
                         if auto_id:
-                    item["automation_id"] = auto_id
+                            elem["automation_id"] = auto_id
-                if value:
+                        # Try to get value
-                    item["value"] = value
+                        try:
-                if rect and rect.width() > 0:
+                            vp = ctrl.GetValuePattern()
-                    item["rect"] = {"x": rect.left, "y": rect.top, "w": rect.width(), "h": rect.height()}
+                            if vp:
-
+                                val = vp.Value
-                children = []
+                                if val:
-                if depth < max_d:
+                                    elem["value"] = val
-                    for child in ctrl.GetChildren():
+                        except Exception:
-                        children.extend(walk(child, depth + 1, max_d))
+                            pass
-                if children:
+                        elements.append(elem)
-                    item["children"] = children
+            except Exception:
-
+                pass
-                items.append(item)
+            try:
+                for child in ctrl.GetChildren():
+                    collect(child, depth + 1)
             except Exception:
                 pass
-            return items
 
-        tree = walk(root, 0, max_depth)
+        collect(root, 0)
-        return {"window": root.Name or "Desktop", "tree": tree}
+        _last_ui_elements = elements
+        return {"window": root.Name or "Desktop", "element_count": len(elements), "elements": elements}
     except Exception as e:
         return {"error": str(e)}
 
@@ -201,7 +240,8 @@ def _get_ui_feedback() -> dict[str, Any]:
 
 
 async def tool_click_element(args: dict[str, Any]) -> dict[str, Any]:
-    """Click a UI element by name, automation_id, or coordinates."""
+    """Click a UI element by index (from UI tree), name, automation_id, or coordinates."""
+    index = args.get("index")
     name = args.get("name")
     auto_id = args.get("automation_id")
     x = args.get("x")
@@ -210,6 +250,16 @@ async def tool_click_element(args: dict[str, Any]) -> dict[str, Any]:
     double = args.get("double", False)
 
     try:
+        # Resolve index to coordinates from cached UI tree
+        if index is not None:
+            if not _last_ui_elements:
+                return {"error": "No UI tree cached. Call get_ui_tree first."}
+            if index < 0 or index >= len(_last_ui_elements):
+                return {"error": f"Index {index} out of range (0-{len(_last_ui_elements) - 1})"}
+            elem = _last_ui_elements[index]
+            x, y = elem["coords"]
+            logger.debug("Index %d resolved to (%d, %d) — %s", index, x, y, elem.get("name", ""))
+
         if x is not None and y is not None:
             import pyautogui
             if double:
@@ -295,11 +345,13 @@ async def tool_press_keys(args: dict[str, Any]) -> dict[str, Any]:
 
 
 async def tool_screenshot(args: dict[str, Any]) -> dict[str, Any]:
-    """Take a screenshot with UI annotations: cursor marker + focused element highlight."""
+    """Take a screenshot with Set of Marks (SoM) annotations from the UI tree."""
-    from PIL import ImageGrab, ImageDraw
+    from PIL import ImageGrab, ImageDraw, ImageFont
 
     try:
         region = args.get("region")
+        annotate = args.get("annotate", True)
+
         if region:
             bbox = (region["x"], region["y"], region["x"] + region["w"], region["y"] + region["h"])
             img = ImageGrab.grab(bbox=bbox)
@@ -343,6 +395,58 @@ async def tool_screenshot(args: dict[str, Any]) -> dict[str, Any]:
         except Exception:
             pass
 
+        # 3. Set of Marks (SoM): draw indexed labels on interactive elements
+        if annotate and _last_ui_elements:
+            # Try to load a small font for labels
+            try:
+                font = ImageFont.truetype("arial.ttf", 11)
+            except Exception:
+                try:
+                    font = ImageFont.truetype("C:\\Windows\\Fonts\\arial.ttf", 11)
+                except Exception:
+                    font = ImageFont.load_default()
+
+            region_ox = region["x"] if region else 0
+            region_oy = region["y"] if region else 0
+
+            for elem in _last_ui_elements:
+                try:
+                    r = elem["rect"]
+                    ex = r["x"] - region_ox
+                    ey = r["y"] - region_oy
+                    ew = r["w"]
+                    eh = r["h"]
+
+                    # Skip elements outside visible area
+                    if ex + ew < 0 or ey + eh < 0 or ex > img.width or ey > img.height:
+                        continue
+
+                    color = _SOM_COLORS.get(elem["type"], "#AAAAAA")
+                    idx = elem["id"]
+
+                    # Draw bounding box (thin outline)
+                    draw.rectangle([ex, ey, ex + ew, ey + eh], outline=color, width=1)
+
+                    # Draw index label (small colored background with white text)
+                    label_text = str(idx)
+                    # Get text size
+                    text_bbox = draw.textbbox((0, 0), label_text, font=font)
+                    tw = text_bbox[2] - text_bbox[0] + 4
+                    th = text_bbox[3] - text_bbox[1] + 2
+
+                    # Position label at top-left of element
+                    lx = ex
+                    ly = ey - th - 1
+                    if ly < 0:
+                        ly = ey  # put inside if no room above
+
+                    # Draw label background
+                    draw.rectangle([lx, ly, lx + tw, ly + th], fill=color)
+                    # Draw label text
+                    draw.text((lx + 2, ly + 1), label_text, fill="white", font=font)
+                except Exception:
+                    continue
+
         # Resize for efficiency (max 1280px wide)
         if img.width > 1280:
             ratio = 1280 / img.width