orchestrata_ultimate.drivers.browser_driver

Reference › orchestrata_ultimate › drivers › browser_driver

browser_driver generates code by driving a free-tier LLM web interface with Playwright's headless Chromium (SPEC-025, REQ-025-051). BrowserDriver is configured with a workspace, a provider name (default qwen_chat) and a headless flag, and takes its selectors, URLs and limits from browser_providers, so the class itself stays generic. execute builds a prompt that includes the contents of the given files and format instructions (_build_prompt, _read_file_context), launches the browser if needed (_ensure_browser, _launch_browser), waits for the provider's rate limit (_enforce_rate_limit), types and submits the prompt and waits for the reply (_submit_and_wait). It then extracts FILE: blocks, or fallback markdown code blocks, with extraction and writes them into the workspace (_process_response, _write_file_blocks).

Sessions are managed: _is_session_expired checks the provider's maximum session length, _save_storage_state saves the login state for reuse, and close closes the browser. _recover_session applies three escalating levels of recovery (page navigation, browser restart and full re-authentication), up to a fixed number of attempts with a short wait between them. should_inject tells the driver factory whether a provider name maps to a browser driver. Tests are in tests/drivers/test_browser_driver.py and tests/drivers/test_browser_driver_extended.py.

At a glance

Class diagram

Dependency diagram

API

Classes

BrowserDriver

Decorators: traces(spec='025', reqs=['REQ-025-051']) · Line 34

Async code gen driver using Playwright headless Chromium.

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/drivers/browser_driver.py

# mypy: ignore-errors
"""Async browser-based code generation driver.

Uses Playwright headless Chromium to interact with free-tier LLM
web UIs (Qwen Chat, DeepSeek Chat, Kimi). Supports 3-tier session
recovery: page navigation, browser restart, and full re-auth.
"""

import asyncio
import logging
import os
import time

from orchestrata_ultimate.core.traceability import traces
from orchestrata_ultimate.drivers.browser_providers import (
    BrowserProviderConfig,
    get_provider_config,
)
from orchestrata_ultimate.drivers.extraction import (
    extract_file_blocks,
    extract_markdown_code_blocks,
)

_logger = logging.getLogger(__name__)

_RECOVERY_WAIT_SECONDS = 3
_MAX_RECOVERY_ATTEMPTS = 3


