""" EasyBot Windows Agent — main entry point. Connects to the EasyBot runtime via WebSocket, receives tasks, and executes them using the `windows-use` library for GUI automation. 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 json import logging import os import platform import signal import sys import traceback 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 # seconds RECONNECT_DELAY_MAX = 120 # seconds HEARTBEAT_TIMEOUT = 45 # seconds — must be > server's 15s interval # ── 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_task_id: str | None = None self.cancel_event = asyncio.Event() self.llm_config: dict[str, Any] = {} self.max_steps = 25 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, additional_headers=extra_headers, ping_interval=None, # we handle heartbeats manually 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 await self._send({ "type": "client_info", "hostname": platform.node(), "os_version": f"{platform.system()} {platform.version()}", "python_version": platform.python_version(), "client_version": __version__, }) try: async for raw in ws: msg = json.loads(raw) await self._handle_message(msg) finally: self.ws = None self.current_task_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 == "config": self.llm_config = msg.get("llm", {}) self.max_steps = msg.get("max_steps", 25) logger.info("Config received: provider=%s model=%s max_steps=%d", self.llm_config.get("provider"), self.llm_config.get("model"), self.max_steps) elif msg_type == "task": task_id = msg["task_id"] instruction = msg["instruction"] logger.info("Task received: %s — %s", task_id, instruction[:100]) # Run task in background so we keep processing messages asyncio.create_task(self._execute_task(task_id, instruction)) elif msg_type == "cancel": task_id = msg.get("task_id") if self.current_task_id == task_id: logger.info("Cancel requested for task %s", task_id) self.cancel_event.set() else: logger.debug("Unknown message type: %s", msg_type) async def _execute_task(self, task_id: str, instruction: str) -> None: """Execute a windows-use task and report results.""" self.current_task_id = task_id self.cancel_event.clear() try: await self._send({"type": "task_started", "task_id": task_id}) result = await self._run_windows_use(task_id, instruction) await self._send({ "type": "task_completed", "task_id": task_id, "result": result.get("result", "Task completed"), "steps_count": result.get("steps_count", 0), }) logger.info("Task %s completed successfully", task_id) except asyncio.CancelledError: await self._send({ "type": "task_error", "task_id": task_id, "error": "Task was cancelled", }) except Exception as e: error_msg = f"{type(e).__name__}: {e}" logger.error("Task %s failed: %s", task_id, error_msg) await self._send({ "type": "task_error", "task_id": task_id, "error": error_msg, }) finally: self.current_task_id = None async def _run_windows_use(self, task_id: str, instruction: str) -> dict[str, Any]: """Run the windows-use agent with the given instruction.""" from windows_use import Agent, AgentConfig # type: ignore config_kwargs: dict[str, Any] = {} if self.llm_config.get("api_key"): config_kwargs["api_key"] = self.llm_config["api_key"] if self.llm_config.get("model"): config_kwargs["model"] = self.llm_config["model"] if self.llm_config.get("provider"): config_kwargs["provider"] = self.llm_config["provider"] config = AgentConfig( max_steps=self.max_steps, **config_kwargs, ) agent = Agent(config=config) step_count = 0 # Callback for step reporting def on_step(step_info: dict[str, Any]) -> None: nonlocal step_count step_count += 1 asyncio.get_event_loop().call_soon_threadsafe( asyncio.ensure_future, self._send({ "type": "task_step", "task_id": task_id, "step_number": step_count, "action": step_info.get("action", "step"), "details": step_info.get("details", ""), }) ) # Run agent (windows-use is sync internally — run in thread) loop = asyncio.get_event_loop() result_text = await loop.run_in_executor( None, lambda: agent.run(instruction) ) return { "result": str(result_text) if result_text else "Task completed", "steps_count": step_count, } async def _send(self, msg: dict[str, Any]) -> None: """Send a JSON message to the server.""" if self.ws and self.ws.open: try: await self.ws.send(json.dumps(msg)) except websockets.ConnectionClosed: 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 "***") client = AgentClient(server_url, token) # Graceful shutdown 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()