Mujoco_WASM/decision_server/protocol.py

"""Strict data-only validation of the same versioned contract used by the browser."""

import json
import math
from copy import deepcopy
from pathlib import Path

SCHEMA = json.loads(
    (Path(__file__).resolve().parents[1] / "contracts/lekiwi-agent-v1.schema.json").read_text()
)
VERSION = "lekiwi-agent-v1"
LANGUAGE_VERSION = "lekiwi-language-v2"
LANGUAGE_SCHEMA = json.loads(
    (
        Path(__file__).resolve().parents[1] / "contracts/lekiwi-language-task-v2.schema.json"
    ).read_text()
)


def schema_for(version):
    if version == VERSION:
        return SCHEMA
    if version == LANGUAGE_VERSION:
        return LANGUAGE_SCHEMA
    raise DecisionError("unsupported_version")


PRECONDITIONS = dict(
    zip(
        SCHEMA["$defs"]["Plan"]["properties"]["steps"]["items"]["properties"]["skill"]["enum"],
        [
            "scene-ready",
            "base-stopped",
            "tcp-above",
            "aligned",
            "dual-contact",
            "verified-grasp",
            "verified-grasp",
            "transported",
            "supported",
            "released",
            "retreat",
        ],
        strict=True,
    )
)
SKILLS = list(PRECONDITIONS)
LANGUAGE_PRECONDITIONS = {"stow": "scene-ready", "approach": "arm-safe", **PRECONDITIONS}
LANGUAGE_SKILLS = list(LANGUAGE_PRECONDITIONS)


class DecisionError(Exception):
    """Public error code only: upstream bodies/headers must never become logs or UI errors."""

    def __init__(self, code, status=400):
        super().__init__(code)
        self.code = code
        self.status = status


def loads(text):
    def pairs(items):
        result = {}
        for key, value in items:
            if key in result:
                raise DecisionError("duplicate_json_key")
            result[key] = value
        return result

    def invalid(_):
        raise DecisionError("nonfinite_json")

    try:
        return json.loads(text, object_pairs_hook=pairs, parse_constant=invalid)
    except (ValueError, TypeError, RecursionError) as exc:
        raise DecisionError("invalid_json") from exc


def check(node, value, schema=SCHEMA):
    def fail():
        raise DecisionError("contract_mismatch")

    if "$ref" in node:
        return check(schema["$defs"][node["$ref"].removeprefix("#/$defs/")], value, schema)
    if "const" in node and (type(value) is not type(node["const"]) or value != node["const"]):
        fail()
    if "enum" in node and value not in node["enum"]:
        fail()
    kind = node.get("type")
    if kind in ("number", "integer"):
        if type(value) not in (int, float) or not math.isfinite(value):
            fail()
        if kind == "integer" and value != int(value):
            fail()
        if not node.get("minimum", -math.inf) <= value <= node.get("maximum", math.inf):
            fail()
    elif kind == "boolean":
        if type(value) is not bool:
            fail()
    elif kind == "string":
        if not isinstance(value, str):
            fail()
        if not node.get("minLength", 0) <= len(value) <= node.get("maxLength", math.inf):
            fail()
    elif kind == "array":
        if not isinstance(value, list):
            fail()
        if not node.get("minItems", 0) <= len(value) <= node.get("maxItems", math.inf):
            fail()
        if node.get("uniqueItems") and len({json.dumps(v) for v in value}) != len(value):
            fail()
        for item in value:
            check(node["items"], item, schema)
    elif kind == "object":
        if not isinstance(value, dict):
            fail()
        props = node.get("properties", {})
        if set(node.get("required", [])) - value.keys():
            fail()
        if node.get("additionalProperties") is False and value.keys() - props.keys():
            fail()
        for key in value.keys() & props.keys():
            check(props[key], value[key], schema)


def validate(name, value, version=VERSION):
    try:
        if len(json.dumps(value, ensure_ascii=False, allow_nan=False)) > 65536:
            raise DecisionError("message_too_large", 413)
        schema = schema_for(version)
        check(schema["$defs"][name], value, schema)
    except (ValueError, TypeError, OverflowError, RecursionError) as exc:
        raise DecisionError("contract_mismatch") from exc
    return deepcopy(value)


def output_schema(name, version=VERSION):
    """Equivalent schema with explicit string types for strict API implementations."""

    def expand(node):
        result = deepcopy(node)
        if "type" not in result and ("const" in result or "enum" in result):
            result["type"] = "string"  # All enum/const nodes in v1 are strings.
        for key, value in result.items():
            if isinstance(value, dict):
                result[key] = {k: expand(v) if isinstance(v, dict) else v for k, v in value.items()}
        if isinstance(result.get("items"), dict):
            result["items"] = expand(node["items"])
        return result

    return expand(schema_for(version)["$defs"][name])


def remaining_skills(value, version=VERSION):
    schema_for(version)
    skills = LANGUAGE_SKILLS if version == LANGUAGE_VERSION else SKILLS
    if (
        not isinstance(value, list)
        or not value
        or value not in [skills[i:] for i in range(len(skills))]
    ):
        raise DecisionError("invalid_remaining_skills")
    return value


def validate_plan(value, remaining, version=VERSION):
    value = validate("Plan", value, version)
    if [step["skill"] for step in value["steps"]] != remaining_skills(remaining, version):
        raise DecisionError("invalid_plan_order")
    preconditions = LANGUAGE_PRECONDITIONS if version == LANGUAGE_VERSION else PRECONDITIONS
    if any(step["precondition"] != preconditions[step["skill"]] for step in value["steps"]):
        raise DecisionError("invalid_precondition")
    return value


def candidates(value, version=VERSION):
    schema_for(version)
    skills = LANGUAGE_SKILLS if version == LANGUAGE_VERSION else SKILLS
    if (
        not isinstance(value, list)
        or not 1 <= len(value) <= 4
        or any(type(v) is not str or v not in [*skills, "stop"] for v in value)
        or len(set(value)) != len(value)
    ):
        raise DecisionError("invalid_candidates")
    return value


def validate_jev(value, choices, version=VERSION):
    value = validate("JevDecision", value, version)
    if value["choice"] not in candidates(choices, version):
        raise DecisionError("invalid_choice")
    return value


def fields(value, required, optional=()):
    if (
        not isinstance(value, dict)
        or set(required) - value.keys()
        or value.keys() - set(required) - set(optional)
    ):
        raise DecisionError("invalid_fields")
    return value