easybot-windows-agent/easybot_windows_agent/main.py

"""
EasyBot Windows Agent — tool-based remote automation client.

The PC exposes granular tools (run_command, get_ui_tree, click, type, etc.)
and the EasyBot LLM agent controls everything directly. No LLM on client side.

Configuration via environment variables:
  EASYBOT_SERVER    WebSocket URL (e.g., wss://app.example.com/ws/windows-agent)
  EASYBOT_TOKEN     Authentication token (eb_wag_...)
  LOG_LEVEL         DEBUG | INFO | WARNING (default: INFO)
"""

import asyncio
import base64
import io
import json
import logging
import os
import platform
import signal
import subprocess
import sys
import time
from typing import Any

import websockets
from websockets.client import WebSocketClientProtocol

from . import __version__

logger = logging.getLogger("easybot-windows-agent")

# ── Configuration ────────────────────────────────────────────────────────────

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 ────────────────────────────────────────────────────────────

async def tool_run_command(args: dict[str, Any]) -> dict[str, Any]:
    """Execute a shell command and return output."""
    command = args.get("command", "")
    shell = args.get("shell", "powershell")

    if not command:
        return {"error": "No command provided"}

    try:
        if shell == "cmd":
            proc = await asyncio.create_subprocess_exec(
                "cmd", "/c", command,
                stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                creationflags=subprocess.CREATE_NO_WINDOW if sys.platform == "win32" else 0,
            )
        else:
            proc = await asyncio.create_subprocess_exec(
                "powershell", "-NoProfile", "-NonInteractive", "-Command", command,
                stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                creationflags=subprocess.CREATE_NO_WINDOW if sys.platform == "win32" else 0,
            )

        stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=TOOL_TIMEOUT)
        return {
            "stdout": stdout.decode("utf-8", errors="replace").strip(),
            "stderr": stderr.decode("utf-8", errors="replace").strip(),
            "exit_code": proc.returncode,
        }
    except asyncio.TimeoutError:
        return {"error": "Command timed out", "exit_code": -1}
    except Exception as e:
        return {"error": str(e)}


