trycua--cua/libs/python/bench-ui/bench_ui/child.py

import asyncio
import json
import os
import random
import socket
import sys
import threading
from pathlib import Path
from typing import Optional

import webview
from aiohttp import web


def _get_free_port() -> int:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("127.0.0.1", 0))
        return s.getsockname()[1]


def _start_http_server(
    window: webview.Window,
    port: int,
    ready_event: threading.Event,
    html_content: str | None = None,
    folder_path: str | None = None,
):
    async def rect_handler(request: web.Request):
        try:
            data = await request.json()
        except Exception:
            return web.json_response({"error": "invalid_json"}, status=400)
        selector = data.get("selector")
        space = data.get("space", "window")
        if not isinstance(selector, str):
            return web.json_response({"error": "selector_required"}, status=400)

        # Ensure window content is loaded
        if not ready_event.is_set():
            # give it a short chance to finish loading
            ready_event.wait(timeout=2.0)
        if not ready_event.is_set():
            return web.json_response({"error": "window_not_ready"}, status=409)

        # Safely embed selector into JS
        selector_js = json.dumps(selector)
        if space == "screen":
            # Compute approximate screen coordinates using window metrics
            js = (
                "(function(){"
                f"const s = {selector_js};"
                "const el = document.querySelector(s);"
                "if(!el){return null;}"
                "const r = el.getBoundingClientRect();"
                "const sx = (window.screenX ?? window.screenLeft ?? 0);"
                "const syRaw = (window.screenY ?? window.screenTop ?? 0);"
                "const frameH = (window.outerHeight - window.innerHeight) || 0;"
                "const sy = syRaw + frameH;"
                "return {x:sx + r.left, y:sy + r.top, width:r.width, height:r.height};"
                "})()"
            )
        else:
            js = (
                "(function(){"
                f"const s = {selector_js};"
                "const el = document.querySelector(s);"
                "if(!el){return null;}"
                "const r = el.getBoundingClientRect();"
                "return {x:r.left,y:r.top,width:r.width,height:r.height};"
                "})()"
            )
        try:
            # Evaluate JS on the target window; this call is thread-safe in pywebview
            result = window.evaluate_js(js)
        except Exception as e:
            return web.json_response({"error": str(e)}, status=500)
        return web.json_response({"rect": result})

    async def eval_handler(request: web.Request):
        try:
            data = await request.json()
        except Exception:
            return web.json_response({"error": "invalid_json"}, status=400)
        code = data.get("javascript") or data.get("code")
        if not isinstance(code, str):
            return web.json_response({"error": "javascript_required"}, status=400)

        if not ready_event.is_set():
            ready_event.wait(timeout=2.0)
        if not ready_event.is_set():
            return web.json_response({"error": "window_not_ready"}, status=409)

        try:
            result = window.evaluate_js(code)
        except Exception as e:
            return web.json_response({"error": str(e)}, status=500)
        return web.json_response({"result": result})

    async def index_handler(request: web.Request):
        if html_content is None:
            return web.json_response({"status": "ok", "message": "bench-ui control server"})
        return web.Response(text=html_content, content_type="text/html")

    app = web.Application()

    # If serving a folder, add static file routes
    if folder_path:
        app.router.add_static("/", folder_path, show_index=True)
    else:
        app.router.add_get("/", index_handler)

    app.router.add_post("/rect", rect_handler)
    app.router.add_post("/eval", eval_handler)

    loop = asyncio.new_event_loop()

    def run_loop():
        asyncio.set_event_loop(loop)
        runner = web.AppRunner(app)
        loop.run_until_complete(runner.setup())
        site = web.TCPSite(runner, "127.0.0.1", port)
        loop.run_until_complete(site.start())
        loop.run_forever()

    t = threading.Thread(target=run_loop, daemon=True)
    t.start()


def main():
    if len(sys.argv) < 2:
        print("Usage: python -m bench_ui.child <config.json>", file=sys.stderr)
        sys.exit(2)

    cfg_path = Path(sys.argv[1])
    cfg = json.loads(cfg_path.read_text(encoding="utf-8"))

    html: Optional[str] = cfg.get("html") or ""
    url: Optional[str] = cfg.get("url")
    folder: Optional[str] = cfg.get("folder")
    title: str = cfg.get("title", "Window")
    x: Optional[int] = cfg.get("x")
    y: Optional[int] = cfg.get("y")
    width: int = int(cfg.get("width", 600))
    height: int = int(cfg.get("height", 400))
    icon: Optional[str] = cfg.get("icon")
    use_inner_size: bool = bool(cfg.get("use_inner_size", False))
    title_bar_style: str = cfg.get("title_bar_style", "default")

    # Choose port early so we can point the window to it when serving inline HTML or folder
    port = _get_free_port()

    # Create window
    if url:
        window = webview.create_window(
            title,
            url=url,
            width=width,
            height=height,
            x=x,
            y=y,
            confirm_close=False,
            text_select=True,
            background_color="#FFFFFF",
        )
        html_for_server = None
        folder_for_server = None
    elif folder:
        # Serve static folder at control server root and point window to index.html
        resolved_url = f"http://127.0.0.1:{port}/index.html"
        window = webview.create_window(
            title,
            url=resolved_url,
            width=width,
            height=height,
            x=x,
            y=y,
            confirm_close=False,
            text_select=True,
            background_color="#FFFFFF",
        )
        html_for_server = None
        folder_for_server = folder
    else:
        # Serve inline HTML at control server root and point window to it
        resolved_url = f"http://127.0.0.1:{port}/"
        window = webview.create_window(
            title,
            url=resolved_url,
            width=width,
            height=height,
            x=x,
            y=y,
            confirm_close=False,
            text_select=True,
            background_color="#FFFFFF",
        )
        html_for_server = html
        folder_for_server = None

    # Track when the page is loaded so JS execution succeeds
    window_ready = threading.Event()

    def _on_loaded():
        window_ready.set()

    window.events.loaded += _on_loaded  # type: ignore[attr-defined]

    # Start HTTP server for control (and optionally serve inline HTML or static folder)
    _start_http_server(
        window, port, window_ready, html_content=html_for_server, folder_path=folder_for_server
    )

    # Print startup info for parent to read
    print(json.dumps({"pid": os.getpid(), "port": port}), flush=True)

    # Start GUI (blocking)
    webview.start(debug=os.environ.get("CUA_BENCH_UI_DEBUG", "false").lower() in ("true", "1"))


if __name__ == "__main__":
    main()