synth/generators/playwright/run_eval.py

"""End-to-end evaluation runner for synthetic *playwright* tasks.

Self-contained: lives entirely in ``synth/generators/playwright/`` and does NOT
touch the shared pipeline (``pipeline.py``) or ``src/mcp_services/playwright/``.
It reuses MCPMark's *real* agent (``src.agents.mcpmark_agent.MCPMarkAgent``) so
every rollout uses the exact same hyperparameters as a stock playwright eval:

    npx -y @playwright/mcp@latest --headless --isolated --no-sandbox \\
        --browser chromium --viewport-size 1280,720
    MAX_TURNS = 100, timeout = 3600s, compaction disabled, reasoning "default"
    model deepseek-v3.2-instruct -> litellm "deepseek/deepseek-chat"

What this runner adds (the glue the no-op playwright state manager does not do):
  * generates the synthetic tasks (build -> description.md / verify.py / index.html)
  * serves every task's index.html on a localhost HTTP server and injects the URL
    into the instruction so the browser agent can actually reach the page
  * rotates a randomly-chosen DeepSeek API key **per task** (per question)
  * runs verify.py with PLAYWRIGHT_TEST_DIR (the page) + MCP_MESSAGES (the chat
    transcript) set, exactly the two channels the verifier reads
  * writes results under ``synth/generators/playwright/runs/<exp>/``

Usage (from repo root, in the dedicated ``mcpmark`` conda env):

    conda run -n mcpmark python -m synth.generators.playwright.run_eval \\
        --n 8 --k 4 --keys-file synth/generators/playwright/deepseek_keys.txt

Keys are loaded (first source that yields any) from: --keys "k1,k2,...",
--keys-file (one per line, '#' comments ok), $DEEPSEEK_API_KEY, or .mcp_env.
"""

import argparse
import datetime
import functools
import http.server
import json
import os
import random
import shutil
import socket
import socketserver
import subprocess
import sys
import threading
import time
from pathlib import Path

# .../mcpmark/synth/generators/playwright/run_eval.py -> mcpmark
REPO_ROOT = Path(__file__).resolve().parents[3]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from synth.llm import LLMClient  # noqa: E402

from . import REGISTRY  # noqa: E402
from .base import start_state_server  # noqa: E402

MCP = "playwright"
LITELLM_MODEL = "openai/deepseek-v4-pro"  # dimcode OpenAI-compatible endpoint
MODEL_NAME = "deepseek-v4-pro"
DEEPSEEK_BASE_URL = "https://dimcode.cn/v1"
# deepseek-v4-pro is a thinking model; litellm rejects a top-level reasoning_effort,
# so these go through extra_body (matching browser_automation.py's reference call).
DEEPSEEK_EXTRA_BODY = {"thinking": {"type": "enabled"}, "reasoning_effort": "high"}


def patch_litellm_extra_body():
    """Inject deepseek-v4-pro's thinking/reasoning_effort via extra_body on every
    litellm call (the agent itself never sets it). Idempotent; call per process."""
    import litellm
    if getattr(litellm, "_synth_eb_patched", False):
        return
    _orig = litellm.acompletion

    async def _patched(*a, **kw):
        if "deepseek-v4" in str(kw.get("model", "")):
            eb = dict(kw.get("extra_body") or {})
            for k, v in DEEPSEEK_EXTRA_BODY.items():
                eb.setdefault(k, v)
            kw["extra_body"] = eb
        return await _orig(*a, **kw)

    litellm.acompletion = _patched
    litellm._synth_eb_patched = True

# Hyperparameters mirrored from a stock MCPMark playwright eval ----------------
MCPMARK_DEFAULTS = dict(
    timeout=3600,                 # pipeline.py --timeout default
    compaction_token=999_999_999,  # compaction disabled
    reasoning_effort="default",
    service_config={              # PlaywrightStateManager defaults -> @playwright/mcp args
        "browser": "chromium",
        "headless": True,
        "viewport_width": 1280,
        "viewport_height": 720,
    },
)


