diff --git a/pyproject.toml b/pyproject.toml index 3ad9c39..f455ed5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,6 +5,7 @@ description = "Dataclass for trame UI binding" authors = [{ name = "Kitware Inc." }] dependencies = [ "trame_client>=3.10", + "trame-common>=1.3.0", "loguru", "typing_extensions; python_version == '3.13'", ] diff --git a/src/trame_dataclass/mcp.py b/src/trame_dataclass/mcp.py new file mode 100644 index 0000000..307e70c --- /dev/null +++ b/src/trame_dataclass/mcp.py @@ -0,0 +1,102 @@ +"""What agents (MCP) see of a model class's fields. + +``StateDataModel``'s ``_mcp_fields`` hook (trame-common's ``MCPExposable``) +calls this module, importing it the first time it is used. It applies the +rules only models have: the field descriptors (``Sync``, ``ClientOnly``, +``ServerOnly``) and ``all_fields``. The trame server handles the rest, such +as schemas, reads and writes, from what the classes declare. +""" + +from trame_common.obj import mcp as mcp_options +from trame_common.obj.mcp import Field + +from trame_dataclass.v2 import mcp_model_options + +__all__ = [ + "fields", +] + +_MISSING = object() + + +def _descriptor(cls, name): + for klass in cls.__mro__: + if name in vars(klass): + return vars(klass)[name] + return None + + +def _option(cls, name): + """The nearest ``mcp=`` value given for field ``name`` along the MRO: a + subclass redeclaring a field (to change its default) keeps its options.""" + for klass in cls.__mro__: + option = getattr(vars(klass).get(name), "_mcp", None) + if option is not None: + return option + return None + + +def _exposed(cls, name, own, listed, all_fields, descriptor_type): + """What agents see of field ``name``: its options, None when it is not + exposed, or False when it is hidden on purpose. The class's + ``properties=`` entry goes over the field's own ``mcp=`` value.""" + if listed is False or (listed is None and own is False): + return False + if listed is not None or own is not None: + options = {**(own or {}), **(listed or {})} + elif all_fields and descriptor_type != "server_only": + options = {} + else: + return None + if "read_only" not in options: + if descriptor_type == "server_only": + msg = ( + f"{cls.__name__}.{name} is a ServerOnly field: exposing it requires " + f"read_only=True or False" + ) + raise ValueError(msg) + options["read_only"] = descriptor_type == "client_only" + return options + + +def fields(cls): + """``{name: Field}`` for every field of ``cls``, in definition order (a + base's first). A field is exposed when the class lists it in + ``properties=`` or the field has an ``mcp=`` value, and otherwise, with + ``all_fields``, when it is ``Sync`` or ``ClientOnly``. ``ClientOnly`` + fields are read-only unless their options say otherwise, and an exposed + ``ServerOnly`` field must say.""" + names = getattr(cls, "FIELD_NAMES", ()) + if not names: + return {} # StateDataModel itself + listed = mcp_options.class_options(cls)["properties"] + all_fields = mcp_model_options(cls)["all_fields"] + ordered = {} + for klass in reversed(cls.__mro__): + ordered.update((name, None) for name in vars(klass) if name in names) + result = {} + for name in [*ordered, *sorted(set(names) - ordered.keys())]: + descriptor = _descriptor(cls, name) + if name in cls.CLIENT_ONLY_NAMES: + descriptor_type = "client_only" + elif name in cls.CLIENT_NAMES: + descriptor_type = "sync" + else: + descriptor_type = "server_only" + checker = cls.TYPE_CHECKING.get(name) + result[name] = Field( + name=name, + type_def=checker.type_def if checker else None, + options=_exposed( + cls, + name, + _option(cls, name), + listed.get(name), + all_fields, + descriptor_type, + ), + nullable=getattr(descriptor, "_default", _MISSING) is None, + holds_objects=name in cls.DATACLASS_NAMES + or bool(getattr(descriptor, "_has_dataclass", False)), + ) + return result diff --git a/src/trame_dataclass/v2.py b/src/trame_dataclass/v2.py index c8aa2c4..216c573 100644 --- a/src/trame_dataclass/v2.py +++ b/src/trame_dataclass/v2.py @@ -1,7 +1,6 @@ import asyncio import contextlib import inspect -import string import sys import types import weakref @@ -11,6 +10,9 @@ from typing import Any, ForwardRef, Union, get_args, get_origin from loguru import logger +from trame_common.obj import mcp as mcp_options +from trame_common.obj.mcp import MCPExposable +from trame_common.utils.ids import next_id from trame_dataclass import module as dataclass_module from trame_dataclass.widgets.dataclass import Provider @@ -21,26 +23,10 @@ from typing import evaluate_forward_ref # ----------------------------------------------------------------------------- -# Id generator +# Instances, by id (ids come from trame-common, shared with trame components) # ----------------------------------------------------------------------------- INSTANCES = weakref.WeakValueDictionary() -_INSTANCE_COUNT = 0 -_INSTANCE_ID_CHARS = string.digits + string.ascii_letters - - -def _next_id(): - global _INSTANCE_COUNT # noqa: PLW0603 - _INSTANCE_COUNT += 1 - - result = [] - value = _INSTANCE_COUNT - size = len(_INSTANCE_ID_CHARS) - while value > 0: - remainder = value % size - result.append(_INSTANCE_ID_CHARS[remainder]) - value //= size - - return "".join(result[::-1]) +_next_id = next_id # ----------------------------------------------------------------------------- @@ -269,7 +255,7 @@ def validate(self, instance, value): logger.warning(msg) -class StateDataModel: +class StateDataModel(MCPExposable): """Base class for reactive data models that synchronize state between server and client. Declare fields as class-level descriptors (``Sync``, ``ServerOnly``, ``ClientOnly``) to @@ -284,8 +270,27 @@ class MyModel(StateDataModel): model = MyModel(trame_server=server) model.count = 42 # triggers async push to client + + Agents (MCP) see a model class when it declares something: the ``mcp=`` + class keyword, a field's ``mcp=`` value, or the ``@mcp.method``, + ``@mcp.property`` and ``@mcp.constructor`` markers. The keyword takes + True (every field by mode: ``Sync`` writable, ``ClientOnly`` read-only, + ``ServerOnly`` only when marked; methods only when marked), False + (hidden), or a dict of ``description``, ``all_fields``, ``all_methods``, + ``properties``, ``methods`` and ``constructor``. """ + _mcp_api_packages = ("trame_dataclass",) + + def __init_subclass__(cls, mcp=None, **kwargs): + label = f"{cls.__name__}(mcp=...)" + common, own = _model_declaration(label, mcp) + _check_model_markers(cls) + _check_server_only(cls, common, label) + super().__init_subclass__(mcp=common, **kwargs) + if own is not None: + setattr(cls, _MODEL_KEYWORD_ATTR, own) + def __init__(self, trame_server=None, enable_collaboration=False, **kwargs): """Initialize the model, apply defaults, and register with the server if provided. @@ -351,6 +356,47 @@ def generate_gui(cls, trame_server=None) -> str: # noqa: ARG003 """ return getattr(cls, "TEMPLATE", "") + @classmethod + def mcp_describe(cls, declare=None, /, **options): + """Set what agents see of this model class, when it comes from another + library, with the options of the ``mcp=`` class keyword, or False to hide + it. The settings apply to the whole process and to subclasses, and only + one call per class may change them. + + Every call returns the class's effective options as plain data; + ``server.mcp.describe(cls)`` shows what agents see. + """ + if declare is not None and options: + msg = "mcp_describe() takes True, False or options, not both" + raise TypeError(msg) + value = declare if declare is not None else (options or None) + label = f"{cls.__name__}.mcp_describe()" + common, own = _model_declaration(label, value) + _check_server_only(cls, common, label) + if isinstance(common, dict): + result = super().mcp_describe(**common) + else: + result = super().mcp_describe(common) + if own is not None: + setattr(cls, _MODEL_DESCRIBED_ATTR, own) + mcp_options.changed() + return {**result, **mcp_model_options(cls)} + + # What agents (MCP) see of fields (MCPExposable's hooks) + + @classmethod + def _mcp_field_names(cls): + return tuple(getattr(cls, "FIELD_NAMES", ())) + + @classmethod + def _mcp_fields(cls): + from trame_dataclass import mcp # noqa: PLC0415 (it imports this module) + + return mcp.fields(cls) + + async def _mcp_settle(self): + await self.completion() + def _register_server(self, **_): self.server.protocol_call("trame.dataclass.register", self) @@ -571,6 +617,120 @@ def flush(self, dirty_set: set[str] | None = None, force_push=False): self._flush_impl(msg) +# ----------------------------------------------------------------------------- +# Agents (MCP) +# ----------------------------------------------------------------------------- +MCP_FIELD_OPTIONS = ("description", "read_only", "schema") + + +def _field_mcp_option(value, needs_read_only): + """A field's ``mcp=`` value as None, False or a dict of options.""" + if value is None or value is False: + return value + if value is True: + value = {} + elif isinstance(value, str): + value = {"description": value} + elif isinstance(value, dict): + unknown = sorted(set(value) - set(MCP_FIELD_OPTIONS)) + if unknown: + msg = f"mcp: unknown options {unknown}, allowed: {list(MCP_FIELD_OPTIONS)}" + raise ValueError(msg) + value = dict(value) + if "read_only" in value and not isinstance(value["read_only"], bool): + msg = "mcp: read_only must be True or False" + raise TypeError(msg) + if "schema" in value and not isinstance(value["schema"], dict): + msg = "mcp: schema must be a dict of JSON Schema keywords" + raise TypeError(msg) + else: + msg = f"mcp must be True, False, a description or a dict, got {value!r}" + raise TypeError(msg) + if needs_read_only and "read_only" not in value: + msg = ( + "A ServerOnly field exposed to agents must state read_only=True or " + "False: mcp={'read_only': ...}" + ) + raise ValueError(msg) + return value + + +def _check_server_only(cls, value, label): + """A ServerOnly field listed in ``properties`` must state read_only.""" + layer = mcp_options.normalize_options(value, label) + if not isinstance(layer, dict): + return + server_only = getattr(cls, "FIELD_NAMES", set()) - getattr( + cls, "CLIENT_NAMES", set() + ) + for name, options in layer.get("properties", {}).items(): + if name in server_only and options is not False and "read_only" not in options: + msg = ( + f"{cls.__name__}.{name} is a ServerOnly field: listing it requires " + f"read_only=True or False" + ) + raise ValueError(msg) + + +# What a model class declares besides the options every class has +MCP_MODEL_OPTIONS = ("all_fields",) +_MODEL_KEYWORD_ATTR = "_trame_mcp_model_keyword" +_MODEL_DESCRIBED_ATTR = "_trame_mcp_model_described" + + +def _model_declaration(label, value): + """``value``, a model class's ``mcp=`` keyword or ``mcp_describe()`` + declaration, as ``(common, own)``: the part trame-common handles, and + the model's own options (``all_fields``) or None.""" + if value is None or value is False: + return value, None + if value is True: # every field, by mode + return True, {"all_fields": True} + allowed = (*mcp_options.CLASS_OPTIONS, *MCP_MODEL_OPTIONS) + if not isinstance(value, dict): + msg = f"{label} must be True, False or a dict of {list(allowed)}, got {value!r}" + raise TypeError(msg) + unknown = sorted(set(value) - set(allowed)) + if unknown: + msg = f"{label}: unknown options {unknown}, allowed: {list(allowed)}" + raise ValueError(msg) + own = {} + if value.get("all_fields") is not None: + if not isinstance(value["all_fields"], bool): + msg = f"{label}: all_fields must be True or False" + raise TypeError(msg) + own["all_fields"] = value["all_fields"] + common = {k: v for k, v in value.items() if k not in MCP_MODEL_OPTIONS} + # Options of its own still declare the class + return (common if common or not own else True), (own or None) + + +def _check_model_markers(cls): + """Reject first-class tools on a model: they are for components built + once.""" + for name, _, _ in mcp_options.first_class_members(cls): + msg = ( + f"{cls.__name__}.{name}: @mcp.tool, @mcp.resource and @mcp.prompt " + f"are for components built once; a model's methods use @mcp.method" + ) + raise TypeError(msg) + + +def mcp_model_options(cls): + """ + The model's own options, combined like the options every class has: + ``all_fields``, from the base classes down, looking on each class at its + keyword, then its ``mcp_describe()``. + """ + all_fields = False + for klass in reversed(cls.__mro__): + own = vars(klass) + for layer in (own.get(_MODEL_KEYWORD_ATTR), own.get(_MODEL_DESCRIBED_ATTR)): + if layer is not None and "all_fields" in layer: + all_fields = layer["all_fields"] + return {"all_fields": all_fields} + + # ----------------------------------------------------------------------------- # Generic encoder/decoder # ----------------------------------------------------------------------------- @@ -680,8 +840,16 @@ class ServerOnly: or a container of them; automatically sets up the appropriate encoder/decoder. client_deep_reactive: reserved for subclass ``Sync`` — ignored here. type_checking: how type mismatches are reported (default: ``WARNING``). + mcp: what agents see of the field. None follows the class, True/False + expose/hide it, a string exposes it with that description, and a + dict of ``description``, ``read_only`` and ``schema`` (JSON Schema + keywords enforced on every agent write) sets options. Agents may + write a ``Sync`` field and only read a ``ClientOnly`` one; a + ``ServerOnly`` field must state ``read_only``. """ + _MCP_NEEDS_READ_ONLY = True # a ServerOnly field has no access to default to + def __init__( self, _type, @@ -690,7 +858,9 @@ def __init__( has_dataclass: bool = False, client_deep_reactive: bool = False, type_checking: TypeValidation = TypeValidation.WARNING, + mcp=None, ): + self._mcp = _field_mcp_option(mcp, self._MCP_NEEDS_READ_ONLY) self._type_checker = TypeChecker(_type, type_checking) self._client_deep_reactive = client_deep_reactive self._default = default @@ -747,6 +917,8 @@ class Sync(ServerOnly): template needs to observe at a property level (e.g. a reactive list of objects). """ + _MCP_NEEDS_READ_ONLY = False + def __set_name__(self, owner, name): _setup_class_fields(owner) @@ -778,6 +950,8 @@ class ClientOnly(ServerOnly): Inherits all constructor parameters from :class:`ServerOnly`. """ + _MCP_NEEDS_READ_ONLY = False + def __set_name__(self, owner, name): _setup_class_fields(owner) diff --git a/tests/test_mcp.py b/tests/test_mcp.py new file mode 100644 index 0000000..fbf0602 --- /dev/null +++ b/tests/test_mcp.py @@ -0,0 +1,1069 @@ +"""Models as agents (MCP) see them, through trame-server's object tools.""" + +import asyncio +import gc +import itertools +import json +import logging +from typing import Annotated + +import pytest +from trame_common.decorators import mcp +from trame_common.obj.mcp import Constructor + +from trame.app import TrameComponent, get_server +from trame_dataclass.v2 import ( + ClientOnly, + ServerOnly, + StateDataModel, + Sync, + watch, +) + +_counter = itertools.count() + + +@pytest.fixture +def server(): + server = get_server(f"test_dataclass_mcp_{next(_counter)}") + server.mcp.enabled = True + return server + + +def rpc(rpc_id, name, arguments): + return { + "jsonrpc": "2.0", + "id": rpc_id, + "method": "tools/call", + "params": {"name": name, "arguments": arguments}, + } + + +async def call(server, name, **arguments): + """``(value, is_error)`` of a tools/call, as an agent reads it.""" + response = await server.mcp.protocol.handle(rpc(1, name, arguments)) + result = response["result"] + text = result["content"][0]["text"] if result["content"] else None + try: + value = json.loads(text) if text is not None else None + except json.JSONDecodeError: + value = text + return value, bool(result.get("isError")) + + +async def ok(server, name, **arguments): + value, error = await call(server, name, **arguments) + assert not error, value + return value + + +def read(reply): + """The object a reply's result refers to, as its entry plus its id.""" + obj_id = reply["result"]["id"] + return {"id": obj_id, **reply["objects"][obj_id]} + + +async def fails(server, name, **arguments): + value, error = await call(server, name, **arguments) + assert error, value + return value + + +# ----------------------------------------------------------------------------- +# Models (agent-facing type names are process-wide: keep them unique) +# ----------------------------------------------------------------------------- + + +class Label(StateDataModel, mcp={"all_fields": True, "constructor": "default"}): + """A text label in the scene. It follows the camera.""" + + text = Sync(str, "") + position = Sync( + tuple[float, float, float], + (0.0, 0.0, 0.0), + mcp="[x, y, z] in world coordinates", + ) + hint = ClientOnly(str, "") + cache = ServerOnly(dict, dict) + + +class Layer(StateDataModel): + opacity = Sync(float, 1.0, mcp={"schema": {"minimum": 0, "maximum": 1}}) + path = ServerOnly( + str, mcp={"read_only": True, "description": "File it was loaded from"} + ) + cache = ServerOnly(dict) + + @mcp.property + def point_count(self) -> int: + return 42 + + +class Scene(StateDataModel): + """Everything the viewers show.""" + + layers = Sync(list[Layer], list, has_dataclass=True, mcp={"read_only": True}) + labels = Sync(list[Label], list, has_dataclass=True, mcp=True) + title = Sync(str, "") # unmarked, and no keyword: hidden + focus = Sync(Label, None, has_dataclass=True) # unmarked reference: hidden + + @mcp.method(requires="files") + def load_file(self, path: str) -> Layer: + """Load a file as a new layer.""" + layer = Layer(self.server, path=path) + self.layers = [*self.layers, layer] + return layer + + @mcp.method + def add_label(self, label: Label) -> None: + """Show a label.""" + if label in self.labels: + msg = f"{label.text} is already shown" + raise ValueError(msg) + self.labels = [*self.labels, label] + + @mcp.method(read_only=True) + def count(self) -> int: + return len(self.labels) + + def unmarked(self): ... + + +class Note(Label): + """A label with a body.""" + + body = Sync(str, "") + + +class Hidden(Label, mcp=False): + pass + + +class Station(StateDataModel): + name = Sync(str, "", mcp=True) + + @mcp.constructor + @classmethod + def create(cls, server, name: str) -> "Station": + """A station named so.""" + return cls(server, name=f"station {name}") + + +class Substation(Station): + """Inherits the constructor, called with its own class.""" + + +class Closed(Station, mcp={"constructor": False}): + pass + + +class Archive(Station): + @mcp.constructor(requires="files") + @classmethod + def load(cls, server, path: str) -> "Archive": + """An archive read from a file.""" + return cls(server, name=path) + + +class Picky(StateDataModel, mcp={"all_fields": True, "constructor": "default"}): + value = Sync(int, 0) + + def __init__(self, server, value): + super().__init__(server, value=value) + + +# A graph with cycles: plugs point back at their box, wires at their plug +class Wire(StateDataModel, mcp=True): + source = Sync("Plug", None, has_dataclass=True) + + +class Plug(StateDataModel, mcp=True): + owner = Sync("Box", None, has_dataclass=True) + wires = Sync(list[Wire], list, has_dataclass=True) + note = Sync(str, "") + + +class Box(StateDataModel, mcp=True): + label = Sync(str, "") + plugs = Sync(list[Plug], list, has_dataclass=True) + + +class Dial(StateDataModel): + """Marked only by a field's mcp=True.""" + + value = Sync(int, 0, mcp=True) + hidden = Sync(int, 0) + + +class Thermostat(StateDataModel, mcp=True): + target = Sync(int, 20) + heating = ClientOnly(bool, False) + + @watch("target") + def _on_target(self, target): + self.heating = target > 20 + + @mcp.method + def warm_up(self) -> int: + self.target += 2 + return self.target + + +class Hub(TrameComponent): + """A component holding models.""" + + def __init__(self, server, scene): + super().__init__(server) + self._scene = scene + + @mcp.property + def scene(self) -> Scene: + return self._scene + + +# ----------------------------------------------------------------------------- +# Declarations +# ----------------------------------------------------------------------------- + + +def test_fields_as_agents_see_them(): + fields = Label._mcp_fields() + assert fields["text"].options == {"read_only": False} # Sync + assert fields["hint"].options == {"read_only": True} # ClientOnly + assert fields["cache"].options is None # ServerOnly: only when marked + assert fields["position"].options == { + "description": "[x, y, z] in world coordinates", + "read_only": False, + } + scene = Scene._mcp_fields() + assert scene["labels"].holds_objects is True + assert scene["focus"].nullable is True + assert scene["focus"].options is None # unmarked + assert Label._mcp_constructor() == Constructor(attr=None) + assert "missing a required argument: 'value'" in Picky._mcp_constructor().problem + + +def test_a_nullable_field_keeps_its_description(server): + class Memo(StateDataModel, mcp=True): + text = Sync(Annotated[str, "What the memo says"], None) + + text = server.mcp.describe(Memo)["properties"]["text"] + assert text["description"] == "What the memo says" + assert text["anyOf"] == [{"type": "string"}, {"type": "null"}] + + +def test_fields_come_in_definition_order(): + class Ordered(StateDataModel): + zeta = Sync(int, 0) + alpha = ServerOnly(int, 0) + mid = ClientOnly(int, 0) + + class More(Ordered): + beta = Sync(int, 0) + zeta = Sync(int, 1) # redeclared: keeps its place + + assert list(More._mcp_fields()) == ["zeta", "alpha", "mid", "beta"] + + +def test_redeclared_fields_keep_their_options(): + class Shape(StateDataModel, mcp=True): + shape = ClientOnly(str, "", mcp="Object shape") + items = Sync(list[str], list, mcp={"read_only": True}) + secret = Sync(str, "", mcp=False) + + class Cone(Shape): + shape = ClientOnly(str, "cone") + items = Sync(list[str], lambda: ["a"]) + secret = Sync(str, "s") + + fields = Cone._mcp_fields() + assert fields["shape"].options == {"description": "Object shape", "read_only": True} + assert fields["items"].options == {"read_only": True} + assert fields["secret"].options is False + + +def test_field_options_are_checked(): + with pytest.raises(ValueError, match="must state read_only"): + ServerOnly(str, mcp=True) + with pytest.raises(ValueError, match="must state read_only"): + ServerOnly(str, mcp="described") + with pytest.raises(ValueError, match="unknown options"): + Sync(str, mcp={"hidden": True}) + with pytest.raises(TypeError): + Sync(str, mcp=3) + assert Sync(str, mcp=False)._mcp is False + assert ServerOnly(str, mcp={"read_only": False})._mcp == {"read_only": False} + + +def test_class_options_are_checked(): + with pytest.raises(ValueError, match="has no property 'nope'"): + + class Typo(StateDataModel, mcp={"properties": ["nope"]}): + a = Sync(int, 0) + + with pytest.raises(ValueError, match="ServerOnly field: listing it requires"): + + class Unstated(StateDataModel, mcp={"properties": ["secret"]}): + secret = ServerOnly(int, 0) + + with pytest.raises(ValueError, match="not a classmethod"): + + class NoCtor(StateDataModel, mcp={"constructor": "build"}): + a = Sync(int, 0) + + def build(self): ... + + +def test_mcp_describe_returns_the_options(): + options = Label.mcp_describe() + assert options["declared"] is True + assert options["all_fields"] is True + assert options["constructor"] == {"name": "default"} + assert options["description"] is None # the docstring is the server's default + assert Scene.mcp_describe()["declared"] is None # visible through its markers + assert Scene.mcp_describe()["constructor"] is None + assert Hidden.mcp_describe()["declared"] is False + marked = {"name": "create", "description": None, "requires": ()} + assert Station.mcp_describe()["constructor"] == marked + assert Substation.mcp_describe()["constructor"] == marked + assert Closed.mcp_describe()["constructor"] is False + + +def test_a_dict_keyword_exposes_what_is_marked(server): + class Marked(StateDataModel, mcp={"description": "Only what is marked."}): + shown = Sync(int, 0, mcp=True) + unmarked = Sync(int, 0) + + class Every(StateDataModel, mcp=True): + a = Sync(int, 0) + b = ClientOnly(int, 0) + + assert set(server.mcp.describe(Marked)["properties"]) == {"shown"} + assert set(server.mcp.describe(Every)["properties"]) == {"a", "b"} + + +def test_constructor_marker(): + class Made(StateDataModel): + a = Sync(int, 0, mcp=True) + + @mcp.constructor(description="A new one", requires="files") + @classmethod + def new(cls, server): ... + + class Other(StateDataModel): + a = Sync(int, 0, mcp=True) + + @classmethod + @mcp.constructor + def other(cls, server): ... + + assert Made.new.__func__._trame_mcp_constructor == { + "description": "A new one", + "requires": ("files",), + } + assert Other.mcp_describe()["constructor"] == { + "name": "other", + "description": None, + "requires": (), + } + assert Made._mcp_constructor() == Constructor( + attr="new", description="A new one", requires=("files",) + ) + + +def test_model_options_are_the_models_own(): + class OnlyAConstructor(StateDataModel, mcp={"constructor": "default"}): + a = Sync(int, 0) + + options = OnlyAConstructor.mcp_describe() + assert options["declared"] is True # options of its own still declare it + assert options["constructor"] == {"name": "default"} + assert options["all_fields"] is False + + with pytest.raises(ValueError, match=r"unknown options \['fields'\].*all_fields"): + + class Typo(StateDataModel, mcp={"fields": ["a"]}): + a = Sync(int, 0) + + class Described(StateDataModel): + a = Sync(int, 0) + + assert Described.mcp_describe(all_fields=True)["all_fields"] is True + with pytest.raises(ValueError, match="already set"): + Described.mcp_describe(constructor="default") + + +def test_constructor_markers_are_checked(): + with pytest.raises(ValueError, match="marks several constructors"): + + class Twice(StateDataModel): + a = Sync(int, 0, mcp=True) + + @mcp.constructor + @classmethod + def one(cls, server): ... + + @mcp.constructor + @classmethod + def two(cls, server): ... + + with pytest.raises(TypeError, match="marks a classmethod"): + + class Plain(StateDataModel): + a = Sync(int, 0, mcp=True) + + @mcp.constructor + def build(self): ... + + with pytest.raises(TypeError, match="not a staticmethod"): + mcp.constructor(staticmethod(lambda: None)) + + +def test_first_class_markers_are_for_components(): + with pytest.raises(TypeError, match="for components built once"): + + class Misplaced(StateDataModel): + value = Sync(int, 0) + + @mcp.tool + def run(self): ... + + +def test_ids_are_shared_with_components(server): + label = Label(server) + hub = Hub(server, None) + ids = {label._id, hub._id, Label(server)._id} + assert len(ids) == 3 + assert not any(i[0].isalpha() for i in ids) + + +# ----------------------------------------------------------------------------- +# What agents see +# ----------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_a_field_marked_true_makes_its_class_visible(server): + dial = Dial(server) + dial.mcp_expose("dial") + assert server.mcp.describe(Dial)["properties"] == {"value": {"type": "integer"}} + assert await ok(server, "object_update", target="dial", values={"value": 3}) + assert dial.value == 3 + + +@pytest.mark.asyncio +async def test_a_child_server_translates_the_name(server): + child = server.create_child_server(prefix="left_") + dial = Dial(child) + dial.mcp_expose("dial") + assert await ok(server, "object_update", target="left_dial", values={"value": 3}) + assert dial.value == 3 + + +@pytest.mark.asyncio +async def test_fields_by_mode_and_marks(server): + server.mcp.permissions.declare("files", "Read files") + scene = Scene(server) + scene.mcp_expose("scene") + described = await ok(server, "object_types", type_names=["Label", "Layer", "Scene"]) + types = described["types"] + + assert types["Label"]["properties"] == { + "hint": {"type": "string", "read_only": True}, + "position": { + "type": "array", + "minItems": 3, + "maxItems": 3, + "items": {"type": "number"}, + "description": "[x, y, z] in world coordinates", + }, + "text": {"type": "string"}, + } + assert types["Label"]["constructor"]["parameters"]["properties"].keys() == { + "position", + "text", + } + assert types["Layer"]["properties"]["opacity"] == { + "type": "number", + "minimum": 0, + "maximum": 1, + } + assert types["Layer"]["properties"]["path"]["read_only"] is True + assert types["Layer"]["properties"]["point_count"] == { + "type": "integer", + "read_only": True, + } + assert "constructor" not in types["Layer"] + assert set(types["Scene"]["properties"]) == {"labels", "layers"} + assert set(types["Scene"]["methods"]) == {"add_label", "count", "load_file"} + assert types["Scene"]["methods"]["load_file"]["requires"] == ["files"] + assert types["Scene"]["properties"]["labels"]["items"]["object"] == "Label" + assert types["Scene"]["properties"]["layers"]["read_only"] is True + + +@pytest.mark.asyncio +async def test_index_and_subclasses(server): + scene = Scene(server) + scene.mcp_expose("scene") + assert await ok(server, "object_list") == {"objects": {"scene": {"type": "Scene"}}} + index = await ok(server, "object_types") + # Note extends Label (a reference typed with Label may hold one); Hidden never shows + assert index["types"]["Note"] == { + "about": "A label with a body.", + "properties": ["body"], + "constructor": True, + "extends": "Label", + } + assert index["types"]["Label"]["about"] == "A text label in the scene." + assert "Hidden" not in index["types"] + assert "Layer" in index["types"] + + +@pytest.mark.asyncio +async def test_library_models_with_mcp_describe(server): + class Camera(StateDataModel): + position = Sync(list[float], lambda: [0.0, 0.0, 1.0]) + cache_key = Sync(str, "") + + @property + def distance(self) -> float: + return 1.0 + + def reset(self) -> None: + """Back to the default position.""" + self.position = [0.0, 0.0, 1.0] + + Camera.mcp_describe( + description="The camera of the 3D view.", + properties={"position": "[x, y, z]", "cache_key": False, "distance": True}, + methods=["reset"], + ) + with pytest.raises(ValueError, match="already set"): + Camera.mcp_describe(False) + + camera = Camera(server, position=[1.0, 2.0, 3.0]) + camera.mcp_expose("camera") + described = server.mcp.describe(Camera) + assert described["description"] == "The camera of the 3D view." + assert set(described["properties"]) == {"position", "distance"} + assert set(described["methods"]) == {"reset"} + await ok(server, "object_call", target="camera", method="reset") + assert camera.position == [0.0, 0.0, 1.0] + + +# ----------------------------------------------------------------------------- +# Reachability +# ----------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_reachability(server): + scene = Scene(server) + label, loose = Label(server, text="a"), Label(server, text="loose") + scene.labels = [label] + scene.focus = loose # a hidden reference is no path + scene.mcp_expose("scene") + + value = await ok(server, "object_get", target="scene") + assert value == { + "result": {"id": scene._id}, + "objects": { + scene._id: { + "type": "Scene", + "names": ["scene"], + "values": { + "labels": [{"id": label._id, "type": "Label"}], + "layers": [], + }, + } + }, + } + value = await ok(server, "object_get", target=label._id) + assert read(value)["values"]["text"] == "a" + assert "Unknown object" in await fails(server, "object_get", target=loose._id) + + value = await ok(server, "object_get", target="scene", depth=2) + assert read(value)["values"]["labels"] == [{"id": label._id}] + assert value["objects"][label._id]["values"]["text"] == "a" + + scene.mcp_unexpose() + assert "Unknown object" in await fails(server, "object_get", target=label._id) + + +@pytest.mark.asyncio +async def test_component_properties_reach_models(server): + scene = Scene(server) + label = Label(server, text="through a component") + scene.labels = [label] + hub = Hub(server, scene) + hub.mcp_expose("hub") + value = read(await ok(server, "object_get", target=hub._id)) + assert value["values"]["scene"] == {"id": scene._id, "type": "Scene"} + value = read(await ok(server, "object_get", target=label._id)) + assert value["values"]["text"] == "through a component" + + +@pytest.mark.asyncio +async def test_read_only_names(server): + scene = Scene(server) + label = Label(server, text="a") + scene.labels = [label] + scene.mcp_expose("scene", read_only=True) + + value = read(await ok(server, "object_get", target=label._id)) + assert value["read_only"] is True + message = await fails( + server, "object_update", target=label._id, values={"text": "b"} + ) + assert message == f"Nothing was applied: Label {label._id} is exposed read-only." + assert await ok(server, "object_call", target="scene", method="count") == 1 + assert "read-only" in await fails( + server, + "object_call", + target="scene", + method="add_label", + args={"label": label._id}, + ) + + # A writable name wins + label.mcp_expose("label") + await ok(server, "object_update", target=label._id, values={"text": "b"}) + assert label.text == "b" + + +@pytest.mark.asyncio +async def test_a_model_name_has_a_description(server): + Label(server).mcp_expose("title", description="The plot's title.") + listed = (await ok(server, "object_list"))["objects"] + assert listed["title"] == {"type": "Label", "description": "The plot's title."} + + +def test_models_need_a_server(): + with pytest.raises(RuntimeError, match="has no trame server"): + Label().mcp_expose("label") + + +def test_hidden_models_cannot_be_exposed(server): + with pytest.raises(TypeError, match="declares nothing"): + Hidden(server).mcp_expose("hidden") + + +# ----------------------------------------------------------------------------- +# Writes and calls +# ----------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_update(server): + label = Label(server) + label.mcp_expose("label") + assert await ok( + server, "object_update", target="label", values={"position": [0, 2, 5]} + ) == {"id": label._id, "updated": ["position"]} + assert label.position == (0.0, 2.0, 5.0) # the declared tuple + + message = await fails( + server, "object_update", target="label", values={"text": "ok", "position": [1]} + ) + assert message == "Nothing was applied: position: expected at least 3 items, got 1" + assert label.text == "" + message = await fails(server, "object_update", target="label", values={"hint": "x"}) + assert "hint (read-only)" in message + + +@pytest.mark.asyncio +async def test_reference_fields_are_written_by_id(server): + scene = Scene(server) + scene.mcp_expose("scene") + first, second = Label(server, text="1"), Label(server, text="2") + first.mcp_expose("first") + second.mcp_expose("second") + layer = Layer(server) + scene.layers = [layer] + + await ok(server, "object_update", target="scene", values={"labels": [second._id]}) + assert scene.labels == [second] + message = await fails( + server, "object_update", target="scene", values={"labels": ["first", "x"]} + ) + assert 'labels[1]: "x" is not the id or name of an object' in message + assert scene.labels == [second] # all or nothing + + message = await fails( + server, "object_update", target="scene", values={"labels": [layer._id]} + ) + assert f'labels[0]: "{layer._id}" has type Layer, not Label' in message + message = await fails( + server, "object_update", target="scene", values={"layers": []} + ) + assert "layers (read-only)" in message + + +@pytest.mark.asyncio +async def test_writes_return_once_watchers_ran(server): + thermostat = Thermostat(server) + thermostat.mcp_expose("thermostat") + await ok(server, "object_update", target="thermostat", values={"target": 22}) + assert thermostat.heating is True + thermostat.target = 18 + await thermostat.completion() + assert await ok(server, "object_call", target="thermostat", method="warm_up") == 20 + assert await ok(server, "object_call", target="thermostat", method="warm_up") == 22 + assert thermostat.heating is True + + +@pytest.mark.asyncio +async def test_a_failing_watcher_does_not_fail_the_update(server, caplog): + class Fragile(StateDataModel, mcp=True): + level = Sync(int, 0) + + @watch("level") + def _on_level(self, level): + if level > 5: + msg = "too high" + raise RuntimeError(msg) + + fragile = Fragile(server) + fragile.mcp_expose("fragile") + with caplog.at_level(logging.ERROR, logger="trame_server.mcp"): + result = await ok( + server, "object_update", target="fragile", values={"level": 9} + ) + assert result["updated"] == ["level"] + assert fragile.level == 9 + assert "A watcher of Fragile" in caplog.text + + +@pytest.mark.asyncio +async def test_calls_and_permissions(server): + server.mcp.permissions.declare("files", "Read files on this computer") + scene = Scene(server) + scene.mcp_expose("scene") + value = await ok( + server, + "object_call", + target="scene", + method="load_file", + args={"path": "/a.vti"}, + ) + layer = scene.layers[0] + assert value == { + "result": {"id": layer._id}, + "objects": { + layer._id: { + "type": "Layer", + "values": {"opacity": 1.0, "path": "/a.vti", "point_count": 42}, + } + }, + } + + server.mcp.permissions["files"] = False + described = await ok(server, "object_types", type_names=["Scene"]) + assert "load_file" not in described["types"]["Scene"]["methods"] + message = await fails( + server, "object_call", target="scene", method="load_file", args={"path": "/b"} + ) + assert 'permission "files"' in message + assert len(scene.layers) == 1 + + +# ----------------------------------------------------------------------------- +# Constructors and batches +# ----------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_constructors(server, caplog): + things = [Station(server), Substation(server), Closed(server), Picky(server, 1)] + for i, thing in enumerate(things): + thing.mcp_expose(f"thing_{i}") + described = await ok(server, "object_types", type_names=["Station"]) + assert described["types"]["Station"]["constructor"] == { + "parameters": { + "type": "object", + "properties": {"name": {"type": "string"}}, + "required": ["name"], + "additionalProperties": False, + }, + "description": "A station named so.", + } + with caplog.at_level(logging.WARNING, logger="trame_server.mcp"): + described = await ok(server, "object_types", type_names=["Picky"]) + assert "constructor" not in described["types"]["Picky"] + assert "Picky has no constructor for agents" in caplog.text + + value = await ok( + server, + "object_batch", + steps=[ + {"new": "Station", "args": {"name": "x"}}, + {"new": "Substation", "args": {"name": "y"}, "as": "sub"}, + ], + results="full", + ) + station, substation = (value["objects"][r["id"]] for r in value["results"]) + assert station["values"] == {"name": "station x"} + assert substation["type"] == "Substation" + assert len(value["unattached"]) == 2 + assert "Closed" in await fails(server, "object_batch", steps=[{"new": "Closed"}]) + + +@pytest.mark.asyncio +async def test_a_subclass_without_its_base_constructor(server): + server.mcp.permissions.declare("files", "Read files", allowed=False) + for i, thing in enumerate([Station(server), Archive(server), Closed(server)]): + thing.mcp_expose(f"thing_{i}") + described = await ok(server, "object_types", type_names=["Archive", "Closed"]) + # Still described as what they change: their constructor is gone + assert described["types"] == { + "Archive": { + "extends": "Station", + "properties": {}, + "methods": {}, + "constructor": False, + }, + "Closed": { + "extends": "Station", + "properties": {}, + "methods": {}, + "constructor": False, + }, + } + index = (await ok(server, "object_types"))["types"] + assert index["Archive"] == {"constructor": False, "extends": "Station"} + + server.mcp.permissions["files"] = True + described = await ok(server, "object_types", type_names=["Archive"]) + assert described["types"]["Archive"]["constructor"]["requires"] == ["files"] + + +@pytest.mark.asyncio +async def test_batch_builds_attaches_and_links(server): + scene = Scene(server) + scene.mcp_expose("scene") + value = await ok( + server, + "object_batch", + steps=[ + { + "new": "Label", + "args": {"text": "Skull", "position": [0, 0, 120]}, + "as": "skull", + }, + {"call": "scene", "method": "add_label", "args": {"label": "@skull"}}, + {"update": "@skull", "values": {"text": "Skull!"}}, + {"new": "Note", "args": {"body": "never attached"}}, + ], + ) + skull = scene.labels[0] + assert skull.text == "Skull!" + assert skull.position == (0.0, 0.0, 120.0) + note = value["results"][3] + assert value == { + "results": [ + {"id": skull._id}, + None, + {"id": skull._id, "updated": ["text"]}, + {"id": note["id"]}, + ], + "unattached": [{"id": note["id"]}], + "objects": { + skull._id: {"type": "Label"}, # links only, and none + note["id"]: {"type": "Note"}, + }, + } + + full = await ok( + server, + "object_batch", + steps=[{"call": "scene", "method": "count", "as": "n"}], + results="full", + ) + assert full["results"] == [1] + + # References follow properties + value = await ok( + server, + "object_batch", + steps=[ + {"new": "Label", "args": {"text": "source"}, "as": "a"}, + {"new": "Label", "args": {"text": "@a.text"}, "as": "b"}, + {"update": "scene", "values": {"labels": ["@a", "@b"]}}, + ], + ) + assert [label.text for label in scene.labels] == ["source", "source"] + assert "unattached" not in value + assert "labels[0]" in await fails( + server, + "object_batch", + steps=[ + {"call": "scene", "method": "count", "as": "n"}, + {"update": "scene", "values": {"labels": ["@n"]}}, + ], + ) + + +@pytest.mark.asyncio +async def test_batch_results_are_cut_last(server): + server.mcp.max_response_tokens = 500 # 2000 characters + Scene(server).mcp_expose("scene") # Labels can be built for it + steps = [{"new": "Label", "args": {"text": f"label {i}"}} for i in range(80)] + value = await ok(server, "object_batch", steps=steps) + assert len(json.dumps(value, separators=(",", ":"))) <= 2000 + # The objects go first: object_get reads them again, the batch can't rerun + assert "objects" not in value + assert value["more"]["objects"].startswith("0 of 80 objects read") + for key in ("results", "unattached"): + assert value["more"][key] == ( + f"{len(value[key])} of 80 items shown to fit the response budget; " + f"run fewer steps per batch" + ) + + +@pytest.mark.asyncio +async def test_batch_failures(server): + scene = Scene(server) + scene.mcp_expose("scene") + message = await fails( + server, + "object_batch", + steps=[ + {"new": "Label", "args": {"text": "a"}, "as": "a"}, + {"call": "scene", "method": "add_label", "args": {"label": "@a"}}, + {"call": "scene", "method": "add_label", "args": {"label": "@a"}}, + ], + ) + assert message.startswith( + "steps[2] (call add_label on scene) failed: ValueError: a is already shown" + ) + assert "Steps 0-1 stay applied" in message + assert len(scene.labels) == 1 + + server.mcp.permissions.declare("files", "Read files", allowed=False) + message = await fails( + server, + "object_batch", + steps=[ + {"new": "Label", "as": "b"}, + {"call": "scene", "method": "add_label", "args": {"label": "@b"}}, + {"call": "scene", "method": "load_file", "args": {"path": "/x"}}, + ], + ) + assert message.startswith("Nothing was applied:") + assert len(scene.labels) == 1 + message = await fails(server, "object_batch", steps=[{"new": "Layer"}]) + assert 'no type "Layer" can be built' in message + + +@pytest.mark.asyncio +async def test_a_cancel_stops_a_batch_between_steps(server): + from trame_server.mcp.protocol import REQUEST_CANCELLED + + release = asyncio.Event() + + class Waiter(StateDataModel, mcp=True): + done = Sync(int, 0) + + @mcp.method + async def wait(self) -> None: + await release.wait() + + @mcp.method + def finish(self) -> None: + self.done += 1 + + waiter = Waiter(server) + waiter.mcp_expose("waiter") + running = asyncio.ensure_future( + server.mcp.protocol.handle( + rpc( + 7, + "object_batch", + { + "steps": [ + {"call": "waiter", "method": "wait"}, + {"call": "waiter", "method": "finish"}, + ] + }, + ) + ) + ) + await asyncio.sleep(0.01) + assert server.mcp.protocol.cancel(7, reason="changed my mind") + response = await running + assert response["error"]["code"] == REQUEST_CANCELLED + assert waiter.done == 0 + + +@pytest.mark.asyncio +async def test_depth_reads_each_object_once(server): + box = Box(server, label="b") + plug = Plug(server, owner=box) + wire = Wire(server, source=plug) + box.plugs = [plug] + plug.wires = [wire] + box.mcp_expose("box") + + value = await ok(server, "object_get", target="box", depth=3) + assert value == { + "result": {"id": box._id}, + "objects": { + box._id: { + "type": "Box", + "names": ["box"], + "values": {"label": "b", "plugs": [{"id": plug._id}]}, + }, + plug._id: { + "type": "Plug", + "values": { + "owner": {"id": box._id}, + "wires": [{"id": wire._id}], + "note": "", + }, + }, + wire._id: {"type": "Wire", "values": {"source": {"id": plug._id}}}, + }, + } + value = await ok(server, "object_get", target="box", depth=2) + assert value["objects"][plug._id]["values"]["wires"] == [ + {"id": wire._id, "type": "Wire"} # not read: with its type + ] + assert list(value["objects"]) == [box._id, plug._id] + + # The objects a capped list leaves out are not read; the farthest objects + # are left out to fit the budget + box.plugs = [Plug(server, owner=box, note="n" * 100) for _ in range(60)] + value = await ok(server, "object_get", target="box", depth=2) + assert len(value["objects"]) == 21 + server.mcp.max_response_tokens = 500 + value = await ok(server, "object_get", target="box", depth=2) + assert len(json.dumps(value, separators=(",", ":"))) <= 2000 # as sent + kept = len(value["objects"]) + assert 1 < kept < 21 + assert value["more"]["objects"] == ( + f"{kept} of 21 objects read to fit the response budget; " + "object_get(id) reads one" + ) + + +@pytest.mark.asyncio +async def test_models_disappear_when_collected(server): + scene = Scene(server) + scene.labels = [Label(server, text="x")] + _id = scene.labels[0]._id + holder = [scene] + server.mcp.expose_object(lambda: holder[0], "scene") + await ok(server, "object_get", target=_id) + holder[0] = None + del scene + gc.collect() + assert "currently empty" in await fails(server, "object_get", target="scene") + assert "Unknown object" in await fails(server, "object_get", target=_id) + + +def test_audit_reports_hidden_references(server, caplog): + scene = Scene(server) + scene.mcp_expose("scene") + with caplog.at_level(logging.WARNING, logger="trame_server.mcp"): + server.mcp._objects.audit() + assert "Scene.focus holds objects agents see (Label, Note)" in caplog.text