computer_use.py

"""
RUBRA — computer_use.py
Real browser automation via Playwright. $0 cost, fully open-source.

This is what was missing for "I can't fetch links" — browse_url() in
tools.py is plain HTTP+regex (no JS execution), which fails on anything
client-side-rendered (Google Business shares, SPAs, etc). This engine
launches an actual headless Chromium, executes real page actions, and
returns screenshots + DOM for the agent to reason over visually.

Every `goto` is checked against link_security.evaluate_link() BEFORE
navigation — a "high" risk verdict blocks the navigation outright.
"""

import asyncio
import base64
import logging
import threading
from dataclasses import dataclass, field
from typing import Optional, Literal, Any

from playwright.async_api import async_playwright, Browser, BrowserContext, Page, Playwright

try:
    from link_security import evaluate_link
    LINK_SECURITY_OK = True
except ImportError:
    LINK_SECURITY_OK = False
    async def evaluate_link(url):  # fail open on heuristics, never crash the tool
        return {"risk": "safe", "reasons": [], "block": False, "safe_browsing_checked": False}

log = logging.getLogger("rubra.computer_use")

ActionType = Literal["goto", "click", "type", "scroll", "screenshot"]

DOM_SNAPSHOT_MAX_CHARS = 4000
SCREENSHOT_MAX_BYTES   = 2_000_000  # ~2MB cap before downscaling
NAV_TIMEOUT_MS         = 20_000
ACTION_TIMEOUT_MS      = 8_000


@dataclass
class ActionRequest:
    """Schema the agent router must populate. One of selector/coordinates for click."""
    action:      ActionType
    url:         Optional[str] = None                    # goto
    selector:    Optional[str] = None                    # click / type — CSS or XPath
    coordinates: Optional[tuple[float, float]] = None     # click — viewport x,y fallback
    text:        Optional[str] = None                     # type
    direction:   Optional[Literal["up", "down"]] = None    # scroll
    scroll_to:   Optional[int] = None                      # scroll — absolute y position


@dataclass
class ExecutionPayload:
    """Unified return shape for every action — what the agent actually sees."""
    success:        bool
    action:         str
    current_url:    str = ""
    screenshot_b64: str = ""
    dom_snapshot:   str = ""
    error:          Optional[str] = None
    blocked_reason: Optional[str] = None


# ── Interactive DOM extraction — keeps the agent's context window sane ───────
_DOM_EXTRACT_JS = """
() => {
    const out = [];
    const seen = new Set();

    const pushEl = (tag, el) => {
        const text = (el.innerText || el.value || el.getAttribute('aria-label') || el.getAttribute('placeholder') || '').trim().slice(0, 80);
        if (!text) return;
        const key = tag + ':' + text;
        if (seen.has(key)) return;
        seen.add(key);
        const rect = el.getBoundingClientRect();
        if (rect.width === 0 || rect.height === 0) return;
        out.push(`[${tag}] "${text}" @ (${Math.round(rect.x + rect.width/2)},${Math.round(rect.y + rect.height/2)})`);
    };

    document.querySelectorAll('button, a[href], [role="button"]').forEach(el => pushEl('button', el));
    document.querySelectorAll('input, textarea, select').forEach(el => pushEl('input', el));
    document.querySelectorAll('h1, h2, h3').forEach(el => pushEl('heading', el));

    return out.slice(0, 80).join('\\n');
}
"""