# --------------------------------------------------------------------------- #
# API key loading + per-task rotation
# --------------------------------------------------------------------------- #
def load_keys(args) -> list:
    keys = []
    if args.keys:
        keys = [k.strip() for k in args.keys.split(",") if k.strip()]
    elif args.keys_file:
        p = Path(args.keys_file)
        if p.exists():
            keys = [ln.strip() for ln in p.read_text().splitlines()
                    if ln.strip() and not ln.strip().startswith("#")]
    if not keys and os.environ.get("DEEPSEEK_API_KEY"):
        keys = [os.environ["DEEPSEEK_API_KEY"].strip()]
    if not keys:
        env_file = REPO_ROOT / ".mcp_env"
        if env_file.exists():
            for ln in env_file.read_text().splitlines():
                if ln.strip().startswith("DEEPSEEK_API_KEY"):
                    keys = [ln.split("=", 1)[1].strip().strip('"').strip("'")]
    return [k for k in keys if k]


def mask(key: str) -> str:
    return f"...{key[-4:]}" if len(key) >= 4 else "????"


# The localhost server (base.start_state_server) serves the generated pages over
# GET and records form submissions (write tasks) into <task>/state.json over POST.


# --------------------------------------------------------------------------- #
# Task generation (build the synthetic environment + task files)
# --------------------------------------------------------------------------- #
def build_tasks(args, rng, llm, pages_dir: Path, tasks_dir: Path) -> list:
    """Build n tasks; return list of dicts describing each generated task."""
    types = [t.strip() for t in args.types.split(",") if t.strip()]
    for t in types:
        if t not in REGISTRY:
            raise SystemExit(f"unknown type '{t}'. Choices: {', '.join(REGISTRY)}")

    tasks = []
    for i in range(args.n):
        key = types[i % len(types)]
        gen = REGISTRY[key](difficulty=args.difficulty)
        category = f"{args.prefix}_{key}_{i + 1:02d}"

        env_dir = pages_dir / category          # holds index.html (served + PLAYWRIGHT_TEST_DIR)
        if env_dir.exists():
            shutil.rmtree(env_dir)
        env_dir.mkdir(parents=True, exist_ok=True)

        try:
            spec = gen.build(env_dir, llm, rng)  # may scrape the web (real-web tasks)
        except Exception as exc:
            print(f"  ✗ {category}: build failed — {type(exc).__name__}: {exc} (skipped)")
            shutil.rmtree(env_dir, ignore_errors=True)
            continue

        # Optional oracle self-check (drives a real browser, writes answer.txt,
        # confirms verify accepts it and rejects the untouched env). Needs the
        # python `playwright` package + chromium.
        if not args.skip_self_check:
            passed, detail = self_check(gen, spec, env_dir)
            if not passed:
                print(f"  ✗ {category}: SELF-CHECK FAILED — {detail} (skipped)")
                shutil.rmtree(env_dir, ignore_errors=True)
                continue

        task_dir = tasks_dir / category
        task_dir.mkdir(parents=True, exist_ok=True)
        (task_dir / "description.md").write_text(gen.description(spec), encoding="utf-8")
        (task_dir / "verify.py").write_text(gen.verify_src(spec), encoding="utf-8")
        (task_dir / "meta.json").write_text(json.dumps({
            "task_id": key, "category_id": category,
            "task_name": gen.CATEGORY_NAME, "category_name": gen.CATEGORY_NAME,
            "difficulty": gen.DIFFICULTY, "tags": gen.TAGS, "mcp": [MCP],
        }, indent=2), encoding="utf-8")

        entry = spec.get("page") or (spec.get("list_pages") or spec.get("pages")
                                     or ["index.html"])[0]
        tasks.append({
            "category": category, "key": key, "spec": spec,
            "env_dir": env_dir, "task_dir": task_dir,
            "page": entry,
            "url": spec.get("url"),
            "needs_server": getattr(gen, "NEEDS_SERVER", False),
            "needs_net": getattr(gen, "NEEDS_NET", False),
            "instruction": gen.description(spec),
        })
        print(f"  ✓ {category}: built")
    return tasks


