easybot-windows-agent/easybot_windows_agent/main.py

808 lines
30 KiB
Python

"""
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)}
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 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)
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 UI annotations: cursor marker + focused element highlight."""
from PIL import ImageGrab, ImageDraw
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()
# 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
# 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()