synth/generators/playwright/base.py

"""Task-type generators for synthetic playwright benchmarks.

Each generator produces a self-consistent web-automation task. The design
mirrors the filesystem generators: the verifier *recomputes* the correct answer
from the resulting environment instead of trusting a stored answer key, and the
state the model produces is **persisted to the test directory** so the oracle
self-check (solve -> verify) can actually observe it.

* ``build(env_dir, llm, rng)`` writes the source ``index.html`` (with a
  machine-readable data layer) into the test environment and returns a ``spec``
  dict describing how to parse it. The ``spec`` deliberately does NOT contain the
  answer.
* ``description(spec)`` renders the natural-language ``description.md`` the model
  will see. The model browses the page with Playwright and records its answer by
  writing ``answer.txt`` into the test directory (falling back to stating the
  answer as the final line of its reply when it cannot write files).
* ``verify_src(spec)`` renders a self-contained ``verify.py`` that re-parses the
  source HTML to recompute the correct answer, reads the answer the model
  submitted (``answer.txt`` first, then the chat transcript via ``MCP_MESSAGES``),
  and compares them. No external answer key.
* ``solve(work_dir, spec)`` is the oracle: it drives a real Playwright browser to
  read the page, computes the answer, and writes ``answer.txt`` so the pipeline
  can prove ``verify.py`` accepts the intended answer.

The description and the verifier are written together so they always agree.
"""

import asyncio
import json
import random
from pathlib import Path
from typing import Dict, List

# --------------------------------------------------------------------------- #
# Shared helpers
# --------------------------------------------------------------------------- #

ANSWER_FILE = "answer.txt"