def self_check(gen, spec, env_dir) -> tuple:
    import tempfile
    tmp = Path(tempfile.mkdtemp(prefix="synth_pw_check_"))
    try:
        verify_src = gen.verify_src(spec)
        vfile = tmp / "verify.py"
        vfile.write_text(verify_src, encoding="utf-8")

        neg = tmp / "neg"
        shutil.copytree(env_dir, neg)
        if _run_verify(vfile, neg)[0] == 0:
            return False, "verifier passes on UNSOLVED environment (too weak)"

        pos = tmp / "pos"
        shutil.copytree(env_dir, pos)
        httpd = None
        try:
            if getattr(gen, "NEEDS_SERVER", False):
                # Write task: bring the stateful server up so the oracle's form
                # submission persists to pos/state.json, then verify reads it.
                httpd, port = start_state_server(tmp)
                gen.solve(pos, spec, f"http://127.0.0.1:{port}/pos/")
            else:
                gen.solve(pos, spec)
        except Exception as exc:  # pragma: no cover
            return False, f"oracle solve() raised {type(exc).__name__}: {exc}"
        finally:
            if httpd is not None:
                httpd.shutdown()
        if _run_verify(vfile, pos)[0] != 0:
            return False, "verifier REJECTS the oracle solution"
        return True, "oracle passes, unsolved fails"
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


def _run_verify(verify_file: Path, test_dir: Path, messages_path: Path = None) -> int:
    env = {"PLAYWRIGHT_TEST_DIR": str(test_dir), "PATH": os.environ.get("PATH", "")}
    if messages_path:
        env["MCP_MESSAGES"] = str(messages_path)
        env["PLAYWRIGHT_WORK_DIR"] = str(messages_path.parent)
    proc = subprocess.run(
        [sys.executable, str(verify_file)], env=env,
        capture_output=True, text=True, timeout=120,
    )
    return proc.returncode, proc.stdout, proc.stderr


# --------------------------------------------------------------------------- #
# Real-web tasks: give the agent's @playwright/mcp browser an egress proxy, while
# bypassing localhost so local-synthetic tasks stay direct. Patches the agent at
# runtime — does NOT modify src/.
# --------------------------------------------------------------------------- #
# Heavy image/media origins to block in the agent browser. The real-web pages we
# anchor to (Wikipedia "List of..." tables) carry hundreds of flag/photo images
# served from upload.wikimedia.org; pumping those through one SOCKS proxy blows
# past the 60s browser_navigate timeout while contributing nothing to the table
# TEXT the task needs. Blocking the image host keeps the HTML/a11y tree intact and
# makes these pages load in seconds. arXiv pages are light and unaffected.
BLOCKED_ORIGINS = os.environ.get(
    "SYNTH_BLOCKED_ORIGINS",
    "https://upload.wikimedia.org;https://maps.wikimedia.org",
)


def install_browser_proxy(proxy, bypass="127.0.0.1,localhost,::1"):
    if not proxy:
        return
    from src.agents import mcpmark_agent as _m
    orig = _m.MCPMarkAgent._create_mcp_server
    if getattr(orig, "_proxy_patched", False):
        return

    def patched(self):
        srv = orig(self)
        try:
            if self.mcp_service in ("playwright", "playwright_webarena"):
                extra = ["--proxy-server", proxy, "--proxy-bypass", bypass]
                if BLOCKED_ORIGINS:
                    extra += ["--blocked-origins", BLOCKED_ORIGINS]
                srv.params.args = list(srv.params.args) + extra
        except Exception:
            pass
        return srv

    patched._proxy_patched = True
    _m.MCPMarkAgent._create_mcp_server = patched


