Mujoco_WASM/control_bridge/src/mujoco_control_bridge/client.py

"""Synchronous, generic robot SDK; stdlib only on the client path."""

import contextlib
import json
import os
import time
from urllib.error import HTTPError, URLError
from urllib.parse import urlparse
from urllib.request import HTTPRedirectHandler, ProxyHandler, Request, build_opener

from . import protocol as p
from .protocol import RobotError


class NoRedirect(HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None  # Never forward the bearer token to a redirected endpoint.


class SimRobotClient:
    def __init__(self, endpoint="http://127.0.0.1:8766", token=None, timeout=2.0):
        parsed = urlparse(endpoint)
        if (
            parsed.scheme != "http"
            or parsed.hostname not in {"127.0.0.1", "localhost"}
            or parsed.username
            or parsed.password
            or parsed.query
            or parsed.fragment
            or parsed.path not in {"", "/"}
        ):
            raise ValueError("V1 仅支持本机 http://127.0.0.1:port 或 localhost")
        self.endpoint = endpoint.rstrip("/")
        self.token = token or os.environ.get("MUJOCO_CONTROL_TOKEN", "")
        if not self.token:
            raise ValueError("需要 token 或 MUJOCO_CONTROL_TOKEN")
        if (
            not isinstance(self.token, str)
            or not 16 <= len(self.token) <= 4096
            or not all(33 <= ord(c) <= 126 for c in self.token)
        ):
            raise ValueError("控制 token 必须为16–4096字符的可打印 ASCII，且不含空白")
        self.timeout = timeout
        self._opener = build_opener(ProxyHandler({}), NoRedirect())
        self._descriptor = None
        self._identity = None
        self._action_seq = 0
        self._last_action_at = 0.0

    def _request(self, method, path, payload=None, lease=None):
        headers = {"Authorization": f"Bearer {self.token}"}
        if lease:
            headers["X-Control-Lease"] = lease["leaseId"]
        raw = None
        if payload is not None:
            raw = json.dumps(payload, allow_nan=False).encode()
            if len(raw) > 65536:
                raise RobotError("INVALID_MESSAGE", "消息超过64KiB")
            headers["Content-Type"] = "application/json"
        request = Request(
            f"{self.endpoint}/api/control/v1{path}", data=raw, headers=headers, method=method
        )
        try:
            with self._opener.open(request, timeout=self.timeout) as response:
                data = response.read(65537)
            if len(data) > 65536:
                raise RobotError("INVALID_MESSAGE", "响应超过64KiB")
            return p.loads(data)
        except HTTPError as exc:
            try:
                error = p.loads(exc.read(65537))["error"]
                code, message = error["code"], error["message"]
            except (RobotError, KeyError, TypeError):
                code, message = "DISCONNECTED", f"桥接 HTTP {exc.code}"
            raise RobotError(code, message) from None
        except (URLError, TimeoutError, ConnectionError, OSError) as exc:
            raise RobotError("DISCONNECTED", "无法访问本机控制桥接") from exc

    @property
    def is_connected(self):
        return self._identity is not None and time.monotonic() - self._last_action_at <= 0.5

    def describe(self):
        self._descriptor = p.descriptor(self._request("GET", "/robot"))
        return self._descriptor

    def connect(self):
        if self._identity:
            raise RobotError("CONFLICT", "客户端已连接；请先 disconnect")
        desc = self.describe()
        obs = p.observation(self._request("GET", "/observation"), desc)
        claim = {
            "sessionId": obs["sessionId"],
            "modelEpoch": obs["modelEpoch"],
            "modelFingerprint": desc["modelFingerprint"],
        }
        lease = p.record(self._request("POST", "/lease", claim), p.IDENTITY)
        p.identity(lease)
        p.same_identity(lease, obs, ("sessionId", "modelEpoch"))
        self._identity = lease
        self._action_seq = 0
        self._last_action_at = time.monotonic()
        return desc

    def _connected(self):
        if not self.is_connected:
            raise RobotError("DISCONNECTED", "没有有效租约；请重新在浏览器授权并连接")
        return dict(self._identity)

    def get_observation(self):
        identity = self._connected()
        obs = p.observation(self._request("GET", "/observation"), self._descriptor)
        p.same_identity(obs, identity, ("sessionId", "modelEpoch"))
        return obs

    def send_action(self, values):
        identity = self._connected()
        values = p.values(values, self._descriptor["actionChannels"])
        self._action_seq += 1
        packet = {"protocolVersion": 1, **identity, "actionSeq": self._action_seq, "values": values}
        try:
            result = p.action_result(
                self._request("POST", "/action", packet, lease=identity), self._descriptor
            )
            p.same_identity(result, packet, (*p.IDENTITY, "actionSeq"))
            self._last_action_at = time.monotonic()
            return result
        except RobotError:
            self.disconnect()
            raise

    def reset(self):
        identity = self._connected()
        try:
            observed = p.observation(
                self._request("POST", "/reset", {}, lease=identity), self._descriptor
            )
            if (
                observed["sessionId"] != identity["sessionId"]
                or observed["modelEpoch"] <= identity["modelEpoch"]
                or not observed["paused"]
            ):
                raise RobotError("STALE", "reset 没有返回新代次")
            return observed
        finally:
            self.disconnect()

    def disconnect(self):
        identity, self._identity = self._identity, None
        if identity:
            # The server/browser watchdog remains the final safety barrier.
            with contextlib.suppress(RobotError):
                self._request("DELETE", "/lease", lease=identity)

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, *_):
        self.disconnect()