easybot-windows-agent/easybot_windows_agent/main.py

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