# --------------------------------------------------------------------------- #
# Worker: one rollout in its own process. `payload` must be picklable (spawn).
# --------------------------------------------------------------------------- #
def run_one(payload: dict) -> dict:
    task = payload["task"]
    out_dir = Path(payload["out_dir"])
    proxy = payload["proxy"]
    t0 = time.time()

    # Docker reset for write tasks before each rollout
    container_name = task.get("docker_container")
    if container_name:
        if not _docker_restart(container_name):
            return {"category": task["category"], "key_type": task["key"],
                    "run": payload["run"], "success": False,
                    "error": f"docker restart failed for {container_name}"}

    if proxy:
        os.environ["SYNTH_PROXY"] = proxy
        install_browser_proxy(proxy)
    patch_litellm_extra_body()
    from src.agents.mcpmark_agent import MCPMarkAgent

    if task.get("is_webarena"):
        instruction = _webarena_instruction(task)
    elif task.get("needs_net"):
        instruction = task["instruction"] + "\n\nUse Playwright MCP tools to complete this web task."
    else:
        url = f"http://127.0.0.1:{payload['port']}/{task['category']}/{task['page']}"
        instruction = _local_instruction(task, url)

    base = {"category": task["category"], "key_type": task["key"],
            "run": payload["run"], "api_key": mask(payload["api_key"])}

    agent = MCPMarkAgent(
        litellm_input_model_name=LITELLM_MODEL,
        api_key=payload["api_key"],
        base_url=payload["base_url"],
        mcp_service=MCP,
        timeout=MCPMARK_DEFAULTS["timeout"],
        service_config=MCPMARK_DEFAULTS["service_config"],
        reasoning_effort="default",  # extra_body carries thinking/reasoning_effort
        compaction_token=MCPMARK_DEFAULTS["compaction_token"],
    )

    out_dir.mkdir(parents=True, exist_ok=True)
    try:
        result = agent.execute_sync(instruction, str(out_dir / "execution.log"))
    except Exception as exc:
        return {**base, "success": False, "error": f"{type(exc).__name__}: {exc}",
                "agent_time": time.time() - t0}
    agent_time = time.time() - t0

    messages_path = out_dir / "messages.json"
    messages_path.write_text(
        json.dumps(result.get("output", []), indent=2, ensure_ascii=False),
        encoding="utf-8",
    )
    vf = Path(task["task_dir"]) / "verify.py"
    rc, vout, verr = (_run_verify(vf, Path(task["env_dir"]), messages_path)
                      if vf.exists() else (1, "", "no verify.py"))

    res = {**base, "success": rc == 0, "verify_rc": rc,
           "verify_stdout": vout.strip()[-300:], "verify_stderr": verr.strip()[-200:],
           "turn_count": result.get("turn_count"), "token_usage": result.get("token_usage"),
           "litellm_run_model_name": result.get("litellm_run_model_name"),
           "agent_time": agent_time}
    (out_dir / "result.json").write_text(json.dumps(res, indent=2), encoding="utf-8")
    return res


def run_pool(payloads, workers, records, label=""):
    """Run payloads through a spawn-based process pool; collect into `records`."""
    import concurrent.futures as cf
    import multiprocessing as mp

    ctx = mp.get_context("spawn")
    done = 0
    total = len(payloads)
    with cf.ProcessPoolExecutor(max_workers=workers, mp_context=ctx) as ex:
        futs = {ex.submit(run_one, p): p for p in payloads}
        for fut in cf.as_completed(futs):
            try:
                res = fut.result()
            except Exception as exc:
                p = futs[fut]
                res = {"category": p["task"]["category"], "run": p["run"],
                       "success": False, "error": f"worker crashed: {exc}"}
            records.append(res)
            done += 1
            flag = "PASS" if res.get("success") else "FAIL"
            print(f"  [{label} {done}/{total}] {flag} {res['category']} run-{res.get('run')} "
                  f"turns={res.get('turn_count')} {str(res.get('error',''))[:80]}", flush=True)
    return records


