Spaces:
Sleeping
Sleeping
| import asyncio | |
| import base64 | |
| import gzip | |
| import json | |
| import re | |
| import time | |
| import uuid | |
| from collections.abc import Awaitable, Callable | |
| from pathlib import Path | |
| from playwright.async_api import ( | |
| Browser, | |
| BrowserContext, | |
| Error as PlaywrightError, | |
| Page, | |
| Playwright, | |
| TimeoutError as PlaywrightTimeoutError, | |
| async_playwright, | |
| ) | |
| from .config import Settings | |
| class ArenaBridgeError(RuntimeError): | |
| def __init__(self, message: str, *, code: str, status_code: int = 502) -> None: | |
| super().__init__(message) | |
| self.code = code | |
| self.status_code = status_code | |
| class ArenaPlaywrightBridge: | |
| """Drive Arena's normal web UI with Playwright and expose text results to FastAPI.""" | |
| def __init__(self, settings: Settings) -> None: | |
| self.settings = settings | |
| self._playwright: Playwright | None = None | |
| self._browser: Browser | None = None | |
| self._context: BrowserContext | None = None | |
| self._semaphore = asyncio.Semaphore(settings.max_concurrency) | |
| self._model_lock = asyncio.Lock() | |
| self._models: list[str] = [] | |
| self._models_cached_at = 0.0 | |
| self._storage_state: dict | None = self._decode_storage_state( | |
| settings.arena_storage_state_b64 | |
| ) | |
| def ready(self) -> bool: | |
| return bool(self._browser and self._browser.is_connected() and self._context) | |
| def _decode_storage_state(encoded: str) -> dict | None: | |
| if not encoded.strip(): | |
| return None | |
| try: | |
| raw = base64.b64decode(encoded.strip(), validate=True) | |
| if raw.startswith(b"\x1f\x8b"): | |
| raw = gzip.decompress(raw) | |
| state = json.loads(raw.decode("utf-8")) | |
| except Exception as exc: | |
| raise RuntimeError("ARENA_STORAGE_STATE_B64 is not valid base64 JSON") from exc | |
| if not isinstance(state, dict) or not isinstance(state.get("cookies", []), list): | |
| raise RuntimeError("ARENA_STORAGE_STATE_B64 does not contain Playwright storage_state") | |
| return state | |
| def _storage_state_without_auth_cookie(self) -> dict | None: | |
| """Keep account-owner Local Storage consent while forcing a fresh Arena login.""" | |
| if not self._storage_state: | |
| return None | |
| return { | |
| **self._storage_state, | |
| "cookies": [ | |
| cookie | |
| for cookie in self._storage_state.get("cookies", []) | |
| if not str(cookie.get("name") or "").startswith("arena-auth-prod-v1") | |
| ], | |
| } | |
| async def start(self) -> None: | |
| if not self.settings.bridge_api_key: | |
| raise RuntimeError("BRIDGE_API_KEY must be configured") | |
| has_cookie = bool(self.settings.arena_auth_token) | |
| has_storage_state = bool(self._storage_state) | |
| has_credentials = bool(self.settings.arena_email and self.settings.arena_password) | |
| if not has_cookie and not has_storage_state and not has_credentials: | |
| raise RuntimeError( | |
| "Configure ARENA_STORAGE_STATE_B64, ARENA_AUTH_TOKEN, or both " | |
| "ARENA_EMAIL and ARENA_PASSWORD" | |
| ) | |
| self._playwright = await async_playwright().start() | |
| self._browser = await self._playwright.chromium.launch( | |
| headless=True, | |
| args=["--disable-dev-shm-usage"], | |
| ) | |
| try: | |
| if has_storage_state and has_credentials: | |
| # Re-login with a fresh Cookie while retaining the account owner's | |
| # manually established Local Storage/consent state. | |
| await self.login_with_credentials() | |
| elif has_storage_state: | |
| self._context = await self._create_context(storage_state=self._storage_state) | |
| if has_cookie and not await self._read_auth_cookie(self._context): | |
| await self._add_auth_cookie(self._context, self.settings.arena_auth_token) | |
| elif has_credentials: | |
| await self.login_with_credentials() | |
| else: | |
| self._context = await self._create_context() | |
| await self._add_auth_cookie(self._context, self.settings.arena_auth_token) | |
| except Exception: | |
| await self.stop() | |
| raise | |
| async def _create_context(self, *, storage_state: dict | None = None) -> BrowserContext: | |
| if not self._browser or not self._browser.is_connected(): | |
| raise ArenaBridgeError( | |
| "Chromium is not available", code="browser_not_ready", status_code=503 | |
| ) | |
| kwargs = { | |
| "viewport": { | |
| "width": self.settings.viewport_width, | |
| "height": self.settings.viewport_height, | |
| }, | |
| "device_scale_factor": 1.0, | |
| "locale": "en-US", | |
| } | |
| if storage_state: | |
| kwargs["storage_state"] = storage_state | |
| return await self._browser.new_context(**kwargs) | |
| async def _add_auth_cookie(context: BrowserContext, token: str) -> None: | |
| # Large auth values may be split into numbered cookie chunks by web frameworks. | |
| chunk_size = 3500 | |
| values = [token[index : index + chunk_size] for index in range(0, len(token), chunk_size)] | |
| cookies = [] | |
| for index, value in enumerate(values): | |
| name = "arena-auth-prod-v1" if len(values) == 1 else f"arena-auth-prod-v1.{index}" | |
| cookies.append( | |
| { | |
| "name": name, | |
| "value": value, | |
| "url": "https://arena.ai/", | |
| "httpOnly": True, | |
| "secure": True, | |
| "sameSite": "Lax", | |
| } | |
| ) | |
| if cookies: | |
| await context.add_cookies(cookies) | |
| async def _read_auth_cookie(context: BrowserContext) -> str: | |
| cookies = await context.cookies(["https://arena.ai/"]) | |
| for cookie in cookies: | |
| if cookie.get("name") == "arena-auth-prod-v1" and cookie.get("value"): | |
| return str(cookie["value"]) | |
| chunks: list[tuple[int, str]] = [] | |
| for cookie in cookies: | |
| match = re.fullmatch(r"arena-auth-prod-v1\.(\d+)", str(cookie.get("name") or "")) | |
| if match and cookie.get("value"): | |
| chunks.append((int(match.group(1)), str(cookie["value"]))) | |
| return "".join(value for _, value in sorted(chunks)) | |
| async def _first_visible(scope, selector: str): | |
| candidates = scope.locator(selector) | |
| for index in range(await candidates.count()): | |
| candidate = candidates.nth(index) | |
| if await candidate.is_visible(): | |
| return candidate | |
| return None | |
| async def _button_by_accessible_names(cls, scope, names: list[str]): | |
| """Prefer a semantically named button instead of an unrelated submit icon.""" | |
| for name in names: | |
| candidates = scope.get_by_role( | |
| "button", name=re.compile(rf"^\s*{re.escape(name)}\s*$", re.IGNORECASE) | |
| ) | |
| for index in range(await candidates.count()): | |
| candidate = candidates.nth(index) | |
| if await candidate.is_visible(): | |
| return candidate | |
| return None | |
| async def _scope_near_input(page: Page, input_locator): | |
| """Use the input's form/dialog before ever searching the entire Arena page.""" | |
| form = input_locator.locator("xpath=ancestor::form[1]") | |
| if await form.count(): | |
| return form | |
| dialog = input_locator.locator("xpath=ancestor::*[@role='dialog'][1]") | |
| if await dialog.count(): | |
| return dialog | |
| # Some React auth widgets do not use <form> or role="dialog". Walk from | |
| # the field outward and choose the first local container with a visible | |
| # button instead of falling back to an unrelated submit icon on the page. | |
| for depth in range(1, 9): | |
| ancestor = input_locator.locator(f"xpath=ancestor::*[{depth}]") | |
| if not await ancestor.count(): | |
| continue | |
| buttons = ancestor.locator( | |
| 'button[type="submit"], button:has-text("Continue"), ' | |
| 'button:has-text("Next"), button:has-text("Sign in"), ' | |
| 'button:has-text("Log in"), [role="button"]:has-text("Continue")' | |
| ) | |
| for index in range(await buttons.count()): | |
| if await buttons.nth(index).is_visible(): | |
| return ancestor | |
| return page | |
| async def _safe_button_label(button) -> str: | |
| """Return only non-sensitive visible/accessibility metadata for diagnostics.""" | |
| try: | |
| text = " ".join((await button.inner_text()).split()) | |
| if text: | |
| return text[:80] | |
| for attribute in ("aria-label", "title", "name", "type"): | |
| value = str(await button.get_attribute(attribute) or "").strip() | |
| if value: | |
| return value[:80] | |
| except Exception: | |
| pass | |
| return "" | |
| async def login_with_credentials(self) -> str: | |
| """Use Arena's normal non-OAuth email/password form and capture its auth Cookie.""" | |
| if not self.settings.arena_email or not self.settings.arena_password: | |
| raise ArenaBridgeError( | |
| "ARENA_EMAIL and ARENA_PASSWORD are required for automatic login", | |
| code="arena_credentials_missing", | |
| status_code=503, | |
| ) | |
| new_context = await self._create_context( | |
| storage_state=self._storage_state_without_auth_cookie() | |
| ) | |
| page = await new_context.new_page() | |
| page.set_default_timeout(self.settings.navigation_timeout_ms) | |
| deadline = time.monotonic() + self.settings.login_timeout_seconds | |
| stage = "open_arena" | |
| try: | |
| # "commit" is more reliable than waiting for every third-party script | |
| # on a login page. Individual fields are waited for below. | |
| await page.goto( | |
| self.settings.arena_url, | |
| wait_until="commit", | |
| timeout=self.settings.navigation_timeout_ms, | |
| ) | |
| stage = "wait_for_page_body" | |
| await page.locator("body").wait_for(state="attached", timeout=30_000) | |
| stage = "inspect_login_page" | |
| await self._raise_if_blocked(page) | |
| email_input = await self._first_visible(page, self.settings.email_selector) | |
| if email_input is None: | |
| login_control = await self._first_visible(page, self.settings.login_button_selector) | |
| if login_control is None: | |
| sidebar_control = await self._first_visible( | |
| page, self.settings.sidebar_open_selector | |
| ) | |
| if sidebar_control is not None: | |
| stage = "open_sidebar" | |
| await sidebar_control.click() | |
| await asyncio.sleep(0.75) | |
| login_control = await self._first_visible( | |
| page, self.settings.login_button_selector | |
| ) | |
| if login_control is None: | |
| raise ArenaBridgeError( | |
| "Arena login control was not found. Update ARENA_LOGIN_BUTTON_SELECTOR if " | |
| "the third-party page changed.", | |
| code="arena_login_control_not_found", | |
| status_code=503, | |
| ) | |
| stage = "click_login_control" | |
| await login_control.click() | |
| stage = "wait_for_email_field" | |
| while time.monotonic() < deadline: | |
| await self._raise_if_blocked(page) | |
| email_input = await self._first_visible(page, self.settings.email_selector) | |
| if email_input is not None: | |
| break | |
| await asyncio.sleep(0.25) | |
| if email_input is None: | |
| raise ArenaBridgeError( | |
| "Arena email field did not appear", | |
| code="arena_email_field_not_found", | |
| status_code=503, | |
| ) | |
| stage = "fill_email" | |
| # Some React auth forms enable their Continue button only after real | |
| # keyboard/input/blur events, not after a single value assignment. | |
| await email_input.fill("") | |
| await email_input.press_sequentially( | |
| self.settings.arena_email.strip(), delay=30 | |
| ) | |
| await email_input.press("Tab") | |
| await asyncio.sleep(0.5) | |
| stage = "find_password_field" | |
| password_input = await self._first_visible(page, self.settings.password_selector) | |
| if password_input is None: | |
| # Also supports a two-step email -> password form. Prefer the submit | |
| # button in the email field's own form so unrelated page buttons are | |
| # never clicked. | |
| email_scope = await self._scope_near_input(page, email_input) | |
| continue_button = await self._button_by_accessible_names( | |
| email_scope, | |
| ["Continue with email", "Continue", "Next", "Sign in", "Log in"], | |
| ) | |
| if continue_button is None: | |
| continue_button = await self._first_visible( | |
| email_scope, self.settings.login_submit_selector | |
| ) | |
| if continue_button is not None: | |
| stage = "wait_for_email_submit_enabled" | |
| enabled_deadline = min(deadline, time.monotonic() + 15) | |
| while time.monotonic() < enabled_deadline: | |
| if await continue_button.is_enabled(): | |
| break | |
| await asyncio.sleep(0.25) | |
| if not await continue_button.is_enabled(): | |
| disabled_label = await self._safe_button_label(continue_button) | |
| # Report only field metadata, never values. This reveals whether | |
| # the form expects a checkbox/name/etc. without exposing credentials. | |
| field_descriptions: list[str] = [] | |
| diagnostic_scope = email_scope | |
| required_fields = diagnostic_scope.locator( | |
| 'input[required], select[required], textarea[required]' | |
| ) | |
| for field_index in range(min(await required_fields.count(), 12)): | |
| field = required_fields.nth(field_index) | |
| if not await field.is_visible(): | |
| continue | |
| field_descriptions.append( | |
| "/".join( | |
| part | |
| for part in [ | |
| str(await field.get_attribute("type") or "field"), | |
| str(await field.get_attribute("name") or "unnamed"), | |
| str(await field.get_attribute("aria-label") or "no-label"), | |
| ] | |
| if part | |
| ) | |
| ) | |
| safe_fields = ", ".join(field_descriptions) or "none detected" | |
| button_descriptions: list[str] = [] | |
| visible_buttons = diagnostic_scope.locator('button, [role="button"]') | |
| for button_index in range(min(await visible_buttons.count(), 12)): | |
| candidate = visible_buttons.nth(button_index) | |
| if not await candidate.is_visible(): | |
| continue | |
| label = await self._safe_button_label(candidate) | |
| enabled = await candidate.is_enabled() | |
| button_descriptions.append( | |
| f"{label or '<unlabeled>'}({'enabled' if enabled else 'disabled'})" | |
| ) | |
| safe_buttons = ", ".join(button_descriptions) or "none detected" | |
| print( | |
| "Arena email-step button remained disabled: " | |
| f"label={disabled_label!r}, visible required fields={safe_fields}, " | |
| f"visible buttons={safe_buttons}", | |
| flush=True, | |
| ) | |
| raise ArenaBridgeError( | |
| "Arena's email-step submit button stayed disabled after normal keyboard " | |
| "entry and blur events. The form may require another visible field or " | |
| f"account-owner action. Required field metadata: {safe_fields}", | |
| code="arena_email_submit_disabled", | |
| status_code=409, | |
| ) | |
| # The label and enabled state are safe diagnostics; credentials are never logged. | |
| button_label = await self._safe_button_label(continue_button) | |
| print( | |
| f"Arena email-step submit control: label={button_label!r}, enabled=True", | |
| flush=True, | |
| ) | |
| stage = "submit_email_step" | |
| try: | |
| # Keep this action timeout short. Some auth pages start a navigation that | |
| # never reaches Playwright's normal readiness condition even though the | |
| # click itself succeeded. | |
| await continue_button.click(timeout=15_000, no_wait_after=True) | |
| except PlaywrightTimeoutError: | |
| password_after_click = await self._first_visible( | |
| page, self.settings.password_selector | |
| ) | |
| if password_after_click is None: | |
| stage = "submit_email_step_enter_fallback" | |
| try: | |
| if await email_input.is_visible(): | |
| await email_input.press("Enter", timeout=15_000) | |
| except PlaywrightError as exc: | |
| raise ArenaBridgeError( | |
| "Arena's email step could not be submitted by clicking the " | |
| "button or pressing Enter.", | |
| code="arena_email_submit_failed", | |
| status_code=503, | |
| ) from exc | |
| stage = "wait_for_password_field" | |
| while time.monotonic() < deadline: | |
| await self._raise_if_blocked(page) | |
| password_input = await self._first_visible( | |
| page, self.settings.password_selector | |
| ) | |
| if password_input is not None: | |
| break | |
| await asyncio.sleep(0.25) | |
| if password_input is None: | |
| raise ArenaBridgeError( | |
| "Arena did not present a password field. The selected account may use a magic " | |
| "link, OTP, Google OAuth, or a changed login flow that requires user interaction.", | |
| code="arena_password_field_not_found", | |
| status_code=409, | |
| ) | |
| stage = "fill_password" | |
| await password_input.fill(self.settings.arena_password) | |
| stage = "find_login_submit" | |
| password_scope = await self._scope_near_input(page, password_input) | |
| submit = await self._button_by_accessible_names( | |
| password_scope, | |
| ["Sign in", "Log in", "Continue", "Submit"], | |
| ) | |
| if submit is None: | |
| submit = await self._first_visible( | |
| password_scope, self.settings.login_submit_selector | |
| ) | |
| if submit is None: | |
| raise ArenaBridgeError( | |
| "Arena login submit button was not found", | |
| code="arena_login_submit_not_found", | |
| status_code=503, | |
| ) | |
| stage = "submit_credentials" | |
| try: | |
| await submit.click(timeout=15_000, no_wait_after=True) | |
| except PlaywrightTimeoutError: | |
| token_after_click = await self._read_auth_cookie(new_context) | |
| if not token_after_click: | |
| stage = "submit_credentials_enter_fallback" | |
| try: | |
| if await password_input.is_visible(): | |
| await password_input.press("Enter", timeout=15_000) | |
| except PlaywrightError as exc: | |
| raise ArenaBridgeError( | |
| "Arena's credential form could not be submitted by clicking the button " | |
| "or pressing Enter.", | |
| code="arena_credentials_submit_failed", | |
| status_code=503, | |
| ) from exc | |
| stage = "wait_for_auth_cookie" | |
| token = "" | |
| while time.monotonic() < deadline: | |
| await self._raise_if_blocked(page) | |
| token = await self._read_auth_cookie(new_context) | |
| if token: | |
| break | |
| await asyncio.sleep(0.5) | |
| if not token: | |
| raise ArenaBridgeError( | |
| "Arena login did not create arena-auth-prod-v1 before timeout. Check the " | |
| "credentials and whether the account requires MFA, OTP, email confirmation, " | |
| "or interactive verification.", | |
| code="arena_login_failed", | |
| status_code=401, | |
| ) | |
| old_context = self._context | |
| self._context = new_context | |
| self.settings.arena_auth_token = token | |
| self._storage_state = await new_context.storage_state() | |
| self._models = [] | |
| self._models_cached_at = 0.0 | |
| if old_context: | |
| async def close_old_context_later() -> None: | |
| await asyncio.sleep(60) | |
| try: | |
| await old_context.close() | |
| except Exception: | |
| pass | |
| asyncio.create_task(close_old_context_later()) | |
| return token | |
| except ArenaBridgeError: | |
| await new_context.close() | |
| raise | |
| except PlaywrightTimeoutError as exc: | |
| await new_context.close() | |
| raise ArenaBridgeError( | |
| f"Arena login timed out during stage '{stage}'. The page may have changed, " | |
| "loaded too slowly, or displayed an interaction that requires the account owner.", | |
| code=f"arena_login_timeout_{stage}", | |
| status_code=504, | |
| ) from exc | |
| except PlaywrightError as exc: | |
| await new_context.close() | |
| raise ArenaBridgeError( | |
| f"Playwright failed during Arena login stage '{stage}': {type(exc).__name__}", | |
| code=f"arena_login_playwright_error_{stage}", | |
| status_code=503, | |
| ) from exc | |
| except Exception: | |
| await new_context.close() | |
| raise | |
| finally: | |
| if not page.is_closed(): | |
| await page.close() | |
| async def refresh_cookie_with_current_session(self) -> str: | |
| """Reload Arena and capture a Cookie refreshed by the site's own session logic.""" | |
| if not self._context: | |
| raise ArenaBridgeError( | |
| "No Arena browser context is active", code="browser_not_ready", status_code=503 | |
| ) | |
| page = await self._context.new_page() | |
| try: | |
| await page.goto( | |
| self.settings.arena_url, | |
| wait_until="domcontentloaded", | |
| timeout=self.settings.navigation_timeout_ms, | |
| ) | |
| await self._raise_if_blocked(page) | |
| await asyncio.sleep(5) | |
| token = await self._read_auth_cookie(self._context) | |
| if not token: | |
| raise ArenaBridgeError( | |
| "Arena did not provide an updated authentication Cookie", | |
| code="arena_cookie_refresh_failed", | |
| status_code=503, | |
| ) | |
| self.settings.arena_auth_token = token | |
| self._storage_state = await self._context.storage_state() | |
| return token | |
| finally: | |
| await page.close() | |
| async def stop(self) -> None: | |
| if self._context: | |
| await self._context.close() | |
| if self._browser: | |
| await self._browser.close() | |
| if self._playwright: | |
| await self._playwright.stop() | |
| self._context = None | |
| self._browser = None | |
| self._playwright = None | |
| async def _new_page(self) -> Page: | |
| if not self.ready or not self._context: | |
| raise ArenaBridgeError( | |
| "Playwright browser is not ready", code="browser_not_ready", status_code=503 | |
| ) | |
| page = await self._context.new_page() | |
| page.set_default_timeout(self.settings.navigation_timeout_ms) | |
| try: | |
| await page.goto( | |
| self.settings.arena_url, | |
| wait_until="domcontentloaded", | |
| timeout=self.settings.navigation_timeout_ms, | |
| ) | |
| await self._raise_if_blocked(page) | |
| await page.locator(self.settings.input_selector).first.wait_for(state="visible") | |
| return page | |
| except ArenaBridgeError: | |
| await page.close() | |
| raise | |
| except PlaywrightTimeoutError as exc: | |
| current_url = page.url | |
| await self._save_failure_artifacts(page) | |
| await page.close() | |
| raise ArenaBridgeError( | |
| "Arena's chat input did not become available. The authentication Cookie may be " | |
| "expired/invalid, or the third-party page layout may have changed. " | |
| f"Current page: {current_url}", | |
| code="arena_page_not_ready", | |
| status_code=503, | |
| ) from exc | |
| except PlaywrightError as exc: | |
| await self._save_failure_artifacts(page) | |
| await page.close() | |
| raise ArenaBridgeError( | |
| "Playwright could not load Arena's chat page", | |
| code="arena_navigation_error", | |
| status_code=503, | |
| ) from exc | |
| async def _raise_if_blocked(self, page: Page) -> None: | |
| title = (await page.title()).strip().lower() | |
| if "just a moment" in title: | |
| raise ArenaBridgeError( | |
| "Arena displayed an interactive security checkpoint. Complete it on the official " | |
| "site with your own browser, then refresh ARENA_AUTH_TOKEN.", | |
| code="interactive_verification_required", | |
| status_code=503, | |
| ) | |
| dialogs = page.locator('[role="dialog"]:visible') | |
| for index in range(await dialogs.count()): | |
| text = (await dialogs.nth(index).inner_text()).strip() | |
| lowered = text.lower() | |
| if "terms of use & privacy policy" in lowered and "agree" in lowered: | |
| raise ArenaBridgeError( | |
| "Arena requires the account owner to review and accept its terms on the official " | |
| "website. This adapter will not accept legal terms automatically.", | |
| code="account_consent_required", | |
| status_code=409, | |
| ) | |
| if "recaptcha" in lowered and ("verification" in lowered or "verify" in lowered): | |
| raise ArenaBridgeError( | |
| "Arena requires interactive reCAPTCHA verification. Complete it manually on the " | |
| "official website; this adapter does not solve CAPTCHA challenges.", | |
| code="interactive_verification_required", | |
| status_code=503, | |
| ) | |
| verification_text = page.get_by_text( | |
| re.compile(r"recaptcha requires verification", re.IGNORECASE) | |
| ) | |
| if await verification_text.count(): | |
| for index in range(await verification_text.count()): | |
| if await verification_text.nth(index).is_visible(): | |
| raise ArenaBridgeError( | |
| "Arena requires interactive reCAPTCHA verification. Complete it manually on " | |
| "the official website.", | |
| code="interactive_verification_required", | |
| status_code=503, | |
| ) | |
| async def _open_model_dialog(self, page: Page): | |
| button = page.locator(f"{self.settings.model_button_selector}:visible").first | |
| if not await button.count(): | |
| raise ArenaBridgeError( | |
| "Arena model selector was not found; the third-party page layout may have changed", | |
| code="model_selector_not_found", | |
| ) | |
| await button.click() | |
| dialog = page.locator('[role="dialog"]:visible').last | |
| try: | |
| await dialog.wait_for(state="visible") | |
| await dialog.locator('[role="listbox"]').wait_for(state="visible") | |
| except PlaywrightTimeoutError as exc: | |
| raise ArenaBridgeError( | |
| "Arena model dialog did not open", code="model_dialog_not_found" | |
| ) from exc | |
| return dialog | |
| def _option_label(raw_text: str) -> str: | |
| return next((line.strip() for line in raw_text.splitlines() if line.strip()), "") | |
| async def list_models(self, *, force: bool = False) -> list[str]: | |
| now = time.monotonic() | |
| if ( | |
| not force | |
| and self._models | |
| and now - self._models_cached_at < self.settings.model_cache_seconds | |
| ): | |
| return list(self._models) | |
| async with self._model_lock: | |
| now = time.monotonic() | |
| if ( | |
| not force | |
| and self._models | |
| and now - self._models_cached_at < self.settings.model_cache_seconds | |
| ): | |
| return list(self._models) | |
| async with self._semaphore: | |
| page = await self._new_page() | |
| try: | |
| dialog = await self._open_model_dialog(page) | |
| options = dialog.locator('[role="option"]') | |
| models: list[str] = [] | |
| for index in range(await options.count()): | |
| label = self._option_label(await options.nth(index).inner_text()) | |
| if label and label.casefold() not in {item.casefold() for item in models}: | |
| models.append(label) | |
| await page.keyboard.press("Escape") | |
| if not models: | |
| raise ArenaBridgeError( | |
| "Arena returned an empty model selector", code="empty_model_list" | |
| ) | |
| self._models = models | |
| self._models_cached_at = time.monotonic() | |
| return list(models) | |
| finally: | |
| await page.close() | |
| async def _select_model(self, page: Page, requested_model: str) -> str: | |
| button = page.locator(f"{self.settings.model_button_selector}:visible").first | |
| if not await button.count(): | |
| raise ArenaBridgeError( | |
| "Arena model selector was not found", code="model_selector_not_found" | |
| ) | |
| current = self._option_label(await button.inner_text()) | |
| if current.casefold() == requested_model.casefold(): | |
| return current | |
| dialog = await self._open_model_dialog(page) | |
| search = dialog.locator('[role="combobox"]').last | |
| if await search.count(): | |
| await search.fill(requested_model) | |
| options = dialog.locator('[role="option"]') | |
| matched = None | |
| available: list[str] = [] | |
| for index in range(await options.count()): | |
| option = options.nth(index) | |
| label = self._option_label(await option.inner_text()) | |
| if label: | |
| available.append(label) | |
| if label.casefold() == requested_model.casefold(): | |
| matched = option | |
| requested_model = label | |
| break | |
| if matched is None: | |
| await page.keyboard.press("Escape") | |
| raise ArenaBridgeError( | |
| f"Model '{requested_model}' was not found in Arena's current model selector", | |
| code="model_not_found", | |
| status_code=404, | |
| ) | |
| await matched.click() | |
| return requested_model | |
| async def _install_response_observer(self, page: Page, prompt: str) -> None: | |
| await page.evaluate( | |
| """ | |
| (prompt) => { | |
| if (window.__arenaBridgeObserver) window.__arenaBridgeObserver.disconnect(); | |
| const state = { elements: new Set(), prompt }; | |
| window.__arenaBridgeState = state; | |
| const main = document.querySelector('main') || document.body; | |
| const excluded = 'button, textarea, input, form, nav, header, [role="dialog"]'; | |
| const observer = new MutationObserver((mutations) => { | |
| for (const mutation of mutations) { | |
| let element = mutation.target.nodeType === Node.TEXT_NODE | |
| ? mutation.target.parentElement | |
| : mutation.target; | |
| if (!(element instanceof Element) || !main.contains(element)) continue; | |
| if (element.closest(excluded)) continue; | |
| const block = element.closest('p, pre, blockquote, li, h1, h2, h3, h4, h5, h6') || element; | |
| const text = (block.innerText || block.textContent || '').trim(); | |
| // User prompts can be rendered as many smaller DOM blocks. Exclude | |
| // both the complete prompt and sufficiently long fragments of it. | |
| if (!text || text === prompt || text.includes(prompt)) continue; | |
| if (text.length >= 20 && prompt.includes(text)) continue; | |
| state.elements.add(block); | |
| } | |
| }); | |
| observer.observe(main, { subtree: true, childList: true, characterData: true }); | |
| window.__arenaBridgeObserver = observer; | |
| } | |
| """, | |
| prompt, | |
| ) | |
| def _looks_like_prompt_fragment(text: str, prompt: str) -> bool: | |
| normalized_text = re.sub(r"\s+", " ", text.replace("\u200b", " ")).strip() | |
| normalized_prompt = re.sub(r"\s+", " ", prompt.replace("\u200b", " ")).strip() | |
| if len(normalized_text) < 20: | |
| return False | |
| return normalized_text in normalized_prompt or normalized_prompt in normalized_text | |
| def _clean_response_text(text: str, prompt: str) -> str: | |
| text = text.replace("\u200b", "").strip() | |
| if prompt and prompt in text: | |
| text = text.rsplit(prompt, 1)[-1].strip() | |
| ignored_lines = { | |
| "Add files", | |
| "Code", | |
| "Image", | |
| "Search", | |
| "Send message", | |
| "Stop generating", | |
| } | |
| lines = [line.rstrip() for line in text.splitlines()] | |
| while lines and (not lines[-1].strip() or lines[-1].strip() in ignored_lines): | |
| lines.pop() | |
| return "\n".join(lines).strip() | |
| async def _extract_response_text(self, page: Page, prompt: str) -> str: | |
| if self.settings.response_selector: | |
| candidates = page.locator(f"{self.settings.response_selector}:visible") | |
| if await candidates.count(): | |
| return self._clean_response_text(await candidates.last.inner_text(), prompt) | |
| generic_selectors = ( | |
| '[data-message-author-role="assistant"]:visible, ' | |
| '[data-role="assistant"]:visible, ' | |
| '[data-testid*="assistant" i]:visible, ' | |
| 'article [class*="markdown" i]:visible, ' | |
| 'article [class*="prose" i]:visible' | |
| ) | |
| generic = page.locator(generic_selectors) | |
| if await generic.count(): | |
| raw_text = await generic.last.inner_text() | |
| if not self._looks_like_prompt_fragment(raw_text, prompt): | |
| text = self._clean_response_text(raw_text, prompt) | |
| if text and not self._looks_like_prompt_fragment(text, prompt): | |
| return text | |
| observed = await page.evaluate( | |
| """ | |
| () => { | |
| const state = window.__arenaBridgeState; | |
| if (!state) return ''; | |
| const seen = new Set(); | |
| const blocks = []; | |
| for (const element of state.elements) { | |
| if (!element?.isConnected) continue; | |
| const style = getComputedStyle(element); | |
| if (style.display === 'none' || style.visibility === 'hidden') continue; | |
| const text = (element.innerText || element.textContent || '').trim(); | |
| if (!text || text === state.prompt || text.includes(state.prompt) || seen.has(text)) continue; | |
| if (text.length >= 20 && state.prompt.includes(text)) continue; | |
| seen.add(text); | |
| blocks.push({ element, text }); | |
| } | |
| blocks.sort((a, b) => { | |
| const relation = a.element.compareDocumentPosition(b.element); | |
| return relation & Node.DOCUMENT_POSITION_FOLLOWING ? -1 : 1; | |
| }); | |
| return blocks.map(item => item.text).join(String.fromCharCode(10)).trim(); | |
| } | |
| """ | |
| ) | |
| text = self._clean_response_text(str(observed or ""), prompt) | |
| if text and not self._looks_like_prompt_fragment(text, prompt): | |
| return text | |
| # Last-resort fallback for future DOM changes. A configurable | |
| # ARENA_RESPONSE_SELECTOR is preferred when this fallback is inaccurate. | |
| main_text = await page.locator("main").first.inner_text() | |
| if prompt in main_text: | |
| candidate = self._clean_response_text(main_text.rsplit(prompt, 1)[-1], prompt) | |
| if candidate and not self._looks_like_prompt_fragment(candidate, prompt): | |
| return candidate | |
| return "" | |
| async def _is_generating(page: Page) -> bool: | |
| selectors = ( | |
| 'button[aria-label*="stop" i]:visible, ' | |
| 'button[title*="stop" i]:visible, ' | |
| 'button[aria-label*="cancel" i]:visible' | |
| ) | |
| return bool(await page.locator(selectors).count()) | |
| async def _save_failure_artifacts(self, page: Page) -> None: | |
| if not self.settings.save_failure_artifacts: | |
| return | |
| directory = Path(self.settings.artifact_dir) | |
| directory.mkdir(parents=True, exist_ok=True) | |
| artifact_id = uuid.uuid4().hex | |
| try: | |
| await page.screenshot(path=str(directory / f"{artifact_id}.png"), full_page=True) | |
| (directory / f"{artifact_id}.html").write_text( | |
| await page.content(), encoding="utf-8" | |
| ) | |
| except Exception: | |
| pass | |
| async def chat( | |
| self, | |
| *, | |
| model: str, | |
| prompt: str, | |
| on_delta: Callable[[str], Awaitable[None]] | None = None, | |
| ) -> tuple[str, str]: | |
| async with self._semaphore: | |
| stage = "open_chat_page" | |
| page = await self._new_page() | |
| try: | |
| stage = "select_model" | |
| selected_model = await self._select_model(page, model) | |
| stage = "fill_prompt" | |
| input_box = page.locator(self.settings.input_selector).first | |
| await input_box.fill(prompt) | |
| # Install once for same-page submissions. Arena currently changes from | |
| # /text/direct to /c/<conversation>, so it is installed again below. | |
| stage = "install_initial_response_observer" | |
| await self._install_response_observer(page, prompt) | |
| stage = "find_send_button" | |
| send = page.locator(f"{self.settings.send_selector}:visible").first | |
| if not await send.count(): | |
| raise ArenaBridgeError( | |
| "Arena send button was not found", code="send_button_not_found" | |
| ) | |
| initial_url = page.url | |
| stage = "click_send_button" | |
| await send.click(timeout=15_000, no_wait_after=True) | |
| # Sending a first message commonly performs a client-side navigation. | |
| # Waiting briefly and reinstalling the observer prevents | |
| # "execution context was destroyed" errors on the old document. | |
| stage = "wait_for_conversation_navigation" | |
| try: | |
| await page.wait_for_url( | |
| lambda url: str(url) != initial_url, | |
| timeout=10_000, | |
| wait_until="commit", | |
| ) | |
| except PlaywrightTimeoutError: | |
| pass | |
| await asyncio.sleep(0.5) | |
| stage = "install_conversation_response_observer" | |
| observer_installed = False | |
| for _ in range(20): | |
| try: | |
| await self._install_response_observer(page, prompt) | |
| observer_installed = True | |
| break | |
| except PlaywrightError: | |
| if page.is_closed(): | |
| break | |
| await asyncio.sleep(0.25) | |
| if page.is_closed(): | |
| raise ArenaBridgeError( | |
| "Arena closed the conversation page after the message was sent", | |
| code="arena_conversation_page_closed", | |
| status_code=502, | |
| ) | |
| if not observer_installed: | |
| raise ArenaBridgeError( | |
| "Arena's conversation document did not become stable after navigation", | |
| code="arena_conversation_navigation_unstable", | |
| status_code=502, | |
| ) | |
| started_at = time.monotonic() | |
| stable_since = started_at | |
| transient_error_since: float | None = None | |
| previous = "" | |
| while True: | |
| now = time.monotonic() | |
| if now - started_at > self.settings.request_timeout_seconds: | |
| raise ArenaBridgeError( | |
| "Timed out while waiting for Arena's web UI response", | |
| code="upstream_timeout", | |
| status_code=504, | |
| ) | |
| try: | |
| stage = "check_post_send_blockers" | |
| await self._raise_if_blocked(page) | |
| stage = "extract_assistant_response" | |
| current = await self._extract_response_text(page, prompt) | |
| if current != previous: | |
| stable_since = now | |
| if on_delta and current.startswith(previous): | |
| delta = current[len(previous) :] | |
| if delta: | |
| await on_delta(delta) | |
| previous = current | |
| stage = "check_generation_state" | |
| generating = await self._is_generating(page) | |
| transient_error_since = None | |
| except PlaywrightError as exc: | |
| # React/Next.js can briefly replace the execution context while | |
| # establishing the conversation route. Retry transient errors, | |
| # but retain a bounded failure window for real page breakage. | |
| if page.is_closed(): | |
| raise ArenaBridgeError( | |
| "Arena closed the conversation page while generating", | |
| code="arena_conversation_page_closed", | |
| status_code=502, | |
| ) from exc | |
| if transient_error_since is None: | |
| transient_error_since = now | |
| if now - transient_error_since > 10: | |
| safe_detail = str(exc).splitlines()[0][:240] | |
| raise ArenaBridgeError( | |
| f"Playwright repeatedly failed during chat stage '{stage}': " | |
| f"{safe_detail}", | |
| code=f"playwright_page_error_{stage}", | |
| status_code=502, | |
| ) from exc | |
| await asyncio.sleep(0.25) | |
| continue | |
| if ( | |
| previous | |
| and not generating | |
| and now - stable_since >= self.settings.response_stable_seconds | |
| ): | |
| return previous, selected_model | |
| await asyncio.sleep(0.25) | |
| except ArenaBridgeError: | |
| await self._save_failure_artifacts(page) | |
| raise | |
| except PlaywrightError as exc: | |
| await self._save_failure_artifacts(page) | |
| safe_detail = str(exc).splitlines()[0][:240] | |
| print( | |
| f"Playwright chat failure at stage '{stage}': {type(exc).__name__}: " | |
| f"{safe_detail}", | |
| flush=True, | |
| ) | |
| raise ArenaBridgeError( | |
| f"Playwright failed during chat stage '{stage}': {safe_detail}", | |
| code=f"playwright_page_error_{stage}", | |
| status_code=502, | |
| ) from exc | |
| finally: | |
| if not page.is_closed(): | |
| await page.close() | |