codex-app-server/cloud/playwright_mcp.py

"""Stdio MCP server exposing headless Chromium to codex.

Protocol: JSON-RPC 2.0 over stdio, one JSON object per line.
Logs go to stderr only -- anything on stdout would corrupt the
JSON-RPC stream.

This file is deliberately thin: all browser logic lives in
``playwright_browser.py``. It can run two ways:
- inside the repo: ``python3 -m cloud.playwright_mcp`` (imports
  ``cloud.playwright_browser``);
- copied flat into a tenant codex-home next to
  ``playwright_browser.py`` (imports plain ``playwright_browser``
  from its own directory).
The ``ensure_codex_home_playwright`` helper does the copying;
proxy vars and PLAYWRIGHT_BROWSERS_PATH arrive via the MCP env
block in config.toml, never via tool arguments.

One Browser instance lives for the whole MCP process: navigate
once, then snapshot/click/fill/screenshot against the current
page. browser_new_page / browser_list_pages /
browser_switch_page / browser_close_page manage multiple pages;
each new page gets a fresh Playwright context, so cookies are
isolated between pages.
"""

import importlib
import json
import os
import sys

# cloud 包优先、租户 flat copy 回落：同一段代码在两种部署形态下
# 都要跑。importlib 动态取属性而不是两次 from-import：两次
# from-import 会让 mypy 报 no-redef（cloud.* 与顶层模块名不同），
# 而属性绑定在两条路径下都是同一组名字；cloud 导入本身的
# ImportError（含内层缺依赖）同样回落，与原来 except 语义一致。
try:
    _mcp_backend = importlib.import_module("cloud.playwright_browser")
except ImportError:  # flat copy in a tenant codex-home
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    _mcp_backend = importlib.import_module("playwright_browser")
Browser = _mcp_backend.Browser
BrowserError = _mcp_backend.BrowserError

SERVER_NAME = "playwright"
SERVER_VERSION = "0.1.0"