def _local_instruction(task, url):
    """Build the instruction for a local-synthetic task (served on localhost)."""
    if task.get("needs_server"):
        # Each rollout starts from a clean slate (no leftover writes from a prior run).
        (task["env_dir"] / "state.json").unlink(missing_ok=True)
        return (
            task["instruction"]
            + f"\n\nThe site is served at: {url}\n"
            + "Navigate there with the browser. Complete the task by filling in and "
            + "submitting the on-page form — your answer is the record you create, "
            + "not a chat reply.\n\n"
            + "Use Playwright MCP tools to complete this web automation task."
        )
    return (
        task["instruction"]
        + f"\n\nThe page is served at: {url}\n"
        + "Navigate to that URL with the browser to view the page.\n"
        + "You do not have any file-writing tools available; provide the required "
        + "answer line(s) as the FINAL lines of your reply.\n\n"
        + "Use Playwright MCP tools to complete this web automation task."
    )


def _webarena_instruction(task):
    """Build the instruction for a WebArena docker task (real container URL)."""
    return (
        task["instruction"]
        + "\n\nUse Playwright MCP tools to complete this web automation task."
    )


def _docker_restart(container_name: str) -> bool:
    """Restart a Docker container and wait for it to be ready."""
    import subprocess, time
    try:
        subprocess.run(["docker", "restart", container_name], check=True,
                       capture_output=True, timeout=30)
        # Wait for container to be ready (max 30s)
        for _ in range(60):
            result = subprocess.run(
                ["docker", "ps", "--filter", f"name=^{container_name}$",
                 "--format", "{{.Status}}"],
                capture_output=True, text=True, timeout=5
            )
            if "Up" in result.stdout:
                time.sleep(1)  # Extra grace period
                return True
            time.sleep(0.5)
        return False
    except Exception as e:
        print(f"Docker restart failed for {container_name}: {e}")
        return False


