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


# ── 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 the accessibility tree of a window or the desktop."""
    import uiautomation as auto

    window_title = args.get("window")
    max_depth = args.get("depth", 3)

    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()

        def walk(ctrl: Any, depth: int, max_d: int) -> list[dict]:
            if depth > max_d:
                return []
            items = []
            try:
                name = ctrl.Name or ""
                ctrl_type = ctrl.ControlTypeName or ""
                auto_id = ctrl.AutomationId or ""
                rect = ctrl.BoundingRectangle
                value = ""
                try:
                    value = ctrl.GetValuePattern().Value if ctrl.GetValuePattern() else ""
                except Exception:
                    pass

                item: dict[str, Any] = {
                    "type": ctrl_type,
                    "name": name,
                }
                if auto_id:
                    item["automation_id"] = auto_id
                if value:
                    item["value"] = value
                if rect and rect.width() > 0:
                    item["rect"] = {"x": rect.left, "y": rect.top, "w": rect.width(), "h": rect.height()}

                children = []
                if depth < max_d:
                    for child in ctrl.GetChildren():
                        children.extend(walk(child, depth + 1, max_d))
                if children:
                    item["children"] = children

                items.append(item)
            except Exception:
                pass
            return items

        tree = walk(root, 0, max_depth)
        return {"window": root.Name or "Desktop", "tree": tree}
    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)}


async def tool_click_element(args: dict[str, Any]) -> dict[str, Any]:
    """Click a UI element by name, automation_id, or coordinates."""
    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:
        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)
            return {"success": True, "method": "coordinates", "x": x, "y": y}

        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()
            return {"success": True, "method": "ui_automation", "element": ctrl.Name}
        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")

    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:
                    ctrl.GetValuePattern().SetValue(text)
                    return {"success": True, "method": "set_value", "element": element_name}
                except Exception:
                    ctrl.Click()
                    await asyncio.sleep(0.2)

        import pyautogui
        pyautogui.typewrite(text, interval=0.02) if text.isascii() else pyautogui.write(text)
        return {"success": True, "method": "keyboard"}
    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])
        return {"success": True, "keys": keys}
    except Exception as e:
        return {"error": str(e)}


async def tool_screenshot(args: dict[str, Any]) -> dict[str, Any]:
    """Take a screenshot and return as base64 PNG."""
    from PIL import ImageGrab

    try:
        region = args.get("region")
        if region:
            bbox = (region["x"], region["y"], region["x"] + region["w"], region["y"] + region["h"])
            img = ImageGrab.grab(bbox=bbox)
        else:
            img = ImageGrab.grab()

        # 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)}


# ── 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,
}


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

    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:
                    msg = json.loads(raw)
                    await self._handle_message(msg)
            finally:
                self.ws = None
                self.current_call_id = None

    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()

        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

    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()