diff --git a/easybot_windows_agent/main.py b/easybot_windows_agent/main.py new file mode 100644 index 0000000..5bf78a8 --- /dev/null +++ b/easybot_windows_agent/main.py @@ -0,0 +1,264 @@ +""" +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()