async def tool_get_ui_tree(args: dict[str, Any]) -> dict[str, Any]:
    """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", 8)

    try:
        if window_title:
            root = auto.WindowControl(searchDepth=1, Name=window_title)
            if not root.Exists(maxSearchSeconds=3):
                # Try substring match
                root = auto.WindowControl(searchDepth=1, SubName=window_title)
                if not root.Exists(maxSearchSeconds=2):
                    return {"error": f"Window '{window_title}' not found"}
        else:
            root = auto.GetForegroundControl()
            if not root:
                root = auto.GetRootControl()

        elements: list[dict] = []

        def collect(ctrl: Any, depth: int) -> None:
            if depth > max_depth:
                return
            try:
                ctrl_type = ctrl.ControlTypeName or ""
                if ctrl_type in _INTERACTIVE_TYPES:
                    name = ctrl.Name or ""
                    auto_id = ctrl.AutomationId or ""
                    rect = ctrl.BoundingRectangle
                    if rect and rect.width() > 0 and rect.height() > 0:
                        cx = rect.left + rect.width() // 2
                        cy = rect.top + rect.height() // 2
                        elem: dict[str, Any] = {
                            "id": len(elements),
                            "type": ctrl_type,
                            "name": name,
                            "coords": [cx, cy],
                            "rect": {"x": rect.left, "y": rect.top, "w": rect.width(), "h": rect.height()},
                        }
                        if auto_id:
                            elem["automation_id"] = auto_id
                        # Try to get value
                        try:
                            vp = ctrl.GetValuePattern()
                            if vp:
                                val = vp.Value
                                if val:
                                    elem["value"] = val
                        except Exception:
                            pass
                        elements.append(elem)
            except Exception:
                pass
            try:
                for child in ctrl.GetChildren():
                    collect(child, depth + 1)
            except Exception:
                pass

        collect(root, 0)
        _last_ui_elements = elements
        return {"window": root.Name or "Desktop", "element_count": len(elements), "elements": elements}
    except Exception as e:
        return {"error": str(e)}


async def tool_list_windows(args: dict[str, Any]) -> dict[str, Any]:
    """List all open top-level windows."""
    import uiautomation as auto

    try:
        windows = []
        for win in auto.GetRootControl().GetChildren():
            try:
                if win.ControlTypeName == "WindowControl" or win.ClassName == "Window":
                    name = win.Name or ""
                    if name and name != "Program Manager":
                        windows.append({
                            "title": name,
                            "class": win.ClassName or "",
                            "process_id": win.ProcessId,
                        })
            except Exception:
                continue
        return {"windows": windows}
    except Exception as e:
        return {"error": str(e)}


def _get_ui_feedback() -> dict[str, Any]:
    """Collect UI state feedback after an action: cursor position, focused element, active window."""
    feedback: dict[str, Any] = {}
    try:
        import pyautogui
        mx, my = pyautogui.position()
        feedback["cursor"] = {"x": mx, "y": my}
    except Exception:
        pass
    try:
        import uiautomation as auto
        fw = auto.GetForegroundControl()
        if fw:
            feedback["active_window"] = fw.Name or ""
        focused = auto.GetFocusedControl()
        if focused:
            fe: dict[str, Any] = {
                "type": focused.ControlTypeName or "",
                "name": focused.Name or "",
            }
            if focused.AutomationId:
                fe["automation_id"] = focused.AutomationId
            try:
                vp = focused.GetValuePattern()
                if vp:
                    fe["value"] = vp.Value
            except Exception:
                pass
            try:
                rect = focused.BoundingRectangle
                if rect and rect.width() > 0:
                    fe["rect"] = {"x": rect.left, "y": rect.top, "w": rect.width(), "h": rect.height()}
            except Exception:
                pass
            feedback["focused_element"] = fe
    except Exception:
        pass
    return feedback


async def tool_click_element(args: dict[str, Any]) -> dict[str, Any]:
    """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")
    y = args.get("y")
    button = args.get("button", "left")
    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:
                pyautogui.doubleClick(x, y)
            elif button == "right":
                pyautogui.rightClick(x, y)
            else:
                pyautogui.click(x, y)
            await asyncio.sleep(0.3)
            return {"success": True, "method": "coordinates", "x": x, "y": y, **_get_ui_feedback()}

        import uiautomation as auto
        ctrl = None
        if auto_id:
            ctrl = auto.Control(searchDepth=8, AutomationId=auto_id)
        elif name:
            ctrl = auto.Control(searchDepth=8, Name=name)

        if ctrl and ctrl.Exists(maxSearchSeconds=3):
            if double:
                ctrl.DoubleClick()
            else:
                ctrl.Click()
            await asyncio.sleep(0.3)
            return {"success": True, "method": "ui_automation", "element": ctrl.Name, **_get_ui_feedback()}
        else:
            return {"error": f"Element not found: name={name}, automation_id={auto_id}"}
    except Exception as e:
        return {"error": str(e)}


async def tool_type_text(args: dict[str, Any]) -> dict[str, Any]:
    """Type text, optionally into a specific element."""
    text = args.get("text", "")
    element_name = args.get("element_name")
    clear_first = args.get("clear_first", False)

    if not text:
        return {"error": "No text provided"}

    try:
        if element_name:
            import uiautomation as auto
            ctrl = auto.Control(searchDepth=8, Name=element_name)
            if ctrl.Exists(maxSearchSeconds=3):
                try:
                    if clear_first:
                        ctrl.GetValuePattern().SetValue("")
                    ctrl.GetValuePattern().SetValue(text)
                    return {"success": True, "method": "set_value", "element": element_name, **_get_ui_feedback()}
                except Exception:
                    ctrl.Click()
                    await asyncio.sleep(0.2)

        import pyautogui
        if clear_first:
            pyautogui.hotkey("ctrl", "a")
            await asyncio.sleep(0.05)
        pyautogui.typewrite(text, interval=0.02) if text.isascii() else pyautogui.write(text)
        await asyncio.sleep(0.1)
        return {"success": True, "method": "keyboard", **_get_ui_feedback()}
    except Exception as e:
        return {"error": str(e)}