_TOOLS = [
    {
        "name": "browser_navigate",
        "description": (
            "Open a URL in headless Chromium and return its final "
            "URL and title. Use for pages that need JavaScript to "
            "render, which web_fetch cannot see. Network goes "
            "through the egress proxy, so domains outside the "
            "Tinyproxy whitelist fail with a proxy 403 -- report "
            "that plainly, do not bypass it."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "url": {"type": "string",
                        "description": "http(s) URL, or file:// URL "
                        "for a page already in the worktree"},
            },
            "required": ["url"],
        },
    },
    {
        "name": "browser_snapshot",
        "description": (
            "Return the current page's title, URL and body text "
            "(truncated at 8000 characters with a flag). Call "
            "after navigate, and after click/fill, to see what "
            "actually rendered."
        ),
        "inputSchema": {"type": "object", "properties": {}},
    },
    {
        "name": "browser_click",
        "description": (
            "Click an element by CSS selector or Playwright "
            "selector (e.g. 'button.submit', 'text=Next'). "
            "Then snapshot to see the result. When "
            "BROWSER_ALLOWED_HOSTS is set, an anchor click to "
            "a non-listed host is refused BEFORE the click; a "
            "click whose target cannot be resolved beforehand "
            "(button, JS handler) is reported after it lands "
            "on a blocked host instead."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "selector": {"type": "string"},
            },
            "required": ["selector"],
        },
    },
    {
        "name": "browser_fill",
        "description": (
            "Fill an input/textarea by CSS selector with text. "
            "Then snapshot or click to submit."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "selector": {"type": "string"},
                "text": {"type": "string"},
            },
            "required": ["selector", "text"],
        },
    },
    {
        "name": "browser_evaluate",
        "description": (
            "Run a JavaScript expression in the page and return "
            "its string value (truncated at 4000 characters). "
            "Same sandbox privilege as the shell tool; prefer "
            "snapshot/click/fill for ordinary reads and actions."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "expression": {"type": "string"},
            },
            "required": ["expression"],
        },
    },
    {
        "name": "browser_click_at",
        "description": (
            "Click at viewport coordinates (CSS pixels, origin "
            "top-left of the 1280x800 viewport). This is the tool "
            "a live viewer drives: the user clicks on a screenshot "
            "and the coordinates arrive here. Then snapshot to see "
            "the result. Coordinates identify no element, so "
            "there is no pre-click allowlist gate here: when "
            "BROWSER_ALLOWED_HOSTS is set, a landing on a "
            "blocked host is reported after the click, and the "
            "egress proxy 403 is the backstop that stops the load."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "x": {"type": "number"},
                "y": {"type": "number"},
            },
            "required": ["x", "y"],
        },
    },
    {
        "name": "browser_live_shot",
        "description": (
            "Screenshot to the FIXED live-view file "
            "browser-live.png plus a browser-live.json status "
            "sidecar (url, title, timestamp), so a watcher can "
            "poll one stable path. Call after every navigate / "
            "click / fill while a live view is wanted. Returns "
            "path, bytes, url, title and timestamp."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "output_filename": {"type": "string",
                                    "description": "Plain basename "
                                    "to save as. Default: "
                                    "browser-live.png"},
            },
        },
    },
    {
        "name": "browser_snapshot_refs",
        "description": (
            "Accessibility-style snapshot: interactive elements "
            "with stable refs (e.g. e3) plus tag, role and name. "
            "Use browser_click_ref to click one; refs are valid "
            "until the next snapshot_refs call."
        ),
        "inputSchema": {"type": "object", "properties": {}},
    },
    {
        "name": "browser_click_ref",
        "description": (
            "Click an element by its ref from browser_snapshot_refs. "
            "More robust than a CSS selector across re-renders. "
            "Same allowlist rule as browser_click: an anchor ref "
            "to a non-listed host is refused before the click "
            "when BROWSER_ALLOWED_HOSTS is set."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {"ref": {"type": "string"}},
            "required": ["ref"],
        },
    },
    {
        "name": "browser_mouse_move",
        "description": "Move the mouse to viewport coordinates (1280x800).",
        "inputSchema": {
            "type": "object",
            "properties": {
                "x": {"type": "number"},
                "y": {"type": "number"},
            },
            "required": ["x", "y"],
        },
    },
    {
        "name": "browser_key_press",
        "description": "Press a key (Playwright key name, e.g. Enter, Tab).",
        "inputSchema": {
            "type": "object",
            "properties": {"key": {"type": "string"}},
            "required": ["key"],
        },
    },
    {
        "name": "browser_type_text",
        "description": "Type text at the current focus.",
        "inputSchema": {
            "type": "object",
            "properties": {"text": {"type": "string"}},
            "required": ["text"],
        },
    },
    {
        "name": "browser_screenshot",
        "description": (
            "Save a PNG screenshot of the current page into the "
            "working directory and return its absolute path and "
            "byte size. Verify the returned path with shell/PIL "
            "before reporting success."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "output_filename": {"type": "string",
                                    "description": "Plain basename "
                                    "to save as, e.g. 'shot.png'. "
                                    "Default: shot-<epoch>.png"},
            },
        },
    },
    {
        "name": "browser_new_page",
        "description": (
            "Open a new page (tab) and make it current. Every "
            "new page gets a fresh Playwright context, so "
            "cookies and storage are isolated from other pages. "
            "Later navigate/snapshot/click/fill/screenshot act "
            "on this page until browser_switch_page."
        ),
        "inputSchema": {"type": "object", "properties": {}},
    },
    {
        "name": "browser_list_pages",
        "description": (
            "List all open pages with index, URL, title and "
            "which one is current. Does not launch a browser "
            "when none is open."
        ),
        "inputSchema": {"type": "object", "properties": {}},
    },
    {
        "name": "browser_switch_page",
        "description": (
            "Make page ``index`` (from browser_list_pages) "
            "current. Later single-page tools act on it. "
            "Switching to a missing index is an error, not a "
            "crash; refs from the previous page are dropped."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "index": {"type": "integer"},
            },
            "required": ["index"],
        },
    },
    {
        "name": "browser_close_page",
        "description": (
            "Close one page (``index``, or the current page "
            "when omitted). Closing the last page is an error; "
            "use browser_close for the whole browser. Closing "
            "the current page switches to the previous page."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "index": {"type": "integer"},
            },
        },
    },
    {
        "name": "browser_close",
        "description": "Close the browser page. Safe to call twice.",
        "inputSchema": {"type": "object", "properties": {}},
    },
]


def _err(code, message):
    return {"code": code, "message": message}


_BROWSER = None
_BRIDGE = None


def _get_browser(browser=None):
    global _BROWSER
    if browser is not None:
        return browser
    if _BROWSER is None:
        _BROWSER = Browser()
    return _BROWSER


