From 1b18f0a6bd6ff7e1b4f42a87cd06c8bd2284805d Mon Sep 17 00:00:00 2001 From: ldm0 Date: Mon, 5 Oct 2026 16:42:33 +0800 Subject: [PATCH] refactor(playground): share native WebMCP session client --- moli-playground/webmcp/webmcp_cdp_demo.py | 275 +------------- moli-playground/webmcp/webmcp_client.py | 440 ++++++++++++++++++++++ 2 files changed, 449 insertions(+), 266 deletions(-) create mode 100644 moli-playground/webmcp/webmcp_client.py diff --git a/moli-playground/webmcp/webmcp_cdp_demo.py b/moli-playground/webmcp/webmcp_cdp_demo.py index cbeb6bfec2..6837248b37 100755 --- a/moli-playground/webmcp/webmcp_cdp_demo.py +++ b/moli-playground/webmcp/webmcp_cdp_demo.py @@ -10,27 +10,21 @@ from __future__ import annotations import argparse import asyncio import json -import socket -import subprocess import sys -import tempfile -import time -import urllib.request -from collections import deque -from collections.abc import Callable -from contextlib import contextmanager from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Any -try: - from websockets.asyncio.client import ClientConnection, connect - from websockets.exceptions import ConnectionClosed -except ImportError as error: - raise SystemExit("Run this script with `uv run webmcp_cdp_demo.py`.") from error +from webmcp_client import ( + CdpPage, + DemoError, + discovery, + moli_server, + require, +) +from websockets.asyncio.client import connect +from websockets.exceptions import ConnectionClosed - -REPO_ROOT = Path(__file__).resolve().parents[2] DEMO_BASE = "https://googlechromelabs.github.io/webmcp-tools/demos/" DEMOS = { "booking": ("explainer/", {"getAvailability", "bookSlot", "cancelBooking"}), @@ -50,257 +44,6 @@ DEMOS = { } -class DemoError(RuntimeError): - pass - - -def require(condition: Any, message: str) -> None: - if not condition: - raise DemoError(message) - - -def discovery(endpoint: str) -> dict[str, Any]: - # CDP discovery must reach the local server even when an HTTP proxy is set. - opener = urllib.request.build_opener(urllib.request.ProxyHandler({})) - with opener.open(endpoint.rstrip("/") + "/json/version", timeout=2) as response: - return json.load(response) - - -@contextmanager -def moli_server(args: argparse.Namespace): - if args.cdp_endpoint: - yield args.cdp_endpoint - return - if args.moli_bin: - binary = Path(args.moli_bin).expanduser().resolve() - else: - binary = (REPO_ROOT / "target" / "release" / "moli").resolve() - require( - "debug" not in binary.parts, - "Debug builds are not allowed; use a release Moli binary.", - ) - require( - binary.is_file(), - "Build Moli first: `cargo build --release -p moli`, or pass --moli-bin with a release binary.", - ) - with socket.socket() as sock: - sock.bind(("127.0.0.1", 0)) - port = sock.getsockname()[1] - endpoint = f"http://127.0.0.1:{port}" - with ( - tempfile.TemporaryDirectory(prefix="moli-webmcp-") as directory, - tempfile.TemporaryFile(mode="w+t") as logs, - ): - command = [ - str(binary), - "serve", - "--host", - "127.0.0.1", - "--port", - str(port), - "--profile-dir", - str(Path(directory) / "profile"), - "--http-timeout", - str(int(args.timeout * 1000)), - ] - if args.http_proxy: - command.extend(["--http-proxy", args.http_proxy]) - process = subprocess.Popen(command, cwd=REPO_ROOT, stdout=logs, stderr=logs) - try: - deadline = time.monotonic() + args.startup_timeout - while time.monotonic() < deadline and process.poll() is None: - try: - discovery(endpoint) - break - except (OSError, ValueError): - time.sleep(0.05) - else: - logs.seek(0) - raise DemoError(f"Moli did not start:\n{logs.read()[-4000:]}") - print(f"Moli: {binary}\nCDP: {endpoint}", file=sys.stderr) - yield endpoint - finally: - if process.poll() is None: - process.terminate() - try: - process.wait(timeout=3) - except subprocess.TimeoutExpired: - process.kill() - process.wait(timeout=3) - - -class CdpPage: - def __init__(self, websocket: ClientConnection, timeout: float): - self.websocket = websocket - self.timeout = timeout - self.next_id = 1 - self.session_id: str | None = None - self.target_id: str | None = None - self.frame_id: str | None = None - self.dom_ready = False - self.tools: dict[tuple[str, str], dict[str, Any]] = {} - self.invoked: set[str] = set() - self.responded: dict[str, dict[str, Any]] = {} - self.errors: deque[dict[str, Any]] = deque(maxlen=5) - self.steps: list[dict[str, Any]] = [] - - async def receive(self) -> dict[str, Any]: - message = json.loads(await self.websocket.recv()) - if message.get("sessionId") != self.session_id: - return message - method = message.get("method") - params = message.get("params", {}) - if method == "WebMCP.toolsAdded": - for tool in params["tools"]: - self.tools[(tool["frameId"], tool["name"])] = tool - elif method == "WebMCP.toolsRemoved": - for tool in params["tools"]: - self.tools.pop((tool["frameId"], tool["name"]), None) - elif method == "WebMCP.toolInvoked": - self.invoked.add(params["invocationId"]) - elif method == "WebMCP.toolResponded": - # A result may arrive before the invokeTool command reply. - self.responded[params["invocationId"]] = params - elif method == "Page.domContentEventFired": - self.dom_ready = True - elif method == "Runtime.exceptionThrown": - self.errors.append(params.get("exceptionDetails", params)) - return message - - async def command( - self, - method: str, - params: dict[str, Any] | None = None, - *, - browser: bool = False, - ) -> dict[str, Any]: - message_id = self.next_id - self.next_id += 1 - message: dict[str, Any] = { - "id": message_id, - "method": method, - "params": params or {}, - } - if self.session_id and not browser: - message["sessionId"] = self.session_id - await self.websocket.send(json.dumps(message, ensure_ascii=False)) - try: - async with asyncio.timeout(self.timeout): - while True: - response = await self.receive() - if response.get("id") == message_id: - if "error" in response: - raise DemoError(f"{method}: {response['error']}") - return response.get("result", {}) - except TimeoutError as error: - raise DemoError(f"Timed out waiting for {method}") from error - - async def until(self, predicate: Callable[[], bool], label: str) -> None: - try: - async with asyncio.timeout(self.timeout): - while not predicate(): - await self.receive() - except TimeoutError as error: - raise DemoError( - f"Timed out waiting for {label}; page errors: {list(self.errors)}" - ) from error - - async def evaluate(self, expression: str) -> Any: - response = await self.command( - "Runtime.evaluate", - { - "expression": expression, - "returnByValue": True, - }, - ) - if "exceptionDetails" in response: - raise DemoError(f"Runtime.evaluate: {response['exceptionDetails']}") - return response.get("result", {}).get("value") - - async def open(self, url: str, expected: set[str]) -> None: - target = await self.command( - "Target.createTarget", {"url": "about:blank"}, browser=True - ) - self.target_id = target["targetId"] - attached = await self.command( - "Target.attachToTarget", - { - "targetId": self.target_id, - "flatten": True, - }, - browser=True, - ) - self.session_id = attached["sessionId"] - for method in ("Page.enable", "Runtime.enable", "WebMCP.enable"): - await self.command(method) - self.dom_ready = False - navigation = await self.command("Page.navigate", {"url": url}) - require(not navigation.get("errorText"), f"Navigation failed: {navigation}") - self.frame_id = navigation["frameId"] - await self.until(lambda: self.dom_ready, "DOMContentLoaded") - require( - await self.evaluate( - "typeof ModelContext !== 'undefined' && document.modelContext instanceof ModelContext" - ), - "This page does not expose the native ModelContext API.", - ) - await self.until( - lambda: ( - expected.issubset(self.catalog()) if expected else bool(self.catalog()) - ), - f"native tools {sorted(expected)} (received: {sorted(self.catalog())})", - ) - - def catalog(self) -> dict[str, dict[str, Any]]: - return { - name: tool - for (frame, name), tool in self.tools.items() - if frame == self.frame_id - } - - async def begin(self, name: str, arguments: dict[str, Any]) -> str: - tool = self.catalog().get(name) - require(tool, f"Native tool {name!r} is not registered.") - print( - f" invoke {name} {json.dumps(arguments, ensure_ascii=False)}", - file=sys.stderr, - ) - response = await self.command( - "WebMCP.invokeTool", - { - "frameId": tool["frameId"], - "toolName": name, - "input": arguments, - }, - ) - invocation_id = response["invocationId"] - self.steps.append( - {"tool": name, "input": arguments, "invocationId": invocation_id} - ) - return invocation_id - - async def finish(self, invocation_id: str) -> Any: - await self.until( - lambda: invocation_id in self.responded, f"tool response {invocation_id}" - ) - response = self.responded.pop(invocation_id) - self.invoked.discard(invocation_id) - self.steps[-1]["response"] = response - require( - response.get("status") == "Completed", f"Tool invocation failed: {response}" - ) - return response.get("output") - - async def invoke(self, name: str, arguments: dict[str, Any]) -> Any: - return await self.finish(await self.begin(name, arguments)) - - async def close(self) -> None: - if self.target_id: - await self.command( - "Target.closeTarget", {"targetId": self.target_id}, browser=True - ) - - async def booking(page: CdpPage) -> dict[str, Any]: today = datetime.now(timezone.utc).date() availability = await page.invoke( diff --git a/moli-playground/webmcp/webmcp_client.py b/moli-playground/webmcp/webmcp_client.py new file mode 100644 index 0000000000..27f9f6491b --- /dev/null +++ b/moli-playground/webmcp/webmcp_client.py @@ -0,0 +1,440 @@ +"""Shared release Moli server and native WebMCP CDP client.""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import socket +import subprocess +import sys +import tempfile +import time +import urllib.request +from collections import deque +from collections.abc import Callable +from contextlib import contextmanager, suppress +from pathlib import Path +from typing import Any + +try: + from websockets.asyncio.client import ClientConnection + from websockets.exceptions import ConnectionClosed +except ImportError as error: + raise SystemExit( + "Run with `uv run webmcp.py` or `uv run webmcp_cdp_demo.py`." + ) from error + + +REPO_ROOT = Path(__file__).resolve().parents[2] + + +class DemoError(RuntimeError): + pass + + +def require(condition: Any, message: str) -> None: + if not condition: + raise DemoError(message) + + +def discovery(endpoint: str) -> dict[str, Any]: + # CDP discovery must reach the local server even when an HTTP proxy is set. + opener = urllib.request.build_opener(urllib.request.ProxyHandler({})) + with opener.open(endpoint.rstrip("/") + "/json/version", timeout=2) as response: + return json.load(response) + + +@contextmanager +def moli_server(args: argparse.Namespace): + if args.cdp_endpoint: + yield args.cdp_endpoint + return + if args.moli_bin: + binary = Path(args.moli_bin).expanduser().resolve() + else: + binary = (REPO_ROOT / "target" / "release" / "moli").resolve() + require( + "debug" not in binary.parts, + "Debug builds are not allowed; use a release Moli binary.", + ) + require( + binary.is_file(), + "Build Moli first: `cargo build --release -p moli`, or pass --moli-bin with a release binary.", + ) + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + endpoint = f"http://127.0.0.1:{port}" + with ( + tempfile.TemporaryDirectory(prefix="moli-webmcp-") as directory, + tempfile.TemporaryFile(mode="w+t") as logs, + ): + command = [ + str(binary), + "serve", + "--host", + "127.0.0.1", + "--port", + str(port), + "--profile-dir", + str(Path(directory) / "profile"), + "--http-timeout", + str(int(args.timeout * 1000)), + ] + if args.http_proxy: + command.extend(["--http-proxy", args.http_proxy]) + process = subprocess.Popen(command, cwd=REPO_ROOT, stdout=logs, stderr=logs) + try: + deadline = time.monotonic() + args.startup_timeout + while time.monotonic() < deadline and process.poll() is None: + try: + discovery(endpoint) + break + except (OSError, ValueError): + time.sleep(0.05) + else: + logs.seek(0) + raise DemoError(f"Moli did not start:\n{logs.read()[-4000:]}") + print(f"Moli: {binary}\nCDP: {endpoint}", file=sys.stderr) + yield endpoint + finally: + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=3) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=3) + + +class CdpPage: + def __init__(self, websocket: ClientConnection, timeout: float): + self.websocket = websocket + self.timeout = timeout + self.next_id = 1 + self.session_id: str | None = None + self.target_id: str | None = None + self.frame_id: str | None = None + self.dom_ready = False + self.tools: dict[tuple[str, str], dict[str, Any]] = {} + self.invoked: set[str] = set() + self.responded: dict[str, dict[str, Any]] = {} + self.errors: deque[dict[str, Any]] = deque(maxlen=5) + self.steps: list[dict[str, Any]] = [] + self.replies: dict[int, asyncio.Future[dict[str, Any]]] = {} + self.changed = asyncio.Event() + self.reader_error: Exception | None = None + self.reader = asyncio.create_task(self.read_messages()) + + async def read_messages(self) -> None: + try: + while True: + message = await self.receive() + reply = self.replies.get(message.get("id")) + if reply is not None and not reply.done(): + reply.set_result(message) + self.changed.set() + except (OSError, ValueError, ConnectionClosed) as error: + self.reader_error = error + for reply in self.replies.values(): + if not reply.done(): + reply.set_exception(error) + self.changed.set() + + async def receive(self) -> dict[str, Any]: + message = json.loads(await self.websocket.recv()) + if message.get("sessionId") != self.session_id: + return message + method = message.get("method") + params = message.get("params", {}) + if method == "WebMCP.toolsAdded": + for tool in params["tools"]: + self.tools[(tool["frameId"], tool["name"])] = tool + elif method == "WebMCP.toolsRemoved": + for tool in params["tools"]: + self.tools.pop((tool["frameId"], tool["name"]), None) + elif method == "WebMCP.toolInvoked": + self.invoked.add(params["invocationId"]) + elif method == "WebMCP.toolResponded": + # A result may arrive before the invokeTool command reply. + self.responded[params["invocationId"]] = params + elif method == "Page.domContentEventFired": + self.dom_ready = True + elif method == "Page.frameNavigated" and not params["frame"].get("parentId"): + self.frame_id = params["frame"]["id"] + elif method == "Page.frameDetached": + self.tools = { + key: tool + for key, tool in self.tools.items() + if key[0] != params["frameId"] + } + elif method == "Runtime.exceptionThrown": + self.errors.append(params.get("exceptionDetails", params)) + return message + + async def command( + self, + method: str, + params: dict[str, Any] | None = None, + *, + browser: bool = False, + ) -> dict[str, Any]: + message_id = self.next_id + self.next_id += 1 + message: dict[str, Any] = { + "id": message_id, + "method": method, + "params": params or {}, + } + if self.session_id and not browser: + message["sessionId"] = self.session_id + if self.reader_error is not None: + raise self.reader_error + reply = asyncio.get_running_loop().create_future() + self.replies[message_id] = reply + try: + async with asyncio.timeout(self.timeout): + await self.websocket.send(json.dumps(message, ensure_ascii=False)) + response = await reply + if "error" in response: + raise DemoError(f"{method}: {response['error']}") + return response.get("result", {}) + except TimeoutError as error: + raise DemoError(f"Timed out waiting for {method}") from error + finally: + self.replies.pop(message_id, None) + if not reply.done(): + reply.cancel() + elif not reply.cancelled(): + reply.exception() + + async def until(self, predicate: Callable[[], bool], label: str) -> None: + try: + async with asyncio.timeout(self.timeout): + while not predicate(): + if self.reader_error is not None: + raise self.reader_error + self.changed.clear() + await self.changed.wait() + except TimeoutError as error: + raise DemoError( + f"Timed out waiting for {label}; page errors: {list(self.errors)}" + ) from error + + async def evaluate(self, expression: str) -> Any: + response = await self.command( + "Runtime.evaluate", + { + "expression": expression, + "returnByValue": True, + }, + ) + if "exceptionDetails" in response: + raise DemoError(f"Runtime.evaluate: {response['exceptionDetails']}") + return response.get("result", {}).get("value") + + async def open( + self, + url: str | None, + expected: set[str] | None = None, + *, + target_id: str | None = None, + ) -> None: + self.owns_target = target_id is None + if target_id is None: + target = await self.command( + "Target.createTarget", {"url": "about:blank"}, browser=True + ) + target_id = target["targetId"] + self.target_id = target_id + attached = await self.command( + "Target.attachToTarget", + { + "targetId": self.target_id, + "flatten": True, + }, + browser=True, + ) + self.session_id = attached["sessionId"] + for method in ("Page.enable", "Runtime.enable", "WebMCP.enable"): + await self.command(method) + if url is not None: + self.dom_ready = False + navigation = await self.command("Page.navigate", {"url": url}) + require(not navigation.get("errorText"), f"Navigation failed: {navigation}") + self.frame_id = navigation["frameId"] + await self.until(lambda: self.dom_ready, "DOMContentLoaded") + else: + await self.refresh() + if expected is not None: + require( + await self.evaluate( + "typeof ModelContext !== 'undefined' && document.modelContext instanceof ModelContext" + ), + "This page does not expose the native ModelContext API.", + ) + await self.until( + lambda: ( + expected.issubset(self.catalog()) + if expected + else bool(self.catalog()) + ), + f"native tools {sorted(expected)}", + ) + + async def settle(self, seconds: float) -> None: + await asyncio.sleep(seconds) + await self.refresh() + + async def refresh(self) -> None: + # A command reply is a barrier for already queued tool-change events. + tree = await self.command("Page.getFrameTree") + self.frame_id = tree["frameTree"]["frame"]["id"] + + def all_tools(self, frame_id: str | None = None) -> list[dict[str, Any]]: + return [ + tool + for (frame, _), tool in sorted(self.tools.items()) + if frame_id is None or frame == frame_id + ] + + def select_tool(self, name: str, frame_id: str | None = None) -> dict[str, Any]: + matches = [tool for tool in self.all_tools(frame_id) if tool["name"] == name] + require( + matches, + f"Native tool {name!r} is not registered; list tools to see the current catalog.", + ) + require( + len(matches) == 1, + f"Tool {name!r} exists in multiple frames: {[tool['frameId'] for tool in matches]}; select --frame-id or use the shell's frame command.", + ) + return matches[0] + + def catalog(self) -> dict[str, dict[str, Any]]: + return { + name: tool + for (frame, name), tool in self.tools.items() + if frame == self.frame_id + } + + async def begin( + self, name: str, arguments: dict[str, Any], *, frame_id: str | None = None + ) -> str: + tool = self.tools.get((frame_id or self.frame_id, name)) + require(tool, f"Native tool {name!r} is not registered.") + print(f" invoke {name} (frame {tool['frameId']})", file=sys.stderr) + response = await self.command( + "WebMCP.invokeTool", + { + "frameId": tool["frameId"], + "toolName": name, + "input": arguments, + }, + ) + invocation_id = response["invocationId"] + step = { + "tool": name, + "frameId": tool["frameId"], + "input": arguments, + "invocationId": invocation_id, + } + if "backendNodeId" in tool: + step["backendNodeId"] = tool["backendNodeId"] + self.steps.append(step) + return invocation_id + + def invocation(self, invocation_id: str) -> dict[str, Any]: + step = next( + (step for step in self.steps if step["invocationId"] == invocation_id), None + ) + require(step, f"Unknown invocation {invocation_id!r} in this session.") + return step + + async def response(self, invocation_id: str) -> dict[str, Any]: + step = self.invocation(invocation_id) + await self.until( + lambda: invocation_id in self.responded, f"tool response {invocation_id}" + ) + response = self.responded[invocation_id] + self.invoked.discard(invocation_id) + step["response"] = response + return response + + async def finish(self, invocation_id: str) -> Any: + response = await self.response(invocation_id) + require( + response.get("status") == "Completed", f"Tool invocation failed: {response}" + ) + return response.get("output") + + async def invoke(self, name: str, arguments: dict[str, Any]) -> Any: + return await self.finish(await self.begin(name, arguments)) + + async def confirm(self, invocation_id: str) -> None: + step = self.invocation(invocation_id) + require( + "backendNodeId" in step, + "This invocation does not belong to a declarative form.", + ) + await self.until( + lambda: invocation_id in self.invoked or invocation_id in self.responded, + f"form invocation {invocation_id}", + ) + require( + invocation_id not in self.responded, "This invocation has already finished." + ) + node = await self.command( + "DOM.resolveNode", {"backendNodeId": step["backendNodeId"]} + ) + object_id = node["object"]["objectId"] + try: + result = await self.command( + "Runtime.callFunctionOn", + { + "objectId": object_id, + "functionDeclaration": """function() { + if (!this.matches(':tool-form-active')) throw new Error('The tool form is no longer active.'); + const button = Array.from(this.elements).find(el => el.matches(':tool-submit-active')); + if (button) { + if (button.matches(':disabled')) throw new Error('The submit button is disabled.'); + button.click(); + } else this.requestSubmit(); + }""", + }, + ) + require( + "exceptionDetails" not in result, f"Form confirmation failed: {result}" + ) + finally: + try: + await self.command("Runtime.releaseObject", {"objectId": object_id}) + except (DemoError, OSError, ConnectionClosed): + pass + + async def cancel(self, invocation_id: str) -> dict[str, Any]: + self.invocation(invocation_id) + await self.command("WebMCP.cancelInvocation", {"invocationId": invocation_id}) + return await self.response(invocation_id) + + async def close(self) -> None: + try: + if self.target_id and not self.owns_target: + for step in self.steps: + if step["invocationId"] not in self.responded: + with suppress(DemoError): + await self.cancel(step["invocationId"]) + if self.target_id and self.owns_target: + await self.command( + "Target.closeTarget", {"targetId": self.target_id}, browser=True + ) + elif self.session_id: + await self.command( + "Target.detachFromTarget", + {"sessionId": self.session_id}, + browser=True, + ) + finally: + self.reader.cancel() + with suppress(asyncio.CancelledError): + await self.reader