async def tool_press_keys(args: dict[str, Any]) -> dict[str, Any]:
    """Press keyboard keys/shortcuts (e.g. 'ctrl+s', 'enter', 'alt+f4')."""
    keys = args.get("keys", "")
    if not keys:
        return {"error": "No keys provided"}

    try:
        import pyautogui
        parts = [k.strip() for k in keys.lower().split("+")]
        if len(parts) > 1:
            pyautogui.hotkey(*parts)
        else:
            pyautogui.press(parts[0])
        await asyncio.sleep(0.3)
        return {"success": True, "keys": keys, **_get_ui_feedback()}
    except Exception as e:
        return {"error": str(e)}


async def tool_screenshot(args: dict[str, Any]) -> dict[str, Any]:
    """Take a screenshot with Set of Marks (SoM) annotations from the UI tree."""
    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)
        else:
            img = ImageGrab.grab()

        # Draw annotations on the screenshot
        draw = ImageDraw.Draw(img)

        # 1. Draw cursor position (red crosshair)
        try:
            import pyautogui
            mx, my = pyautogui.position()
            # Adjust if region capture
            if region:
                mx -= region["x"]
                my -= region["y"]
            r = 12
            draw.ellipse([mx - r, my - r, mx + r, my + r], outline="red", width=3)
            draw.line([mx - r - 4, my, mx + r + 4, my], fill="red", width=2)
            draw.line([mx, my - r - 4, mx, my + r + 4], fill="red", width=2)
        except Exception:
            pass

        # 2. Highlight focused element (green rectangle + label)
        try:
            import uiautomation as auto
            focused = auto.GetFocusedControl()
            if focused:
                rect = focused.BoundingRectangle
                if rect and rect.width() > 0:
                    fx, fy = rect.left, rect.top
                    fw, fh = rect.width(), rect.height()
                    if region:
                        fx -= region["x"]
                        fy -= region["y"]
                    draw.rectangle([fx, fy, fx + fw, fy + fh], outline="lime", width=2)
                    label = focused.Name or focused.ControlTypeName or ""
                    if label:
                        draw.text((fx + 2, fy - 14), label[:40], fill="lime")
        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
            img = img.resize((1280, int(img.height * ratio)))

        buf = io.BytesIO()
        img.save(buf, format="PNG", optimize=True)
        b64 = base64.b64encode(buf.getvalue()).decode("ascii")
        return {"image_base64": b64, "width": img.width, "height": img.height}
    except Exception as e:
        return {"error": str(e)}


async def tool_read_clipboard(args: dict[str, Any]) -> dict[str, Any]:
    """Read text from the clipboard."""
    try:
        import pyperclip
        text = pyperclip.paste()
        return {"text": text}
    except ImportError:
        # Fallback via PowerShell
        proc = await asyncio.create_subprocess_exec(
            "powershell", "-NoProfile", "-Command", "Get-Clipboard",
            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        )
        stdout, _ = await proc.communicate()
        return {"text": stdout.decode("utf-8", errors="replace").strip()}
    except Exception as e:
        return {"error": str(e)}


async def tool_set_clipboard(args: dict[str, Any]) -> dict[str, Any]:
    """Set clipboard text content."""
    text = args.get("text", "")
    try:
        import pyperclip
        pyperclip.copy(text)
        return {"success": True}
    except ImportError:
        proc = await asyncio.create_subprocess_exec(
            "powershell", "-NoProfile", "-Command", f"Set-Clipboard -Value '{text}'",
            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        )
        await proc.communicate()
        return {"success": True}
    except Exception as e:
        return {"error": str(e)}


async def tool_focus_window(args: dict[str, Any]) -> dict[str, Any]:
    """Bring a window to the foreground."""
    title = args.get("title", "")
    if not title:
        return {"error": "No window title provided"}

    try:
        import uiautomation as auto
        win = auto.WindowControl(searchDepth=1, SubName=title)
        if win.Exists(maxSearchSeconds=3):
            win.SetFocus()
            win.SetActive()
            return {"success": True, "window": win.Name}
        else:
            return {"error": f"Window '{title}' not found"}
    except Exception as e:
        return {"error": str(e)}


async def tool_file_read(args: dict[str, Any]) -> dict[str, Any]:
    """Read a text file."""
    path = args.get("path", "")
    if not path:
        return {"error": "No path provided"}
    try:
        with open(path, "r", encoding="utf-8", errors="replace") as f:
            content = f.read(100_000)  # limit to 100KB
        return {"content": content, "size": os.path.getsize(path)}
    except Exception as e:
        return {"error": str(e)}