def _write(path: Path, content: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(content, encoding="utf-8")


# Header shared by every generated verify.py. Pure standard library so the
# verifier is deterministic and never depends on launching a browser; the
# Playwright *browsing* is exercised by the model and by solve().
_PLAYWRIGHT_VERIFY_HEADER = '''#!/usr/bin/env python3
"""Auto-generated verifier (synthetic playwright task). Recomputes the answer
from the source page; do not edit."""
import json
import os
import re
import sys
from pathlib import Path

ANSWER_FILE = "answer.txt"


def get_test_dir() -> Path:
    """Directory that holds the source page (and where the model writes answer.txt).

    Prefers PLAYWRIGHT_TEST_DIR, falls back to FILESYSTEM_TEST_DIR (the env var the
    synth self-check harness sets) and then PLAYWRIGHT_WORK_DIR."""
    for var in ("PLAYWRIGHT_TEST_DIR", "FILESYSTEM_TEST_DIR", "PLAYWRIGHT_WORK_DIR"):
        d = os.environ.get(var)
        if d:
            return Path(d)
    raise ValueError("PLAYWRIGHT_TEST_DIR (or FILESYSTEM_TEST_DIR) is required")


def read_page(filename: str) -> str:
    return (get_test_dir() / filename).read_text(encoding="utf-8")


def get_model_response():
    """Last completed assistant output_text from MCP_MESSAGES (real-eval channel)."""
    messages_path = os.environ.get("MCP_MESSAGES")
    if not messages_path or not Path(messages_path).exists():
        return None
    try:
        with open(messages_path, "r", encoding="utf-8") as f:
            messages = json.load(f)
    except Exception:
        return None
    for message in reversed(messages):
        if message.get("role") != "assistant":
            continue
        content = message.get("content", "")
        if isinstance(content, list):
            for item in content:
                if isinstance(item, dict) and item.get("type") in ("text", "output_text"):
                    return item.get("text", "")
        elif isinstance(content, str) and content:
            return content
    return None


def get_submitted_answer():
    """The answer the model produced: answer.txt if present, else the chat reply."""
    try:
        p = get_test_dir() / ANSWER_FILE
        if p.exists():
            return p.read_text(encoding="utf-8")
    except Exception:
        pass
    return get_model_response()


def check_answer(expected: str):
    """Pass iff the submitted answer equals (file) or contains (chat) `expected`."""
    submitted = get_submitted_answer()
    if submitted is None:
        fail(f"no answer found (neither {ANSWER_FILE} nor a chat reply)")
    s = submitted.strip()
    if s == expected or expected in s:
        ok(f"correct answer: {expected}")
        return
    fail(f"expected '{expected}', got '{s[:200]}'")


def fail(msg):
    print("\\u274c " + msg)
    sys.exit(1)


def ok(msg):
    print("\\u2705 " + msg)


def _norm(s):
    return " ".join(str(s).split()).strip()


def parse_answer_block(text):
    """Parse the multi-field answer the model reports.

    Accepts a ``<answer>...</answer>`` block (real-task format) or, as a
    fallback, a bare set of ``Key|Value`` lines. Tolerates markdown bullets
    (``- Key|Value``). Returns a dict, or None if nothing parseable."""
    if not text:
        return None
    m = re.search(r"<answer>(.*?)</answer>", text, re.IGNORECASE | re.DOTALL)
    body = m.group(1) if m else text
    fields = {}
    for line in body.splitlines():
        line = line.strip().lstrip("-").strip()
        if "|" in line:
            k, v = line.split("|", 1)
            fields[k.strip()] = v.strip()
    return fields or None


def check_fields_text(text, expected, numeric_keys=()):
    """Compare a multi-field answer parsed from `text` against `expected`.

    `numeric_keys` are compared as integers; the rest as normalized text."""
    got = parse_answer_block(text)
    if not got:
        fail("could not find/parse a Key|Value answer block in the submission")
    mism = []
    for k, exp in expected.items():
        actual = got.get(k, "")
        if k in numeric_keys:
            try:
                if int(str(actual).strip()) != int(str(exp).strip()):
                    mism.append(f"{k}: expected {exp}, got {actual!r}")
            except (ValueError, TypeError):
                mism.append(f"{k}: expected numeric {exp}, got {actual!r}")
        elif _norm(actual) != _norm(exp):
            mism.append(f"{k}: expected {exp!r}, got {actual!r}")
    if mism:
        for x in mism:
            print("\\u274c " + x)
        fail(f"{len(mism)} of {len(expected)} field(s) mismatched")
    ok("all fields match: " + ", ".join(f"{k}={v}" for k, v in expected.items()))


def check_fields(expected, numeric_keys=()):
    """Pass iff every expected Key|Value is present in the submitted answer block
    (read from answer.txt or the chat reply)."""
    submitted = get_submitted_answer()
    if submitted is None:
        fail(f"no answer found (neither {ANSWER_FILE} nor a chat reply)")
    check_fields_text(submitted, expected, numeric_keys)


def read_state():
    """Read state.json (what the model created via the stateful server). {} if absent."""
    p = get_test_dir() / "state.json"
    if not p.exists():
        return {}
    try:
        return json.loads(p.read_text(encoding="utf-8"))
    except Exception:
        return {}


def state_events(kind=None):
    """The recorded write events, optionally filtered by their `kind` field."""
    evs = read_state().get("events", [])
    return [e for e in evs if kind is None or e.get("kind") == kind]


def read_pages(filenames):
    """Read several page files from the test dir (for multi-page recompute).

    Returns a list of (filename, html); missing files are skipped."""
    d = get_test_dir()
    out = []
    for fn in filenames:
        p = d / fn
        if p.exists():
            out.append((fn, p.read_text(encoding="utf-8")))
    return out


'''


class Generator:
    KEY = "base"
    CATEGORY_NAME = "Base"
    DIFFICULTY = "L2"
    TAGS: List[str] = []
    # Write tasks need the stateful server up during solve()/rollout so the
    # browser's form submissions are persisted to state.json.
    NEEDS_SERVER = False

    def __init__(self, difficulty: str = "medium"):
        self.difficulty = difficulty

    def build(self, env_dir: Path, llm, rng: random.Random) -> Dict:
        raise NotImplementedError

    def description(self, spec: Dict) -> str:
        raise NotImplementedError

    def verify_src(self, spec: Dict) -> str:
        raise NotImplementedError

    def solve(self, work_dir: Path, spec: Dict) -> None:
        raise NotImplementedError


def _render_verify(body: str, consts: dict) -> str:
    return _PLAYWRIGHT_VERIFY_HEADER + body.replace("__CONSTS__", json.dumps(json.dumps(consts)))


def answer_block(fields: dict) -> str:
    """Render an ordered multi-field answer the way the real tasks expect it."""
    lines = "\n".join(f"{k}|{v}" for k, v in fields.items())
    return f"<answer>\n{lines}\n</answer>"


def kv_lines(fields: dict) -> str:
    """Plain Key|Value lines (used as a submitted post/form body, not a chat answer)."""
    return "\n".join(f"{k}|{v}" for k, v in fields.items())


def diversify_question(llm, base_question, must_include=(), forbid=(), attempts=3):
    """Paraphrase ``base_question`` with an LLM for wording variety, WITHOUT ever
    leaking the answer. The answer is computed elsewhere (deterministic); the LLM
    only sees the question + parameters, never the ground truth.

    Guardrails (any failure -> retry, then fall back to ``base_question``):
      1. every string in ``must_include`` must still appear (URLs, column names,
         conditions) — so the task stays solvable and unambiguous;
      2. no string in ``forbid`` (the ground-truth answer / answer-revealing
         tokens) may appear — so the question can't give the answer away;
      3. output is a single non-empty line of reasonable length.
    """
    if llm is None or not getattr(llm, "enabled", False):
        return base_question
    must = [m for m in must_include if str(m).strip()]
    bad = [f for f in forbid if str(f).strip()]
    prompt = (
        "Rephrase the following web task instruction to vary the wording and make "
        "it sound natural. Keep the MEANING identical — same target, same asked "
        "attribute, same conditions. Do NOT add any new facts, hints, examples, or "
        "the answer. Output ONLY the rephrased instruction as a single line.\n\n"
        + (("You MUST keep these exact strings verbatim: "
            + " | ".join(must) + "\n") if must else "")
        + f"\nInstruction:\n{base_question}"
    )
    for _ in range(attempts):
        out = llm.complete(prompt, temperature=1.0, max_tokens=300)
        if not out:
            break
        cand = " ".join(out.split()).strip().strip('"').strip()
        if not cand or len(cand) < 8 or len(cand) > len(base_question) + 400:
            continue
        low = cand.lower()
        if any(str(m).lower() not in low for m in must):
            continue
        if any(str(b).lower() in low for b in bad):
            continue
        return cand
    return base_question


# --------------------------------------------------------------------------- #
# Stateful localhost server for *write* tasks. Serves the static site (GET) and
# records form submissions (POST to ``<task>/__action__``) into ``<task>/state.json``
# as ``{"events": [ {<form fields>}, ... ]}``. Used by both the oracle self-check
# and the live runner so the browser's writes are actually persisted.
# --------------------------------------------------------------------------- #
def start_state_server(root_dir):
    import functools
    import http.server
    import socket
    import socketserver
    import threading
    import urllib.parse

    root = Path(root_dir)
    lock = threading.Lock()

    class Handler(http.server.SimpleHTTPRequestHandler):
        def log_message(self, *a):  # silence per-request logging
            pass

        def do_POST(self):
            parts = [p for p in self.path.split("?")[0].split("/") if p]
            if not parts or parts[-1] != "__action__":
                self.send_error(404, "unknown action endpoint")
                return
            task = "/".join(parts[:-1])
            length = int(self.headers.get("Content-Length", 0) or 0)
            raw = self.rfile.read(length).decode("utf-8", "replace")
            form = {k: v[0] for k, v in
                    urllib.parse.parse_qs(raw, keep_blank_values=True).items()}
            sp = root / task / "state.json"
            with lock:
                state = {"events": []}
                if sp.exists():
                    try:
                        state = json.loads(sp.read_text(encoding="utf-8"))
                    except Exception:
                        state = {"events": []}
                state.setdefault("events", []).append(form)
                sp.parent.mkdir(parents=True, exist_ok=True)
                sp.write_text(json.dumps(state, indent=2), encoding="utf-8")
            body = b"<!DOCTYPE html><html><body><h1>Submitted</h1>" \
                   b"<p>Your submission was recorded.</p></body></html>"
            self.send_response(200)
            self.send_header("Content-Type", "text/html; charset=utf-8")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)

    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.bind(("127.0.0.1", 0))
    port = sock.getsockname()[1]
    sock.close()
    httpd = socketserver.ThreadingTCPServer(
        ("127.0.0.1", port), functools.partial(Handler, directory=str(root))
    )
    httpd.daemon_threads = True
    threading.Thread(target=httpd.serve_forever, daemon=True).start()
    return httpd, port