# --------------------------------------------------------------------------- #
def main():
    ap = argparse.ArgumentParser(description="Run synthetic playwright tasks with DeepSeek.")
    ap.add_argument("--n", type=int, default=8, help="number of tasks to generate")
    ap.add_argument("--k", type=int, default=4, help="rollouts per task (pass@k); mcpmark default 4")
    ap.add_argument("--types", default=",".join(REGISTRY),
                    help="comma-separated task types (default all): " + ", ".join(REGISTRY))
    ap.add_argument("--difficulty", default="medium",
                    choices=["easy", "medium", "hard", "none", "all"])
    ap.add_argument("--seed", type=int, default=1)
    ap.add_argument("--prefix", default="synth")
    ap.add_argument("--no-llm", action="store_true", help="offline page content (no LLM calls)")
    ap.add_argument("--llm-model", default="openai/deepseek-v4-pro",
                    help="model for page-content generation (build phase)")
    ap.add_argument("--keys", default="", help="comma-separated DeepSeek API keys")
    ap.add_argument("--keys-file", default=str(Path(__file__).parent / "deepseek_keys.txt"),
                    help="file with one DeepSeek API key per line")
    ap.add_argument("--base-url", default=DEEPSEEK_BASE_URL, help="DeepSeek base_url")
    ap.add_argument("--exp-name", default=None, help="results subfolder name")
    ap.add_argument("--out", default=str(Path(__file__).parent / "runs"))
    ap.add_argument("--skip-self-check", action="store_true",
                    help="skip the oracle browser self-check during generation")
    ap.add_argument("--proxy", default="socks5://127.0.0.1:7893",
                    help="egress proxy for real-web (NEEDS_NET) tasks; '' to disable")
    ap.add_argument("--workers", type=int, default=32, help="concurrent processes for local tasks")
    ap.add_argument("--net-workers", type=int, default=4,
                    help="concurrent processes for real-web tasks (kept low to avoid throttling)")
    args = ap.parse_args()

    # Real-web tasks scrape (build) and browse (rollout) through this proxy; local
    # tasks bypass it (localhost). Set the env before building so build() can scrape.
    if args.proxy:
        os.environ["SYNTH_PROXY"] = args.proxy
        install_browser_proxy(args.proxy)
    patch_litellm_extra_body()  # deepseek-v4-pro thinking/reasoning_effort

    keys = load_keys(args)
    if not keys:
        raise SystemExit(
            "No DeepSeek API keys found. Provide --keys 'k1,k2', a --keys-file, "
            "$DEEPSEEK_API_KEY, or a DEEPSEEK_API_KEY line in .mcp_env."
        )
    print(f"Loaded {len(keys)} DeepSeek key(s): {', '.join(mask(k) for k in keys)}")

    rng = random.Random(args.seed)
    key_rng = random.Random(args.seed + 1000)  # independent stream for key rotation
    llm = LLMClient(model=args.llm_model, enabled=not args.no_llm, seed=args.seed,
                    base_url=args.base_url)

    exp = args.exp_name or ("run-" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
    exp_dir = Path(args.out) / exp
    pages_dir = exp_dir / "pages"
    tasks_dir = exp_dir / "tasks"
    pages_dir.mkdir(parents=True, exist_ok=True)
    tasks_dir.mkdir(parents=True, exist_ok=True)

    print(f"\n=== Building {args.n} task(s) ===")
    tasks = build_tasks(args, rng, llm, pages_dir, tasks_dir)
    if not tasks:
        raise SystemExit("No tasks survived generation/self-check.")

    httpd, port = start_state_server(pages_dir)
    print(f"Serving pages at http://127.0.0.1:{port}/  ({len(tasks)} task(s))")

    # Build the payload list (k rollouts per task, a random key each), then run
    # local tasks at high concurrency and real-web tasks at low concurrency.
    local_payloads, net_payloads = [], []
    for task in tasks:
        for run_i in range(1, args.k + 1):
            api_key = key_rng.choice(keys)
            payload = {
                "task": task, "api_key": api_key, "base_url": args.base_url,
                "port": port, "proxy": args.proxy,
                "run": run_i, "out_dir": str(exp_dir / f"run-{run_i}" / task["category"]),
            }
            (net_payloads if task.get("needs_net") else local_payloads).append(payload)

    records = []
    try:
        if local_payloads:
            print(f"\n=== Local tasks: {len(local_payloads)} rollouts @ {args.workers} workers ===",
                  flush=True)
            run_pool(local_payloads, args.workers, records, label="local")
        if net_payloads:
            print(f"\n=== Real-web tasks: {len(net_payloads)} rollouts @ {args.net_workers} workers ===",
                  flush=True)
            run_pool(net_payloads, args.net_workers, records, label="web")
    finally:
        httpd.shutdown()

    # Aggregate: pass@1 (avg over rollouts) and pass@k (any rollout passes per task)
    by_task = {}
    for r in records:
        by_task.setdefault(r["category"], []).append(r["success"])
    pass_at_1 = sum(r["success"] for r in records) / max(1, len(records))
    pass_at_k = sum(any(v) for v in by_task.values()) / max(1, len(by_task))
    summary = {
        "model": MODEL_NAME, "litellm_model": LITELLM_MODEL,
        "n_tasks": len(tasks), "k": args.k, "n_rollouts": len(records),
        "n_keys": len(keys), "hyperparameters": MCPMARK_DEFAULTS,
        "pass@1": round(pass_at_1, 4), f"pass@{args.k}": round(pass_at_k, 4),
        "per_task": {c: {"pass_rate": round(sum(v) / len(v), 4), "passes": sum(v), "n": len(v)}
                     for c, v in by_task.items()},
        "created_at": datetime.datetime.now().isoformat(),
    }
    (exp_dir / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")

    print(f"\n=== Summary ===")
    print(f"  pass@1 = {summary['pass@1']:.3f}   pass@{args.k} = {summary[f'pass@{args.k}']:.3f}"
          f"   ({len(records)} rollouts over {len(tasks)} tasks)")
    print(f"  results: {exp_dir}")


if __name__ == "__main__":
    main()