async def tool_file_write(args: dict[str, Any]) -> dict[str, Any]:
    """Write content to a text file."""
    path = args.get("path", "")
    content = args.get("content", "")
    if not path:
        return {"error": "No path provided"}
    try:
        os.makedirs(os.path.dirname(path), exist_ok=True) if os.path.dirname(path) else None
        with open(path, "w", encoding="utf-8") as f:
            f.write(content)
        return {"success": True, "bytes_written": len(content.encode("utf-8"))}
    except Exception as e:
        return {"error": str(e)}


# ── Browser CDP Tool ─────────────────────────────────────────────────────────

CDP_PORT = 9222
_browser_process: subprocess.Popen | None = None


async def tool_open_browser(args: dict[str, Any]) -> dict[str, Any]:
    """Launch Edge/Chrome with CDP enabled. Returns the CDP port for tunnel."""
    global _browser_process
    url = args.get("url", "about:blank")
    port = args.get("port", CDP_PORT)

    # Find browser executable
    browser_paths = [
        r"C:\Program Files (x86)\Microsoft\Edge\Application\msedge.exe",
        r"C:\Program Files\Microsoft\Edge\Application\msedge.exe",
        r"C:\Program Files\Google\Chrome\Application\chrome.exe",
        r"C:\Program Files (x86)\Google\Chrome\Application\chrome.exe",
    ]
    browser_exe = None
    for p in browser_paths:
        if os.path.exists(p):
            browser_exe = p
            break

    if not browser_exe:
        return {"error": "No Chromium browser found (Edge or Chrome)"}

    try:
        # Kill existing CDP browser if any
        if _browser_process and _browser_process.poll() is None:
            _browser_process.terminate()
            await asyncio.sleep(0.5)

        _browser_process = subprocess.Popen(
            [browser_exe, f"--remote-debugging-port={port}",
             "--no-first-run", "--no-default-browser-check", url],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
            creationflags=subprocess.CREATE_NO_WINDOW if sys.platform == "win32" else 0,
        )
        # Wait a bit for CDP to be ready
        await asyncio.sleep(2)

        # Verify CDP is listening
        try:
            reader, writer = await asyncio.open_connection("127.0.0.1", port)
            writer.close()
            await writer.wait_closed()
        except OSError:
            return {"error": f"Browser started but CDP port {port} not responding"}

        return {
            "success": True,
            "browser": os.path.basename(browser_exe),
            "cdp_port": port,
            "pid": _browser_process.pid,
            "message": f"Browser ready with CDP on port {port}. Tunnel available.",
        }
    except Exception as e:
        return {"error": str(e)}


async def tool_close_browser(args: dict[str, Any]) -> dict[str, Any]:
    """Close the CDP browser."""
    global _browser_process
    if _browser_process and _browser_process.poll() is None:
        _browser_process.terminate()
        _browser_process = None
        return {"success": True}
    return {"success": True, "message": "No browser was running"}


# ── Tool Registry ────────────────────────────────────────────────────────────

TOOLS: dict[str, Any] = {
    "run_command": tool_run_command,
    "get_ui_tree": tool_get_ui_tree,
    "list_windows": tool_list_windows,
    "click_element": tool_click_element,
    "type_text": tool_type_text,
    "press_keys": tool_press_keys,
    "screenshot": tool_screenshot,
    "read_clipboard": tool_read_clipboard,
    "set_clipboard": tool_set_clipboard,
    "focus_window": tool_focus_window,
    "file_read": tool_file_read,
    "file_write": tool_file_write,
    "open_browser": tool_open_browser,
    "close_browser": tool_close_browser,
}


# ── Agent Client ─────────────────────────────────────────────────────────────