# --------------------------------------------------------------------------- #
# Oracle helper: drive a real browser to read the page, then write answer.txt.
# `extract(page)` is an async callable returning the answer string.
# --------------------------------------------------------------------------- #
def solve_via_browser(work_dir: Path, page_file: str, extract) -> None:
    from playwright.async_api import async_playwright

    async def _solve():
        async with async_playwright() as p:
            browser = await p.chromium.launch(headless=True)
            page = await browser.new_page()
            try:
                await page.goto(work_dir.joinpath(page_file).resolve().as_uri())
                answer = await extract(page)
            finally:
                await browser.close()
        _write(work_dir / ANSWER_FILE, answer + "\n")

    asyncio.run(_solve())


_WORDS = (
    "system module config network buffer kernel thread cache socket render "
    "matrix tensor sample dataset gradient logging parser schema invoice client"
).split()

_NAMES = ["Alice", "Bob", "Charlie", "Diana", "Eve", "Frank", "Grace", "Henry", "Ivy", "Jack"]
_LAST_NAMES = ["Smith", "Johnson", "Williams", "Brown", "Jones", "Garcia", "Miller", "Davis"]
_COLORS = ["red", "blue", "green", "yellow", "purple", "orange", "pink", "cyan", "magenta", "lime"]
_PRODUCTS = ["Laptop", "Mouse", "Keyboard", "Monitor", "Headset", "Webcam", "Tablet", "Speaker", "Phone", "Camera"]
_ROLES = ["Developer", "Designer", "Manager", "Analyst", "Tester", "Writer", "Engineer", "Consultant"]


def _para(rng: random.Random, n_lines: int, inject=None) -> list:
    lines = []
    for _ in range(n_lines):
        line = " ".join(rng.choice(_WORDS) for _ in range(rng.randint(4, 9))).capitalize() + "."
        lines.append(line)
    if inject:
        for idx, text in inject:
            lines[idx % len(lines)] = text
    return lines