class ComputerUseEngine:
    """
    Persistent Playwright engine — one browser, one context, kept alive
    across consecutive agent turns. Do NOT instantiate this per-call;
    use the module-level `engine` singleton below.
    """

    def __init__(self):
        self._playwright: Optional[Playwright]     = None
        self._browser:    Optional[Browser]         = None
        self._context:    Optional[BrowserContext]  = None
        self._page:       Optional[Page]             = None
        self._lock = asyncio.Lock()
        self._started = False

    async def _ensure_started(self):
        if self._started:
            return
        async with self._lock:
            if self._started:  # re-check inside lock — avoid double-launch race
                return
            self._playwright = await async_playwright().start()
            self._browser = await self._playwright.chromium.launch(
                headless=True,
                args=[
                    "--no-sandbox",
                    "--disable-setuid-sandbox",
                    "--disable-dev-shm-usage",
                    "--disable-gpu",
                ],
            )
            self._context = await self._browser.new_context(
                viewport={"width": 1280, "height": 800},
                user_agent="Mozilla/5.0 (compatible; RUBRA/4.0; +https://rubra.ai/bot)",
            )
            self._context.set_default_navigation_timeout(NAV_TIMEOUT_MS)
            self._context.set_default_timeout(ACTION_TIMEOUT_MS)
            self._page = await self._context.new_page()
            self._started = True
            log.info("[COMPUTER_USE] Browser launched and persistent context ready")

    async def _capture(self, action_name: str) -> ExecutionPayload:
        """Shared tail for every action: screenshot + DOM + current URL."""
        try:
            raw_png = await self._page.screenshot(type="png")
            if len(raw_png) > SCREENSHOT_MAX_BYTES:
                # Downscale viewport temporarily for a lighter capture rather
                # than shipping multi-MB base64 into the agent's context.
                raw_png = await self._page.screenshot(type="jpeg", quality=60)
            screenshot_b64 = base64.b64encode(raw_png).decode("ascii")

            dom_snapshot = ""
            try:
                dom_snapshot = await self._page.evaluate(_DOM_EXTRACT_JS)
                dom_snapshot = dom_snapshot[:DOM_SNAPSHOT_MAX_CHARS]
            except Exception as e:
                dom_snapshot = f"[DOM extraction failed: {e}]"

            return ExecutionPayload(
                success=True, action=action_name,
                current_url=self._page.url,
                screenshot_b64=screenshot_b64,
                dom_snapshot=dom_snapshot,
            )
        except Exception as e:
            log.error(f"[COMPUTER_USE] Capture failed: {e}")
            return ExecutionPayload(success=False, action=action_name,
                                     current_url=self._page.url if self._page else "",
                                     error=str(e)[:300])

    # ── Public actions ────────────────────────────────────────────────────────

    async def goto(self, url: str) -> ExecutionPayload:
        # Security check happens BEFORE the browser even launches — a blocked
        # URL shouldn't cost a Chromium spin-up, and on first-ever call this
        # also avoids launching a browser at all if the very first link is bad.
        verdict = await evaluate_link(url)
        if verdict["block"]:
            reason = "; ".join(verdict["reasons"]) or "Flagged as high-risk"
            log.warning(f"[COMPUTER_USE] Blocked navigation to {url}: {reason}")
            return ExecutionPayload(
                success=False, action="goto",
                current_url=self._page.url if self._page else "",
                error="Navigation blocked by security filter",
                blocked_reason=reason,
            )

        await self._ensure_started()

        try:
            await self._page.goto(url, wait_until="domcontentloaded")
            # Give SPA/client-rendered pages a brief settle window —
            # this is exactly the case (Google Business shares, etc.)
            # that the old HTTP-only fetch could never handle.
            await self._page.wait_for_timeout(1200)
        except Exception as e:
            return ExecutionPayload(success=False, action="goto",
                                     current_url=self._page.url if self._page else url,
                                     error=f"Navigation failed: {str(e)[:200]}")

        return await self._capture("goto")

    async def click(self, selector: Optional[str] = None,
                     coordinates: Optional[tuple] = None) -> ExecutionPayload:
        await self._ensure_started()
        try:
            if selector:
                await self._page.click(selector, timeout=ACTION_TIMEOUT_MS)
            elif coordinates:
                x, y = coordinates
                await self._page.mouse.click(x, y)
            else:
                return ExecutionPayload(success=False, action="click",
                                         current_url=self._page.url,
                                         error="No selector or coordinates provided")
            await self._page.wait_for_timeout(500)
        except Exception as e:
            return ExecutionPayload(success=False, action="click",
                                     current_url=self._page.url,
                                     error=f"Click failed: {str(e)[:200]}")
        return await self._capture("click")

    async def type_text(self, selector: str, text: str) -> ExecutionPayload:
        await self._ensure_started()
        try:
            await self._page.fill(selector, text, timeout=ACTION_TIMEOUT_MS)
        except Exception as e:
            return ExecutionPayload(success=False, action="type",
                                     current_url=self._page.url,
                                     error=f"Type failed: {str(e)[:200]}")
        return await self._capture("type")

    async def scroll(self, direction: Optional[str] = None,
                      scroll_to: Optional[int] = None) -> ExecutionPayload:
        await self._ensure_started()
        try:
            if scroll_to is not None:
                await self._page.evaluate(f"window.scrollTo(0, {int(scroll_to)})")
            elif direction == "down":
                await self._page.evaluate("window.scrollBy(0, window.innerHeight * 0.8)")
            elif direction == "up":
                await self._page.evaluate("window.scrollBy(0, -window.innerHeight * 0.8)")
            await self._page.wait_for_timeout(300)
        except Exception as e:
            return ExecutionPayload(success=False, action="scroll",
                                     current_url=self._page.url,
                                     error=f"Scroll failed: {str(e)[:200]}")
        return await self._capture("scroll")

    async def screenshot(self) -> ExecutionPayload:
        await self._ensure_started()
        return await self._capture("screenshot")

    async def shutdown(self):
        """Call on app shutdown — not between agent turns."""
        try:
            if self._context: await self._context.close()
            if self._browser: await self._browser.close()
            if self._playwright: await self._playwright.stop()
        except Exception as e:
            log.warning(f"[COMPUTER_USE] Shutdown error: {e}")
        finally:
            self._started = False
            log.info("[COMPUTER_USE] Browser shut down")


# ── Module-level singleton — persists across agent turns, per the spec ───────
engine = ComputerUseEngine()


async def execute_action(req: ActionRequest) -> dict:
    """
    Single entry point the agent router calls. Decodes the action type,
    dispatches to the engine, returns a plain dict (JSON-serializable for SSE).
    """
    try:
        if req.action == "goto":
            if not req.url:
                payload = ExecutionPayload(success=False, action="goto", error="Missing url")
            else:
                payload = await engine.goto(req.url)

        elif req.action == "click":
            payload = await engine.click(selector=req.selector, coordinates=req.coordinates)

        elif req.action == "type":
            if not req.selector or req.text is None:
                payload = ExecutionPayload(success=False, action="type", error="Missing selector or text")
            else:
                payload = await engine.type_text(req.selector, req.text)

        elif req.action == "scroll":
            payload = await engine.scroll(direction=req.direction, scroll_to=req.scroll_to)

        elif req.action == "screenshot":
            payload = await engine.screenshot()

        else:
            payload = ExecutionPayload(success=False, action=req.action,
                                        error=f"Unknown action: {req.action}")

    except Exception as e:
        log.exception("[COMPUTER_USE] execute_action crashed")
        payload = ExecutionPayload(success=False, action=req.action, error=str(e)[:300])

    return {
        "success":        payload.success,
        "action":         payload.action,
        "current_url":    payload.current_url,
        "screenshot_b64": payload.screenshot_b64,
        "dom_snapshot":   payload.dom_snapshot,
        "error":          payload.error,
        "blocked_reason": payload.blocked_reason,
    }