class AgentClient:
    def __init__(self, server_url: str, token: str):
        self.server_url = server_url
        self.token = token
        self.ws: WebSocketClientProtocol | None = None
        self.running = True
        self.current_call_id: str | None = None
        self.cancel_event = asyncio.Event()
        self._reconnect_delay = RECONNECT_DELAY_BASE
        # TCP tunnel state: conn_id → (reader, writer)
        self._tunnel_connections: dict[int, tuple[asyncio.StreamReader, asyncio.StreamWriter]] = {}

    async def run(self) -> None:
        """Main loop: connect → handle messages → reconnect on failure."""
        while self.running:
            try:
                await self._connect_and_handle()
            except (websockets.ConnectionClosed, OSError, ConnectionRefusedError) as e:
                logger.warning("Connection lost: %s — reconnecting in %ds", e, self._reconnect_delay)
            except Exception:
                logger.exception("Unexpected error — reconnecting in %ds", self._reconnect_delay)

            if not self.running:
                break
            await asyncio.sleep(self._reconnect_delay)
            self._reconnect_delay = min(self._reconnect_delay * 2, RECONNECT_DELAY_MAX)

    async def _connect_and_handle(self) -> None:
        """Single connection lifecycle."""
        extra_headers = {"Authorization": f"Bearer {self.token}"}

        async with websockets.connect(
            self.server_url,
            extra_headers=extra_headers,
            ping_interval=None,
            max_size=16 * 1024 * 1024,
            close_timeout=10,
        ) as ws:
            self.ws = ws
            self._reconnect_delay = RECONNECT_DELAY_BASE
            logger.info("Connected to %s", self.server_url)

            # Send client info with list of available tools
            await self._send({
                "type": "client_info",
                "hostname": platform.node(),
                "os_version": f"{platform.system()} {platform.version()}",
                "python_version": platform.python_version(),
                "client_version": __version__,
                "tools": list(TOOLS.keys()),
            })

            try:
                async for raw in ws:
                    if isinstance(raw, bytes):
                        # Binary frame: tunnel data [4-byte connId BE][payload]
                        await self._handle_tunnel_data(raw)
                    else:
                        msg = json.loads(raw)
                        await self._handle_message(msg)
            finally:
                self.ws = None
                self.current_call_id = None
                await self._close_all_tunnels()

    async def _handle_message(self, msg: dict[str, Any]) -> None:
        """Dispatch incoming messages from the server."""
        msg_type = msg.get("type")

        if msg_type == "ping":
            await self._send({"type": "pong"})

        elif msg_type == "tool_call":
            call_id = msg["call_id"]
            tool_name = msg["tool_name"]
            arguments = msg.get("arguments", {})
            logger.info("Tool call: %s (call_id=%s)", tool_name, call_id[:8])
            asyncio.create_task(self._execute_tool(call_id, tool_name, arguments))

        elif msg_type == "cancel":
            call_id = msg.get("call_id")
            if self.current_call_id == call_id:
                logger.info("Cancel requested for %s", call_id)
                self.cancel_event.set()

        elif msg_type == "tunnel_open":
            conn_id = msg["connection_id"]
            port = msg.get("port", CDP_PORT)
            asyncio.create_task(self._tunnel_open(conn_id, port))

        elif msg_type == "tunnel_close":
            conn_id = msg["connection_id"]
            await self._tunnel_close(conn_id)

        else:
            logger.debug("Unknown message type: %s", msg_type)

    async def _execute_tool(self, call_id: str, tool_name: str, arguments: dict[str, Any]) -> None:
        """Execute a tool and return the result."""
        self.current_call_id = call_id
        self.cancel_event.clear()
        start = time.time()

        try:
            handler = TOOLS.get(tool_name)
            if not handler:
                await self._send({
                    "type": "tool_result",
                    "call_id": call_id,
                    "error": f"Unknown tool: {tool_name}",
                    "duration_ms": 0,
                })
                return

            result = await asyncio.wait_for(handler(arguments), timeout=TOOL_TIMEOUT)
            duration_ms = int((time.time() - start) * 1000)

            if "error" in result:
                await self._send({
                    "type": "tool_result",
                    "call_id": call_id,
                    "error": result["error"],
                    "duration_ms": duration_ms,
                })
            else:
                await self._send({
                    "type": "tool_result",
                    "call_id": call_id,
                    "result": result,
                    "duration_ms": duration_ms,
                })

            logger.info("Tool %s completed in %dms", tool_name, duration_ms)

        except asyncio.TimeoutError:
            await self._send({
                "type": "tool_result",
                "call_id": call_id,
                "error": f"Tool '{tool_name}' timed out after {TOOL_TIMEOUT}s",
                "duration_ms": int((time.time() - start) * 1000),
            })
        except Exception as e:
            logger.error("Tool %s failed: %s", tool_name, e)
            await self._send({
                "type": "tool_result",
                "call_id": call_id,
                "error": f"{type(e).__name__}: {e}",
                "duration_ms": int((time.time() - start) * 1000),
            })
        finally:
            self.current_call_id = None

    # ── TCP Tunnel Methods ──────────────────────────────────────────────────

    async def _tunnel_open(self, conn_id: int, port: int) -> None:
        """Open a TCP connection to localhost:port and start forwarding."""
        try:
            reader, writer = await asyncio.open_connection("127.0.0.1", port)
            self._tunnel_connections[conn_id] = (reader, writer)
            await self._send({"type": "tunnel_open_ok", "connection_id": conn_id})
            # Start reading from TCP and forwarding to WS
            asyncio.create_task(self._tunnel_read_loop(conn_id, reader))
            logger.debug("Tunnel connection %d opened to port %d", conn_id, port)
        except Exception as e:
            await self._send({"type": "tunnel_open_err", "connection_id": conn_id, "error": str(e)})

    async def _tunnel_read_loop(self, conn_id: int, reader: asyncio.StreamReader) -> None:
        """Read from TCP socket and send as binary frames to server."""
        try:
            while True:
                data = await reader.read(65536)
                if not data:
                    break
                # Binary frame: [4-byte connId BE][payload]
                frame = conn_id.to_bytes(4, "big") + data
                if self.ws:
                    await self.ws.send(frame)
        except (ConnectionError, asyncio.IncompleteReadError):
            pass
        finally:
            await self._tunnel_close(conn_id)
            await self._send({"type": "tunnel_close", "connection_id": conn_id})

    async def _handle_tunnel_data(self, raw: bytes) -> None:
        """Handle incoming binary frame: forward to TCP socket."""
        if len(raw) < 5:
            return
        conn_id = int.from_bytes(raw[:4], "big")
        payload = raw[4:]
        conn = self._tunnel_connections.get(conn_id)
        if conn:
            _, writer = conn
            try:
                writer.write(payload)
                await writer.drain()
            except (ConnectionError, OSError):
                await self._tunnel_close(conn_id)

    async def _tunnel_close(self, conn_id: int) -> None:
        """Close a tunnel connection."""
        conn = self._tunnel_connections.pop(conn_id, None)
        if conn:
            _, writer = conn
            try:
                writer.close()
                await writer.wait_closed()
            except Exception:
                pass

    async def _close_all_tunnels(self) -> None:
        """Close all tunnel connections."""
        for conn_id in list(self._tunnel_connections.keys()):
            await self._tunnel_close(conn_id)

    # ── Communication ────────────────────────────────────────────────────────

    async def _send(self, msg: dict[str, Any]) -> None:
        """Send a JSON message to the server."""
        if self.ws:
            try:
                await self.ws.send(json.dumps(msg))
            except (websockets.ConnectionClosed, Exception):
                pass

    def stop(self) -> None:
        """Signal the client to stop."""
        self.running = False
        self.cancel_event.set()