@traces(
    spec="025",
    reqs=["REQ-025-051"],
)
class BrowserDriver:
    """Async code gen driver using Playwright headless Chromium.

    Connects to free-tier LLM web UIs, submits coding prompts,
    waits for responses, and extracts generated code.

    Attributes:
        workspace: Directory where generated code is written.
        provider_name: Name key of the browser provider.
        headless: Whether to run the browser in headless mode.
    """

    @traces(spec="025", story="US-025-001", reqs=["REQ-025-052", "REQ-025-053"])
    def __init__(
        self,
        workspace: str | None = None,
        provider: str = "qwen_chat",
        headless: bool = True,
        environ: dict[str, str] | None = None,
    ) -> None:
        """Initialize with target workspace and provider config.

        Args:
            workspace: Directory to write generated files into.
            provider: Browser provider key from BROWSER_PROVIDERS.
            headless: Run browser headless (True) or visible (False).
            environ: Optional environment mapping for testing.
        """
        env = environ or os.environ
        self.workspace = workspace or env.get(
            "ORCHESTRATA_WORKSPACE", "/home/mohamed/workspace"
        )
        os.makedirs(self.workspace, exist_ok=True)
        self.provider_name = provider
        self.headless = headless
        self._config: BrowserProviderConfig = get_provider_config(provider)
        self._browser = None
        self._page = None
        self._session_start: float = 0.0
        self._last_request: float = 0.0
        self._storage_path = os.path.join(
            self.workspace,
            ".browser_state",
            f"{provider}_state.json",
        )

    @traces(
        spec="025",
        story="US-025-001",
        reqs=["REQ-025-058", "REQ-025-060", "REQ-025-064"],
    )
    async def execute(
        self,
        task: str,
        files: list[str] | None = None,
        is_fix: bool = False,
    ) -> tuple[bool, str]:
        """Submit a coding task to the browser-based LLM.

        Args:
            task: The coding task/prompt.
            files: Optional file paths for context.
            is_fix: True if this is a repair prompt.

        Returns:
            Tuple of (success, output_text).
        """
        try:
            await self._ensure_browser()
        except Exception as exc:
            return False, f"Browser init failed: {exc}"

        if self._is_session_expired():
            await self._recover_session(level=2)

        await self._enforce_rate_limit()

        prompt = self._build_prompt(task, files, is_fix)

        for attempt in range(_MAX_RECOVERY_ATTEMPTS):
            try:
                response = await self._submit_and_wait(prompt)
                self._last_request = time.monotonic()
                return self._process_response(response)
            except Exception as exc:
                level = min(attempt + 1, 2)
                _logger.warning(
                    "Attempt %d failed: %s — recovering (level %d)",
                    attempt + 1,
                    exc,
                    level,
                )
                await self._recover_session(level=level)

        return False, "All recovery attempts exhausted"

    @traces(spec="025", story="US-025-001", reqs=["REQ-025-055"])
    async def close(self) -> None:
        """Close browser and save login state."""
        if self._page and self._browser:
            await self._save_storage_state()
        if self._browser:
            try:
                await self._browser.close()
            except Exception as exc:
                _logger.warning("Browser close error: %s", exc)
            self._browser = None
            self._page = None

    @traces(
        spec="025",
        story="US-025-001",
        reqs=["REQ-025-051", "REQ-025-054", "REQ-025-059"],
    )
    async def _ensure_browser(self) -> None:
        """Launch browser if not already running.

        Raises:
            ImportError: If playwright is not installed.
        """
        if self._browser and self._page:
            return
        await self._launch_browser()

    @traces(
        spec="025",
        story="US-025-001",
        reqs=["REQ-025-051", "REQ-025-054", "REQ-025-059"],
    )
    async def _launch_browser(self) -> None:
        """Launch Chromium and navigate to the provider URL.

        Raises:
            ImportError: If playwright is not installed.
        """
        from playwright.async_api import async_playwright

        pw = await async_playwright().start()
        browser = await pw.chromium.launch(headless=self.headless)

        # Reuse a previously-saved storage_state (cookies/localStorage) via
        # new_context(), the counterpart of context.storage_state(path=...)
        # in _save_storage_state(). NOTE: launch_persistent_context() takes
        # its state from `user_data_dir` (a real Chromium profile dir), not
        # from a storage_state JSON export — passing both `storage_state`
        # and an empty `user_data_dir=""` to it (as this used to do) is not
        # a real Playwright parameter and always raised TypeError, silently
        # breaking session recovery whenever a saved state file existed.
        if os.path.exists(self._storage_path):
            context = await browser.new_context(storage_state=self._storage_path)
        else:
            context = await browser.new_context()
        self._page = await context.new_page()
        self._browser = browser

        await self._page.goto(
            self._config.base_url,
            timeout=self._config.timeout_seconds * 1000,
        )
        self._session_start = time.monotonic()

    @traces(spec="025", story="US-025-001", reqs=["REQ-025-057"])
    def _is_session_expired(self) -> bool:
        """Check if the current session has exceeded max duration.

        Returns:
            True if the session has expired.
        """
        if self._session_start == 0.0:
            return False
        elapsed = time.monotonic() - self._session_start
        return elapsed > (self._config.max_session_minutes * 60)

    @traces(spec="025", story="US-025-001", reqs=["REQ-025-056"])
    async def _enforce_rate_limit(self) -> None:
        """Wait if needed to respect provider rate limits."""
        if self._last_request == 0.0:
            return
        elapsed = time.monotonic() - self._last_request
        remaining = self._config.rate_limit_seconds - elapsed
        if remaining > 0:
            await asyncio.sleep(remaining)

    async def _submit_and_wait(self, prompt: str) -> str:
        """Type prompt, click submit, and wait for response.

        Args:
            prompt: The full prompt text to submit.

        Returns:
            Raw response text from the LLM web UI.

        Raises:
            TimeoutError: If response does not appear in time.
        """
        cfg = self._config
        timeout_ms = cfg.timeout_seconds * 1000

        # execute() always calls _ensure_browser() (which launches the
        # browser and sets self._page) before this method runs.
        assert self._page is not None
        await self._page.fill(cfg.input_selector, prompt)
        await self._page.click(cfg.submit_selector)

        await self._page.wait_for_selector(
            cfg.completion_indicator,
            timeout=timeout_ms,
            state="attached",
        )

        elements = await self._page.query_selector_all(cfg.response_selector)
        if not elements:
            return ""
        last_element = elements[-1]
        return await last_element.inner_text()

    @traces(
        spec="025",
        story="US-025-001",
        reqs=["REQ-025-065", "REQ-025-066"],
    )
    def _build_prompt(
        self,
        task: str,
        files: list[str] | None,
        is_fix: bool,
    ) -> str:
        """Build full prompt with file context and format instructions.

        Args:
            task: The coding task.
            files: Optional file paths for context.
            is_fix: True if this is a repair prompt.

        Returns:
            The complete prompt string.
        """
        context = ""
        if not is_fix and files:
            context = self._read_file_context(files)

        format_hint = (
            "\n\nOutput ALL files using this format:\n"
            "FILE: path/to/file.py\n```python\n<content>\n```"
        )
        if context:
            return f"Existing files:\n{context}\n\nTask:\n{task}{format_hint}"
        return task + format_hint

    @traces(
        spec="025",
        story="US-025-001",
        reqs=["REQ-025-067", "REQ-025-068"],
    )
    def _read_file_context(self, files: list[str]) -> str:
        """Read file contents for prompt context.

        Args:
            files: List of relative file paths.

        Returns:
            Formatted file context string.
        """
        parts: list[str] = []
        for f in files:
            filepath = os.path.join(self.workspace, f)
            if not os.path.exists(filepath):
                continue
            try:
                with open(filepath, encoding="utf-8") as fh:
                    content = fh.read()
                if len(content) > 5000:
                    content = content[:5000] + "\n# ... truncated"
                parts.append(f"=== {f} ===\n{content}")
            except OSError, UnicodeDecodeError:
                continue
        return "\n".join(parts)

    @traces(
        spec="025",
        story="US-025-001",
        reqs=["REQ-025-069", "REQ-025-070", "REQ-025-071"],
    )
    def _process_response(
        self,
        response: str,
    ) -> tuple[bool, str]:
        """Extract code from response and write files to workspace.

        Args:
            response: Raw response text from the web UI.

        Returns:
            Tuple of (success, output_text).
        """
        file_blocks = extract_file_blocks(response)
        if not file_blocks:
            md_blocks = extract_markdown_code_blocks(response)
            if md_blocks:
                return True, response
            return False, "No code blocks found in response"

        written = self._write_file_blocks(file_blocks)
        if not written:
            return False, "No files written"

        for f in written:
            _logger.info("Wrote %s", f)
        return True, response

    @traces(spec="025", story="US-025-001", reqs=["REQ-025-072"])
    def _write_file_blocks(
        self,
        blocks: list[tuple[str, str]],
    ) -> list[str]:
        """Write extracted FILE: blocks to disk.

        Args:
            blocks: List of (filename, content) tuples.

        Returns:
            List of written file paths.
        """
        written: list[str] = []
        for filename, content in blocks:
            filepath = os.path.join(self.workspace, filename)
            os.makedirs(os.path.dirname(filepath), exist_ok=True)
            try:
                with open(filepath, "w", encoding="utf-8") as fh:
                    fh.write(content)
                written.append(filename)
            except OSError as exc:
                _logger.warning("Failed to write %s: %s", filename, exc)
        return written

    @traces(
        spec="025",
        story="US-025-001",
        reqs=["REQ-025-061", "REQ-025-062", "REQ-025-063"],
    )
    async def _recover_session(self, level: int) -> None:
        """Recover browser session with escalating strategies.

        Level 1: Navigate back to the base URL.
        Level 2: Close browser and relaunch.

        Args:
            level: Recovery level (1 = navigate, 2 = restart).
        """
        _logger.info("Session recovery level %d", level)
        if level == 1 and self._page:
            try:
                await self._page.goto(
                    self._config.base_url,
                    timeout=self._config.timeout_seconds * 1000,
                )
                return
            except Exception as exc:
                _logger.warning(
                    "Level 1 recovery failed: %s — escalating",
                    exc,
                )

        await self.close()
        await asyncio.sleep(_RECOVERY_WAIT_SECONDS)
        await self._launch_browser()

    @traces(spec="025", story="US-025-001", reqs=["REQ-025-055"])
    async def _save_storage_state(self) -> None:
        """Persist login cookies and storage for session reuse."""
        os.makedirs(os.path.dirname(self._storage_path), exist_ok=True)
        if self._page is None:
            return
        try:
            context = self._page.context
            await context.storage_state(path=self._storage_path)
            _logger.info("Saved browser state to %s", self._storage_path)
        except Exception as exc:
            _logger.warning("Failed to save storage state: %s", exc)


def should_inject(provider: str) -> bool:
    """Check if a provider name maps to a browser driver.

    Args:
        provider: Provider string from configuration.

    Returns:
        True if this provider should use BrowserDriver.
    """
    return provider in (
        "qwen_browser",
        "deepseek_browser",
        "kimi_browser",
    )