def _get_bridge(browser=None):
    """File-bridge remote control, lazily created. Keyed by browser
    identity so tests passing different fake browsers don't share
    one holder flag."""
    global _BRIDGE
    target = _get_browser(browser)
    if _BRIDGE is not None and _BRIDGE.browser is target:
        return _BRIDGE
    try:
        from playwright_remote import RemoteBridge
    except ImportError:
        try:
            from cloud.playwright_remote import RemoteBridge
        except ImportError:
            return None
    _BRIDGE = RemoteBridge(target, getattr(target, "cwd", None))
    return _BRIDGE


def handle_request(request, browser=None):
    """Handle one JSON-RPC request dict. Returns a response dict or
    None for notifications (no id). Tests pass a fake browser."""
    if not isinstance(request, dict):
        return {"jsonrpc": "2.0", "id": None,
                "error": _err(-32700, "request is not an object")}
    method = request.get("method")
    req_id = request.get("id")
    is_notification = "id" not in request

    if method == "initialize":
        if is_notification:
            return None
        return {"jsonrpc": "2.0", "id": req_id,
                "result": {
                    "protocolVersion": request.get("params", {}).get(
                        "protocolVersion", "2024-11-05"),
                    "capabilities": {"tools": {}},
                    "serverInfo": {"name": SERVER_NAME,
                                   "version": SERVER_VERSION}}}
    if method in ("notifications/initialized", "notifications/cancelled"):
        return None
    if method == "ping":
        if is_notification:
            return None
        return {"jsonrpc": "2.0", "id": req_id, "result": {}}
    if method == "tools/list":
        if is_notification:
            return None
        return {"jsonrpc": "2.0", "id": req_id,
                "result": {"tools": _TOOLS}}
    if method == "tools/call":
        if is_notification:
            return None
        params = request.get("params") or {}
        name = params.get("name")
        args = params.get("arguments") or {}
        target = _get_browser(browser)
        remote = _get_bridge(browser)
        if remote is not None and remote.get_holder() == "user":
            return {"jsonrpc": "2.0", "id": req_id,
                    "result": {"content": [{"type": "text",
                                            "text": "error: user holds "
                                            "browser control; wait for "
                                            "release"}],
                               "isError": True}}
        try:
            if name == "browser_navigate":
                result = target.navigate(args.get("url"))
            elif name == "browser_snapshot":
                result = target.snapshot()
            elif name == "browser_click":
                result = target.click(args.get("selector"))
            elif name == "browser_fill":
                result = target.fill(args.get("selector"),
                                     args.get("text"))
            elif name == "browser_click_at":
                result = target.click_at(args.get("x"), args.get("y"))
            elif name == "browser_mouse_move":
                result = target.mouse_move(args.get("x"), args.get("y"))
            elif name == "browser_key_press":
                result = target.key_press(args.get("key"))
            elif name == "browser_type_text":
                result = target.type_text(args.get("text"))
            elif name == "browser_snapshot_refs":
                result = target.snapshot_refs()
            elif name == "browser_click_ref":
                result = target.click_ref(args.get("ref"))
            elif name == "browser_live_shot":
                result = target.live_shot(
                    args.get("output_filename"))
            elif name == "browser_evaluate":
                result = target.evaluate(args.get("expression"))
            elif name == "browser_screenshot":
                result = target.screenshot(
                    args.get("output_filename"))
            elif name == "browser_new_page":
                result = target.new_page()
            elif name == "browser_list_pages":
                result = target.list_pages()
            elif name == "browser_switch_page":
                result = target.switch_page(args.get("index"))
            elif name == "browser_close_page":
                result = target.close_page(args.get("index"))
            elif name == "browser_close":
                result = target.close()
            else:
                return {"jsonrpc": "2.0", "id": req_id,
                        "error": _err(-32602, "unknown tool: %r" % name)}
        except Exception as exc:  # noqa: BLE001 - tool errors are data
            return {"jsonrpc": "2.0", "id": req_id,
                    "result": {"content": [{"type": "text",
                                            "text": "error: %s" % exc}],
                               "isError": True}}
        return {"jsonrpc": "2.0", "id": req_id,
                "result": {"content": [{"type": "text",
                                        "text": json.dumps(
                                            result, ensure_ascii=False)}],
                           "isError": False}}
    if is_notification:
        return None
    return {"jsonrpc": "2.0", "id": req_id,
            "error": _err(-32601, "method not found: %r" % method)}