# ── Entry Point ──────────────────────────────────────────────────────────────

def main() -> None:
    server_url = os.environ.get("EASYBOT_SERVER")
    token = os.environ.get("EASYBOT_TOKEN")

    if not server_url or not token:
        print("Error: EASYBOT_SERVER and EASYBOT_TOKEN environment variables are required.")
        print()
        print("Usage:")
        print("  set EASYBOT_SERVER=wss://app.example.com/ws/windows-agent")
        print("  set EASYBOT_TOKEN=eb_wag_...")
        print("  easybot-windows-agent")
        sys.exit(1)

    log_level = os.environ.get("LOG_LEVEL", "INFO").upper()
    logging.basicConfig(
        level=getattr(logging, log_level, logging.INFO),
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )

    logger.info("EasyBot Windows Agent v%s starting", __version__)
    logger.info("Server: %s", server_url)
    logger.info("Token: %s...", token[:12] if len(token) > 12 else "***")
    logger.info("Available tools: %s", ", ".join(TOOLS.keys()))

    client = AgentClient(server_url, token)

    def handle_signal(*_: Any) -> None:
        logger.info("Shutdown signal received")
        client.stop()

    signal.signal(signal.SIGINT, handle_signal)
    signal.signal(signal.SIGTERM, handle_signal)

    asyncio.run(client.run())
    logger.info("Agent stopped")


if __name__ == "__main__":
    main()