browser_pool.py

"""
Google AI Mode browser pool — URL-based query, warm-page design.

Key decisions
-------------
* --single-process REMOVED  → was crashing headless Chrome on Linux containers.
* Each query navigates directly to the full search URL with the prompt encoded —
  no typing into the search box needed (faster, no selector fragility).
* Response is extracted from DOM using a ranked list of selectors covering
  AI Overviews, Knowledge Panel, Featured Snippets, and standard results.
* Pages are recycled in the background; the HTTP response is not delayed.
* Sequential page init (not parallel) prevents OOM on free-tier hardware.
"""

import asyncio
import logging
import urllib.parse
from collections import deque

from playwright.async_api import async_playwright

logger = logging.getLogger(__name__)

BASE_URL = "https://www.google.com/search?udm=50&hl=en"

DEFAULT_PRELUDE = (
    "Respond without em dashes. Write in plain, natural, humanized language. "
    "Do not use bold text, headers, or markdown formatting unless explicitly asked. "
    "Be direct and clear."
)

# DOM selectors for Google AI content — tried in order, most specific first
RESPONSE_SELECTORS = [
    ".ILfuVd",                               # AI Overview main text
    ".ezO2md",                               # AI Overview alternate
    ".VkpGBb",                               # AI Overview block
    '[data-attrid="wa:/description"] span',  # Knowledge panel
    ".kno-rdesc span",                       # Knowledge panel alt
    ".hgKElc",                               # Short direct answer
    ".LGOjhe span",                          # Featured snippet
    ".BNeawe span",                          # Short answer card
    ".VwiC3b span",                          # Standard snippet (fallback)
]

# Copy-button selectors (Google's class names rotate; try all)
COPY_BTN_SELECTORS = [
    "button[aria-label='Copy text'].bKxaof",
    "button[aria-label='Copy'].bKxaof",
    "button[aria-label='Copy text']",
    "button[aria-label='Copy']",
]

LOADED_SIGNAL = ",".join(RESPONSE_SELECTORS + COPY_BTN_SELECTORS)

BROWSER_ARGS = [
    "--no-sandbox",
    "--disable-setuid-sandbox",
    "--disable-dev-shm-usage",
    "--disable-gpu",
    "--no-zygote",                              # safe on Linux; no --single-process
    "--disable-blink-features=AutomationControlled",
    "--disable-extensions",
    "--mute-audio",
]


class BrowserPool:
    def __init__(self, max_pages: int = 2):
        self.max_pages = max_pages
        self._playwright = None
        self._browser = None
        self._free: deque = deque()
        self._lock = asyncio.Lock()
        self._available = asyncio.Event()
        self._init_lock = asyncio.Lock()
        self._initialized = False

    # ── Browser helpers ────────────────────────────────────────────────────

    async def _start_browser(self):
        self._playwright = await async_playwright().start()
        self._browser = await self._playwright.chromium.launch(
            headless=True,
            args=BROWSER_ARGS,
        )

    async def _new_context(self):
        return await self._browser.new_context(
            user_agent=(
                "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                "AppleWebKit/537.36 (KHTML, like Gecko) "
                "Chrome/124.0.0.0 Safari/537.36"
            ),
            locale="en-US",
            permissions=["clipboard-read", "clipboard-write"],
            extra_http_headers={"Accept-Language": "en-US,en;q=0.9"},
        )

    async def _open_ready_page(self):
        ctx = await self._new_context()
        page = await ctx.new_page()
        await page.goto(BASE_URL, wait_until="domcontentloaded", timeout=30000)
        logger.info("Google page warmed")
        return page, ctx

    # ── Pool lifecycle ─────────────────────────────────────────────────────

    async def initialize(self):
        async with self._init_lock:
            if self._initialized:
                return
            await self._start_browser()

            # Sequential init — prevents OOM from parallel browser spawning
            for i in range(self.max_pages):
                try:
                    pair = await self._open_ready_page()
                    self._free.append(pair)
                except Exception as e:
                    logger.warning("Page %d warm failed: %s", i, e)

            self._initialized = True
            if self._free:
                self._available.set()
            logger.info("Pool ready: %d/%d pages", len(self._free), self.max_pages)

    # ── Acquisition / release ──────────────────────────────────────────────

    async def _acquire(self):
        while True:
            async with self._lock:
                if self._free:
                    pair = self._free.popleft()
                    if not self._free:
                        self._available.clear()
                    return pair
            await self._available.wait()

    async def _release(self, page, ctx, healthy: bool):
        if healthy:
            try:
                # Quick reset — just navigate back to the base URL
                await page.goto(BASE_URL, wait_until="domcontentloaded", timeout=15000)
                async with self._lock:
                    self._free.append((page, ctx))
                self._available.set()
                return
            except Exception as e:
                logger.warning("Recycle failed, replacing page: %s", e)

        try:
            await ctx.close()
        except Exception:
            pass
        try:
            new_page, new_ctx = await self._open_ready_page()
            async with self._lock:
                self._free.append((new_page, new_ctx))
            self._available.set()
        except Exception as e:
            logger.error("Replacement page failed: %s", e)

    # ── Public API ─────────────────────────────────────────────────────────

    async def run_query(self, system_prompt: str, question: str) -> str:
        if not self._initialized:
            await self.initialize()
        page, ctx = await self._acquire()
        success = False
        try:
            result = await self._execute(page, system_prompt, question)
            success = True
            return result
        finally:
            asyncio.ensure_future(self._release(page, ctx, healthy=success))

    # ── Execution ──────────────────────────────────────────────────────────

    async def _execute(self, page, system_prompt: str, question: str) -> str:
        parts = [DEFAULT_PRELUDE]
        if system_prompt.strip():
            parts.append(system_prompt.strip())
        parts.append(question.strip())
        full_prompt = "\n\n".join(parts)

        # Navigate directly — query in URL, no typing needed
        encoded = urllib.parse.quote(full_prompt)
        await page.goto(
            f"https://www.google.com/search?q={encoded}&udm=50&hl=en",
            wait_until="domcontentloaded",
            timeout=30000,
        )

        # Wait for AI content to appear (up to 60 s for streaming overviews)
        try:
            await page.wait_for_selector(LOADED_SIGNAL, timeout=60000, state="visible")
        except Exception:
            logger.warning("No AI signal within timeout — extracting whatever is there")

        # Small buffer for streaming to settle
        await asyncio.sleep(1.2)

        # Primary: pull text directly from DOM
        result = await page.evaluate(
            """
            (selectors) => {
                for (const sel of selectors) {
                    const els = document.querySelectorAll(sel);
                    for (const el of els) {
                        const text = (el.innerText || el.textContent || '').trim();
                        if (text.length > 40) return text;
                    }
                }
                return '';
            }
            """,
            RESPONSE_SELECTORS,
        )
        if result:
            return result.strip()

        # Fallback: copy-button click → clipboard
        for sel in COPY_BTN_SELECTORS:
            try:
                btns = page.locator(sel)
                n = await btns.count()
                if n:
                    await btns.nth(n - 1).click()
                    await asyncio.sleep(0.5)
                    text = await page.evaluate("navigator.clipboard.readText()")
                    if text and len(text) > 10:
                        return text.strip()
            except Exception:
                continue

        # Last resort: dump visible text from the main results block
        fallback = await page.evaluate(
            """
            () => {
                const root = document.querySelector('#main')
                          || document.querySelector('#search')
                          || document.querySelector('[role="main"]');
                return root ? root.innerText.trim().slice(0, 3000) : '';
            }
            """
        )
        return fallback.strip()