def _drain_remote(browser):
    """Main-thread only: poll the file bridge (control, inputs,
    frame, state). Called on every loop iteration so remote input
    lands even while no JSON-RPC request is in flight."""
    bridge = _get_bridge(browser)
    if bridge is None:
        return
    try:
        bridge.poll(browser)
    except Exception as exc:  # noqa: BLE001 - bridge best-effort
        print("playwright_mcp: bridge poll error: %s" % exc,
              file=sys.stderr)


class _LineReader:
    """Non-blocking stdin line pump (starvation-free).

    The previous loop used select() + sys.stdin.readline(). That
    starves on pipelined input: one readline() can buffer several
    lines in userspace, after which select() sees an empty fd
    forever and the buffered lines are never read. A client that
    writes notifications/initialized and its next request
    back-to-back (every MCP client does this after initialize)
    hangs the server exactly this way. Reading the raw fd into
    our own buffer delivers every line in arrival order.
    """

    def __init__(self, fd):
        self._fd = fd
        os.set_blocking(fd, False)
        self._buf = b""
        self.eof = False

    def poll(self):
        """Return (lines, drained).

        lines are decoded strings without the newline. A trailing
        partial line is returned once at EOF. drained is True only
        when EOF was reached and nothing is left buffered.
        """
        if not self.eof:
            try:
                chunk = os.read(self._fd, 65536)
            except BlockingIOError:
                chunk = None
            except OSError:
                self.eof = True
                chunk = None
            if chunk == b"":
                self.eof = True
            elif chunk:
                self._buf += chunk
        lines = []
        while b"\n" in self._buf:
            raw, self._buf = self._buf.split(b"\n", 1)
            lines.append(raw.decode("utf-8", "replace"))
        if self.eof and self._buf:
            lines.append(self._buf.decode("utf-8", "replace"))
            self._buf = b""
        return lines, (self.eof and not self._buf)


def main():
    import time
    browser = _get_browser()
    # Frame sink must be set before any navigate so pump_screencast
    # has frames for the bridge to publish. No bridge (missing
    # module) means no sink either, so frames are dropped instead
    # of accumulating unboundedly in a queue nobody drains.
    if _get_bridge(browser) is not None:
        browser.set_frame_sink(_frame_queue())
    reader = _LineReader(sys.stdin.fileno())
    last_drain = 0.0
    while True:
        lines, drained = reader.poll()
        for line in lines:
            line = line.strip()
            if not line:
                continue
            _handle_line(line, browser)
        # Pending requests go first: a tool call used to pay a
        # full drain round of file I/O before executing. Drain
        # on idle iterations as before, plus a 200ms floor while
        # busy so a request burst can never starve remote input
        # (takeover/click) or the state heartbeat for long.
        now = time.time()
        if not lines or now - last_drain >= 0.2:
            _drain_remote(browser)
            last_drain = now
        if drained:
            break
        if not lines:
            # 10ms idle cadence: a pending request or remote input
            # is picked up within ~10ms, and a drain round costs
            # ~0.3ms (change-gated writes since #187), so 100
            # idle iterations/s burn ~3% of one CPU.
            time.sleep(0.01)


def _handle_line(line, browser):
    """Handle one JSON-RPC line; responses go to stdout."""
    try:
        request = json.loads(line)
    except json.JSONDecodeError as exc:
        sys.stdout.write(json.dumps(
            {"jsonrpc": "2.0", "id": None,
             "error": _err(-32700, "parse error: %s" % exc)}) + "\n")
        sys.stdout.flush()
        return
    try:
        response = handle_request(request, browser)
    except Exception as exc:  # noqa: BLE001 - server must not die
        print("playwright_mcp: handler error: %s" % exc,
              file=sys.stderr)
        if "id" in request:
            response = {"jsonrpc": "2.0", "id": request["id"],
                        "error": _err(-32603, "internal error")}
        else:
            response = None
    if response is not None:
        sys.stdout.write(json.dumps(response) + "\n")
        sys.stdout.flush()


def _frame_queue():
    try:
        from playwright_remote import SimpleQueue
    except ImportError:
        try:
            from cloud.playwright_remote import SimpleQueue
        except ImportError:
            import collections
            import threading

            class SimpleQueue:  # noqa: D101 - last-resort fallback
                def __init__(self):
                    self._deque = collections.deque()
                    self._lock = threading.Lock()

                def put(self, item):
                    with self._lock:
                        self._deque.append(item)

                def get_nowait(self):
                    with self._lock:
                        if not self._deque:
                            raise IndexError("queue empty")
                        return self._deque.popleft()
    return SimpleQueue()


if __name__ == "__main__":
    main()